[None][feat] Disable mamba replay by default

[Wanli-Jiang]

Summary by CodeRabbit

• Bug Fixes

  • Fixed handling of tensor memory layouts in Mamba operations to prevent potential errors with non-contiguous data.

• Chores

  • Updated default configuration for Mamba replay feature in KV cache management to be disabled by default.

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-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.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8bf3cefd-e3a8-45ee-a168-56df80742bb4

📥 Commits

Reviewing files that changed from the base of the PR and between 25e5a99 and 5586966.

📒 Files selected for processing (2)

• tensorrt_llm/_torch/modules/mamba/selective_state_update.py
• tensorrt_llm/_torch/pyexecutor/_util.py

---

📝 Walkthrough

Walkthrough

The PR contains two independent fixes: ensuring the Mamba state tensor is contiguous before CUDA operations and updating the default Mamba replay configuration from enabled to disabled.

Changes

Mamba State Update Tensor Contiguity

Layer / File(s)	Summary
State tensor contiguity handling in selective_state_update_mtp_ssm_cache_trtllm tensorrt_llm/_torch/modules/mamba/selective_state_update.py	The function now creates a contiguous staging tensor (state_in) when the input state is non-contiguous, passes the staging tensor to the CUDA op, and copies updated results back to the original tensor only when a staging copy was made and state updates are enabled.

Mamba Replay Default Behavior

Layer / File(s)	Summary
TRTLLM_USE_MAMBA_REPLAY default value flip tensorrt_llm/_torch/pyexecutor/_util.py	The environment variable default changed from '1' (replay enabled) to '0' (replay disabled) in the Nemotron-hybrid KV cache manager setup.

---

🎯 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. It contains only the repository template with placeholder comments and no actual description content, missing key sections like Description and Test Coverage.	Fill in the Description section explaining why mamba replay is disabled by default and the Test Coverage section documenting relevant tests for these changes.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main feature change: disabling mamba replay by default, which aligns with the changeset modifications.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
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.}

[Wanli-Jiang]

/bot run --disable-fail-fast

[JadoTu]

LGTM as a quick fix.
Can use a batch stride to cut runtime tensor copy in future.

[tensorrt-cicd]

PR_Github #49853 [ run ] triggered by Bot. Commit: 02139aa Link to invocation

[2ez4bz]

Nit: isn't contiguous() a no-op if it already is contiguous? https://docs.pytorch.org/docs/2.12/generated/torch.Tensor.contiguous.html