[TRTLLM-14540][perf] Skip fp32 state round-trip in FlashInfer GDN pre…

[nv-guomingz]

…fill on SM100/SM103

FlashInfer's SM100/SM103 GDN prefill kernel carries the recurrent state in fp32 in TMEM regardless of the initial/output-state I/O dtype, so the adapter no longer needs to up-cast the gathered state to fp32 and down-cast the final state back. Keep the state in its native pool dtype (bf16/fp16) across the FlashInfer boundary there, halving the state bytes moved by the gather/scatter and dropping the fp32 scratch buffer. SM90/SM120, whose kernels still require fp32 state, are unchanged; the recurrent accumulation stays in fp32, so the result is numerically identical.

Dev Engineer Review

• Updated FlashInfer GDN prefill state handling for SM100/SM103 to retain bf16/fp16 state dtype, reducing gather/scatter memory traffic and eliminating unnecessary fp32 scratch storage.
• Preserved fp32 state handling for SM90/SM120 and fp32 recurrent accumulation.
• Added an optional out_dtype parameter to gather_cast_vk_to_fp32_vk, maintaining fp32 as the default for API compatibility.
• Documentation and state packing behavior were clarified. No configuration or test-list changes were identified.

QA Engineer Review

No test changes.

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.

[nv-guomingz]

/bot run --diable-fail-fast

[coderabbitai]

Walkthrough

The FlashInfer chunk wrapper now selects state dtype by GPU architecture, passes it through gathered initial state, and uses it for optional state buffering. The gather helper supports configurable output dtype while defaulting to fp32.

Changes

SSM state dtype handling

Layer / File(s)	Summary
Configurable gathered-state output dtype tensorrt_llm/_torch/modules/fla/fused_state_io.py	gather_cast_vk_to_fp32_vk accepts an optional output dtype and allocates gathered state using it, defaulting to torch.float32.
GPU-specific FlashInfer state policy tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py	The wrapper selects native state dtypes on SM100/SM103, uses fp32 elsewhere, and applies the selection to initial-state gathering and state-buffer allocation.

Estimated code review effort: 2 (Simple) | ~10 minutes

Suggested reviewers: asfiyab-nvidia

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 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 clearly summarizes the FlashInfer GDN fp32 state round-trip optimization and includes the ticket reference.
Description check	✅ Passed	The PR description explains the SM100/SM103 change and rationale, and is mostly complete despite the blank Test Coverage section.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[nv-guomingz]

/bot run --disable-fail-fast

[ZhanruiSunCh]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60929 [ run ] triggered by Bot. Commit: 1a4eef1 Link to invocation

[tensorrt-cicd]

PR_Github #60929 [ run ] completed with state FAILURE. Commit: 1a4eef1
/LLM/main/L0_MergeRequest_PR pipeline #49194 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

[nv-guomingz]

/bot run --diable-fail-fast

[tensorrt-cicd]

PR_Github #60980 Bot args parsing error: usage: /bot [-h]
{run,kill,skip,submit,reviewers,reuse-pipeline,reuse-review} ...
/bot: error: unrecognized arguments: --diable-fail-fast

Link to invocation

[nv-guomingz]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60999 [ run ] triggered by Bot. Commit: 1a4eef1 Link to invocation

[tensorrt-cicd]

PR_Github #60999 [ run ] completed with state SUCCESS. Commit: 1a4eef1
/LLM/main/L0_MergeRequest_PR pipeline #49257 completed with status: 'SUCCESS'

CI Report

Link to invocation