[None][feat] Tune mamba config by env variables#14730
Conversation
ddd31a0 to
2a2242b
Compare
|
/bot run --disable-fail-fast |
📝 WalkthroughWalkthroughTwo Mamba kernel control paths are updated: FlashInfer SSD eligibility becomes environment-configurable via ChangesMamba Kernel Feature Toggles
Estimated code review effort🎯 2 (Simple) | ⏱️ ~12 minutes Possibly related PRs
Suggested reviewers
🚥 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.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
1317-1325:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winValidate
TRTLLM_USE_MAMBA_REPLAYvalues and fix the log message.Any value other than
"0"currently enables replay, but the log says it was"1". This can hide misconfiguration.Suggested change
- enforce_disable_replay = os.environ.get('TRTLLM_USE_MAMBA_REPLAY', - '1') == '0' + replay_env = os.environ.get('TRTLLM_USE_MAMBA_REPLAY', '1') + if replay_env not in {'0', '1'}: + raise ValueError( + f"Invalid value for TRTLLM_USE_MAMBA_REPLAY: {replay_env}") + enforce_disable_replay = replay_env == '0' @@ - logger.info( - "Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY=1") + logger.info( + f"Replay kernel is not changed since TRTLLM_USE_MAMBA_REPLAY={replay_env}")🤖 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 1317 - 1325, The current logic treats any value other than "0" as enabling replay but always logs "1", which conceals misconfiguration; update the TRTLLM_USE_MAMBA_REPLAY handling so you first read the raw value (e.g., env_val = os.environ.get('TRTLLM_USE_MAMBA_REPLAY')), then explicitly branch: if env_val == '0' set enforce_disable_replay True, set use_replay False and log that replay is disabled; elif env_val == '1' set enforce_disable_replay False, leave use_replay unchanged and log that replay is enabled; else log a warning showing the actual env_val and its fallback behavior (decide and document whether you default to enabled or disabled) so the message and behavior reflect the real value; update references to enforce_disable_replay, use_replay, and logger.info/logger.warning accordingly.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/_util.py (1)
1307-1312: ⚡ Quick winAlign replay-Philox SM gating with the stated support range.
The condition currently special-cases only 120/121, but the adjacent comment documents
100 <= sm < 120. Please make code and contract consistent to avoid enabling replay on unsupported future SM values.Suggested change
- if (stochastic_rounding - and mamba_params.mamba_ssm_cache_dtype == torch.float16 - and (sm < 100 or sm in (120, 121))): + if (stochastic_rounding + and mamba_params.mamba_ssm_cache_dtype == torch.float16 + and (sm < 100 or sm >= 120)):🤖 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 1307 - 1312, The replay-Philox SM gating is inconsistent with the comment: update the condition that checks SM so replay is only allowed when 100 <= sm < 120; specifically modify the existing if that references stochastic_rounding and mamba_params.mamba_ssm_cache_dtype to treat sm values outside [100,120) (i.e., sm < 100 or sm >= 120) as unsupported, and keep the same logger.info/error path (symbols: stochastic_rounding, mamba_params.mamba_ssm_cache_dtype, sm, logger.info) so the replay kernel is disabled/logged for any SM not in the stated supported range.
🤖 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.
Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 1317-1325: The current logic treats any value other than "0" as
enabling replay but always logs "1", which conceals misconfiguration; update the
TRTLLM_USE_MAMBA_REPLAY handling so you first read the raw value (e.g., env_val
= os.environ.get('TRTLLM_USE_MAMBA_REPLAY')), then explicitly branch: if env_val
== '0' set enforce_disable_replay True, set use_replay False and log that replay
is disabled; elif env_val == '1' set enforce_disable_replay False, leave
use_replay unchanged and log that replay is enabled; else log a warning showing
the actual env_val and its fallback behavior (decide and document whether you
default to enabled or disabled) so the message and behavior reflect the real
value; update references to enforce_disable_replay, use_replay, and
logger.info/logger.warning accordingly.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 1307-1312: The replay-Philox SM gating is inconsistent with the
comment: update the condition that checks SM so replay is only allowed when 100
<= sm < 120; specifically modify the existing if that references
stochastic_rounding and mamba_params.mamba_ssm_cache_dtype to treat sm values
outside [100,120) (i.e., sm < 100 or sm >= 120) as unsupported, and keep the
same logger.info/error path (symbols: stochastic_rounding,
mamba_params.mamba_ssm_cache_dtype, sm, logger.info) so the replay kernel is
disabled/logged for any SM not in the stated supported range.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 456e0a20-b091-421c-847b-19d817243b3b
📒 Files selected for processing (2)
tensorrt_llm/_torch/modules/mamba/ssd_combined.pytensorrt_llm/_torch/pyexecutor/_util.py
|
PR_Github #50972 [ run ] triggered by Bot. Commit: |
|
PR_Github #50972 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51089 [ run ] triggered by Bot. Commit: |
|
PR_Github #51089 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51110 [ run ] triggered by Bot. Commit: |
|
PR_Github #51110 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51126 [ run ] triggered by Bot. Commit: |
|
PR_Github #51126 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51136 [ run ] triggered by Bot. Commit: |
|
PR_Github #51136 [ run ] completed with state
|
|
/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 run --stage-list "H100_PCIe-AutoDeploy-1,DGX_B200-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #51153 [ run ] triggered by Bot. Commit: |
2a2242b to
e4d7c5d
Compare
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1,DGX_B200-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #51155 [ run ] triggered by Bot. Commit: |
|
PR_Github #51153 [ run ] completed with state |
|
PR_Github #51155 [ run ] completed with state
|
Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
…kernel * Default is 0, using triton prefill kernel. Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>
e4d7c5d to
9e85d78
Compare
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1,DGX_B200-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #51189 [ run ] triggered by Bot. Commit: |
|
PR_Github #51189 [ run ] completed with state
|
|
/bot run --stage-list "H100_PCIe-AutoDeploy-1,DGX_B200-4_GPUs-AutoDeploy-1" --disable-fail-fast |
|
PR_Github #51201 [ run ] triggered by Bot. Commit: |
|
PR_Github #51201 [ run ] completed with state
|
|
/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 run --disable-fail-fast |
|
PR_Github #51223 [ run ] triggered by Bot. Commit: |
|
PR_Github #51223 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51269 [ run ] triggered by Bot. Commit: |
|
PR_Github #51269 [ run ] completed with state |
5 participants
Features
TRTLLM_USE_MAMBA_REPLAY=0to disable replay if needed.TRTLLM_USE_MAMBA_FI_SSD=0to use triton kernel.Summary by CodeRabbit
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.