Add power-of-2 page sizes >= 128 (128/256/512/1024) to the FlashInfer
backend's supported kernel block sizes. These are served only by the
trtllm-gen dynamic kernel, which requires Blackwell + GQA/MQA
(num_qo_heads // num_kv_heads > 1), not MHA.
Fail fast in FlashInferMetadataBuilder.__init__ when page_size >= 128 is
requested without a usable trtllm-gen path (trtllm disabled, not
Blackwell, or MHA) instead of silently hitting the native wrappers, and
force the prefill path to trtllm so it does not fall back to FA2.
Regenerate docs/design/attention_backends.md for the new block sizes.
AI assistance (Claude) was used for this change.
Co-authored-by: Claude
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Move the MiniMax M3 sparse-attention backend and Triton kernels out of the
shared attention tree into the model definition (mirroring deepseek_v4), split
the kernels by role, and tighten the backend, then port the dedicated split-K
decode top-k from the sglang reference.
Structure:
- vllm/v1/attention/backends/minimax_m3_sparse.py
-> vllm/models/minimax_m3/common/sparse_attention.py
- vllm/v1/attention/ops/minimax_m3_sparse_ops.py split into
common/ops/index_topk.py (index-score + top-k kernels)
common/ops/sparse_attn.py (block-sparse GQA attention kernels)
Pure-Triton/cross-platform, so under common/ (not nvidia/).
- registry enum + doc generator RELEVANT_PATTERNS repointed; the auto-generated
attention_backends.md now has a dedicated "MiniMax M3 Sparse Attention" section.
Backend cleanups:
- Drop redundant metadata fields (sparse-selection params, decode max_query_len,
decode cu_seqlens_q/context_lens) and read them from the layer/impl instead.
- Impl ctor takes explicit named args instead of kwargs.get lookups.
- Cudagraph support UNIFORM_SINGLE_TOKEN_DECODE -> UNIFORM_BATCH; reorder
threshold and decode buffer scale by 1 + num_speculative_tokens.
- Unify the per-request context-length buffer (max_num_batched_tokens).
- Declare bf16-only KV cache.
Decode top-k (ported from sglang minimax_sparse_ops/decode):
- Split-K top-k: per-chunk partial top-k (_topk_index_partial_kernel) + merge
(_topk_index_merge_kernel), replacing the single-program prefill top-k reuse.
- init/local block forcing moved into the decode score kernel (matches sglang).
- Verified against a torch reference (selected-block sets match across varied
seq lengths, with/without init+local forcing, and short sequences).
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Scaffold the lightning-indexer sparse-attention path:
- MiniMaxM3SparseBackend (registered as MINIMAX_M3_SPARSE): block-sparse GQA
backend; get_kv_cache_shape serves both the main K/V cache and the
single-vector index-key side cache.
- MiniMaxM3IndexerCache: side KV cache for per-token index keys, key-only so it
uses a single-vector MLAAttentionSpec rather than a K+V FullAttentionSpec.
- MiniMaxM3SparseMetadata (+ prefill/decode sub-metadata) and its builder,
splitting the batch via split_decodes_and_prefills.
- MiniMaxM3SparseImpl: subclasses AttentionImplBase so it can take a custom
forward(query, index_query, kv_cache, index_kv_cache); no alibi / sliding
window / logits soft cap. forward is a stub pending the kernel port.
MiniMaxM3SparseAttention is merged into a single AttentionLayerBase: it owns the
projections, per-head QK norm and RoPE, binds the backend + impl, registers the
main K/V cache, and holds the index cache. Its forward computes q/k/v and the
index q/k, pre-inserts K/V and index-K into their caches, then calls the sparse
impl with only the queries.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Yongye Zhu <yongye@inferact.ai>
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
Signed-off-by: Jee Jee Li <pandaleefree@gmail.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: Roger Wang <hey@rogerw.io>
