[https://nvbugs/6550132][fix] Preserve disaggregated Mamba cache routing - #17224
Open
trtllm-agent wants to merge 3 commits into
Open
[https://nvbugs/6550132][fix] Preserve disaggregated Mamba cache routing#17224trtllm-agent wants to merge 3 commits into
trtllm-agent wants to merge 3 commits into
Conversation
…mba override Hybrid Mamba models default `use_kv_cache_manager_v2` to True for the hybrid state layout. `TRTLLM_USE_PY_MAMBA=1` and `TLLM_MAMBA_MANAGER_PREFERENCE` in (CPP, MIXED) select a mutually exclusive V1-route manager, and `get_kv_cache_manager_cls` rejects that pair. A user who set only the env var left the setting at "auto", which resolved to the model default of True, so the factory raised about an "explicit" V2 the user never selected -- inside the MPI worker, surfacing as "RuntimeError: Executor worker returned error". Suppress the default at the producer instead of correcting it afterwards: the hybrid `get_model_defaults` now propose V2 only when no V1-route override is active, so the contradicted value is never proposed and nothing is downgraded. The predicate lives next to `use_py_mamba_cache_manager()`, which already owns these env vars, and returns False for disaggregated serving because both overrides are aggregated-mode-only. An explicit `use_kv_cache_manager_v2=True` still raises -- that is a genuine user contradiction, which finally makes the "explicit" wording truthful. `get_model_defaults` is called with None (perf-sanity transceiver precheck) and with a plain dict (config-database test) as well as with `TorchLlmArgs`, so the predicate reads `llm_args` via `getattr`; both loose forms now get the correct answer rather than the precheck's silent "assuming V1" fallback, and are covered by a new regression test. Removes the now-stale waiver for the aggregated H100 parametrization. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
…aggregated serving TLLM_MAMBA_MANAGER_PREFERENCE selects a compatibility (V1) Mamba manager, so like TRTLLM_USE_PY_MAMBA it only applies in aggregated serving -- in disagg the manager follows the transceiver configuration. It was missing that guard, and get_kv_cache_manager_cls's is_disagg block does not return on a supported V2 route, so a hybrid model whose auto-resolved default is V2 fell through to the preference branch and raised about an 'explicit' V2 the user never set. Give the preference branch the same 'and not is_disagg' guard the sibling env var already has. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com>
trtllm-agent
requested review from
2ez4bz,
BowenFu,
QiJune,
SimengLiu-nv,
StanleySun639,
ZhanruiSunCh,
allisonlim-nv,
jieli-matrix,
tburt-nv,
yizhang-nv and
zhaoyangwang-nvidia
August 3, 2026 20:45
Contributor
WalkthroughChangesThe change centralizes Mamba manager preference handling, limits environment overrides to aggregated serving, and updates Nemotron-H and Qwen3-Next defaults for disaggregation-aware V1/V2 selection. Tests cover preference variants, argument types, conflicts, and the restored Nemotron integration test. Mamba cache-manager routing
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant ModelDefaults
participant MambaPreference
participant PyExecutor
participant Transceiver
ModelDefaults->>MambaPreference: evaluate TLLM_MAMBA_MANAGER_PREFERENCE
MambaPreference-->>ModelDefaults: V1 override status
ModelDefaults->>PyExecutor: select KV-cache manager default
PyExecutor->>Transceiver: inspect disaggregated configuration
Transceiver-->>PyExecutor: transceiver routing state
PyExecutor-->>ModelDefaults: apply V1 or V2 compatibility route
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tests/unittest/_torch/executor/test_mamba_cache_manager.py (1)
518-668: 📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick winCover disaggregated
model_dump()mappings.
mamba_manager_override_forces_v1()uses attribute access, so{"cache_transceiver_config": {"backend": "NIXL"}}is treated as aggregated and incorrectly proposes V1. Add this case and update the helper to read mapping keys. Expect_proposed_use_v2(...) == {True}.Test coverage: Added tests cover routing, defaults, overrides, and loose inputs. All changed tests are included through
unittest/_torch/executorin the applicabletest-dblists. Coverage remains insufficient until the populated mapping case is covered.🤖 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/_torch/executor/test_mamba_cache_manager.py` around lines 518 - 668, Update mamba_manager_override_forces_v1() to inspect cache_transceiver_config from both object attributes and mapping keys, including nested mapping values, so a disaggregated model_dump() with backend "NIXL" is recognized and remains on V2. Extend test_hybrid_defaults_accept_loose_llm_args() or a nearby helper test with the populated mapping case and assert _proposed_use_v2(...) == {True}.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/pyexecutor/mamba_cache_manager.py`:
- Around line 235-252: Update mamba_manager_override_forces_v1 to read
cache_transceiver_config from both object and dictionary llm_args, then read
backend from either object or dictionary configuration. Ensure a populated
dictionary with a non-None backend is treated as disaggregated and does not
force the V1 route; add regression coverage for this case.
---
Outside diff comments:
In `@tests/unittest/_torch/executor/test_mamba_cache_manager.py`:
- Around line 518-668: Update mamba_manager_override_forces_v1() to inspect
cache_transceiver_config from both object attributes and mapping keys, including
nested mapping values, so a disaggregated model_dump() with backend "NIXL" is
recognized and remains on V2. Extend
test_hybrid_defaults_accept_loose_llm_args() or a nearby helper test with the
populated mapping case and assert _proposed_use_v2(...) == {True}.
🪄 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: 7fb942fc-e102-4d8b-9427-e0f38b241fbf
📒 Files selected for processing (6)
tensorrt_llm/_torch/models/modeling_nemotron_h.pytensorrt_llm/_torch/models/modeling_qwen3_next.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytests/integration/test_lists/waives.txttests/unittest/_torch/executor/test_mamba_cache_manager.py
💤 Files with no reviewable changes (1)
- tests/integration/test_lists/waives.txt
StanleySun639
approved these changes
Aug 4, 2026
StanleySun639
left a comment
StanleySun639
left a comment
Collaborator
There was a problem hiding this comment.
LGTM if the target test case can pass.
zhaoyangwang-nvidia
approved these changes
Aug 4, 2026
zhaoyangwang-nvidia
left a comment
zhaoyangwang-nvidia
left a comment
Collaborator
There was a problem hiding this comment.
Please update related doc under docs/ maybe in docs/source/features/kvcache.md
…g (By Agent) Signed-off-by: Xiwen Yu <13230610+VALLIS-NERIA@users.noreply.github.com>
Collaborator
|
/bot run |
Collaborator
|
PR_Github #63730 [ run ] triggered by Bot. Commit: |
VALLIS-NERIA
changed the title
and not is_disagg to the TLLM_MAMBA_MANAGER_PREFERENCE branch…
1 task
Collaborator
|
PR_Github #63730 [ run ] completed with state
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
5 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
TorchLlmArgs.model_dump()nested dictionary was misclassified as aggregated serving.Test plan
Links
Dev Engineer Review
TLLM_MAMBA_MANAGER_PREFERENCEto aggregated serving.QA Engineer Review
CPPandMIXEDrouting.test_hybrid_models_default_to_v2_and_python_transceiver.tests/integration/test_lists/.waives.txtentry covered the original nvbugs/6550132 regression.