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

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

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)