Skip to content

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

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

Closed
tensorrt-cicd wants to merge 1 commit into from
Closed

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

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 @@ -1129,7 +1129,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;
// Override mMaxSeqLenKv with the max cache capacity so FMHA picks the same kernel as
// CUDA graph warmup and avoids the eager-mode JIT miss/recompile. The trtllm-gen FMHA
// JIT warmup grid added by PR #14851 does not cover the (HVPerCta256, MultiCtasKv-no-Cga)
// variant that the autotuner picks for DeepSeek-R1 MLA at runtime kv-len, so without this
// override an eager JIT miss costs ~20s mid-benchmark on every cache miss. Safe for PagedKv:
// strides do not depend on mMaxSeqLenKv, and extra KV CTAs exit early through seqLensKvPtr.
// TODO: mirror the is_swa + W+1 logic from xqaDispatcher.cpp when MLA gains SWA support
// (also requires adding Sliding cubins to the MLA gen kernel set).
tllmRunnerParams.mMaxSeqLenKv = generation_params.max_attention_window_size;
tllmRunnerParams.mJITWarmup = generation_params.trtllm_gen_jit_warmup;
tllmRunnerParams.mJITWarmupMaxNumRequests = mMaxNumRequests;
tllmRunnerParams.mJITWarmupMaxSeqLenQ = mMaxContextLength;
Expand Down
Loading