diff --git a/vllm/models/minimax_m3/common/indexer.py b/vllm/models/minimax_m3/common/indexer.py index c66e7ce5267..fea22dbdacd 100644 --- a/vllm/models/minimax_m3/common/indexer.py +++ b/vllm/models/minimax_m3/common/indexer.py @@ -424,6 +424,9 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): # (decode at [:, :nd], prefill at [:, nd:]) and return views into it; the # kernels' out= writes out[:, :total_q]. None -> allocate fresh. buf = self.topk_indices_buffer + buf_htk = ( + buf if buf is None or current_platform.is_rocm() else buf.transpose(0, 1) + ) decode_topk: torch.Tensor | None = None prefill_topk: torch.Tensor | None = None if index_md.num_decodes > 0: @@ -441,7 +444,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.num_kv_heads, d.decode_query_len, d.max_decode_query_len, - out=buf, + out=buf_htk, ) if index_md.num_prefills > 0: p = index_md.prefill @@ -465,7 +468,7 @@ class MiniMaxM3IndexerTritonImpl(MiniMaxM3IndexerImpl): self.topk_blocks, self.init_blocks, self.local_blocks, - out=buf[:, nd:, :] if buf is not None else None, + out=buf_htk[:, nd:, :] if buf_htk is not None else None, ) return decode_topk, prefill_topk diff --git a/vllm/models/minimax_m3/common/sparse_attention.py b/vllm/models/minimax_m3/common/sparse_attention.py index f887f14b643..04aaef7c50f 100644 --- a/vllm/models/minimax_m3/common/sparse_attention.py +++ b/vllm/models/minimax_m3/common/sparse_attention.py @@ -384,7 +384,14 @@ class MiniMaxM3SparseTritonImpl(MiniMaxM3SparseImpl): nd = main_md.num_decode_tokens num_tokens = main_md.num_actual_tokens # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:]. - topk = layer.topk_indices_buffer # type: ignore[attr-defined] + topk_buffer = layer.topk_indices_buffer # type: ignore[attr-defined] + assert topk_buffer is not None + + topk = ( + topk_buffer + if current_platform.is_rocm() + else topk_buffer[:num_tokens].transpose(0, 1) + ) assert topk is not None hd = self.head_size q = query[:num_tokens].view(-1, self.num_heads, hd)