[CUDA] Add sliding-window support to XQA decode

docs/contrib_ops/cuda/gqa.md

@@ -215,7 +215,7 @@ order and the first eligible backend wins:
 
 | Priority | Backend | Selected when (summary) |
 |----------|---------|-------------------------|
-| 1 | **XQA** | Single-token global decode (`seq_len == 1`), shared KV buffer. Fastest decode path; the only backend that serves a quantized KV cache. |
+| 1 | **XQA** | Single-token decode (`seq_len == 1`), shared KV buffer. Supports sliding-window attention on both the non-quantized and quantized (INT8/FP8) paths (attention sink remains non-quantized-only). Fastest decode path; the only backend that serves a quantized KV cache. |
 | 2 | **cuDNN SDPA** | Non-quantized FP16/BF16 causal attention. Auto-preferred on SM≥90 (Hopper/Blackwell). |
 | 3 | **Flash Attention** | General FP16/BF16 prompt and decode, including local window, softcap, and packed QKV. |
 | 4 | **Memory Efficient Attention (MEA)** | Fallback for FP16/FP32 (and BF16 on SM80+). |
@@ -233,8 +233,9 @@ enabled (see §10).
 
 ### 6.1 XQA
 
-Checked first. Used only for single-token global decode under the conditions detailed in §7. When
-XQA is selected, no other backend is considered.
+Checked first. Used only for single-token decode under the conditions detailed in §7 (global or
+sliding-window decode; sliding window is supported on both the non-quantized and quantized paths).
+When XQA is selected, no other backend is considered.
 
 ### 6.2 cuDNN SDPA
 
@@ -294,20 +295,15 @@ XQA (a highly optimized cross/decode attention kernel) is used only when **all**
 4. Past and present KV cache share the same buffer.
 5. No softcap.
 6. Standard softmax, **or** smooth softmax expressed via a `head_sink` tensor (non-quantized KV cache).
-7. No local (sliding) window attention — global attention only.
+7. Global attention, **or** local (sliding) window attention (`local_window_size > 0`), supported on
+   both the non-quantized and quantized (INT8/FP8) paths.
 8. Supported `head_size` (64, 128, or 256) and group size.
 
 `head_sink` (attention sink) is supported on the non-quantized XQA path only. Quantized KV cache
 (int8 / fp8) paths explicitly reject a non-null attention sink, so a GQA node with both `head_sink`
 and a quantized cache falls back to Flash/Flash-Decoding.
 
-XQA selection defaults are:
-
-- **Quantized KV cache (int8 / fp8):** on by default.
-- **Non-quantized with a `head_sink` input:** on by default (GPT-OSS style decode).
-- **Non-quantized without `head_sink`:** opt-in via `ORT_ENABLE_XQA=1`.
-
-Setting `ORT_ENABLE_XQA=0` disables XQA for the non-quantized path regardless of `head_sink`.
+XQA selection is on by default. Setting `ORT_ENABLE_XQA=0` disables XQA.
 
 ## 8. XQA `head_sink` PrePack
 
@@ -354,7 +350,7 @@ sess = ort.InferenceSession(
 
 | Variable | Effect |
 |----------|--------|
-| `ORT_ENABLE_XQA` | `1` enables the XQA decode path for the non-quantized KV cache; `0` disables XQA entirely (including the quantized and `head_sink` default-on paths). Unset: on for quantized / `head_sink`, off otherwise (see §7). |
+| `ORT_ENABLE_XQA` | `1` enables the XQA decode; `0` disables XQA entirely. Unset: on by default. |
 | `ORT_DISABLE_FLASH_ATTENTION` | `1` disables Flash Attention. |
 | `ORT_DISABLE_FLASH_DECODE` | `1` disables the Flash-Decoding split-KV optimization. |
 | `ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION` | `1` disables Memory Efficient Attention. |
@@ -456,10 +452,11 @@ popular LLMs. Listed roughly by impact.
    fused per-head RMSNorm to Q and K before RoPE when `q_norm_weight` / `k_norm_weight` are provided
     (see §3), matching **Qwen3, Gemma 2/3, OLMo2, SmolLM3**, etc. Remaining limitation: QK-Norm
     disables Flash-Decoding, and quantized-cache QK-Norm does not yet get the XQA fast path.
-2. **Sliding-window + attention-sink on the fused decode path.** XQA requires global attention
-   (`local_window_size == -1`), so **GPT-OSS / Mistral / Gemma 2** layers that combine a sliding
-   window with a `head_sink` fall back to Flash / Flash-Decoding instead of XQA. Unifying sliding
-   window and sink in XQA would close this gap.
+2. **Sliding-window on the quantized fused decode path.** *Implemented.* The XQA decode path now
+   serves sliding-window attention (`local_window_size > 0`) on both the non-quantized and the
+   quantized (INT8/FP8) paths. Remaining limitation: the attention sink (`head_sink`) is still
+   non-quantized-only, so quantized **GPT-OSS / Mistral / Gemma 2** sliding-window layers that also
+   use an attention sink fall back to Flash / Flash-Decoding.
 3. **Softcap on the fastest kernels.** Logit soft-capping (**Gemma 2**) disables both XQA and cuDNN
    SDPA, forcing the Flash / MEA / unfused paths. Adding softcap support to XQA and cuDNN would
    recover decode throughput.

onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc

@@ -4,6 +4,7 @@
 #include <algorithm>
 #include <cmath>
 #include <cstdint>
+#include <limits>
 #include <vector>
 #include "core/common/safeint.h"
 #include "core/providers/cuda/cuda_common.h"
@@ -100,7 +101,12 @@ GroupQueryAttention<T, U>::GroupQueryAttention(const OpKernelInfo& info)
   kv_num_heads_ = static_cast<int>(kv_num_heads);
   is_past_bsnh_ = false;
   is_unidirectional_ = true;
-  local_window_size_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("local_window_size", -1));
+  const int64_t local_window_size_attr = info.GetAttrOrDefault<int64_t>("local_window_size", -1);
+  // Validate before narrowing to int so an out-of-range attribute cannot wrap to a valid-looking
+  // small window (e.g. 2^32 + 128) and silently run a different window than the model specifies.
+  ORT_ENFORCE(local_window_size_attr == -1 || (local_window_size_attr > 0 && local_window_size_attr <= std::numeric_limits<int>::max()),
+              "local_window_size must be -1 or greater than 0 (and not exceed INT_MAX).");
+  local_window_size_ = static_cast<int>(local_window_size_attr);
   do_rotary_ = info.GetAttrOrDefault<int64_t>("do_rotary", 0) == 1;
   rotary_interleaved_ = info.GetAttrOrDefault<int64_t>("rotary_interleaved", 0) == 1;
   scale_ = info.GetAttrOrDefault<float>("scale", 0.0f);
@@ -414,26 +420,29 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
   // 4. Past and Present KV cache share the same buffer (required for XQA specific memory access).
   // 5. No Softcap (XQA doesn't support softcap).
   // 6. Standard Softmax, or smooth softmax represented by a head_sink tensor.
-  // 7. No local window attention (global attention only).
+  // 7. Local window (sliding window) attention is supported on both the non-quantized and the
+  //    quantized (INT8/FP8) XQA paths (local_window_size == -1 means global attention).
   // QK-Norm can use XQA for the non-quantized KV-cache path: ExtremeDecoding runs the same
   // UnpackRoPEAppend preprocess before XQA, so Q/K can be normalized before the XQA kernel consumes
   // Q and the appended cache. Keep quantized QK-Norm off the XQA route until scale correctness is
   // validated for normalized K before quantized-cache append.
   const bool xqa_qk_norm_ok = !parameters.use_qk_norm || !is_inputs_quantized;
-  const bool use_xqa_attention_sinks = parameters.use_smooth_softmax && head_sink != nullptr && !is_inputs_quantized;
-  const bool xqa_smooth_softmax_ok = !parameters.use_smooth_softmax || (head_sink != nullptr && !is_inputs_quantized);
+  const bool use_xqa_attention_sinks = head_sink != nullptr && !is_inputs_quantized;
+  const bool is_xqa_smooth_softmax_supported = !parameters.use_smooth_softmax || use_xqa_attention_sinks;
+  // XQA is enabled when enable_xqa_=true; ineligible shapes/group sizes fall back via data.use_xqa below.
   if (enable_xqa_ &&
       (device_prop.major >= 8) &&
       !parameters.is_first_prompt &&
       parameters.sequence_length == 1 &&
       parameters.kv_sequence_length > 0 &&  // Shared KV (kv_seq=0) has no new K/V to append
       parameters.past_present_share_buffer &&
       parameters.softcap == 0.0f &&
-      parameters.local_window_size == -1 &&
       xqa_qk_norm_ok &&
-      xqa_smooth_softmax_ok) {
+      is_xqa_smooth_softmax_supported) {
     int group_size = parameters.num_heads / parameters.kv_num_heads;
 
+    // Sliding window (local_window_size > 0) is wired through to the quantized XQA kernels as well,
+    // so the INT8/FP8 variants no longer need to be restricted to global attention.
     bool is_int8_quantized_supported = is_int8 &&
                                        (k_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                         v_quant_type_ == KVQuantizationType::PER_TENSOR &&

onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu

@@ -767,6 +767,7 @@ Status ExtremeDecoding(
       parameters.head_size,
       parameters.seqlen_present_kv_cache,  // max_seq_len (Capacity)
       scale,
+      parameters.local_window_size,  // -1 means global attention; >0 enables sliding window
       past_bsnh,
       data.past_seq_lens,
       data.xqa_head_sink,

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh

@@ -59,7 +59,10 @@ inline Status Launch(
     [[maybe_unused]] const float* attention_sinks,
     [[maybe_unused]] const float* kv_cache_scale,
     [[maybe_unused]] void* workspace,
-    [[maybe_unused]] size_t workspace_size) {
+    [[maybe_unused]] size_t workspace_size,
+    // No default: every caller must thread local_window_size through explicitly so a future
+    // caller that forgets it fails to compile instead of silently running global attention.
+    [[maybe_unused]] const int local_window_size) {
 #ifdef XQA_HAS_SM80_TARGET
   const InputHead* q_ptr = reinterpret_cast<const InputHead*>(query);
   GMemKVCacheHead* k_ptr = reinterpret_cast<GMemKVCacheHead*>(const_cast<void*>(key_cache));
@@ -92,9 +95,26 @@ inline Status Launch(
     cudaMemsetAsync(semaphores, 0, semaphore_size, stream);
   }
 
+#if SLIDING_WINDOW
+  // SLIDING_WINDOW is #defined to 1 only by the fp16/bf16/int8/fp8 xqa_loader_*_impl.cuh headers
+  // that include this file; it defaults to 0 in defines.h, so this block is dead (and
+  // local_window_size is unused) in any translation unit that does not opt in.
+  // ORT local_window_size semantics: -1 => global attention; >0 => each query attends to the
+  // last local_window_size tokens (including the current one). XQA's slidingWinSize uses the
+  // same "last N tokens incl. current" definition, so pass it through directly. For global
+  // attention, use max_seq_len so the kernel's runtime guard (cacheSeqLen > slidingWinSize) is
+  // never taken and no masking work is performed (numerically identical to the global kernel).
+  uint32_t const sliding_win_size = (local_window_size > 0)
+                                        ? static_cast<uint32_t>(local_window_size)
+                                        : static_cast<uint32_t>(max_seq_len);
+#endif
+
   launchMHA(
       device_prop,
       static_cast<uint32_t>(kv_num_heads),
+#if SLIDING_WINDOW
+      sliding_win_size,
+#endif
       scale,
       out_ptr,
       q_ptr,

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h

@@ -34,6 +34,7 @@ Status LaunchXQAKernel(
     const int head_size,
     const int max_seq_len,  // Max sequence length of cache
     const float scale,
+    const int local_window_size,   // Sliding window size; -1 means global attention (no sliding window)
     const bool is_bsnh,            // Layout of KV cache
     const int* past_seq_lens,      // Past sequence lengths [BatchSize]
     const float* attention_sinks,  // Attention sink per query head, nullptr if not used

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu

@@ -24,6 +24,7 @@ Status LaunchXQAKernelImpl(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,
@@ -48,6 +49,7 @@ Status LaunchXQAKernelImpl(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,
@@ -72,6 +74,7 @@ Status LaunchXQAKernelImpl(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,
@@ -95,6 +98,7 @@ Status LaunchXQAInt8KernelBF16(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* kv_cache_scale,
@@ -119,6 +123,7 @@ Status LaunchXQAKernel<__nv_bfloat16>(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,
@@ -129,22 +134,22 @@ Status LaunchXQAKernel<__nv_bfloat16>(
   // Dispatch to INT8 path if requested
   if (kv_quant_type == XqaQuantType::kInt8) {
     ORT_RETURN_IF(attention_sinks != nullptr, "XQA attention sinks are not supported with INT8 KV cache.");
-    return LaunchXQAInt8KernelBF16(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, kv_cache_scale, workspace, workspace_size);
+    return LaunchXQAInt8KernelBF16(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, local_window_size, is_bsnh, past_seq_lens, kv_cache_scale, workspace, workspace_size);
   }
 
   // Dispatch based on head_size
   if (head_size == 256) {
     return H256::LaunchXQAKernelImpl<__nv_bfloat16>(
         device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size,
-        max_seq_len, scale, is_bsnh, past_seq_lens, attention_sinks, kv_cache_scale, kv_quant_type, workspace, workspace_size);
+        max_seq_len, scale, local_window_size, is_bsnh, past_seq_lens, attention_sinks, kv_cache_scale, kv_quant_type, workspace, workspace_size);
   } else if (head_size == 128) {
     return H128::LaunchXQAKernelImpl<__nv_bfloat16>(
         device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size,
-        max_seq_len, scale, is_bsnh, past_seq_lens, attention_sinks, kv_cache_scale, kv_quant_type, workspace, workspace_size);
+        max_seq_len, scale, local_window_size, is_bsnh, past_seq_lens, attention_sinks, kv_cache_scale, kv_quant_type, workspace, workspace_size);
   } else if (head_size == 64) {
     return H64::LaunchXQAKernelImpl<__nv_bfloat16>(
         device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size,
-        max_seq_len, scale, is_bsnh, past_seq_lens, attention_sinks, kv_cache_scale, kv_quant_type, workspace, workspace_size);
+        max_seq_len, scale, local_window_size, is_bsnh, past_seq_lens, attention_sinks, kv_cache_scale, kv_quant_type, workspace, workspace_size);
   } else {
     return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA only supports head_size=64, 128, or 256. Input has ", head_size);
   }

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_128.cu

@@ -24,6 +24,7 @@ template Status HEAD_DIM_NAMESPACE::LaunchXQAKernelImpl<__nv_bfloat16>(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_256.cu

@@ -24,6 +24,7 @@ template Status HEAD_DIM_NAMESPACE::LaunchXQAKernelImpl<__nv_bfloat16>(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_64.cu

@@ -24,6 +24,7 @@ template Status HEAD_DIM_NAMESPACE::LaunchXQAKernelImpl<__nv_bfloat16>(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* attention_sinks,

onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_fp8_impl.cuh

@@ -21,6 +21,10 @@
 #define TOKENS_PER_PAGE 0
 #define INPUT_FP16 0  // Q is BF16
 #define ALLOW_MULTI_BLOCK_MODE 1
+// Compile the FP8 KV-cache XQA kernels with sliding-window support so the same kernels
+// can serve both global attention (local_window_size == -1, mapped to a window >= max_seq_len
+// -> zero masking overhead) and sliding-window models.
+#define SLIDING_WINDOW 1
 
 #pragma nv_diag_suppress 177
 #pragma nv_diag_suppress 20012
@@ -94,6 +98,7 @@ Status LaunchXQAFp8KernelBF16(
     const int head_size,
     const int max_seq_len,
     const float scale,
+    const int local_window_size,
     const bool is_bsnh,
     const int* past_seq_lens,
     const float* kv_cache_scale,
@@ -102,13 +107,13 @@ Status LaunchXQAFp8KernelBF16(
   int group_size = num_heads / kv_num_heads;
   switch (group_size) {
     case 4:
-      return grp4_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size);
+      return grp4_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 8:
-      return grp8_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size);
+      return grp8_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 16:
-      return grp16_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size);
+      return grp16_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 32:
-      return grp32_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size);
+      return grp32_bf16_fp8::Launch<__nv_bfloat16>(device_prop, stream, query, key_cache, value_cache, output, batch_size, num_heads, kv_num_heads, head_size, max_seq_len, scale, is_bsnh, past_seq_lens, nullptr, kv_cache_scale, workspace, workspace_size, local_window_size);
     default:
       return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA FP8 only supports group_size 4, 8, 16, 32. Input has ", group_size);
   }