Fixes mypy call-arg errors: two fp8 call sites pass sm_scale (by
keyword and positionally) that the reference signature lacked. Add
sm_scale (default 1.0) and apply it to the score; top-k selection is
invariant to a positive scalar, so existing callers are unchanged.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 522f262265)
Route the MSA (SM100) indexer decode through the Triton fused
minimax_m3_index_decode kernel (the same kernel the Triton indexer impl
uses) instead of fmha_sm100's OnlyScore path. For q_len==1 decode the
Triton kernel is a purpose-built vector x matrix score with a 256-way
split-K and fused split-K top-k, which beats fmha's OnlyScore (wasted MMA
tiles on a single query, 64-split cap) by ~1.1-3.7x in benchmarks. It is
cudagraph-safe by construction (shape-constant split grids) and writes the
shared topk_indices_buffer via out=. Prefill keeps fmha OnlyScore + the
single-pass Triton top-k, where fmha is ~3-5x faster for the wide score.
This drops the persistent fmha decode plan buffers, the num_kv_splits
estimate, and the now-unused VLLM_M3_INDEXER_CONTEXT_LEN env.
Add fp8 (e4m3) index KV cache support to the Triton decode score kernel:
make the QK MMA accumulate in fp32 (out_dtype=tl.float32) so the per-block
max score is exact for the e4m3 cache. Top-k is invariant to positive
scalar scaling, so fp8 needs no scale tuning.
Tests: tests/kernels/attention/test_minimax_m3.py (44 passed), including a
new test_decode_index_topk_fp8 validating fp8 decode top-k vs a
dequantized-fp32 reference. mypy + ruff clean. AI assistance was used.
Co-authored-by: Claude <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
(cherry picked from commit 0a102417bb)
Unify the Triton indexer impl onto the same persistent top-k buffer as the MSA
impl, and thread the buffer through the AMD model too.
- minimax_m3_index_decode gains an out= param (writes out[:, :total_q]); the merge
kernel already writes via strides, so a buffer view works.
- MiniMaxM3IndexerTritonImpl.forward writes decode ([:, :nd]) and prefill ([:, nd:])
into the shared topk_indices_buffer and returns views into it (no fresh per-step
top-k allocations), matching the MSA impl.
- amd/model.py: allocate the model-level topk_indices_buffer and thread it
model -> decoder layer -> sparse attention -> indexer (mirrors nvidia). AMD keeps
its eager break, so this is purely allocation reuse there.
test_msa_indexer_impl_matches_triton now gives each impl its own buffer and asserts
both decode/prefill outputs are views into it. 42/42 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>
(cherry picked from commit 06324dd3df)
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>
(cherry picked from commit d2fbaf73c1)
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>
(cherry picked from commit bb4844dba3)
