[None][fix] Fix a numerical stability issue for XQA with spec dec #7114
Conversation
📝 WalkthroughWalkthroughRefines kernel launch bounds under an NDEBUG/OPTIMIZE_FOR_LATENCY path, adds gated debug prints for XV GEMM row statistics, and reduces a numerical sentinel constant (safeInitRowMax) with explanatory comments. No public APIs or signatures changed. Changes
Sequence Diagram(s)(omitted) Estimated code review effort🎯 3 (Moderate) | ⏱️ ~20 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 (2)
🚧 Files skipped from review as they are similar to previous changes (2)
✨ 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
|
![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
🧹 Nitpick comments (6)
cpp/kernels/xqa/utils.cuh (3)
32-36: Tighten literal type and fix comment typo (‘log2d’ → ‘log2e’).
- log2e is declared as float but initialized with a double literal; add the F suffix to avoid implicit narrowing.
- The comment uses log2d, which is a typo and can confuse readers.
Apply this diff:
-inline constexpr float log2e = 1.4426950408889634; // std::log2(M_E) +inline constexpr float log2e = 1.4426950408889634F; // std::log2(M_E) -/* bias = rowMax * log2d // shared for the whole row +/* bias = rowMax * log2e // shared for the whole row exp(x-rowMax) = exp(x * log2e - bias) */
37-39: Value change to safeInitRowMax: good intent; suggest naming and documentation tweak.Lowering the magnitude to -1e5F aligns with the stability note and should reduce FMA-induced issues. To match the guidelines (“Constants use kPREFIXED_UPPER_SNAKE”), keep the existing name for ABI/back-compat but introduce a named constant.
Apply this diff:
-// this reason, don't set a huge safeInitRowMax. -inline constexpr float safeInitRowMax = -1e+5F; +// This optimization can be numerically unstable when (x * log2e - bias) is fused into an FMA for large |x|. +// To mitigate, avoid an excessively large sentinel for initialization. +inline constexpr float kSAFE_INIT_ROW_MAX = -1e5F; +inline constexpr float safeInitRowMax = kSAFE_INIT_ROW_MAX; // Backward-compatible alias.Optionally, add a one-liner explaining the selection rationale, e.g., “Chosen to be far below plausible QK logits while avoiding extreme products with log2e.”
2116-2139: Avoid FMA fusion in softmax exponent to further improve stability.Current form elem = exp2f(elem * log2e - bias) invites FMA. Pre-subtract then multiply to prevent fusion and lower catastrophic cancellation risk, while maintaining perf. Same change applies to both SWAP_AB and non-SWAP_AB variants.
Apply this diff for both warpGrpOnlineSoftmax() overloads:
- float const maxVal = colMax[n][j]; - float const bias = maxVal * log2e; + float const maxVal = colMax[n][j]; ... - float& elem = acc(m, n)(i, j); - assert(maxVal >= elem); - elem = exp2f(elem * log2e - bias); + float& elem = acc(m, n)(i, j); + assert(maxVal >= elem); + float const delta = elem - maxVal; + elem = exp2f(delta * log2e);And similarly in the row-wise variant:
- float const maxVal = rowMax[m][i]; - float const bias = maxVal * log2e; + float const maxVal = rowMax[m][i]; ... - float& elem = acc(m, n)(i, j); - assert(maxVal >= elem); - elem = exp2f(elem * log2e - bias); + float& elem = acc(m, n)(i, j); + assert(maxVal >= elem); + float const delta = elem - maxVal; + elem = exp2f(delta * log2e);Also applies to: 2440-2463
cpp/kernels/xqa/mha_sm90.cu (3)
1094-1110: Debug prints for non-SWAP_AB row stats: looks good; minor gate suggestion.Gating on blockIdx.y == 1 and threadIdx.x == 0 keeps noise down and the barrier maintains ordering. Consider switching to blockIdx.y == 0 to ensure a print even when nbSubSeq == 1; or add a short comment explaining why y == 1 is preferred.
2950-2954: Optional: clamp normalization denominator to avoid division by tiny sums in corner cases.Extremely adversarial inputs could yield near-zero sums before adding attention sinks. A tiny clamp protects against denorm/Inf without impacting normal runs.
Example:
- auto const regOutScale = __frcp_rn(regColSum) * xvoScale; + auto const regOutScale = __frcp_rn(max(regColSum, Vec<float, decltype(regColSum)::size>::filled(1e-20f))) * xvoScale;Same consideration applies to the row-wise variant.
840-846: Request targeted validation for NaN/Inf eradication on SPEC_DEC.Given the PR goal, please run a stress A/B on SPEC_DEC with adversarial logits to confirm NaN/Inf are gone and that the safeInitRowMax tweak plus current exponent path holds.
If helpful, I can craft a focused harness that:
- Forces all-masked rows and very negative tiles to check softmax behavior.
- Sweeps large-magnitude pre-softmax values to check for NaN/Inf and overflow in exp2f paths.
Also applies to: 1094-1110
📜 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/kernels/xqa/mha_sm90.cu(2 hunks)cpp/kernels/xqa/utils.cuh(1 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/kernels/xqa/utils.cuhcpp/kernels/xqa/mha_sm90.cu
**/*.{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/kernels/xqa/utils.cuhcpp/kernels/xqa/mha_sm90.cu
**/*.{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/kernels/xqa/utils.cuh
**/*.{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/kernels/xqa/utils.cuhcpp/kernels/xqa/mha_sm90.cu
**/*.{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/kernels/xqa/mha_sm90.cu
🔇 Additional comments (2)
cpp/kernels/xqa/mha_sm90.cu (2)
832-836: Sanity: reinit of per-CTA maxima with new sentinel is consistent across passes.Initializing gemm0CurrentSeqMax and gemm1AccColMax with the updated sentinel aligns the two-stage softmax flow. Good.
Also applies to: 1014-1019
633-637: Verify occupancy impact under OPTIMIZE_FOR_LATENCY and eliminate magic numbersI wasn’t able to capture ptxas register/SMEM stats in this environment, so please run the following locally to ensure we don’t regress throughput on mid-/long-sequence shapes and to avoid hard-coding
128 * 3:• File: cpp/kernels/xqa/mha_sm90.cu
• Lines: ~633–637Apply this diff:
- __launch_bounds__(128 * 3, headElems* ctaNbQHeads <= 128 * 16 ? 3 : 2) + __launch_bounds__(warp_size * gmmaWarpsPerGrp * 3, headElems * ctaNbQHeads <= 128 * 16 ? 3 : 2) #else - __launch_bounds__(128 * 3) + __launch_bounds__(warp_size * gmmaWarpsPerGrp * 3)Then locally compile with ptxas statistics enabled:
nvcc -std=c++17 -O3 -DNDEBUG -DOPTIMIZE_FOR_LATENCY=1 -arch=sm_90 \ -Xptxas -v -c cpp/kernels/xqa/mha_sm90.cu -o /dev/null 2>&1 | tee ptxas.txt grep -E "ptxas info.*Used" ptxas.txtVerify:
- Registers per thread
- Shared memory usage
- Achieved occupancy (max resident warps/SM)
across representative sequence lengths and batch sizes to confirm no regression in latency‐optimized builds.
|
/bot run |
|
/bot run |
|
PR_Github #16031 [ run ] triggered by Bot |
|
PR_Github #16035 [ run ] triggered by Bot |
|
PR_Github #16031 [ run ] completed with state |
|
PR_Github #16035 [ run ] completed with state |
|
/bot run |
|
PR_Github #16094 [ run ] triggered by Bot |
|
PR_Github #16094 [ run ] completed with state |
…VIDIA#6282 NVIDIA#6279 * [None][infra] Pin the version for triton to 3.3.1 (NVIDIA#6508) Signed-off-by: qqiao <[email protected]> * [None][infra] Pin the version for triton to 3.3.1 (NVIDIA#6508) (NVIDIA#6519) (NVIDIA#6549) Signed-off-by: Yanchao Lu <[email protected]> * [fix]: use safeInitRowMax instead of fp32_lowest to avoid NaN (NVIDIA#7087) Signed-off-by: Yao Yao <[email protected]> * [None][fix] Fix a numerical stability issue for XQA with spec dec Signed-off-by: Yao Yao <[email protected]> * fix typo Signed-off-by: Jhao-Ting Chen <[email protected]> * fix precompiled multi_query_token kernel not having is_fp8_out hash key (NVIDIA#6279) Signed-off-by: Jhao-Ting Chen <[email protected]> * [fix] Fix missing fields in xqa kernel cache key (NVIDIA#6282) Signed-off-by: Yao Yao <[email protected]> --------- Signed-off-by: qqiao <[email protected]> Signed-off-by: Yanchao Lu <[email protected]> Signed-off-by: Yao Yao <[email protected]> Signed-off-by: Jhao-Ting Chen <[email protected]> Co-authored-by: Emma Qiao <[email protected]> Co-authored-by: Yanchao Lu <[email protected]> Co-authored-by: Yao Yao <[email protected]>
Signed-off-by: Yao Yao <[email protected]>
|
/bot run |
|
PR_Github #17471 [ run ] triggered by Bot |
|
PR_Github #17471 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #17561 [ run ] triggered by Bot |
|
PR_Github #17561 [ run ] completed with state |
…IDIA#7114) Signed-off-by: Yao Yao <[email protected]>
Summary by CodeRabbit