[https://nvbugs/6185446][fix] Add warmup for trtllm-gen fmha JIT kernels

[pengbowang-nv]

Revert c37992c to remove padding.

Added warmup with a grid that should be dense enough to catch almost every possible JIT kernels. Note that this PR only enables warmup and have no shape binning at runtime. I'd like to see how it performs and if it can warmup most kernels then we may not need further works.

Added a message when a possible JIT happened, to give a clear clue why performance dropped.

Some warmup data collected on our unit tests.

The kernel compilation times are similar, generally falling within the 7-9s range. The increase on warmup time can be estimated as $t_{kernel Compile} \cdot (N_{WarmupKernels} - N_{OriginalWarmupKernel})$

Model and Configuration	Warmup Enabled	Number of Kernels in FMHA Warmup Phase	Number of JIT Compilations in Original Warmup Phase	Number of JIT Compilations During Runtime
dsv3-bf16-mtp - single GPU - without cudagraph	On	16	0	0
dsv3-bf16-mtp - single GPU - without cudagraph	Off	0	4	8
dsv3-bf16-mtp - single GPU - with cudagraph	On	16	0	0
dsv3-bf16-mtp - single GPU - with cudagraph	Off	0	13	1
dsr1-fp4-mtp - 4 GPUs - with cudagraph	On	12	0	0
dsr1-fp4-mtp - 4 GPUs - with cudagraph	Off	0	4	8
gptoss-eagle3 - 4 GPUs - with cudagraph	On	16	0	0
gptoss-eagle3 - 4 GPUs - with cudagraph	Off	0	15	1
Qwen3.5-mtp - single GPU - with cudagraph	On	9	0	0
Qwen3.5-mtp - single GPU - with cudagraph	Off	0	6	0
Qwen3.5 - single GPU - with cudagraph	On	3	0	0
Qwen3.5 - single GPU - with cudagraph	Off	0	2	0
dsv3.2 - 4 GPUs - with cudagraph	On	0	0	0
dsv3.2 - 4 GPUs - with cudagraph	Off	0	0	0

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.

Summary by CodeRabbit

Release Notes

• Refactor
  • Enhanced attention computation paths to support optional JIT warmup configuration for improved FMHA kernel selection and dynamic compilation.
  • Updated sparse attention kernel selection behavior for prefill operations.
  • Streamlined attention mask type configuration to reduce complexity.
  • Replaced sliding window parameters with JIT warmup configuration for more flexible kernel warmup strategies.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR introduces an optional JIT warmup mechanism for TensorRT-LLM Generation FMHA kernels to reduce first-dispatch latency. It adds warmup configuration fields across parameter structs, removes is_sliding_window simplifying mask selection, implements a warmup grid in fmhaKernels.h with compilation logging, wires parameters through dispatchers, extends the C++ and Python APIs, and adds warmup orchestration to model initialization.

Changes

TRTLLM-Gen FMHA JIT Warmup Integration

Layer / File(s)	Summary
JIT Warmup Configuration Structures cpp/tensorrt_llm/common/attentionOp.h, cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h, cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h, cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h	JIT warmup configuration fields (trtllmGenFmhaJITWarmup, max request/Q/KV sequence-length limits) are added to EnqueueParams, MHARunnerParams, XQAParams, and TllmGenFmhaRunnerParams with default values; is_sliding_window member removed from XQAParams.
Dispatcher Logic Simplification and Parameter Routing cpp/tensorrt_llm/common/attentionOp.cpp, cpp/tensorrt_llm/kernels/xqaDispatcher.cpp	is_sliding_window computation block removed from mask-type calculation; mask-type selection simplified to use only Custom (when spec-dec-tree) or Causal; mMaxSeqLenKv assignment consolidated to consistently use max_past_kv_length.
JIT Warmup Parameter Forwarding Through Dispatchers cpp/tensorrt_llm/common/attentionOp.cpp, cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp	XQAParams and MHARunnerParams are populated with JIT warmup fields from EnqueueParams; sparse-attention prefill kernel mode is configured via mUseGenKernelForPrefill; MLA generation path updated to use max_past_kv_length for kernel selection.
JIT Warmup Grid Execution and Compilation Instrumentation cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h	Warmup candidate constants and grid-iteration helpers added; runJITWarmupGridIfRequested iterates batch/sequence combinations and compiles candidate kernels via NVRTC; generateAndCompileKernel call instrumented with timing to warn when compilation exceeds 1s.
C++ torch_ext::attention API Extension and Runner Wiring cpp/tensorrt_llm/thop/attentionOp.h, cpp/tensorrt_llm/thop/attentionOp.cpp, cpp/tensorrt_llm/nanobind/thop/bindings.cpp	RunnerBase::run and Runner::run virtual/override signatures extended with JIT warmup parameters; exported attention function and nanobind binding updated with trtllmGenFmhaJITWarmup and trtllmGenFmhaJITWarmupMaxSeqLenKv; warmup controls propagated into common_enqueue_params.
Python Attention Backend and Model Engine Integration tensorrt_llm/_torch/attention_backend/trtllm.py, tensorrt_llm/_torch/pyexecutor/model_engine.py	TrtllmAttentionMetadata gains trtllm_gen_fmha_jit_warmup field; attention execution path forwards warmup flag to thop.attention; ModelEngine.forward extended with warmup parameter; _run_attention_warmup helper orchestrates generation/context warmup requests with CUDA graphs disabled.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~60 minutes

Suggested reviewers

• lowsfer
• brb-nv
• tburt-nv

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 29.03% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	⚠️ Warning	PR description is incomplete. Required sections 'Description' and 'Test Coverage' are empty; only checklist item is checked.	Fill in the 'Description' section explaining the warmup mechanism, grid strategy, and rationale. Add 'Test Coverage' section listing tests that validate the warmup behavior and JIT compilation avoidance.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main change: adding warmup for trtllm-gen FMHA JIT kernels, with proper ticket reference and fix type.
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]

Actionable comments posted: 5

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/pyexecutor/model_engine.py (1)

84-91: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Fix forward signature mismatch (positional arg binding) between ModelEngine and PyTorchModelEngine.

ModelEngine.forward places trtllm_gen_fmha_jit_warmup immediately after num_accepted_tokens_device, but PyTorchModelEngine.forward inserts req_id_to_old_request before it—so positional callers following the abstract interface will bind the warmup flag to req_id_to_old_request and leave trtllm_gen_fmha_jit_warmup at its default.

🤖 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/pyexecutor/model_engine.py` around lines 84 - 91, The
abstract/parent signature and subclass signature are misaligned:
ModelEngine.forward currently places trtllm_gen_fmha_jit_warmup immediately
after num_accepted_tokens_device while PyTorchModelEngine.forward inserts
req_id_to_old_request before the warmup flag, causing incorrect positional
binding; fix by adding the req_id_to_old_request parameter (matching
PyTorchModelEngine) into ModelEngine.forward between num_accepted_tokens_device
and trtllm_gen_fmha_jit_warmup (with the same Optional type/default), and update
the type hints/docstring for ModelEngine.forward to match
PyTorchModelEngine.forward so both signatures are identical.

🤖 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.

Inline comments:
In `@cpp/tensorrt_llm/common/attentionOp.cpp`:
- Around line 297-300: When trtllmGenFmhaJITWarmup is enabled, reject negative
warmup limits before forwarding: change the assignments around
xqaParams.trtllmGenFmhaJITWarmupMaxNumRequests,
xqaParams.trtllmGenFmhaJITWarmupMaxSeqLenQ and
xqaParams.trtllmGenFmhaJITWarmupMaxSeqLenKv so they are only copied from
generationsParams if generationsParams.trtllmGenFmhaJITWarmup is true AND the
corresponding generationsParams.trtllmGenFmhaJITWarmupMax* value > 0; otherwise
set the xqaParams field to 0 (or another safe non-negative sentinel). Apply the
same guard pattern wherever these fields are set (including the other
occurrences around the blocks referenced: the assignments at the other
occurrences for the same symbols).

In `@cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h`:
- Around line 370-423: runJITWarmupGridIfRequested currently re-runs the full
warmup grid every call; add memoization to avoid redundant warmups by recording
completed "envelopes" and skipping any grid point already covered by an
equal-or-larger prior warm. Implement a compact key type (e.g., tuple or encoded
int) representing (useGenKernelForPrefill, batchSize, seqLenQ, seqLenKv) as a
member (or shared static) cache in the TllmGenFmhaKernel/ TllmFmhaKernelFactory
scope, check this cache before calling warmupOneKernel, and when a warmup
finishes insert the envelope; ensure comparison treats a prior warmed envelope
with batch>=batchSize and seqQ>=seqLenQ and seqKv>=seqLenKv as covering the
current point and skip warmup accordingly. Use the existing symbols
runJITWarmupGridIfRequested, warmupOneKernel, RunnerParams,
mJITWarmupMaxNumRequests, mJITWarmupMaxSeqLenQ, mJITWarmupMaxSeqLenKv and
mUseGenKernelForPrefill to locate where to add the checks and cache updates.
- Around line 390-391: The runtime check currently uses != 0 which allows
negative warmup limits; update the validation in the TLLM_CHECK_WITH_INFO call
to require strictly positive bounds (e.g., maxBatchSize > 0 && maxSeqLenKv > 0
&& (!useGenKernelForPrefill || maxSeqLenQ > 0)) so negative values are rejected
before they reach makeWarmupCandidateSizes() and warmupOneKernel(); ensure the
error message still references the same context (TRTLLM-Gen Fmha Warmup Param is
invalid) so failures fail fast with correct validation.

In `@cpp/tensorrt_llm/thop/attentionOp.h`:
- Around line 86-87: The API in attentionOp.h allows callers to set
trtllmGenFmhaJITWarmup=true while leaving trtllmGenFmhaJITWarmupMaxSeqLenKv
defaulting to 0, which is rejected by the warmup path; change the function
signature/behavior so enabling warmup cannot pick an invalid default—either give
trtllmGenFmhaJITWarmupMaxSeqLenKv a sensible non-zero default (e.g. derive from
a constant or existing max KV size), remove the default so callers must supply a
valid max when trtllmGenFmhaJITWarmup is true, or add an overload/validation
that auto-populates a valid max when trtllmGenFmhaJITWarmup is true; update
callers and any validation logic that checks trtllmGenFmhaJITWarmupMaxSeqLenKv
accordingly.

In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 903-930: The TRT FMHA JIT warmup in _run_attention_warmup may
raise torch.OutOfMemoryError and currently can abort initialization; wrap the
forward+torch.cuda.synchronize() call inside a try/except that catches
torch.OutOfMemoryError, logs a warning (including the exception), calls
torch.cuda.empty_cache() to free memory, and then continues to the next warmup
shape rather than re-raising; keep the existing no_cuda_graph() and
_release_batch_context(...) usage and only catch torch.OutOfMemoryError so other
exceptions still surface.

---

Outside diff comments:
In `@tensorrt_llm/_torch/pyexecutor/model_engine.py`:
- Around line 84-91: The abstract/parent signature and subclass signature are
misaligned: ModelEngine.forward currently places trtllm_gen_fmha_jit_warmup
immediately after num_accepted_tokens_device while PyTorchModelEngine.forward
inserts req_id_to_old_request before the warmup flag, causing incorrect
positional binding; fix by adding the req_id_to_old_request parameter (matching
PyTorchModelEngine) into ModelEngine.forward between num_accepted_tokens_device
and trtllm_gen_fmha_jit_warmup (with the same Optional type/default), and update
the type hints/docstring for ModelEngine.forward to match
PyTorchModelEngine.forward so both signatures are identical.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: da358839-ed5b-47d2-9984-c0f7fea9a4a7

📥 Commits

Reviewing files that changed from the base of the PR and between 260b80f and a8297f6.

📒 Files selected for processing (13)

• cpp/tensorrt_llm/common/attentionOp.cpp
• cpp/tensorrt_llm/common/attentionOp.h
• cpp/tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h
• cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h
• cpp/tensorrt_llm/kernels/fmhaDispatcher.cpp
• cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaKernels.h
• cpp/tensorrt_llm/kernels/trtllmGenKernels/fmha/fmhaRunnerParams.h
• cpp/tensorrt_llm/kernels/xqaDispatcher.cpp
• cpp/tensorrt_llm/nanobind/thop/bindings.cpp
• cpp/tensorrt_llm/thop/attentionOp.cpp
• cpp/tensorrt_llm/thop/attentionOp.h
• tensorrt_llm/_torch/attention_backend/trtllm.py
• tensorrt_llm/_torch/pyexecutor/model_engine.py

[pengbowang-nv]

/bot run

[tensorrt-cicd]

PR_Github #51541 [ run ] triggered by Bot. Commit: 699ae71 Link to invocation

[pengbowang-nv]

/bot run

[tensorrt-cicd]

PR_Github #51541 [ run ] completed with state FAILURE. Commit: 699ae71
/LLM/main/L0_MergeRequest_PR pipeline #40937 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

[tensorrt-cicd]

PR_Github #51549 [ run ] triggered by Bot. Commit: 699ae71 Link to invocation

[tensorrt-cicd]

PR_Github #51549 [ run ] completed with state SUCCESS. Commit: 699ae71
/LLM/main/L0_MergeRequest_PR pipeline #40944 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 #51612 [ run ] triggered by Bot. Commit: 699ae71 Link to invocation

[tensorrt-cicd]

PR_Github #51612 [ run ] completed with state SUCCESS. Commit: 699ae71
/LLM/main/L0_MergeRequest_PR pipeline #40996 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 #51750 [ run ] triggered by Bot. Commit: 8c5419e Link to invocation

[tensorrt-cicd]

PR_Github #51750 [ run ] completed with state SUCCESS. Commit: 8c5419e
/LLM/main/L0_MergeRequest_PR pipeline #41123 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 #51940 [ run ] triggered by Bot. Commit: 8c5419e Link to invocation

[tensorrt-cicd]

PR_Github #51940 [ run ] completed with state FAILURE. Commit: 8c5419e
/LLM/main/L0_MergeRequest_PR pipeline #41292 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 #51999 [ run ] triggered by Bot. Commit: 8c5419e Link to invocation

[tensorrt-cicd]

PR_Github #51999 [ run ] completed with state SUCCESS. Commit: 8c5419e
/LLM/main/L0_MergeRequest_PR pipeline #41342 completed with status: 'SUCCESS'

CI Report

Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52650 [ run ] triggered by Bot. Commit: 1a0b6e0 Link to invocation

[tensorrt-cicd]

PR_Github #52647 [ run ] completed with state ABORTED. Commit: 4ec118b

Link to invocation

[pengbowang-nv]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #52733 [ run ] triggered by Bot. Commit: 1a0b6e0 Link to invocation

[tensorrt-cicd]

PR_Github #52650 [ run ] completed with state ABORTED. Commit: 1a0b6e0
/LLM/main/L0_MergeRequest_PR pipeline #41927 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

[tensorrt-cicd]

PR_Github #52733 [ run ] completed with state SUCCESS. Commit: 1a0b6e0
/LLM/main/L0_MergeRequest_PR pipeline #41996 completed with status: 'SUCCESS'

CI Report

Link to invocation