Skip to content

[None][perf] Address decode-first assumption in MSA - #16923

Merged
pcicotti merged 1 commit into
NVIDIA:feat/m3_with_msafrom
brb-nv:user/brb/split-mixed-batch
Jul 28, 2026
Merged

[None][perf] Address decode-first assumption in MSA#16923
pcicotti merged 1 commit into
NVIDIA:feat/m3_with_msafrom
brb-nv:user/brb/split-mixed-batch

Conversation

@brb-nv

@brb-nv brb-nv commented Jul 28, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Description

Problem:
In a mixed batch (context + generation requests in one step), fmha_sm100_plan recognized only decode-first ordering, but TensorRT-LLM always orders context requests first; so has_mixed was always false and the whole batch was planned as a single sparse-prefill call.

Fix:
_mixed_batch_split now also recognizes prefill-first ordering, restricted to sparse plans so the proxy and dense plans stay byte-identical (the proxy's max-score output drives top-k block selection). The sparse GQA plan then covers only the context rows, and generation rows route to the paged decode kernel.

For ISL=8192, OSL=128, c=320, request latency in ms:

Before:

[Latency] P50    : 60801.3481
[Latency] P90    : 192000.5032
[Latency] P95    : 199494.0752
[Latency] P99    : 205562.1021
[Latency] MINIMUM: 45813.1489
[Latency] MAXIMUM: 206961.5720
[Latency] AVERAGE: 105663.2407

After:

[Latency] P50    : 54283.6902
[Latency] P90    : 66217.9545
[Latency] P95    : 72916.6434
[Latency] P99    : 78360.5503
[Latency] MINIMUM: 25553.0081
[Latency] MAXIMUM: 79614.9843
[Latency] AVERAGE: 52492.6185

Test Coverage

$ pytest tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestMiniMaxM3::test_nvfp4[use_msa=True-eval_mode=default] -s -v

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

@brb-nv brb-nv changed the title User/brb/split mixed batch [None][perf] Address decode-first assumption in MSA Jul 28, 2026
@brb-nv
[None][perf] Address decode-first assumption in MSA
fa7279f 
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
@brb-nv
brb-nv force-pushed the user/brb/split-mixed-batch branch from af3e4dc to fa7279f Compare July 28, 2026 03:00
@brb-nv
brb-nv marked this pull request as ready for review July 28, 2026 03:01
@brb-nv
brb-nv requested review from a team as code owners July 28, 2026 03:01
@brb-nv
brb-nv requested review from PerkzZheng, pcicotti, peihu-nv, yunruis and zheyuf and removed request for a team, PerkzZheng and yunruis July 28, 2026 03:01
@brb-nv

brb-nv commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

tensorrt-cicd commented Jul 28, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #62080 [ run ] triggered by Bot. Commit: fa7279f Link to invocation

@pcicotti
pcicotti merged commit b8a5e8e into NVIDIA:feat/m3_with_msa Jul 28, 2026
10 of 11 checks passed
@peihu-nv peihu-nv mentioned this pull request Jul 28, 2026
Merged
2 tasks
ZhanruiSunCh pushed a commit that referenced this pull request Jul 29, 2026
@brb-nv
[None][perf] Address decode-first assumption in MSA (#16923)
6e9c670 
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
@tensorrt-cicd

tensorrt-cicd commented Jul 29, 2026

Copy link
Copy Markdown
Collaborator

PR_Github #62080 [ run ] completed with state ABORTED. Commit: fa7279f

Link to invocation

@brb-nv brb-nv mentioned this pull request Aug 19, 2026
Merged
1 task
brb-nv added a commit to brb-nv/TensorRT-LLM that referenced this pull request Aug 20, 2026
@brb-nv
[None][perf] Address decode-first assumption in MSA
d0afa22 
In a mixed batch (context plus generation requests in one step)
fmha_sm100_plan recognized only decode-first ordering, but TensorRT-LLM
always orders context requests first. has_mixed was therefore always false
and the whole batch was planned as a single sparse-prefill call, dragging
every generation row through the MM-SA-Nv prefill kernel.

_mixed_batch_split now also recognizes prefill-first ordering, restricted to
sparse plans so the proxy and dense plans stay byte-identical (the proxy's
max-score output drives top-k block selection). The sparse GQA plan then
covers only the context rows and generation rows route to the paged decode
kernel. Since each half of a split plan now writes its own rows of the
caller's output buffer, out= is plumbed through sparse_fmha,
sparse_atten_func and _sparse_atten_csr_varlen_forward, which removes a
full-size concatenate plus copy from every split call.

Cherry-picked from NVIDIA#16923 on feat/m3_with_msa, whose changes outside the
submodule patch do not apply here: the msa_backend.py hunk is a comment
inside an on_update_kv_lens override main's MSA metadata does not have, and
main carries no MSA patch-marker check for the new fmha_sm100.api symbol to
extend. Rebased onto main's copy of the patch, which carries a later
revision of the _prepare_paged_hnd_input hunk and its tests; the interface.py
hunk offsets are adjusted for those three extra lines.

Verified with `patch -p1 --forward --batch` against MSA e2ebe76 (the pinned
submodule commit): all four files apply with no fuzz or offset, the reverse
dry run still succeeds, and the patched sources parse.

(cherry picked from commit 6e9c670)
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv added a commit to brb-nv/TensorRT-LLM that referenced this pull request Aug 21, 2026
@brb-nv
[None][perf] Address decode-first assumption in MSA
5200ffc 
In a mixed batch (context plus generation requests in one step)
fmha_sm100_plan recognized only decode-first ordering, but TensorRT-LLM
always orders context requests first. has_mixed was therefore always false
and the whole batch was planned as a single sparse-prefill call, dragging
every generation row through the MM-SA-Nv prefill kernel.

_mixed_batch_split now also recognizes prefill-first ordering, restricted to
sparse plans so the proxy and dense plans stay byte-identical (the proxy's
max-score output drives top-k block selection). The sparse GQA plan then
covers only the context rows and generation rows route to the paged decode
kernel. Since each half of a split plan now writes its own rows of the
caller's output buffer, out= is plumbed through sparse_fmha,
sparse_atten_func and _sparse_atten_csr_varlen_forward, which removes a
full-size concatenate plus copy from every split call.

Cherry-picked from NVIDIA#16923 on feat/m3_with_msa, whose changes outside the
submodule patch do not apply here: the msa_backend.py hunk is a comment
inside an on_update_kv_lens override main's MSA metadata does not have, and
main carries no MSA patch-marker check for the new fmha_sm100.api symbol to
extend. Rebased onto main's copy of the patch, which carries a later
revision of the _prepare_paged_hnd_input hunk and its tests; the interface.py
hunk offsets are adjusted for those three extra lines.

Verified with `patch -p1 --forward --batch` against MSA e2ebe76 (the pinned
submodule commit): all four files apply with no fuzz or offset, the reverse
dry run still succeeds, and the patched sources parse.

(cherry picked from commit 6e9c670)
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv added a commit to brb-nv/TensorRT-LLM that referenced this pull request Aug 24, 2026
@brb-nv
[None][perf] Address decode-first assumption in MSA
f4561a4 
In a mixed batch (context plus generation requests in one step)
fmha_sm100_plan recognized only decode-first ordering, but TensorRT-LLM
always orders context requests first. has_mixed was therefore always false
and the whole batch was planned as a single sparse-prefill call, dragging
every generation row through the MM-SA-Nv prefill kernel.

_mixed_batch_split now also recognizes prefill-first ordering, restricted to
sparse plans so the proxy and dense plans stay byte-identical (the proxy's
max-score output drives top-k block selection). The sparse GQA plan then
covers only the context rows and generation rows route to the paged decode
kernel. Since each half of a split plan now writes its own rows of the
caller's output buffer, out= is plumbed through sparse_fmha,
sparse_atten_func and _sparse_atten_csr_varlen_forward, which removes a
full-size concatenate plus copy from every split call.

Cherry-picked from NVIDIA#16923 on feat/m3_with_msa, whose changes outside the
submodule patch do not apply here: the msa_backend.py hunk is a comment
inside an on_update_kv_lens override main's MSA metadata does not have, and
main carries no MSA patch-marker check for the new fmha_sm100.api symbol to
extend. Rebased onto main's copy of the patch, which carries a later
revision of the _prepare_paged_hnd_input hunk and its tests; the interface.py
hunk offsets are adjusted for those three extra lines.

Verified with `patch -p1 --forward --batch` against MSA e2ebe76 (the pinned
submodule commit): all four files apply with no fuzz or offset, the reverse
dry run still succeeds, and the patched sources parse.

(cherry picked from commit 6e9c670)
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
brb-nv added a commit to brb-nv/TensorRT-LLM that referenced this pull request Aug 25, 2026
@brb-nv
[None][perf] Address decode-first assumption in MSA (NVIDIA#16923)
3ef5d90 
Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@pcicotti pcicotti pcicotti approved these changes
@zheyuf zheyuf Awaiting requested review from zheyuf
@peihu-nv peihu-nv Awaiting requested review from peihu-nv

Assignees

@brb-nv brb-nv

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@brb-nv @tensorrt-cicd @pcicotti