[https://nvbugs/6248837][fix] Densify trtllm-gen fmha warmup grid to catch missing kernels

[pengbowang-nv]

Description

Background and Analysis

After previous warmup modification PR at #14851 , most JIT in the middle of execution issue is resolved. However we met a problem on some particular DSR1 case, where a kernel is not covered in the warmup grid and cause performance issue afterwards. In this special case the kernel would be selected only if seqlen between a very narrow area (8193-9316).

This PR densified the grid based on learning from autotuner, and should protect this case and other possible narrow cases without introducing much overhead.

Why the soultion in this PR should work (even for future cases)

Given other parameters, the autotuner would select kernel based on batchSize, seqLenQ, seqLenKv. For decode case, seqLenQ is given. As attention is done in tiles, these parameters only affect autotuner by affecting number of total tile and tile per sequence.

When tile number is in the range of 1 - 24, the autotuner logic would be rather delicate so that we require dense cover of this area. Outside this area the autotuner logic will be much less delicate and we can just cover some of them.

Each tile can be of size 128, 256 and 512. The tile number per sequence is capped at multiprocessorCount and the total tile number also mostly work by comparing with multiprocessorCount multiplied by a factor.

So for seqLenKv we use such range: [*dense_tile_num_range, *sprase_tile_num_range] multiplied by [128, 256, 512]. dense_tile_num_range = [1,2, ..., 24] covers full 1-24 tile num range, and sparese_tile_num_range = [26, 28, 32, 40, 48, 64, 96, 128, 192, 256] would loosely cover from 24 - 256 (upper cap of multiprocessor count)

For batchsize we want total number of tiles cover 1 - 24 and loosely cover other cases. Remember we have full 1-24 tile for seqLenKv range, we would like to cover most cases for each seqLenKv (effectively tile number per sequence). For tileNumPerSeq = 1, we want batchsize = 1, ..., 24; for tileNumPerSeq = 2, we want batchsize = 1, ..., 12, etc. So the final coverage would be similar to tile range of seqLenKv, but I added more batchsizes just to be sure about real cases.

Overhead

As the warmup happens on C++ fmhaRunner level, the overhead would be much easier to handle. Overhead = Tuning Overhead + Kernel Compilation Overhead. For a given model config, as the warmup should catch every posssible kernel in the run, this can be seen as a constant. Before and after this PR this should not change much (it will increase as we did catch a previously missing kernel), likely an increase of 8s. And the Tuning Overhead = mesh point * tuning time * model layer number * number of execution, tuning time is typically less than 20us, interpreting to less than 6s for most models. In actual test environment the tuning overhead is measured less than 1s. I believe this is due to the fact we are overestimating the tuning time and mesh point.

Follow up tasks

1. Change to tune only on a new layer, so we can remove model layer number factor from tuning overhead and allow for even more warmup.
2. Reduce kernel compilation overhead.

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.

Summary by CodeRabbit

• Refactor
  • Updated generation kernel warmup configuration to enhance compilation efficiency and provide optimized performance across various batch and sequence length scenarios.

[pengbowang-nv]

/bot run --stage-list "DGX_B200-8_GPUs-PyTorch-PerfSanity-Post-Merge-3"

[tensorrt-cicd]

PR_Github #53861 [ run ] triggered by Bot. Commit: 7d2a9d3 Link to invocation

[tensorrt-cicd]

PR_Github #53861 [ run ] completed with state SUCCESS. Commit: 7d2a9d3
/LLM/main/L0_MergeRequest_PR pipeline #42967 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[pengbowang-nv]

/bot run --stage-list "PerfSanity"

[tensorrt-cicd]

PR_Github #53898 [ run ] triggered by Bot. Commit: 7d2a9d3 Link to invocation

[tensorrt-cicd]

PR_Github #53898 [ run ] completed with state FAILURE. Commit: 7d2a9d3
/LLM/main/L0_MergeRequest_PR pipeline #42997 (Partly Tested) 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

[pengbowang-nv]

/bot run --stage-list "PerfSanity" --disable-fail-fast

[tensorrt-cicd]

PR_Github #54012 [ run ] triggered by Bot. Commit: 7d2a9d3 Link to invocation

[tensorrt-cicd]

PR_Github #54012 [ run ] completed with state FAILURE. Commit: 7d2a9d3
/LLM/main/L0_MergeRequest_PR pipeline #43096 (Partly Tested) 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

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54138 [ run ] triggered by Bot. Commit: 7d2a9d3 Link to invocation

[tensorrt-cicd]

PR_Github #54138 [ run ] completed with state SUCCESS. Commit: 7d2a9d3
/LLM/main/L0_MergeRequest_PR pipeline #43221 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

Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54319 [ run ] triggered by Bot. Commit: 34fa064 Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54433 [ run ] triggered by Bot. Commit: 34fa064 Link to invocation

[tensorrt-cicd]

PR_Github #54319 [ run ] completed with state ABORTED. Commit: 34fa064

Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54452 [ run ] triggered by Bot. Commit: 34fa064 Link to invocation

[tensorrt-cicd]

PR_Github #54433 [ run ] completed with state ABORTED. Commit: 34fa064

Link to invocation

[tensorrt-cicd]

PR_Github #54452 [ run ] completed with state FAILURE. Commit: 34fa064
/LLM/main/L0_MergeRequest_PR pipeline #43516 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

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54467 [ run ] triggered by Bot. Commit: 34fa064 Link to invocation

[coderabbitai]

📝 Walkthrough

Walkthrough

In fmhaKernels.h, the JIT warmup grid candidate constants are replaced: the old combined Q/KV and prefill-specific batch vectors are removed and new separate default vectors for batch size, KV sequence length, and generation-kernel Q sequence length are introduced. The runJITWarmupGridIfRequested method is updated to always use these new vectors, with Q candidates switching on useGenKernelForPrefill.

Changes

JIT Warmup Grid Candidate Update

Layer / File(s)	Summary
Warmup candidate constants and selection logic cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h	Removes prior default batch and combined Q/KV candidate vectors; adds new separate default vectors for batch size, KV sequence length, and gen-kernel Q sequence length (fixed 128). Updates runJITWarmupGridIfRequested to always draw batch from the unified default list, switch Q candidates based on useGenKernelForPrefill, and always draw KV candidates from the new KV list.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 0.00% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: densifying the FMHA warmup grid to catch missing kernels, directly addressing the performance issue described in the PR.
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.
Description check	✅ Passed	The PR description provides comprehensive background, technical rationale, and analysis of the solution with detailed explanations of the warmup grid densification strategy and overhead estimates.

_{✏️ 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)

cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h (1)

304-326: 💤 Low value

Well-designed warmup grid candidates with clear rationale.

The new constant vectors provide good coverage based on autotuner sensitivity analysis. The comments clearly explain the selection strategy (dense coverage for tile counts 1-24, sparse coverage beyond).

Minor readability enhancement

Consider using digit separators for the larger literals to improve readability:

-        10752, 11264, 11776, 12288, 13312, 14336, 16384, 20480, 24576, 32768, 40960, 49152, 65536, 98304, 131072};
+        10'752, 11'264, 11'776, 12'288, 13'312, 14'336, 16'384, 20'480, 24'576, 32'768, 40'960, 49'152, 65'536, 98'304, 131'072};

🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h` around lines
304 - 326, Add digit separators to improve readability of large numeric literals
in the warmup candidate vectors. Use C++ digit separator syntax (single quotes)
to separate groups of three digits in numbers within
kDefaultWarmupBatchSizeCandidates, kDefaultWarmupSeqLenKvCandidates, and
kDefaultWarmupSeqLenQCandidates constants. For example, replace numbers like
98304 with 98'304 and 131072 with 131'072 to make the large values easier to
read at a glance.

🤖 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 `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h`:
- Around line 304-326: Add digit separators to improve readability of large
numeric literals in the warmup candidate vectors. Use C++ digit separator syntax
(single quotes) to separate groups of three digits in numbers within
kDefaultWarmupBatchSizeCandidates, kDefaultWarmupSeqLenKvCandidates, and
kDefaultWarmupSeqLenQCandidates constants. For example, replace numbers like
98304 with 98'304 and 131072 with 131'072 to make the large values easier to
read at a glance.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6e8a4636-c199-4177-b196-ace462903d9c

📥 Commits

Reviewing files that changed from the base of the PR and between 130ae82 and 34fa064.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h

[tensorrt-cicd]

PR_Github #54521 [ run ] triggered by Bot. Commit: 34fa064 Link to invocation

[tensorrt-cicd]

PR_Github #54467 [ run ] completed with state ABORTED. Commit: 34fa064

Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54567 [ run ] triggered by Bot. Commit: a2a9c60 Link to invocation

[tensorrt-cicd]

PR_Github #54521 [ run ] completed with state ABORTED. Commit: 34fa064

Link to invocation

[tensorrt-cicd]

PR_Github #54567 [ run ] completed with state SUCCESS. Commit: a2a9c60
/LLM/main/L0_MergeRequest_PR pipeline #43611 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

Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54719 [ run ] triggered by Bot. Commit: a2a9c60 Link to invocation

[tensorrt-cicd]

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

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54747 [ run ] triggered by Bot. Commit: a2a9c60 Link to invocation

[tensorrt-cicd]

PR_Github #54747 [ run ] completed with state SUCCESS. Commit: a2a9c60
/LLM/main/L0_MergeRequest_PR pipeline #43766 completed with status: 'SUCCESS'

CI Report

Link to invocation