[None][perf] Address inter-iter idle times by brb-nv · Pull Request #16875 · NVIDIA/TensorRT-LLM

brb-nv · 2026-07-26T04:13:44Z

Description

This MR does multiple perf optimizations to address inter-iter idle times:

Page table and out_cache_loc are now built on the host from the block ids, replacing per-request device kernels and per-token .item() reads
Block ids, length mirrors and query-row staging tensors are pinned and copied non-blocking instead of landing in pageable memory.
on_update_kv_lens patches plan lengths entirely on device for every step, dropping the prefill/mixed host copy-back and plan rebuild.
The MSA planner keeps its lengths host-side and loses its blocking copies and .item(), and sparse plans now patch seqused_k - the length their kernels actually mask with.
Two correctness guards: corrected kv_lens are clamped to the extent prepare()` planned for, and the stale host empty-block flag is gone.

At TP=4, ISL=8192, OSL=128, c=320, request latency in ms.

Baseline

[Latency] P50    : 72969.3661
[Latency] P90    : 87000.6669
[Latency] P95    : 96023.0939
[Latency] P99    : 103330.9747
[Latency] MINIMUM: 32364.5672
[Latency] MAXIMUM: 105018.2742
[Latency] AVERAGE: 69344.3009

Feature

[Latency] P50    : 60563.0345
[Latency] P90    : 72547.7751
[Latency] P95    : 80037.2101
[Latency] P99    : 86102.7175
[Latency] MINIMUM: 27184.2698
[Latency] MAXIMUM: 87501.6225
[Latency] AVERAGE: 57888.3157

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.

get_block_ids_per_seq built one tensor per request and joined them with pad_sequence. That is a per-request Python loop, and pad_sequence returns pageable memory, so the copy that stages the result to the device blocks until the CUDA queue drains. At batch 128 the loop alone shows up as a 2.1 ms host-only gap in every decode step. Fill a single zero-filled pinned int32 tensor through a numpy view instead. The loop cost disappears and callers can stage the result with a non-blocking copy. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

build_kv_page_indices recovered each request's page ids by gathering the first slot of every page out of req_to_token on the device. That is five tiny kernels per request per step: at batch 128 it is one contiguous burst of 644 launches spanning 10.6 ms, during which the GPU starves. The work is redundant. Both callers use the manager's tokens_per_block as the page size, and req_to_token is built as block_id * tokens_per_block + offset, so req_to_token[b, p * page_size] // page_size is just block_ids[b, p] -- a value get_block_ids_per_seq already returned on the host. Build the table there instead and hand it back through the slot mapping, which now also exposes block_ids_cpu so no second manager query or device round trip is needed. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

out_cache_loc was built by reading req_to_token one new token at a time with .item(). At 6.6k context tokens that is a 17.4 ms window with no CUDA activity at all, and it scales with context length. It also forced a ~4.3 MB device-to-host copy of the whole req_to_token grid that existed only to feed the loop. The same slot ids follow from the host block ids by expanding the per-request lengths, so compute them there and stage the result with one asynchronous copy. Optimistic offsets can overhang the last allocated slot and a zero-length padding row lands at -1, so positions are clamped into the row; those slots are placeholders that on_update_kv_lens re-derives before any forward reads them. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

The per-request length mirrors and the query-row staging tensors were plain pageable allocations. Every copy_ that ships them to the device therefore ends in a cudaStreamSynchronize that drains the whole queue, which is what stops the host from running ahead of the GPU. Pin them where pinning helps and stage them non-blocking. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

sparse_fmha_plan moved its length inputs to the device with blocking copies and read total_q back through cu_seqlens_q[-1].item(). Planning sits on the critical path of input preparation, so each of those drains the CUDA queue. Stage the lengths to the device once, asynchronously, and take every scalar from the host copy instead. kv_segment_lens now stays on the host in the returned plan: only the page-table builder consumes it, and it reads per-request page counts on the host, while the kernels take their lengths from cu_seqlens_k and seqused_k. That subsumes _stage_sparse_plan_kv_lens_host, which patched the same field from the backend after the fact to dodge a per-layer D2H, so it is dropped. The MSA change ships as an additive update to the existing cumulative patch against submodule commit e2ebe76; the prior hunks are unchanged. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

_build_step_plans computed, on the host, whether an eager step selects no KV blocks at all, and run_indexer returned an all -1 selection when it did. The flag is derived from the optimistic lengths prepare() staged. The overlap scheduler then corrects kv_lens on device, and on_update_kv_lens cannot recompute a host bool without a device read, so the flag can be consulted after it has gone stale. Remove it. The staged per-token counts are clamped to at least one block, matching the decode path, and the kernel already masks each query to its own valid-block extent, so the empty case needs no special handling. The clamp is also what keeps a fully-masked row from producing a NaN. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

on_update_kv_lens rests on the correction only ever shrinking lengths: that is what lets the host-baked plan worklists and the page-table layout survive it. Nothing enforced that. A length that grew instead would drive the kernels past the extent prepare() planned for and read slots that were never mapped, silently. Snapshot the staged lens in prepare() and clamp the corrected lens against them. The snapshot is a device-to-device copy and the clamp is a torch.minimum, so the check costs no synchronization. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

Prefill and mixed steps handled the corrected lengths by copying kv_lens_cuda back to the host and re-running _build_msa_fields and _build_step_plans. The copy targets pageable memory, so it is a synchronizing 636-byte transfer that costs ~64 ms of host time per step, and the rebuild redoes work the correction cannot invalidate. Decode already patched the same state with device ops only. Extend that to every step: iterate whichever plan set is live, and patch each sub-plan's length rows in place from the corrected lengths. A mixed batch is split at the first long request, so each sub-plan covers its own request range and the row count says whether it is indexed per request or per query token. The token offsets come from a host-side tuple built in prepare(), so no device read is needed to slice a range. This also removes the corrected-host-lens override on msa_kv_lens_cpu, which existed only to feed the rebuild. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

The length mirroring wrote kv_segment_lens and qo_offset into every plan. For sparse plans that is wrong in both directions. kv_segment_lens is host-resident there -- only the page-table builder reads it, and it did so at plan time -- so writing a device tensor into it is a pointless copy that synchronizes the host, which is the stream sync left over after the D2H removal. Meanwhile seqused_k, the length sparse kernels actually mask with, was never patched and kept the optimistic value. Select the keys by plan flavour off the MM-SA-Nv tag instead of writing a fixed pair and letting a device check silently skip the host entries. seqused_k joins the graph-stable key set, since it is now a patch target and must keep a stable device buffer across replays. A missing mirror raises rather than being skipped: it means corrected lengths never reach the kernel, which is a silent accuracy bug. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

Two run_indexer tests raise AttributeError before reaching an assertion. The FP8 index-query test's fake predates head-major routing and never defines num_contexts or num_generations, which run_indexer now reads unconditionally to pick the output layout. It stands in for a two-request decode, so it gets 0 and 2. The head-major routing test's fake returns None from msa_idx_k_cache until msa_write_idx_k populates it, but run_indexer reads the cache before it writes, so the local is always None when the FP8 check dereferences it. The production cache always exists, so the fake holds a tensor from the start and the write copies into it. Also drops msa_eager_all_blocks_empty from that fake; nothing reads it. Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

…wiring The MiniMax-M3 decode path currently runs its generation rows through fmha_sm100, which schedules a generation row like a context row. Three kernels ported from vLLM replace that: a CuTe DSL indexer scoring kernel, a Triton sparse block decode kernel and a trtllm-gen dense decode kernel. This lands the kernels, their custom-op registration and their correctness tests. Nothing dispatches to them yet: the MSA backend, indexer and cache manager are untouched, so the decode path is byte-for-byte what it was and the kernels are reachable only from the tests and the microbenchmark. The dispatch is a follow-up, since it rests on the device-side length patching and the fused per-layer cache writes that are still landing on feat/m3_with_msa. Split out of NVIDIA#17268 on feat/m3_with_msa, which carries the same kernels plus that wiring. Two additions differ from it. msa_indexer gains only cutedsl_score_runner and _cutedsl_score, the self-contained entry points the scorer test drives, and not the run_indexer dispatch that calls them. The tests reach fmha_sm100 through a local _flat_page_table helper, because build_kv_page_indices does not take a block table until NVIDIA#16875; the helper feeds it the slot map that block table implies, so the A/B comparisons still run against the production page-table builder rather than a test-local copy. The CuTe DSL indexer decode kernel and its tests were originally contributed to vLLM by Thien Tran (vllm-project/vllm#48582), as were the CuTe utilities (vllm-project/vllm#43273). The Triton sparse decode kernel and its tests were originally contributed to vLLM by Kaichao You (vllm-project/vllm#45381). Thanks to both. No test-list change: l0_b300 already collects unittest/_torch/attention wholesale, so the three new files are picked up there, and each skips itself off SM100/SM103. (cherry picked from commit 727c683) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

Input preparation for a MiniMax-M3 MSA step spent most of its host time in work that either blocked on the CUDA queue or launched a burst of tiny kernels, leaving the GPU idle between iterations: * get_block_ids_per_seq built one tensor per request and joined them with pad_sequence, a per-request Python loop landing in pageable memory. It now fills a single zero-filled pinned int32 tensor through a numpy view. * build_kv_page_indices recovered each request's page ids by gathering the first slot of every page out of req_to_token on the device, five tiny kernels per request per step. The work is redundant: req_to_token is built as block_id * tokens_per_block + offset, so req_to_token[b, p * page_size] // page_size is just block_ids[b, p], a value the host already has. The table is built there instead and the slot mapping hands the host block ids back so no second manager query or device round trip is needed. * out_cache_loc was built by reading req_to_token one new token at a time with .item(), a window that scales with context length. The same slot ids follow from the host block ids by expanding the per-request lengths, so they are computed there and staged with one asynchronous copy. * The per-request length mirrors are pinned and staged non-blocking, so their copies no longer end in a cudaStreamSynchronize. * sparse_fmha_plan staged its length inputs with blocking copies and read total_q back through cu_seqlens_q[-1].item(). It now stages once, asynchronously, and takes every scalar from the host copy. kv_segment_lens stays on the host in the returned plan: only the page-table builder consumes it, and it reads per-request page counts on the host, while the kernels take their lengths from cu_seqlens_k and seqused_k. That subsumes _stage_sparse_plan_kv_lens_host, which patched the same field from the backend after the fact, so it is dropped. Cherry-picked from NVIDIA#16875 on feat/m3_with_msa. The parts of that PR that turn on an MSA on_update_kv_lens override are dropped: main's MSA metadata has no such override, nor the device mirrors (msa_req_to_token, msa_q_batch_row, msa_q_intra, msa_qo_lens_dev) it patches through, so there is nothing here to make sync-free, clamp against a staged snapshot, or redirect from kv_segment_lens to seqused_k. The eager empty-block short-circuit that PR removes likewise does not exist on main, so its replacement clamp is left out and the eager valid-block counts keep main's unclamped staging. The two run_indexer test fakes it repairs cover head-major routing and the FP8 index-K producer, neither of which main carries. (cherry picked from commit a69907a) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

…wiring The MiniMax-M3 decode path currently runs its generation rows through fmha_sm100, which schedules a generation row like a context row. Three kernels ported from vLLM replace that: a CuTe DSL indexer scoring kernel, a Triton sparse block decode kernel and a trtllm-gen dense decode kernel. This lands the kernels, their custom-op registration and their correctness tests. Nothing dispatches to them yet: the MSA backend, indexer and cache manager are untouched, so the decode path is byte-for-byte what it was and the kernels are reachable only from the tests and the microbenchmark. The dispatch is a follow-up, since it rests on the device-side length patching and the fused per-layer cache writes that are still landing on feat/m3_with_msa. Split out of NVIDIA#17268 on feat/m3_with_msa, which carries the same kernels plus that wiring. Two additions differ from it. msa_indexer gains only cutedsl_score_runner and _cutedsl_score, the self-contained entry points the scorer test drives, and not the run_indexer dispatch that calls them. The tests reach fmha_sm100 through a local _flat_page_table helper, because build_kv_page_indices does not take a block table until NVIDIA#16875; the helper feeds it the slot map that block table implies, so the A/B comparisons still run against the production page-table builder rather than a test-local copy. The CuTe DSL indexer decode kernel and its tests were originally contributed to vLLM by Thien Tran (vllm-project/vllm#48582), as were the CuTe utilities (vllm-project/vllm#43273). The Triton sparse decode kernel and its tests were originally contributed to vLLM by Kaichao You (vllm-project/vllm#45381). Thanks to both. No test-list change: l0_b300 already collects unittest/_torch/attention wholesale, so the three new files are picked up there, and each skips itself off SM100/SM103. (cherry picked from commit 727c683) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

Input preparation for a MiniMax-M3 MSA step spent most of its host time in work that either blocked on the CUDA queue or launched a burst of tiny kernels, leaving the GPU idle between iterations: * get_block_ids_per_seq built one tensor per request and joined them with pad_sequence, a per-request Python loop landing in pageable memory. It now fills a single zero-filled pinned int32 tensor through a numpy view. * build_kv_page_indices recovered each request's page ids by gathering the first slot of every page out of req_to_token on the device, five tiny kernels per request per step. The work is redundant: req_to_token is built as block_id * tokens_per_block + offset, so req_to_token[b, p * page_size] // page_size is just block_ids[b, p], a value the host already has. The table is built there instead and the slot mapping hands the host block ids back so no second manager query or device round trip is needed. * out_cache_loc was built by reading req_to_token one new token at a time with .item(), a window that scales with context length. The same slot ids follow from the host block ids by expanding the per-request lengths, so they are computed there and staged with one asynchronous copy. * The per-request length mirrors are pinned and staged non-blocking, so their copies no longer end in a cudaStreamSynchronize. * sparse_fmha_plan staged its length inputs with blocking copies and read total_q back through cu_seqlens_q[-1].item(). It now stages once, asynchronously, and takes every scalar from the host copy. kv_segment_lens stays on the host in the returned plan: only the page-table builder consumes it, and it reads per-request page counts on the host, while the kernels take their lengths from cu_seqlens_k and seqused_k. That subsumes _stage_sparse_plan_kv_lens_host, which patched the same field from the backend after the fact, so it is dropped. Cherry-picked from NVIDIA#16875 on feat/m3_with_msa. The parts of that PR that turn on an MSA on_update_kv_lens override are dropped: main's MSA metadata has no such override, nor the device mirrors (msa_req_to_token, msa_q_batch_row, msa_q_intra, msa_qo_lens_dev) it patches through, so there is nothing here to make sync-free, clamp against a staged snapshot, or redirect from kv_segment_lens to seqused_k. The eager empty-block short-circuit that PR removes likewise does not exist on main, so its replacement clamp is left out and the eager valid-block counts keep main's unclamped staging. The two run_indexer test fakes it repairs cover head-major routing and the FP8 index-K producer, neither of which main carries. (cherry picked from commit a69907a) Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

…wiring The MiniMax-M3 decode path currently runs its generation rows through fmha_sm100, which schedules a generation row like a context row. Three kernels ported from vLLM replace that: a CuTe DSL indexer scoring kernel, a Triton sparse block decode kernel and a trtllm-gen dense decode kernel. This lands the kernels, their custom-op registration and their correctness tests. Nothing dispatches to them yet: the MSA backend, indexer and cache manager are untouched, so the decode path is byte-for-byte what it was and the kernels are reachable only from the tests and the microbenchmark. The dispatch is a follow-up, since it rests on the device-side length patching and the fused per-layer cache writes that are still landing on feat/m3_with_msa. Split out of NVIDIA#17268 on feat/m3_with_msa, which carries the same kernels plus that wiring. Two additions differ from it. msa_indexer gains only cutedsl_score_runner and _cutedsl_score, the self-contained entry points the scorer test drives, and not the run_indexer dispatch that calls them. The tests reach fmha_sm100 through a local _flat_page_table helper, because build_kv_page_indices does not take a block table until NVIDIA#16875; the helper feeds it the slot map that block table implies, so the A/B comparisons still run against the production page-table builder rather than a test-local copy. The CuTe DSL indexer decode kernel and its tests were originally contributed to vLLM by Thien Tran (vllm-project/vllm#48582), as were the CuTe utilities (vllm-project/vllm#43273). The Triton sparse decode kernel and its tests were originally contributed to vLLM by Kaichao You (vllm-project/vllm#45381). Thanks to both. No test-list change: l0_b300 already collects unittest/_torch/attention wholesale, so the three new files are picked up there, and each skips itself off SM100/SM103. (cherry picked from commit 727c683) 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 added 10 commits July 25, 2026 19:34

brb-nv requested review from a team as code owners July 26, 2026 04:13

brb-nv requested review from PerkzZheng, pcicotti, peihu-nv, pengbowang-nv and zheyuf and removed request for a team, PerkzZheng and pengbowang-nv July 26, 2026 04:13

github-actions Bot assigned brb-nv Jul 26, 2026

pcicotti approved these changes Jul 26, 2026

View reviewed changes

pcicotti merged commit d8b09fc into NVIDIA:feat/m3_with_msa Jul 26, 2026
8 checks passed

ZhanruiSunCh pushed a commit that referenced this pull request Jul 29, 2026

[None][perf] Address inter-iter idle times (#16875)

a69907a

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

This was referenced Aug 17, 2026

[None][feat] Add the ported MiniMax-M3 decode kernels ahead of their wiring #17842

Merged

[None][perf] Address inter-idle times and decode-first assumption in MSA #17986

Merged

brb-nv added a commit to brb-nv/TensorRT-LLM that referenced this pull request Aug 25, 2026

[None][perf] Address inter-iter idle times (NVIDIA#16875)

ef89dc9

Signed-off-by: Balaram Buddharaju <169953907+brb-nv@users.noreply.github.com>

Navigation Menu

Uh oh!

[None][perf] Address inter-iter idle times - #16875

[None][perf] Address inter-iter idle times#16875
pcicotti merged 10 commits into
NVIDIA:feat/m3_with_msafrom
brb-nv:user/brb/address-inter-iter-delays-redo

brb-nv commented Jul 26, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

brb-nv commented Jul 26, 2026

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants