[CUDA] Add sliding-window support to XQA decode

docs/contrib_ops/cuda/gqa.md

@@ -188,7 +188,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 the non-quantized path; quantized (INT8/FP8) variants remain global-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+). |
@@ -199,8 +199,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 decode,
+or sliding-window decode on the non-quantized path). When XQA is selected, no other backend is
+considered.
 
 ### 6.2 cuDNN SDPA
 
@@ -260,20 +261,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`) on the
+   non-quantized path. The quantized (INT8/FP8) variants remain global-only.
 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
 
@@ -320,7 +316,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. |
@@ -417,10 +413,12 @@ popular LLMs. Listed roughly by impact.
 1. **Fused QK-Norm (per-head Q/K RMSNorm prologue).** The CUDA kernel rejects `q_norm_weight` /
    `k_norm_weight`; only the WebGPU EP implements the fused prologue. Required by **Qwen3,
    Gemma 2/3, OLMo2, SmolLM3**, etc., which otherwise run the normalization unfused.
-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.** The non-quantized XQA decode path now
+   serves sliding-window attention (`local_window_size > 0`), including in combination with a
+   `head_sink`. The quantized (INT8/FP8) XQA variants still require global attention
+   (`local_window_size == -1`), so quantized **GPT-OSS / Mistral / Gemma 2** sliding-window layers
+   fall back to Flash / Flash-Decoding. Wiring sliding window into the quantized kernels would close
+   this gap.
 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

@@ -100,6 +100,8 @@ GroupQueryAttention<T, U>::GroupQueryAttention(const OpKernelInfo& info)
   is_past_bsnh_ = false;
   is_unidirectional_ = true;
   local_window_size_ = static_cast<int>(info.GetAttrOrDefault<int64_t>("local_window_size", -1));
+  ORT_ENFORCE(local_window_size_ == -1 || local_window_size_ > 0,
+              "local_window_size must be -1 or greater than 0.");
   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);
@@ -388,7 +390,8 @@ 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 the non-quantized path; the
+  //    quantized (INT8/FP8) paths remain global-only (see the per-variant checks below).
   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.
@@ -399,19 +402,22 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
       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 &&
-      is_xqa_smooth_softmax_supported &&
-      parameters.local_window_size == -1) {
+      is_xqa_smooth_softmax_supported) {
     int group_size = parameters.num_heads / parameters.kv_num_heads;
 
-    bool is_int8_quantized_supported = is_int8 &&
+    // Sliding window (local_window_size > 0) is only wired into the non-quantized XQA kernels.
+    // The quantized variants stay restricted to global attention.
+    const bool is_global_attention = parameters.local_window_size == -1;
+
+    bool is_int8_quantized_supported = is_int8 && is_global_attention &&
                                        (k_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                         v_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                         data.k_scale == data.v_scale &&  // XQA requires k_scale and v_scale to be the same. Here requires k_scale and v_scale are same tensor.
                                         (parameters.head_size == 256 || parameters.head_size == 128 || parameters.head_size == 64) &&
                                         (group_size == 4 || group_size == 8 || group_size == 16 || group_size == 32));
 
 #ifdef USE_FP8_KV_CACHE
-    bool is_fp8_quantized_supported = is_fp8 &&
+    bool is_fp8_quantized_supported = is_fp8 && is_global_attention &&
                                       (k_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                        v_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                        data.k_scale == data.v_scale &&

onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu

@@ -698,6 +698,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,8 @@ 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,
+    [[maybe_unused]] const int local_window_size = -1) {
 #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 +93,23 @@ inline Status Launch(
     cudaMemsetAsync(semaphores, 0, semaphore_size, stream);
   }
 
+#if SLIDING_WINDOW
+  // 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,
@@ -119,6 +122,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,
@@ -136,15 +140,15 @@ Status LaunchXQAKernel<__nv_bfloat16>(
   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_impl.cuh

@@ -20,6 +20,11 @@
 #define TOKENS_PER_PAGE 0
 #define INPUT_FP16 0  // Set to 0 for BFloat16
 #define ALLOW_MULTI_BLOCK_MODE 1
+// Compile the non-quantized bf16 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 (GPT-OSS / Mistral /
+// Gemma2). Attention sinks (head_sink) already work and compose with the window in-kernel.
+#define SLIDING_WINDOW 1
 
 #pragma nv_diag_suppress 177
 #pragma nv_diag_suppress 20012
@@ -164,6 +169,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,
@@ -186,6 +192,7 @@ Status 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,
@@ -218,19 +225,19 @@ Status LaunchXQAKernelImpl<__nv_bfloat16>(
   int group_size = num_heads / kv_num_heads;
   switch (group_size) {
     case 1:
-      return grp1_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp1_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 2:
-      return grp2_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp2_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 4:
-      return grp4_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp4_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 5:
-      return grp5_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp5_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 8:
-      return grp8_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp8_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 16:
-      return grp16_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp16_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     case 32:
-      return grp32_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size);
+      return grp32_bf16::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, attention_sinks, kv_cache_scale, workspace, workspace_size, local_window_size);
     default:
       return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA supports group_size 1, 2, 4, 5, 8, 16, 32. Input has ", group_size);
   }