[Diffusion] Fuse Qwen-Image residual norm and NVFP4 quantization (#37129)

Co-authored-by: Cursor <cursoragent@cursor.com>

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

@@ -22,6 +22,14 @@
 #include <sgl_kernel/vec.cuh>    // For AlignedVector
 #include <sgl_kernel/warp.cuh>   // For warp::reduce_sum
 
+#if defined(ENABLE_FP4) && ENABLE_FP4
+// FlashInfer's TensorRT-LLM quantization helper uses the C macro directly.
+#ifndef FLT_MAX
+#define FLT_MAX __FLT_MAX__
+#endif
+#include <tensorrt_llm/kernels/quantization_utils.cuh>
+#endif
+
 #include <cstdint>
 
 namespace sglang {
@@ -41,6 +49,7 @@ struct NormScaleShiftParams {
   void* y;
   void* res_out;
   void* quantized;
+  void* quant_scales;
   const void* x;
   const void* input_bias;
   const void* residual;
@@ -48,6 +57,8 @@ struct NormScaleShiftParams {
   const void* scale;
   const void* shift;
   const void* input_scale;
+  const void* global_scale;
+  uint32_t num_rows;
   float eps;
 };
 
@@ -78,13 +89,24 @@ SGL_DEVICE float triton_scale_reciprocal(float scale) {
   return reciprocal;
 }
 
-template <bool kHasResidual, bool kHasInputBias = false, bool kQuantizeFp8 = false>
+template <bool kHasResidual, bool kHasInputBias = false, bool kQuantizeFp8 = false, bool kQuantizeNvfp4 = false>
 __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_constant__ params) {
+  static_assert(!(kQuantizeFp8 && kQuantizeNvfp4));
   using namespace device;
   using Vec = AlignedVector<bf16_t, kVecElems>;
 
   const int row = blockIdx.x;
   const int tid = threadIdx.x;
+#if defined(ENABLE_FP4) && ENABLE_FP4
+  if constexpr (kQuantizeNvfp4) {
+    if (row >= params.num_rows) {
+      auto* scales = static_cast<uint8_t*>(params.quant_scales);
+      const int64_t scale_offset = tensorrt_llm::kernels::get_sf_out_offset_128x4(row, tid, kThreads);
+      scales[scale_offset] = 0;
+      return;
+    }
+  }
+#endif
   const int lane = tid & int(kWarpThreads - 1);
   const int warp = tid >> 5;
   const int row_offset = row * kHidden;
@@ -168,9 +190,29 @@ __global__ void norm_scale_shift_kernel(const NormScaleShiftParams __grid_consta
       qv[i] = static_cast<fp8_e4m3_t>(clamped);
     }
   }
-  yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
-  if constexpr (kQuantizeFp8) {
-    qv.store(static_cast<fp8_e4m3_t*>(params.quantized) + row_offset + elem_offset);
+  if constexpr (kQuantizeNvfp4) {
+#if defined(ENABLE_FP4) && ENABLE_FP4
+    tensorrt_llm::kernels::PackedVec<__nv_bfloat16, kVecElems> quant_vec;
+    auto* quant_values = reinterpret_cast<__nv_bfloat16*>(&quant_vec);
+#pragma unroll
+    for (int i = 0; i < kVecElems; ++i) {
+      quant_values[i] = static_cast<__nv_bfloat16>(yv[i]);
+    }
+
+    auto* scales = static_cast<uint8_t*>(params.quant_scales);
+    const int64_t scale_offset = tensorrt_llm::kernels::get_sf_out_offset_128x4(row, tid, kThreads);
+    const float global_scale = *static_cast<const float*>(params.global_scale);
+    const uint64_t packed = tensorrt_llm::kernels::cvt_warp_fp16_to_fp4<__nv_bfloat16, kVecElems, kVecElems, false>(
+        quant_vec, global_scale, scales + scale_offset);
+    static_cast<uint64_t*>(params.quantized)[int64_t(row) * kThreads + tid] = packed;
+#else
+    static_assert(!kQuantizeNvfp4);
+#endif
+  } else {
+    yv.store(static_cast<bf16_t*>(params.y) + row_offset + elem_offset);
+    if constexpr (kQuantizeFp8) {
+      qv.store(static_cast<fp8_e4m3_t*>(params.quantized) + row_offset + elem_offset);
+    }
   }
 }
 
@@ -226,6 +268,7 @@ struct NormScaleShiftKernel {
         .y = y.data_ptr(),
         .res_out = nullptr,
         .quantized = nullptr,
+        .quant_scales = nullptr,
         .x = x.data_ptr(),
         .input_bias = nullptr,
         .residual = nullptr,
@@ -233,6 +276,8 @@ struct NormScaleShiftKernel {
         .scale = scale.data_ptr(),
         .shift = shift.data_ptr(),
         .input_scale = nullptr,
+        .global_scale = nullptr,
+        .num_rows = grid,
         .eps = static_cast<float>(eps),
     };
     LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<false>, params);
@@ -268,6 +313,7 @@ struct ScaleResidualNormScaleShiftKernel {
         .y = y.data_ptr(),
         .res_out = res_out.data_ptr(),
         .quantized = nullptr,
+        .quant_scales = nullptr,
         .x = x.data_ptr(),
         .input_bias = nullptr,
         .residual = residual.data_ptr(),
@@ -275,6 +321,8 @@ struct ScaleResidualNormScaleShiftKernel {
         .scale = scale.data_ptr(),
         .shift = shift.data_ptr(),
         .input_scale = nullptr,
+        .global_scale = nullptr,
+        .num_rows = grid,
         .eps = static_cast<float>(eps),
     };
     LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true>, params);
@@ -306,6 +354,7 @@ struct NormScaleShiftFp8Kernel {
         .y = y.data_ptr(),
         .res_out = nullptr,
         .quantized = quantized.data_ptr(),
+        .quant_scales = nullptr,
         .x = x.data_ptr(),
         .input_bias = nullptr,
         .residual = nullptr,
@@ -313,6 +362,8 @@ struct NormScaleShiftFp8Kernel {
         .scale = scale.data_ptr(),
         .shift = shift.data_ptr(),
         .input_scale = input_scale.data_ptr(),
+        .global_scale = nullptr,
+        .num_rows = grid,
         .eps = static_cast<float>(eps),
     };
     LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<false, false, true>, params);
@@ -353,6 +404,7 @@ struct ScaleResidualNormScaleShiftFp8Kernel {
         .y = y.data_ptr(),
         .res_out = res_out.data_ptr(),
         .quantized = quantized.data_ptr(),
+        .quant_scales = nullptr,
         .x = x.data_ptr(),
         .input_bias = nullptr,
         .residual = residual.data_ptr(),
@@ -360,6 +412,8 @@ struct ScaleResidualNormScaleShiftFp8Kernel {
         .scale = scale.data_ptr(),
         .shift = shift.data_ptr(),
         .input_scale = input_scale.data_ptr(),
+        .global_scale = nullptr,
+        .num_rows = grid,
         .eps = static_cast<float>(eps),
     };
     LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true, false, true>, params);
@@ -402,6 +456,7 @@ struct BiasScaleResidualNormScaleShiftKernel {
         .y = y.data_ptr(),
         .res_out = res_out.data_ptr(),
         .quantized = nullptr,
+        .quant_scales = nullptr,
         .x = x.data_ptr(),
         .input_bias = input_bias.data_ptr(),
         .residual = residual.data_ptr(),
@@ -409,6 +464,8 @@ struct BiasScaleResidualNormScaleShiftKernel {
         .scale = scale.data_ptr(),
         .shift = shift.data_ptr(),
         .input_scale = nullptr,
+        .global_scale = nullptr,
+        .num_rows = grid,
         .eps = static_cast<float>(eps),
     };
     LaunchKernel(grid, kThreads, device.unwrap())(norm_scale_shift_kernel<true, true>, params);
@@ -435,6 +492,7 @@ struct BiasMulAddKernel {
         .y = y.data_ptr(),
         .res_out = nullptr,
         .quantized = nullptr,
+        .quant_scales = nullptr,
         .x = x.data_ptr(),
         .input_bias = input_bias.data_ptr(),
         .residual = residual.data_ptr(),
@@ -442,12 +500,70 @@ struct BiasMulAddKernel {
         .scale = nullptr,
         .shift = nullptr,
         .input_scale = nullptr,
+        .global_scale = nullptr,
+        .num_rows = grid,
         .eps = 0.0f,
     };
     LaunchKernel(grid, kThreads, device.unwrap())(bias_mul_add_kernel, params);
   }
 };
 
+#if defined(ENABLE_FP4) && ENABLE_FP4
+struct ScaleResidualNormScaleShiftNvfp4Kernel {
+  static void
+  run(tvm::ffi::TensorView quantized,
+      tvm::ffi::TensorView quant_scales,
+      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,
+      tvm::ffi::TensorView global_scale,
+      double eps) {
+    using namespace host;
+    auto N = SymbolicSize{"num_rows"};
+    auto NP = SymbolicSize{"num_rows_padded"};
+    auto device = SymbolicDevice{};
+    device.set_options<kDLCUDA>();
+
+    TensorMatcher({N, kHidden}).with_dtype<bf16_t>().with_device(device).verify(x).verify(residual).verify(res_out);
+    TensorMatcher({kHidden})
+        .with_dtype<bf16_t>()
+        .with_device(device)
+        .verify(input_bias)
+        .verify(gate)
+        .verify(scale)
+        .verify(shift);
+    TensorMatcher({N, kHidden / 2}).with_dtype<uint8_t>().with_device(device).verify(quantized);
+    TensorMatcher({NP, kThreads}).with_dtype<uint8_t>().with_device(device).verify(quant_scales);
+    TensorMatcher({1}).with_dtype<fp32_t>().with_device(device).verify(global_scale);
+
+    const uint32_t num_rows = verify_nss_geometry(N);
+    const uint32_t num_rows_padded = div_ceil(num_rows, uint32_t(128)) * 128;
+    RuntimeCheck(NP.unwrap() == num_rows_padded, "quant scale rows must be padded to 128");
+    const auto params = NormScaleShiftParams{
+        .y = nullptr,
+        .res_out = res_out.data_ptr(),
+        .quantized = quantized.data_ptr(),
+        .quant_scales = quant_scales.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(),
+        .input_scale = nullptr,
+        .global_scale = global_scale.data_ptr(),
+        .num_rows = num_rows,
+        .eps = static_cast<float>(eps),
+    };
+    LaunchKernel(num_rows_padded, kThreads, device.unwrap())(norm_scale_shift_kernel<true, true, false, true>, params);
+  }
+};
+#endif
+
 }  // namespace norm_scale_shift
 
 }  // namespace sglang

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

@@ -103,6 +103,13 @@ _SPECS: tuple[tuple[str, KernelBackend, str, frozenset, str], ...] = (
         _HIP,
         "FlyDSL (ROCm gfx950) residual + norm + scale/shift.",
     ),
+    (
+        "diffusion.scale_residual_norm_scale_shift_nvfp4",
+        KernelBackend.JIT,
+        "norm.norm_scale_shift_jit:try_fused_scale_residual_norm_scale_shift_nvfp4",
+        _CUDA,
+        "Qwen residual LayerNorm/modulation + NVFP4 quantization.",
+    ),
     (
         "diffusion.norm_scale_shift",
         KernelBackend.CUTE_DSL,
@@ -405,6 +412,7 @@ _EXPORTS: dict[str, str] = {
     "try_fused_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
     "try_fused_scale_residual_norm_scale_shift_fp8": "norm.norm_scale_shift_jit",
     "validate_scale_shift": "norm.scale_residual_norm_cutedsl",
+    "try_fused_scale_residual_norm_scale_shift_nvfp4": "norm.norm_scale_shift_jit",
     "can_use_wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton",
     "wan_rmsnorm_silu": "norm.wan_rmsnorm_silu_triton",
     "can_use_qk_rmsnorm_native": "norm.zimage_qk_rmsnorm_triton",

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

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 from typing import TYPE_CHECKING
 
 import torch
@@ -31,6 +32,12 @@ def _sm103(device: torch.device) -> bool:
     )
 
 
+def _blackwell_sm10x(device: torch.device) -> bool:
+    return (
+        torch.cuda.is_available() and torch.cuda.get_device_capability(device)[0] == 10
+    )
+
+
 def _nss_activation(t, like=None) -> bool:
     return (
         isinstance(t, torch.Tensor)
@@ -107,6 +114,25 @@ def norm_scale_shift_module() -> Module:
 _module = norm_scale_shift_module
 
 
+@cache_once
+def norm_scale_shift_nvfp4_module() -> Module:
+    return load_jit(
+        "norm_scale_shift_nvfp4_native",
+        cuda_files=["diffusion/norm_scale_shift.cuh"],
+        cuda_wrappers=[
+            (
+                "srnss_nvfp4_row",
+                "norm_scale_shift::ScaleResidualNormScaleShiftNvfp4Kernel::run",
+            ),
+        ],
+        extra_cuda_cflags=["-DENABLE_BF16", "-DENABLE_FP4"],
+        extra_dependencies=["flashinfer", "flashinfer_nv_internal"],
+    )
+
+
+_nvfp4_module = norm_scale_shift_nvfp4_module
+
+
 def fused_norm_scale_shift_fp8(x, scale, shift, input_scale, eps):
     """Return exact BF16 modulation output and its static E4M3 quantization."""
     normalized = torch.empty_like(x)
@@ -257,6 +283,81 @@ def _fp8_input_scale(t, device: torch.device) -> bool:
     )
 
 
+def _env_enabled(name: str) -> bool:
+    return os.getenv(name, "").strip().lower() not in {"", "0", "false", "off", "no"}
+
+
+def try_fused_scale_residual_norm_scale_shift_nvfp4(
+    residual,
+    x,
+    input_bias,
+    gate,
+    weight,
+    bias,
+    scale,
+    shift,
+    global_scale,
+    norm_type,
+    eps,
+):
+    """Fuse Qwen residual LayerNorm/modulation with FC1-input NVFP4 quantization."""
+    if (
+        torch.compiler.is_compiling()
+        or _env_enabled("FLASHINFER_DISABLE_FP4_QUANT_FAST_MATH")
+        or _env_enabled("TRTLLM_DISABLE_FP4_QUANT_FAST_MATH")
+        or _env_enabled("FLASHINFER_NVFP4_4OVER6")
+    ):
+        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 _blackwell_sm10x(x.device)
+    ):
+        return None
+    if torch.cuda.is_current_stream_capturing():
+        return None
+
+    gate = _row_bf16(gate, x.device)
+    input_bias = _row_bf16(input_bias, x.device)
+    scale = _row_bf16(scale, x.device)
+    shift = _row_bf16(shift, x.device)
+    if input_bias is None or gate is None or scale is None or shift is None:
+        return None
+    if not (
+        isinstance(global_scale, torch.Tensor)
+        and global_scale.is_cuda
+        and global_scale.device == x.device
+        and global_scale.dtype == torch.float32
+        and global_scale.numel() == 1
+        and global_scale.is_contiguous()
+    ):
+        return None
+
+    rows = x.numel() // _HIDDEN
+    padded_rows = (rows + 127) // 128 * 128
+    quantized = torch.empty((rows, _HIDDEN // 2), dtype=torch.uint8, device=x.device)
+    quant_scales = torch.empty(
+        (padded_rows, _HIDDEN // 16), dtype=torch.uint8, device=x.device
+    )
+    residual_out = torch.empty_like(x)
+    _nvfp4_module().srnss_nvfp4_row(
+        quantized,
+        quant_scales,
+        residual_out.view(-1, _HIDDEN),
+        residual.view(-1, _HIDDEN),
+        x.view(-1, _HIDDEN),
+        input_bias,
+        gate,
+        scale,
+        shift,
+        global_scale.reshape(1),
+        float(eps),
+    )
+    return (quantized, quant_scales), residual_out
+
+
 def try_fused_bias_scale_residual_norm_scale_shift(
     residual, x, input_bias, gate, weight, bias, scale, shift, norm_type, eps
 ):

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

@@ -29,6 +29,7 @@ from sglang.kernels.ops.diffusion import (
     try_fused_norm_scale_shift_fp8,
     try_fused_qwen_qkv_epilogue,
     try_fused_scale_residual_norm_scale_shift_fp8,
+    try_fused_scale_residual_norm_scale_shift_nvfp4,
 )
 from sglang.multimodal_gen.configs.models.dits.qwenimage import QwenImageDitConfig
 from sglang.multimodal_gen.configs.models.fsdp import is_transformer_block
@@ -78,7 +79,9 @@ from sglang.multimodal_gen.runtime.layers.quantization.configs.nunchaku_config i
     is_nunchaku_available,
 )
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
+    ModelOptFp4LinearMethod,
     ModelOptFp8LinearMethod,
+    apply_nvfp4_gemm_prequantized,
 )
 from sglang.multimodal_gen.runtime.layers.rotary_embedding import (
     apply_flashinfer_rope_qk_inplace,
@@ -1095,12 +1098,23 @@ class QwenImageGELU(nn.Module):
         # epilogue. Off by default; mounted per batch by the denoising stage.
         mark_fused_gelu_site(self, "proj")
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        if fused_gelu_active(self) and can_use_linear_gelu(self.proj, hidden_states):
+    def forward(
+        self,
+        hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+    ) -> torch.Tensor:
+        if isinstance(hidden_states, tuple):
+            hidden_states = apply_nvfp4_gemm_prequantized(
+                self.proj,
+                *hidden_states,
+                output_dtype=self.proj.params_dtype,
+                bias=self.proj.bias,
+            )
+        elif fused_gelu_active(self) and can_use_linear_gelu(self.proj, hidden_states):
             return fused_linear_gelu_tanh(
                 hidden_states, self.proj.weight, self.proj.bias
             )
-        hidden_states, _ = self.proj(hidden_states)
+        else:
+            hidden_states, _ = self.proj(hidden_states)
         return F.gelu(hidden_states, approximate="tanh")
 
 
@@ -1154,13 +1168,15 @@ class QwenImageFeedForward(nn.Module):
         )
 
     def forward_with_bias(
-        self, hidden_states: torch.Tensor
+        self, hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
     ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
         hidden_states = self.net[0](hidden_states)
         hidden_states = self.net[1](hidden_states)
         return self.net[2](hidden_states)
 
-    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    def forward(
+        self, hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor]
+    ) -> torch.Tensor:
         hidden_states, bias = self.forward_with_bias(hidden_states)
         return hidden_states if bias is None else hidden_states + bias
 
@@ -1294,6 +1310,20 @@ class QwenImageTransformerBlock(nn.Module):
             self.img_mlp = NunchakuFeedForward(self.img_mlp, **nunchaku_kwargs)
             self.txt_mlp = NunchakuFeedForward(self.txt_mlp, **nunchaku_kwargs)
 
+        self._enable_nvfp4_resnorm_quant = False
+        capability = current_platform.get_device_capability()
+        if (
+            not nunchaku_enabled
+            and dim == 3072
+            and capability is not None
+            and capability.major == 10
+        ):
+            img_fc1 = self.img_mlp.net[0].proj
+            txt_fc1 = self.txt_mlp.net[0].proj
+            self._enable_nvfp4_resnorm_quant = isinstance(
+                img_fc1.quant_method, ModelOptFp4LinearMethod
+            ) and isinstance(txt_fc1.quant_method, ModelOptFp4LinearMethod)
+
         self._fp8_img_attn_norm_quant = False
         self._fp8_txt_attn_norm_quant = False
         self._fp8_img_mlp_norm_quant = False
@@ -1428,6 +1458,52 @@ class QwenImageTransformerBlock(nn.Module):
         )
         return img_mod_params, txt_mod_params
 
+    def _try_nvfp4_resnorm_quant(
+        self,
+        norm_module: ScaleResidualLayerNormScaleShift,
+        mlp: QwenImageFeedForward,
+        *,
+        residual: torch.Tensor,
+        x: torch.Tensor,
+        x_bias: Optional[torch.Tensor],
+        residual_gate: torch.Tensor,
+        mod_params: torch.Tensor,
+        modulate_index: Optional[torch.Tensor],
+        use_bcg_helpers: bool,
+    ) -> Optional[
+        tuple[
+            tuple[torch.Tensor, torch.Tensor],
+            torch.Tensor,
+            torch.Tensor,
+        ]
+    ]:
+        if (
+            not self._enable_nvfp4_resnorm_quant
+            or modulate_index is not None
+            or use_bcg_helpers
+        ):
+            return None
+
+        shift, scale, gate = mod_params.chunk(3, dim=-1)
+        fc1 = mlp.net[0].proj
+        result = try_fused_scale_residual_norm_scale_shift_nvfp4(
+            residual,
+            x,
+            x_bias,
+            residual_gate,
+            getattr(norm_module.norm, "weight", None),
+            getattr(norm_module.norm, "bias", None),
+            scale.unsqueeze(1),
+            shift.unsqueeze(1),
+            fc1.input_scale_inv,
+            norm_module.norm_type,
+            norm_module.eps,
+        )
+        if result is None:
+            return None
+        packed, residual_out = result
+        return packed, residual_out, gate.unsqueeze(1)
+
     def _try_fp8_norm_quant(
         self,
         norm_module: LayerNormScaleShift,
@@ -1742,40 +1818,55 @@ class QwenImageTransformerBlock(nn.Module):
             txt_attn_bias,
         ) = attn_output
         # Process image stream - norm2 + MLP
-        img_fp8_mlp = self._try_fp8_residual_norm_quant(
+        img_nvfp4 = self._try_nvfp4_resnorm_quant(
             self.img_norm2,
+            self.img_mlp,
             residual=hidden_states,
             x=img_attn_output,
+            x_bias=img_attn_bias,
             residual_gate=img_gate1,
             mod_params=img_mod2,
-            input_scale=(
-                self.img_mlp.net[0].proj.input_scale
-                if self._fp8_img_mlp_norm_quant
-                else None
-            ),
-            enabled=self._fp8_img_mlp_norm_quant,
             modulate_index=modulate_index,
             use_bcg_helpers=use_bcg_helpers,
         )
-        if img_fp8_mlp is None:
-            img_modulated2, hidden_states, img_gate2 = self._modulate(
-                img_attn_output,
-                img_mod2,
-                self.img_norm2,
-                modulate_index,
-                gate_x=img_gate1,
-                residual_x=hidden_states,
-                x_bias=img_attn_bias,
-                use_bcg_helpers=use_bcg_helpers,
-            )
+        if img_nvfp4 is not None:
+            img_modulated2, hidden_states, img_gate2 = img_nvfp4
             img_modulated2_bf16 = None
         else:
-            (
-                img_modulated2,
-                hidden_states,
-                img_gate2,
-                img_modulated2_bf16,
-            ) = img_fp8_mlp
+            img_fp8_mlp = self._try_fp8_residual_norm_quant(
+                self.img_norm2,
+                residual=hidden_states,
+                x=img_attn_output,
+                residual_gate=img_gate1,
+                mod_params=img_mod2,
+                input_scale=(
+                    self.img_mlp.net[0].proj.input_scale
+                    if self._fp8_img_mlp_norm_quant
+                    else None
+                ),
+                enabled=self._fp8_img_mlp_norm_quant,
+                modulate_index=modulate_index,
+                use_bcg_helpers=use_bcg_helpers,
+            )
+            if img_fp8_mlp is None:
+                img_modulated2, hidden_states, img_gate2 = self._modulate(
+                    img_attn_output,
+                    img_mod2,
+                    self.img_norm2,
+                    modulate_index,
+                    gate_x=img_gate1,
+                    residual_x=hidden_states,
+                    x_bias=img_attn_bias,
+                    use_bcg_helpers=use_bcg_helpers,
+                )
+                img_modulated2_bf16 = None
+            else:
+                (
+                    img_modulated2,
+                    hidden_states,
+                    img_gate2,
+                    img_modulated2_bf16,
+                ) = img_fp8_mlp
         if isinstance(self.img_mlp, QwenImageFeedForward):
             img_mlp_output, img_mlp_bias = self.img_mlp.forward_with_bias(
                 img_modulated2
@@ -1797,63 +1888,79 @@ class QwenImageTransformerBlock(nn.Module):
 
         # Process text stream - norm2 + MLP
         txt_shift2, txt_scale2, txt_gate2_raw = txt_mod2.chunk(3, dim=-1)
-        txt_fp8_mlp = self._try_fp8_residual_norm_quant(
+        txt_nvfp4 = self._try_nvfp4_resnorm_quant(
             self.txt_norm2,
+            self.txt_mlp,
             residual=encoder_hidden_states,
             x=txt_attn_output,
+            x_bias=txt_attn_bias,
             residual_gate=txt_gate1,
             mod_params=txt_mod2,
-            input_scale=(
-                self.txt_mlp.net[0].proj.input_scale
-                if self._fp8_txt_mlp_norm_quant
-                else None
-            ),
-            enabled=self._fp8_txt_mlp_norm_quant,
             modulate_index=modulate_index,
             use_bcg_helpers=use_bcg_helpers,
         )
-        if txt_fp8_mlp is not None:
-            (
-                txt_modulated2,
-                encoder_hidden_states,
-                txt_gate2,
-                txt_modulated2_bf16,
-            ) = txt_fp8_mlp
-        elif use_bcg_helpers:
+        txt_fp8_mlp = None
+        if txt_nvfp4 is not None:
+            txt_modulated2, encoder_hidden_states, txt_gate2 = txt_nvfp4
             txt_modulated2_bf16 = None
-            if txt_attn_bias is not None:
-                txt_attn_output = txt_attn_output + txt_attn_bias
-            (
-                txt_modulated2,
-                encoder_hidden_states,
-            ) = self._scale_residual_norm_scale_shift(
-                self.txt_norm2,
-                residual=encoder_hidden_states,
-                x=txt_attn_output,
-                gate=txt_gate1,
-                shift=txt_shift2,
-                scale=txt_scale2,
-            )
-        elif txt_attn_bias is not None:
-            txt_modulated2_bf16 = 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_bf16 = None
-            txt_modulated2, encoder_hidden_states = self.txt_norm2(
+            txt_fp8_mlp = self._try_fp8_residual_norm_quant(
+                self.txt_norm2,
                 residual=encoder_hidden_states,
                 x=txt_attn_output,
-                gate=txt_gate1,
-                shift=txt_shift2,
-                scale=txt_scale2,
+                residual_gate=txt_gate1,
+                mod_params=txt_mod2,
+                input_scale=(
+                    self.txt_mlp.net[0].proj.input_scale
+                    if self._fp8_txt_mlp_norm_quant
+                    else None
+                ),
+                enabled=self._fp8_txt_mlp_norm_quant,
+                modulate_index=modulate_index,
+                use_bcg_helpers=use_bcg_helpers,
             )
-        if txt_fp8_mlp is None:
+            if txt_fp8_mlp is not None:
+                (
+                    txt_modulated2,
+                    encoder_hidden_states,
+                    txt_gate2,
+                    txt_modulated2_bf16,
+                ) = txt_fp8_mlp
+            elif use_bcg_helpers:
+                txt_modulated2_bf16 = None
+                if txt_attn_bias is not None:
+                    txt_attn_output = txt_attn_output + txt_attn_bias
+                (
+                    txt_modulated2,
+                    encoder_hidden_states,
+                ) = self._scale_residual_norm_scale_shift(
+                    self.txt_norm2,
+                    residual=encoder_hidden_states,
+                    x=txt_attn_output,
+                    gate=txt_gate1,
+                    shift=txt_shift2,
+                    scale=txt_scale2,
+                )
+            elif txt_attn_bias is not None:
+                txt_modulated2_bf16 = 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_bf16 = None
+                txt_modulated2, encoder_hidden_states = self.txt_norm2(
+                    residual=encoder_hidden_states,
+                    x=txt_attn_output,
+                    gate=txt_gate1,
+                    shift=txt_shift2,
+                    scale=txt_scale2,
+                )
+        if txt_nvfp4 is None and txt_fp8_mlp is None:
             txt_gate2 = txt_gate2_raw.unsqueeze(1)
         if isinstance(self.txt_mlp, QwenImageFeedForward):
             txt_mlp_output, txt_mlp_bias = self.txt_mlp.forward_with_bias(

test/registered/kernels/benchmark/diffusion/bench_qwen_image_resnorm_nvfp4_quant.py

@@ -0,0 +1,97 @@
+import flashinfer
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    fused_scale_residual_norm_scale_shift,
+    try_fused_scale_residual_norm_scale_shift_nvfp4,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=30,
+    stage="base-b-kernel-benchmark",
+    runner_config="1-gpu-large",
+    disabled="standalone Qwen-Image NVFP4 residual-norm benchmark",
+)
+
+
+def _benchmark(fn, iterations: int = 100) -> float:
+    for _ in range(10):
+        fn()
+    torch.cuda.synchronize()
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    start.record()
+    for _ in range(iterations):
+        fn()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) * 1000 / iterations
+
+
+def _run_case(token_count: int) -> None:
+    hidden_size = 3072
+    generator = torch.Generator(device="cuda")
+    generator.manual_seed(20260830 + token_count)
+
+    def randn(shape):
+        return torch.randn(
+            shape, device="cuda", dtype=torch.bfloat16, generator=generator
+        ).contiguous()
+
+    x = randn((1, token_count, hidden_size))
+    residual = randn(x.shape)
+    input_bias = randn((hidden_size,))
+    gate = randn((1, 1, hidden_size))
+    scale = randn((1, 1, hidden_size))
+    shift = randn((1, 1, hidden_size))
+    global_scale = torch.tensor(0.625, device="cuda", dtype=torch.float32)
+
+    def baseline():
+        modulated, residual_out = fused_scale_residual_norm_scale_shift(
+            residual, x + input_bias, gate, None, None, scale, shift, "layer", 1e-6
+        )
+        quantized, quant_scales = flashinfer.fp4_quantize(
+            modulated.view(-1, hidden_size), global_scale
+        )
+        return quantized, quant_scales, residual_out
+
+    def fused():
+        result = try_fused_scale_residual_norm_scale_shift_nvfp4(
+            residual,
+            x,
+            input_bias,
+            gate,
+            None,
+            None,
+            scale,
+            shift,
+            global_scale,
+            "layer",
+            1e-6,
+        )
+        assert result is not None
+        (quantized, quant_scales), residual_out = result
+        return quantized, quant_scales, residual_out
+
+    expected = baseline()
+    actual = fused()
+    exact = [torch.equal(lhs, rhs) for lhs, rhs in zip(actual, expected)]
+    baseline_us = _benchmark(baseline)
+    fused_us = _benchmark(fused)
+    print(
+        {
+            "tokens": token_count,
+            "baseline_us": baseline_us,
+            "fused_us": fused_us,
+            "speedup": baseline_us / fused_us,
+            "exact": exact,
+        }
+    )
+
+
+if __name__ == "__main__":
+    if not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] != 10:
+        raise RuntimeError("This benchmark requires an NVIDIA Blackwell SM10x GPU")
+    for tokens in (17, 1024, 4096, 4608):
+        _run_case(tokens)

test/registered/kernels/ops/quantization/test_diffusion_nvfp4_scaled_mm.py

@@ -4,6 +4,10 @@ import flashinfer
 import pytest
 import torch
 
+from sglang.kernels.ops.diffusion import (
+    fused_scale_residual_norm_scale_shift,
+    try_fused_scale_residual_norm_scale_shift_nvfp4,
+)
 from sglang.multimodal_gen.runtime.layers.quantization import (
     modelopt_quant as diffusion_modelopt_quant,
 )
@@ -37,6 +41,62 @@ def _nvfp4_supported() -> bool:
     return torch.cuda.is_available() and torch.cuda.get_device_capability() >= (10, 0)
 
 
+def _qwen_resnorm_nvfp4_supported() -> bool:
+    return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 10
+
+
+@pytest.mark.skipif(
+    not _qwen_resnorm_nvfp4_supported(),
+    reason="Qwen-Image fused residual norm + NVFP4 quantization requires SM10x",
+)
+@pytest.mark.parametrize("token_count", [17, 1024])
+def test_qwen_image_fused_resnorm_nvfp4_quant_is_exact(token_count: int) -> None:
+    hidden_size = 3072
+    generator = torch.Generator(device=DEVICE)
+    generator.manual_seed(20260830 + token_count)
+
+    def randn(shape):
+        return torch.randn(
+            shape, device=DEVICE, dtype=DTYPE, generator=generator
+        ).contiguous()
+
+    residual = randn((1, token_count, hidden_size))
+    x = randn((1, token_count, hidden_size))
+    input_bias = randn((hidden_size,))
+    gate = randn((1, 1, hidden_size))
+    scale = randn((1, 1, hidden_size))
+    shift = randn((1, 1, hidden_size))
+    global_scale = torch.tensor(512.0, device=DEVICE, dtype=torch.float32)
+
+    expected_modulated, expected_residual = fused_scale_residual_norm_scale_shift(
+        residual, x + input_bias, gate, None, None, scale, shift, "layer", 1e-6
+    )
+    expected_quantized, expected_scales = flashinfer.fp4_quantize(
+        expected_modulated.view(-1, hidden_size), global_scale
+    )
+
+    actual = try_fused_scale_residual_norm_scale_shift_nvfp4(
+        residual,
+        x,
+        input_bias,
+        gate,
+        None,
+        None,
+        scale,
+        shift,
+        global_scale,
+        "layer",
+        1e-6,
+    )
+    assert actual is not None
+    (actual_quantized, actual_scales), actual_residual = actual
+    assert torch.equal(actual_quantized, expected_quantized)
+    assert torch.equal(
+        actual_scales.view(torch.uint8), expected_scales.view(torch.uint8)
+    )
+    assert torch.equal(actual_residual, expected_residual)
+
+
 def _make_global_scale(x: torch.Tensor) -> torch.Tensor:
     max_abs = torch.amax(x.abs()).clamp_min_(1e-6)
     return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / max_abs).to(torch.float32)