[None][chore] Clean deprecated CppMambaCacheManager#15533
Conversation
bo-nv
requested review from
chuangz0,
liji-nv,
schetlur-nv,
shaharmor98 and
symphonylyh
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #55128 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThe PR removes the slot-based ChangesUnified KV Cache Pool: Remove RnnStateManager Slot Path
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (4)
cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp (2)
69-85: 🚀 Performance & Scalability | 🟠 Major | ⚡ Quick winHonor
maxNumTokenswhen sizing real-block capacity.The helper accepts
maxNumTokensfor buffer sizing but sizes fromwsMeta.maxTokenNumonly, somax_tokens_in_bufferno longer caps RNN preallocation for long-context unified-pool transfers.Proposed fix
auto const recurrentWs = static_cast<SizeType32>(kv_cache_manager::LinearAttentionMetadata::LinearCacheType::kRecurrentStates); auto const wsMeta = blockManager.getWindowSizeMetadata(recurrentWs); - SizeType32 maxTokenNum = wsMeta.maxTokenNum; + auto const effectiveMaxTokenNum = maxNumTokens.has_value() + ? std::min(static_cast<size_t>(wsMeta.maxTokenNum), *maxNumTokens) + : static_cast<size_t>(wsMeta.maxTokenNum); + SizeType32 const maxTokenNum = static_cast<SizeType32>(effectiveMaxTokenNum); // Number of interval snapshots + end-of-prompt block + saveLastSnapshot. maxRealBlocksPerSeq = maxTokenNum / linearMeta->statesSnapshotInterval + 1 + (linearMeta->saveLastSnapshot ? 1 : 0);🤖 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 `@cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp` around lines 69 - 85, The buffer sizing logic uses wsMeta.maxTokenNum directly to calculate maxRealBlocksPerSeq, but ignores the maxNumTokens parameter that was passed to this function. The fix is to honor the maxNumTokens constraint by taking the minimum of wsMeta.maxTokenNum and maxNumTokens when assigning to the maxTokenNum variable, ensuring that long-context unified-pool transfers respect the configured token buffer cap.
30-34: 🩺 Stability & Availability | 🔴 Critical | ⚡ Quick winValidate
kvCacheManagerbefore the base initializer dereferences it.
computeTransferBufferSizeFromPool(...)dereferenceskvCacheManagerat Line 71 during theBaseTransBufferManagerinitializer, so the constructor-body check at Line 100 runs too late for nullptr callers.Proposed fix
size_t RnnCacheTransBufferManager::computeTransferBufferSizeFromPool( kv_cache_manager::BaseKVCacheManager* kvCacheManager, executor::kv_cache::CacheState const& cacheState, std::optional<size_t> maxNumTokens) { + TLLM_CHECK(kvCacheManager != nullptr); TLLM_CHECK(cacheState.hasRnnConfig()); @@ { - TLLM_CHECK(kvCacheManager != nullptr); TLLM_LOG_INFO("RnnCacheTransBufferManager created for unified pool RNN cache"); }Also applies to: 71-72, 95-100
🤖 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 `@cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp` around lines 30 - 34, The kvCacheManager parameter is dereferenced in the BaseTransBufferManager initializer at lines 71-72, but the null pointer check currently occurs in the constructor body at lines 95-100, which runs after the initializer list. Add a null validation for kvCacheManager at the very beginning of the RnnCacheTransBufferManager constructor, before the initializer list executes, to catch nullptr inputs before they cause a dereference error. This ensures the TLLM_CHECK or similar validation happens early enough to prevent undefined behavior when kvCacheManager is null.cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp (1)
269-270: 🎯 Functional Correctness | 🔴 CriticalRNN models using
CacheTransceiverFactory::createCacheTransceiver()will not initialize RNN config.The factory function does not compute or pass
rnnLayerNumPerPPto theCacheTransceiverconstructor, unlikeattentionLayerNumPerPPwhich is properly distributed across pipeline parallel ranks. SincernnLayerNumPerPPdefaults to an empty vector, the RNN config initialization at line 350 (if (!rnnLayerNumPerPP.empty())) is skipped. Add logic to computernnLayerNumPerPPfrommodelConfig.getNbRnnLayers()similar to theattentionLayerNumPerPPcomputation, or ensure RNN models use an alternative initialization path that explicitly passes non-emptyrnnLayerNumPerPP.🤖 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 `@cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp` around lines 269 - 270, In the CacheTransceiverFactory::createCacheTransceiver() factory function, compute rnnLayerNumPerPP from modelConfig.getNbRnnLayers() following the same pattern used to compute attentionLayerNumPerPP (distributing the RNN layers across pipeline parallel ranks), and pass this computed rnnLayerNumPerPP to the CacheTransceiver constructor call alongside the existing parameters. This ensures that RNN config initialization in the CacheTransceiver constructor is not skipped due to an empty rnnLayerNumPerPP vector.tensorrt_llm/_torch/pyexecutor/_util.py (1)
100-116: 🎯 Functional Correctness | 🔴 CriticalFix hybrid manager routing to prevent AttributeError when block reuse + PYTHON transceiver are combined.
The routing logic at lines 110–111 returns
CppMambaHybridCacheManagerifenable_block_reuse=True, before checkingtransceiver_runtime="PYTHON". This allows an unsupported configuration: if bothenable_block_reuse=Trueandtransceiver_runtime="PYTHON"are set, the code returnsCppMambaHybridCacheManager, but downstream disaggregation code (kv_cache_transceiver.py:194, kv_extractor.py:94, transceiver.py:125) assumesmanager._impl.mamba_layer_offsetsandmanager._impl.mamba_cacheexist.CppMambaHybridCacheManagerhas no_implattribute—it storesmamba_layer_offsetsdirectly—causing anAttributeErrorat runtime.Add an explicit guard at line 110–111 to reject the incompatible combination, or reorder checks so PYTHON transceiver always receives
MixedMambaHybridCacheManager.🤖 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/_util.py` around lines 100 - 116, The routing logic for selecting the hybrid cache manager returns CppMambaHybridCacheManager when enable_block_reuse=True before checking whether transceiver_runtime="PYTHON", allowing an unsupported configuration that causes downstream AttributeError. Fix this by either adding an explicit guard that raises a ValueError when both enable_block_reuse=True and transceiver_runtime="PYTHON" are set together, or reorder the conditional checks so that the PYTHON transceiver condition is evaluated before the enable_block_reuse check to ensure MixedMambaHybridCacheManager is always returned when PYTHON transceiver is detected.
🧹 Nitpick comments (3)
tests/unittest/others/test_kv_cache_transceiver.py (1)
558-558: 📐 Maintainability & Code Quality | 🔵 TrivialQA coverage status: needs follow-up outside this PR for route-matrix completeness.
tests/unittest/others/test_kv_cache_transceiver.pycontinues to cover hybrid transfer/cancel behavior, but a follow-up should explicitly add a Cpp-hybrid manager route case (and expected behavior assertion under PYTHON transceiver-related config) so the manager-selection contract is locked down.As per path instructions, this QA feedback flags
tests/unittest/others/test_kv_cache_transceiver.pyas needing follow-up coverage outside the PR.Also applies to: 657-657
🤖 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 `@tests/unittest/others/test_kv_cache_transceiver.py` at line 558, The test function test_hybrid_cache_transceiver_single_process currently covers hybrid transfer and cancel behavior but is missing explicit coverage for a Cpp-hybrid manager route case. Add a test case or extend the existing test to explicitly cover the Cpp-hybrid manager route scenario with expected behavior assertions under PYTHON transceiver-related config to lock down the manager-selection contract and ensure route-matrix completeness. This follow-up coverage should be added to fully validate the hybrid cache transceiver behavior across all required manager routing paths.Source: Path instructions
tests/unittest/disaggregated/test_mamba_transfer.py (1)
332-358: 📐 Maintainability & Code Quality | 🔵 TrivialQA coverage status: sufficient for this change.
tests/unittest/disaggregated/test_mamba_transfer.pystill validates the dummy-slot allocation behavior after removing the deprecated env-var hook, so no extra in-PR coverage is required for this specific edit.As per path instructions, this QA assessment is scoped to
tests/unittest/disaggregated/test_mamba_transfer.pyand marks coverage as sufficient for the changed segment.🤖 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 `@tests/unittest/disaggregated/test_mamba_transfer.py` around lines 332 - 358, Verify that the test function test_mamba_disagg_attention_dp_dummy_with_batch_size_one properly validates the dummy-slot allocation behavior without depending on any deprecated environment variable hooks. Ensure that the test assertions (checking mamba_cache_free_blocks, mamba_cache_index values, and the comparison with _attention_dp_dummy_slot) adequately validate the expected behavior after the removal of the deprecated env-var hook, and confirm that no environment variable configuration is required for the test to pass.Source: Path instructions
tests/integration/defs/accuracy/test_llm_api_pytorch.py (1)
6842-6853: 📐 Maintainability & Code Quality | 🔵 TrivialCoverage status: sufficient for this file’s scope after legacy-path removal.
For
tests/integration/defs/accuracy/test_llm_api_pytorch.py, removinguse_cpp_mambaparametrization is appropriate with the deprecated CppMamba path removed, andtest_fp8_4gpusstill exercises bothattention_dpmodes.As per path instructions,
tests/**feedback should explicitly state coverage status and file scope; this file is sufficient for the current PR objective.🤖 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 `@tests/integration/defs/accuracy/test_llm_api_pytorch.py` around lines 6842 - 6853, Ensure that the test_fp8_4gpus method in this file does not have a use_cpp_mamba parametrization decorator, as the deprecated CppMamba path has been removed and the attention_dp parametrization alone provides sufficient coverage for the test's scope. If a use_cpp_mamba parametrization exists on the test_fp8_4gpus method, remove it entirely and keep only the attention_dp parametrization that exercises both on and off modes.Source: Path instructions
🤖 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/modules/mamba/gdn_mixer.py`:
- Around line 881-883: The Qwen3NextGatedDeltaNet instantiation in
tensorrt_llm/_torch/auto_deploy/models/custom/modeling_qwen3_next.py at line 496
is missing the required aux_stream parameter. Currently only config and
layer_idx are passed as positional arguments, but the __init__ method expects
model_config, aux_stream, and layer_idx in that order. This causes layer_idx to
be incorrectly assigned to aux_stream, leaving self.layer_idx as None, which
then causes the mamba_layer_cache call at lines 881-883 in gdn_mixer.py to fail
since it cannot accept None as a layer index. Locate the Qwen3NextGatedDeltaNet
construction at line 496 in the custom modeling file and add the missing
aux_stream argument between config and layer_idx to match the correct
implementation found at line 329 in
tensorrt_llm/_torch/models/modeling_qwen3_next.py.
In `@tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`:
- Around line 853-860: The logger.info message at line 859 incorrectly hardcodes
"Using CppMambaHybridCacheManager" without checking the actual manager selection
logic that depends on cache_transceiver_config.transceiver_runtime. Update the
else block to accurately log which cache manager will actually be used based on
the transceiver_runtime configuration value. The message should reflect that
when transceiver_runtime is set to 'PYTHON', PythonMambaCacheManager is used,
otherwise CppMambaHybridCacheManager is used, rather than always reporting
CppMambaHybridCacheManager.
---
Outside diff comments:
In `@cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp`:
- Around line 269-270: In the CacheTransceiverFactory::createCacheTransceiver()
factory function, compute rnnLayerNumPerPP from modelConfig.getNbRnnLayers()
following the same pattern used to compute attentionLayerNumPerPP (distributing
the RNN layers across pipeline parallel ranks), and pass this computed
rnnLayerNumPerPP to the CacheTransceiver constructor call alongside the existing
parameters. This ensures that RNN config initialization in the CacheTransceiver
constructor is not skipped due to an empty rnnLayerNumPerPP vector.
In `@cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp`:
- Around line 69-85: The buffer sizing logic uses wsMeta.maxTokenNum directly to
calculate maxRealBlocksPerSeq, but ignores the maxNumTokens parameter that was
passed to this function. The fix is to honor the maxNumTokens constraint by
taking the minimum of wsMeta.maxTokenNum and maxNumTokens when assigning to the
maxTokenNum variable, ensuring that long-context unified-pool transfers respect
the configured token buffer cap.
- Around line 30-34: The kvCacheManager parameter is dereferenced in the
BaseTransBufferManager initializer at lines 71-72, but the null pointer check
currently occurs in the constructor body at lines 95-100, which runs after the
initializer list. Add a null validation for kvCacheManager at the very beginning
of the RnnCacheTransBufferManager constructor, before the initializer list
executes, to catch nullptr inputs before they cause a dereference error. This
ensures the TLLM_CHECK or similar validation happens early enough to prevent
undefined behavior when kvCacheManager is null.
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 100-116: The routing logic for selecting the hybrid cache manager
returns CppMambaHybridCacheManager when enable_block_reuse=True before checking
whether transceiver_runtime="PYTHON", allowing an unsupported configuration that
causes downstream AttributeError. Fix this by either adding an explicit guard
that raises a ValueError when both enable_block_reuse=True and
transceiver_runtime="PYTHON" are set together, or reorder the conditional checks
so that the PYTHON transceiver condition is evaluated before the
enable_block_reuse check to ensure MixedMambaHybridCacheManager is always
returned when PYTHON transceiver is detected.
---
Nitpick comments:
In `@tests/integration/defs/accuracy/test_llm_api_pytorch.py`:
- Around line 6842-6853: Ensure that the test_fp8_4gpus method in this file does
not have a use_cpp_mamba parametrization decorator, as the deprecated CppMamba
path has been removed and the attention_dp parametrization alone provides
sufficient coverage for the test's scope. If a use_cpp_mamba parametrization
exists on the test_fp8_4gpus method, remove it entirely and keep only the
attention_dp parametrization that exercises both on and off modes.
In `@tests/unittest/disaggregated/test_mamba_transfer.py`:
- Around line 332-358: Verify that the test function
test_mamba_disagg_attention_dp_dummy_with_batch_size_one properly validates the
dummy-slot allocation behavior without depending on any deprecated environment
variable hooks. Ensure that the test assertions (checking
mamba_cache_free_blocks, mamba_cache_index values, and the comparison with
_attention_dp_dummy_slot) adequately validate the expected behavior after the
removal of the deprecated env-var hook, and confirm that no environment variable
configuration is required for the test to pass.
In `@tests/unittest/others/test_kv_cache_transceiver.py`:
- Line 558: The test function test_hybrid_cache_transceiver_single_process
currently covers hybrid transfer and cancel behavior but is missing explicit
coverage for a Cpp-hybrid manager route case. Add a test case or extend the
existing test to explicitly cover the Cpp-hybrid manager route scenario with
expected behavior assertions under PYTHON transceiver-related config to lock
down the manager-selection contract and ensure route-matrix completeness. This
follow-up coverage should be added to fully validate the hybrid cache
transceiver behavior across all required manager routing paths.
🪄 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: 98e15294-3765-4d86-b749-e423d9ad7500
📒 Files selected for processing (25)
cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.hcpp/include/tensorrt_llm/batch_manager/rnnCacheFormatter.hcpp/tensorrt_llm/batch_manager/cacheTransceiver.cppcpp/tensorrt_llm/batch_manager/cacheTransferLayer.cppcpp/tensorrt_llm/batch_manager/rnnCacheFormatter.cppcpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cppcpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.hcpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cucpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cppcpp/tests/unit_tests/batch_manager/rnnCacheFormatterTest.cpptensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/modules/mamba/gdn_mixer.pytensorrt_llm/_torch/modules/mamba/mamba2_mixer.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_transceiver.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytests/integration/defs/accuracy/test_llm_api_pytorch.pytests/integration/test_lists/qa/llm_function_core.txttests/integration/test_lists/test-db/l0_dgx_h100.ymltests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/disaggregated/test_mamba_transfer.pytests/unittest/others/test_kv_cache_transceiver.py
💤 Files with no reviewable changes (2)
- tensorrt_llm/_torch/model_config.py
- tests/unittest/_torch/executor/test_mamba_cache_manager.py
|
PR_Github #55128 [ run ] completed with state
|
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #55357 [ run ] triggered by Bot. Commit: |
|
PR_Github #55357 [ run ] completed with state
|
bo-nv
force-pushed
the
main-clean-CppMambaManager
branch
from
bc39287 to
f69845a
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #55723 [ run ] triggered by Bot. Commit: |
|
PR_Github #55723 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57884 [ run ] triggered by Bot. Commit: |
|
PR_Github #57884 [ run ] completed with state
|
|
/bot run |
|
PR_Github #57907 [ run ] triggered by Bot. Commit: |
bo-nv
force-pushed
the
main-clean-CppMambaManager
branch
from
5ee5c27 to
296e83d
Compare
|
/bot run |
|
PR_Github #57940 [ run ] triggered by Bot. Commit: |
|
PR_Github #57907 [ run ] completed with state |
|
PR_Github #57940 [ run ] completed with state
|
Signed-off-by: Bo Deng <deemod@nvidia.com>
Signed-off-by: Bo Deng <deemod@nvidia.com>
bo-nv
force-pushed
the
main-clean-CppMambaManager
branch
from
296e83d to
15ca2b3
Compare
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #57992 [ run ] triggered by Bot. Commit: |
|
PR_Github #57992 [ run ] completed with state
|
|
/bot run --add-multi-gpu-test --disable-fail-fast |
|
PR_Github #58172 [ run ] triggered by Bot. Commit: |
|
PR_Github #58172 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58209 [ run ] triggered by Bot. Commit: |
|
PR_Github #58209 [ run ] completed with state
|
|
/bot run |
|
PR_Github #58367 [ run ] triggered by Bot. Commit: |
|
PR_Github #58367 [ run ] completed with state |
Signed-off-by: Bo Deng <deemod@nvidia.com>
Signed-off-by: Bo Deng <deemod@nvidia.com>
|
/bot run --skip-test --extra-stage "DGX_H100-4_GPUs-PyTorch-Others-1,DGX_H100-4_GPUs-PyTorch-Others-2" |
|
PR_Github #58627 [ run ] triggered by Bot. Commit: |
|
PR_Github #58627 [ run ] completed with state
|
6 participants
Summary by CodeRabbit
Release Notes
TRTLLM_USE_CPP_MAMBAenvironment variable; Python Mamba cache manager is now the default implementation.Description
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.