Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
chenfeiz0326 wants to merge 1 commit into
base: main
Choose a base branch
from
+18 −2
Open

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 9 additions & 1 deletion cpp/tensorrt_llm/common/attentionOp.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1133,7 +1133,15 @@ int AttentionOp::mlaGeneration(
tllmRunnerParams.mMaxSeqLenCacheKv = generation_params.max_attention_window_size;
// This should be set to numDraftTokens + 1.
tllmRunnerParams.mMaxSeqLenQ = params.acc_q_len / batch_beam;
tllmRunnerParams.mMaxSeqLenKv = generation_params.max_past_kv_length;
// Pin mMaxSeqLenKv to the static max cache capacity so trtllm-gen FMHA picks the
// same kernel as CUDA-graph warmup and the PR #14851 JIT warmup grid. Using a
// dynamic max_past_kv_length here cycles through buckets that the warmup grid
// (which uses mMaxSeqLen) never covered, causing 1-15 NVRTC re-compiles of
// 7-9 s each on the first MTP-decode iterations and blowing TTFT P99 from
// 882 ms to 25 s on DeepSeek-R1-FP8 8K1K. Safe for PagedKv: strides do not
// depend on mMaxSeqLenKv, and extra KV CTAs early-exit through seqLensKvPtr
// (rationale from PR #13505 / commit c37992c that #14851 inadvertently reverted).
tllmRunnerParams.mMaxSeqLenKv = generation_params.max_attention_window_size;
tllmRunnerParams.mJITWarmup = generation_params.trtllm_gen_jit_warmup;
tllmRunnerParams.mJITWarmupMaxNumRequests = mMaxNumRequests;
tllmRunnerParams.mJITWarmupMaxSeqLenQ = mMaxContextLength;
Expand Down
10 changes: 9 additions & 1 deletion cpp/tensorrt_llm/kernels/xqaDispatcher.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -506,7 +506,15 @@ void XqaDispatcher::runImpl(
}
else
{
tllmRunnerParams.mMaxSeqLenKv = params.max_past_kv_length;
// Pin mMaxSeqLenKv to a static per-layer value so warmup and runtime pick the
// same FMHA kernel and avoid the eager-mode NVRTC JIT misses re-introduced by
// PR #14851. For PagedKv, strides do not depend on mMaxSeqLenKv and extra KV
// CTAs early-exit via seqLensKvPtr (rationale from the original protection that
// #14851 reverted). ContiguousKv must keep its true past-kv length because its
// strides do depend on it.
tllmRunnerParams.mMaxSeqLenKv = (tllmRunnerParams.mQkvLayout == QkvLayout::PagedKv)
? params.max_attention_window_size
: params.max_past_kv_length;
}
tllmRunnerParams.mJITWarmup = params.trtllm_gen_jit_warmup;
tllmRunnerParams.mJITWarmupMaxNumRequests = params.trtllm_gen_jit_warmup_max_num_requests;
Expand Down
Loading