diff --git a/python/sglang/kernels/ops/diffusion/rope/hunyuan_qkv_pack_triton.py b/python/sglang/kernels/ops/diffusion/rope/hunyuan_qkv_pack_triton.py index 0561dac95..4a0859e00 100644 --- a/python/sglang/kernels/ops/diffusion/rope/hunyuan_qkv_pack_triton.py +++ b/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 diff --git a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py b/test/registered/kernels/ops/diffusion/test_model_fast_paths.py index f971997bf..0b4702fca 100644 --- a/test/registered/kernels/ops/diffusion/test_model_fast_paths.py +++ b/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()