@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9021589498

[Minimax-M3] Using tok_sparse_select from MSA instead of triton kernels (#47502)

Signed-off-by: Yongye Zhu <zyy1102000@gmail.com> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Yongye Zhu committed Jul 7, 2026 at 21:01 UTC 902158949803be74ad6941ee4bd4b08741297fa0

9 files changed +233 -72

cmake/external_projects/fmha_sm100.cmake

+1 -1

@@ -17,7 +17,7 @@ else()

17 FetchContent_Declare(

18 fmha_sm100

19 GIT_REPOSITORY https://github.com/vllm-project/MSA.git

20 - GIT_TAG fee783153f3efe57e3e933c5cb7e267a7cebcfb5

20 + GIT_TAG 2e63ec37a0fc29bc20f39cd1a52e0f5affc33a73

21 GIT_PROGRESS TRUE

22 CONFIGURE_COMMAND ""

23 BUILD_COMMAND ""

tests/kernels/attention/test_minimax_m3.py

+4 -4

@@ -150,7 +150,7 @@ def _reference_index_topk(

150 num_blocks = (seq_len + BLOCK_SIZE - 1) // BLOCK_SIZE

151 pages = block_table[req_id, :num_blocks]

152 k = index_kv_cache[pages].reshape(num_blocks * BLOCK_SIZE, -1)

153 - score = torch.einsum("qhd,kd->hqk", q.float(), k.float()) * sm_scale

153 + score = sm_scale * torch.einsum("qhd,kd->hqk", q.float(), k.float())

154

155 q_pos = prefix_len + torch.arange(q_len, device=idx_q.device)

156 k_pos = torch.arange(k.shape[0], device=idx_q.device)

@@ -621,11 +621,12 @@ def test_decode_index_topk_fp8(num_idx_heads: int):

621 init_blocks=init_blocks,

622 local_blocks=local_blocks,

623 num_kv_heads=num_idx_heads,

624 - sm_scale=head_dim**-0.5,

624 decode_query_len=decode_query_len,

625 + max_decode_query_len=decode_query_len,

626 )

627 # Reference from the DEQUANTIZED fp8 values (the kernel computes the fp8 QK

628 - # in fp32, so it must match an fp32 matmul of the same e4m3 values).

628 + # in fp32 with no scaling, so it must match an unscaled fp32 matmul of the

629 + # same e4m3 values).

630 expected = _reference_index_topk(

631 idx_q.float(),

632 index_kv_cache.float(),

@@ -636,7 +637,6 @@ def test_decode_index_topk_fp8(num_idx_heads: int):

637 topk,

638 init_blocks,

639 local_blocks,

639 - head_dim**-0.5,

640 )

641 _assert_topk_indices_equal_unordered(actual, expected)

642

vllm/models/minimax_m3/common/indexer.py

@@ -256,6 +256,13 @@ class MiniMaxM3IndexerMetadataBuilder(

256 dtype=torch.int32,

257 device=device,

258 )

259 + # Stable per-token causal page-count buffer for decode cudagraph replays

260 + # (consumed by the MSA top-k path's sparse_topk_select num_valid_pages).

261 + self.num_valid_pages_buffer = torch.empty(

262 + vllm_config.scheduler_config.max_num_batched_tokens,

263 + dtype=torch.int32,

264 + device=device,

265 + )

266

267

268 class MiniMaxM3IndexerTritonMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

vllm/models/minimax_m3/common/ops/__init__.py

@@ -4,6 +4,7 @@

5 from .index_topk import (

6 minimax_m3_index_decode,

7 + minimax_m3_index_decode_score,

8 minimax_m3_index_score,

9 minimax_m3_index_topk,

10 )

@@ -11,6 +12,7 @@ from .sparse_attn import minimax_m3_sparse_attn, minimax_m3_sparse_attn_decode

13 __all__ = [

14 "minimax_m3_index_decode",

15 + "minimax_m3_index_decode_score",

16 "minimax_m3_index_score",

17 "minimax_m3_index_topk",

18 "minimax_m3_sparse_attn",

vllm/models/minimax_m3/common/ops/index_topk.py

+67 -15

@@ -757,25 +757,26 @@ def minimax_m3_index_topk(

757

758

759 @torch.no_grad()

760 -def minimax_m3_index_decode(

760 +def minimax_m3_index_decode_score(

761 idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]

762 index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]

763 block_table: torch.Tensor, # [num_reqs, max_blocks]

764 seq_lens: torch.Tensor, # [num_reqs] int32

765 max_seq_len: int,

766 - topk: int,

766 init_blocks: int,

767 local_blocks: int,

768 num_kv_heads: int,

769 decode_query_len: int,

770 max_decode_query_len: int,

772 - out: torch.Tensor | None = None,

771 + score_out: torch.Tensor | None = None,

772 ) -> torch.Tensor:

774 - """Decode index block-score + top-k, both split-K (cudagraph-safe).

773 + """Decode index block-score (split-K, cudagraph-safe); no top-k.

774

776 - Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).

777 - When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into

778 - ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.

775 + Returns score [num_kv_heads, total_q, >=max_block] (fp32; init/local blocks

776 + forced to 1e30/1e29). When ``score_out`` is given the scores are written into

777 + it (read/written by strides, so a transposed view of a unified buffer is

778 + accepted) instead of a fresh tensor -- used to share a unified score buffer

779 + with the prefill side and run a single top-k over both.

780 """

781 total_q, num_idx_heads, head_dim = idx_q.shape

782 assert num_idx_heads == num_kv_heads, (

@@ -783,7 +784,6 @@ def minimax_m3_index_decode(

784 )

785 assert decode_query_len <= max_decode_query_len

786 assert total_q == seq_lens.shape[0] * decode_query_len

786 - batch = total_q

787 max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)

788 use_pdl = current_platform.is_arch_support_pdl()

789 # `launch_pdl` is a Triton runtime kwarg only some backends accept (CUDA

@@ -800,13 +800,16 @@ def minimax_m3_index_decode(

800 if num_idx_heads > 1 and max_decode_query_len > 1:

801 score_kwargs.update({"num_warps": 4, "num_stages": 2})

802

803 - # Keep score strides 16-divisible to avoid Triton recompiles.

804 - score_block_stride = round_up(max_block, 16)

805 - score = torch.empty(

806 - (num_idx_heads, total_q, score_block_stride),

807 - dtype=torch.float32,

808 - device=idx_q.device,

809 - )

803 + if score_out is not None:

804 + score = score_out

805 + else:

806 + # Keep score strides 16-divisible to avoid Triton recompiles.

807 + score_block_stride = round_up(max_block, 16)

808 + score = torch.empty(

809 + (num_idx_heads, total_q, score_block_stride),

810 + dtype=torch.float32,

811 + device=idx_q.device,

812 + )

813 # split-K over seq blocks; chunk count depends only on shape constants so

814 # the grid is fixed within a cuda graph.

815 TARGET_GRID = 512

@@ -848,6 +851,55 @@ def minimax_m3_index_decode(

851 USE_PDL=use_pdl,

852 **score_kwargs,

853 )

854 + return score

855 +

856 +

857 +@torch.no_grad()

858 +def minimax_m3_index_decode(

859 + idx_q: torch.Tensor, # [total_q, num_idx_heads, head_dim]

860 + index_kv_cache: torch.Tensor, # [num_blocks, 128, head_dim]

861 + block_table: torch.Tensor, # [num_reqs, max_blocks]

862 + seq_lens: torch.Tensor, # [num_reqs] int32

863 + max_seq_len: int,

864 + topk: int,

865 + init_blocks: int,

866 + local_blocks: int,

867 + num_kv_heads: int,

868 + decode_query_len: int,

869 + max_decode_query_len: int,

870 + out: torch.Tensor | None = None,

871 + score_out: torch.Tensor | None = None,

872 +) -> torch.Tensor:

873 + """Decode index block-score + top-k, both split-K (cudagraph-safe).

874 +

875 + Returns topk_idx [num_kv_heads, total_q, topk] (0-indexed block ids, -1 pad).

876 + When ``out`` ([num_kv_heads, >=total_q, topk]) is given, writes into

877 + ``out[:, :total_q, :]`` (stable address for cudagraph) instead of allocating.

878 + When ``score_out`` ([num_kv_heads, total_q, >=max_block]) is given, the block

879 + scores are written into it (read back by the top-k) instead of a fresh

880 + tensor -- used to share a unified score buffer with the prefill side. Reads

881 + via strides, so a transposed view of a block-major buffer is accepted.

882 + """

883 + total_q, num_idx_heads, _ = idx_q.shape

884 + batch = total_q

885 + max_block = triton.cdiv(max_seq_len, SPARSE_BLOCK_SIZE)

886 + use_pdl = current_platform.is_arch_support_pdl()

887 + pdl_kwargs: dict[str, bool | int] = {}

888 + if use_pdl:

889 + pdl_kwargs.update({"launch_pdl": True})

890 + score = minimax_m3_index_decode_score(

891 + idx_q,

892 + index_kv_cache,

893 + block_table,

894 + seq_lens,

895 + max_seq_len,

896 + init_blocks,

897 + local_blocks,

898 + num_kv_heads,

899 + decode_query_len,

900 + max_decode_query_len,

901 + score_out=score_out,

902 + )

903

904 if out is not None:

905 topk_idx = out[:, :total_q, :]

vllm/models/minimax_m3/common/ops/sparse_attn.py

+7 -11

@@ -43,7 +43,6 @@ _FP8_DTYPES = (

43 {

44 "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),

45 "BLOCK_SIZE_H": lambda args: triton.next_power_of_2(args["gqa_group_size"]),

46 - "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),

46 "BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"]

47 * triton.next_power_of_2(args["gqa_group_size"]),

48 }

@@ -84,7 +83,6 @@ def _gqa_sparse_fwd_kernel(

83 BLOCK_SIZE_K: tl.constexpr, # == SPARSE_BLOCK_SIZE (128)

84 BLOCK_SIZE_D: tl.constexpr,

85 BLOCK_SIZE_H: tl.constexpr,

87 - BLOCK_SIZE_T: tl.constexpr,

86 BLOCK_SIZE_QH: tl.constexpr,

87 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load

88 ):

@@ -109,9 +107,10 @@ def _gqa_sparse_fwd_kernel(

107 for j in range(real_q_loop):

108 pid_q_j = pid_q * num_q_loop + j

109 t_ptr_j = t_ptr + (q_block_start + pid_q_j) * stride_tn + pid_kh * stride_th

112 - off_t = tl.arange(0, BLOCK_SIZE_T)

113 - topk_idx = tl.load(t_ptr_j + off_t * stride_tk, mask=off_t < max_topk, other=-1)

114 - real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)

110 + # Valid block count from seq position (no sentinel): block_size_q == 1.

111 + q_abs = prefix_len + pid_q_j * BLOCK_SIZE_Q

112 + valid_blocks = (q_abs + BLOCK_SIZE_K) // BLOCK_SIZE_K

113 + real_topk = tl.minimum(max_topk, valid_blocks)

114 q_ptrs = tl.make_block_ptr(

115 base=q_ptr + q_start * stride_qn + pid_h * stride_qh,

116 shape=(q_len, gqa_group_size, head_dim),

@@ -202,7 +201,6 @@ def _gqa_sparse_fwd_kernel(

201 16, triton.next_power_of_2(args["gqa_group_size"])

202 ),

203 "BLOCK_SIZE_D": lambda args: triton.next_power_of_2(args["head_dim"]),

205 - "BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),

204 }

205 )

206 @triton.jit(do_not_specialize=["decode_query_len"])

@@ -243,7 +241,6 @@ def _gqa_sparse_decode_kernel(

241 NUM_TOPK_CHUNKS: tl.constexpr,

242 BLOCK_SIZE_H: tl.constexpr,

243 BLOCK_SIZE_D: tl.constexpr,

246 - BLOCK_SIZE_T: tl.constexpr,

244 USE_FP8: tl.constexpr, # fp8 KV cache: dequantize K/V to q.dtype on load

245 USE_PDL: tl.constexpr,

246 ):

@@ -268,11 +265,10 @@ def _gqa_sparse_decode_kernel(

265 # attention range instead of letting padded rows produce negative lengths.

266 kv_len = tl.maximum(query_pos + 1, 0)

267

271 - # number of valid (non-padded) selected blocks for this query token

272 - off_t = tl.arange(0, BLOCK_SIZE_T)

268 + # Valid block count from seq_len (no sentinel): min(topk, cdiv(kv_len, blk)).

269 idx_base = t_ptr + pid_kh * stride_th + pid_b * stride_tn

274 - topk_idx = tl.load(idx_base + off_t * stride_tk, mask=off_t < max_topk, other=-1)

275 - real_topk = tl.sum((topk_idx >= 0).to(tl.int32), axis=0)

270 + num_blocks = (kv_len + BLOCK_SIZE_K - 1) // BLOCK_SIZE_K

271 + real_topk = tl.minimum(max_topk, num_blocks)

272 chunk_end_topk = tl.minimum(chunk_end_compiletime, real_topk)

273

274 off_n = tl.arange(0, BLOCK_SIZE_K)

vllm/models/minimax_m3/nvidia/indexer_msa.py

+135 -35

@@ -2,16 +2,21 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """MSA (SM100/Blackwell) indexer impl for MiniMax M3.

5 -Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path then selects

6 -top-k blocks with the Triton ``minimax_m3_index_topk`` kernel -- fmha is much

7 -faster than Triton for the wide prefill score (benchmarked ~3-5x).

8 -

9 -Decode uses the Triton fused ``minimax_m3_index_decode`` (the same kernel the

10 -Triton indexer impl uses): for q_len==1 it is a purpose-built vector x matrix

11 -score (no wasted tensor-core tiles) with a 256-way split-K and a fused split-K

12 -top-k, which beats fmha's OnlyScore (wasted MMA on a single query, 64-split cap)

13 -by ~1.1-3.7x. It is cudagraph-safe by construction (shape-constant split grids)

14 -and writes the shared ``topk_indices_buffer`` via ``out=``.

5 +Both sides write block scores into one unified token-major buffer

6 +``[total_q, H, max_k_tiles]``, then a single ``fmha_sm100.sparse_topk_select``

7 +selects the top-k blocks for the whole batch (decode ``[:nd]`` + prefill

8 +``[nd:]``) into the shared ``topk_indices_buffer``. It bounds each row by its

9 +causal page count and force-includes the init/local blocks, so the unwritten

10 +tail of the buffer is pre-filled with ``-inf``.

11 +

12 +Prefill scores with ``fmha_sm100``'s score-only (``OnlyScore``) path (much faster

13 +than Triton for the wide prefill score, benchmarked ~3-5x), writing its

14 +``max_score`` straight into the buffer's prefill region (stride-aware, no copy).

15 +

16 +Decode scores with the Triton split-K ``minimax_m3_index_decode_score`` (a

17 +purpose-built vector x matrix score, no wasted tensor-core tiles, 256-way

18 +split-K, cudagraph-safe by shape-constant grids), writing into the decode

19 +region. Its tuning heuristics are kept; only the top-k is shared with prefill.

21 ``fmha_sm100`` imports are function-local so this module is import-safe on

22 AMD / non-SM100.

@@ -22,6 +27,7 @@ from typing import ClassVar

28 import torch

30 +from vllm.config import VllmConfig

31 from vllm.forward_context import get_forward_context

32 from vllm.models.minimax_m3.common.indexer import (

33 MiniMaxM3IndexerBackend,

@@ -31,8 +37,7 @@ from vllm.models.minimax_m3.common.indexer import (

37 MiniMaxM3IndexerMetadataBuilder,

38 )

39 from vllm.models.minimax_m3.common.ops.index_topk import (

34 - minimax_m3_index_decode,

35 - minimax_m3_index_topk,

40 + minimax_m3_index_decode_score,

41 )

42 from vllm.v1.attention.backend import (

43 AttentionBackend,

@@ -40,10 +45,22 @@ from vllm.v1.attention.backend import (

45 CommonAttentionMetadata,

46 )

47 from vllm.v1.attention.backends.utils import split_decodes_and_prefills

48 +from vllm.v1.kv_cache_interface import AttentionSpec

50 # Page size == sparse block size == index-K block; fmha tile id == M3 block id.

51 PAGE_SIZE = 128

53 +# Fill for unwritten score tiles: -inf so they never win the top-k (score kernels

54 +# only write causally-valid blocks).

55 +_SCORE_SENTINEL = float("-inf")

56 +

57 +# Tile (KV-block) dim of the unified score buffer, hardcoded as a cudagraph

58 +# capture-time constant so the decode score kernel's buffer shape is frozen

59 +# across replays. 8192 tiles == 1M tokens of context; -inf padding +

60 +# num_valid_pages bound each row to its causal range, so shorter replays reuse

61 +# the same buffer safely.

62 +MAX_K_TILES = 8192

63 +

65 class MiniMaxM3IndexerMSABackend(MiniMaxM3IndexerBackend):

66 """Indexer side-cache backend selecting the MSA builder."""

@@ -71,6 +88,25 @@ class MiniMaxM3IndexerMSAMetadata(MiniMaxM3IndexerMetadata):

88 (the base ``prefill`` field is unused on this path)."""

90 prefill_msa: MiniMaxM3IndexerMSAPrefillMetadata | None = None

91 + # Per-forward view (``[:num_tokens]``) of the builder's persistent unified

92 + # score buffer ``[total_q, H, MAX_K_TILES]``, shared by decode and prefill

93 + # and reused across all layers. Pre-filled with the -inf sentinel in

94 + # ``build()``; each layer overwrites its valid tiles before its own top-k.

95 + unified_scores: torch.Tensor | None = None

96 + # Tile (KV-block) dim of the unified score buffer (== ``MAX_K_TILES``).

97 + # Forced as the fmha plan's ``max_k_tiles`` so prefill writes its max_score

98 + # straight into the shared buffer.

99 + max_k_tiles: int = 0

100 + # Batch-wide inputs for the single top-k over the unified buffer (decode +

101 + # prefill in one call). ``cu_seqlens_q`` is the per-request query-start

102 + # offsets (== query_start_loc, a stable view for cudagraph); ``prefix_lens``

103 + # is the per-request context length (== context_lens).

104 + topk_cu_seqlens_q: torch.Tensor | None = None

105 + topk_prefix_lens: torch.Tensor | None = None

106 + topk_max_query_len: int = 0

107 + # Per-token causal page count cdiv(seq_pos+1, PAGE_SIZE), [total_q] int32:

108 + # drives sparse_topk_select force_end_blocks + -1 out-of-range clamp.

109 + topk_num_valid_pages: torch.Tensor | None = None

110

111

112 class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

@@ -79,6 +115,28 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

115

116 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

117

118 + def __init__(

119 + self,

120 + kv_cache_spec: AttentionSpec,

121 + layer_names: list[str],

122 + vllm_config: VllmConfig,

123 + device: torch.device,

124 + ) -> None:

125 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

126 + # Persistent unified score buffer [T, H, MAX_K_TILES] shared by all

127 + # indexer layers and reused across forwards. Stable address (required

128 + # for the captured decode path) + fixed tile dim so the decode score

129 + # kernel's shape is frozen at capture. Filled with -inf per forward in

130 + # build(); the valid/padding partition is a per-forward constant, so

131 + # every layer overwrites the same valid tiles before its own top-k.

132 + self.unified_scores_buffer = torch.empty(

133 + vllm_config.scheduler_config.max_num_batched_tokens,

134 + self.num_index_heads,

135 + MAX_K_TILES,

136 + dtype=torch.float32,

137 + device=device,

138 + )

139 +

140 def build(

141 self,

142 common_prefix_len: int,

@@ -107,6 +165,19 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

165 common_attn_metadata.compute_num_computed_tokens(), non_blocking=True

166 )

167

168 + # Per-token causal page count for the top-k, into the stable cg buffer.

169 + positions = common_attn_metadata.positions

170 + assert positions is not None

171 + num_valid_pages = self.num_valid_pages_buffer[:num_tokens]

172 + num_valid_pages.copy_(positions[:num_tokens] // PAGE_SIZE + 1)

173 +

174 + # Unified score buffer: a per-forward view of the persistent buffer,

175 + # reset to the -inf sentinel once here and shared by every layer (the

176 + # tile dim is the capture-time constant MAX_K_TILES).

177 + max_k_tiles = MAX_K_TILES

178 + unified_scores = self.unified_scores_buffer[:num_tokens]

179 + unified_scores.fill_(_SCORE_SENTINEL)

180 +

181 decode: MiniMaxM3IndexerDecodeMetadata | None = None

182 if num_decodes > 0:

183 qsl_cpu = common_attn_metadata.query_start_loc_cpu

@@ -148,6 +219,11 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

219 causal=True,

220 num_kv_splits=1,

221 )

222 + # Force the plan's tile dim to the unified buffer's so prefill writes

223 + # its max_score straight into unified[:, :, nd:] (the stride-aware

224 + # binding shape-matches the tile dim exactly). max_k_tiles >= the

225 + # plan's natural value, so the extra tiles are simply never written.

226 + plan["max_k_tiles"] = max_k_tiles

227 cols = torch.arange(block_table.shape[1], device=block_table.device)

228 valid = cols[None, :] < nvp[lo:hi].to(block_table.device)[:, None]

229 prefill = MiniMaxM3IndexerMSAPrefillMetadata(

@@ -171,6 +247,12 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder):

247 num_prefill_tokens=num_prefill_tokens,

248 decode=decode,

249 prefill_msa=prefill,

250 + unified_scores=unified_scores,

251 + max_k_tiles=max_k_tiles,

252 + topk_cu_seqlens_q=query_start_loc[: num_reqs + 1],

253 + topk_prefix_lens=context_lens,

254 + topk_max_query_len=common_attn_metadata.max_query_len,

255 + topk_num_valid_pages=num_valid_pages,

256 )

257

258

@@ -195,36 +277,49 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):

277 -1, self.num_index_heads, self.index_head_dim

278 )

279 kv = self.index_cache.kv_cache

198 - # Both sides write into the single shared persistent topk_indices_buffer:

199 - # decode at [:, :nd], prefill at [:, nd:] (each kernel writes [:, :total_q]).

280 + # Shared persistent top-k output buffer; the unified top-k below writes

281 + # the selected block ids into buf[:num_tokens].

282 buf = self.topk_indices_buffer

201 -

202 - decode_topk: torch.Tensor | None = None

283 + assert buf is not None

284 +

285 + # Unified token-major score buffer [total_q, H, MAX_K_TILES]: the tile

286 + # dim is innermost/contiguous, so both fmha writes (native [T,H,K]) and

287 + # the block-iterating top-k reads hit contiguous tiles. Each side gets a

288 + # contiguous slice on dim 0: decode [:nd], prefill [nd:]; the kernels

289 + # read/write by strides. The builder allocates it once (persistent,

290 + # shared by all layers) and resets it to the sentinel each forward, so

291 + # the top-k never picks an unwritten tile.

292 + unified_scores = md.unified_scores

293 + assert unified_scores is not None

294 +

295 + # Decode scores -> unified[:nd] (transposed [H, nd, MK] view; the kernel

296 + # writes by strides). Top-k is deferred to the single unified call below.

297 if md.decode is not None:

298 d = md.decode

205 - decode_topk = minimax_m3_index_decode(

299 + minimax_m3_index_decode_score(

300 index_q[:nd],

301 kv,

302 d.block_table,

303 d.seq_lens,

304 d.max_seq_len,

211 - self.topk_blocks,

305 self.init_blocks,

306 self.local_blocks,

307 self.num_kv_heads,

308 d.decode_query_len,

309 d.max_decode_query_len,

217 - out=buf,

310 + score_out=unified_scores[:nd].transpose(0, 1),

311 )

312

220 - prefill_topk: torch.Tensor | None = None

313 if md.prefill_msa is not None:

314 from vllm.third_party.fmha_sm100.api import _fmha_sm100

315

316 p = md.prefill_msa

317 # Index-K cache (num_blocks, 128, D) -> paged MQA (num_blocks,1,128,D).

318 k_pages = kv.view(kv.shape[0], 1, PAGE_SIZE, self.index_head_dim)

227 - _, max_score = _fmha_sm100(

319 + # fmha writes its max_score natively as [nnz_p, H, max_k_tiles] into

320 + # the prefill region (stride-aware; plan max_k_tiles forced to

321 + # md.max_k_tiles so the shape matches exactly -> no copy).

322 + _fmha_sm100(

323 index_q[nd:],

324 k_pages,

325 k_pages, # V placeholder; not read in OnlyScore

@@ -233,19 +328,24 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl):

328 output_o=False,

329 output_maxscore=True,

330 sm_scale=self.scale,

331 + max_score=unified_scores[nd:],

332 )

237 - # Triton top-k wants [num_index_heads, num_tokens, max_block]; the

238 - # transpose is a strided view (the kernel reads via strides).

239 - out = buf[:, nd:, :] if buf is not None else None

240 - prefill_topk = minimax_m3_index_topk(

241 - max_score.transpose(1, 2),

242 - p.cu_seqlens_q,

243 - p.prefix_lens,

244 - p.max_query_len,

245 - self.topk_blocks,

246 - self.init_blocks,

247 - self.local_blocks,

248 - out=out,

249 - )

333

251 - return decode_topk, prefill_topk

334 + # Single top-k over the unified buffer via fmha_sm100 sparse_topk_select

335 + # (THK, no transpose) into ``buf``. num_valid_pages drives force_end_blocks

336 + # (always keep each token's local block; fmha OnlyScore won't) + the -1

337 + # out-of-range clamp.

338 + from vllm.third_party.fmha_sm100.api import sparse_topk_select

339 +

340 + sparse_topk_select(

341 + unified_scores,

342 + self.topk_blocks,

343 + num_valid_pages=md.topk_num_valid_pages,

344 + force_begin_blocks=self.init_blocks,

345 + force_end_blocks=self.local_blocks,

346 + output=buf[:num_tokens],

347 + max_score_layout="THK",

348 + )

349 +

350 + # The attend reads ``buf`` directly; this return is vestigial.

351 + return None, None

vllm/models/minimax_m3/nvidia/model.py

+4 -3

@@ -769,8 +769,9 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

769

770 # Reserved top-k indices buffer shared by all sparse-attention indexer

771 # layers (mirrors DeepseekV4); kept at a stable address so the indexer's

772 - # top-k output survives cudagraph capture/replay. Shape matches the

773 - # per-head index top-k output [num_index_heads, total_q, topk].

772 + # top-k output survives cudagraph capture/replay. Token-major

773 + # [total_q, num_index_heads, topk] so the indexer writes its native

774 + # [token, head, topk] top-k; the attend transposes to [H, tokens, topk].

775 sparse_cfg = getattr(config, "sparse_attention_config", None)

776 if sparse_cfg is not None:

777 tp_size = get_tensor_model_parallel_world_size()

@@ -780,8 +781,8 @@ class MiniMaxM3Model(nn.Module, EagleModelMixin):

781 max_num_batched_tokens = vllm_config.scheduler_config.max_num_batched_tokens

782 padded_num_tokens = (max_num_batched_tokens + 3) // 4 * 4

783 self.topk_indices_buffer = torch.empty(

783 - num_index_heads,

784 padded_num_tokens,

785 + num_index_heads,

786 sparse_cfg["sparse_topk_blocks"],

787 dtype=torch.int32,

788 )

vllm/models/minimax_m3/nvidia/sparse_attention_msa.py

+6 -3

@@ -39,7 +39,8 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):

40 nd = main_md.num_decode_tokens

41 num_tokens = main_md.num_actual_tokens

42 - # Indexer top-k from the shared buffer: decode [:, :nd], prefill [:, nd:].

42 + # Indexer top-k from the shared token-major buffer [total_q, H, MK]; the

43 + # kernels want [H, tokens, MK], so slice tokens on dim 0 then transpose.

44 topk = layer.topk_indices_buffer # type: ignore[attr-defined]

45 assert topk is not None

46 hd = self.head_size

@@ -56,7 +57,7 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):

57 minimax_m3_sparse_attn_decode(

58 q[:nd],

59 kv_cache,

59 - topk[:, :nd, :],

60 + topk[:nd].transpose(0, 1),

61 d.block_table,

62 d.seq_lens,

63 self.num_kv_heads,

@@ -74,7 +75,9 @@ class MiniMaxM3SparseMSAImpl(MiniMaxM3SparseImpl):

76 p = main_md.prefill

77 assert p is not None

77 - prefill_topk = topk[:, nd:num_tokens, :]

78 + # [H, prefill, MK] transposed view; build_k2q_csr consumes the

79 + # strided view directly (topK stays innermost-contiguous).

80 + prefill_topk = topk[nd:num_tokens].transpose(0, 1)

81 qp = q[nd:]

82 k_cache = kv_cache[:, 0].transpose(1, 2)

83 v_cache = kv_cache[:, 1].transpose(1, 2)