[None][fix] Stabilize Mamba replay state update - #14509
Conversation
Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #50125 [ run ] triggered by Bot. Commit: |
📝 WalkthroughWalkthroughThis PR hardens the Mamba speculative replay path against numerical instability by introducing per-slot stochastic random seeding and defensive sanitization throughout the Triton replay kernel. It allocates and manages a new ChangesMamba Speculative Replay with Hardened Randomization
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 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/mamba_cache_manager.py (1)
1399-1403:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't clear restored conv/SSM snapshots for every context slot.
Lines 1399-1400 can already queue recurrent-state restores for prefix-cache hits / chunked-prefill contexts. Lines 1431-1434 then zero those same live state stripes unconditionally, so prefill loses the restored initial state and falls back to all-zero conv/SSM state. Keep resetting the replay scratch buffers, but gate the live-state clear to truly fresh slots only.
Also applies to: 1431-1434
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py (1)
545-547: ⚡ Quick winAvoid a host sync in the slot-allocation loop.
Lines 546-547 allocate a CPU tensor and call
.item(), which forces a sync for every newly assigned request slot. Seed the selected CUDA element in place instead, or batch-generate seeds outside the loop.♻️ Suggested change
if self.mamba_cache.mamba_ssm_rand_seed is not None: - self.mamba_cache.mamba_ssm_rand_seed[block] = torch.randint( - 1, 2**62, (1, ), dtype=torch.int64).item() + self.mamba_cache.mamba_ssm_rand_seed[block].random_( + 1, 2**62 + )Based on learnings: "In files under tensorrt_llm/_torch/pyexecutor, avoid accessing torch.Tensor objects inside for-loops when iterating over requests."
🤖 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/mamba_cache_manager.py` around lines 545 - 547, The loop currently calls torch.randint(...).item() and stores a Python int into self.mamba_cache.mamba_ssm_rand_seed[block], causing a host sync for each slot; instead generate seeds on the tensor device and write them in-place (or batch-generate them before the loop). Concretely, replace the per-iteration .item() pattern for self.mamba_cache.mamba_ssm_rand_seed (referenced as mamba_cache.mamba_ssm_rand_seed and index variable block) with a device-aware assignment using torch.randint(..., device=self.mamba_cache.mamba_ssm_rand_seed.device, dtype=self.mamba_cache.mamba_ssm_rand_seed.dtype) and assign directly (or call torch.randint once to create a 1-D tensor of seeds for all blocks and then index/assign from that), avoiding any .item() calls inside the allocation loop.
🤖 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/mamba_cache_manager.py`:
- Around line 545-547: The loop currently calls torch.randint(...).item() and
stores a Python int into self.mamba_cache.mamba_ssm_rand_seed[block], causing a
host sync for each slot; instead generate seeds on the tensor device and write
them in-place (or batch-generate them before the loop). Concretely, replace the
per-iteration .item() pattern for self.mamba_cache.mamba_ssm_rand_seed
(referenced as mamba_cache.mamba_ssm_rand_seed and index variable block) with a
device-aware assignment using torch.randint(...,
device=self.mamba_cache.mamba_ssm_rand_seed.device,
dtype=self.mamba_cache.mamba_ssm_rand_seed.dtype) and assign directly (or call
torch.randint once to create a 1-D tensor of seeds for all blocks and then
index/assign from that), avoiding any .item() calls inside the allocation loop.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: f5bd5139-e083-4aca-89b1-48ef6088dc07
📒 Files selected for processing (3)
tensorrt_llm/_torch/modules/mamba/mamba2_mixer.pytensorrt_llm/_torch/modules/mamba/replay_selective_state_update.pytensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py
|
PR_Github #50125 [ run ] completed with state
|
Signed-off-by: qgai <qgai@nvidia.com>
|
/bot run --disable-fail-fast |
|
PR_Github #50274 [ run ] triggered by Bot. Commit: |
|
PR_Github #50274 [ run ] completed with state
|
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <qgai@nvidia.com>
Signed-off-by: qgai <sunnyqgg@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #50533 [ run ] triggered by Bot. Commit: |
|
PR_Github #50533 [ run ] completed with state
|
Signed-off-by: sunnyqgg <159101675+sunnyqgg@users.noreply.github.com>
|
/bot run --disable-fail-fast |
|
PR_Github #50992 [ run ] triggered by Bot. Commit: |
|
PR_Github #50992 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #51293 [ run ] triggered by Bot. Commit: |
|
PR_Github #51293 [ run ] completed with state |
Funatiq
left a comment
Funatiq
left a comment
There was a problem hiding this comment.
Executor changes look fine.
5 participants
Description
This PR stabilizes the Mamba replay selective state-update path for MTP decoding with stochastic rounding and fp16 SSM cache.
The failure mode seen in long/high-concurrency replay runs was that stale or non-finite replay/cache values could poison the SSM state, then propagate through downstream layers into
<unk>/token-id-0 collapse or over-generation. The fix keeps replay enabled and makes the replay path more robust instead of falling back to the non-replay path.Key changes:
dt/dA_cumsum, SSM state, and SSM output before they can persist or escape the kernel.Test Coverage
python3 -m py_compile tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pygit diff --check refs/tmp/nvidia_main_latest...HEAD -- tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.pypre-commit run --files tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py tensorrt_llm/_torch/modules/mamba/replay_selective_state_update.py tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py<unk>/token-id-0 collapse; replay-off remained clean as a control.PR Checklist
Summary by CodeRabbit
Bug Fixes
Improvements