[DSV4] Use int64 for compressor out_loc tensors (#27973)

2026-06-11 17:45:34 -07:00
parent 97a0031799
commit 2e74ff192c
4 changed files with 19 additions and 22 deletions
@@ -148,7 +148,7 @@ def test_fp4_fused_norm_rope_store_layout(num_tokens: int) -> None:
    plan = CompressorDecodePlan.generate_legacy(
        compress_ratio, req_pool_indices, seq_lens
    )
-    loc = torch.arange(num_tokens, device="cuda", dtype=torch.int32)
+    loc = torch.arange(num_tokens, device="cuda", dtype=torch.int64)
    freqs_cis = precompute_freqs_cis(
        64, int(seq_lens.max().item()) + 1, 0, 10000, 1, 32, 1
    ).to("cuda")