[TRTLLM-13977][fix] Fix NT3 NVFP4 perf regression in Blackwell

[xuantengh]

This PR fixes the perf regression of NT3 in Blackwell introduced by #13476. That PR adds constraints about head_group_ratio to enable flashinfer. In this PR, we relax the constraints ti align upstream flashinfer so that NT model family could harness faster kernels.

Summary by CodeRabbit

• Bug Fixes
  • Improved compatibility for an optimized state-update path so it can be used in more valid model configurations.
  • Reduced unnecessary fallbacks in supported setups, which may improve runtime behavior and consistency.

[xuantengh]

/bot run --disable-fail-fast

[coderabbitai]

📝 Walkthrough

Walkthrough

The condition determining when _use_flashinfer is enabled in Mamba2Mixer.__init__ is changed. It now checks head_dim against supported_head_dims and d_state against supported_d_states, removing the previous head_group_ratio / supported_head_group_ratios check.

Changes

Flashinfer Selection Update

Layer / File(s)	Summary
Update flashinfer usage condition tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py	_use_flashinfer now depends on head_dim in supported_head_dims and d_state in supported_d_states, removing the head_group_ratio constraint.

Estimated code review effort: 1 (Trivial) | ~3 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description explains the fix, but it omits the required Test Coverage and PR Checklist sections from the template.	Add the missing template sections, especially Test Coverage and PR Checklist, and briefly document the issue, solution, and validation steps.

✅ Passed checks (4 passed)

Check name	Status	Explanation
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.
Title check	✅ Passed	The title matches the change and follows the required ticket/type/summary format.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[tensorrt-cicd]

PR_Github #57923 [ run ] triggered by Bot. Commit: a1d339c Link to invocation

[tensorrt-cicd]

PR_Github #57923 [ run ] completed with state FAILURE. Commit: a1d339c
/LLM/main/L0_MergeRequest_PR pipeline #46613 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

[xuantengh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57963 [ run ] triggered by Bot. Commit: a1d339c Link to invocation

[tensorrt-cicd]

PR_Github #57963 [ run ] completed with state SUCCESS. Commit: a1d339c
/LLM/main/L0_MergeRequest_PR pipeline #46640 completed with status: 'SUCCESS'

CI Report

Link to invocation

[xuantengh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57990 [ run ] triggered by Bot. Commit: 357b993 Link to invocation

[tensorrt-cicd]

PR_Github #57990 [ run ] completed with state SUCCESS. Commit: 357b993
/LLM/main/L0_MergeRequest_PR pipeline #46660 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

[xuantengh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58026 [ run ] triggered by Bot. Commit: 357b993 Link to invocation

[tensorrt-cicd]

PR_Github #58026 [ run ] completed with state SUCCESS. Commit: 357b993
/LLM/main/L0_MergeRequest_PR pipeline #46696 completed with status: 'SUCCESS'

CI Report

Link to invocation

[xuantengh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58186 [ run ] triggered by Bot. Commit: 01b1611 Link to invocation

[tensorrt-cicd]

PR_Github #58186 [ run ] completed with state SUCCESS. Commit: 01b1611
/LLM/main/L0_MergeRequest_PR pipeline #46833 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

[xuantengh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58335 [ run ] triggered by Bot. Commit: 01b1611 Link to invocation

[tensorrt-cicd]

PR_Github #58335 [ run ] completed with state SUCCESS. Commit: 01b1611
/LLM/main/L0_MergeRequest_PR pipeline #46964 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

[xuantengh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58373 [ run ] triggered by Bot. Commit: 01b1611 Link to invocation

[tensorrt-cicd]

PR_Github #58373 [ run ] completed with state SUCCESS. Commit: 01b1611
/LLM/main/L0_MergeRequest_PR pipeline #46999 completed with status: 'SUCCESS'

CI Report

Link to invocation