[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

The initial functional support of GDN MTP replay doesn't get beneficials under all cases, e.g. the low latency case. #16464.
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.
• The runtime path skips unused replay-work-item construction and fuses PNAT/cache bookkeeping for small batches, reducing launch and scheduling overhead.
• The replay kernel now accepts row-strided g and beta tensors directly, eliminating two redundant contiguous copies without changing computation or precision.
• Now the GDN replay mode is on by default.

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%

Freshed e2e results

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
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 follows the required ticket/type format and clearly summarizes the kernel improvement.
Description check	✅ Passed	The description covers the problem, solution, and checklist well enough, with the required sections present.

✨ 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

[nvpohanh]

[by Codex] @VALLIS-NERIA Could you please review this PR? Thank you!

[JadoTu]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61553 [ run ] triggered by Bot. Commit: f954542 Link to invocation

[tensorrt-cicd]

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