[TRTLLM-14054][perf] Reduce per-step host preparation overhead in the PyTorch executor decode path#16313
Conversation
kaiyux
changed the title
📝 WalkthroughWalkthroughChangesThe PR adds a steady-state generation fast path, bounds KV-cache block-offset staging to the batch’s required width, and replaces per-element Mamba state-index processing with bulk conversions and active-request mappings. Runtime preparation updates
Estimated code review effort: 3 (Moderate) | ~30 minutes Sequence Diagram(s)sequenceDiagram
participant PyTorchModelEngine
participant steady_gen_cache
participant attention_metadata
participant device_input_buffers
PyTorchModelEngine->>steady_gen_cache: check generation eligibility
PyTorchModelEngine->>attention_metadata: refresh per-step metadata
PyTorchModelEngine->>device_input_buffers: copy next input tokens
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
2151-2242: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueCROSS +
beam_width > 1branch doesn't forwardmax_blocks.
_stage_block_offsets_for_copy(num_rows)in the cross-attention branch (line 2177) omitsmax_blocks, so it always stages the fullmax_blocks_per_seqwidth even though the caller intrtllm.pynow computes and passes a tighter cap. This isn't a correctness issue (same "read-only-allocated-prefix" invariant still holds) but leaves the cross-attention/beam-width>1 combination without the intended H2D footprint reduction.♻️ Proposed fix
num_rows = num_context + num_gen_requests - host_block_offsets = self._stage_block_offsets_for_copy(num_rows) + host_block_offsets = self._stage_block_offsets_for_copy( + num_rows, max_blocks)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 2151 - 2242, Update the CROSS and beam_width > 1 path in copy_batch_block_offsets to pass the caller-provided max_blocks value into _stage_block_offsets_for_copy, matching the non-cross path. Preserve the existing row expansion and asynchronous copy behavior while ensuring staging uses the tighter block-width cap.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py`:
- Around line 1971-1974: Update the `_setup_state_indices` dummy-mask refresh to
use only the current `requests` batch; remove the fallback that rebuilds the
mask from `self.requests`. Preserve the existing `_refresh_dummy_request_mask`
behavior and rely on `get_state_indices()` for reordered batches.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 2151-2242: Update the CROSS and beam_width > 1 path in
copy_batch_block_offsets to pass the caller-provided max_blocks value into
_stage_block_offsets_for_copy, matching the non-cross path. Preserve the
existing row expansion and asynchronous copy behavior while ensuring staging
uses the tighter block-width cap.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: fc111ef9-2173-4dde-b0c7-4c08d1d97b24
📒 Files selected for processing (6)
tensorrt_llm/_torch/attention_backend/trtllm.pytensorrt_llm/_torch/modules/mamba/mamba2_metadata.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/resource_manager.py
|
/bot run --disable-fail-fast |
|
PR_Github #58981 [ run ] triggered by Bot. Commit: |
|
PR_Github #58981 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59066 [ run ] triggered by Bot. Commit: |
|
PR_Github #59066 [ run ] completed with state
|
pengbowang-nv
left a comment
pengbowang-nv
left a comment
There was a problem hiding this comment.
Attention modification LGTM
|
/bot run |
9aad9cb to
78c363c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59622 [ run ] triggered by Bot. Commit: |
|
[by Codex] @VALLIS-NERIA Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks! |
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot kill |
|
PR_Github #59680 [ kill ] triggered by Bot. Commit: |
|
PR_Github #59622 [ run ] completed with state |
The bulk copy_ raises on most length mismatches, but a length-1 sequence would broadcast silently. Assert the one-index-per-request contract explicitly so out-of-tree cache managers fail loudly. Addresses review feedback on NVIDIA#16313. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
|
PR_Github #59680 [ kill ] completed with state |
|
Need to rebase and relaunch CI pipeline after #16444 gets merged. |
… executor decode path At high concurrency, per-step host preparation in the overlap decode loop is exposed on the critical path. This change removes three sources of that overhead: - KVCacheManager.copy_batch_block_offsets staged block offsets through a freshly allocated pinned buffer sized for max_seq_len worth of block columns (~33.5 MB per step at batch size 256 with max_seq_len=262144), while the batch actually needed ~0.9 MB. Bound the staged/H2D copy width by the batch's maximum KV length (new optional max_blocks argument); kernels only read each sequence's allocated block prefix. - _prepare_tp_inputs re-walked the whole batch in Python every step even when the scheduled generation batch was unchanged. Cache the batch-composition-invariant state for non-speculative overlap decode (steady-state generation-only batches) and advance positions with one vectorized op, refreshing only per-step metadata. Any full prepare pass invalidates and re-records the cache. - Mamba cache-manager bookkeeping did one small tensor read or write per request (.item() calls and per-element stores). Replace them with one bulk tolist()/copy_ per step. Measured with trtllm-serve + benchmark_serving on Qwen3.6-35B-A3B-NVFP4 (hybrid gated-delta-net + MoE) on 1x B200, ISL=1024/OSL=6144, over a concurrency sweep 1..256: output throughput improves 3.17% on average across the 9-point curve (+1.58% to +4.88%, positive at every point). nsys at concurrency 256 shows _prepare_inputs going from 2.999 ms to 0.646 ms per decode step and per-step GPU idle from 8.6% to 2.5%. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
- Accept max_blocks in DeepseekV4CacheManager.copy_batch_block_offsets for signature parity with the base managers; TrtllmAttentionMetadata.prepare now always passes it. - Advance the device position buffer in place on fast steps instead of staging a mutated pinned buffer through an async H2D whose previous-step copy may still be pending under the overlap scheduler (the nvbug 6293536 hazard class); the pinned counter is host-side bookkeeping only. - Keep the (3,1,N) mrope position layout on the fast path: text-only batches on mrope models advance mrope_position_ids_cuda in place and return the mrope view, matching the full pass in eager, piecewise, and CUDA-graph modes. - Record the steady-gen cached-token counts from a pre-prepare snapshot so in-place clamping by sparse backends (RocketKV) is re-applied from true values each step instead of compounding into the recorded state. - Invalidate the steady-gen cache on the incremental-update early return so the two fast paths can never interleave. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…ve decoding The batch-max cap on the staged block-table width assumed the host kv_lens snapshot bounds every block-table access. That holds only for the non-speculative path: speculative decoding advances kv_lens_cuda on device past the host snapshot (one-model draft/tree sub-steps, overlap-scheduler kv-len offsets), and draft-token blocks are allocated ahead of the host kv_lens, so its kernels dereference block columns the cap left unstaged -- uninitialized values in the device buffer. CI hit this as an illegal memory access during EAGLE3 warmup (test_llama_eagle3_dynamic_tree on B200, test_eagle3_cdl_sampling on H100). Keep the cap only when no speculative-decoding state is attached to the metadata (no draft KV cache manager, spec-dec attention mode off, no extra KV tokens, no draft engine or drafter in the runtime); speculative setups stage the full pre-allocated width as before. Verified on B200: test_llama_eagle3_dynamic_tree[True-False] fails at the parent commit and passes with this change. New unit tests pin the staging width gate and the max_blocks copy semantics. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
…th cap Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
The bulk copy_ raises on most length mismatches, but a length-1 sequence would broadcast silently. Assert the one-index-per-request contract explicitly so out-of-tree cache managers fail loudly. Addresses review feedback on NVIDIA#16313. Signed-off-by: Kaiyu Xie <26294424+kaiyux@users.noreply.github.com>
e4212a2 to
190c73c
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #59701 [ run ] triggered by Bot. Commit: |
|
PR_Github #59701 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #59805 [ run ] triggered by Bot. Commit: |
|
PR_Github #59805 [ run ] completed with state |
8 participants
Description
At high concurrency, the per-step host preparation work in the PyTorch executor's overlap decode loop is exposed on the critical path: on Qwen3.6-35B-A3B-NVFP4 (hybrid gated-delta-net + MoE) at concurrency 256, nsys shows a recurring GPU-idle gap between decode steps dominated by
_prepare_inputs. This PR removes three sources of that overhead:Bound the KV-cache block-table staging to the batch's real width.
KVCacheManager.copy_batch_block_offsetsstaged block offsets through a freshly allocated pinned buffer sized formax_seq_lenworth of block columns — ~33.5 MB of pinned allocation + host memcpy + H2D copy per step at batch 256 withmax_seq_len=262144, of which only ~0.9 MB was live. A new optionalmax_blocksargument bounds the staged/H2D width by the batch's maximum KV length; kernels only read each sequence's allocated block prefix, and the persistent host buffer (used by synchronous CPU readers) is untouched. Default behavior is unchanged, andKVCacheManagerV2accepts the argument for signature parity (its copy already scales with allocated blocks only).Stop re-walking unchanged generation batches in
_prepare_tp_inputs. In non-speculative overlap decode, when the scheduled batch is exactly the same set of generation requests as the previous step, every request advanced by exactly one committed token. A full prepare pass now records the batch-invariant state (request ids, prompt lens, seq-lens, pinned position buffer) when the batch qualifies — generation-only, non-dummy, all requests carrying a previous-step device tensor, no spec decode / beam search / LoRA / multimodal / mrope work / attention-DP — and subsequent steps advance positions with one vectorized increment, gather input tokens from the previous step's device sample buffer, and refresh only per-step metadata. Any full pass invalidates and re-records the cache.Vectorize mamba cache-manager bookkeeping. Replace per-request
.item()reads and per-element pinned-buffer stores with one bulktolist()/copy_()per step.nsys at concurrency 256:
_prepare_inputsdrops from 2.999 ms to 0.646 ms per decode step, and per-step GPU idle from 8.6% to 2.5%.Test Coverage
_torchexecutor and attention-backend suites; the steady-state fast path engages for any non-speculative overlap-scheduler decode workload and is invalidated by any batch-composition change.trtllm-serve+benchmark_servingon Qwen3.6-35B-A3B-NVFP4, 1x B200, ISL 1024 / OSL 6144, random dataset with--ignore-eos, one point per concurrency (paired num_prompts 8...1024), measured at base8ef7190ebf:Output throughput improves +3.17% on average across the 9-point curve, positive at every point. Independent of #16314; on the same harness the two compose to about +5.6%.
Update (2026-07-13): hardening fixes + steady-state nsys A/B
The branch now carries a follow-up commit (
12ceb00216) hardening the fast prepare path from review findings:DeepseekV4CacheManager.copy_batch_block_offsetsacceptsmax_blocksfor signature parity (the baseprepare()now always passes it).(3,1,N)position layout on the fast path: text-only batches advancemrope_position_ids_cudain place and return the mrope view, matching the full pass in eager, piecewise, and CUDA-graph modes (Qwen3.5/3.6 normalizerope_parameterswithmrope_sectionintorope_scaling.type='mrope', so they take this path even for pure-text serving).num_cached_tokens_per_seqbeforeattn_metadata.prepare(), so backends that clamp the aliased list in place (RocketKV) re-clamp from true values every step instead of compounding into the recorded state.Fresh nsys A/B at the head of this branch vs its merge base
53d2bde62f— same GPU, same prebuilt binaries (python-only diff), Qwen3.6-35B-A3B-NVFP4 on 1x B200,trtllm-serve(CUDA graphsmax_batch_size: 256+ padding, TRTLLM MoE backend) +benchmark_servingrandom ISL 1024 / OSL 6144 with--ignore-eos, c=256 / num_prompts=1024, capture windowTLLM_PROFILE_START_STOP=100-150(-t cuda,nvtx,python-gil -c cudaProfilerApi --cuda-graph-trace node). Both captures contain exactly 50 forward steps, every one0 ctx reqs, 0 ctx tokens, 256 gen reqs— pure steady-state decode at the full batch:_prepare_inputshost time / stepThe H2D reduction is the bounded block-table staging; the prepare-time and idle reduction is the steady-state fast prepare, which engaged on all 50 steps — on an mrope-normalized model, so the window also exercises the mrope in-place advance under CUDA-graph replay. Correctness probe: a fixed greedy completion (temperature 0, 64 new tokens, fresh server per side) is byte-identical between the two sides.
Update (2026-07-14): CI fix — full-width block-table staging for speculative decoding
Pre-merge CI caught an illegal memory access during EAGLE3 warmup (
test_llama_eagle3_dynamic_tree[True-False]on B200,test_eagle3_cdl_sampling[True]on H100, plus follow-on thread-leak failures intest_eagle3_cuda_graph_padding). Root cause: the batch-max cap on the staged block-table width assumed the hostkv_lenssnapshot bounds every block-table access. That holds only for the non-speculative path — speculative decoding advanceskv_lens_cudaon device past the host snapshot (one-model draft/tree sub-steps, overlap-scheduler kv-len offsets), and draft-token blocks are allocated ahead of the host kv_lens, so spec kernels dereferenced block columns the cap left unstaged (uninitialized memory in the device buffer).The new commit keeps the cap only when no speculative-decoding state is attached to the metadata (no draft KV cache manager, spec-dec attention mode off, no extra KV tokens, no draft engine or drafter in the runtime); speculative setups stage the full pre-allocated width exactly as before this PR. The non-speculative serving path that motivated the optimization keeps the bounded staging and its measured wins.
Verified on 1x B200 against
LLM_MODELS_ROOTcheckpoints:test_llama_eagle3_dynamic_tree[True-False]fails at the parent commit and passes with the fix;test_eagle3_cdl_sampling,test_eagle3_cuda_graph_padding,test_llama_eagle3[True-TRTLLM-False-True-True-False-True-False-False-False](spec-vs-ref equality with the overlap-scheduler ref exercising the steady-state fast prepare), andtest_kv_block_offset_overlap_race.pyall pass. New unit tests pin the staging-width gate (test_block_offsets_staging_width_spec_gate) and themax_blockscopy semantics (test_copy_batch_block_offsets_max_blocks_staging_width).PR Checklist
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.🤖 Generated with Claude Code
Summary by CodeRabbit
Performance Improvements
Bug Fixes