From 0512421f255f2427f66c0640e3f3f72f242c1c72 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 9 Sep 2026 21:34:41 +0000 Subject: [PATCH 1/2] Normalize the PER_CHANNEL K scale folded into Q for XQA XQA folds the per-channel K scale into the query and stores the product in T, while the portable kernel keeps it in FP32. A large scale saturated FP16 there, and a zero cache code turned that infinity into NaN, so the two backends could disagree on valid input. INT4 is the most exposed because its scale spans max|K| / 7 rather than max|K| / 127, but the fold is shared with INT8 and FP8. Divide the fold by max|k_scale| and pass that maximum to XQA as its scalar K scale, which it multiplies back into qkScale once per CTA. The correction is exact, bounds the folded query by max|q| at any scale, and adds no inner-loop work: only a single-block max reduction over the scale table, launched on device so the step stays capturable. INT4 dequantizes into FP16 shared memory, so isKVCacheQuantized is false and the scalar scale was ignored; gate that on a separate predicate instead. NOT BUILT: no CUDA toolkit on the dev machine. Needs a GPU build plus the perf comparison against the previous head before this is merged. --- docs/contrib_ops/cuda/paged_attention.md | 15 +++--- .../contrib_ops/cuda/bert/attention_data.h | 3 ++ .../contrib_ops/cuda/bert/paged_attention.cc | 3 ++ .../cuda/bert/paged_attention_impl.cu | 51 ++++++++++++++++--- .../contrib_ops/cuda/bert/xqa/mha_impl.cuh | 11 +++- .../transformers/test_paged_attention_int4.py | 21 ++++---- 6 files changed, 77 insertions(+), 27 deletions(-) diff --git a/docs/contrib_ops/cuda/paged_attention.md b/docs/contrib_ops/cuda/paged_attention.md index 779fa3bd93546..ce9a8b9815306 100644 --- a/docs/contrib_ops/cuda/paged_attention.md +++ b/docs/contrib_ops/cuda/paged_attention.md @@ -830,13 +830,14 @@ when reusing a directory configured with the feature disabled. INT8 kernels are > softmax denominator. > - The kernel reads pages in place at their stored width, so a decode step touches the KV cache once > at `int8`/`fp8` bandwidth instead of gathering and dequantizing the whole live context. -> - **XQA folds the `PER_CHANNEL` K scale into the query in `T`, not FP32.** `PagedFoldChannelScaleKernel` -> multiplies in FP32 but stores `q_c * k_scale_c` back as FP16/BF16, whereas the portable kernel keeps -> that product in an FP32 shared-memory tile. The two agree only while `max|q_c * k_scale_c|` is -> representable in `T`; beyond that FP16 saturates to infinity, and a zero cache code then yields NaN. -> For FP16 the bound is 65504, which INT4 reaches ~18x sooner than INT8 at equal data, because an INT4 -> scale covers `max|K| / 7` instead of `max|K| / 127`. Removing the bound requires applying the scale -> during the cache load, where `DequantizeInt4CacheGrain` currently dequantizes at unit scale. +> - **XQA normalizes a `PER_CHANNEL` K scale before folding it into the query.** +> `PagedFoldChannelScaleKernel` stores `q_c * k_scale_c` back in `T`, so a large scale would saturate +> FP16 and a zero cache code would then yield NaN, while the portable kernel keeps that product in +> FP32. The fold therefore divides by `max|k_scale|`, which `PagedMaxAbsScaleKernel` computes on +> device so the step stays capturable, and hands that maximum to XQA as its scalar K scale. XQA +> multiplies it back into `qkScale` once per CTA, outside the K/V loop, so the correction is exact +> and adds no inner-loop work. The folded query is then bounded by `max|q|` at any scale magnitude. +> INT4 needs this more than INT8 because its scale spans `max|K| / 7` rather than `max|K| / 127`. > - `softcap` matches FlashAttention bit-for-bit: `softcap * tanh(qk_raw * scale / softcap)`, which is > what `flash_api.cc` produces from `params.softcap = softmax_scale / softcap` and > `params.scale_softmax = softcap`. diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index d8b121de912dd..9ce2d3d53139b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -290,11 +290,14 @@ struct PagedAttentionData { // xqa_page_table_scratch : mutable destination for expansion when block_size is greater than 128. // xqa_query : scratch for Q pre-scaled by a PER_CHANNEL k_scale; unused otherwise. // xqa_head_sink : head_sink converted to fp32, which is what XQA consumes. + // xqa_k_scale_norm : max|k_scale| divided out of that pre-scaled Q and handed to XQA as its + // scalar K scale, so the FP16 copy of Q cannot overflow on a large scale. void* xqa_workspace = nullptr; size_t xqa_workspace_size = 0; int* xqa_page_table_scratch = nullptr; T* xqa_query = nullptr; float* xqa_head_sink = nullptr; + float* xqa_k_scale_norm = nullptr; uint32_t* xqa_spec_dec_mask = nullptr; // Output Tensors diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc index 427969e302bad..d21fd99bf4852 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention.cc @@ -733,6 +733,7 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons IAllocatorUniquePtr xqa_workspace_buffer; IAllocatorUniquePtr xqa_page_table_buffer; IAllocatorUniquePtr xqa_query_buffer; + IAllocatorUniquePtr xqa_k_scale_norm_buffer; IAllocatorUniquePtr xqa_head_sink_buffer; IAllocatorUniquePtr xqa_spec_dec_mask_buffer; size_t xqa_workspace_bytes = 0; @@ -762,6 +763,7 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons xqa_query_buffer = GetScratchBuffer( sizeof(T) * static_cast(parameters.token_count) * parameters.num_heads * parameters.head_size, GetComputeStream(context)); + xqa_k_scale_norm_buffer = GetScratchBuffer(sizeof(float), GetComputeStream(context)); } if (parameters.use_smooth_softmax && head_sink != nullptr) { xqa_head_sink_buffer = GetScratchBuffer(sizeof(float) * parameters.num_heads, @@ -868,6 +870,7 @@ Status PagedAttention::ComputeInternal(OpKernelContext* context) cons data.xqa_workspace_size = xqa_workspace_bytes; data.xqa_page_table_scratch = reinterpret_cast(xqa_page_table_buffer.get()); data.xqa_query = reinterpret_cast(xqa_query_buffer.get()); + data.xqa_k_scale_norm = reinterpret_cast(xqa_k_scale_norm_buffer.get()); data.xqa_head_sink = reinterpret_cast(xqa_head_sink_buffer.get()); data.xqa_spec_dec_mask = reinterpret_cast(xqa_spec_dec_mask_buffer.get()); } diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu index 40d918695d3fd..19b15ebf7bcbf 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu @@ -1572,10 +1572,13 @@ __global__ void ExpandBlockTableToPages(const int* __restrict__ block_table, // Multiply every head vector by a PER_CHANNEL scale indexed [kv_head, channel]. Used to fold // k_scale into Q before XQA and v_scale into XQA's output afterwards. dst may alias src (the // output scaling is done in place), so neither pointer is marked __restrict__. +// When scale_norm is set, the scale is divided by it first; XQA multiplies the same value back into +// qkScale, which keeps the folded product inside T's range without changing the result. template __global__ void PagedFoldChannelScaleKernel(T* dst, const T* src, const float* __restrict__ channel_scale, + const float* __restrict__ scale_norm, const int num_heads, const int head_size, const int group_size, const int64_t total_elements) { const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; @@ -1584,7 +1587,38 @@ __global__ void PagedFoldChannelScaleKernel(T* dst, } const int h = static_cast(i / head_size) % num_heads; const int c = static_cast(i % head_size); - dst[i] = static_cast(static_cast(src[i]) * channel_scale[(h / group_size) * head_size + c]); + const float norm = (scale_norm == nullptr) ? 1.0f : (1.0f / scale_norm[0]); + dst[i] = static_cast(static_cast(src[i]) * channel_scale[(h / group_size) * head_size + c] * norm); +} + +// Largest magnitude in a PER_CHANNEL scale table, computed in one block so the XQA path stays +// capturable. A table of all zeros would make the normalized fold 0/0, so it reports 1 instead. +__global__ void PagedMaxAbsScaleKernel(float* __restrict__ out, const float* __restrict__ scale, const int count) { + constexpr int kWarpSize = 32; + __shared__ float warp_max[kWarpSize]; + float local = 0.0f; + for (int i = threadIdx.x; i < count; i += blockDim.x) { + local = fmaxf(local, fabsf(scale[i])); + } + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + local = fmaxf(local, __shfl_down_sync(0xffffffffu, local, offset)); + } + const int lane = threadIdx.x % kWarpSize; + const int warp = threadIdx.x / kWarpSize; + if (lane == 0) { + warp_max[warp] = local; + } + __syncthreads(); + if (warp == 0) { + const int num_warps = (blockDim.x + kWarpSize - 1) / kWarpSize; + local = lane < num_warps ? warp_max[lane] : 0.0f; + for (int offset = kWarpSize / 2; offset > 0; offset >>= 1) { + local = fmaxf(local, __shfl_down_sync(0xffffffffu, local, offset)); + } + if (lane == 0) { + out[0] = local > 0.0f ? local : 1.0f; + } + } } template @@ -1675,9 +1709,13 @@ Status PagedXqaDecodeAttention( if (k_per_channel) { // Q may point straight at the (const) graph input when there is no packed-QKV / rotary // prologue, so the scaled copy always goes to a dedicated scratch buffer. + ORT_RETURN_IF_NOT(data.xqa_k_scale_norm, "XQA k_scale normalizer scratch was not allocated."); + PagedMaxAbsScaleKernel<<<1, 256, 0, stream>>>(data.xqa_k_scale_norm, data.k_scale, kv_num_heads * head_size); + CUDA_RETURN_IF_ERROR(cudaGetLastError()); const int blocks = static_cast((q_elements + max_threads_per_block - 1) / max_threads_per_block); PagedFoldChannelScaleKernel<<>>( - data.xqa_query, query, data.k_scale, num_heads, head_size, num_heads / kv_num_heads, q_elements); + data.xqa_query, query, data.k_scale, data.xqa_k_scale_norm, num_heads, head_size, + num_heads / kv_num_heads, q_elements); CUDA_RETURN_IF_ERROR(cudaGetLastError()); query = data.xqa_query; } @@ -1703,9 +1741,9 @@ Status PagedXqaDecodeAttention( kIsInt4Cache ? XqaQuantType::kInt4 : kIsFp8Cache ? XqaQuantType::kFp8 : (kIsInt8Cache ? XqaQuantType::kInt8 : XqaQuantType::kNone); - // A PER_CHANNEL scale has already been folded into Q / will be applied to the output, so XQA - // receives a null scalar scale (which means one). - const float* xqa_k_scale = k_per_channel ? nullptr : data.k_scale; + // A PER_CHANNEL K scale is folded into Q up to max|k_scale|, which XQA reapplies as its scalar + // scale; a PER_CHANNEL V scale is applied to the output below, so XQA sees a null scale (one). + const float* xqa_k_scale = k_per_channel ? data.xqa_k_scale_norm : data.k_scale; const float* xqa_v_scale = v_per_channel ? nullptr : data.v_scale; if (data.use_xqa_spec_dec) { ORT_RETURN_IF_NOT(data.xqa_spec_dec_mask, "Speculative XQA mask scratch was not allocated."); @@ -1745,7 +1783,8 @@ Status PagedXqaDecodeAttention( if (v_per_channel) { const int blocks = static_cast((q_elements + max_threads_per_block - 1) / max_threads_per_block); PagedFoldChannelScaleKernel<<>>( - data.output, data.output, data.v_scale, num_heads, head_size, num_heads / kv_num_heads, q_elements); + data.output, data.output, data.v_scale, nullptr, num_heads, head_size, + num_heads / kv_num_heads, q_elements); CUDA_RETURN_IF_ERROR(cudaGetLastError()); } diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh index 7ba172630e83c..478309a8e1e77 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh @@ -1489,6 +1489,13 @@ CUBIN_EXPORT __global__ const uint32_t seqStrideIters = nbSubSeqPerSeq; constexpr bool isKVCacheQuantized = (cacheElemSize < 2); +#if defined(XQA_PAGED_INT4) + // INT4 dequantizes into FP16 shared memory, so cacheElemSize is 2 and isKVCacheQuantized is false, + // yet the packed codes still carry the caller's scalar dequant factor. + constexpr bool hasScalarCacheScale = true; +#else + constexpr bool hasScalarCacheScale = isKVCacheQuantized; +#endif const uint32_t seqIterInit = nbSkipLeadingTiles + idxSubSeqInSeq; #if BEAM_WIDTH > 1 const uint32_t nbCtxCtaTiles = beamSearchParams.ctxLenList[idxReq * beamWidth] / ctaTile.x; @@ -1502,7 +1509,7 @@ CUBIN_EXPORT __global__ }; if (warpIdx.z == 0) { // qkScale is applied onto Q*K.T before softmax. A null kCacheScale means the scale is already in Q. - const float qkScale = qScale * ((isKVCacheQuantized && kCacheScale != nullptr) ? kCacheScale[0] : 1.f); + const float qkScale = qScale * ((hasScalarCacheScale && kCacheScale != nullptr) ? kCacheScale[0] : 1.f); CircIdx idxCurrSMemKBuf{nbKBuffers - 1}; const auto getSMemKTile = [&](uint32_t idx) -> SharedMem::KSmemBuffer& { return smem.k[warpIdx.x][idx]; }; #if BEAM_WIDTH > 1 @@ -2195,7 +2202,7 @@ CUBIN_EXPORT __global__ } // A null vCacheScale means the caller rescales the output itself (per-channel V scale). - float voScale = ((isKVCacheQuantized && vCacheScale != nullptr) ? vCacheScale[0] : 1.F); + float voScale = ((hasScalarCacheScale && vCacheScale != nullptr) ? vCacheScale[0] : 1.F); if (seqIterInit < nbSeqIters) { // otherwise rcpRowSum will be NAN. // The attention sinks are moved to the multi-block reduction part if the multi-block is enabled. if (!isMultiBlock && attentionSinks != nullptr) { diff --git a/onnxruntime/test/python/transformers/test_paged_attention_int4.py b/onnxruntime/test/python/transformers/test_paged_attention_int4.py index bef633c884a8b..8c2db84dcd3cd 100644 --- a/onnxruntime/test/python/transformers/test_paged_attention_int4.py +++ b/onnxruntime/test/python/transformers/test_paged_attention_int4.py @@ -721,23 +721,20 @@ def test_int4_per_channel_xqa_matches_portable(self): @unittest.skipUnless(has_sm80_cuda(), "XQA requires an SM80 or newer GPU") def test_int4_xqa_large_per_channel_k_scale_matches_portable(self): - # XQA folds the K scale into the query and stores it as FP16, while the portable kernel - # keeps that product in FP32. The folded query must therefore stay within the FP16 range; - # this pins the largest scale that does, over channels holding both zero and nonzero codes. + # XQA folds the K scale into the query and stores it in T. A large PER_CHANNEL scale used to + # saturate FP16 there, and a zero cache code then turned that infinity into NaN. The fold now + # divides by max|k_scale| and the kernel reapplies it. Every K code is zero here, so the logits + # are uniform and the expected output stays well conditioned no matter how large the scale is. heads, kv_heads, width = 24, 4, 256 model, feeds, _ = make_case(width=width, heads=heads, kv_heads=kv_heads, past=(513, 138), block_size=256) - zero_channels = width // 2 - for side in ("key", "value"): - feeds[f"{side}_cache"][..., : zero_channels // 2] = 0x88 # two zero codes per byte + feeds["key"][:] = 0 + feeds["key_cache"][:] = 0x88 # two zero codes per byte query = np.abs(feeds["query"].reshape(-1, heads, width).astype(np.float32)) - k_scale = np.full((kv_heads, 1, width), 0.1, dtype=np.float32) - for kv_head in range(kv_heads): - group = query[:, kv_head * (heads // kv_heads) : (kv_head + 1) * (heads // kv_heads), :zero_channels] - k_scale[kv_head, 0, :zero_channels] = 30000.0 / np.maximum(group.max(axis=(0, 1)), 1e-3) + k_scale = np.tile(np.linspace(0.5, 1.0, width, dtype=np.float32), (kv_heads, 1)).reshape(kv_heads, 1, width) + k_scale *= np.float32(1.0e6 / (query * k_scale[:, 0, :].repeat(heads // kv_heads, axis=0)).max()) replace_input(model, feeds, "k_scale", k_scale) - folded = query[:, :, :zero_channels] * k_scale[:, 0, :zero_channels].repeat(heads // kv_heads, axis=0) - self.assertLess(folded.max(), np.finfo(np.float16).max) + self.assertGreater((query * k_scale[:, 0, :].repeat(heads // kv_heads, axis=0)).max(), np.finfo(np.float16).max) with patch.dict(os.environ, {"ORT_ENABLE_XQA": "0"}): portable = run_with_kernel(model, feeds, "DECODER_ATTENTION")[0] From 8b52ae2af171b6e5223215aa1ea6b11a2f3c75d6 Mon Sep 17 00:00:00 2001 From: Tianlei WU Date: Wed, 9 Sep 2026 17:09:19 -0700 Subject: [PATCH 2/2] address feedbacks --- .../cuda/bert/paged_attention_impl.cu | 5 ++- .../contrib_ops/cuda/bert/xqa/mha_impl.cuh | 8 ++-- .../cuda/bert/xqa/xqa_paged_loader.cu | 2 +- .../cuda/bert/xqa/xqa_paged_loader.h | 11 ++--- .../transformers/test_paged_attention.py | 41 ++++++++++++++++++- 5 files changed, 53 insertions(+), 14 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu index 19b15ebf7bcbf..4ebbb8c82e4e7 100644 --- a/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu @@ -1587,8 +1587,9 @@ __global__ void PagedFoldChannelScaleKernel(T* dst, } const int h = static_cast(i / head_size) % num_heads; const int c = static_cast(i % head_size); - const float norm = (scale_norm == nullptr) ? 1.0f : (1.0f / scale_norm[0]); - dst[i] = static_cast(static_cast(src[i]) * channel_scale[(h / group_size) * head_size + c] * norm); + const float scale = channel_scale[(h / group_size) * head_size + c]; + const float normalized_scale = (scale_norm == nullptr) ? scale : (scale / scale_norm[0]); + dst[i] = static_cast(static_cast(src[i]) * normalized_scale); } // Largest magnitude in a PER_CHANNEL scale table, computed in one block so the XQA path stays diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh index 478309a8e1e77..f587bc8beeb06 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/mha_impl.cuh @@ -1282,7 +1282,7 @@ CUBIN_EXPORT __global__ #endif #endif const uint32_t batchSize, - // Device memory scalars, used only for int8/fp8 KV cache. K and V have independent scales: + // Device memory scalars for quantized KV cache. K and V have independent scales: // kCacheScale is folded into qkScale (applied to Q*K.T before softmax) and vCacheScale into // voScale (applied to the P*V accumulator). Both are read once per CTA, outside the K/V loop. // Either may be null, meaning "scale is 1": the caller has already folded a non-scalar @@ -2466,7 +2466,7 @@ CUBIN_EXPORT __global__ __launch_bounds__(256, nbCtaPerSM) void kernel_mha( const BeamSearchParams beamSearchParams, #endif const uint32_t batchSize, - // Device memory scalars, used only for int8/fp8 KV cache. See kernel_mha_impl. + // Device memory scalars for quantized KV cache. See kernel_mha_impl. const float* __restrict__ kCacheScale, const float* __restrict__ vCacheScale, uint32_t* __restrict__ semaphores = nullptr, void* __restrict__ scratch = nullptr) { @@ -2547,8 +2547,8 @@ void launchMHA(const cudaDeviceProp& prop, uint32_t nbKHeads, const BeamSearchParams& beamSearchParams, #endif uint32_t batchSize, - // Device memory scalars, used only for int8/fp8 KV cache. K and V may have different - // scales; both are per-tensor (a single float each). + // Device memory scalars for quantized KV cache. K and V may have different scales; + // each is either a per-tensor scale or a normalizer for a folded per-channel scale. const float* __restrict__ kCacheScale, const float* __restrict__ vCacheScale, #if SPEC_DEC diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu index 6910e99c4c92c..929d1073b232f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.cu @@ -144,7 +144,7 @@ Status LaunchXQAPagedKernel( #ifdef USE_INT4_KV_CACHE if (kv_quant_type == XqaQuantType::kInt4) { - // Null scales are the PER_CHANNEL case: the caller folded them into Q and the output. + // The caller passes the K folding normalizer and applies the folded V scale to the output. ORT_RETURN_IF_NOT(head_size == 256 && !is_bf16 && kv_num_heads > 0 && num_heads == 6 * kv_num_heads, "INT4 paged XQA requires FP16 queries, head_size 256, and group size 6."); return H256::LaunchXQAPagedInt4Kernel(XQA_PAGED_ARGS); diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h index acd6d2320dc4c..e933a23ce0e37 100644 --- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h +++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_paged_loader.h @@ -26,9 +26,10 @@ constexpr int kXqaTokensPerPage = 128; // Paged-KV XQA decode launcher. Unlike LaunchXQAKernel (contiguous per-request cache) this reads // K and V from a shared block pool addressed through a page table. // kInt4 uses packed UINT8 heads with static FP32 PER_CHANNEL scales folded into Q and the output -// by the caller; both scale pointers must be null. It supports FP16 query/output, head_size 256, -// and group_size 6 only. Other quantized types use -// FP32 per-tensor scales. The INT4 shared-memory and scratch layouts match native FP16 XQA. +// by the caller. The K fold is divided by max|k_scale| and k_cache_scale carries that normalizer; +// v_cache_scale is null because the V scale is applied to the output. It supports FP16 query/output, +// head_size 256, and group_size 6 only. Other quantized types use FP32 per-tensor scales or the same +// normalized PER_CHANNEL folding. The INT4 shared-memory and scratch layouts match native FP16 XQA. // // Preconditions: one query token per sequence, head_size in {64, 128, 256}, group_size in // {4, 6, 8, 16, 32}, supported FP16/INT8/FP8/INT4 cache, block_size % kXqaTokensPerPage == 0. @@ -50,8 +51,8 @@ Status LaunchXQAPagedKernel( const int local_window_size, // -1 => global attention const int* past_seq_lens, // [batch_size]; the kernel attends to past_seq_lens[i] + 1 tokens const float* attention_sinks, // [num_heads] fp32, nullptr if unused - const float* k_cache_scale, // per-tensor dequant scale; nullptr means "1" (folded into Q) - const float* v_cache_scale, // per-tensor dequant scale; nullptr means "1" (applied to output) + const float* k_cache_scale, // per-tensor scale or folded-scale normalizer; nullptr means "1" + const float* v_cache_scale, // per-tensor scale; nullptr means "1" (applied to output) const XqaQuantType kv_quant_type, const bool is_bf16, // dtype of query and output void* workspace, diff --git a/onnxruntime/test/python/transformers/test_paged_attention.py b/onnxruntime/test/python/transformers/test_paged_attention.py index f9803b318fce0..2bb7b18e0a269 100644 --- a/onnxruntime/test/python/transformers/test_paged_attention.py +++ b/onnxruntime/test/python/transformers/test_paged_attention.py @@ -869,6 +869,7 @@ def parity_check_paged_attention( new_seqlens_override=None, local_window_size_override=None, past_seqlens_override=None, + k_scale_max_override=None, ): # Generate padded inputs q = torch.randn( @@ -1005,6 +1006,9 @@ def parity_check_paged_attention( k_scale = compute_kv_scale( [k_cache_paged, k_ro], config.k_quant_type, config.kv_cache_type, config.kv_num_heads, config.head_size ) + if k_scale_max_override is not None: + assert config.k_quant_type == "PER_CHANNEL" + k_scale *= k_scale_max_override / k_scale.max() v_scale = compute_kv_scale( [v_cache_paged, v_new], config.v_quant_type, config.kv_cache_type, config.kv_num_heads, config.head_size ) @@ -2289,7 +2293,16 @@ def _config(self, **overrides): setattr(config, key, value) return config - def _check_xqa(self, quant_type="PER_TENSOR", kv_cache_type="int8", rtol=5e-3, atol=5e-3, **overrides): + def _check_xqa( + self, + quant_type="PER_TENSOR", + kv_cache_type="int8", + rtol=5e-3, + atol=5e-3, + k_scale_max_override=None, + require_xqa=False, + **overrides, + ): if kv_cache_type == "fp8": if not has_fp8_kv_cache(): self.skipTest("FP8 KV cache kernels are not built") @@ -2302,7 +2315,21 @@ def _check_xqa(self, quant_type="PER_TENSOR", kv_cache_type="int8", rtol=5e-3, a v_quant_type=quant_type, **overrides, ) - parity_check_paged_attention(config, rtol=rtol, atol=atol) + + def run(): + parity_check_paged_attention( + config, rtol=rtol, atol=atol, k_scale_max_override=k_scale_max_override + ) + + if require_xqa: + with patch.dict( + os.environ, + {"ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO": "1", "ORT_ENABLE_XQA": "1"}, + ): + debug_output = capture_native_stdout(run) + self.assertIn("SdpaKernel=XQA", debug_output) + return + run() def _capture_xqa_debug(self, config): with patch.dict( @@ -2486,6 +2513,16 @@ def test_xqa_context_not_page_aligned(self): def test_xqa_quant_type(self, _, kv_cache_type, quant_type): self._check_xqa(kv_cache_type=kv_cache_type, quant_type=quant_type) + @parameterized.expand([("int8", "int8"), ("fp8", "fp8")]) + def test_xqa_large_per_channel_k_scale(self, _, kv_cache_type): + # A finite channel scale can overflow FP32 when multiplied by Q before normalization. + self._check_xqa( + kv_cache_type=kv_cache_type, + quant_type="PER_CHANNEL", + k_scale_max_override=torch.finfo(torch.float32).max, + require_xqa=True, + ) + def test_xqa_mixed_granularity(self): # k PER_CHANNEL folds into Q, v PER_TENSOR stays a kernel argument: the two scales take # different routes, so an asymmetric config catches a mix-up between them.