[Diffusion] Fuse Qwen-Image final adaptive LayerNorm (#37144)

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

@@ -7,20 +7,23 @@ from sglang.kernels.jit.benchmark.utils import run_benchmark_no_cudagraph
 from sglang.kernels.ops.diffusion import (
     fuse_layernorm_scale_shift_gate_select01_kernel,
     fuse_residual_layernorm_scale_shift_gate_select01_kernel,
+    fused_layernorm_modulate_raw,
     norm_infer,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.utils import is_in_ci
 
 register_cuda_ci(
-    est_time=13, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+    est_time=15, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
 )
 register_amd_ci(est_time=13, stage="jit-kernel-benchmark", runner_config="amd")
 
 if is_in_ci():
     B_RANGE, S_RANGE, D_RANGE = [1], [128], [3072]
+    NORM_OUT_S_RANGE = [128]
 else:
     B_RANGE, S_RANGE, D_RANGE = [1, 2], [128, 512, 2048], [1024, 1536, 3072]
+    NORM_OUT_S_RANGE = [128, 512, 2048, 4096, 4608]
 
 DTYPE = torch.bfloat16
 DEVICE = "cuda"
@@ -29,6 +32,7 @@ LINE_VALS = ["split", "fused"]
 LINE_NAMES = ["Triton Norm + Torch Select", "Fused Triton"]
 STYLES = [("red", "-"), ("blue", "--")]
 CONFIG = [(b, s, d) for b in B_RANGE for s in S_RANGE for d in D_RANGE]
+NORM_OUT_CONFIG = [(1, s, 3072) for s in NORM_OUT_S_RANGE]
 
 
 def _make_common_inputs(batch_size: int, seq_len: int, hidden_size: int):
@@ -62,6 +66,40 @@ def _apply_select01_modulation(
     return x * (1 + scale) + shift, gate
 
 
+@triton.testing.perf_report(
+    triton.testing.Benchmark(
+        x_names=["B", "S", "D"],
+        x_vals=NORM_OUT_CONFIG,
+        line_arg="provider",
+        line_vals=LINE_VALS,
+        line_names=["Torch LayerNorm + Modulate", "Fused Triton"],
+        styles=STYLES,
+        ylabel="us",
+        plot_name="qwen_image_norm_out",
+        args={},
+    )
+)
+def bench_norm_out(B: int, S: int, D: int, provider: str) -> Tuple[float, float, float]:
+    x = torch.randn(B, S, D, dtype=DTYPE, device=DEVICE)
+    scale = torch.randn(B, D, dtype=DTYPE, device=DEVICE)
+    shift = torch.randn(B, D, dtype=DTYPE, device=DEVICE)
+    norm = torch.nn.LayerNorm(
+        D, eps=EPS, elementwise_affine=False, device=DEVICE, dtype=DTYPE
+    )
+
+    if provider == "split":
+
+        def fn():
+            return norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+
+    else:
+
+        def fn():
+            return fused_layernorm_modulate_raw(x, scale, shift, EPS)
+
+    return run_benchmark_no_cudagraph(fn)
+
+
 @triton.testing.perf_report(
     triton.testing.Benchmark(
         x_names=["B", "S", "D"],
@@ -176,6 +214,14 @@ def bench_residual_layernorm_scale_shift_gate_select01(
 
 
 if __name__ == "__main__":
+    # The bit-exact LayerNorm kernel uses CUDA inline PTX; keep the shared
+    # Qwen modulation benchmark runnable on its registered ROCm lane.
+    if torch.version.hip is None:
+        print(f"\n{'=' * 80}")
+        print("Benchmark: qwen_image norm_out")
+        print(f"{'=' * 80}\n")
+        bench_norm_out.run(print_data=True)
+
     print(f"\n{'=' * 80}")
     print("Benchmark: qwen_image layernorm + scale_shift_gate_select01")
     print(f"{'=' * 80}\n")

test/registered/kernels/ops/diffusion/test_model_fast_paths.py

@@ -36,6 +36,7 @@ import sglang.multimodal_gen.runtime.models.dits.flux_2 as flux2
 import sglang.multimodal_gen.runtime.models.dits.glm_image as glm_image
 import sglang.multimodal_gen.runtime.models.dits.longcat_image as longcat_image
 import sglang.multimodal_gen.runtime.models.dits.ltx_2 as ltx2_module
+import sglang.multimodal_gen.runtime.models.dits.qwen_image as qwen_image
 import sglang.multimodal_gen.runtime.models.dits.sana as sana
 from sglang.kernels.ops.diffusion import (
     can_use_fused_layernorm_modulate,
@@ -98,6 +99,7 @@ from sglang.multimodal_gen.runtime.models.dits.ltx_2 import _ltx2_rms_norm_modul
 from sglang.multimodal_gen.runtime.models.dits.qwen_image import (
     QwenImageTransformerBlock,
     _qwen_modulation_cache_key,
+    _qwen_norm_out,
 )
 from sglang.multimodal_gen.runtime.models.dits.sana import (
     _eager_ln_modulate as _sana_eager_ln_modulate,
@@ -151,6 +153,81 @@ def test_bitexact_norm_guards_follow_platform():
     assert can_use_fused_rmsnorm_scale_shift(x, weight, vec, vec) is is_cuda()
 
 
+# -------------------------------------------------------------------------
+# Qwen-Image -- final LayerNorm + adaLN scale/shift
+# -------------------------------------------------------------------------
+
+
+@requires_inline_ptx
+def test_qwen_norm_out_matches_adaln_reference():
+    qwen_image._QWEN_NORM_OUT.disabled = False
+    qwen_image._QWEN_NORM_OUT.verified = False
+    qwen_image._QWEN_NORM_OUT_SIGS.clear()
+    torch.manual_seed(0)
+    norm_out = (
+        qwen_image.AdaLayerNormContinuous(
+            3072, 3072, elementwise_affine=False, eps=1e-6
+        )
+        .cuda()
+        .bfloat16()
+    )
+    hidden_states = torch.randn(1, 257, 3072, device="cuda", dtype=torch.bfloat16)
+    conditioning = torch.randn(1, 3072, device="cuda", dtype=torch.bfloat16)
+
+    expected = norm_out(hidden_states, conditioning)
+    actual = _qwen_norm_out(norm_out, hidden_states, conditioning)
+
+    assert torch.equal(actual, expected)
+    assert qwen_image._QWEN_NORM_OUT.verified
+    assert not qwen_image._QWEN_NORM_OUT.disabled
+
+
+def test_qwen_norm_out_preserves_compile_path(monkeypatch):
+    norm_out = (
+        qwen_image.AdaLayerNormContinuous(16, 16, elementwise_affine=False, eps=1e-6)
+        .cuda()
+        .bfloat16()
+    )
+    hidden_states = torch.randn(1, 3, 16, device="cuda", dtype=torch.bfloat16)
+    conditioning = torch.randn(1, 16, device="cuda", dtype=torch.bfloat16)
+    expected = norm_out(hidden_states, conditioning)
+
+    monkeypatch.setattr(torch.compiler, "is_compiling", lambda: True)
+    monkeypatch.setattr(
+        qwen_image,
+        "fused_layernorm_modulate_raw",
+        lambda *args, **kwargs: pytest.fail("compile path must not dispatch kernel"),
+    )
+
+    assert torch.equal(_qwen_norm_out(norm_out, hidden_states, conditioning), expected)
+
+
+def test_qwen_norm_out_does_not_verify_during_graph_capture(monkeypatch):
+    qwen_image._QWEN_NORM_OUT.disabled = False
+    qwen_image._QWEN_NORM_OUT.verified = False
+    qwen_image._QWEN_NORM_OUT_SIGS.clear()
+    norm_out = (
+        qwen_image.AdaLayerNormContinuous(
+            3072, 3072, elementwise_affine=False, eps=1e-6
+        )
+        .cuda()
+        .bfloat16()
+    )
+    hidden_states = torch.randn(1, 17, 3072, device="cuda", dtype=torch.bfloat16)
+    conditioning = torch.randn(1, 3072, device="cuda", dtype=torch.bfloat16)
+    expected = norm_out(hidden_states, conditioning)
+
+    monkeypatch.setattr(torch.cuda, "is_current_stream_capturing", lambda: True)
+    monkeypatch.setattr(
+        qwen_image,
+        "fused_layernorm_modulate_raw",
+        lambda *args, **kwargs: pytest.fail("capture must not verify a new layout"),
+    )
+
+    assert torch.equal(_qwen_norm_out(norm_out, hidden_states, conditioning), expected)
+    assert not qwen_image._QWEN_NORM_OUT_SIGS
+
+
 # -------------------------------------------------------------------------
 # FLUX.1 -- LayerNorm + adaLN modulate, and the shared-FF GELU site
 # -------------------------------------------------------------------------