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 <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
The SM100 (MSA) lightning indexer declared AttentionCGSupport.NEVER and ran
eager: its fmha_sm100 score plan allocated fresh buffers every build() and the
run allocated a fresh max_score / top-k each call.
Make the decode side cudagraph-replay-safe:
- Reserve persistent plan buffers in the builder __init__ (sized over a scan of
decode sizes [1, max_num_seqs]); build() fills them in place and calls the plan
kernel directly with a fixed, batch-size-only num_kv_splits
(estimate_num_kv_splits, a uniform-context replica of the planner's auto-split
math; tunable via VLLM_M3_INDEXER_CONTEXT_LEN). A positive split count takes the
deterministic plan path with no device->host sync.
- workspace_o / workspace_lse / cute_workspace are builder-owned dedicated tensors
(not the shared global _alloc_workspace_buf cache) so a larger fmha call
elsewhere can't realloc and move an address a captured graph baked.
- max_score is a 1-D-backed contiguous [H, max_k_tiles, nnz_qo] view (a sliced 3-D
buffer is non-contiguous; the kernel assumes contiguous); max_k_tiles pinned for
a stable shape.
- Top-k output goes to a model-level topk_indices_buffer (DeepSeek-V3.2 pattern),
threaded model -> decoder layer -> sparse attention -> indexer; index_topk gains
an out= param.
- Narrow _run_attention's eager break so the indexer runs in the captured segment
and only the sparse attention is eager-broken. Builder reports UNIFORM_BATCH;
prefill stays eager.
Verified: tests/kernels/attention/test_minimax_m3.py 42/42 pass (incl. impl-level
parity with num_kv_splits > 1 over short context); GPQA accuracy matches eager.
AI assistance (Claude Code) was used for this change.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Add an SM100/Blackwell lightning-indexer impl that computes the per-128-block
QK max-scores with fmha_sm100's score-only (OnlyScore) path and selects the
top-k blocks with the existing Triton minimax_m3_index_topk kernel, mirroring
how the main MSA attention pairs the SM100 attend with Triton. Decode and
prefill requests are split manually (decode-first batch) and each side gets its
own _fmha_sm100_plan / _fmha_sm100 call. Auto-selected on SM100 when
topk_blocks in (4, 8, 16, 32) for both bf16 and fp8 index caches; falls back to
the Triton indexer otherwise. The builder declares AttentionCGSupport.NEVER
(eager; the attention is broken out of the graph by _run_attention).
Extend the fused qknorm+rope+kv-insert kernel to optionally emit fp8 (e4m3) for
the index-K cache and index-Q via a direct cast with no scale tensors (RMSNorm
outputs are O(1) and scalar scales do not change top-k ordering). Only the index
outputs go fp8; q/k/v and q_out stay bf16 and bit-identical to the existing
path. MiniMaxM3IndexerCache now accepts fp8 caches and the model allocates
index_q in the cache dtype.
Tests: test_fmha_sm100_indexer_matches_reference (bf16/fp8 x prefill/decode) and
test_msa_indexer_impl_matches_triton (full impl parity vs the Triton indexer
through the real metadata builders); fp8 fused-kernel parity is covered in
test_fused_minimax_m3_qknorm_rope_kv_insert.
AI assistance (Claude Code) was used for this change.
Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
