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

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

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)