[None][perf] Fuse FlashInfer GDN prefill state I/O into Triton kernels

[nv-guomingz]

Summary by CodeRabbit

• New Features

  • Added fused GPU-accelerated utilities for efficient state layout conversions in tensor operations.

• Improvements

  • Enhanced numerical precision in gating operations by computing key operations in FP32.
  • Optimized state handling through fused conversion paths.

• Tests

  • Added comprehensive GPU test suite for state I/O operations.

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 --disable-fail-fast

[tensorrt-cicd]

PR_Github #50277 [ run ] triggered by Bot. Commit: 314e38d Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces fused Triton kernels for efficient bidirectional tensor layout conversion between TRT-LLM's (K, V) and FlashInfer's (V, K) formats. The new kernels handle optional indexed gather and scatter in a single GPU operation, eliminating intermediate allocations. The FlashInfer wrapper is updated to use these fused utilities. Supporting optimization ensures beta tensors are allocated in fp32 to avoid redundant casting downstream.

Changes

Fused State I/O and Integration

Layer / File(s)	Summary
Fused state I/O Triton kernels and wrappers tensorrt_llm/_torch/modules/fla/fused_state_io.py	Implements _gather_cast_transpose_kv_to_fp32_vk_kernel and _transpose_cast_scatter_fp32_vk_to_kv_kernel Triton kernels with tiled (K, V) block processing and masked edge handling. Public wrappers gather_cast_transpose_kv_to_fp32_vk and transpose_cast_scatter_fp32_vk_to_kv validate shapes, compute grid/block parameters, and launch kernels with explicit strides. Forward path gathers from [N_pool, H, K, V] into [num_seqs, H, V, K] fp32 with optional index mapping; reverse path scatters from [num_seqs, H, V, K] fp32 back into [N_pool, H, K, V] with optional indexed writes and casting.
FlashInfer chunk wrapper integration tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py	Imports and integrates fused state I/O utilities: replaces manual initial_state gather+cast+transpose with single gather_cast_transpose_kv_to_fp32_vk(initial_state, initial_state_indices) call; replaces manual out_state transpose+cast+assignment with transpose_cast_scatter_fp32_vk_to_kv(out_state, initial_state, initial_state_indices) for in-place indexed updates or with allocation and fused conversion for output_final_state=True branch. Kernel selection and layout assertions remain unchanged.
FP32 beta tensor allocation optimization tensorrt_llm/_torch/modules/mamba/gdn_mixer.py	In fused_gdn_gating_with_sigmoid, allocates beta_out explicitly as fp32 (dtype=torch.float32) instead of inheriting dtype from input b. Comments document that fused kernel computes sigmoid in fp32 and downstream FlashInfer GDN prefill and Triton chunked delta-rule paths accept fp32 beta, avoiding redundant casting on the hot path.
Fused kernel comprehensive tests tests/unittest/_torch/modules/mamba/test_fused_state_io.py	New test module validates both fused kernels with bit-exact PyTorch reference implementations. Forward path tests cover indexed gather+cast+transpose (multiple shapes and dtypes) and full-pool scenarios; reverse path tests verify in-place scatter preserving untouched slots and full-write mode; roundtrip test confirms gather-then-scatter restores original slots bit-exactly. SM90/SM100 architecture gating gates GPU-only tests. Import smoke test validates symbol availability.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#13644: Both PRs update tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py's chunk_gated_delta_rule to change how SSM state tensors are converted between TRT-LLM's (K, V) layout and FlashInfer's (V, K) layout (with the main PR adding fused gather/cast/transpose + transpose/cast/scatter helpers).

Suggested reviewers

• tomeras91
• symphonylyh
• QiJune
• xinhe-nv
• Wanli-Jiang
• HuiGao-NV
• yuxianq

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description contains only the template with placeholders and no actual implementation details, rationale, or test coverage information filled in.	Fill in the Description section explaining the optimization approach, the Test Coverage section listing the new test file and existing tests, and verify all PR Checklist items are addressed.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main optimization: fusing FlashInfer GDN prefill state I/O operations into Triton kernels for performance improvement.
Docstring Coverage	✅ Passed	Docstring coverage is 86.67% 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.}

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/modules/fla/fused_state_io.py (1)

1-24: 💤 Low value

Add __all__ to document the public interface.

Per coding guidelines, despite the prohibition on wildcard imports, __all__ should be kept updated to document the public interface. This module exports two public functions.

Suggested addition after imports

 import triton
 import triton.language as tl
 
+__all__ = [
+    "gather_cast_transpose_kv_to_fp32_vk",
+    "transpose_cast_scatter_fp32_vk_to_kv",
+]
+
 
 `@triton.jit`

As per coding guidelines: "Despite the prohibition on wildcard imports, keep __all__ updated to document the public interface"

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/modules/fla/fused_state_io.py` around lines 1 - 24, Add a
module-level __all__ immediately after the imports to document the public API:
include "gather_cast_transpose_kv_to_fp32_vk" and the other exported function
name(s) from this module (i.e., the two public functions implemented in this
file) as string entries in the __all__ list so tooling and readers know the
intended public interface and keep it updated when exports change.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/modules/fla/fused_state_io.py`:
- Around line 1-24: Add a module-level __all__ immediately after the imports to
document the public API: include "gather_cast_transpose_kv_to_fp32_vk" and the
other exported function name(s) from this module (i.e., the two public functions
implemented in this file) as string entries in the __all__ list so tooling and
readers know the intended public interface and keep it updated when exports
change.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 7b0abdba-9f30-40d1-80c7-7897200842f7

📥 Commits

Reviewing files that changed from the base of the PR and between 526d7ee and 314e38d.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/modules/fla/flashinfer_chunk.py
• tensorrt_llm/_torch/modules/fla/fused_state_io.py
• tensorrt_llm/_torch/modules/mamba/gdn_mixer.py
• tests/unittest/_torch/modules/mamba/test_fused_state_io.py

[tensorrt-cicd]

PR_Github #50277 [ run ] completed with state SUCCESS. Commit: 314e38d
/LLM/main/L0_MergeRequest_PR pipeline #39807 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.

CI Report

Link to invocation

[HuiGao-NV]

LGTM