[None][fix] Fix CppMambaHybridCacheManager functional and perf issues by VALLIS-NERIA · Pull Request #14003 · NVIDIA/TensorRT-LLM

VALLIS-NERIA · 2026-05-11T15:34:40Z

Summary

Fixes a set of functional and performance issues in CppMambaHybridCacheManager and the recurrent-state cache path.

Functional fixes

Zero local mamba layers under PP sharding (ac30ded849)
Under pipeline-parallel sharding, a rank can end up with only full-attention layers and no mamba layers. The previous constructor unconditionally allocated SSM/conv state and set up linear-attention metadata, leaving such ranks in an invalid state. Added an early-exit after super().__init__ when local_num_mamba_layers == 0, forwarded the union of mamba+attention layer masks to the parent KVCacheManager, and guarded prepare_resources / update_mamba_states to no-op on those ranks.
Reserve recurrent-state slots for CUDA graph padding and draft-len sentinels (992ec1bd08)
The recurrent-state snapshot pool (_calculate_max_num_blocks_for_linear_attention → max_snapshots) was sized at exactly max_batch_size, which doesn't reserve room for:
- The CUDA graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID) — needs 1 extra slot.
- The per-draft-len sentinels that CUDAGraphRunner::_get_padded_batch issues during speculative decoding — needs one extra slot per draft length.
Under load, this caused the padding sentinels to evict live recurrent state and surface as block-allocation failures. Added +1 for the CUDA graph padding sentinel and + spec_config.max_draft_tokens when a spec_config is present. The block-reuse override (max_tokens // interval) still wins when active.
Follow-ups on the slot-reservation fix
- 3310de0a0e — fixed a max_draft_tokens typo and added unit tests covering the slot reservation logic.
- 5a858312ce — store spec_config on self in __init__ so the override path can read it.
- 5aad39a889 — addressed CodeRabbit review comments on the above.

Performance fixes (`516ea84eaa`)

Several host/device sync points were eliminated from the mamba-hybrid prep path. With Qwen3.5-A17B-NVFP4 + MTP Eagle one-model + CUDA graphs, nsys showed numGenReq × (1 + max_draft_len) cudaMemcpyAsync + cudaStreamSynchronize pairs per iteration sitting inside _prepare_inputs, plus a refresh_blocks stream sync at the tail of prepare_resources that blocked subsequent prep work.

Mamba2Metadata.state_indices: alias the cache manager's CUDA buffer instead of per-element D2H.
The previous code did for i, idx in enumerate(indices): state_indices_cpu[i] = idx, where indices is CppMambaHybridCacheManager.cuda_state_indices (a CUDA tensor). Iterating a CUDA tensor produces 0-d CUDA slices, and assigning each to a CPU tensor issues a cudaMemcpyAsync + cudaStreamSynchronize per element. Replaced with a direct reference assignment (self.state_indices = indices) when the source is on CUDA, guarded by a data_ptr() invariant assertion to catch future buffer reallocations that would silently break CUDA graph replays. CPU/list paths keep the original behavior.
KVCacheTransferManager::copyBlock layer-first path: one cudaMemcpy2DAsync instead of N async memcpys.
Pool layout is {numLayers, numBlocks, kvFactor, blockSize}. For a fixed block index, per-layer slices are equal-length rows separated by stride numBlocks * rowBytes. Replaced the per-layer cudaMemcpyAsync loop with a single pitched 2D copy.
Defer refresh_blocks() (stream sync) until just before forward.
Split CppMambaHybridCacheManager._prepare_resources into two phases: (1) fire all copy_linear_attention_block onboards asynchronously and record whether any actually scheduled a transfer, and (2) a new flush_state_transfers() method that calls refresh_blocks() only when there are pending transfers. The flush is invoked at the end of Mamba2Metadata.prepare(), so the rest of _prepare_tp_inputs (host work + H2D copies) overlaps with the in-flight onboard transfers.

To support the "skip if nothing scheduled" optimization, KVCacheManager::copyLinearAttentionBlock now returns bool indicating whether at least one onboard was actually issued. The 3 layers of the call (KVCacheManager → BlockManager → WindowBlockManager) all return bool; pybind reflects the new return type automatically.

Test plan

Unit tests: tests/unittest/_torch/executor/test_mamba_cache_manager.py covers slot reservation across CUDA-graph + draft-len combinations.
Integration: end-to-end CUDA graph + MTP on a mamba hybrid model no longer hits "no free recurrent-state blocks".
Perf: nsys on Qwen3.5-A17B-NVFP4 + MTP Eagle one-model + CUDA graphs shows the per-(gen_req × (draft+1)) cudaMemcpyAsync + sync pairs gone, and the refresh_blocks stream sync moved out of the critical path.

Summary by CodeRabbit

Bug Fixes
- Improved memory allocation calculation for linear attention operations to better account for CUDA graph padding and speculative decoding configurations, enhancing overall resource management efficiency.

…mamba layers (By Agent) Under PP sharding, a rank can end up with only full-attention layers and no mamba layers. The previous constructor unconditionally allocated SSM / conv state and set up linear-attention metadata, leading to invalid state on these ranks. Take an early-exit after super().__init__ when local_num_mamba_layers == 0: forward num_layers (not mamba_num_layers + num_layers) and the union of mamba+attention layer masks to the parent KVCacheManager, skip all mamba-only setup. Guard prepare_resources / update_mamba_states / _setup_state_indices with the same condition so they no-op on these ranks. Add a regression test that constructs the manager with a real parent KVCacheManager and verifies: early-exit invariants on mgr state, no mamba-only attributes set, and the three guarded methods no-op without raising. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

…aph padding and draft-len sentinels (By Agent) The recurrent-state snapshot pool was sized at exactly max_batch_size, which leaves no room for the CUDA graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID) and the per-draft-len sentinels that CUDAGraphRunner::_get_padded_batch issues during speculative decoding. Reserve one extra slot for the CUDA graph padding dummy, and when a spec_config is present, reserve one additional slot per draft length so all distinct sentinel request IDs can coexist without evicting live recurrent state. The block-reuse path (max_tokens // interval) still overrides this when active. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-11T15:34:51Z

/bot run

coderabbitai · 2026-05-11T15:38:20Z

📝 Walkthrough

Walkthrough

Resource manager updates the recurrent-state block pool size calculation to account for CUDA graph padding (+1) and speculative decoding draft tokens. The revised max_snapshots value is used to derive block allocations unless overwritten by existing block-reuse logic.

Changes

Recurrent-State Block Pool Sizing

Layer / File(s)	Summary
Recurrent-State Max Snapshots Calculation `tensorrt_llm/_torch/pyexecutor/resource_manager.py`	`max_snapshots` for recurrent-state blocks is initialized with `self.max_batch_size + 1` (CUDA graph padding), then conditionally increased by `self.spec_config.max_draft_tokens` when speculative decoding is configured, before the existing block-reuse interval logic may further modify it.

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The PR title references 'CppMambaHybridCacheManager' but the actual changes are in the recurrent-state snapshot pool calculation for LinearCacheType, not specifically about CppMambaHybridCacheManager functionality as the title suggests.	Update the title to reflect the actual change: something like '[None][fix] Reserve recurrent-state slots for CUDA graph padding and draft-length sentinels' to match the detailed objectives.
Description check	⚠️ Warning	The PR description is largely incomplete and does not follow the required template structure.	Rewrite the description using the template: add explicit sections for Description, Test Coverage, and a completed PR Checklist with all required items verified.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

🤖 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/resource_manager.py`:
- Around line 1560-1563: The code uses self.spec_config in the snapshot sizing
branch but KVCacheManager.__init__ never assigns the incoming spec_config to the
instance; update KVCacheManager.__init__ (the constructor) to store the argument
(e.g., self.spec_config = spec_config) before any logic that may reference it so
the branch that adjusts max_snapshots using self.spec_config.max_draft_tokens
will not raise AttributeError.

🪄 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: 7bb47e0c-e115-4317-a5f3-42c778f8b88e

📥 Commits

Reviewing files that changed from the base of the PR and between f3570a8 and 992ec1b.

📒 Files selected for processing (1)

tensorrt_llm/_torch/pyexecutor/resource_manager.py

tensorrt-cicd · 2026-05-11T15:40:26Z

PR_Github #47767 [ run ] triggered by Bot. Commit: 992ec1b Link to invocation

…servation tests (By Agent) Follow-up to the recurrent-state slot reservation: spec_config has no max_draft_tokens attribute (DecodingBaseConfig uses StrictBaseModel with extra="forbid"). The correct field is max_draft_len, which is also the count of additional dummy request ids issued by CUDAGraphRunner._get_padded_batch. Add two unit tests that construct a real CppMambaHybridCacheManager and inspect blocks_per_window[RECURRENT_STATES]: - without spec_config: pool reserves max_batch_size + 1 slots - with spec_config(max_draft_len=N): pool reserves max_batch_size + 1 + N Tests disable block reuse so the override (max_tokens // interval) does not mask the addition under test. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-11T15:48:11Z

/bot run

tensorrt-cicd · 2026-05-11T15:55:03Z

PR_Github #47768 [ run ] triggered by Bot. Commit: 3310de0 Link to invocation

tensorrt-cicd · 2026-05-11T18:51:13Z

PR_Github #47768 [ run ] completed with state SUCCESS. Commit: 3310de0
/LLM/main/L0_MergeRequest_PR pipeline #37660 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-05-12T01:05:54Z

/bot run

tensorrt-cicd · 2026-05-12T01:12:03Z

PR_Github #47812 [ run ] triggered by Bot. Commit: 3310de0 Link to invocation

… Agent) Address CodeRabbit review on NVIDIA#14003: the linear-attention sizing branch reads self.spec_config, but KVCacheManager.__init__ never assigned the incoming spec_config argument to self. The code worked only by accident, because CppMambaHybridCacheManager's _setup_mtp_intermediate_states runs before super().__init__ and writes self.spec_config on the partially initialized instance. Any other KVCacheManager path that exercises the linear-attention branch with a spec_config would AttributeError. Assign self.spec_config = spec_config at the top of KVCacheManager.__init__ so the attribute is always present regardless of subclass. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-12T01:41:51Z

Addressed CodeRabbit feedback: store spec_config on self in KVCacheManager.__init__ so the linear-attention sizing branch doesn't rely on a subclass setting it before super().__init__.

/bot run

VALLIS-NERIA · 2026-05-12T02:00:47Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-12T02:06:25Z

PR_Github #47826 [ run ] triggered by Bot. Commit: 5a85831 Link to invocation

… Agent) Two fixes from CodeRabbit's review on NVIDIA#13999: 1. Handle layer_mask=None defensively. The signature has always typed it as Optional[List[bool]], and the new constructor was unconditionally calling layer_mask.copy() / indexing before the zero-mamba guard ran. Treat None as an all-False full-attention mask and validate length against mamba_layer_mask. 2. Seed externally visible mamba fields before the early-return guard. get_mamba_ssm_cache_dtype() reads self.ssm_state_dtype, and the use_replay_state_update property reads self._use_replay_state_update; both were only assigned on the non-early-exit path, so any caller that hit those accessors on a zero-local-mamba rank would have AttributeError'd. Move both assignments above the early-return. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

tensorrt-cicd · 2026-05-12T06:59:16Z

PR_Github #47826 [ run ] completed with state SUCCESS. Commit: 5a85831
/LLM/main/L0_MergeRequest_PR pipeline #37711 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…recurrent-cuda-graph-snapshots (By Agent) Brings in the CppMambaHybridCacheManager zero-local-mamba-layers fix and its CodeRabbit follow-ups (PR NVIDIA#13999). Conflict in tests/unittest/_torch/executor/test_mamba_cache_manager.py is resolved by keeping both test sections (recurrent-state pool sizing from this branch, zero-local-mamba early-exit from the merged branch) and unioning the import set. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

… Agent) Address CodeRabbit review on NVIDIA#14003: the linear-attention sizing branch reads self.spec_config, but KVCacheManager.__init__ never assigned the incoming spec_config argument to self. The code worked only by accident, because CppMambaHybridCacheManager's _setup_mtp_intermediate_states runs before super().__init__ and writes self.spec_config on the partially initialized instance. Any other KVCacheManager path that exercises the linear-attention branch with a spec_config would AttributeError. Assign self.spec_config = spec_config at the top of KVCacheManager.__init__ so the attribute is always present regardless of subclass. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

tensorrt-cicd · 2026-05-18T06:30:13Z

PR_Github #48803 [ run ] completed with state SUCCESS. Commit: 7a6e90f
/LLM/main/L0_MergeRequest_PR pipeline #38567 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…dels (By Agent) `_calculate_max_num_blocks_for_linear_attention` sizes the recurrent-state slot pool as `max_batch_size + 1`, ignoring `pp_size`. With pipeline parallelism (e.g. `TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]`), multiple microbatches are in-flight on the same rank simultaneously, each holding up to `max_batch_size` sequences' Mamba state — so the live-state slot count must be `max_batch_size * pp_size + 1`. The KVCacheManagerV2 path already does this (`max_num_sequences = max_batch_size * pp_size`); this commit aligns the linear-attention sizing path with that invariant. Three call sites are fixed: 1. `KVCacheManager._calculate_max_num_blocks_for_linear_attention` — both `intercept` (memory reservation in the affine model) and `max_snapshots` (slot count) now scale with `pp_size`. 2. The `is_linear_attention` branch in the estimation dry-run code path, which previously used a bare `max_batch_size`. 3. `CppMambaHybridCacheManager.get_cache_size_per_token` — the per-request fixed-cost intercept that drives upstream `_tokens_for_budget`. Without this fix, `PP > 1` runs trip `No free block found. This shouldn't happen!` in `evictionPolicy::getFreeBlock` once requests beyond the first microbatch arrive at the pool. Added a regression unit test that asserts `recurrent_primary >= max_batch_size * pp_size + 1` for `pp_size ∈ {2, 4}`, and verified the fix end-to-end with the TP4_PP2 model run on a TP2_PP2 configuration locally (test passed in 8m53s, no `No free block` error). Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA · 2026-05-18T15:46:26Z

/bot run --disable-fail-fast --stage-list "DGX_B200-4_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-2"

tensorrt-cicd · 2026-05-18T15:53:07Z

PR_Github #48927 [ run ] triggered by Bot. Commit: c9933a8 Link to invocation

tensorrt-cicd · 2026-05-18T17:24:01Z

PR_Github #48927 [ run ] completed with state FAILURE. Commit: c9933a8
/LLM/main/L0_MergeRequest_PR pipeline #38675 (Partly Tested) completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

VALLIS-NERIA · 2026-05-20T02:31:59Z

/bot run --disable-fail-fast --stage-list "DGX_B200-4_GPUs-PyTorch-2, DGX_B200-8_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-1, GB200-4_GPUs-PyTorch-2"

tensorrt-cicd · 2026-05-20T02:38:15Z

PR_Github #49315 [ run ] triggered by Bot. Commit: c9933a8 Link to invocation

tensorrt-cicd · 2026-05-20T13:16:01Z

PR_Github #49315 [ run ] completed with state SUCCESS. Commit: c9933a8
/LLM/main/L0_MergeRequest_PR pipeline #38976 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

VALLIS-NERIA · 2026-05-21T02:00:09Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T02:06:36Z

PR_Github #49535 [ run ] triggered by Bot. Commit: c9933a8 Link to invocation

yuxianq

tensorrt_llm/_torch/pyexecutor/_util.py LGTM

commit c9933a8 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 18 17:58:20 2026 +0800 [None][fix] Scale recurrent-state pool by pp_size for hybrid Mamba models (By Agent) `_calculate_max_num_blocks_for_linear_attention` sizes the recurrent-state slot pool as `max_batch_size + 1`, ignoring `pp_size`. With pipeline parallelism (e.g. `TestNemotronV3Super::test_nvfp4_parallelism[TP4_PP2]`), multiple microbatches are in-flight on the same rank simultaneously, each holding up to `max_batch_size` sequences' Mamba state — so the live-state slot count must be `max_batch_size * pp_size + 1`. The KVCacheManagerV2 path already does this (`max_num_sequences = max_batch_size * pp_size`); this commit aligns the linear-attention sizing path with that invariant. Three call sites are fixed: 1. `KVCacheManager._calculate_max_num_blocks_for_linear_attention` — both `intercept` (memory reservation in the affine model) and `max_snapshots` (slot count) now scale with `pp_size`. 2. The `is_linear_attention` branch in the estimation dry-run code path, which previously used a bare `max_batch_size`. 3. `CppMambaHybridCacheManager.get_cache_size_per_token` — the per-request fixed-cost intercept that drives upstream `_tokens_for_budget`. Without this fix, `PP > 1` runs trip `No free block found. This shouldn't happen!` in `evictionPolicy::getFreeBlock` once requests beyond the first microbatch arrive at the pool. Added a regression unit test that asserts `recurrent_primary >= max_batch_size * pp_size + 1` for `pp_size ∈ {2, 4}`, and verified the fix end-to-end with the TP4_PP2 model run on a TP2_PP2 configuration locally (test passed in 8m53s, no `No free block` error). Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 7a6e90f Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 18 01:45:33 2026 +0800 [None][fix] Cap CppMamba warmup tokens by recurrent-state pool capacity (By Agent) `_create_cuda_graph_warmup_request` constructs the longest dummy request using `available_tokens` returned by `kv_cache_manager.get_num_available_tokens`. For `CppMambaHybridCacheManager`, the parent's implementation only bounds this by the attention KV cache, so the long dummy's `token_num` could balloon to near `max_seq_len` (262K for Qwen3.5-397B-A17B). With block reuse on and `mamba_state_cache_interval=256`, that one dummy then needs ~1024 recurrent-state blocks, exhausting the RS pool when stacked with the (batch_size - 1) short dummies in the same warmup batch and tripping `No free block found` in `evictionPolicy::getFreeBlock`. Override `get_num_available_tokens` in CppMambaHybridCacheManager to also clamp by `(rs_free_blocks - 1) * states_snapshot_interval` whenever the snapshot interval is positive, leaving one block of headroom for the always-allocated trailing snapshot. The override is a no-op when `enable_block_reuse=False` (interval == 0). Verified with `tests/integration/defs/accuracy/test_llm_api_pytorch.py::TestQwen3_5_397B_A17B::test_nvfp4[tep4_block_reuse]`. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 717edeb Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Sat May 16 16:27:02 2026 +0800 fix schedule issue, improve runtime perf, allow change manager Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 1813212 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Fri May 15 08:47:09 2026 +0800 [None][fix] Address recurrent-state pool sizing and layer-first transfer bugs (By Agent) Concern 1 (kvCacheTransferManager.cpp): In the layer-first cudaMemcpy2DAsync path, source and destination pitches were both derived from pool.primaryPtr. For host-offload transfers, the secondary pool has mNumSecondaryBlocks which can differ from mNumPrimaryBlocks, giving a different per-layer stride. Fix: compute srcLayerStrideBytes and dstLayerStrideBytes independently from srcPool and dstPool, and use them as separate pitches in cudaMemcpy2DAsync. Concern 2 (resource_manager.py): In _calculate_max_num_blocks_for_linear_attention, the block-reuse branch unconditionally overwrote max_snapshots with max_tokens // interval, discarding the live-state + CUDA-graph-padding floor (max_batch_size + 1 + max_draft_len) set earlier. Fix: use max(max_tokens // interval, max_snapshots) to preserve the floor. Add a regression test (enable_block_reuse=True) that would have caught the dropped floor: with max_batch_size=4 and max_tokens=512/interval=256, the naive branch returns 2 slots (< required floor of 5). Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit e00c230 Merge: 69899c5 be4f57e Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed May 13 18:12:00 2026 +0800 Merge branch 'main' into user/xiweny/mamba-recurrent-cuda-graph-snapshots Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 69899c5 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed May 13 18:01:04 2026 +0800 fix replay check on sm120 Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 3f0e36c Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Wed May 13 17:57:48 2026 +0800 fix missing is_draft Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 516ea84 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 23:16:54 2026 +0800 reduce sync overhead Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit ace03a2 Merge: 5a85831 5aad39a Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 21:27:53 2026 +0800 Merge branch 'user/xiweny/mamba-hybrid-zero-mamba-layers' into mamba-recurrent-cuda-graph-snapshots (By Agent) Brings in the CppMambaHybridCacheManager zero-local-mamba-layers fix and its CodeRabbit follow-ups (PR NVIDIA#13999). Conflict in tests/unittest/_torch/executor/test_mamba_cache_manager.py is resolved by keeping both test sections (recurrent-state pool sizing from this branch, zero-local-mamba early-exit from the merged branch) and unioning the import set. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 5aad39a Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 11:15:21 2026 +0800 [None][fix] CppMambaHybridCacheManager: address CodeRabbit review (By Agent) Two fixes from CodeRabbit's review on NVIDIA#13999: 1. Handle layer_mask=None defensively. The signature has always typed it as Optional[List[bool]], and the new constructor was unconditionally calling layer_mask.copy() / indexing before the zero-mamba guard ran. Treat None as an all-False full-attention mask and validate length against mamba_layer_mask. 2. Seed externally visible mamba fields before the early-return guard. get_mamba_ssm_cache_dtype() reads self.ssm_state_dtype, and the use_replay_state_update property reads self._use_replay_state_update; both were only assigned on the non-early-exit path, so any caller that hit those accessors on a zero-local-mamba rank would have AttributeError'd. Move both assignments above the early-return. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 5a85831 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Tue May 12 09:41:25 2026 +0800 [None][fix] KVCacheManager: store spec_config on self in __init__ (By Agent) Address CodeRabbit review on NVIDIA#14003: the linear-attention sizing branch reads self.spec_config, but KVCacheManager.__init__ never assigned the incoming spec_config argument to self. The code worked only by accident, because CppMambaHybridCacheManager's _setup_mtp_intermediate_states runs before super().__init__ and writes self.spec_config on the partially initialized instance. Any other KVCacheManager path that exercises the linear-attention branch with a spec_config would AttributeError. Assign self.spec_config = spec_config at the top of KVCacheManager.__init__ so the attribute is always present regardless of subclass. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 3310de0 Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 11 23:46:42 2026 +0800 [None][fix] KVCacheManager: fix max_draft_tokens typo and add slot reservation tests (By Agent) Follow-up to the recurrent-state slot reservation: spec_config has no max_draft_tokens attribute (DecodingBaseConfig uses StrictBaseModel with extra="forbid"). The correct field is max_draft_len, which is also the count of additional dummy request ids issued by CUDAGraphRunner._get_padded_batch. Add two unit tests that construct a real CppMambaHybridCacheManager and inspect blocks_per_window[RECURRENT_STATES]: - without spec_config: pool reserves max_batch_size + 1 slots - with spec_config(max_draft_len=N): pool reserves max_batch_size + 1 + N Tests disable block reuse so the override (max_tokens // interval) does not mask the addition under test. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit 992ec1b Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 11 23:34:04 2026 +0800 [None][fix] KVCacheManager: reserve recurrent-state slots for CUDA graph padding and draft-len sentinels (By Agent) The recurrent-state snapshot pool was sized at exactly max_batch_size, which leaves no room for the CUDA graph padding sentinel (CUDA_GRAPH_DUMMY_REQUEST_ID) and the per-draft-len sentinels that CUDAGraphRunner::_get_padded_batch issues during speculative decoding. Reserve one extra slot for the CUDA graph padding dummy, and when a spec_config is present, reserve one additional slot per draft length so all distinct sentinel request IDs can coexist without evicting live recurrent state. The block-reuse path (max_tokens // interval) still overrides this when active. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> commit ac30ded Author: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Date: Mon May 11 20:26:14 2026 +0800 [None][fix] CppMambaHybridCacheManager: handle ranks with zero local mamba layers (By Agent) Under PP sharding, a rank can end up with only full-attention layers and no mamba layers. The previous constructor unconditionally allocated SSM / conv state and set up linear-attention metadata, leading to invalid state on these ranks. Take an early-exit after super().__init__ when local_num_mamba_layers == 0: forward num_layers (not mamba_num_layers + num_layers) and the union of mamba+attention layer masks to the parent KVCacheManager, skip all mamba-only setup. Guard prepare_resources / update_mamba_states / _setup_state_indices with the same condition so they no-op on these ranks. Add a regression test that constructs the manager with a real parent KVCacheManager and verifies: early-exit invariants on mgr state, no mamba-only attributes set, and the three guarded methods no-op without raising. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

tensorrt-cicd · 2026-05-21T09:01:27Z

PR_Github #49535 [ run ] completed with state SUCCESS. Commit: c9933a8
/LLM/main/L0_MergeRequest_PR pipeline #39166 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

…LM_DISABLE_MPI=1 (By Agent) test_cpp_hybrid_recurrent_pool_scales_with_pp_size builds a Mapping with world_size=pp_size>1. When pytest runs with --run-ray (conftest sets TLLM_DISABLE_MPI=1), Mapping factory resolves to DeviceMeshTopology whose pp_rank requires torch.distributed initialisation that is not available in this single-process unit test. The PP-pool-sizing behaviour exercised by the test is orthogonal to the Ray orchestrator, so skip cleanly in that session mode rather than mocking the distributed stack. Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

…urrent-cuda-graph-snapshots Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/_util.py

VALLIS-NERIA · 2026-05-21T11:06:14Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T11:11:41Z

PR_Github #49684 [ run ] triggered by Bot. Commit: 23ffaf3 Link to invocation

tensorrt-cicd · 2026-05-21T18:22:34Z

PR_Github #49684 [ run ] completed with state SUCCESS. Commit: 23ffaf3
/LLM/main/L0_MergeRequest_PR pipeline #39292 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

pamelap-nvidia · 2026-05-21T18:42:42Z

/bot run

tensorrt-cicd · 2026-05-21T18:48:22Z

PR_Github #49752 [ run ] triggered by Bot. Commit: 23ffaf3 Link to invocation

tensorrt-cicd · 2026-05-21T19:21:43Z

PR_Github #49752 [ run ] completed with state FAILURE. Commit: 23ffaf3
/LLM/main/L0_MergeRequest_PR pipeline #39353 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

pamelap-nvidia · 2026-05-21T20:23:53Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-21T20:30:16Z

PR_Github #49760 [ run ] triggered by Bot. Commit: 23ffaf3 Link to invocation

tensorrt-cicd · 2026-05-21T21:15:32Z

PR_Github #49760 [ run ] completed with state SUCCESS. Commit: 23ffaf3
/LLM/main/L0_MergeRequest_PR pipeline #39361 completed with status: 'SUCCESS'

CI Report

Link to invocation

…NVIDIA#14003) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA added 2 commits May 11, 2026 20:31

VALLIS-NERIA requested a review from a team as a code owner May 11, 2026 15:34

github-actions Bot assigned VALLIS-NERIA May 11, 2026

coderabbitai Bot reviewed May 11, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/pyexecutor/resource_manager.py Outdated

VALLIS-NERIA added 2 commits May 12, 2026 21:27

reduce sync overhead

516ea84

Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

VALLIS-NERIA requested review from a team as code owners May 12, 2026 15:17

VALLIS-NERIA requested review from syuoni and tomeras91 May 12, 2026 15:17

VALLIS-NERIA changed the title ~~[None][fix] KVCacheManager: reserve recurrent-state slots for CUDA graph padding and draft-len sentinels~~ [None][fix] Fix CppMambaHybridCacheManager functional and perf issues May 12, 2026

nv-guomingz approved these changes May 12, 2026

View reviewed changes

Wanli-Jiang mentioned this pull request May 13, 2026

[None][feat] Stack PRs for sweep perfing and accuracy checking #13967

Closed

1 task

Wanli-Jiang mentioned this pull request May 20, 2026

[None][feat] Do not REVIEW: Stack PRs for sweep perfing and accuracy checking (v6) #14356

Closed

pamelap-nvidia mentioned this pull request May 20, 2026

[None][fix] Fix hybrid model prefix caching + MTP for sm120 #14071

Closed

1 task

yuxianq approved these changes May 21, 2026

View reviewed changes

VALLIS-NERIA added 2 commits May 21, 2026 19:01

Merge remote-tracking branch 'origin/main' into user/xiweny/mamba-rec…

23ffaf3

…urrent-cuda-graph-snapshots Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com> # Conflicts: # tensorrt_llm/_torch/pyexecutor/_util.py

VALLIS-NERIA merged commit 79ede08 into NVIDIA:main May 22, 2026
12 of 13 checks passed

coderabbitai Bot mentioned this pull request May 25, 2026

[https://nvbugs/6210714][fix] Fix mamba block calculation #14524

Merged

4 tasks

KleinBlueC pushed a commit to KleinBlueC/TensorRT-LLM that referenced this pull request May 26, 2026

[None][fix] Fix CppMambaHybridCacheManager functional and perf issues (…

4ae0a99

…NVIDIA#14003) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

bmarimuthu-nv pushed a commit to nv-auto-deploy/TensorRT-LLM that referenced this pull request May 28, 2026

[None][fix] Fix CppMambaHybridCacheManager functional and perf issues (…

743667e

…NVIDIA#14003) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>

Uh oh!

Conversation

VALLIS-NERIA commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Functional fixes

Performance fixes (516ea84eaa)

Test plan

Summary by CodeRabbit

Uh oh!

VALLIS-NERIA commented May 11, 2026

Uh oh!

coderabbitai Bot commented May 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

❌ Failed checks (2 warnings)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

VALLIS-NERIA commented May 11, 2026

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

tensorrt-cicd commented May 11, 2026

Uh oh!

VALLIS-NERIA commented May 12, 2026

Uh oh!

tensorrt-cicd commented May 12, 2026

Uh oh!

VALLIS-NERIA commented May 12, 2026

Uh oh!

VALLIS-NERIA commented May 12, 2026

Uh oh!

tensorrt-cicd commented May 12, 2026

Uh oh!

tensorrt-cicd commented May 12, 2026

Uh oh!

tensorrt-cicd commented May 18, 2026

Uh oh!

VALLIS-NERIA commented May 18, 2026

Uh oh!

tensorrt-cicd commented May 18, 2026

Uh oh!

tensorrt-cicd commented May 18, 2026

Uh oh!

VALLIS-NERIA commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

tensorrt-cicd commented May 20, 2026

Uh oh!

VALLIS-NERIA commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

yuxianq left a comment

Choose a reason for hiding this comment

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

VALLIS-NERIA commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

pamelap-nvidia commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

Uh oh!

pamelap-nvidia commented May 21, 2026

Uh oh!

tensorrt-cicd commented May 21, 2026

VALLIS-NERIA commented May 11, 2026 •

edited

Loading

Performance fixes (`516ea84eaa`)

coderabbitai Bot commented May 11, 2026 •

edited

Loading