[CUDA] Normalize the PER_CHANNEL K scale folded into Q for paged XQA

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`.

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

onnxruntime/contrib_ops/cuda/bert/paged_attention.cc

@@ -733,6 +733,7 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
   IAllocatorUniquePtr<void> xqa_workspace_buffer;
   IAllocatorUniquePtr<void> xqa_page_table_buffer;
   IAllocatorUniquePtr<void> xqa_query_buffer;
+  IAllocatorUniquePtr<void> xqa_k_scale_norm_buffer;
   IAllocatorUniquePtr<void> xqa_head_sink_buffer;
   IAllocatorUniquePtr<void> xqa_spec_dec_mask_buffer;
   size_t xqa_workspace_bytes = 0;
@@ -762,6 +763,7 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
       xqa_query_buffer = GetScratchBuffer<void>(
           sizeof(T) * static_cast<size_t>(parameters.token_count) * parameters.num_heads * parameters.head_size,
           GetComputeStream(context));
+      xqa_k_scale_norm_buffer = GetScratchBuffer<void>(sizeof(float), GetComputeStream(context));
     }
     if (parameters.use_smooth_softmax && head_sink != nullptr) {
       xqa_head_sink_buffer = GetScratchBuffer<void>(sizeof(float) * parameters.num_heads,
@@ -868,6 +870,7 @@ Status PagedAttention<T, TCACHE>::ComputeInternal(OpKernelContext* context) cons
     data.xqa_workspace_size = xqa_workspace_bytes;
     data.xqa_page_table_scratch = reinterpret_cast<int*>(xqa_page_table_buffer.get());
     data.xqa_query = reinterpret_cast<CudaT*>(xqa_query_buffer.get());
+    data.xqa_k_scale_norm = reinterpret_cast<float*>(xqa_k_scale_norm_buffer.get());
     data.xqa_head_sink = reinterpret_cast<float*>(xqa_head_sink_buffer.get());
     data.xqa_spec_dec_mask = reinterpret_cast<uint32_t*>(xqa_spec_dec_mask_buffer.get());
   }

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 <typename T>
 __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<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
@@ -1584,7 +1587,39 @@ __global__ void PagedFoldChannelScaleKernel(T* dst,
   }
   const int h = static_cast<int>(i / head_size) % num_heads;
   const int c = static_cast<int>(i % head_size);
-  dst[i] = static_cast<T>(static_cast<float>(src[i]) * channel_scale[(h / group_size) * head_size + c]);
+  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<T>(static_cast<float>(src[i]) * normalized_scale);
+}
+
+// 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 <typename T>
@@ -1675,9 +1710,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<int>((q_elements + max_threads_per_block - 1) / max_threads_per_block);
     PagedFoldChannelScaleKernel<T><<<blocks, max_threads_per_block, 0, stream>>>(
-        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 +1742,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 +1784,8 @@ Status PagedXqaDecodeAttention(
   if (v_per_channel) {
     const int blocks = static_cast<int>((q_elements + max_threads_per_block - 1) / max_threads_per_block);
     PagedFoldChannelScaleKernel<T><<<blocks, max_threads_per_block, 0, stream>>>(
-        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());
   }
 

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
@@ -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<nbKBuffers> 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) {
@@ -2459,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) {
@@ -2540,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

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);

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,

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.