[None][perf] Custom decode kernels for MinimaxM3 - #17171
Closed
brb-nv wants to merge 10 commits into
Closed
Conversation
The indexer's per-block max score currently comes from an fmha_sm100 output_maxscore pass on the context schedule, which spends a 128-row Q tile on a single decode token. Port vLLM's dedicated mma.sync + TMA scorer, which computes the same per-block causal max of Q.K with a grid sized to the batch. The kernel writes [head, token, block] while the block selector consumes [head, block, token]; since the selector op takes arbitrary strides, the transposed view of the existing graph-stable score buffer is passed straight through, so no copy or selector change is needed and the stores become coalesced across tokens. Both decode kernels being ported need a per-request 2-D page table, which the MSA metadata did not have (fmha_sm100 takes a flattened one), so add msa_block_table and msa_seq_lens_cuda to the graph buffer pool alongside it. Off by default; set TLLM_M3_INDEXER_SCORE=cutedsl to enable. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The MSA context-schedule FMHA kernel spends a 128-row Q tile on a single decode token. Add a flash-decoding Triton kernel over the indexer's selected blocks, gated behind TLLM_M3_SPARSE_DECODE=triton, and route uniform-length pure-decode steps through it. Ported from vLLM's minimax_m3 sparse_attn ops, adapted to separate HND K/V views, fp32 split-K partials from the persistent buffer arena for CUDA graph stability, and zeroed (not NaN) output for graph padding rows. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Layers 0-2 attend the whole page table and only run on MSA because MsaSparseGqaFmha claims every M3 layer, so decode pays for a context-schedule kernel spending a 128-row Q tile on one token. Add a trtllm-gen decode path for them behind TLLM_M3_DENSE_DECODE=trtllm_gen. FlashInferTrtllmGenFmha cannot be reused: build_trtllm_gen_kv_cache_metadata assumes one layer contributes exactly K+V to a pool slot, while M3 packs the whole layer group plus per-sparse-layer index-K sub-pages, leaving no uniform per-layer stride. Expose that geometry as a flat sub-page pool on the cache manager instead, and derive the kernel's separate K and V block-table rows from M3's slot table. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The MiniMax-M3 decode kernels and their test oracles are derived from vLLM (Apache-2.0), and their headers named the upstream file plus a version string but no permalink, so the exact source revision could not be checked. Record the full source URL at commit 6f91edf96d3f3272945809c04702380053bff4de instead, matching how the rest of the tree cites borrowed code. Attribution is otherwise already in place: both projects are Apache-2.0, the vLLM copyright line is retained on every affected file, the docstrings list the downstream changes, and LICENSE lists vLLM. vLLM ships no NOTICE file, so there is none to propagate. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
The Triton sparse, trtllm-gen dense and CuteDSL indexer decode kernels only ran when a step contained no context request, so any batch that picked up a prefill sent every decode row back to fmha_sm100. Generation requests are the row suffix of a batch, so they form a span the ported kernels can own on a mixed step exactly as on a pure decode step. Resolve that span once per step (uniform query length over the generation rows, max KV length over those rows alone so a long context row cannot inflate the scheduling bound) and let fmha_sm100 keep the context prefix under a plan built over just those rows. The CuteDSL indexer stays whole-batch on mixed steps: max_score carries the query tokens in its last dimension, so the context and generation ranges are not separable sub-blocks two scorers could share. Splitting it is a follow-up. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv
requested review from
QiJune,
allisonlim-nv,
asfiyab-nvidia,
crazydemo,
dhansen-nvidia,
dpitman-nvda,
eopXD,
hyukn,
kris1025,
liji-nv,
schetlur-nv,
tburt-nv and
yuxianq
August 2, 2026 22:26
brb-nv
requested review from
pcicotti,
peihu-nv and
zheyuf
and removed request for
a team,
QiJune,
allisonlim-nv,
asfiyab-nvidia,
crazydemo,
dhansen-nvidia,
dpitman-nvda,
eopXD,
hyukn,
kris1025,
liji-nv,
schetlur-nv,
tburt-nv and
yuxianq
August 2, 2026 22:28
brb-nv
changed the title
Collaborator
Author
|
/bot run --disable-fail-fast |
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Collaborator
Author
|
Closing in favor of #17268. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
This MR adds a few custom decode kernels for MinimaxM3 ported from vLLM.
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.