[None][feat] Disable mamba replay by default by tijyojwad · Pull Request #14471 · NVIDIA/TensorRT-LLM

tijyojwad · 2026-05-22T20:32:59Z

Summary by CodeRabbit

Bug Fixes
- Improved tensor memory handling in Mamba2 mixer operations for CUDA compatibility
- Enhanced state tensor handling in selective state update operations
- Changed Mamba MTP replay feature to be disabled by default

Description

This PR overrides the PR from #14439 with an additional bug fix

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-compatible or api-breaking. For api-breaking, include BREAKING in 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.

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

The CUDA kernel mamba2_mtp_ssm_cache_update requires contiguous dense inputs, but x_d_4d, B_d_4d, and C_d_4d arrive as non-contiguous views from torch.split + rearrange. Previously this was masked because the replay path (checked first) handled non-contiguous inputs via explicit Triton strides. With replay disabled by default, the custom op path is now reachable and fails on non-contiguous x. Add .contiguous() to x, B, C at the custom op call site (and keep the existing calls in the legacy/flashinfer path). Signed-off-by: tijyojwad <1127155+tijyojwad@users.noreply.github.com>

coderabbitai · 2026-05-22T20:34:53Z

📝 Walkthrough

Walkthrough

This PR enforces tensor memory layout requirements for the Mamba 2 MTP decoder by introducing contiguous staging in the selective state update function, coordinating caller-side contiguity in the Mamba decoder, and adjusting the replay kernel default configuration to disabled.

Changes

Mamba tensor contiguity and state management

Layer / File(s)	Summary
State tensor contiguity enforcement in CUDA op `tensorrt_llm/_torch/modules/mamba/selective_state_update.py`	The `selective_state_update_mtp_ssm_cache_trtllm` function stages non-contiguous `state` inputs into a contiguous `state_in` copy before calling the CUDA op, then copies the updated result back to the original `state` when staging was required and state updates are enabled.
Caller-side tensor contiguity for Mamba decoder `tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py`	The MTP decode path passes `.contiguous()` versions of `x_d_4d`, `B_d_4d`, and `C_d_4d` to the selective state update function; an inline comment clarifies that the Triton + flashinfer path requires contiguous/alignment-friendly tensors.
Mamba replay configuration default `tensorrt_llm/_torch/pyexecutor/_util.py`	The `TRTLLM_USE_MAMBA_REPLAY` environment variable default is changed from `'1'` to `'0'`, disabling replay by default when unset.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~8 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is incomplete and lacks required sections. The Title field is empty (showing only template placeholder), Description section is minimal without explaining the actual changes or rationale, and Test Coverage section is completely empty.	Add a proper PR title following the template format (e.g., [TRTLLM-XXXX][feat] Disable mamba replay by default), expand the Description section to explain what changes were made and why, and document the test coverage for the replay-disabling and contiguity fix changes.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: disabling mamba replay by default, which aligns with the primary objective of the pull request.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

tijyojwad · 2026-05-22T20:42:00Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T20:47:34Z

PR_Github #49984 [ run ] triggered by Bot. Commit: 731d779 Link to invocation

Tabrizian

approving executor changes.

tensor.contiguous() already returns self when contiguous, so the explicit is_contiguous() check is unnecessary. Signed-off-by: tijyojwad <1127155+tijyojwad@users.noreply.github.com>

tijyojwad · 2026-05-22T22:05:02Z

/bot run --disable-fail-fast

github-actions · 2026-05-22T22:05:12Z

⚠️ Bot command ignored: The /bot command must appear at the very beginning of the comment (no leading blank lines or spaces). Please post a new comment with /bot as the first character.

tijyojwad · 2026-05-22T22:05:28Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-22T22:11:01Z

PR_Github #49993 [ run ] triggered by Bot. Commit: 8fed4c8 Link to invocation

tensorrt-cicd · 2026-05-22T22:16:37Z

PR_Github #49984 [ run ] completed with state ABORTED. Commit: 731d779

Link to invocation

tensorrt-cicd · 2026-05-23T05:22:38Z

PR_Github #49993 [ run ] completed with state FAILURE. Commit: 8fed4c8
/LLM/main/L0_MergeRequest_PR pipeline #39558 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

Wanli-Jiang · 2026-05-23T07:20:29Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-05-23T07:26:55Z

PR_Github #50033 [ run ] triggered by Bot. Commit: 8fed4c8 Link to invocation

tensorrt-cicd · 2026-05-23T10:27:02Z

PR_Github #50033 [ run ] completed with state SUCCESS. Commit: 8fed4c8
/LLM/main/L0_MergeRequest_PR pipeline #39595 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com> Signed-off-by: tijyojwad <1127155+tijyojwad@users.noreply.github.com> Co-authored-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

[None][feat] Disable mamba replay by default

02139aa

Signed-off-by: Wanli Jiang <35160485+Wanli-Jiang@users.noreply.github.com>

tijyojwad requested review from a team as code owners May 22, 2026 20:33

tijyojwad requested review from danielafrimi, lancelly and symphonylyh May 22, 2026 20:33

github-actions Bot assigned tijyojwad May 22, 2026

tijyojwad force-pushed the user/williamj/disable-ssu-replay branch from e9f0b5d to 731d779 Compare May 22, 2026 20:34

Tabrizian approved these changes May 22, 2026

View reviewed changes

2ez4bz approved these changes May 22, 2026

View reviewed changes

2ez4bz reviewed May 22, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/modules/mamba/selective_state_update.py Outdated

[None][fix] Simplify contiguous guard in MTP SSM cache wrapper

8fed4c8

tensor.contiguous() already returns self when contiguous, so the explicit is_contiguous() check is unnecessary. Signed-off-by: tijyojwad <1127155+tijyojwad@users.noreply.github.com>

Wanli-Jiang enabled auto-merge (squash) May 23, 2026 10:12

Wanli-Jiang merged commit d741a66 into NVIDIA:main May 23, 2026
8 checks passed

Wanli-Jiang mentioned this pull request May 25, 2026

[None][feat] Disable mamba replay by default #14439

Closed

1 task

coderabbitai Bot mentioned this pull request May 25, 2026

[None][fix] Stabilize Mamba replay state update #14509

Merged

10 tasks

coderabbitai Bot mentioned this pull request May 29, 2026

[None][feat] Tune mamba config by env variables #14730

Merged

1 task

Uh oh!

Conversation

tijyojwad commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

coderabbitai Bot commented May 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

❌ Failed checks (1 warning)

Uh oh!

tijyojwad commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

Tabrizian left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

tijyojwad commented May 22, 2026

Uh oh!

github-actions Bot commented May 22, 2026

Uh oh!

tijyojwad commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 22, 2026

Uh oh!

tensorrt-cicd commented May 23, 2026

Uh oh!

Wanli-Jiang commented May 23, 2026

Uh oh!

tensorrt-cicd commented May 23, 2026

Uh oh!

tensorrt-cicd commented May 23, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

tijyojwad commented May 22, 2026 •

edited

Loading

coderabbitai Bot commented May 22, 2026 •

edited

Loading