From 2a0602c7ac12eaf05c2e4cd8ae61965bcc39153a Mon Sep 17 00:00:00 2001
From: Mick <mickjagger19@icloud.com>
Date: Fri, 4 Sep 2026 09:17:23 +0800
Subject: [PATCH] [diffusion] optimization: unfused w2 bias on SM12.x (cuBLAS
 16x16 kernel mis-dispatch) for minimax-h3 vae decoder: (#37835)

---
 .../vaes/minimax_h3_video_vae/base_module.py  | 30 ++++++++++++++++++-
 1 file changed, 29 insertions(+), 1 deletion(-)

diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py
index ec6727388..b258bb895 100644
--- a/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py
+++ b/python/sglang/multimodal_gen/runtime/models/vaes/minimax_h3_video_vae/base_module.py
@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # Transformer building blocks for the MiniMax H3 visual VAE ViT decoder.
+import functools
 import math
 from contextlib import nullcontext
 from typing import Optional
@@ -33,6 +34,33 @@ def _scaled_residual_add(residual, x, scale):
     return residual + x * scale if fused is None else fused
 
 
+@functools.lru_cache(maxsize=1)
+def _is_sm120() -> bool:
+    return bool(current_platform.is_sm120())
+
+
+def _unfused_bias_linear(
+    linear: nn.Linear, hidden_states: torch.Tensor
+) -> torch.Tensor:
+    """``linear`` as a plain matmul plus a separate bias add on SM12.x.
+
+    For the decoder's w2 shape ([~1800, 8192] x [8192, 2048], fp16/bf16) the
+    fused ``addmm`` epilogue makes cuBLAS on a GB10 pick a 16x16 wmma kernel
+    that runs at 14 TFLOPS; the same product without the fused bias runs at
+    76-91 TFLOPS (measured: 4.0 ms -> 0.8 ms per call, 3780 calls per decode).
+    """
+    if (
+        linear.bias is None
+        or not hidden_states.is_cuda
+        or hidden_states.dtype != linear.weight.dtype
+        or not _is_sm120()
+    ):
+        return linear(hidden_states)
+    out = torch.matmul(hidden_states, linear.weight.t())
+    out += linear.bias
+    return out
+
+
 class FeedForward(nn.Module):
     def __init__(
         self,
@@ -99,7 +127,7 @@ class FeedForward(nn.Module):
         else:
             hidden_states = self.act_fn(hidden_states)
 
-        hidden_states = self.w2(hidden_states)
+        hidden_states = _unfused_bias_linear(self.w2, hidden_states)
         return hidden_states
 
     def _get_forward_impl(self):