[diffusion] model: support minimax-h3 (#33275)

Co-authored-by: zhenaozhenfu <zhenaozhenfu@minimaxi.com>
Co-authored-by: BBuf <1182563586@qq.com>
Co-authored-by: andyluo7 <andy.luo@amd.com>
Co-authored-by: Zijie Xia <zijie_xia@icloud.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: chao-xue <877184285@qq.com>
Co-authored-by: Cursor <cursoragent@cursor.com>

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

@@ -122,7 +122,7 @@ def test_qknorm_rope(
         if is_neox:
             elems_per_thread = head_dim // 32
             rotary_lanes = rope_dim // elems_per_thread
-            if rotary_lanes < 2 or rotary_lanes & (rotary_lanes - 1):
+            if rotary_lanes < 2 or rotary_lanes % 2:
                 continue
 
         q = torch.randn(batch_size, num_heads, head_dim, device=DEVICE, dtype=DTYPE)
@@ -150,5 +150,60 @@ def test_qknorm_rope(
         triton.testing.assert_close(k_ref, k_fused, atol=ATOL, rtol=RTOL)
 
 
+def test_qknorm_rope_preserves_split_bf16_rounding() -> None:
+    from sgl_kernel import rotary_embedding
+
+    from sglang.kernels.ops.diffusion.qknorm_rope import (
+        fused_inplace_qknorm_rope,
+    )
+    from sglang.kernels.ops.layernorm.norm import fused_inplace_qknorm
+
+    num_tokens, num_heads, head_dim, rope_dim = 257, 28, 128, 96
+    inner_dim = num_heads * head_dim
+    qkv = torch.randn(
+        num_tokens,
+        3 * inner_dim,
+        device=DEVICE,
+        dtype=DTYPE,
+    )
+    q_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    k_weight = torch.randn(head_dim, device=DEVICE, dtype=DTYPE)
+    positions = torch.arange(num_tokens, device=DEVICE, dtype=torch.int64)
+    cos_sin_cache = create_cos_sin_cache(rope_dim, num_tokens).to(DTYPE)
+
+    qkv_ref, qkv_fused = qkv.clone(), qkv.clone()
+    q_ref, k_ref, _ = qkv_ref.split(inner_dim, dim=-1)
+    q_fused, k_fused, _ = qkv_fused.split(inner_dim, dim=-1)
+    q_ref = q_ref.view(num_tokens, num_heads, head_dim)
+    k_ref = k_ref.view(num_tokens, num_heads, head_dim)
+    q_fused = q_fused.view(num_tokens, num_heads, head_dim)
+    k_fused = k_fused.view(num_tokens, num_heads, head_dim)
+
+    fused_inplace_qknorm(q_ref, k_ref, q_weight, k_weight, eps=1e-5)
+    rotary_embedding(
+        positions,
+        q_ref.view(num_tokens, -1),
+        k_ref.view(num_tokens, -1),
+        head_dim,
+        cos_sin_cache,
+        True,
+    )
+    fused_inplace_qknorm_rope(
+        q_fused,
+        k_fused,
+        q_weight,
+        k_weight,
+        cos_sin_cache,
+        positions,
+        is_neox=True,
+        eps=1e-5,
+        rope_dim=rope_dim,
+        round_norm_before_rope=True,
+    )
+
+    assert torch.equal(q_ref, q_fused)
+    assert torch.equal(k_ref, k_fused)
+
+
 if __name__ == "__main__":
     sys.exit(pytest.main([__file__, "-v", "-s"]))

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

@@ -0,0 +1,69 @@
+"""Bitwise tests for the generic Ulysses output head-merge fast path."""
+
+import sys
+from unittest.mock import patch
+
+import pytest
+import torch
+
+from sglang.kernels.ops.diffusion.usp_relayout import (
+    can_use_usp_merge_heads,
+    usp_merge_heads,
+)
+from sglang.test.ci.ci_register import register_cuda_ci
+
+register_cuda_ci(est_time=4, stage="base-b-kernel-unit", runner_config="1-gpu-large")
+
+DEVICE = "cuda"
+
+
+@pytest.mark.parametrize(
+    "world,seq,batch,h_local,head_dim",
+    [
+        (4, 7936, 1, 14, 128),  # H3 768p production shape (Ulysses 4)
+        (2, 64, 3, 4, 64),  # batched
+        (4, 33, 2, 4, 100),  # scalar fallback inside the CUDA kernel
+    ],
+)
+@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32])
+def test_usp_merge_heads_bitwise(dtype, world, seq, batch, h_local, head_dim):
+    generator = torch.Generator(device=DEVICE).manual_seed(4321)
+    x = torch.randn(
+        world,
+        seq,
+        batch,
+        h_local,
+        head_dim,
+        dtype=dtype,
+        device=DEVICE,
+        generator=generator,
+    )
+    assert can_use_usp_merge_heads(x)
+    out = usp_merge_heads(x)
+    ref = x.permute(2, 1, 0, 3, 4).contiguous()
+    assert out.shape == ref.shape
+    assert torch.equal(out, ref)
+
+
+def test_usp_merge_heads_unsupported_inputs_use_exact_fallback():
+    x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
+    unsupported = [x.transpose(0, 1), x[:0]]
+
+    for value in unsupported:
+        assert not can_use_usp_merge_heads(value)
+        assert torch.equal(
+            usp_merge_heads(value), value.permute(2, 1, 0, 3, 4).contiguous()
+        )
+
+    with patch.object(torch.version, "hip", "6.3"):
+        assert not can_use_usp_merge_heads(x)
+        assert torch.equal(usp_merge_heads(x), x.permute(2, 1, 0, 3, 4).contiguous())
+
+
+def test_usp_merge_heads_fast_path_rejects_wrong_rank():
+    x = torch.randn(2, 4, 1, 4, 64, dtype=torch.bfloat16, device=DEVICE)
+    assert not can_use_usp_merge_heads(x[0])
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))