[TRTLLM-14345][feat] Improve the GDN Replay Kernel Under Low Latency

[JadoTu]

Summary by CodeRabbit

• Performance Improvements
  • Improved cached-replay execution for fine-grained head mappings, including 2:1 and 4:1 configurations.
  • Added adaptive tile sizing and dispatch selection for BV16 and BV32 workloads.
  • Optimized warp allocation and value-block handling to better preserve cache layouts and support varied head dimensions.

Description

This change improves the GDN cached-replay Triton kernel at low batch sizes by using finer value-dimension partitioning.
Previously, low-batch workloads used a BV=128, four-warp configuration. This produced too few CTAs to utilize the GPU efficiently. The new shape-aware mapping splits the value dimension into smaller tiles and uses two warps.

Kernel finetune with different tiling

BS	Old	New	Gain
1	5.788 µs	3.364 µs	41.9%
2	5.859 µs	4.222 µs	27.9%
4	5.899 µs	5.240 µs	11.2%
8	7.468 µs	6.878 µs	7.9%

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]

📝 Walkthrough

Walkthrough

Cached replay tiling now supports bounded fine-grained mappings for selected head ratios, with helper-based block_v selection and updated dispatch logic for mapping and warp-count choices.

Changes

Cached replay tiling

Layer / File(s)	Summary
Fine-grained tiling selection tensorrt_llm/_torch/modules/fla/cached_replay.py	Adds head-tile thresholds and helpers that gate fine-grained tiling and choose fixed or power-of-two value tile sizes.
Dispatch mapping integration tensorrt_llm/_torch/modules/fla/cached_replay.py	Separates production, tuned, and fine-grained mappings, then applies the corresponding block_v, small-grid, and num_warps decisions.

Estimated code review effort: 3 (Moderate) | ~15–30 minutes

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title is concise, matches the main change, and includes the required ticket and type format.
Description check	✅ Passed	The description covers the change, impact, and checklist, and it mostly matches the template.
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.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

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

27-40: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

Add boundary tests for the tile-selection helper.

Please cover the 64/65, 128/129, and 512/513 head-tile transitions, plus the power-of-two fallback. These thresholds directly change the replay grid shape and are easy to regress during future tuning.

🤖 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/cached_replay.py` around lines 27 - 40, Add
boundary tests for _default_cached_replay_block_v covering head_tiles values
64/65, 128/129, and 512/513 with tuned BF16 mapping enabled, asserting each
transition selects the expected tile size. Also test the disabled-tuning or
out-of-range path to verify the fallback returns
triton.next_power_of_2(value_dim).

🤖 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/cached_replay.py`:
- Around line 27-40: Add boundary tests for _default_cached_replay_block_v
covering head_tiles values 64/65, 128/129, and 512/513 with tuned BF16 mapping
enabled, asserting each transition selects the expected tile size. Also test the
disabled-tuning or out-of-range path to verify the fallback returns
triton.next_power_of_2(value_dim).

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 3342d260-aad7-422d-92d3-4c81d9e8afd4

📥 Commits

Reviewing files that changed from the base of the PR and between a19410e and c4d931c.

📒 Files selected for processing (1)

• tensorrt_llm/_torch/modules/fla/cached_replay.py