forked from Karylab-cklius/vllm
cb7eb5c9b4
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>