[https://nvbugs/6293823][fix] Restored the one-line MLA-generation override from PR #13505 in…

[tensorrt-cicd]

Summary

• Root cause: PR [https://nvbugs/6185446][fix] Add warmup for trtllm-gen fmha JIT kernels #14851 (commit 6dee167, "Add warmup for trtllm-gen fmha JIT kernels") reverted PR [None][perf] Drop cubin and Eliminate ~6s FMHA JIT recompile in eager generation by aligning kernel selection with CUDA graph warmup #13505's MLA-generation override tllmRunnerParams.mMaxSeqLenKv = generation_params.max_attention_window_size and replaced it with max_past_kv_length plus a startup JIT warmup grid. The new warmup grid covers (HVPerCta128, MultiCtasKv-noCga), (HVPerCta128, MultiCtasKvCga), and (HVPerCta256, MultiCtasKvCga), but NOT the (HVPerCta256, MultiCtasKv-noCga) variant the autotuner picks at runtime for DeepSeek-R1 MLA, FP8/BF16, HQk576/HV512, page=32, batch=1, kv-len≈8199. Each cache miss synchronously JIT-compiles for ~20s mid-benchmark, blowing up mean/p99 TTFT and dropping output_token_throughput by 38.2%.
• Fix: Restored the one-line MLA-generation override from PR [None][perf] Drop cubin and Eliminate ~6s FMHA JIT recompile in eager generation by aligning kernel selection with CUDA graph warmup #13505 in mlaGeneration() so tllmRunnerParams.mMaxSeqLenKv is again pinned to generation_params.max_attention_window_size (the max cache capacity). This collapses runtime kv-len variability onto the single shape the CUDA-graph warmup already covers, eliminating the autotuner's mid-benchmark switch to the uncovered (HVPerCta256, MultiCtasKv-noCga) variant. Kept all of PR [https://nvbugs/6185446][fix] Add warmup for trtllm-gen fmha JIT kernels #14851's JIT warmup framework, plumbing, and the cache-miss warning intact (this is Option B from the triage report and matches the historical PR [None][perf] Drop cubin and Eliminate ~6s FMHA JIT recompile in eager generation by aligning kernel selection with CUDA graph warmup #13505 fix pattern). One-file, ~9-line change.
• Perf metric: output_token_throughput (higher-is-better)
• Bad perf: 615.3
• Good perf: 992
• Perf after fix: 990.2
• Automated fix generated by repair-bot

Test plan

• Verify fix on the same GPU type as the original failure
• Check for regressions in related tests

Links

• Bug: https://nvbugs/6293823

Summary by CodeRabbit

• Bug Fixes
  • Improved attention parameter handling during model generation to ensure proper kernel selection during JIT compilation and avoid potential performance issues.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: f63c4677-4f3b-4368-ab6c-da5394fa76d0

📥 Commits

Reviewing files that changed from the base of the PR and between 1b360ee and 9a74256.

📒 Files selected for processing (1)

• cpp/tensorrt_llm/common/attentionOp.cpp

---

📝 Walkthrough

Walkthrough

This PR updates the FMHA runner parameter in TensorRT-LLM's attention operation. The mMaxSeqLenKv parameter is changed from max_past_kv_length to max_attention_window_size for MLA generation under the trtllm-gen code path, with clarifying comments about JIT warmup kernel selection and paged-KV constraints.

Changes

FMHA Kernel Parameter Configuration

Layer / File(s)	Summary
Max sequence length parameter for FMHA kernel selection cpp/tensorrt_llm/common/attentionOp.cpp	Updated tllmRunnerParams.mMaxSeqLenKv to use generation_params.max_attention_window_size instead of max_past_kv_length in the MLA generation+trtllm-gen flow. Added comments explaining the intent around FMHA kernel selection during JIT warmup and handling of paged-KV behavior.

Estimated code review effort

🎯 1 (Trivial) | ⏱️ ~3 minutes

Suggested reviewers

• chzblych
• yunruis
• dpitman-nvda

🚥 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 pull request title correctly identifies the fix being applied by referencing the restored override from PR #13505 and the bug ID, though it's truncated with an ellipsis.
Description check	✅ Passed	The PR description thoroughly covers the root cause, fix approach, performance metrics, test plan, and relevant links, though it lacks explicit coverage of the checklist items from the template.
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.}