[diffusion] model: support qwen-image-2.1 (#39983)

Co-authored-by: Mick Qian <mickqian@users.noreply.github.com>
Co-authored-by: BBuf <1182563586@qq.com>

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

@@ -0,0 +1,91 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    BitExactFusionGate,
+    can_use_fused_complex_rope,
+    fused_complex_rope,
+)
+from sglang.multimodal_gen.runtime.models.dits import qwen_image21
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is not None,
+    reason="NVIDIA CUDA required",
+)
+
+
+def reference(x, rope):
+    z = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+    return torch.view_as_real(z * rope[None, :, None]).flatten(-2).to(x.dtype)
+
+
+def inputs(shape, dtype):
+    torch.manual_seed(42)
+    x = torch.randn(shape, device="cuda", dtype=dtype)
+    # a contiguous slice retains the nonzero cache offset used by SP ranks
+    angles = torch.randn(shape[1] + 5, shape[-1] // 2, device="cuda") * 20
+    rope = torch.polar(torch.ones_like(angles), angles)[5:]
+    return x, rope
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize(
+    "shape", [(1, 1, 1, 32), (2, 17, 3, 64), (1, 257, 16, 128), (1, 4096, 32, 128)]
+)
+def test_complex_rope_matches_complex_multiplication(dtype, shape):
+    x, rope = inputs(shape, dtype)
+    assert can_use_fused_complex_rope(x, rope)
+    actual = fused_complex_rope(x, rope)
+    torch.testing.assert_close(actual, reference(x, rope), atol=0, rtol=0)
+
+
+def test_complex_rope_layout_guards():
+    x, rope = inputs((2, 17, 3, 64), torch.bfloat16)
+    assert not can_use_fused_complex_rope(x.cpu(), rope.cpu())
+    assert not can_use_fused_complex_rope(x.double(), rope)
+    assert not can_use_fused_complex_rope(x, rope.to(torch.complex128))
+    assert not can_use_fused_complex_rope(x[:, ::2], rope[::2])
+    assert not can_use_fused_complex_rope(x, rope[:-1])
+    assert not can_use_fused_complex_rope(x[:, :0], rope[:0])
+
+
+def test_complex_rope_compile_and_graph_replay():
+    x, rope = inputs((1, 257, 8, 128), torch.bfloat16)
+    compiled = torch.compile(fused_complex_rope, fullgraph=True)
+    torch.testing.assert_close(compiled(x, rope), reference(x, rope), atol=0, rtol=0)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = fused_complex_rope(x, rope)
+    x.normal_()
+    graph.replay()
+    torch.testing.assert_close(out, reference(x, rope), atol=0, rtol=0)
+
+
+def test_qwen21_rope_first_sight_verification(monkeypatch):
+    x, rope = inputs((1, 257, 8, 128), torch.bfloat16)
+    gate = BitExactFusionGate("test complex RoPE")
+    monkeypatch.setattr(qwen_image21, "_ROPE_FUSION", gate)
+    torch.testing.assert_close(
+        qwen_image21.apply_rope(x, rope), reference(x, rope), atol=0, rtol=0
+    )
+    assert gate.verified and not gate.disabled
+
+    gate = BitExactFusionGate("test mismatched RoPE")
+    monkeypatch.setattr(qwen_image21, "_ROPE_FUSION", gate)
+    monkeypatch.setattr(
+        qwen_image21, "fused_complex_rope", lambda x, rope: torch.zeros_like(x)
+    )
+    torch.testing.assert_close(
+        qwen_image21.apply_rope(x, rope), reference(x, rope), atol=0, rtol=0
+    )
+    assert gate.disabled and not gate.verified
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

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

@@ -0,0 +1,90 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+
+import pytest
+import torch
+from torch.nn import functional as F
+
+from sglang.kernels.ops.diffusion import (
+    can_use_fused_layernorm_modulate,
+    fused_layernorm_modulate,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is not None,
+    reason="NVIDIA CUDA required",
+)
+
+
+def reference(x, scale, shift, eps):
+    out = F.layer_norm(x, (x.shape[-1],), eps=eps) * (1 + scale[:, None])
+    return out if shift is None else out + shift[:, None]
+
+
+@pytest.mark.parametrize("shape", [(1, 1, 128), (1, 4359, 4096), (2, 1024, 4096)])
+@pytest.mark.parametrize("amplitude,eps", [(1e-4, 1e-6), (1.0, 1e-6), (100.0, 1e-5)])
+@pytest.mark.parametrize("has_shift", [False, True])
+def test_modulation_preserves_bits(shape, amplitude, eps, has_shift):
+    torch.manual_seed(42)
+    x = torch.randn(shape, device="cuda", dtype=torch.bfloat16) * amplitude
+    modulation = torch.randn(shape[0], 4 * shape[-1], device="cuda", dtype=x.dtype)
+    scale, shift = modulation.chunk(4, dim=-1)[:2]
+    scale[:, :3] = torch.tensor([-1, 0, 1], device=x.device, dtype=x.dtype)
+    if not has_shift:
+        shift = None
+    assert can_use_fused_layernorm_modulate(x, scale, shift)
+    actual = fused_layernorm_modulate(x, scale, shift, eps)
+    expected = reference(x, scale, shift, eps)
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16))
+
+
+def test_scale_only_preserves_signed_zero():
+    x = torch.ones(1, 17, 128, device="cuda", dtype=torch.bfloat16)
+    scale = torch.full((1, 128), -2, device=x.device, dtype=x.dtype)
+    actual = fused_layernorm_modulate(x, scale, None, 1e-6)
+    expected = reference(x, scale, None, 1e-6)
+    assert torch.signbit(expected).all()
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16))
+
+
+def test_scale_only_layout_guards():
+    x = torch.randn(2, 17, 128, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(2, 128, device=x.device, dtype=x.dtype)
+    assert can_use_fused_layernorm_modulate(x, scale, None)
+    assert not can_use_fused_layernorm_modulate(x.cpu(), scale.cpu(), None)
+    assert not can_use_fused_layernorm_modulate(x.float(), scale.float(), None)
+    assert not can_use_fused_layernorm_modulate(x[:, ::2], scale, None)
+    assert not can_use_fused_layernorm_modulate(x, scale.float(), None)
+    assert not can_use_fused_layernorm_modulate(x, scale[:, :-1], None)
+    assert not can_use_fused_layernorm_modulate(x[:, :0], scale, None)
+    assert not can_use_fused_layernorm_modulate(x, scale, scale.float())
+    strided = torch.empty(2, 256, device=x.device, dtype=x.dtype)[:, :128]
+    assert not can_use_fused_layernorm_modulate(x, scale, strided)
+
+
+@pytest.mark.parametrize("has_shift", [False, True])
+def test_compile_and_graph_replay(has_shift):
+    x = torch.randn(2, 17, 128, device="cuda", dtype=torch.bfloat16)
+    modulation = torch.randn(2, 512, device=x.device, dtype=x.dtype)
+    scale, shift = modulation.chunk(4, dim=-1)[:2]
+    if not has_shift:
+        shift = None
+    compiled = torch.compile(fused_layernorm_modulate, fullgraph=True)
+    expected = reference(x, scale, shift, 1e-6)
+    actual = compiled(x, scale, shift, 1e-6)
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16))
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = fused_layernorm_modulate(x, scale, shift, 1e-6)
+    x.normal_()
+    modulation.normal_()
+    graph.replay()
+    expected = reference(x, scale, shift, 1e-6)
+    assert torch.equal(out.view(torch.int16), expected.view(torch.int16))
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

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

@@ -37,8 +37,10 @@ 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.qwen_image21 as qwen_image21
 import sglang.multimodal_gen.runtime.models.dits.sana as sana
 from sglang.kernels.ops.diffusion import (
+    BitExactFusionGate,
     can_use_fused_layernorm_modulate,
     can_use_fused_qk_head_layernorm,
     can_use_fused_rmsnorm_scale_shift,
@@ -1479,5 +1481,97 @@ def test_autoencoder_kl_fastpath_install():
     assert torch.equal(opt.decode(z), ref)
 
 
+@torch.no_grad()
+def test_qwen21_qk_norm_verifies_and_preserves_native_fallback(monkeypatch):
+    x = torch.randn(1, 257, 8, 128, device="cuda", dtype=torch.bfloat16)
+    norm = qwen_image21.RMSNorm(
+        128, 1e-6, cast_x_before_out_mul=True, force_native=True
+    ).to(device=x.device, dtype=x.dtype)
+    norm.weight.normal_()
+    expected = norm(x)
+    gate = BitExactFusionGate("test Q/K norm")
+    monkeypatch.setattr(qwen_image21, "_QK_NORM_FUSION", gate)
+    assert torch.equal(qwen_image21.apply_qk_norm(x, norm), expected)
+    assert gate.verified and not gate.disabled
+    x.normal_()
+    assert torch.equal(qwen_image21.apply_qk_norm(x, norm), norm(x))
+
+    gate = BitExactFusionGate("test mismatched Q/K norm")
+    monkeypatch.setattr(qwen_image21, "_QK_NORM_FUSION", gate)
+    monkeypatch.setattr(
+        qwen_image21,
+        "rmsnorm_preserve_reduction",
+        lambda x, weight, eps: torch.zeros_like(x),
+    )
+    assert torch.equal(qwen_image21.apply_qk_norm(x, norm), norm(x))
+    assert gate.disabled and not gate.verified
+
+
+@torch.no_grad()
+def test_qwen21_qk_norm_does_not_verify_during_capture(monkeypatch):
+    x = torch.randn(1, 17, 2, 128, device="cuda", dtype=torch.bfloat16)
+    norm = qwen_image21.RMSNorm(
+        128, 1e-6, cast_x_before_out_mul=True, force_native=True
+    ).to(device=x.device, dtype=x.dtype)
+    norm(x)
+    gate = BitExactFusionGate("test captured Q/K norm")
+    monkeypatch.setattr(qwen_image21, "_QK_NORM_FUSION", gate)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = qwen_image21.apply_qk_norm(x, norm)
+    assert not gate.verified and not gate.disabled
+    x.normal_()
+    graph.replay()
+    assert torch.equal(out, norm(x))
+
+
+@torch.no_grad()
+def test_qwen21_modulation_verifies_and_preserves_native_fallback(monkeypatch):
+    x = torch.randn(1, 257, 4096, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 1, 4096, device=x.device, dtype=x.dtype)
+    norm = torch.nn.LayerNorm(4096, eps=1e-6, elementwise_affine=False).cuda()
+    gate = BitExactFusionGate("test scale-only modulation")
+    monkeypatch.setattr(qwen_image21, "_MODULATION_FUSION", gate)
+    expected = norm(x) * (1 + scale)
+    actual = qwen_image21.apply_modulation(x, norm, scale)
+    assert torch.equal(actual.view(torch.int16), expected.view(torch.int16))
+    assert gate.verified and not gate.disabled
+    x.normal_()
+    scale.normal_()
+    assert torch.equal(
+        qwen_image21.apply_modulation(x, norm, scale), norm(x) * (1 + scale)
+    )
+
+    gate = BitExactFusionGate("test mismatched modulation")
+    monkeypatch.setattr(qwen_image21, "_MODULATION_FUSION", gate)
+    monkeypatch.setattr(
+        qwen_image21,
+        "fused_layernorm_modulate",
+        lambda x, scale, shift, eps: torch.zeros_like(x),
+    )
+    assert torch.equal(
+        qwen_image21.apply_modulation(x, norm, scale), norm(x) * (1 + scale)
+    )
+    assert gate.disabled and not gate.verified
+
+
+@torch.no_grad()
+def test_qwen21_modulation_does_not_verify_during_capture(monkeypatch):
+    x = torch.randn(1, 17, 128, device="cuda", dtype=torch.bfloat16)
+    scale = torch.randn(1, 1, 128, device=x.device, dtype=x.dtype)
+    norm = torch.nn.LayerNorm(128, eps=1e-6, elementwise_affine=False).cuda()
+    norm(x)
+    gate = BitExactFusionGate("test captured modulation")
+    monkeypatch.setattr(qwen_image21, "_MODULATION_FUSION", gate)
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = qwen_image21.apply_modulation(x, norm, scale)
+    assert not gate.verified and not gate.disabled
+    x.normal_()
+    scale.normal_()
+    graph.replay()
+    assert torch.equal(out, norm(x) * (1 + scale))
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v"]))

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

@@ -0,0 +1,74 @@
+# SPDX-License-Identifier: Apache-2.0
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_rmsnorm_preserve_reduction,
+    rmsnorm_preserve_reduction,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.version.hip is not None,
+    reason="NVIDIA CUDA required",
+)
+
+
+def reference(x, weight, eps):
+    value = x.float()
+    variance = value.pow(2).mean(dim=-1, keepdim=True)
+    return weight * (value * torch.rsqrt(variance + eps)).to(x.dtype)
+
+
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
+@pytest.mark.parametrize("shape", [(1, 128), (2, 17, 3, 128), (1, 4096, 32, 128)])
+@pytest.mark.parametrize("scale,eps", [(1e-4, 1e-6), (1.0, 1e-6), (100.0, 1e-5)])
+def test_preserves_native_reduction_and_rounding(dtype, shape, scale, eps):
+    torch.manual_seed(42)
+    x = (torch.randn(shape, device="cuda") * scale).to(dtype)
+    weight = torch.randn(shape[-1], device="cuda", dtype=dtype)
+    assert can_use_rmsnorm_preserve_reduction(x, weight)
+    actual = rmsnorm_preserve_reduction(x, weight, eps)
+    torch.testing.assert_close(actual, reference(x, weight, eps), atol=0, rtol=0)
+
+
+def test_layout_guards_and_offset():
+    x = torch.randn(259, 128, device="cuda", dtype=torch.bfloat16)[2:]
+    weight = torch.randn(128, device="cuda", dtype=x.dtype)
+    assert can_use_rmsnorm_preserve_reduction(x, weight)
+    torch.testing.assert_close(
+        rmsnorm_preserve_reduction(x, weight, 1e-6),
+        reference(x, weight, 1e-6),
+        atol=0,
+        rtol=0,
+    )
+    assert not can_use_rmsnorm_preserve_reduction(x.cpu(), weight.cpu())
+    assert not can_use_rmsnorm_preserve_reduction(x.float(), weight.float())
+    assert not can_use_rmsnorm_preserve_reduction(x[:, ::2], weight[::2])
+    assert not can_use_rmsnorm_preserve_reduction(x, weight.float())
+    assert not can_use_rmsnorm_preserve_reduction(x, weight[:-1])
+    assert not can_use_rmsnorm_preserve_reduction(x[:0], weight)
+
+
+def test_compile_and_graph_replay():
+    x = torch.randn(257, 128, device="cuda", dtype=torch.bfloat16)
+    weight = torch.randn(128, device="cuda", dtype=x.dtype)
+    compiled = torch.compile(rmsnorm_preserve_reduction, fullgraph=True)
+    torch.testing.assert_close(
+        compiled(x, weight, 1e-6), reference(x, weight, 1e-6), atol=0, rtol=0
+    )
+    graph = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(graph):
+        out = rmsnorm_preserve_reduction(x, weight, 1e-6)
+    x.normal_()
+    weight.normal_()
+    graph.replay()
+    torch.testing.assert_close(out, reference(x, weight, 1e-6), atol=0, rtol=0)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))