[https://nvbugs/6422332][fix] Keep SSM cache in weights dtype when ma… - #16065
Conversation
|
/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-5" |
|
PR_Github #57994 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR refactors mamba SSM cache dtype resolution in ChangesSSM Cache Dtype Resolution
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Warning Review ran into problems🔥 ProblemsGit: Failed to clone repository. Please run the Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/model_loader.py (1)
54-70: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winAdd test coverage for the new checkpoint-dtype mismatch warning.
This is the core of the new warning logic, but no test asserts it actually fires (or is suppressed when dtypes match).
test_modeling_qwen3_5_vl_moe.py's fixture already hasmamba_ssm_dtype="float32"vs resolvedbfloat16— the exact mismatch scenario — so add acaplog/mock-based assertion there (or a new unit test) verifyinglogger.warningis invoked with the mismatch message, and ideally a companion case where checkpoint and resolved dtypes match (no warning).As per path instructions, "Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM ... suggest concrete list file names and whether coverage is sufficient".
🤖 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/model_loader.py` around lines 54 - 70, Add test coverage for the new checkpoint-dtype mismatch warning in model_loader.py by asserting the warning path in the relevant model test or a new unit test. Use the existing mismatch fixture in test_modeling_qwen3_5_vl_moe.py (mamba_ssm_dtype float32 vs resolved bfloat16) and verify logger.warning is called with the mismatch message when checkpoint_ssm_dtype differs from mamba_ssm_cache_dtype in the auto branch. Also add a matching-dtype case to confirm no warning is emitted when resolve_checkpoint_ssm_dtype and the resolved cache dtype agree.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.
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 54-70: Add test coverage for the new checkpoint-dtype mismatch
warning in model_loader.py by asserting the warning path in the relevant model
test or a new unit test. Use the existing mismatch fixture in
test_modeling_qwen3_5_vl_moe.py (mamba_ssm_dtype float32 vs resolved bfloat16)
and verify logger.warning is called with the mismatch message when
checkpoint_ssm_dtype differs from mamba_ssm_cache_dtype in the auto branch. Also
add a matching-dtype case to confirm no warning is emitted when
resolve_checkpoint_ssm_dtype and the resolved cache dtype agree.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: ca9aab03-e417-428d-ac18-8506eb37902a
📒 Files selected for processing (4)
tensorrt_llm/_torch/pyexecutor/config_utils.pytensorrt_llm/_torch/pyexecutor/model_loader.pytests/unittest/_torch/modeling/test_modeling_qwen3_5_vl_moe.pytests/unittest/_torch/modeling/test_qwen_image_bench_modeling.py
d472023 to
603e5fa
Compare
|
PR_Github #57994 [ run ] completed with state |
|
/bot run |
|
PR_Github #58008 [ run ] triggered by Bot. Commit: |
|
PR_Github #58008 [ run ] completed with state
|
…mba_ssm_cache_dtype is auto PR NVIDIA#14599 made the 'auto' resolution of mamba_ssm_cache_dtype also honor the checkpoint's mamba_ssm_dtype field. Qwen3.5 checkpoints declare mamba_ssm_dtype=float32 (SSM compute intent), so the GDN state cache silently flipped from bf16 to fp32. An fp32 state cache disables the FlashInfer bf16-state GDN decode kernel (falls back to the slower Triton recurrent kernel) and doubles state memory traffic and footprint, regressing qwen3_5_397b_fp4_dep4_1k1k perf-sanity output throughput by ~19.5% (25435.9 -> 20479.4 tok/s) on B200. Restore the pre-NVIDIA#14599 behavior: 'auto' resolves from the explicit mamba_ssm_cache_dtype field (top-level or text_config) and otherwise falls back to the model weights dtype. The checkpoint's mamba_ssm_dtype is no longer used for cache allocation; users can still opt in explicitly via kv_cache_config.mamba_ssm_cache_dtype. Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
603e5fa to
693c414
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #58030 [ run ] triggered by Bot. Commit: |
moraxu
left a comment
moraxu
left a comment
There was a problem hiding this comment.
Thanks for updating these snippets
|
PR_Github #58030 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58099 [ run ] triggered by Bot. Commit: |
|
PR_Github #58099 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #58150 [ run ] triggered by Bot. Commit: |
|
PR_Github #58150 [ run ] completed with state |
|
/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-5" |
|
PR_Github #58185 [ ] completed with state |
NVIDIA#16065) Signed-off-by: nv-guomingz <137257613+nv-guomingz@users.noreply.github.com>
…emantics (#16264) Signed-off-by: Michal Guzek <mguzek@nvidia.com>
4 participants
…mba_ssm_cache_dtype is auto
Summary by CodeRabbit
"auto"settings.float32.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.