[https://nvbugs/6315845][fix] Dual-location revert of #14851's protective-code removal: (1) pin…

[chenfeiz0326]

Summary

• Root cause: PR [https://nvbugs/6185446][fix] Add warmup for trtllm-gen fmha JIT kernels #14851 (commit 6dee167) reverted static mMaxSeqLenKv pinning in BOTH mlaGeneration (cpp/tensorrt_llm/common/attentionOp.cpp:1133) and XqaDispatcher::runImpl (cpp/tensorrt_llm/kernels/xqaDispatcher.cpp:~509), replacing them with per-iteration max_past_kv_length. The dynamic value cycles through seqLenKv buckets the warmup grid never covered, triggering 1-15 NVRTC re-compiles of 7-9 s each and blowing TTFT P99 from 882 ms to 25 s while decode (TPOT/ITL) stays within 1%. Prior attempts addressed only mlaGeneration (or only the warmup grid).
• Fix: Dual-location revert of [https://nvbugs/6185446][fix] Add warmup for trtllm-gen fmha JIT kernels #14851's protective-code removal: (1) pin mMaxSeqLenKv = max_attention_window_size in mlaGeneration (safe for PagedKv — strides independent, KV CTAs early-exit via seqLensKvPtr); (2) pin mMaxSeqLenKv = (mQkvLayout == PagedKv) ? max_attention_window_size : max_past_kv_length in XqaDispatcher::runImpl non-spec-dec-tree branch — QkvLayout-aware so ContiguousKv keeps its true past-kv length. C++-only, two files, 18 insertions / 2 deletions, no Python or warmup-grid changes.
• Perf metric: output_token_throughput (higher-is-better)
• Bad perf: 609
• Good perf: 992.5
• Perf after fix: 988.9
• 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/6315845

Summary by CodeRabbit

• Bug Fixes
  • Improved attention window size handling during generation to better manage KV cache sequence lengths and kernel selection behavior.

[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: 8bdaeed3-fc52-4415-b228-fcad084302a9

📥 Commits

Reviewing files that changed from the base of the PR and between c390d2f and 63b8d01.

📒 Files selected for processing (2)

• cpp/tensorrt_llm/common/attentionOp.cpp
• cpp/tensorrt_llm/kernels/xqaDispatcher.cpp

---

📝 Walkthrough

Walkthrough

In two TRTLLM-GEN FMHA runner configuration sites, tllmRunnerParams.mMaxSeqLenKv is changed from max_past_kv_length to max_attention_window_size. In AttentionOp::mlaGeneration the change is unconditional; in XqaDispatcher::runImpl it applies only to the PagedKv layout, while other layouts retain max_past_kv_length.

Changes

mMaxSeqLenKv upper-bound pinning for TRTLLM-GEN FMHA

Layer / File(s)	Summary
mMaxSeqLenKv pinned to max_attention_window_size in MLA and XQA paths cpp/tensorrt_llm/common/attentionOp.cpp, cpp/tensorrt_llm/kernels/xqaDispatcher.cpp	AttentionOp::mlaGeneration unconditionally assigns generation_params.max_attention_window_size to tllmRunnerParams.mMaxSeqLenKv. XqaDispatcher::runImpl conditionally assigns params.max_attention_window_size for PagedKv layout, keeping params.max_past_kv_length for all other layouts.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• NVIDIA/TensorRT-LLM#15305: Also modifies how seqLenKv/mMaxSeqLenKv is determined for TRTLLM-GEN FMHA warmup and kernel selection, directly overlapping with the parameter being changed here.

Suggested reviewers

• yunruis
• yuxianq

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 33.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	The title begins with the bug reference and type tag, but is truncated and incompletely summarizes the fix.	Complete the title to clearly convey the dual-location fix: pin mMaxSeqLenKv in both attentionOp.cpp and xqaDispatcher.cpp to address TTFT P99 regression.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description check	✅ Passed	The description is comprehensive with root cause analysis, fix explanation, and performance metrics, meeting template requirements.
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.}