From d2fbaf73c19e2bf519e4e4fc954ec8c14ba1e1fc Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 16 Jun 2026 20:28:45 +0000 Subject: [PATCH] [Model] M3 MSA indexer: unify top-k buffer, drop numpy, GPU-only decode plan Follow-ups on the cudagraph-capturable MSA indexer: - Top-k: both decode and prefill now write into the single shared, persistent topk_indices_buffer (decode at [:, :nd], prefill at [:, nd:]) and return views into it -- no fresh per-step top-k allocations. - Build the decode plan + flat page table entirely with torch on-GPU: drop numpy and CpuGpuBuffer; segment offsets/lengths are computed via torch.cumsum into the persistent int32 buffers, and the request-major page table is scattered into the buffer via the on-GPU page indptr (the run bounds reads by indptr, so the full buffer is passed and no host page count is needed). - No GPU->CPU sync on the decode path: scalars come from host ints (num_decode_tokens // num_decodes), and seq_lens.cpu() is confined to the eager prefill branch. The impl forward (fmha OnlyScore + Triton top-k) was already sync-free. test_msa_indexer_impl_matches_triton now also asserts both outputs are views into the persistent buffer. 42/42 in test_minimax_m3.py pass. Co-Authored-By: Claude Opus 4.8 Signed-off-by: Yongye Zhu Signed-off-by: Yongye Zhu --- tests/kernels/attention/test_minimax_m3.py | 11 ++ vllm/models/minimax_m3/nvidia/indexer_msa.py | 125 +++++++++++-------- 2 files changed, 83 insertions(+), 53 deletions(-) diff --git a/tests/kernels/attention/test_minimax_m3.py b/tests/kernels/attention/test_minimax_m3.py index 07be65649b4..ef0958edebc 100644 --- a/tests/kernels/attention/test_minimax_m3.py +++ b/tests/kernels/attention/test_minimax_m3.py @@ -475,6 +475,13 @@ def test_msa_indexer_impl_matches_triton(topk, monkeypatch): msa_impl.index_cache.kv_cache = index_cache triton_impl.index_cache.kv_cache = index_cache + # Exercise the shared persistent top-k buffer: the MSA impl must write both + # decode ([:, :nd]) and prefill ([:, nd:]) into it and return views of it. + nd = sum(q for q in batch.query_lens if q <= 1) + msa_impl.topk_indices_buffer = torch.full( + (num_idx_heads, num_tokens, topk), -2, dtype=torch.int32, device=device + ) + attn_metadata = { msa_impl.index_cache.prefix: msa_builder.build(0, common), triton_impl.index_cache.prefix: triton_builder.build(0, common), @@ -487,6 +494,10 @@ def test_msa_indexer_impl_matches_triton(topk, monkeypatch): assert msa_prefill is not None and tri_prefill is not None _assert_topk_indices_equal_unordered(msa_decode, tri_decode) _assert_topk_indices_equal_unordered(msa_prefill, tri_prefill) + # decode/prefill outputs are views into the one persistent buffer. + buf = msa_impl.topk_indices_buffer + assert msa_decode.data_ptr() == buf[:, :nd, :].data_ptr() + assert msa_prefill.data_ptr() == buf[:, nd:, :].data_ptr() @pytest.mark.parametrize( diff --git a/vllm/models/minimax_m3/nvidia/indexer_msa.py b/vllm/models/minimax_m3/nvidia/indexer_msa.py index 535671427c1..a03eb5dd3e0 100644 --- a/vllm/models/minimax_m3/nvidia/indexer_msa.py +++ b/vllm/models/minimax_m3/nvidia/indexer_msa.py @@ -25,7 +25,6 @@ AMD / non-SM100. from dataclasses import dataclass from typing import ClassVar -import numpy as np import torch import vllm.envs as envs @@ -38,7 +37,6 @@ from vllm.models.minimax_m3.common.indexer import ( MiniMaxM3IndexerMetadataBuilder, ) from vllm.models.minimax_m3.common.ops.index_topk import minimax_m3_index_topk -from vllm.utils.platform_utils import is_pin_memory_available from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, @@ -46,7 +44,6 @@ from vllm.v1.attention.backend import ( ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.kv_cache_interface import AttentionSpec -from vllm.v1.utils import CpuGpuBuffer # Page size == sparse block size == index-K block; fmha tile id == M3 block id. PAGE_SIZE = 128 @@ -140,7 +137,6 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): max_model_len = vllm_config.model_config.max_model_len pages_per_req = (max_model_len + PAGE_SIZE - 1) // PAGE_SIZE self._max_k_tiles = ((pages_per_req + 127) // 128) * 128 - pin = not vllm_config.use_v2_model_runner and is_pin_memory_available() max_reqs = vllm_config.scheduler_config.max_num_seqs @@ -173,10 +169,10 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): ) max_total_qo = max_reqs * dql * self._pack_factor - def seg(n: int) -> CpuGpuBuffer: - return CpuGpuBuffer(n, dtype=torch.int32, device=device, pin_memory=pin) + def seg(n: int) -> torch.Tensor: + return torch.empty(n, dtype=torch.int32, device=device) - # Host->device segment buffers (filled per build()). + # Segment buffers, computed on-GPU per build() (no host staging needed). self._qo_seg_off = seg(max_reqs + 1) self._kv_seg_off = seg(max_reqs + 1) self._kv_page_indptr = seg(max_reqs + 1) @@ -218,39 +214,47 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): def _plan_decode( self, - qo_lens: np.ndarray, - kv_lens: np.ndarray, - qo_offset: np.ndarray, + qsl_dec: torch.Tensor, + seq_lens_dec: torch.Tensor, + num_decode_tokens: int, num_kv_splits: int, ) -> dict: - """Fill the persistent plan buffers and launch the plan kernel directly, - returning a PlanInfo over those (stable-address) buffers.""" + """Fill the persistent plan buffers on-GPU and launch the plan kernel + directly, returning a PlanInfo over those (stable-address) buffers. + + ``qsl_dec`` is the decode slice of ``query_start_loc`` ([n+1]); the + per-request q/kv lengths and their cumsums are derived with torch (no + host sync) and written into the persistent segment buffers in place. + """ from vllm.third_party.fmha_sm100.api import _call_plan, _make_plan_info - n = qo_lens.shape[0] + n = seq_lens_dec.shape[0] pf = self._pack_factor hq = self._packed_heads - packed_qo = (qo_lens * pf).astype(np.int32) - nvp = (kv_lens + PAGE_SIZE - 1) // PAGE_SIZE - qo_off = np.zeros(n + 1, dtype=np.int32) - np.cumsum(packed_qo, out=qo_off[1:]) - kv_off = np.zeros(n + 1, dtype=np.int32) - np.cumsum(kv_lens, out=kv_off[1:]) - page_indptr = np.zeros(n + 1, dtype=np.int32) - np.cumsum(nvp, out=page_indptr[1:]) - total_qo_len = int(qo_off[n]) - max_qo_len = int(packed_qo.max()) + # Uniform decode -> packed total/max known from host scalars (no sync). + total_qo_len = num_decode_tokens * pf + max_qo_len = (num_decode_tokens // n) * pf - def fill(buf: CpuGpuBuffer, data: np.ndarray, m: int) -> torch.Tensor: - buf.np[:m] = data - return buf.copy_to_gpu(m) + qo = (qsl_dec[1:] - qsl_dec[:-1]).to(torch.int32) + kv = seq_lens_dec.to(torch.int32) + packed_qo = qo * pf + nvp = (kv + PAGE_SIZE - 1) // PAGE_SIZE + self._qo_seg_lens[:n] = packed_qo + self._kv_seg_lens[:n] = kv + self._qo_offset[:n] = kv - qo # bottom-right causal + self._qo_seg_off[0] = 0 + self._kv_seg_off[0] = 0 + self._kv_page_indptr[0] = 0 + torch.cumsum(packed_qo, 0, out=self._qo_seg_off[1 : n + 1]) + torch.cumsum(kv, 0, out=self._kv_seg_off[1 : n + 1]) + torch.cumsum(nvp, 0, out=self._kv_page_indptr[1 : n + 1]) - qo_seg_off = fill(self._qo_seg_off, qo_off, n + 1) - kv_seg_off = fill(self._kv_seg_off, kv_off, n + 1) - kv_page_indptr = fill(self._kv_page_indptr, page_indptr, n + 1) - qo_offset_gpu = fill(self._qo_offset, qo_offset.astype(np.int32), n) - kv_seg_lens = fill(self._kv_seg_lens, kv_lens.astype(np.int32), n) - qo_seg_lens = fill(self._qo_seg_lens, packed_qo, n) + qo_seg_off = self._qo_seg_off[: n + 1] + kv_seg_off = self._kv_seg_off[: n + 1] + kv_page_indptr = self._kv_page_indptr[: n + 1] + qo_offset_gpu = self._qo_offset[:n] + kv_seg_lens = self._kv_seg_lens[:n] + qo_seg_lens = self._qo_seg_lens[:n] split = num_kv_splits > 1 kv_tile_begin = self._kv_tile_begin if split else None @@ -338,28 +342,32 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): common_attn_metadata.compute_num_computed_tokens(), non_blocking=True ) - qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1] - qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32) - kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32) - nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE cols = torch.arange(block_table.shape[1], device=block_table.device) decode_metadata: MiniMaxM3IndexerMSASubMetadata | None = None if num_decodes > 0: - qo_np = qo_lens_cpu[:num_decodes].numpy() - kv_np = kv_lens_cpu[:num_decodes].numpy() num_kv_splits = estimate_num_kv_splits( num_decodes, num_qo_heads=self._packed_heads, num_sms=self._num_ctas, context_len=self._ctx_len, ) - plan = self._plan_decode(qo_np, kv_np, kv_np - qo_np, num_kv_splits) - # Persistent decode flat page table (request-major). - dec_pages = int(nvp[:num_decodes].sum()) - valid = cols[None, :] < nvp[:num_decodes].to(block_table.device)[:, None] - self._decode_kv_indices[:dec_pages].copy_( - block_table[:num_decodes][valid].to(torch.int32), non_blocking=True + # Fills self._kv_page_indptr (per-request page offsets) on-GPU. + plan = self._plan_decode( + query_start_loc[: num_decodes + 1], + seq_lens[:num_decodes], + num_decode_tokens, + num_kv_splits, + ) + # Flat request-major decode page table, built on-GPU (no host sync): + # scatter block_table[b, :nvp[b]] to _decode_kv_indices[indptr[b] + j]. + # The full buffer is passed to the run; it bounds reads via the page + # indptr, so no host page count is needed. + nvp = (seq_lens[:num_decodes] + PAGE_SIZE - 1) // PAGE_SIZE + valid = cols[None, :] < nvp[:, None] + dest = self._kv_page_indptr[:num_decodes, None] + cols[None, :] + self._decode_kv_indices[dest[valid]] = block_table[:num_decodes][valid].to( + torch.int32 ) self._cu_seqlens_q[: num_decodes + 1].copy_( query_start_loc[: num_decodes + 1] - query_start_loc[0], @@ -369,8 +377,8 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): plan=plan, cu_seqlens_q=self._cu_seqlens_q[: num_decodes + 1], prefix_lens=context_lens[:num_decodes], - max_query_len=int(qo_lens_cpu[:num_decodes].max()), - page_table=self._decode_kv_indices[:dec_pages], + max_query_len=num_decode_tokens // num_decodes, # uniform decode qlen + page_table=self._decode_kv_indices, max_score=self._max_score[ : self.num_index_heads * self._max_k_tiles * num_decode_tokens ].view(self.num_index_heads, self._max_k_tiles, num_decode_tokens), @@ -378,9 +386,16 @@ class MiniMaxM3IndexerMSAMetadataBuilder(MiniMaxM3IndexerMetadataBuilder): prefill_metadata: MiniMaxM3IndexerMSASubMetadata | None = None if num_prefills > 0: + # Prefill is eager (not captured); the host lengths it needs (and the + # _fmha_sm100_plan .tolist() inside) make the D->H sync unavoidable + # here, but it never runs on the captured decode path. from vllm.third_party.fmha_sm100.api import _fmha_sm100_plan lo, hi = num_decodes, num_reqs + qsl_cpu = common_attn_metadata.query_start_loc_cpu[: num_reqs + 1] + qo_lens_cpu = (qsl_cpu[1:] - qsl_cpu[:-1]).to(torch.int32) + kv_lens_cpu = seq_lens[:num_reqs].cpu().to(torch.int32) + nvp = (kv_lens_cpu + PAGE_SIZE - 1) // PAGE_SIZE side_qo = qo_lens_cpu[lo:hi] side_kv = kv_lens_cpu[lo:hi] plan = _fmha_sm100_plan( @@ -469,9 +484,9 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): ) # Triton top-k wants [num_index_heads, num_tokens, max_block]; the # transpose is a strided view (the kernel reads via strides). One - # 128-token KV tile == one M3 sparse block. ``topk_out`` is the - # shared persistent buffer for decode (stable address), None for - # prefill (fresh, eager). + # 128-token KV tile == one M3 sparse block. ``topk_out`` is the slice + # of the shared persistent buffer this side writes into (stable + # address); None -> the kernel allocates fresh. return minimax_m3_index_topk( max_score.transpose(1, 2), meta.cu_seqlens_q, @@ -483,16 +498,20 @@ class MiniMaxM3IndexerMSAImpl(MiniMaxM3IndexerImpl): out=topk_out, ) + # Both sides write into the single persistent topk_indices_buffer (no + # fresh allocations): decode tokens at [:, :nd], prefill at [:, nd:] + # (the index_topk out= writes out[:, :total_q] for each side). + buf = self.topk_indices_buffer + def run_decode() -> torch.Tensor | None: if md.decode_metadata is None: return None - return score_topk( - md.decode_metadata, index_q[:nd], self.topk_indices_buffer - ) + return score_topk(md.decode_metadata, index_q[:nd], buf) def run_prefill() -> torch.Tensor | None: if md.prefill_metadata is None: return None - return score_topk(md.prefill_metadata, index_q[nd:], None) + out = buf[:, nd:, :] if buf is not None else None + return score_topk(md.prefill_metadata, index_q[nd:], out) return run_decode(), run_prefill()