[diffusion] perf: absorb Qwen-Image output projection biases (#37116)

python/sglang/kernels/jit/csrc/diffusion/norm_scale_shift.cuh

@@ -40,6 +40,7 @@ struct NormScaleShiftParams {
   void* y;
   void* res_out;
   const void* x;
+  const void* input_bias;
   const void* residual;
   const void* gate;
   const void* scale;
@@ -65,7 +66,7 @@ SGL_DEVICE float cta_reduce_sum(float v, int warp, int lane, float* scratch) {
   return scratch[kWarps];
 }
 
-template <bool kHasResidual>
+template <bool kHasResidual, bool kHasInputBias = false>
 __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) {
   using namespace device;
   using Vec = AlignedVector<bf16_t, kVecElems>;
@@ -89,6 +90,16 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta
     v[i] = static_cast<float>(xv[i]);
   }
 
+  if constexpr (kHasInputBias) {
+    Vec bv;
+    bv.load(static_cast<const bf16_t*>(params.input_bias) + elem_offset);
+#pragma unroll
+    for (int i = 0; i < kVecElems; ++i) {
+      // Match the standalone BF16 output-projection bias addition.
+      v[i] = static_cast<float>(static_cast<bf16_t>(v[i] + static_cast<float>(bv[i])));
+    }
+  }
+
   if constexpr (kHasResidual) {
     Vec gv;
     Vec rv;
@@ -135,6 +146,31 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta
   yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
 }
 
+__global__ void bias_mul_add_kernel(const NormScaleShiftParams __grid_constant__ params) {
+  using namespace device;
+  using Vec = AlignedVector<bf16_t, kVecElems>;
+
+  const int row_offset = blockIdx.x * kHidden;
+  const int elem_offset = threadIdx.x * kVecElems;
+
+  Vec xv;
+  Vec bv;
+  Vec gv;
+  Vec rv;
+  Vec yv;
+  xv.load(static_cast<const bf16_t*>(params.x) + row_offset + elem_offset);
+  bv.load(static_cast<const bf16_t*>(params.input_bias) + elem_offset);
+  gv.load(static_cast<const bf16_t*>(params.gate) + elem_offset);
+  rv.load(static_cast<const bf16_t*>(params.residual) + row_offset + elem_offset);
+
+#pragma unroll
+  for (int i = 0; i < kVecElems; ++i) {
+    const bf16_t biased = static_cast<bf16_t>(static_cast<float>(xv[i]) + static_cast<float>(bv[i]));
+    yv[i] = __hfma(biased, gv[i], rv[i]);
+  }
+  yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
+}
+
 inline uint32_t verify_nss_geometry(host::SymbolicSize& num_rows) {
   using namespace host;
   RuntimeCheck(num_rows.unwrap() > 0, "num_rows must be positive");
@@ -162,6 +198,7 @@ struct NormScaleShiftKernel {
         .y = y.data_ptr(),
         .res_out = nullptr,
         .x = x.data_ptr(),
+        .input_bias = nullptr,
         .residual = nullptr,
         .gate = nullptr,
         .scale = scale.data_ptr(),
@@ -201,6 +238,7 @@ struct ScaleResidualNormScaleShiftKernel {
         .y = y.data_ptr(),
         .res_out = res_out.data_ptr(),
         .x = x.data_ptr(),
+        .input_bias = nullptr,
         .residual = residual.data_ptr(),
         .gate = gate.data_ptr(),
         .scale = scale.data_ptr(),
@@ -211,6 +249,84 @@ struct ScaleResidualNormScaleShiftKernel {
   }
 };
 
+struct BiasScaleResidualNormScaleShiftKernel {
+  static void
+  run(tvm::ffi::TensorView y,
+      tvm::ffi::TensorView res_out,
+      tvm::ffi::TensorView residual,
+      tvm::ffi::TensorView x,
+      tvm::ffi::TensorView input_bias,
+      tvm::ffi::TensorView gate,
+      tvm::ffi::TensorView scale,
+      tvm::ffi::TensorView shift,
+      double eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_rows"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, kHidden})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(x)
+        .verify(residual)
+        .verify(y)
+        .verify(res_out);
+    TensorMatcher({kHidden})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(input_bias)
+        .verify(gate)
+        .verify(scale)
+        .verify(shift);
+
+    const uint32_t grid = verify_nss_geometry(N);
+    const auto params = NormScaleShiftParams{
+        .y = y.data_ptr(),
+        .res_out = res_out.data_ptr(),
+        .x = x.data_ptr(),
+        .input_bias = input_bias.data_ptr(),
+        .residual = residual.data_ptr(),
+        .gate = gate.data_ptr(),
+        .scale = scale.data_ptr(),
+        .shift = shift.data_ptr(),
+        .eps = static_cast<float>(eps),
+    };
+    LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true, true>, params);
+  }
+};
+
+struct BiasMulAddKernel {
+  static void
+  run(tvm::ffi::TensorView y,
+      tvm::ffi::TensorView x,
+      tvm::ffi::TensorView input_bias,
+      tvm::ffi::TensorView gate,
+      tvm::ffi::TensorView residual) {
+    using namespace host;
+    auto N = SymbolicSize{"num_rows"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, kHidden}).with_dtype<bf16_t>().with_device(device).verify(x).verify(residual).verify(y);
+    TensorMatcher({kHidden}).with_dtype<bf16_t>().with_device(device).verify(input_bias).verify(gate);
+
+    const uint32_t grid = verify_nss_geometry(N);
+    const auto params = NormScaleShiftParams{
+        .y = y.data_ptr(),
+        .res_out = nullptr,
+        .x = x.data_ptr(),
+        .input_bias = input_bias.data_ptr(),
+        .residual = residual.data_ptr(),
+        .gate = gate.data_ptr(),
+        .scale = nullptr,
+        .shift = nullptr,
+        .eps = 0.0f,
+    };
+    LaunchKernel(grid, kThreads, device.unwrap())(bias_mul_add_kernel, params);
+  }
+};
+
 }  // namespace norm_scale_shift
 
 }  // namespace sglang

python/sglang/kernels/ops/diffusion/__init__.py

@@ -373,6 +373,8 @@ _EXPORTS: dict[str, str] = {
     "rmsnorm_tanh_residual": "norm.native_bf16_rmsnorm_triton",
     "norm_infer": "norm.norm_triton",
     "rms_norm_fn": "norm.norm_triton",
+    "try_fused_bias_mul_add": "norm.norm_scale_shift_jit",
+    "try_fused_bias_scale_residual_norm_scale_shift": "norm.norm_scale_shift_jit",
     "triton_one_pass_rms_norm": "norm.rmsnorm_onepass_triton",
     "can_use_fused_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact",
     "can_use_fused_scale_residual_rmsnorm_scale_shift": "norm.rmsnorm_scale_shift_bitexact",

python/sglang/kernels/ops/diffusion/norm/norm_scale_shift_jit.py

@@ -24,6 +24,13 @@ def _blackwell_or_newer(device: torch.device) -> bool:
     )
 
 
+def _sm103(device: torch.device) -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability(device) == (
+        10,
+        3,
+    )
+
+
 def _nss_activation(t, like=None) -> bool:
     return (
         isinstance(t, torch.Tensor)
@@ -72,6 +79,14 @@ def norm_scale_shift_module() -> Module:
                 "srnss_bf16_row",
                 "norm_scale_shift::ScaleResidualNormScaleShiftKernel::run",
             ),
+            (
+                "bias_srnss_bf16_row",
+                "norm_scale_shift::BiasScaleResidualNormScaleShiftKernel::run",
+            ),
+            (
+                "bias_mul_add_bf16_row",
+                "norm_scale_shift::BiasMulAddKernel::run",
+            ),
         ],
     )
 
@@ -128,3 +143,58 @@ def try_fused_scale_residual_norm_scale_shift(
         float(eps),
     )
     return y, residual_out
+
+
+def try_fused_bias_scale_residual_norm_scale_shift(
+    residual, x, input_bias, gate, weight, bias, scale, shift, norm_type, eps
+):
+    if torch.compiler.is_compiling():
+        return None
+    if norm_type != "layer" or weight is not None or bias is not None:
+        return None
+    if not (_nss_activation(x) and _nss_activation(residual, x) and _sm103(x.device)):
+        return None
+
+    input_bias = _row_bf16(input_bias, x.device)
+    gate = _row_bf16(gate, x.device)
+    scale = _row_bf16(scale, x.device)
+    shift = _row_bf16(shift, x.device)
+    if any(tensor is None for tensor in (input_bias, gate, scale, shift)):
+        return None
+
+    y = torch.empty_like(x)
+    residual_out = torch.empty_like(x)
+    _module().bias_srnss_bf16_row(
+        y.view(-1, _HIDDEN),
+        residual_out.view(-1, _HIDDEN),
+        residual.view(-1, _HIDDEN),
+        x.view(-1, _HIDDEN),
+        input_bias,
+        gate,
+        scale,
+        shift,
+        float(eps),
+    )
+    return y, residual_out
+
+
+def try_fused_bias_mul_add(x, input_bias, gate, residual):
+    if torch.compiler.is_compiling():
+        return None
+    if not (_nss_activation(x) and _nss_activation(residual, x) and _sm103(x.device)):
+        return None
+
+    input_bias = _row_bf16(input_bias, x.device)
+    gate = _row_bf16(gate, x.device)
+    if input_bias is None or gate is None:
+        return None
+
+    y = torch.empty_like(x)
+    _module().bias_mul_add_bf16_row(
+        y.view(-1, _HIDDEN),
+        x.view(-1, _HIDDEN),
+        input_bias,
+        gate,
+        residual.view(-1, _HIDDEN),
+    )
+    return y

python/sglang/multimodal_gen/runtime/models/dits/qwen_image.py

@@ -20,6 +20,8 @@ from sglang.kernels.ops.diffusion import (
     fused_gelu_active,
     fused_linear_gelu_tanh,
     mark_fused_gelu_site,
+    try_fused_bias_mul_add,
+    try_fused_bias_scale_residual_norm_scale_shift,
 )
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
 from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block
@@ -76,7 +78,10 @@ from sglang.multimodal_gen.runtime.managers.memory_managers.layerwise_offload im
 )
 from sglang.multimodal_gen.runtime.models.dits.base import CachableDiT
 from sglang.multimodal_gen.runtime.models.dits.common import get_qkv_projections
-from sglang.multimodal_gen.runtime.platforms import AttentionBackendEnum
+from sglang.multimodal_gen.runtime.platforms import (
+    AttentionBackendEnum,
+    current_platform,
+)
 from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
     is_in_breakable_cuda_graph,
@@ -110,6 +115,27 @@ def _local_seq_len(seq_len: int, sp_world_size: int) -> int:
 _get_qkv_projections = get_qkv_projections
 
 
+def _can_defer_modelopt_output_bias(
+    quant_config: Optional[QuantizationConfig], capability: Any
+) -> bool:
+    # Absorbing the bias moves a BF16 rounding point. The resulting image
+    # quality has only been validated on SM103, so other GPUs keep the GEMM
+    # bias epilogue used before this optimization.
+    return (
+        quant_config is not None
+        and hasattr(quant_config, "get_name")
+        and quant_config.get_name() in {"modelopt_fp8", "modelopt_fp4"}
+        and capability is not None
+        and (capability.major, capability.minor) == (10, 3)
+    )
+
+
+def _defer_modelopt_output_bias(quant_config: Optional[QuantizationConfig]) -> bool:
+    return _can_defer_modelopt_output_bias(
+        quant_config, current_platform.get_device_capability()
+    )
+
+
 def _safe_tensor_version(tensor: torch.Tensor) -> Optional[int]:
     """Read a tensor version counter without rejecting inference tensors."""
     return None if tensor.is_inference() else tensor._version
@@ -579,6 +605,7 @@ class QwenImageCrossAttention(nn.Module):
         self.parallel_attention = parallel_attention
         self.added_kv_proj_dim = added_kv_proj_dim
         self.prefix = prefix
+        self.defer_output_bias = _defer_modelopt_output_bias(quant_config)
 
         self.use_fused_qkv = isinstance(quant_config, NunchakuConfig)
 
@@ -672,6 +699,7 @@ class QwenImageCrossAttention(nn.Module):
                 self.dim,
                 bias=out_bias,
                 input_is_parallel=True,
+                skip_bias_add=self.defer_output_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.to_add_out",
             )
@@ -686,6 +714,7 @@ class QwenImageCrossAttention(nn.Module):
                     self.dim,
                     bias=out_bias,
                     input_is_parallel=True,
+                    skip_bias_add=self.defer_output_bias,
                     quant_config=quant_config,
                     prefix=f"{prefix}.to_out.0",
                 )
@@ -857,13 +886,13 @@ class QwenImageCrossAttention(nn.Module):
         )
 
         # Apply output projections
-        img_attn_output, _ = self.to_out[0](img_attn_output)
+        img_attn_output, img_attn_bias = self.to_out[0](img_attn_output)
         if len(self.to_out) > 1:
             (img_attn_output,) = self.to_out[1](img_attn_output)  # dropout
 
-        txt_attn_output, _ = self.to_add_out(txt_attn_output)
+        txt_attn_output, txt_attn_bias = self.to_add_out(txt_attn_output)
 
-        return img_attn_output, txt_attn_output
+        return img_attn_output, txt_attn_output, img_attn_bias, txt_attn_bias
 
 
 class QwenImageGELU(nn.Module):
@@ -918,6 +947,7 @@ class QwenImageFeedForward(nn.Module):
     ) -> None:
         super().__init__()
         inner_dim = dim * mult
+        self.defer_output_bias = _defer_modelopt_output_bias(quant_config)
         if replicated:
             # Keep the whole FFN resident on every rank: no per-block
             # all-reduce. Only worth it when the branch's token count is small
@@ -926,6 +956,7 @@ class QwenImageFeedForward(nn.Module):
                 inner_dim,
                 dim_out,
                 bias=True,
+                skip_bias_add=self.defer_output_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.net.2",
             )
@@ -935,6 +966,7 @@ class QwenImageFeedForward(nn.Module):
                 dim_out,
                 bias=True,
                 input_is_parallel=True,
+                skip_bias_add=self.defer_output_bias,
                 quant_config=quant_config,
                 prefix=f"{prefix}.net.2",
             )
@@ -952,11 +984,16 @@ class QwenImageFeedForward(nn.Module):
             ]
         )
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward_with_bias(
+        self, hidden_states: torch.Tensor
+    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         hidden_states = self.net[0](hidden_states)
         hidden_states = self.net[1](hidden_states)
-        hidden_states, _ = self.net[2](hidden_states)
-        return hidden_states
+        return self.net[2](hidden_states)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states, bias = self.forward_with_bias(hidden_states)
+        return hidden_states if bias is None else hidden_states + bias
 
 
 class QwenImageTransformerBlock(nn.Module):
@@ -1120,6 +1157,25 @@ class QwenImageTransformerBlock(nn.Module):
     ) -> torch.Tensor:
         return self.fuse_mul_add(a, b, c, k)
 
+    def _bias_mul_add(
+        self,
+        a: torch.Tensor,
+        bias: Optional[torch.Tensor],
+        b: torch.Tensor,
+        c: torch.Tensor,
+        *,
+        use_bcg_helpers: bool,
+    ) -> torch.Tensor:
+        if bias is not None and not use_bcg_helpers:
+            fused = try_fused_bias_mul_add(a, bias, b, c)
+            if fused is not None:
+                return fused
+        if bias is not None:
+            a = a + bias
+        if use_bcg_helpers:
+            return self._mul_add(a, b, c)
+        return self.fuse_mul_add(a, b, c)
+
     def _get_modulation_params(
         self,
         temb_img_silu: torch.Tensor,
@@ -1152,6 +1208,7 @@ class QwenImageTransformerBlock(nn.Module):
         index: Optional[torch.Tensor] = None,
         gate_x: Optional[torch.Tensor] = None,
         residual_x: Optional[torch.Tensor] = None,
+        x_bias: Optional[torch.Tensor] = None,
         use_bcg_helpers: bool = False,
     ) -> Union[
         Tuple[torch.Tensor, torch.Tensor],
@@ -1164,6 +1221,8 @@ class QwenImageTransformerBlock(nn.Module):
 
         shift, scale, gate = mod_params.chunk(3, dim=-1)
         if index is not None:
+            if x_bias is not None:
+                x = x + x_bias
             actual_batch = x.shape[0]
             shift0, shift1 = (
                 shift[:actual_batch],
@@ -1214,6 +1273,24 @@ class QwenImageTransformerBlock(nn.Module):
             scale_result = scale.unsqueeze(1)
             gate_result = gate.unsqueeze(1)
         if is_scale_residual:
+            if x_bias is not None and not use_bcg_helpers:
+                fused = try_fused_bias_scale_residual_norm_scale_shift(
+                    residual_x,
+                    x,
+                    x_bias,
+                    gate_x,
+                    getattr(norm_module.norm, "weight", None),
+                    getattr(norm_module.norm, "bias", None),
+                    scale_result,
+                    shift_result,
+                    norm_module.norm_type,
+                    norm_module.eps,
+                )
+                if fused is not None:
+                    modulated, residual_out = fused
+                    return modulated, residual_out, gate_result
+            if x_bias is not None:
+                x = x + x_bias
             if use_bcg_helpers:
                 modulated, residual_out = self._scale_residual_norm_scale_shift(
                     norm_module,
@@ -1324,7 +1401,12 @@ class QwenImageTransformerBlock(nn.Module):
         )
 
         # QwenAttnProcessor2_0 returns (img_output, txt_output) when encoder_hidden_states is provided
-        img_attn_output, txt_attn_output = attn_output
+        (
+            img_attn_output,
+            txt_attn_output,
+            img_attn_bias,
+            txt_attn_bias,
+        ) = attn_output
         # Process image stream - norm2 + MLP
         img_modulated2, hidden_states, img_gate2 = self._modulate(
             img_attn_output,
@@ -1333,20 +1415,32 @@ class QwenImageTransformerBlock(nn.Module):
             modulate_index,
             gate_x=img_gate1,
             residual_x=hidden_states,
+            x_bias=img_attn_bias,
             use_bcg_helpers=use_bcg_helpers,
         )
-        img_mlp_output = self.img_mlp(img_modulated2)
+        if isinstance(self.img_mlp, QwenImageFeedForward):
+            img_mlp_output, img_mlp_bias = self.img_mlp.forward_with_bias(
+                img_modulated2
+            )
+        else:
+            img_mlp_output = self.img_mlp(img_modulated2)
+            img_mlp_bias = None
 
         if img_mlp_output.dim() == 2:
             img_mlp_output = img_mlp_output.unsqueeze(0)
-        if use_bcg_helpers:
-            hidden_states = self._mul_add(img_mlp_output, img_gate2, hidden_states)
-        else:
-            hidden_states = self.fuse_mul_add(img_mlp_output, img_gate2, hidden_states)
+        hidden_states = self._bias_mul_add(
+            img_mlp_output,
+            img_mlp_bias,
+            img_gate2,
+            hidden_states,
+            use_bcg_helpers=use_bcg_helpers,
+        )
 
         # Process text stream - norm2 + MLP
         txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
         if use_bcg_helpers:
+            if txt_attn_bias is not None:
+                txt_attn_output = txt_attn_output + txt_attn_bias
             (
                 txt_modulated2,
                 encoder_hidden_states,
@@ -1358,6 +1452,15 @@ class QwenImageTransformerBlock(nn.Module):
                 shift=txt_shift2,
                 scale=txt_scale2,
             )
+        elif txt_attn_bias is not None:
+            txt_modulated2, encoder_hidden_states, _ = self._modulate(
+                txt_attn_output,
+                txt_mod2,
+                self.txt_norm2,
+                gate_x=txt_gate1,
+                residual_x=encoder_hidden_states,
+                x_bias=txt_attn_bias,
+            )
         else:
             txt_modulated2, encoder_hidden_states = self.txt_norm2(
                 residual=encoder_hidden_states,
@@ -1367,18 +1470,23 @@ class QwenImageTransformerBlock(nn.Module):
                 scale=txt_scale2,
             )
         txt_gate2 = txt_gate2_raw.unsqueeze(1)
-        txt_mlp_output = self.txt_mlp(txt_modulated2)
+        if isinstance(self.txt_mlp, QwenImageFeedForward):
+            txt_mlp_output, txt_mlp_bias = self.txt_mlp.forward_with_bias(
+                txt_modulated2
+            )
+        else:
+            txt_mlp_output = self.txt_mlp(txt_modulated2)
+            txt_mlp_bias = None
 
         if txt_mlp_output.dim() == 2:
             txt_mlp_output = txt_mlp_output.unsqueeze(0)
-        if use_bcg_helpers:
-            encoder_hidden_states = self._mul_add(
-                txt_mlp_output, txt_gate2, encoder_hidden_states
-            )
-        else:
-            encoder_hidden_states = self.fuse_mul_add(
-                txt_mlp_output, txt_gate2, encoder_hidden_states
-            )
+        encoder_hidden_states = self._bias_mul_add(
+            txt_mlp_output,
+            txt_mlp_bias,
+            txt_gate2,
+            encoder_hidden_states,
+            use_bcg_helpers=use_bcg_helpers,
+        )
 
         # Clip to prevent overflow for fp16
         if encoder_hidden_states.dtype == torch.float16: