[diffusion] Fix Hunyuan QKV pack indexing at production video shapes (#36009)

python/sglang/kernels/ops/diffusion/rope/hunyuan_qkv_pack_triton.py

@@ -52,7 +52,10 @@ def _hunyuan_qkv_rope_pack_kernel(
     BLOCK_HEADS: tl.constexpr,
     BLOCK_HALF: tl.constexpr,
 ):
-    token = tl.program_id(0)
+    # Merged Hunyuan projections can have row offsets above INT32_MAX at
+    # production video shapes (for example, 115200 * 21504). Keep all row
+    # address arithmetic in int64, as the other diffusion layout kernels do.
+    token = tl.program_id(0).to(tl.int64)
     head_block = tl.program_id(1)
     total_tokens = img_tokens + txt_tokens
     batch = token // total_tokens

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

@@ -43,6 +43,7 @@ from sglang.kernels.ops.diffusion import (
     can_use_fused_rmsnorm_scale_shift,
     can_use_wan_rmsnorm_silu,
     fused_ltx2_rms_norm_modulate,
+    hunyuan_qkv_rope_pack,
     mark_fused_ln_modulate_site,
     mark_hunyuan_qknorm_site,
     mark_ltx2_rms_norm_modulate_site,
@@ -641,6 +642,47 @@ def test_hunyuan_qkv_rope_pack_is_bit_exact(img_tokens, txt_tokens):
     assert torch.equal(v, v_ref)
 
 
+def test_hunyuan_qkv_rope_pack_uses_int64_row_offsets():
+    if torch.cuda.get_device_properties(0).total_memory < 16 * 2**30:
+        pytest.skip("needs >= 16 GB GPU memory")
+
+    img_tokens, txt_tokens = 115200, 8
+    num_heads, head_dim = 24, 128
+    total_tokens = img_tokens + txt_tokens
+    projection_width = 21504
+
+    projection = torch.zeros(
+        (1, total_tokens, projection_width),
+        device="cuda",
+        dtype=torch.bfloat16,
+    )
+    qkv = projection[..., : 3 * num_heads * head_dim].view(
+        1, total_tokens, 3, num_heads, head_dim
+    )
+    q = qkv[:, :, 0].contiguous()
+    k = qkv[:, :, 1].contiguous()
+    v = qkv[:, :, 2]
+    assert (img_tokens - 1) * v.stride(1) > torch.iinfo(torch.int32).max
+
+    cos = torch.ones((img_tokens, head_dim // 2), device="cuda")
+    sin = torch.zeros_like(cos)
+    packed = hunyuan_qkv_rope_pack(
+        q[:, :img_tokens],
+        k[:, :img_tokens],
+        v[:, :img_tokens],
+        q[:, img_tokens:],
+        k[:, img_tokens:],
+        v[:, img_tokens:],
+        cos,
+        sin,
+    )
+    torch.cuda.synchronize()
+
+    expected_shape = (1, total_tokens, num_heads, head_dim)
+    assert all(x.shape == expected_shape for x in packed)
+    assert all(x[0, img_tokens - 1, -1, -1].item() == 0 for x in packed)
+
+
 def test_hunyuan_quality_qknorm_matches_rmsnorm():
     torch.manual_seed(1)
     site = torch.nn.Module()