[CUDA] GroupQueryAttention: dequantize the K and V caches in one launch - #31480
Merged
Merged
Conversation
The quantized-KV decode fallback dequantized the two caches back to back. Each launch only exposes batch_size * kv_num_heads independent (batch, head) slices -- two for a batch-1 model with two KV heads -- so most of the device sat idle. Add DequantizeKVPairKernel, which selects the tensor with blockIdx.z and is otherwise identical to DequantizeKVVectorizedKernel. That doubles the CTA count, halves the launches, and lets both caches share the memory pipeline. kVecSize is 8 rather than 32 because at decode lengths the kernel is latency bound, not instruction bound, so more threads with narrower vectors wins. Bit-identical to the two-launch path. INT4, non-16-bit outputs and head sizes that do not tile into the fused block shape still take the old path. Measured on Qwen3.6-35B MTP N=3: 21.7 -> 10.8 launches/step and 59.8 -> 24.8 us/step. Opt out with ORT_DISABLE_FUSED_KV_DEQUANT=1.
Tianlei Wu (tianleiwu)
requested review from
Hariharan Seshadri (hariharans29),
kunal-vaishnavi and
Ti-Tai Wang (titaiwangms)
and
a lite review from Copilot
August 6, 2026 02:15
Contributor
There was a problem hiding this comment.
Pull request overview
This PR optimizes the CUDA GroupQueryAttention decode fallback path for quantized KV cache by fusing K-cache and V-cache dequantization into a single kernel launch, reducing per-step launch overhead.
Changes:
- Add a fused CUDA kernel (
DequantizeKVPairKernel) plus a new launcher (LaunchDequantizeKVPair) to dequantize K and V caches in one launch (with fallback to the prior two-launch behavior when unsupported). - Switch the FlashAttention dequantize-fallback path to call the fused launcher instead of launching two separate dequant kernels.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
| onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh | Introduces fused K+V dequantization kernel/launcher and keeps a fallback path for unsupported formats. |
| onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu | Routes quantized-cache FlashAttention fallback dequantization through the new fused launcher. |
Hariharan Seshadri (hariharans29)
approved these changes
Aug 6, 2026
Contributor
Author
Yes. It's covered by existing tests. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
When GroupQueryAttention runs with a quantized KV cache, the K cache and the V cache were dequantized by two separate kernel launches. They have identical shapes and identical per-head scale layouts, so the second launch adds a full grid setup and a second pass over the same index arithmetic for no reason.
This change dequantizes both caches in a single launch. The kernel is moved into a new
group_query_attention_qdq.cuhheader and given a2 *grid in the cache dimension, so one block range covers K and the other covers V; the buffer pointers and scale pointers are selected from the block index.Output is bit-identical to the two-launch path -- the per-element arithmetic is unchanged, only the launch geometry differs.
Motivation and Context
This is on the decode path of speculative-decoding (MTP) workloads, where the launch is issued every layer, every step, and the per-launch fixed cost is a meaningful fraction of a short kernel. Halving the number of launches removes that fixed cost without changing numerics.
Measured on H200 (SM90) with a Qwen3.6-35B-A3B MTP configuration.