[TRTLLM-7348] [feat] Enable Cross-Attention to use XQA kernels for Whisper #7035
Conversation
📝 WalkthroughWalkthroughEnables cross-attention support in XQA during generation by propagating encoder input lengths through AttentionOp to XQAParams, centralizing preprocessing parameter construction, gating generation behavior in unfused kernels, updating decoder info wiring, adding debug logs, and requiring paged KV cache for cross-attention. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant App as Caller
participant Attn as AttentionOp
participant XQAD as XQA Dispatcher
participant Prep as QKV Preprocess
participant DecInfo as BuildDecoderInfo
participant Run as XQA Runner
participant Kern as XQA/Unfused Kernels
App->>Attn: enqueueGeneration(params)
Attn->>Attn: Select XQA (HALF/BF16 + KV cache)
Attn->>XQAD: runImpl(XQAParams{encoder_input_lengths,...})
XQAD->>XQAD: shouldUse(cross_attn, paged_kv?)
alt cross-attention without paged KV
XQAD-->>App: error (require paged KV)
else cross-attention with paged KV
XQAD->>Prep: makeQKVPreprocessingParams(..., encoder_input_lengths, cu_{seq,kv}_lens, gen_phase)
XQAD->>DecInfo: Build(..., seqKVLengths=encoder_input_lengths)
XQAD->>Run: invoke(preprocessing, runnerParams{seqLensKvPtr=encoder_input_lengths})
Run->>Kern: Launch XQA kernels (generation_phase)
note over Kern: Unfused kernels skip encoder KV writes during generation
Kern-->>Run: done
end
XQAD-->>Attn: status
Attn-->>App: return
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. 📜 Recent review detailsConfiguration used: .coderabbit.yaml 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (4)
🚧 Files skipped from review as they are similar to previous changes (2)
🧰 Additional context used📓 Path-based instructions (4)**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Files:
**/*.{cpp,cxx,cc,cu}📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Files:
**/*.{h,hpp,hxx,hh,cuh,cpp,cxx,cc,cu}📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Files:
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Files:
🧠 Learnings (2)📚 Learning: 2025-08-15T06:46:53.780ZApplied to files:
📚 Learning: 2025-08-14T21:04:50.208ZApplied to files:
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
🔇 Additional comments (9)
✨ Finishing Touches
🧪 Generate unit tests
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
Status, Documentation and Community
|
DomBrown
changed the title
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
cpp/tensorrt_llm/common/attentionOp.cpp (1)
1431-1471: Potential OOB write: encoderPaddingOffsets may point to a zero-sized buffer under FMHAWhen mEnableContextFMHA is true, encoder_padding_offset_size is computed as 0, but decoder_params.encoderPaddingOffsets is still set to encoder_padding_offset for cross-attn. If BuildDecoderInfo writes into encoderPaddingOffsets, this can corrupt memory.
Safer is to pass nullptr for encoderPaddingOffsets when FMHA is enabled (consistent with the generation path in XQA). Apply this change:
- int* encoder_padding_offset = (mEnableContextFMHA && !isCrossAttention()) ? nullptr - : reinterpret_cast<int*>(nextWorkspacePtr(workspace_byte_ptr, offset, encoder_padding_offset_size)); + int* encoder_padding_offset = (mEnableContextFMHA && !isCrossAttention()) ? nullptr + : reinterpret_cast<int*>(nextWorkspacePtr(workspace_byte_ptr, offset, encoder_padding_offset_size)); // ... - decoder_params.encoderPaddingOffsets - = isCrossAttention() ? encoder_padding_offset : nullptr; // cross attention takes offsets from encoder inputs + // Only needed in unfused cross-attention; FMHA path should pass nullptr. + decoder_params.encoderPaddingOffsets = (isCrossAttention() && !mEnableContextFMHA) + ? encoder_padding_offset + : nullptr;Optional: For clarity, also set the new flag so logs/kernel gates reflect intent:
+ decoder_params.isCrossAttention = isCrossAttention();
🧹 Nitpick comments (3)
cpp/tensorrt_llm/kernels/gptKernels.h (1)
225-274: Add isCrossAttention to toString() for better debuggabilitytoString() logs many fields but omits the new isCrossAttention flag. Including it will simplify debugging and log triage when cross-attn paths are active.
Apply this diff to include the flag:
ss << "removePadding: " << removePadding << std::endl; + ss << "isCrossAttention: " << (isCrossAttention ? "true" : "false") << std::endl; ss << "attentionMaskType: " << static_cast<int>(attentionMaskType) << std::endl;cpp/tensorrt_llm/common/attentionOp.cpp (1)
1498-1529: Unfused cross-attn mask build: verify host-side T initialization for half/bfloat16The host mask is built with std::vector h_attention_mask(..., 1.). For half/bfloat16, implicit host-side initialization via a double literal may not be portable across toolchains. If you’ve seen sporadic build issues, consider constructing via a typed constant or filling via a small helper that writes the correct bit-pattern for T.
If everything builds cleanly across supported compilers, feel free to ignore.
cpp/tensorrt_llm/kernels/xqaDispatcher.cpp (1)
43-131: Helper extraction is good; remove duplicated assignments in preprocessing paramsmakeQKVPreprocessingParams centralizes construction well. There’s a small duplication: rotary_embedding_inv_freq, rotary_coef_cache_buffer, kvScaleOrigQuant, spec_decoding_position_offsets, and logn_scaling are assigned twice (first in the common section, then again in the cross-attention-only section). This is benign but confusing.
Trim the duplicates in the cross-attention section to reduce maintenance overhead.
// Cross-attention only. - preprocessingParms.cu_kv_seq_lens = cu_kv_seqlens; - preprocessingParms.encoder_seq_lens = params.encoder_input_lengths; - preprocessingParms.rotary_embedding_inv_freq = rotary_inv_freq_buf; - preprocessingParms.rotary_coef_cache_buffer = params.rotary_cos_sin; - - preprocessingParms.kvScaleOrigQuant = params.kv_scale_orig_quant; - preprocessingParms.spec_decoding_position_offsets = nullptr; - preprocessingParms.logn_scaling = params.logn_scaling_ptr; + preprocessingParms.cu_kv_seq_lens = cu_kv_seqlens; + preprocessingParms.encoder_seq_lens = params.encoder_input_lengths; // Not available in generation phase preprocessingParms.mrope_rotary_cos_sin = nullptr;
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (5)
cpp/tensorrt_llm/common/attentionOp.cpp(3 hunks)cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h(2 hunks)cpp/tensorrt_llm/kernels/gptKernels.h(2 hunks)cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h(2 hunks)cpp/tensorrt_llm/kernels/xqaDispatcher.cpp(7 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}: In C++, close namespaces with a comment naming the namespace (e.g., } // namespace foo)
Prefer const/constexpr variables over #define for constants
Declare variables const if not modified after initialization
Use Allman brace style in C++
C++ filenames use lowerCamelCase and must be case-insensitively unique within a build target
C++ type names use UpperCamelCase
Local variables, methods, and namespaces use lowerCamelCase
Global non-static variables not in anonymous namespace use gPrefix lowerCamelCase (e.g., gExample)
Static globals or globals in anonymous namespaces use sPrefix lowerCamelCase
Locally visible static variables start with 's' (e.g., static std::once_flag sFlag;)
Member variables use mPrefix lowerCamelCase; public members may omit but are encouraged to use 'm'
Constants (enums, global/static/function-scope magic numbers) use kPREFIXED_UPPER_SNAKE (e.g., kDIGIT_NUM)
If macros are unavoidable, use UPPER_SNAKE_CASE (prefer constants over #define)
Constructor parameter that conflicts with a public member name gets trailing underscore (foo_)
Literal suffixes should be uppercase (e.g., 1234L not 1234l)
C++: use spaces only; indent 4 spaces
Run clang-format (LLVM style) before submitting; wrap lines at 120 characters
If formatting must be bypassed, use // clang-format off/on around the section
Prefer smart pointers; use unique_ptr for sole ownership, shared_ptr for shared; weak_ptr only in exceptional cases
Do not use deprecated pre-C++11 smart pointers
Use C++ style comments; avoid C comments except special inline cases; prefer // single-line
Capitalize and punctuate full-sentence comments
Follow Doxygen rules: use //! for comments and //!< for members in C++
Disable code with #if/#endif and mnemonic conditions; avoid commented-out code; avoid dead code
Do not throw exceptions across library boundaries
Use least-forceful casts; avoid removing const/volatile; avoid C-style and functional casts (except constructors); p...
Files:
cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.hcpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hcpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
**/*.{h,hpp,hxx,hh,cuh,cpp,cxx,cc,cu}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Parameter names must be consistent between declarations and definitions
Files:
cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.hcpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hcpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
**/*.{h,hpp,hxx,hh,cuh}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Header files must use include guards named TRTLLM__H without underscores prefix/suffix (e.g., TRTLLM_FOO_BAR_HELLO_H)
Files:
cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.hcpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.hcpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hcpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
**/*.{cpp,cxx,cc,cu}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.{cpp,cxx,cc,cu}: Avoid literal values except for 0, nullptr, true, false; use named constexpr for other literals
Place semicolon of empty for/while loop on a new line
Always use brace-delimited bodies for switch/while/do-for/if/else
Use inline C comments in argument lists when parameter meaning is unclear (e.g., /* checkForErrors = */ false)
Do not use assignment in subexpressions (e.g., if (x = y) ... is forbidden)
Switch on enums should enumerate all values and omit default to catch new values at compile time
Structure switch statements; prohibit fallthrough except between empty cases; each case ends with break or throw; return at end of case not allowed; put break inside braces for compound case
Prefer anonymous namespaces over static for internal linkage of functions
Every defined function must be called at least once (no unused methods)
Files:
cpp/tensorrt_llm/common/attentionOp.cppcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
🧠 Learnings (1)
📚 Learning: 2025-08-14T21:04:50.208Z
Learnt from: thorjohnsen
PR: NVIDIA/TensorRT-LLM#6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.208Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Applied to files:
cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.hcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
🔇 Additional comments (10)
cpp/tensorrt_llm/kernels/gptKernels.h (1)
189-221: Good addition: explicit cross-attention gating for BuildDecoderInfoAdding isCrossAttention with a default initializer and short-circuiting isBuildDecoderInfoKernelNeeded() on it is the right call to ensure decoder-info is built for cross-attention scenarios, regardless of other gates.
cpp/tensorrt_llm/kernels/unfusedAttentionKernels/unfusedAttentionKernels_2_template.h (1)
1351-1460: Cross-attn generation gating looks correct; confirm decoder offset semantics
- Setting decoder_seq_len=1 and decoder_seq_offset=batch_idx under generation_phase aligns with the remove-padding layout used in generation paths.
- Gating encoder KV-cache writes with if (!params.generation_phase) prevents redundant writes during gen and is consistent with “encoder KV is static” semantics for cross-attn.
Please double-check that qkv_input layout for generation is indeed [batch_beam, ...] when remove_padding is true so that decoder_seq_offset=batch_idx addresses the correct Q token in all gen modes (including beam search). If not, we should derive the offset from cu_seq_lens for robustness.
cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/xqaParams.h (2)
109-111: Propagating encoder_input_lengths into XQAParams is appropriateThis field is necessary for cross-attention with XQA and aligns with the updates in attentionOp and xqaDispatcher. Good placement and default nullptr.
181-181: Nice: toString() now logs encoder_input_lengthsLogging encoder_input_lengths will help diagnose sequence length mismatches in cross-attention scenarios.
cpp/tensorrt_llm/common/attentionOp.cpp (3)
288-290: Correct propagation of encoder_input_lengths to XQAMapping generationsParams.encoder_input_lengths into xqaParams ensures XQA sees the encoder seqlens for cross-attention. LGTM.
2216-2219: Helpful debug: log when XQA is not selected in generationNon-intrusive and useful for runtime diagnostics.
2738-2740: Relaxed XQA enablement is sensibleEnabling XQA when type is HALF/BF16 and KV cache is used (independent of cross-attn) matches the dispatcher’s capability checks. The final gate still depends on shouldUse() and isSupported(), so low risk.
cpp/tensorrt_llm/kernels/xqaDispatcher.cpp (3)
349-371: Correct: feed encoder lengths and cross-attn metadata into decoder-infoUsing encoder_input_lengths for seqKVLengths and setting isCrossAttention, maxEncoderQSeqLength, and encoderPaddingOffsets=nullptr is consistent with the cross-attention generation path and avoids building a square mask.
392-395: Good reuse: centralized preprocessing param constructionSwapping in makeQKVPreprocessingParams reduces duplication and keeps cross-attn wiring in one place.
444-445: Right choice: pass encoder lengths to TLLM-GEN runnertllmRunnerParams.seqLensKvPtr = encoder_input_lengths in cross-attn ensures the runner sees correct KV lengths.
DomBrown
force-pushed
the
dev/xqa_cross_attention
branch
from
3aad9b2 to
215d5b2
Compare
There was a problem hiding this comment.
Actionable comments posted: 0
🧹 Nitpick comments (6)
cpp/tensorrt_llm/kernels/gptKernels.h (3)
189-189: Good addition: explicit cross-attention flag in params
isCrossAttentionin BuildDecoderInfoParams provides a clear hook to gate decoder-info building for cross-attn paths. Consider adding a short comment to document when it should be set (e.g., generation with encoder KV).Apply this diff to document intent:
- bool isCrossAttention{false}; + // True when running cross-attention (e.g., decoder attends to encoder KV); triggers decoder-info build for KV paths. + bool isCrossAttention{false};
217-220: Ensures decoder-info kernel runs for cross-attentionUnconditionally returning true when
isCrossAttentionis set is the right call; cu_kv_seqlens and related buffers are required in XQA cross-attn paths. If you want to make this more obvious to future readers, consider moving this check to the top of the method.
272-272: toString: consider readability for booleans
isCrossAttentionwill print as 0/1 unless boolalpha is enabled. If logs are for human consumption (and not parsed), consider enablingstd::boolalphaonce at the start oftoString().Apply this diff if you prefer human-readable booleans:
std::string toString() const { std::stringstream ss; + ss << std::boolalpha; ss << "BuildDecoderInfoParams ====================" << std::endl;cpp/tensorrt_llm/kernels/xqaDispatcher.cpp (3)
80-83: Ambiguous duplicate: kvScaleOrigQuant vs kv_scale_quant_origBoth
preprocessingParms.kv_scale_quant_orig(set above) andpreprocessingParms.kvScaleOrigQuant(set here) are being populated fromparams.kv_scale_orig_quant/params.kv_scale_quant_orig(note the swapped word order). This looks like either:
- Two different fields that mean different things (INT8 vs FP8 scales), or
- A naming inconsistency/duplication.
Please confirm the intended fields and unify naming/assignment. Mismatched names tend to compile but wire the wrong buffer at runtime.
372-380: Use gated cu_kv_seqlens consistently to avoid stale/unused pointersYou correctly gate
cu_kv_seqlensto nullptr unless the decoder-info kernel runs. However,tllmRunnerParams.cumSeqLensKvPtrunconditionally useslaunchParams.cu_kv_seq_lens, which may not be built/populated when the kernel is skipped. For consistency with Q, pass the gatedcu_kv_seqlens.This avoids propagating an uninitialized or unnecessary pointer to the runner.
Apply this diff:
- tllmRunnerParams.cumSeqLensKvPtr = reinterpret_cast<int const*>(launchParams.cu_kv_seq_lens); + tllmRunnerParams.cumSeqLensKvPtr = cu_kv_seqlens;Also applies to: 430-433
387-390: Centralized preprocessing param construction improves maintainabilityReplacing inline construction with
makeQKVPreprocessingParamsis a solid refactor. Once the scale wiring is verified (see earlier comments), this will simplify future changes.
📜 Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (2)
cpp/tensorrt_llm/kernels/gptKernels.h(3 hunks)cpp/tensorrt_llm/kernels/xqaDispatcher.cpp(7 hunks)
🧰 Additional context used
📓 Path-based instructions (5)
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh}: In C++, close namespaces with a comment naming the namespace (e.g., } // namespace foo)
Prefer const/constexpr variables over #define for constants
Declare variables const if not modified after initialization
Use Allman brace style in C++
C++ filenames use lowerCamelCase and must be case-insensitively unique within a build target
C++ type names use UpperCamelCase
Local variables, methods, and namespaces use lowerCamelCase
Global non-static variables not in anonymous namespace use gPrefix lowerCamelCase (e.g., gExample)
Static globals or globals in anonymous namespaces use sPrefix lowerCamelCase
Locally visible static variables start with 's' (e.g., static std::once_flag sFlag;)
Member variables use mPrefix lowerCamelCase; public members may omit but are encouraged to use 'm'
Constants (enums, global/static/function-scope magic numbers) use kPREFIXED_UPPER_SNAKE (e.g., kDIGIT_NUM)
If macros are unavoidable, use UPPER_SNAKE_CASE (prefer constants over #define)
Constructor parameter that conflicts with a public member name gets trailing underscore (foo_)
Literal suffixes should be uppercase (e.g., 1234L not 1234l)
C++: use spaces only; indent 4 spaces
Run clang-format (LLVM style) before submitting; wrap lines at 120 characters
If formatting must be bypassed, use // clang-format off/on around the section
Prefer smart pointers; use unique_ptr for sole ownership, shared_ptr for shared; weak_ptr only in exceptional cases
Do not use deprecated pre-C++11 smart pointers
Use C++ style comments; avoid C comments except special inline cases; prefer // single-line
Capitalize and punctuate full-sentence comments
Follow Doxygen rules: use //! for comments and //!< for members in C++
Disable code with #if/#endif and mnemonic conditions; avoid commented-out code; avoid dead code
Do not throw exceptions across library boundaries
Use least-forceful casts; avoid removing const/volatile; avoid C-style and functional casts (except constructors); p...
Files:
cpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
**/*.{h,hpp,hxx,hh,cuh,cpp,cxx,cc,cu}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Parameter names must be consistent between declarations and definitions
Files:
cpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
**/*.{h,hpp,hxx,hh,cuh}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Header files must use include guards named TRTLLM__H without underscores prefix/suffix (e.g., TRTLLM_FOO_BAR_HELLO_H)
Files:
cpp/tensorrt_llm/kernels/gptKernels.h
**/*.{cpp,cxx,cc,cu,h,hpp,hxx,hh,cuh,py}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
Prepend NVIDIA copyright header (current year) to all source files
Files:
cpp/tensorrt_llm/kernels/gptKernels.hcpp/tensorrt_llm/kernels/xqaDispatcher.cpp
**/*.{cpp,cxx,cc,cu}
📄 CodeRabbit Inference Engine (CODING_GUIDELINES.md)
**/*.{cpp,cxx,cc,cu}: Avoid literal values except for 0, nullptr, true, false; use named constexpr for other literals
Place semicolon of empty for/while loop on a new line
Always use brace-delimited bodies for switch/while/do-for/if/else
Use inline C comments in argument lists when parameter meaning is unclear (e.g., /* checkForErrors = */ false)
Do not use assignment in subexpressions (e.g., if (x = y) ... is forbidden)
Switch on enums should enumerate all values and omit default to catch new values at compile time
Structure switch statements; prohibit fallthrough except between empty cases; each case ends with break or throw; return at end of case not allowed; put break inside braces for compound case
Prefer anonymous namespaces over static for internal linkage of functions
Every defined function must be called at least once (no unused methods)
Files:
cpp/tensorrt_llm/kernels/xqaDispatcher.cpp
🧠 Learnings (2)
📓 Common learnings
Learnt from: thorjohnsen
PR: NVIDIA/TensorRT-LLM#6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.208Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
📚 Learning: 2025-08-14T21:04:50.208Z
Learnt from: thorjohnsen
PR: NVIDIA/TensorRT-LLM#6910
File: cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp:0-0
Timestamp: 2025-08-14T21:04:50.208Z
Learning: In KV cache onboarding logic during prefill in cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp, when calculating which blocks fall within the attention window, use getTokensPerBlock() to advance token indices rather than block->getUniqueTokens().size(), because the calculation needs to consider the post-prefill state where blocks will be filled to capacity, not their current token count.
Applied to files:
cpp/tensorrt_llm/kernels/xqaDispatcher.cpp
🔇 Additional comments (7)
cpp/tensorrt_llm/kernels/xqaDispatcher.cpp (7)
19-22: Includes look correct for new helper and typesAdding
decoderXQAImplCommon.hand<cstdint>is appropriate for the new preprocessing helper and fixed-width ints.
63-71: Potential mix-up between Q and KV FP8 scale sources
q_scale_quant_origis initialized fromparams.kv_scale_quant_orig. That looks suspicious: Q’s FP8 quantization scale should typically come from a Q-specific source (e.g.,params.q_scale_quant_orig) rather than KV’s. Also,q_scale_quant_origis only set insidehasFp8KvCache(), which couples Q quantization to KV cache FP8 mode.Please verify:
- Whether there’s a distinct Q scale in
XQAParamsfor FP8 Q preprocessing.- Whether Q’s scale should be independent from KV’s FP8 mode.
If a Q-specific scale exists, route it here to avoid applying KV’s scale to Q.
99-103: Double-check fmha bmm1 scale formula
fmha_host_bmm1_scale = 1 / (sqrt(head_size) * q_scaling)divides byq_scaling. In most FMHA paths, the BMM1 scale is(q_scaling / sqrt(head_size)). Ifq_scalinghere is the usual multiplier, the current formula inverts it.Please verify against your existing MMHA/XQA codepaths to ensure consistency; a wrong scale will silently degrade output quality.
115-123: Cross-attention specifics wired correctly in preprocessingSetting
cu_kv_seq_lensandencoder_seq_lens, and nulling non-generation-phase fields is aligned with the generation-only cross-attn path.
227-231: Good guard: cross-attn requires paged KV cache in TRTLLM-GENThe fallback with a clear debug message helps avoid unsupported configurations.
343-365: Decoder-info params correctly adapted for cross-attention
seqKVLengthssourced fromencoder_input_lengthsunder cross-attn is correct.isCrossAttentionflag andmaxEncoderQSeqLengthset frommax_past_kv_lengthalign with the new gating in BuildDecoderInfo.encoderPaddingOffsets = nullptrfor generation-phase makes sense.LGTM.
439-440: KV length source selection looks correctUsing
encoder_input_lengthsunder cross-attn forseqLensKvPtrmatches the rest of the wiring in this PR.
|
/bot run |
|
PR_Github #15772 [ run ] triggered by Bot |
|
@wu6u3tw do we have model tests (whisper) to guard this ? just want that we won't break it in the future. Thanks. |
|
/bot kill |
DomBrown
force-pushed
the
dev/xqa_cross_attention
branch
from
215d5b2 to
b36910c
Compare
…hisper Signed-off-by: Dom Brown <[email protected]>
Signed-off-by: Dom Brown <[email protected]>
Signed-off-by: Dom Brown <[email protected]>
Signed-off-by: Dom Brown <[email protected]>
DomBrown
force-pushed
the
dev/xqa_cross_attention
branch
from
b36910c to
603173d
Compare
|
/bot kill |
|
PR_Github #15795 [ kill ] triggered by Bot |
|
PR_Github #15772 [ run ] completed with state |
|
PR_Github #15795 [ kill ] completed with state |
|
/bot run |
|
PR_Github #15797 [ run ] triggered by Bot |
|
PR_Github #15797 [ run ] completed with state |
…isper (NVIDIA#7035) Signed-off-by: Dom Brown <[email protected]> Signed-off-by: Yuxin <[email protected]>
Summary by CodeRabbit
New Features
Chores
Description
Test Coverage
GitHub Bot Help
/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...Provide a user friendly way for developers to interact with a Jenkins server.
Run
/bot [-h|--help]to print this help message.See details below for each supported subcommand.
run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]Launch build/test pipelines. All previously running jobs will be killed.
--reuse-test (optional)pipeline-id(OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.--disable-reuse-test(OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.--disable-fail-fast(OPTIONAL) : Disable fail fast on build/tests/infra failures.--skip-test(OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.--stage-list "A10-PyTorch-1, xxx"(OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.--gpu-type "A30, H100_PCIe"(OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.--test-backend "pytorch, cpp"(OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.--only-multi-gpu-test(OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.--disable-multi-gpu-test(OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.--add-multi-gpu-test(OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.--post-merge(OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx"(OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".--detailed-log(OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.--debug(OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in thestage-listparameter to access the appropriate container environment. Note: Does NOT update GitHub check status.For guidance on mapping tests to stage names, see
docs/source/reference/ci-overview.mdand the
scripts/test_to_stage_mapping.pyhelper.kill
killKill all running builds associated with pull request.
skip
skip --comment COMMENTSkip testing for latest commit on pull request.
--comment "Reason for skipping build/test"is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.reuse-pipeline
reuse-pipelineReuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.