[diffusion] model: support VDN-H3 with a hybrid_window_attn_h3 backend (#37903)

Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
Co-authored-by: Haocheng Xi <xihc@berkeley.edu>
Co-authored-by: Mick <mickjagger19@icloud.com>

test/registered/kernel/diffusion/benchmark/bench_vdn_delta_factors.py

@@ -0,0 +1,51 @@
+"""``vdn_delta_factors`` (fused inverse + products) vs the eager Cholesky chain it replaces."""
+
+import torch
+
+from sglang.kernels.jit.benchmark import marker
+from sglang.kernels.ops.diffusion import vdn_delta_factors
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(
+    est_time=10, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
+)
+
+HEAD_DIM = 128
+
+
+def eager_delta_factors(A: torch.Tensor, B: torch.Tensor, alpha: torch.Tensor):
+    eye = torch.eye(HEAD_DIM, device=A.device, dtype=torch.float32).expand_as(A)
+    chol = torch.linalg.cholesky(A + eye)
+    linv = torch.linalg.solve_triangular(chol, eye, upper=False, left=True)
+    inv = linv.transpose(-1, -2) @ linv
+    return alpha.unsqueeze(-1) * inv, B @ inv
+
+
+FN_MAP = {"jit": vdn_delta_factors, "eager": eager_delta_factors}
+
+
+def _inputs(num: int):
+    g = torch.Generator(device="cuda").manual_seed(0)
+    k = torch.nn.functional.normalize(
+        torch.randn(num, 1008, HEAD_DIM, device="cuda", generator=g), dim=-1
+    )
+    v = torch.randn(num, 1008, HEAD_DIM, device="cuda", generator=g)
+    beta = torch.sigmoid(torch.randn(num, 1008, device="cuda", generator=g))
+    A = (k * beta.unsqueeze(-1)).transpose(-1, -2) @ k
+    A = 0.5 * (A + A.transpose(-1, -2))
+    B = (v * beta.unsqueeze(-1)).transpose(-1, -2) @ k
+    alpha = torch.rand(num, HEAD_DIM, device="cuda", generator=g)
+    return A.contiguous(), B.contiguous(), alpha.contiguous()
+
+
+# 707 = 101 frames x 7 heads: the paper workload per rank (8 x B200, Ulysses 8)
+@marker.parametrize("num_matrices", [64, 707], [64])
+@marker.benchmark("impl", ["jit", "eager"], unit="us")
+def benchmark(num_matrices: int, impl: str):
+    A, B, alpha = _inputs(num_matrices)
+    # both eager: graph replay on one side only is not a like-for-like comparison
+    return marker.do_bench(FN_MAP[impl], input_args=(A, B, alpha), use_cuda_graph=False)
+
+
+if __name__ == "__main__":
+    benchmark.run()

test/registered/kernel/diffusion/test_mxfp8_swizzled.py

@@ -0,0 +1,78 @@
+"""MXFP8 producers against ``flashinfer.mxfp8_quantize`` of the bf16 tensor the
+unfused kernel stores: payload and the swizzled E8M0 scale buffer, padding
+included, byte for byte."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_mxfp8_swizzled,
+    can_use_silu_mul_mxfp8,
+    indexed_scale_shift_bf16_,
+    indexed_scale_shift_mxfp8_,
+    mxfp8_quantize_swizzled,
+    silu_mul_mxfp8,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+ROWS, HIDDEN = 333, 5376  # rows and scale columns both need padding
+
+
+def _assert_matches_flashinfer(
+    got: tuple[torch.Tensor, torch.Tensor], bf16: torch.Tensor
+) -> None:
+    import flashinfer
+
+    q, s = flashinfer.mxfp8_quantize(bf16.contiguous(), True)
+    assert torch.equal(got[0].view(torch.uint8), q.view(torch.uint8))
+    assert torch.equal(got[1].view(torch.uint8), s.view(torch.uint8))
+
+
+def test_quantize_swizzled_is_byte_exact() -> None:
+    g = torch.Generator(device="cuda").manual_seed(1)
+    x = (torch.randn((ROWS, HIDDEN), device="cuda", generator=g) * 3).to(torch.bfloat16)
+    x[0, :32] = 0  # an all-zero block takes the minimum exponent
+    assert can_use_mxfp8_swizzled(x)
+    _assert_matches_flashinfer(mxfp8_quantize_swizzled(x), x)
+
+
+def test_silu_mul_mxfp8_is_byte_exact() -> None:
+    g = torch.Generator(device="cuda").manual_seed(2)
+    x = (torch.randn((ROWS, 2 * HIDDEN), device="cuda", generator=g) * 2).to(
+        torch.bfloat16
+    )
+    assert can_use_silu_mul_mxfp8(x)
+    ref = torch.nn.functional.silu(x[:, :HIDDEN]) * x[:, HIDDEN:]
+    _assert_matches_flashinfer(silu_mul_mxfp8(x), ref)
+
+
+@pytest.mark.parametrize("keep_bf16", [True, False])
+def test_indexed_scale_shift_mxfp8_is_byte_exact(keep_bf16: bool) -> None:
+    g = torch.Generator(device="cuda").manual_seed(3)
+    x = torch.randn((ROWS, HIDDEN), device="cuda", generator=g).to(torch.bfloat16)
+    shift = torch.randn((3, HIDDEN), device="cuda", generator=g).to(torch.bfloat16)
+    scale = torch.randn((3, HIDDEN), device="cuda", generator=g).to(torch.bfloat16)
+    indices = torch.randint(0, 3, (ROWS,), device="cuda", generator=g)
+    ref = indexed_scale_shift_bf16_(x.clone(), shift, scale, indices)
+    kept, q, s = indexed_scale_shift_mxfp8_(
+        x, shift, scale, indices, keep_bf16=keep_bf16
+    )
+    _assert_matches_flashinfer((q, s), ref)
+    assert (kept is x and torch.equal(x, ref)) if keep_bf16 else kept is None
+
+
+def test_predicates_reject_unsupported_input() -> None:
+    assert not can_use_mxfp8_swizzled(
+        torch.randn(4, 64, device="cuda", dtype=torch.float16)
+    )
+    assert not can_use_silu_mul_mxfp8(
+        torch.randn(4, 96, device="cuda", dtype=torch.bfloat16)
+    )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernel/diffusion/test_qknorm_rope_out_of_place.py

@@ -0,0 +1,48 @@
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_fused_inplace_qknorm_rope,
+    fused_inplace_qknorm_rope,
+    fused_qknorm_rope_out_of_place,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=20, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+
+def test_out_of_place_qknorm_rope_matches_inplace_and_keeps_inputs() -> None:
+    """The out-of-place variant (strided fused-qkv views in, contiguous copies
+    out) is bit-equal to the in-place kernel and leaves its inputs untouched;
+    VDN-H3's linear branch reads the raw q/k after it."""
+    T, H, D, R = 512, 4, 128, 96
+    if not can_use_fused_inplace_qknorm_rope(
+        D, R, True, torch.bfloat16, torch.bfloat16, True
+    ):
+        pytest.skip("fused qknorm+rope JIT kernel unavailable")
+    g = torch.Generator(device="cpu").manual_seed(0)
+    qkv = torch.randn(T, 3 * H * D, generator=g).to("cuda", torch.bfloat16)
+    q = qkv[:, : H * D].view(T, H, D)
+    k = qkv[:, H * D : 2 * H * D].view(T, H, D)
+    qw = (torch.rand(D, generator=g) + 0.5).to("cuda", torch.bfloat16)
+    kw = (torch.rand(D, generator=g) + 0.5).to("cuda", torch.bfloat16)
+    freqs = torch.randn(T, R // 2, generator=g).to("cuda")
+    cache = torch.cat((freqs.cos(), freqs.sin()), -1).to(torch.bfloat16).contiguous()
+    pos = torch.arange(T, device="cuda")
+    kwargs = dict(
+        is_neox=True, eps=1e-5, head_dim=D, rope_dim=R, round_norm_before_rope=True
+    )
+    q_ref, k_ref = q.clone(), k.clone()
+    fused_inplace_qknorm_rope(q_ref, k_ref, qw, kw, cache, pos, **kwargs)
+    q_out = torch.empty(T, H, D, device="cuda", dtype=torch.bfloat16)
+    k_out = torch.empty_like(q_out)
+    before = qkv.clone()
+    fused_qknorm_rope_out_of_place(q, k, q_out, k_out, qw, kw, cache, pos, **kwargs)
+    assert torch.equal(qkv, before)
+    assert torch.equal(q_out, q_ref) and torch.equal(k_out, k_ref)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))

test/registered/kernel/diffusion/test_vdn_delta_factors.py

@@ -0,0 +1,134 @@
+"""Fused VDN-H3 delta-rule factors against the eager Cholesky chain: both fp32 paths are
+held to the same cond(I + A)-dominated error band vs fp64, on model-shaped inputs."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import can_use_vdn_delta_factors, vdn_delta_factors
+from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as vdn
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=60, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+HEAD_DIM = 128
+
+
+def _inputs(
+    frames: int, heads: int, tokens: int, beta_scale: float = 1.0, seed: int = 0
+):
+    g = torch.Generator(device="cuda").manual_seed(seed)
+    k = torch.nn.functional.normalize(
+        torch.nn.functional.silu(
+            torch.randn(frames, heads, tokens, HEAD_DIM, device="cuda", generator=g)
+        ),
+        dim=-1,
+    )
+    v = torch.nn.functional.silu(
+        torch.randn(frames, heads, tokens, HEAD_DIM, device="cuda", generator=g)
+    )
+    beta = (
+        torch.sigmoid(torch.randn(frames, heads, tokens, device="cuda", generator=g))
+        * beta_scale
+    )
+    alpha = torch.rand(frames, heads, HEAD_DIM, device="cuda", generator=g) * 0.9 + 0.1
+    alpha[0] = 1.0
+    A = (k * beta.unsqueeze(-1)).transpose(-1, -2) @ k
+    A = 0.5 * (A + A.transpose(-1, -2))
+    B = (v * beta.unsqueeze(-1)).transpose(-1, -2) @ k
+    return A.float().contiguous(), B.float().contiguous(), alpha.float().contiguous()
+
+
+def _rel(x: torch.Tensor, ref: torch.Tensor) -> float:
+    return ((x.double() - ref).norm() / ref.norm()).item()
+
+
+def _fp64(A, B, alpha):
+    inv = torch.linalg.inv(
+        torch.eye(HEAD_DIM, device=A.device, dtype=torch.float64) + A.double()
+    )
+    return alpha.double().unsqueeze(-1) * inv, B.double() @ inv
+
+
+@pytest.mark.parametrize("frames,heads,tokens", [(1, 1, 16), (5, 3, 64), (11, 7, 1008)])
+@pytest.mark.parametrize("beta_scale", [1.0, 50.0])
+def test_matches_eager_and_fp64(frames, heads, tokens, beta_scale):
+    A, B, alpha = _inputs(frames, heads, tokens, beta_scale)
+    assert can_use_vdn_delta_factors(A, B, alpha)
+    t_ref, j_ref = _fp64(A, B, alpha)
+    t_eager, j_eager = vdn.delta_factor_apply(
+        "vdn_solve", alpha, A, B, tokens_per_frame=tokens
+    )
+    t_fused, j_fused = vdn_delta_factors(A, B, alpha)
+    assert t_fused.shape == A.shape and j_fused.shape == B.shape
+    assert t_fused.dtype is torch.float32 and j_fused.dtype is torch.float32
+    assert torch.isfinite(t_fused).all() and torch.isfinite(j_fused).all()
+    # same accuracy class as the Cholesky chain (both fp32, cond-dominated)
+    assert _rel(t_fused, t_ref) <= 1.5 * _rel(t_eager, t_ref) + 1e-7
+    assert _rel(j_fused, j_ref) <= 1.5 * _rel(j_eager, j_ref) + 1e-7
+    assert _rel(t_fused, t_ref) < 1e-5 and _rel(j_fused, j_ref) < 3e-5
+    # elementwise the two fp32 paths differ by a few 1e-5 on ill-conditioned inputs (cancellation)
+    torch.testing.assert_close(t_fused, t_eager, rtol=1e-4, atol=1e-5)
+    torch.testing.assert_close(j_fused, j_eager, rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("rule", ["vdn_solve", "vdn_scaled"])
+def test_delta_factor_apply_fused_path(rule):
+    A, B, alpha = _inputs(4, 2, 48)
+    eager = vdn.delta_factor_apply(rule, alpha, A, B, tokens_per_frame=48, fused=False)
+    fused = vdn.delta_factor_apply(rule, alpha, A, B, tokens_per_frame=48, fused=True)
+    for x, y in zip(fused, eager):
+        torch.testing.assert_close(x, y, rtol=2e-5, atol=2e-5)
+
+
+def test_sana_scaled_ignores_fused():
+    A, B, alpha = _inputs(2, 2, 32)
+    eager = vdn.delta_factor_apply(
+        "sana_scaled", alpha, A, B, tokens_per_frame=32, fused=False
+    )
+    fused = vdn.delta_factor_apply(
+        "sana_scaled", alpha, A, B, tokens_per_frame=32, fused=True
+    )
+    for x, y in zip(fused, eager):
+        assert torch.equal(x, y)
+
+
+def _storage_offset_copy(t: torch.Tensor) -> torch.Tensor:
+    # contiguous, but one element past a 16-byte boundary
+    flat = torch.empty(t.numel() + 1, dtype=t.dtype, device=t.device)
+    out = flat[1:].view(t.shape)
+    out.copy_(t)
+    assert out.is_contiguous() and out.data_ptr() % 16 != 0
+    return out
+
+
+@pytest.mark.parametrize("which", ["A", "B", "alpha"])
+def test_storage_offset_input_matches_aligned(which):
+    """A contiguous input with a storage offset must not fault in the float4 loads."""
+    A, B, alpha = _inputs(3, 2, 64)
+    ref = vdn_delta_factors(A, B, alpha)
+    inputs = {"A": A, "B": B, "alpha": alpha}
+    inputs[which] = _storage_offset_copy(inputs[which])
+    assert can_use_vdn_delta_factors(inputs["A"], inputs["B"], inputs["alpha"])
+    out = vdn_delta_factors(inputs["A"], inputs["B"], inputs["alpha"])
+    for got, want in zip(out, ref):
+        assert torch.equal(got, want)
+
+
+def test_can_use_rejects_unsupported():
+    A, B, alpha = _inputs(2, 2, 32)
+    assert can_use_vdn_delta_factors(A, B, alpha)
+    assert not can_use_vdn_delta_factors(
+        A[..., :64, :64].contiguous(),
+        B[..., :64, :64].contiguous(),
+        alpha[..., :64].contiguous(),
+    )
+    assert not can_use_vdn_delta_factors(A.bfloat16(), B, alpha)
+    assert not can_use_vdn_delta_factors(A.transpose(-1, -2), B, alpha)
+    assert not can_use_vdn_delta_factors(A, B, alpha[..., :1].expand_as(alpha))
+    assert not can_use_vdn_delta_factors(A.cpu(), B.cpu(), alpha.cpu())
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v", "-s"]))

test/registered/kernel/diffusion/test_vdn_linear_branch.py

@@ -0,0 +1,139 @@
+"""VDN-H3 linear-branch kernels against the eager chains they replace: the data
+movers bit-exact, the three activation kernels within one bf16 ulp."""
+
+import sys
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion import (
+    can_use_vdn_frame_stats_prep,
+    can_use_vdn_gather_linear_state,
+    can_use_vdn_linear_epilogue,
+    can_use_vdn_silu_l2norm,
+    can_use_vdn_temporal_conv_act,
+    vdn_frame_stats_prep,
+    vdn_linear_epilogue,
+    vdn_silu_l2norm,
+    vdn_temporal_conv_act,
+)
+from sglang.multimodal_gen.configs.models.dits.minimax_h3_vdn import (
+    VDNHybridAttentionArchConfig,
+)
+from sglang.multimodal_gen.runtime.models.dits import minimax_h3_vdn as vdn
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200")
+
+FRAMES, TOKENS, HEADS, HEAD_DIM = 6, 24, 3, 32
+
+
+def _ulp_close(got: torch.Tensor, ref: torch.Tensor) -> bool:
+    scale = max(1.0, ref.float().abs().max().item())
+    return (got.float() - ref.float()).abs().max().item() <= 2e-2 * scale
+
+
+def test_temporal_conv_act_matches_eager_chain() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    x = torch.randn(FRAMES, TOKENS, HEADS * HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    w = (torch.randn(HEADS * HEAD_DIM, 5, generator=g) * 0.4).to("cuda", torch.bfloat16)
+    assert can_use_vdn_temporal_conv_act(x, HEADS, HEAD_DIM)
+    ref = vdn._activate(vdn._temporal_shift(x, w).reshape(-1, HEADS, HEAD_DIM), True)
+    assert _ulp_close(vdn_temporal_conv_act(x, w, HEADS, HEAD_DIM, True), ref)
+    frame_major = vdn_temporal_conv_act(x, w, HEADS, HEAD_DIM, True, frame_major=True)
+    assert (
+        frame_major.shape == (FRAMES, HEADS, TOKENS, HEAD_DIM)
+        and frame_major.is_contiguous()
+    )
+    assert torch.equal(
+        frame_major, ref.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    ) or _ulp_close(
+        frame_major, ref.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    )
+
+
+def test_silu_l2norm_reads_strided_qkv_views() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    tokens = torch.randn(FRAMES * TOKENS, 3 * HEADS * HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    strided = tokens[:, : HEADS * HEAD_DIM].view(FRAMES * TOKENS, HEADS, HEAD_DIM)
+    assert can_use_vdn_silu_l2norm(strided)
+    got = vdn_silu_l2norm(strided, True)
+    assert got.is_contiguous() and _ulp_close(got, vdn._activate(strided, True))
+    got_v = vdn_silu_l2norm(strided, False)
+    assert _ulp_close(got_v, torch.nn.functional.silu(strided))
+    frame_major = vdn_silu_l2norm(strided, True, per_frame=TOKENS)
+    assert frame_major.shape == (FRAMES, HEADS, TOKENS, HEAD_DIM)
+    assert torch.equal(
+        got.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3), frame_major
+    )
+    with pytest.raises(ValueError):
+        vdn_silu_l2norm(strided, True, per_frame=TOKENS + 1)
+
+
+def test_frame_stats_prep_is_bit_exact() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    key = torch.randn(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    value = torch.randn(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    beta = torch.rand(FRAMES * TOKENS, HEADS, generator=g).to("cuda", torch.bfloat16)
+    assert can_use_vdn_frame_stats_prep(key, value)
+    k16, k32, kb32, vb = vdn_frame_stats_prep(key, value, beta, FRAMES, TOKENS)
+    kf = key.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    vf = value.view(FRAMES, TOKENS, HEADS, HEAD_DIM).permute(0, 2, 1, 3)
+    bf = beta.view(FRAMES, TOKENS, HEADS).permute(0, 2, 1)
+    assert torch.equal(k16, kf.contiguous())
+    assert torch.equal(k32, kf.float().contiguous())
+    assert torch.equal(kb32, (kf.float() * bf.unsqueeze(-1).float()).contiguous())
+    assert torch.equal(vb, (vf * bf.unsqueeze(-1).to(vf.dtype)).contiguous())
+
+
+def test_linear_epilogue_matches_eager_chain() -> None:
+    g = torch.Generator(device="cpu").manual_seed(4)
+    readout = torch.randn(FRAMES, HEADS, TOKENS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    weight = (1 + 0.1 * torch.randn(HEAD_DIM, generator=g)).to("cuda", torch.bfloat16)
+    gate = torch.rand(FRAMES * TOKENS, HEADS, HEAD_DIM, generator=g).to(
+        "cuda", torch.bfloat16
+    )
+    assert can_use_vdn_linear_epilogue(readout)
+    got = vdn_linear_epilogue(readout, weight, gate, 1e-6)
+    assert _ulp_close(got, vdn.linear_epilogue(readout, weight, gate, 1e-6))
+
+
+@pytest.mark.parametrize("bridge", ["alpha", "none"])
+@pytest.mark.parametrize("with_text_state", [False, True])
+def test_gather_linear_state_matches_eager(bridge: str, with_text_state: bool) -> None:
+    g = torch.Generator(device="cpu").manual_seed(5)
+    frames, heads, dim = 9, 2, 32
+    hybrid = VDNHybridAttentionArchConfig(chunk=3, radius=1, anchor_frames="none")
+    bounds = hybrid.window_bounds(frames)
+    prefix = torch.randn(frames, heads, dim, dim, generator=g).cuda()
+    suffix = torch.randn(frames, heads, dim, dim, generator=g).cuda()
+    alpha = (torch.rand(frames, heads, dim, generator=g) * 0.5 + 0.5).cuda()
+    text = torch.randn(heads, dim, dim, generator=g).cuda() if with_text_state else None
+    assert can_use_vdn_gather_linear_state(prefix)
+    kwargs = dict(bridge=bridge, text_state=text, out_dtype=torch.float32)
+    ref = vdn.gather_linear_state(prefix, suffix, alpha, bounds, fused=False, **kwargs)
+    got = vdn.gather_linear_state(prefix, suffix, alpha, bounds, **kwargs)
+    torch.testing.assert_close(got, ref, atol=1e-5, rtol=1e-5)
+
+
+def test_predicates_reject_unsupported_inputs() -> None:
+    fp16 = torch.randn(4, 2, 32, device="cuda", dtype=torch.float16)
+    assert not can_use_vdn_silu_l2norm(fp16)
+    odd = torch.randn(4, 2, 48, device="cuda", dtype=torch.bfloat16)
+    assert not can_use_vdn_silu_l2norm(odd)
+    with pytest.raises(ValueError):
+        vdn_silu_l2norm(odd, True)
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))