[None][feature] Add env variables to help debugging mamba modules. - #14170
Conversation
📝 WalkthroughWalkthroughThis pull request adds ChangesPython Mamba Cache Manager Override
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 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
🤖 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/model_config.py`:
- Around line 55-59: The predicate that computes use_split_pool in
model_config.py should enforce mutual exclusivity between TRTLLM_USE_CPP_MAMBA
and TRTLLM_USE_PY_MAMBA: before computing use_split_pool (and before returning
not (use_split_pool or use_spec)), explicitly read both os.environ values for
'TRTLLM_USE_CPP_MAMBA' and 'TRTLLM_USE_PY_MAMBA' and if both are set to '1'
raise a clear exception (or log and abort) so the code does not proceed
silently; keep the rest of the logic (is_disagg, spec_config, use_spec, final
return) unchanged but only after this conflict check.
In `@tensorrt_llm/_torch/pyexecutor/py_executor_creator.py`:
- Around line 838-842: The code currently calls
os.environ.pop("TRTLLM_USE_PY_MAMBA", "0") which mutates global process state
and can cause cross-executor interference; replace the pop with a non-mutating
read (e.g. os.environ.get("TRTLLM_USE_PY_MAMBA", "0") or pass an explicit
per-executor flag) and base the routing decision on that local value, keeping
the existing logger.warning message (the string referencing TRTLLM_USE_PY_MAMBA
and the instruction to use
cache_transceiver_config.transceiver_runtime='PYTHON') so behavior stays
identical but without changing os.environ; update any other occurrence (the
second instance mentioned) similarly.
🪄 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: 5d1698ce-844d-4fdb-99e2-33b7474695fb
📒 Files selected for processing (4)
tensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.py
0e29cca to
6b611ad
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48669 [ run ] triggered by Bot. Commit: |
|
PR_Github #48669 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #48687 [ run ] triggered by Bot. Commit: |
|
PR_Github #48687 [ run ] completed with state
|
6b611ad to
464e3c1
Compare
Wanli-Jiang
changed the title
TRTLLM_USE_PY_MAMBA to use legacy python mamba cache manager
464e3c1 to
a4f67fd
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #48863 [ run ] triggered by Bot. Commit: |
|
PR_Github #48863 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49033 [ run ] triggered by Bot. Commit: |
|
PR_Github #49033 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49191 [ run ] triggered by Bot. Commit: |
|
PR_Github #49191 [ run ] completed with state
|
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
7e081fc to
ba18408
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #49380 [ run ] triggered by Bot. Commit: |
|
PR_Github #49380 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #49572 [ run ] triggered by Bot. Commit: |
|
PR_Github #49572 [ run ] completed with state
|
|
/bot skip --comment "CI passed at different runs" |
|
PR_Github #49616 [ skip ] triggered by Bot. Commit: |
|
PR_Github #49616 [ skip ] completed with state |
…VIDIA#14170) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…VIDIA#14170) Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
5 participants
Features
TRTLLM_USE_PY_MAMBA(default 0) to use legacy python mamba cache manager. Default is not use py mamba cache manager.TRTLLM_USE_MAMBA_REPLAY(default 1) to use replay for mamba. Default is enable replay.Summary by CodeRabbit
New Features
Bug Fixes
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.