[None][fix] Stabilize Mamba replay state update - #14841
Conversation
|
/bot run --disable-fail-fast |
210fe94 to
1b51d44
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #51454 [ run ] triggered by Bot. Commit: |
1b51d44 to
714438a
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #51455 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR hardens Mamba SSM recurrent state handling and softplus numerical stability across hybrid-cache and Triton kernel paths. Fresh context slots now explicitly zero both SSM and conv states in prefill and cache preparation; replay masking prevents NaN propagation from stale cached values; and softplus computation avoids fp32 exp overflow via clamping. ChangesRecurrent State Corruption Prevention
Softplus Numerical Stability
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 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.
Actionable comments posted: 3
🤖 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/softplus.py`:
- Line 4: Update the SPDX copyright year range on the SPDX-FileCopyrightText
header in softplus.py (the line containing "SPDX-FileCopyrightText" and the
current "2022-2024" range) so the end year reflects the latest modification year
(change the end year to 2026); ensure the header now reads the correct updated
range (e.g., "2022-2026").
In `@tests/unittest/_torch/executor/test_mamba_cache_manager.py`:
- Around line 592-598: The test currently checks that prepare_resources didn't
blanket-zero non-target SSM slots by asserting mgr.all_ssm_states[:, control] ==
7.0; extend this preservation check to convolutional state by adding an
equivalent assertion for mgr.all_conv_states (use the same control index
computed from num_blocks and slot) so the test verifies both all_ssm_states and
all_conv_states were not wiped; update the assertion message to reference
all_conv_states similarly and keep the same test setup variables (mgr, slot,
num_blocks, control, prepare_resources) to locate the change.
- Around line 561-582: The current unit test
test_cpp_hybrid_fresh_context_slot_isolated_from_prior_occupant only exercises
the non-replay path of _prepare_resources, so add coverage for the
replay-enabled zeroing path by either (A) creating a new test variant that
builds the manager with replay enabled (use _build_hybrid_with_mamba_layer with
the config/flag that enables replay) and repeats the same flow (use
add_dummy_requests, set mgr.all_ssm_states/all_conv_states to non-zero, create
ScheduledRequests, call mgr._prepare_resources and assert isolation), or (B)
parameterize test_cpp_hybrid_fresh_context_slot_isolated_from_prior_occupant to
run twice (replay on/off) and perform the identical steps for both cases;
reference the same symbols (_prepare_resources, add_dummy_requests,
ScheduledRequests, mgr.all_ssm_states, mgr.all_conv_states) so the replay branch
cannot regress.
🪄 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: fe99a578-1d0f-4105-974c-17429d4b3f11
📒 Files selected for processing (7)
tensorrt_llm/_torch/modules/mamba/mamba2_mixer.pytensorrt_llm/_torch/modules/mamba/replay_selective_state_update.pytensorrt_llm/_torch/modules/mamba/softplus.pytensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pytests/unittest/_torch/executor/test_mamba_cache_manager.pytests/unittest/_torch/modules/mamba/test_mamba_softplus.py
714438a to
29f119b
Compare
|
PR_Github #51454 [ run ] completed with state |
|
PR_Github #51457 [ run ] triggered by Bot. Commit: |
|
PR_Github #51455 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #51471 [ run ] triggered by Bot. Commit: |
|
PR_Github #51457 [ run ] completed with state |
|
PR_Github #51471 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51531 [ run ] triggered by Bot. Commit: |
|
PR_Github #51531 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
29f119b to
5acbebe
Compare
|
PR_Github #51576 [ run ] triggered by Bot. Commit: |
|
/bot run --disable-fail-fast |
|
PR_Github #52047 [ run ] triggered by Bot. Commit: |
|
PR_Github #52047 [ run ] completed with state
|
|
I'm confused by what this PR is fixing. There are two changes, both it seems purely defensive. One is masking off loads to # of accepted tokens vs all of T. That will just stop NaNs from spreading via tl.dot, it won't stop them from getting produced. Upcoming checkpointing replay does seomthign similar (#14203). I am worried that this PR, and also the prior PR with cache clearing, are just making it harder to find the NaNs, but they are still present at a low rate in our production system. IMO, our priority should be isolating what in the MoE or other non-ADP layers cause NaNs, not covering them up. |
…gle-GPU test_single_request_chat_multiple_images[pd_disagg-qwen3_30b_a3b_fp8] in test_mm_encoder_standalone.py fails at setup on the pre-merge DGX_B200 single-GPU stage with a NIXL CacheTransceiver init assertion (status == NIXL_SUCCESS, transferAgent.cpp:614) when the pd_disagg LLM constructs its disaggregated KV-cache transfer agent. This is a fleet-wide failure on the single-GPU pre-merge stage, observed across many unrelated PRs (e.g. NVIDIA#13978, NVIDIA#13925, NVIDIA#14841, NVIDIA#14599, NVIDIA#14524, NVIDIA#14941, NVIDIA#14398); it passes only intermittently (~1/3) depending on node. Waiving until the single-GPU NIXL EPD-disagg path is fixed or the variant is gated to multi-GPU. NVBug: https://nvbugs/6269683 Signed-off-by: venkywonka <23023424+venkywonka@users.noreply.github.com>
As mentioned in the comment, |
|
/bot run --disable-fail-fast |
|
Sorry, why does reset make it more agent friendly? |
|
PR_Github #52261 [ run ] triggered by Bot. Commit: |
During my debugging, the Agent consistently assumes a reset-to-zero is required here. In reality, resetting to zero at this spot has no impact on precision, as the value is already zeroed during kernel execution and resource preparation. However, I’ve noticed the Agent fails to recognize this quickly. Though these two lines of code look redundant, they help the Agent confirm the reset has been done and steer its checks toward more valuable points. Since they incur no performance overhead, I believe this code is worthwhile to keep. |
Hmm. There is a tiny penalty here, clearing it twice? And I feel like a comment could achieve the same goal. but I guess it's ok? Also, are we sure it's safe with disagg? |
|
PR_Github #52261 [ run ] completed with state |
And I feel like a comment could achieve the same goal ==>Good suggestion, |
You're clearing the state at certain times related to ptefill. There's an implicit life cycle in agg that we do ptefill, and then decode. We break that in disagg so it's possible there is some interaction. Not something I thought through, just something to check. But I've slept on it and I really think we should not do this. I do worry they cover up real problems, and they should not be needed. If agents get confused, well, people do to, and that's what comments are for. We shouldn't do redundant compute. The old_* clearances are at least not redundant, but I'd still take them out. IIf something is leaking around, we would rather amplify it so we can catch it. And it does burn a few cycles. I'm not the code owner, but I think we should focus on making sure we don't produce NaNs, or in general covering up bad interactions, which when were lucky express themselves as NaNs. Thinking about that further, cuda graph padding and the ADP dummy slot can use uninitialized memory as it doesn't go through prefill, so we should clear those at init. I mean just don't allocate as torch.empty. Also I'm a little worried about the shared slot interactions for padding. I will look into adding fixes for those to the checkpointing PR, although if you want to change the allocation that'd be great. The tl.dot stuff is fine with me, it's a performance win independent of any the non NaN spreading. |
Signed-off-by: qgai <qgai@nvidia.com>
98d24e7 to
945c532
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #53435 [ run ] triggered by Bot. Commit: |
|
PR_Github #53435 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53501 [ run ] triggered by Bot. Commit: |
|
PR_Github #53501 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #53582 [ run ] triggered by Bot. Commit: |
|
PR_Github #53582 [ run ] completed with state |
5 participants
Summary
Tests
pre-commit run --files tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py: passed after applying reviewer comment wordingpre-commit run --fileson the PR files: passedpython3 -m pytest -q tests/unittest/_torch/modules/mamba/test_mamba_ssm_rand_seed.py tests/unittest/_torch/executor/test_mamba_cache_manager.py: blocked in this shell during collection because the local environment is missing TRT-LLM runtime dependencies (mpi4py; after installing missingnvtx)Excluded from this PR per request:
tensorrt_llm/_torch/pyexecutor/resource_manager.py,examples/llm-api/quickstart_advanced.py,tensorrt_llm/_torch/modules/mamba/softplus.py,tests/unittest/_torch/modules/mamba/test_mamba_softplus.py,tensorrt_llm/_torch/pyexecutor/_util.py,tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py, andtests/unittest/_torch/executor/test_mamba_cache_manager.py.