[ROCm][Performance] Avoid blocking MiniMax M3 scalar upload (#56170)
Signed-off-by: fai <fangzhouai@gmail.com> Co-authored-by: OpenAI Codex <noreply@openai.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Fangzhou Ai committed
Sep 13, 2026 at 12:56 UTC
f09c52a587f205963d3a2b8f1c42820a831ee4b0
1 file changed
+4
-1
vllm/models/minimax_m3/common/sparse_attention.py
+4
-1
@@ -314,7 +314,10 @@ class MiniMaxM3SparseMetadataBuilder(AttentionMetadataBuilder[MiniMaxM3SparseMet
314
prefill_cu_seqlens_k = torch.empty(
315
num_prefills + 1, dtype=torch.int32, device=seq_lens.device
316
)
317
- prefill_cu_seqlens_k[0] = 0
317
+ if current_platform.is_rocm():
318
+ prefill_cu_seqlens_k[:1].zero_()
319
+ else:
320
+ prefill_cu_seqlens_k[0] = 0
321
torch.cumsum(prefill_kv_lens, dim=0, out=prefill_cu_seqlens_k[1:])
322
prefill_cu_seqlens_q = (
323
query_start_loc[num_decodes:] - num_decode_tokens