forked from Karylab-cklius/vllm
Yongye ZhuandClaude Opus 4.8
bb4844dba3
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>