[None][fix] Bypass FlashInfer SSD prefill to fix state dtype precision

[tijyojwad]

The FlashInfer SSD prefill path (added in PR #12731) hardcodes state_dtype=bfloat16 in the SSDCombined kernel, while the Triton path accumulates states in float32 (states_in_fp32=True). This silently drops state accumulation precision on Blackwell (SM100+), which can cause accuracy drift for long sequences.

Disable the FlashInfer dispatch so mamba_chunk_scan_combined always uses the Triton SSD kernels for prefill. The FlashInfer code is preserved for re-enablement once state_dtype is fixed to float32.

Summary by CodeRabbit

• Bug Fixes
  • Modified the SSD prefill execution path to use an alternative implementation instead of the previous optimization route.

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.

[tijyojwad]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR disables the FlashInfer SSD prefill optimization in the Mamba module by introducing an explicit _USE_FLASHINFER_SSD_PREFILL configuration flag set to False. This flag is added to the flashinfer_eligible condition, preventing the function from using the fused FlashInfer fast-path and forcing fallback to the Triton combined implementation for all cases.

Changes

FlashInfer SSD Prefill Disabling

Layer / File(s)	Summary
Disable FlashInfer SSD prefill route tensorrt_llm/_torch/modules/mamba/ssd_combined.py	A disabled _USE_FLASHINFER_SSD_PREFILL = False flag is added and incorporated into the flashinfer_eligible condition, blocking the fused SSD prefill optimization and ensuring the Triton combined implementation is always used.

---

🎯 1 (Trivial) | ⏱️ ~5 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: disabling FlashInfer SSD prefill to fix a state dtype precision issue.
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.
Description check	✅ Passed	PR description clearly explains the issue (state dtype precision mismatch between FlashInfer and Triton paths) and the solution (disabling FlashInfer dispatch), with context about the problem's impact and future re-enablement.

_{✏️ 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.}

[tensorrt-cicd]

PR_Github #50404 [ run ] triggered by Bot. Commit: 32b20db Link to invocation

[2ez4bz]

/bot run

[tensorrt-cicd]

PR_Github #50407 [ run ] triggered by Bot. Commit: 32b20db Link to invocation

[tensorrt-cicd]

PR_Github #50407 [ run ] completed with state SUCCESS. Commit: 32b20db
/LLM/main/L0_MergeRequest_PR pipeline #39931 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #50501 [ run ] triggered by Bot. Commit: 32b20db Link to invocation

[tensorrt-cicd]

PR_Github #50501 [ run ] completed with state SUCCESS. Commit: 32b20db
/LLM/main/L0_MergeRequest_PR pipeline #40009 completed with status: 'SUCCESS'

CI Report

Link to invocation