[Bugfix][MiniMax-M3] Fix token-major top-k buffer handling in Triton … (#49149)
Signed-off-by: rongfu.leng <lenronfu@gmail.com>
rongfu.leng committed
Jul 25, 2026 at 21:45 UTC
d1a8ba63d9d2bb51ebf60dd5ea1463cf61c70cea
2 files changed
+13
-3
vllm/models/minimax_m3/common/indexer.py
+5
-2
@@ -424,6 +424,9 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
424
# (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the
425
# kernels' out= writes out[:, :total_q]. None -> allocate fresh.
426
buf = self.topk_indices_buffer
427
+ buf_htk = (
428
+ buf if buf is None or current_platform.is_rocm() else buf.transpose(0, 1)
429
+ )
430
decode_topk: torch.Tensor | None = None
431
prefill_topk: torch.Tensor | None = None
432
if index_md.num_decodes > 0:
@@ -441,7 +444,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
444
self.num_kv_heads,
445
d.decode_query_len,
446
d.max_decode_query_len,
444
- out=buf,
447
+ out=buf_htk,
448
)
449
if index_md.num_prefills > 0:
450
p = index_md.prefill
@@ -465,7 +468,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl):
468
self.topk_blocks,
469
self.init_blocks,
470
self.local_blocks,
468
- out=buf[:, nd:, :] if buf is not None else None,
471
+ out=buf_htk[:, nd:, :] if buf_htk is not None else None,
472
)
473
return decode_topk, prefill_topk
474
vllm/models/minimax_m3/common/sparse_attention.py
+8
-1
@@ -384,7 +384,14 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl):
384
nd = main_md.num_decode_tokens
385
num_tokens = main_md.num_actual_tokens
386
# Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].
387
- topk = layer.topk_indices_buffer # type: ignore[attr-defined]
387
+ topk_buffer = layer.topk_indices_buffer # type: ignore[attr-defined]
388
+ assert topk_buffer is not None
389
+
390
+ topk = (
391
+ topk_buffer
392
+ if current_platform.is_rocm()
393
+ else topk_buffer[:num_tokens].transpose(0, 1)
394
+ )
395
assert topk is not None
396
hd = self.head_size
397
q = query[:num_tokens].view(-1, self.num_heads, hd)