From 6ebf8638736b48cf3edbc181d28f4b58ccc7d8a1 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Fri, 19 Jun 2026 18:49:00 -0700
Subject: [PATCH 01/12] update doc for cuda contrib op GQA

---
 docs/contrib_ops/cuda/gqa.md | 448 +++++++++++++++++++++++++++++++++++
 docs/contrib_ops/gqa.md      | 173 --------------
 2 files changed, 448 insertions(+), 173 deletions(-)
 create mode 100644 docs/contrib_ops/cuda/gqa.md
 delete mode 100644 docs/contrib_ops/gqa.md

diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md
new file mode 100644
index 0000000000000..9c7c1a5313859
--- /dev/null
+++ b/docs/contrib_ops/cuda/gqa.md
@@ -0,0 +1,448 @@
+# GroupQueryAttention — Operator Documentation
+
+This document describes the `com.microsoft::GroupQueryAttention` (GQA) contrib operator: its schema,
+the CUDA kernel backends and how one is selected, and the attention-sink (`head_sink`) decode path
+that is accelerated by the XQA kernel.
+
+For CPU-specific implementation details (including the quantized KV-cache flash path), see
+[cpu/gqa.md](cpu/gqa.md).
+
+---
+
+## Table of Contents
+
+1. [Overview](#1-overview)
+2. [Operator Schema](#2-operator-schema)
+3. [Input Formats](#3-input-formats)
+4. [KV Cache and Quantization](#4-kv-cache-and-quantization)
+5. [Attention Sink (`head_sink`) and Smooth Softmax](#5-attention-sink-head_sink-and-smooth-softmax)
+6. [CUDA Kernel Backends and Dispatch](#6-cuda-kernel-backends-and-dispatch)
+7. [XQA Decode Path](#7-xqa-decode-path)
+8. [XQA `head_sink` PrePack](#8-xqa-head_sink-prepack)
+9. [Selecting a Kernel: Provider Option and Environment Variables](#9-selecting-a-kernel-provider-option-and-environment-variables)
+10. [Profiling and Benchmarking](#10-profiling-and-benchmarking)
+11. [Fast Build Options](#11-fast-build-options)
+12. [Testing](#12-testing)
+13. [Future Work and Known Limitations](#13-future-work-and-known-limitations)
+
+---
+
+## 1. Overview
+
+GroupQueryAttention implements causal grouped-query attention with KV-cache (past/present) support.
+Grouped-query attention uses fewer key/value heads than query heads: each KV head is shared by a
+group of `num_heads / kv_num_heads` query heads. The operator also supports:
+
+- Rotary positional embeddings (RoPE)
+- Past/present KV cache with optional in-place (shared) buffer
+- Quantized KV cache (int4 / int8 / float8e4m3fn) to reduce memory footprint
+- Optional attention bias and local (sliding) window attention
+- Smooth softmax, including a per-head attention sink (`head_sink`)
+
+The operator schema is defined in
+[onnxruntime/core/graph/contrib_ops/bert_defs.cc](../../onnxruntime/core/graph/contrib_ops/bert_defs.cc).
+The CUDA kernel is implemented in
+[onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc](../../onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc)
+and [group_query_attention_impl.cu](../../onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu).
+
+## 2. Operator Schema
+
+Selected attributes:
+
+| Attribute | Description |
+|-----------|-------------|
+| `num_heads` | Number of query heads. |
+| `kv_num_heads` | Number of key/value heads. `num_heads % kv_num_heads == 0`. |
+| `scale` | Softmax scale. Defaults to `1/sqrt(head_size)`. |
+| `softcap` | Optional logit soft-capping value. `0` disables it. |
+| `local_window_size` | Left window size for local attention. `-1` means global attention. |
+| `do_rotary` / `rotary_interleaved` | Enable RoPE and select interleaved vs. half-rotary layout. |
+| `smooth_softmax` | Add a smooth factor to the softmax denominator. |
+| `k_quant_type` / `v_quant_type` | KV cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
+| `kv_cache_bit_width` | Bit width of the quantized KV cache (`8` or `4`). |
+
+Selected inputs (see the schema for the full list and shapes):
+
+| Index | Name | Notes |
+|-------|------|-------|
+| 0 | `query` | `(batch, seq, hidden)`, or packed QKV. |
+| 1, 2 | `key`, `value` | Optional when QKV is packed into `query`. |
+| 3, 4 | `past_key`, `past_value` | BNSH cache. Shares the buffer with `present_*` when in-place. |
+| 5 | `seqlens_k` | `total_sequence_lengths - 1` per batch entry. |
+| 6 | `total_sequence_length` | Scalar used to distinguish prompt vs. decode. |
+| 7, 8 | `cos_cache`, `sin_cache` | RoPE caches. |
+| 11 | `head_sink` | `(num_heads,)` per-head attention sink (see §5). |
+| 12, 13 | `k_scale`, `v_scale` | FP32 dequant scales for the quantized KV cache. |
+
+Outputs are `output`, `present_key`, `present_value`, and optional `output_qk`.
+
+## 3. Input Formats
+
+GQA accepts query/key/value in two layouts. The layout is inferred from whether `key` (input 1)
+is present.
+
+### Unpacked Q, K, V (`Q_K_V_BSNH`)
+
+`key` and `value` are both provided:
+
+| Tensor | Shape |
+|--------|-------|
+| `query` | `(batch_size, sequence_length, num_heads * head_size)` |
+| `key`   | `(batch_size, sequence_length, kv_num_heads * head_size)` |
+| `value` | `(batch_size, sequence_length, kv_num_heads * head_size)` |
+
+### Packed QKV (`QKV_BS3NH`)
+
+`key` and `value` are omitted (null) and Q, K, V are concatenated along the last dimension of
+`query`:
+
+| Tensor | Shape |
+|--------|-------|
+| `query` | `(batch_size, sequence_length, (num_heads + 2 * kv_num_heads) * head_size)` |
+
+`head_size` is derived as `hidden_size / (num_heads + 2 * kv_num_heads)`.
+
+### KV cache layout
+
+`past_key` / `past_value` / `present_key` / `present_value` always use BNSH:
+`(batch_size, kv_num_heads, cache_sequence_length, head_size)`. For a 4-bit quantized cache the
+last dimension is `(head_size + 1) / 2` because two nibbles are packed per byte.
+
+### Constraints
+
+- `num_heads % kv_num_heads == 0` (each KV head is shared by `num_heads / kv_num_heads` query heads).
+- `head_size == v_head_size` (Q and V share the head size).
+- Q and K/V must have the same `sequence_length` (cross-attention is not supported). The exception
+  is the shared-buffer decode case where `kv_sequence_length == 0` (no new K/V to append — the past
+  buffer already holds the full KV cache).
+- RoPE, packed-QKV unpacking, and KV-head expansion are handled internally (`PrepareQKV`) before the
+  selected backend runs, so every backend sees a consistent layout.
+
+## 4. KV Cache and Quantization
+
+### Layout and shared buffer
+
+The past/present KV cache uses BNSH layout
+`(batch_size, kv_num_heads, cache_sequence_length, head_size)`. When `past_present_share_buffer`
+holds (the past and present tensors alias the same memory), the cache length is the maximum
+sequence length and new keys/values are appended in place. This shared-buffer mode is required by
+the XQA decode path and by the Flash-Decoding fast path.
+
+### Quantized KV cache
+
+To reduce the KV-cache memory footprint, the cache may be stored quantized while `query` stays
+FP16/BF16. Quantization is **symmetric** and configured by three attributes:
+
+| Attribute | Values |
+|-----------|--------|
+| `k_quant_type` / `v_quant_type` | `NONE`, `PER_TENSOR`, `PER_CHANNEL` |
+| `kv_cache_bit_width` | `8` (INT8 / FP8) or `4` (INT4) |
+
+Supported storage types (`T_CACHE`) and their formula:
+
+| Type | Range | Quantize |
+|------|-------|----------|
+| INT8 | `[-128, 127]` | `q = clamp(round(x / scale), -128, 127)` |
+| INT4 | `[-8, 7]`, two nibbles packed per byte | `q = clamp(round(x / scale), -8, 7)` |
+| FP8 E4M3 | `[-448, 448]` | `q = clamp(x / scale, -448, 448)` (SM89+/Ada or SM90+) |
+
+- `k_scale` / `v_scale` (inputs 12, 13) are **always FP32**. For `PER_TENSOR` they are scalars; for
+  `PER_CHANNEL` they have shape `(kv_num_heads, 1, head_size)`.
+- New keys/values are quantized as they are appended to the present cache; the attention kernel
+  dequantizes on the fly while computing scores.
+- Registered type combinations are `T ∈ {float16, bfloat16}` × `T_CACHE ∈ {same as T, int8, FP8E4M3, uint8 (int4)}`.
+
+### How quantized decode is served
+
+The quantized KV-cache path is handled by the **XQA** decode kernel (see §7). XQA requires
+`PER_TENSOR` scaling with `k_scale` and `v_scale` pointing to the **same** FP32 tensor,
+`head_size ∈ {64, 128, 256}`, and a query/KV group size in `{4, 8, 16, 32}`. FP8 additionally
+requires SM89+ (Ada) or SM90+.
+
+INT8 cache kernels are always built; FP8 (`onnxruntime_USE_FP8_KV_CACHE`, default ON) and INT4
+(`onnxruntime_USE_INT4_KV_CACHE`, default OFF) are gated by build options (see §11).
+
+## 5. Attention Sink (`head_sink`) and Smooth Softmax
+
+An attention sink adds a learned per-head bias term to the softmax denominator. With sink value `s_h`
+for head `h`, the attention weights over `T` cached positions become:
+
+$$
+\text{softmax}_i = \frac{e^{x_i - m}}{e^{s_h - m} + \sum_{j} e^{x_j - m}}, \quad m = \max\left(s_h, \max_j x_j\right)
+$$
+
+This is equivalent to appending a single extra logit `s_h` (whose value contributes nothing to the
+output, only to normalization). GPT-OSS style models use this to let a head attend to "nothing".
+
+In the kernel, providing the `head_sink` input is treated as smooth softmax:
+`parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr`. The `head_sink` tensor is
+1D of shape `(num_heads,)` and matches the operator's floating-point type (`float16` or `bfloat16` on
+the XQA path).
+
+## 6. CUDA Kernel Backends and Dispatch
+
+The CUDA EP can route a GQA node to one of five backends. They are evaluated in a fixed priority
+order and the first eligible backend wins:
+
+**XQA → cuDNN SDPA → Flash Attention → Memory Efficient Attention (MEA) → Unfused**
+
+| 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. |
+| 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+). |
+| 5 | **Unfused** | Last-resort fallback (e.g. `head_size > 256`). Any head size, GQA, sliding window, softcap. |
+
+The selected backend is reported in the kernel debug info as `SdpaKernel=...` when debug info is
+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.
+
+### 6.2 cuDNN SDPA
+
+Eligible when **all** of the following hold:
+
+- not already selected for XQA;
+- KV cache is **not** quantized (`T_CACHE == T`);
+- `softcap == 0`, no smooth softmax, and no `head_sink`;
+- no local (sliding) window (`local_window_size == -1`);
+- past/present KV in BNSH (`Q_K_V_BNSH`);
+- cuDNN SDPA is enabled — either explicitly (`ORT_ENABLE_CUDNN_FLASH_ATTENTION=1` or the cuDNN bit of
+  `sdpa_kernel`), or auto-preferred on SM≥90 when no kernel is explicitly pinned;
+- cuDNN ≥ 9.3 (stable) and `is_supported` returns true for the shape.
+
+### 6.3 Flash Attention
+
+Eligible when:
+
+- not XQA and not cuDNN SDPA;
+- FP16/BF16 (`sizeof(T) == 2`) and Flash is enabled (not `ORT_DISABLE_FLASH_ATTENTION`, not disabled
+  via `sdpa_kernel`, and built with `USE_FLASH_ATTENTION`);
+- `flash::is_supported` is true for `head_size` / `num_heads` / `kv_num_heads`.
+
+Flash supports local window, softcap, RoPE, and packed QKV. For decode it additionally uses a
+**Flash-Decoding** split-KV fast path (`seq_len == 1`, shared buffer, non-quantized), unless
+`ORT_DISABLE_FLASH_DECODE=1`.
+
+### 6.4 Memory Efficient Attention (MEA)
+
+Fallback when XQA, cuDNN SDPA, and Flash are all ineligible:
+
+- MEA enabled (not `ORT_DISABLE_MEMORY_EFFICIENT_ATTENTION`, built with `USE_MEMORY_EFFICIENT_ATTENTION`);
+- `has_memory_efficient_attention(sm, is_fp16, is_bf16, head_size)` is true — FP16/FP32 broadly,
+  BF16 on SM80+.
+
+When the query/KV head counts differ, the KV heads are expanded to `num_heads` into a scratch buffer.
+
+### 6.5 Unfused
+
+Last-resort path, activated when XQA / cuDNN / Flash / MEA are all ineligible **and**:
+
+- KV cache is not quantized;
+- no smooth softmax and no `head_sink`;
+- past/present KV in BNSH.
+
+It supports any `head_size` (FP32 QK accumulation), GQA, sliding window, and softcap — for example
+`head_size > 256` with past KV. The unfused (math) path can never be turned off and is always
+available as a fallback.
+
+## 7. XQA Decode Path
+
+XQA (a highly optimized cross/decode attention kernel) is used only when **all** of the following hold:
+
+1. Compute capability SM 8.0+ (Ampere or newer).
+2. Decoding phase (not the first prompt) with `sequence_length == 1`.
+3. `kv_sequence_length > 0` (there is a new K/V to append).
+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.
+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`.
+
+## 8. XQA `head_sink` PrePack
+
+XQA consumes the attention sink as an FP32 buffer, while the model stores `head_sink` as FP16/BF16. To
+avoid converting on every decode step, `GroupQueryAttention::PrePack` converts a **constant-initializer**
+`head_sink` once into a cached FP32 device buffer (`xqa_head_sink_`):
+
+- The cached buffer is reused for every launch when XQA is eligible.
+- A dynamic / non-initializer `head_sink` is **not** prepacked; the kernel instead reserves a small FP32
+  scratch buffer and converts the sink per launch (`xqa_head_sink_needs_conversion = true`).
+- `PrePack` keeps `is_packed = false` so the original FP16/BF16 `head_sink` is still delivered to the
+  Flash/fallback paths when XQA is disabled or ineligible.
+
+## 9. Selecting a Kernel: Provider Option and Environment Variables
+
+### `sdpa_kernel` provider option
+
+The CUDA EP exposes a `sdpa_kernel` provider option (a bitmask defined by `AttentionBackend`) that
+pins which fused attention backends are allowed. It applies to GroupQueryAttention,
+MultiHeadAttention, and Attention nodes.
+
+| Bit value | Backend |
+|-----------|---------|
+| `0` | Default — selection follows heuristics / environment variables (auto-prefers cuDNN SDPA on SM≥90). |
+| `1` | Flash Attention |
+| `2` | Memory Efficient Attention |
+| `8` | cuDNN SDPA |
+| `16` | Unfused (math) — note the unfused fallback can never actually be turned off |
+
+Bits can be OR-ed together. Any positive value is treated as an **explicit** selection: only the
+listed backends are enabled and the automatic cuDNN-on-SM≥90 preference is disabled. **XQA is not
+part of this bitmask** — it is controlled separately by `ORT_ENABLE_XQA`.
+
+```python
+import onnxruntime as ort
+
+sess = ort.InferenceSession(
+    "model.onnx",
+    providers=[("CUDAExecutionProvider", {"sdpa_kernel": "1"})],  # 1 = Flash Attention only
+)
+```
+
+### Environment variables
+
+| 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_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. |
+| `ORT_ENABLE_CUDNN_FLASH_ATTENTION` | `1` enables cuDNN SDPA; `0` disables it and also disables the SM≥90 auto-preference. |
+| `ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO` | `1` prints the selected backend (`SdpaKernel=...`) per node (see §10). |
+
+A positive `sdpa_kernel` value takes precedence over these environment defaults. Environment
+variables are read once when the kernel is constructed.
+
+## 10. Profiling and Benchmarking
+
+### Verify which backend ran
+
+Set `ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO=1`. For each GQA node the kernel prints a line such as:
+
+```
+Operator=GroupQueryAttention Node=<name> DataType=fp16 SdpaKernel=XQA
+```
+
+`SdpaKernel` is one of `XQA`, `FLASH_ATTENTION`, `EFFICIENT_ATTENTION`, `CUDNN_FLASH_ATTENTION`, or
+`MATH` (unfused). Use this to confirm that an env var / `sdpa_kernel` choice took effect.
+
+### Benchmark and profiling scripts
+
+Located in `onnxruntime/test/python/transformers/`:
+
+| Script | Purpose |
+|--------|---------|
+| [profile_gqa.py](../../onnxruntime/test/python/transformers/profile_gqa.py) | Profile GQA (incl. quantized KV cache) with NVTX markers; examples for Nsight Compute (`ncu`) and Nsight Systems (`nsys`). |
+| [benchmark_gqa.py](../../onnxruntime/test/python/transformers/benchmark_gqa.py) | Triton-based throughput comparison across dense / local / packed-QKV and INT4/INT8/FP8 variants. |
+| [benchmark_gqa_windows.py](../../onnxruntime/test/python/transformers/benchmark_gqa_windows.py) | GQA benchmark variant for Windows. |
+| [benchmark_gqa_cpu_flash.py](../../onnxruntime/test/python/transformers/benchmark_gqa_cpu_flash.py) | CPU flash-vs-naive GQA benchmark. |
+
+Example kernel-level and timeline profiling:
+
+```bash
+cd onnxruntime/test/python/transformers
+
+# Kernel-level analysis with Nsight Compute
+ncu --set full -o gqa_int8 python profile_gqa.py --mode int8 --warmup 5 --repeat 1
+
+# Timeline with Nsight Systems, then parse kernel timings
+nsys profile -o gqa_int8 --export=sqlite python profile_gqa.py --mode int8 --warmup 5 --repeat 10
+python parse_nsys.py gqa_int8.sqlite
+```
+
+ONNX Runtime's built-in profiler (`SessionOptions.enable_profiling = True`) also emits a JSON
+timeline with per-node durations.
+
+## 11. Fast Build Options
+
+These CMake options speed up CUDA builds during development. Pass them through
+`--cmake_extra_defines` (see the `ort-build` skill).
+
+| Option | Default | Effect |
+|--------|---------|--------|
+| `onnxruntime_QUICK_BUILD` | `OFF` | Builds only the `hdim128` FP16/BF16 Flash Attention kernels. Greatly reduces compile time, but **changes dispatch**: shapes with `head_size != 128` fall back to Memory Efficient Attention because Flash is no longer compiled for them. Do not use it to characterize Flash-vs-arch behavior. |
+| `onnxruntime_USE_FP8_KV_CACHE` | `ON` | Builds the FP8 (E4M3) quantized KV-cache kernels (`-DUSE_FP8_KV_CACHE=1`). |
+| `onnxruntime_USE_INT4_KV_CACHE` | `OFF` | Builds the INT4 quantized KV-cache kernels (`-DUSE_INT4_KV_CACHE=1`). A `kv_cache_bit_width == 4` node errors out if this is off. |
+
+Other ways to shorten the iteration loop:
+
+- Restrict GPU architectures with `CMAKE_CUDA_ARCHITECTURES` (e.g.
+  `--cmake_extra_defines CMAKE_CUDA_ARCHITECTURES=80`) so kernels are not compiled for unused SMs.
+- Build only the CUDA provider target:
+  `./build.sh --config Release --build --parallel --target onnxruntime_providers_cuda`.
+- Skip `--update` when you only edited existing `.cc` / `.h` / `.cu` files.
+
+```bash
+./build.sh --config Release --parallel --use_cuda \
+  --cuda_home /usr/local/cuda --cudnn_home /usr/local/cuda \
+  --cmake_extra_defines onnxruntime_QUICK_BUILD=ON onnxruntime_USE_INT4_KV_CACHE=ON
+```
+
+## 12. Testing
+
+CUDA parity tests live in
+[onnxruntime/test/python/transformers/test_gqa.py](../../onnxruntime/test/python/transformers/test_gqa.py):
+
+- `TestXQAQuantizedParity` — XQA per-tensor int8 quantized decode parity.
+- `TestXQAHeadSinkParity` — non-quantized XQA decode parity with a `head_sink` (attention sink) input.
+
+`TestXQAQuantizedParity` sets `ORT_ENABLE_XQA=1` to force the XQA path. `TestXQAHeadSinkParity`
+instead clears `ORT_ENABLE_XQA` to validate that XQA is enabled by default when a `head_sink` input
+is present. Both compare against a PyTorch reference (`attention_ref` with `smooth_softmax_ref`).
+
+## 13. Future Work and Known Limitations
+
+The following features are missing or limited in the CUDA GQA kernel and would broaden coverage of
+popular LLMs. Listed roughly by impact.
+
+### High 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.
+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.
+4. **Attention bias / ALiBi.** `attention_bias` is rejected outright. Needed for ALiBi-style models
+   and additive-mask use cases, though less commonly used in current popular decoder-only LLMs.
+
+### Medium impact
+
+5. **Quantized KV cache coverage.** Quantized decode is XQA-only and narrow: `PER_TENSOR` with
+   `k_scale == v_scale`, `head_size ∈ {64, 128, 256}`, group size `{4, 8, 16, 32}`. Gaps worth
+   filling: `PER_CHANNEL` serving, prompt-phase quantized attention, INT4 enabled by default, and
+   `head_sink` combined with a quantized cache (currently rejected).
+6. **Paged KV cache / continuous batching.** GQA uses a contiguous shared buffer; there is a
+   separate `PagedAttention` op, but GQA itself has no paged-cache path. Paged KV is what
+   high-throughput serving (vLLM-style) needs.
+7. **MLA (Multi-head Latent Attention).** **DeepSeek-V2/V3** use latent KV compression with a
+   `v_head_size` that differs from `head_size`; GQA assumes `head_size == v_head_size`. This needs a
+   distinct kernel/op rather than a GQA tweak.
+
+### Lower impact / niche
+
+8. **Returning attention weights (`output_qk`).** Never supported by the CUDA fused kernels. Only
+   relevant for interpretability or speculative-decode scoring.
+9. **Cross-attention (different Q vs KV sequence lengths).** Rejected by the input checker.
+   Encoder-decoder / multimodal cross-attention is not covered by GQA.
diff --git a/docs/contrib_ops/gqa.md b/docs/contrib_ops/gqa.md
deleted file mode 100644
index 08596ff4b5dd9..0000000000000
--- a/docs/contrib_ops/gqa.md
+++ /dev/null
@@ -1,173 +0,0 @@
-# GroupQueryAttention — Operator Documentation
-
-This document describes the `com.microsoft::GroupQueryAttention` (GQA) contrib operator: its schema,
-the CUDA kernel backends and how one is selected, and the attention-sink (`head_sink`) decode path
-that is accelerated by the XQA kernel.
-
-For CPU-specific implementation details (including the quantized KV-cache flash path), see
-[cpu/gqa.md](cpu/gqa.md).
-
----
-
-## Table of Contents
-
-1. [Overview](#1-overview)
-2. [Operator Schema](#2-operator-schema)
-3. [KV Cache and Quantization](#3-kv-cache-and-quantization)
-4. [Attention Sink (`head_sink`) and Smooth Softmax](#4-attention-sink-head_sink-and-smooth-softmax)
-5. [CUDA Kernel Backends and Dispatch](#5-cuda-kernel-backends-and-dispatch)
-6. [XQA Decode Path](#6-xqa-decode-path)
-7. [XQA `head_sink` PrePack](#7-xqa-head_sink-prepack)
-8. [Environment Variables](#8-environment-variables)
-9. [Testing](#9-testing)
-
----
-
-## 1. Overview
-
-GroupQueryAttention implements causal grouped-query attention with KV-cache (past/present) support.
-Grouped-query attention uses fewer key/value heads than query heads: each KV head is shared by a
-group of `num_heads / kv_num_heads` query heads. The operator also supports:
-
-- Rotary positional embeddings (RoPE)
-- Past/present KV cache with optional in-place (shared) buffer
-- Quantized KV cache (int4 / int8 / float8e4m3fn) to reduce memory footprint
-- Optional attention bias and local (sliding) window attention
-- Smooth softmax, including a per-head attention sink (`head_sink`)
-
-The operator schema is defined in
-[onnxruntime/core/graph/contrib_ops/bert_defs.cc](../../onnxruntime/core/graph/contrib_ops/bert_defs.cc).
-The CUDA kernel is implemented in
-[onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc](../../onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc)
-and [group_query_attention_impl.cu](../../onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu).
-
-## 2. Operator Schema
-
-Selected attributes:
-
-| Attribute | Description |
-|-----------|-------------|
-| `num_heads` | Number of query heads. |
-| `kv_num_heads` | Number of key/value heads. `num_heads % kv_num_heads == 0`. |
-| `scale` | Softmax scale. Defaults to `1/sqrt(head_size)`. |
-| `softcap` | Optional logit soft-capping value. `0` disables it. |
-| `local_window_size` | Left window size for local attention. `-1` means global attention. |
-| `do_rotary` / `rotary_interleaved` | Enable RoPE and select interleaved vs. half-rotary layout. |
-| `smooth_softmax` | Add a smooth factor to the softmax denominator. |
-| `k_quant_type` / `v_quant_type` | KV cache quantization mode: `NONE`, `PER_TENSOR`, or `PER_CHANNEL`. |
-| `kv_cache_bit_width` | Bit width of the quantized KV cache (`8` or `4`). |
-
-Selected inputs (see the schema for the full list and shapes):
-
-| Index | Name | Notes |
-|-------|------|-------|
-| 0 | `query` | `(batch, seq, hidden)`, or packed QKV. |
-| 1, 2 | `key`, `value` | Optional when QKV is packed into `query`. |
-| 3, 4 | `past_key`, `past_value` | BNSH cache. Shares the buffer with `present_*` when in-place. |
-| 5 | `seqlens_k` | `total_sequence_lengths - 1` per batch entry. |
-| 6 | `total_sequence_length` | Scalar used to distinguish prompt vs. decode. |
-| 7, 8 | `cos_cache`, `sin_cache` | RoPE caches. |
-| 11 | `head_sink` | `(num_heads,)` per-head attention sink (see §4). |
-| 12, 13 | `k_scale`, `v_scale` | FP32 dequant scales for the quantized KV cache. |
-
-Outputs are `output`, `present_key`, `present_value`, and optional `output_qk`.
-
-## 3. KV Cache and Quantization
-
-The past/present KV cache uses BNSH layout `(batch_size, kv_num_heads, cache_sequence_length, head_size)`.
-When `past_present_share_buffer` holds (the past and present tensors alias the same memory), the cache
-length is the maximum sequence length and new keys/values are appended in place. This shared-buffer mode
-is required by the XQA decode path.
-
-When quantization is enabled, `k_quant_type` and `v_quant_type` select `PER_TENSOR` or `PER_CHANNEL`
-scaling, and `kv_cache_bit_width` selects 8-bit or 4-bit storage. The `k_scale` / `v_scale` inputs are
-always FP32.
-
-## 4. Attention Sink (`head_sink`) and Smooth Softmax
-
-An attention sink adds a learned per-head bias term to the softmax denominator. With sink value `s_h`
-for head `h`, the attention weights over `T` cached positions become:
-
-$$
-\text{softmax}_i = \frac{e^{x_i - m}}{e^{s_h - m} + \sum_{j} e^{x_j - m}}, \quad m = \max\left(s_h, \max_j x_j\right)
-$$
-
-This is equivalent to appending a single extra logit `s_h` (whose value contributes nothing to the
-output, only to normalization). GPT-OSS style models use this to let a head attend to "nothing".
-
-In the kernel, providing the `head_sink` input is treated as smooth softmax:
-`parameters.use_smooth_softmax = use_smooth_softmax_ || head_sink != nullptr`. The `head_sink` tensor is
-1D of shape `(num_heads,)` and matches the operator's floating-point type (`float16` or `bfloat16` on
-the XQA path).
-
-## 5. CUDA Kernel Backends and Dispatch
-
-The CUDA EP can route a GQA node to several backends. At runtime it selects the first eligible one:
-
-| Backend | Typical use |
-|---------|-------------|
-| **XQA** | Single-token global decode (`seq_len == 1`), shared KV buffer. Fastest decode path. |
-| **Flash Attention / Flash Decoding** | General prompt and decode, including local window and softcap. |
-| **cuDNN SDPA** | Preferred on SM≥90 for non-quantized FP16/BF16 causal attention. |
-| **Memory Efficient Attention** | Fallback for FP16/FP32 (and BF16 on SM80+). |
-| **Unfused** | Last-resort fallback (e.g. `head_size > 256` with past KV). |
-
-The selected backend is reported in the kernel debug info as `SdpaKernel=...` when debug info is enabled.
-
-## 6. XQA Decode Path
-
-XQA (a highly optimized cross/decode attention kernel) is used only when **all** of the following hold:
-
-1. Compute capability SM 8.0+ (Ampere or newer).
-2. Decoding phase (not the first prompt) with `sequence_length == 1`.
-3. `kv_sequence_length > 0` (there is a new K/V to append).
-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.
-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`.
-
-## 7. XQA `head_sink` PrePack
-
-XQA consumes the attention sink as an FP32 buffer, while the model stores `head_sink` as FP16/BF16. To
-avoid converting on every decode step, `GroupQueryAttention::PrePack` converts a **constant-initializer**
-`head_sink` once into a cached FP32 device buffer (`xqa_head_sink_`):
-
-- The cached buffer is reused for every launch when XQA is eligible.
-- A dynamic / non-initializer `head_sink` is **not** prepacked; the kernel instead reserves a small FP32
-  scratch buffer and converts the sink per launch (`xqa_head_sink_needs_conversion = true`).
-- `PrePack` keeps `is_packed = false` so the original FP16/BF16 `head_sink` is still delivered to the
-  Flash/fallback paths when XQA is disabled or ineligible.
-
-## 8. Environment Variables
-
-| Variable | Effect |
-|----------|--------|
-| `ORT_ENABLE_XQA` | `1` enables the XQA decode path for the non-quantized KV cache (default off; default on for quantized). |
-| `ORT_DISABLE_FLASH_DECODE` | `1` disables the Flash Decoding split-KV optimization. |
-
-These are read once when the kernel is constructed.
-
-## 9. Testing
-
-CUDA parity tests live in
-[onnxruntime/test/python/transformers/test_gqa.py](../../onnxruntime/test/python/transformers/test_gqa.py):
-
-- `TestXQAQuantizedParity` — XQA per-tensor int8 quantized decode parity.
-- `TestXQAHeadSinkParity` — non-quantized XQA decode parity with a `head_sink` (attention sink) input.
-
-`TestXQAQuantizedParity` sets `ORT_ENABLE_XQA=1` to force the XQA path. `TestXQAHeadSinkParity`
-instead clears `ORT_ENABLE_XQA` to validate that XQA is enabled by default when a `head_sink` input
-is present. Both compare against a PyTorch reference (`attention_ref` with `smooth_softmax_ref`).

From 2f5f9fdeda377083c0494300b642fe8c3f7710bb Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Sat, 20 Jun 2026 03:39:51 +0000
Subject: [PATCH 02/12] gqa slide window

---
 .../cuda/bert/group_query_attention.cc        | 14 +++++++++-----
 .../cuda/bert/group_query_attention_impl.cu   |  1 +
 .../cuda/bert/xqa/xqa_impl_gen.cuh            | 17 ++++++++++++++++-
 .../contrib_ops/cuda/bert/xqa/xqa_loader.h    |  1 +
 .../cuda/bert/xqa/xqa_loader_bf16.cu          | 10 +++++++---
 .../cuda/bert/xqa/xqa_loader_bf16_128.cu      |  1 +
 .../cuda/bert/xqa/xqa_loader_bf16_256.cu      |  1 +
 .../cuda/bert/xqa/xqa_loader_bf16_64.cu       |  1 +
 .../cuda/bert/xqa/xqa_loader_bf16_impl.cuh    | 19 +++++++++++++------
 .../cuda/bert/xqa/xqa_loader_fp16.cu          | 11 ++++++++---
 .../cuda/bert/xqa/xqa_loader_fp16_128.cu      |  1 +
 .../cuda/bert/xqa/xqa_loader_fp16_256.cu      |  1 +
 .../cuda/bert/xqa/xqa_loader_fp16_64.cu       |  1 +
 .../cuda/bert/xqa/xqa_loader_fp16_impl.cuh    | 18 ++++++++++++------
 14 files changed, 73 insertions(+), 24 deletions(-)

diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
index 60235024b9118..43988bc0a0077 100644
--- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
@@ -395,7 +395,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 opt-in for the non-quantized path (ORT_ENABLE_XQA), but a head_sink (attention sink) input
@@ -412,11 +413,14 @@ 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.
@@ -424,7 +428,7 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
                                         (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 &&
diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu
index 6a55f18bd939a..0d2d06aa3d966 100644
--- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu
+++ b/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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh
index cd4088cf757ba..0f38224d4704f 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh
+++ b/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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h
index ee4fbc88982f8..562bbab9a0693 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader.h
+++ b/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
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu
index 4a2d22938d48d..980396be37d11 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu
+++ b/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);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_128.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_128.cu
index a8ea76ab23b8b..cd2436f9da699 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_128.cu
+++ b/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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_256.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_256.cu
index 79ddc2d0d7c34..fc2218e4593ac 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_256.cu
+++ b/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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_64.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_64.cu
index c94f6b5fc0695..daf6a3c98f0f3 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_64.cu
+++ b/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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh
index 6a84d452f1384..cd92d13b81049 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh
+++ b/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
@@ -156,6 +161,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,
@@ -178,6 +184,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,
@@ -210,17 +217,17 @@ 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 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, 8, 16, 32. Input has ", group_size);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu
index b8171392e0f50..b744c23a906a8 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16.cu
@@ -26,6 +26,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,
@@ -51,6 +52,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,
@@ -76,6 +78,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,
@@ -102,6 +105,7 @@ Status LaunchXQAKernel(
     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,
@@ -116,15 +120,15 @@ Status LaunchXQAKernel(
   if (head_size == 256) {
     return H256::LaunchXQAKernelImpl<T>(
         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<T>(
         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<T>(
         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);
   }
@@ -188,6 +192,7 @@ template Status LaunchXQAKernel<half>(
     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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_128.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_128.cu
index 06c8b0ce0ea2a..996d0627983c8 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_128.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_128.cu
@@ -24,6 +24,7 @@ template Status HEAD_DIM_NAMESPACE::LaunchXQAKernelImpl<half>(
     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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_256.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_256.cu
index 756cc61cb9720..c48fb2cbb258b 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_256.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_256.cu
@@ -24,6 +24,7 @@ template Status HEAD_DIM_NAMESPACE::LaunchXQAKernelImpl<half>(
     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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_64.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_64.cu
index 4b5b0fe4f17c9..c2cacff5361af 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_64.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_64.cu
@@ -24,6 +24,7 @@ template Status HEAD_DIM_NAMESPACE::LaunchXQAKernelImpl<half>(
     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,
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh
index 269b7956c0999..e411b268a18a3 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh
@@ -20,6 +20,11 @@
 #define TOKENS_PER_PAGE 0
 #define INPUT_FP16 1
 #define ALLOW_MULTI_BLOCK_MODE 1
+// Compile the non-quantized fp16 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
@@ -156,6 +161,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,
@@ -192,17 +198,17 @@ Status LaunchXQAKernelImpl(
   int group_size = num_heads / kv_num_heads;
   switch (group_size) {
     case 1:
-      return grp1_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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_fp16::Launch<T>(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, 8, 16, 32. Input has ", group_size);
   }

From d90d9428bd2e47a355e86106a09cb30c823424ef Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Sat, 20 Jun 2026 04:51:42 +0000
Subject: [PATCH 03/12] update profile

---
 .../test/python/transformers/profile_gqa.sh   | 70 +++++++++++++++++--
 1 file changed, 64 insertions(+), 6 deletions(-)

diff --git a/onnxruntime/test/python/transformers/profile_gqa.sh b/onnxruntime/test/python/transformers/profile_gqa.sh
index 3fe10eaf35ad2..30d22c1d70df5 100644
--- a/onnxruntime/test/python/transformers/profile_gqa.sh
+++ b/onnxruntime/test/python/transformers/profile_gqa.sh
@@ -13,6 +13,15 @@
 #   ./profile_gqa.sh --bf16 --num-heads 64 --kv-num-heads 8
 #   CUDA_VISIBLE_DEVICES=1 PYTHON=python3 ./profile_gqa.sh --int4
 #
+# Compare the new sliding-window XQA decode kernel against the FlashDecode baseline
+# on a GPT-OSS-20B sliding-window layer shape (head_size=64, 64 query / 8 KV heads,
+# head_sink, local_window_size=128). --compare-xqa profiles each mode twice:
+# once with ORT_ENABLE_XQA=0 (FlashDecode) and once with ORT_ENABLE_XQA=1 (XQA):
+#
+#   ./profile_gqa.sh --gpt-oss --compare-xqa
+#   ./profile_gqa.sh --fp16 --compare-xqa --head-sink --head-size 64 \
+#       --num-heads 64 --kv-num-heads 8 --local-window-size 128 --qkv
+#
 
 set -e
 set -o pipefail
@@ -35,7 +44,13 @@ PACKED_QKV=""
 SHARE_KV_SCALE=""
 NUM_HEADS=""
 KV_NUM_HEADS=""
+HEAD_SIZE=""
+HEAD_SINK=""
 LOCAL_WINDOW_SIZE=""
+
+# When true, profile each mode twice: ORT_ENABLE_XQA=0 (FlashDecode) vs
+# ORT_ENABLE_XQA=1 (XQA, incl. the new sliding-window path).
+COMPARE_XQA=false
 while [[ "$#" -gt 0 ]]; do
     case $1 in
         --fp16)
@@ -58,6 +73,22 @@ while [[ "$#" -gt 0 ]]; do
             RUN_BF16=true
             echo "==== 🚀 BF16 run enabled ===="
             ;;
+        --gpt-oss)
+            # GPT-OSS-20B sliding-window attention layer shape (fp16, head_sink,
+            # 64 query / 8 KV heads, head_size 64, local_window_size 128, packed QKV).
+            RUN_FP16=true
+            NUM_HEADS="--num-heads 64"
+            KV_NUM_HEADS="--kv-num-heads 8"
+            HEAD_SIZE="--head-size 64"
+            HEAD_SINK="--head-sink"
+            LOCAL_WINDOW_SIZE="--local-window-size 128"
+            PACKED_QKV="--is-packed-qkv"
+            echo "==== 🚀 GPT-OSS-20B sliding-window layer preset enabled ===="
+            ;;
+        --compare-xqa)
+            COMPARE_XQA=true
+            echo "==== 🔬 XQA on/off comparison enabled ===="
+            ;;
         --all)
             RUN_FP16=true
             RUN_INT8=true
@@ -99,6 +130,15 @@ while [[ "$#" -gt 0 ]]; do
             echo "==== KV Num Heads: $2 ===="
             shift
             ;;
+        --head-size)
+            HEAD_SIZE="--head-size $2"
+            echo "==== Head size: $2 ===="
+            shift
+            ;;
+        --head-sink)
+            HEAD_SINK="--head-sink"
+            echo "==== Head sink enabled ===="
+            ;;
         -w|--local-window-size)
             LOCAL_WINDOW_SIZE="--local-window-size $2"
             echo "==== Local window size: $2 ===="
@@ -113,7 +153,7 @@ while [[ "$#" -gt 0 ]]; do
 done
 
 # Build extra args string
-EXTRA_ARGS="${BATCH_SIZE} ${SEQUENCE_LENGTH} ${PAST_SEQUENCE_LENGTH} ${PACKED_QKV} ${SHARE_KV_SCALE} ${NUM_HEADS} ${KV_NUM_HEADS} ${LOCAL_WINDOW_SIZE}"
+EXTRA_ARGS="${BATCH_SIZE} ${SEQUENCE_LENGTH} ${PAST_SEQUENCE_LENGTH} ${PACKED_QKV} ${SHARE_KV_SCALE} ${NUM_HEADS} ${KV_NUM_HEADS} ${HEAD_SIZE} ${HEAD_SINK} ${LOCAL_WINDOW_SIZE}"
 
 if ! command -v nsys >/dev/null; then
     echo "Error: nsys not found. Install NVIDIA Nsight Systems or add it to PATH."
@@ -157,22 +197,40 @@ profile_one() {
     fi
 }
 
+# run_mode <mode> <tag> <output_base> [env_var=value ...]
+# When --compare-xqa is set, profile the same config twice: ORT_ENABLE_XQA=0 (the
+# FlashDecode baseline) and ORT_ENABLE_XQA=1 (XQA, including the new sliding-window
+# path), so the kernel latency of the new path is shown side by side.
+run_mode() {
+    local mode="$1"
+    local tag="$2"
+    local base="$3"
+    shift 3
+
+    if [ "$COMPARE_XQA" = true ]; then
+        profile_one "${mode}" "${tag}/XQA=0(FlashDecode)" "${base}_xqa0" ORT_ENABLE_XQA=0 "$@"
+        profile_one "${mode}" "${tag}/XQA=1(XQA)" "${base}_xqa1" ORT_ENABLE_XQA=1 "$@"
+    else
+        profile_one "${mode}" "${tag}" "${base}" "$@"
+    fi
+}
+
 if [ "$RUN_FP16" = true ]; then
-    profile_one fp16 Fp16 gqa_fp16
+    run_mode fp16 Fp16 gqa_fp16
 fi
 
 if [ "$RUN_BF16" = true ]; then
-    profile_one bf16 Bf16 gqa_bf16
+    run_mode bf16 Bf16 gqa_bf16
 fi
 
 if [ "$RUN_INT8" = true ]; then
-    profile_one int8 Int8 gqa_int8 ORT_FLASH_ATTENTION_QUERY_DYNAMIC_QUANT=0
+    run_mode int8 Int8 gqa_int8 ORT_FLASH_ATTENTION_QUERY_DYNAMIC_QUANT=0
 fi
 
 if [ "$RUN_INT8_QUANT" = true ]; then
-    profile_one int8 Int8Q gqa_int8_quant ORT_FLASH_ATTENTION_QUERY_DYNAMIC_QUANT=1
+    run_mode int8 Int8Q gqa_int8_quant ORT_FLASH_ATTENTION_QUERY_DYNAMIC_QUANT=1
 fi
 
 if [ "$RUN_INT4" = true ]; then
-    profile_one int4 Int4 gqa_int4
+    run_mode int4 Int4 gqa_int4
 fi

From 8bbfa27ced085a658346bd0137fd17a503f71ab3 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Sat, 20 Jun 2026 05:08:09 +0000
Subject: [PATCH 04/12] update test

---
 .../test/python/transformers/test_gqa.py      | 73 +++++++++++++++++++
 1 file changed, 73 insertions(+)

diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index 529eae1494e94..412e1cdb83d44 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -2481,6 +2481,79 @@ def test_xqa_head_sink_prepack_parity(self, name, config, torch_type, ort_type):
         )
 
 
+def gqa_xqa_sliding_window_test_cases():
+    # Non-quantized sliding-window (local attention) decode through the XQA kernel.
+    #
+    # The XQA decode path now supports local_window_size > 0 on the non-quantized fp16/bf16
+    # path (GPT-OSS style sliding-window layers). XQA selection requires: decode (q_seq=1),
+    # a shared KV buffer, softcap==0, head_size in {64, 128, 256} and 64 % group_size == 0.
+    #
+    # Two window/past relationships are covered:
+    #   past > window  -> the sliding mask drops the oldest keys (the new code path).
+    #   past <= window -> the window spans the whole cache (parity with global attention).
+    # has_head_sink toggles the GPT-OSS attention-sink input, which composes with the window
+    # in-kernel; both with and without a sink are exercised.
+    for torch_type, ort_type in [(torch.float16, TensorProto.FLOAT16), (torch.bfloat16, TensorProto.BFLOAT16)]:
+        for head_size in [64, 128]:
+            for group_size in [4, 8]:
+                for past_kv_sequence_length, local_window_size in [(512, 128), (64, 128)]:
+                    for has_head_sink in [False, True]:
+                        kv_num_heads = 4
+                        num_heads = kv_num_heads * group_size
+                        config = GQAConfig(
+                            batch_size=2,
+                            q_sequence_length=1,
+                            kv_sequence_length=1,
+                            num_heads=num_heads,
+                            kv_num_heads=kv_num_heads,
+                            head_size=head_size,
+                            past_kv_sequence_length=past_kv_sequence_length,
+                            buffer_sequence_length=past_kv_sequence_length + 128,
+                            local_window_size=local_window_size,
+                            rotary=True,
+                            rotary_interleaved=False,
+                            packed=False,
+                            share_buffer=True,
+                            has_head_sink=has_head_sink,
+                        )
+                        type_str = "bf16" if torch_type == torch.bfloat16 else "fp16"
+                        sink_str = "sink" if has_head_sink else "nosink"
+                        win_str = f"past{past_kv_sequence_length}_win{local_window_size}"
+                        name = f"{type_str}_g{group_size}_h{head_size}_{win_str}_{sink_str}"
+                        yield name, config, torch_type, ort_type
+
+
+@unittest.skipIf(not has_xqa(), "XQA is not available, skipping tests.")
+@unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.")
+class TestXQASlidingWindowParity(unittest.TestCase):
+    """Verify the non-quantized XQA sliding-window (local attention) decode path matches the reference."""
+
+    def tearDown(self):
+        """Clear CUDA cache after each test to prevent memory corruption in batch runs."""
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+        gc.collect()
+
+    @parameterized.expand(gqa_xqa_sliding_window_test_cases())
+    def test_xqa_sliding_window_parity(self, name, config, torch_type, ort_type):
+        """Test XQA non-quantized parity with a sliding (local) attention window."""
+        # ORT_ENABLE_XQA=1 forces XQA selection so the sliding-window path is exercised even when
+        # no head_sink input is present (a head_sink input alone would also enable XQA by default).
+        with scoped_env_var("ORT_ENABLE_XQA", "1"):
+            parity_check_gqa_past(
+                config=config,
+                ep="CUDAExecutionProvider",
+                device="cuda",
+                torch_type=torch_type,
+                ort_type=ort_type,
+                causal=True,
+                rtol=rtol["bf16"] if torch_type == torch.bfloat16 else rtol["fp16"],
+                atol=atol["bf16"] if torch_type == torch.bfloat16 else atol["fp16"],
+                std=0.1,
+            )
+
+
 @unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.")
 @unittest.skipIf(not has_quantized_kv_cache(), "Quantized KV Cache is not available, skipping tests.")
 class TestGQARegressions(unittest.TestCase):

From 8101a08256b9fe6c6b37ea85402bc53815e90a03 Mon Sep 17 00:00:00 2001
From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com>
Date: Sat, 20 Jun 2026 18:24:14 +0000
Subject: [PATCH 05/12] docs: fix relative links in CUDA GQA documentation

---
 docs/contrib_ops/cuda/gqa.md | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md
index 9c7c1a5313859..aeb88dbd07588 100644
--- a/docs/contrib_ops/cuda/gqa.md
+++ b/docs/contrib_ops/cuda/gqa.md
@@ -5,7 +5,7 @@ the CUDA kernel backends and how one is selected, and the attention-sink (`head_
 that is accelerated by the XQA kernel.
 
 For CPU-specific implementation details (including the quantized KV-cache flash path), see
-[cpu/gqa.md](cpu/gqa.md).
+[cpu/gqa.md](../cpu/gqa.md).
 
 ---
 
@@ -40,10 +40,10 @@ group of `num_heads / kv_num_heads` query heads. The operator also supports:
 - Smooth softmax, including a per-head attention sink (`head_sink`)
 
 The operator schema is defined in
-[onnxruntime/core/graph/contrib_ops/bert_defs.cc](../../onnxruntime/core/graph/contrib_ops/bert_defs.cc).
+[onnxruntime/core/graph/contrib_ops/bert_defs.cc](../../../onnxruntime/core/graph/contrib_ops/bert_defs.cc).
 The CUDA kernel is implemented in
-[onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc](../../onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc)
-and [group_query_attention_impl.cu](../../onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu).
+[onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc](../../../onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc)
+and [group_query_attention_impl.cu](../../../onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu).
 
 ## 2. Operator Schema
 
@@ -349,10 +349,10 @@ Located in `onnxruntime/test/python/transformers/`:
 
 | Script | Purpose |
 |--------|---------|
-| [profile_gqa.py](../../onnxruntime/test/python/transformers/profile_gqa.py) | Profile GQA (incl. quantized KV cache) with NVTX markers; examples for Nsight Compute (`ncu`) and Nsight Systems (`nsys`). |
-| [benchmark_gqa.py](../../onnxruntime/test/python/transformers/benchmark_gqa.py) | Triton-based throughput comparison across dense / local / packed-QKV and INT4/INT8/FP8 variants. |
-| [benchmark_gqa_windows.py](../../onnxruntime/test/python/transformers/benchmark_gqa_windows.py) | GQA benchmark variant for Windows. |
-| [benchmark_gqa_cpu_flash.py](../../onnxruntime/test/python/transformers/benchmark_gqa_cpu_flash.py) | CPU flash-vs-naive GQA benchmark. |
+| [profile_gqa.py](../../../onnxruntime/test/python/transformers/profile_gqa.py) | Profile GQA (incl. quantized KV cache) with NVTX markers; examples for Nsight Compute (`ncu`) and Nsight Systems (`nsys`). |
+| [benchmark_gqa.py](../../../onnxruntime/test/python/transformers/benchmark_gqa.py) | Triton-based throughput comparison across dense / local / packed-QKV and INT4/INT8/FP8 variants. |
+| [benchmark_gqa_windows.py](../../../onnxruntime/test/python/transformers/benchmark_gqa_windows.py) | GQA benchmark variant for Windows. |
+| [benchmark_gqa_cpu_flash.py](../../../onnxruntime/test/python/transformers/benchmark_gqa_cpu_flash.py) | CPU flash-vs-naive GQA benchmark. |
 
 Example kernel-level and timeline profiling:
 
@@ -398,7 +398,7 @@ Other ways to shorten the iteration loop:
 ## 12. Testing
 
 CUDA parity tests live in
-[onnxruntime/test/python/transformers/test_gqa.py](../../onnxruntime/test/python/transformers/test_gqa.py):
+[onnxruntime/test/python/transformers/test_gqa.py](../../../onnxruntime/test/python/transformers/test_gqa.py):
 
 - `TestXQAQuantizedParity` — XQA per-tensor int8 quantized decode parity.
 - `TestXQAHeadSinkParity` — non-quantized XQA decode parity with a `head_sink` (attention sink) input.

From 231986930f57819c394cd549923cacaa66a473c0 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Sat, 20 Jun 2026 11:49:26 -0700
Subject: [PATCH 06/12] Address PR feedback: reflect non-quantized XQA
 sliding-window in docs/tests

---
 docs/contrib_ops/cuda/gqa.md                  | 20 +++++++++++--------
 .../test/python/transformers/test_gqa.py      |  1 -
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md
index aeb88dbd07588..fcdfaa94a7195 100644
--- a/docs/contrib_ops/cuda/gqa.md
+++ b/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,7 +261,8 @@ 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
@@ -417,10 +419,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.
diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index 412e1cdb83d44..04e76fe944a49 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -2524,7 +2524,6 @@ def gqa_xqa_sliding_window_test_cases():
 
 
 @unittest.skipIf(not has_xqa(), "XQA is not available, skipping tests.")
-@unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.")
 class TestXQASlidingWindowParity(unittest.TestCase):
     """Verify the non-quantized XQA sliding-window (local attention) decode path matches the reference."""
 

From 07e01d29af6eda20edad283b6add959588a84355 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Sat, 20 Jun 2026 13:51:25 -0700
Subject: [PATCH 07/12] address feedbacks: test kernel selection

---
 .../test/python/transformers/profile_gqa.sh   |  9 ++-
 .../test/python/transformers/test_gqa.py      | 60 ++++++++++++++++++-
 2 files changed, 64 insertions(+), 5 deletions(-)

diff --git a/onnxruntime/test/python/transformers/profile_gqa.sh b/onnxruntime/test/python/transformers/profile_gqa.sh
index 30d22c1d70df5..94e395ccd1b3a 100644
--- a/onnxruntime/test/python/transformers/profile_gqa.sh
+++ b/onnxruntime/test/python/transformers/profile_gqa.sh
@@ -208,8 +208,13 @@ run_mode() {
     shift 3
 
     if [ "$COMPARE_XQA" = true ]; then
-        profile_one "${mode}" "${tag}/XQA=0(FlashDecode)" "${base}_xqa0" ORT_ENABLE_XQA=0 "$@"
-        profile_one "${mode}" "${tag}/XQA=1(XQA)" "${base}_xqa1" ORT_ENABLE_XQA=1 "$@"
+        if [[ "${mode}" == "fp16" || "${mode}" == "bf16" ]]; then
+            profile_one "${mode}" "${tag}/XQA=0(FlashDecode)" "${base}_xqa0" ORT_ENABLE_XQA=0 "$@"
+            profile_one "${mode}" "${tag}/XQA=1(XQA)" "${base}_xqa1" ORT_ENABLE_XQA=1 "$@"
+        else
+            echo "---- ${tag}: skipping XQA on/off comparison for quantized KV-cache mode ----"
+            profile_one "${mode}" "${tag}" "${base}" "$@"
+        fi
     else
         profile_one "${mode}" "${tag}" "${base}" "$@"
     fi
diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index 04e76fe944a49..6f03e8bbecd14 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -14,6 +14,9 @@
 import os
 import platform
 import random
+import re
+import sys
+import threading
 import unittest
 from copy import deepcopy
 from dataclasses import dataclass
@@ -70,6 +73,53 @@
 # #################################################################################################
 
 
+class CaptureStdout:
+    def __init__(self):
+        self.fd = 1
+        self.chunk_size = 1024
+        self.output = b""
+
+    def _capture(self):
+        chunks = []
+        while chunk := os.read(self._pipe_reader, self.chunk_size):
+            chunks.append(chunk)
+        self.output = b"".join(chunks)
+
+    def __enter__(self):
+        sys.stdout.flush()
+        self._duped_fd = os.dup(self.fd)
+        self._pipe_reader, pipe_writer = os.pipe()
+        os.dup2(pipe_writer, self.fd)
+        os.close(pipe_writer)
+        self._capture_thread = threading.Thread(target=self._capture)
+        self._capture_thread.start()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        sys.stdout.flush()
+        os.dup2(self._duped_fd, self.fd)
+        self._capture_thread.join()
+        os.close(self._pipe_reader)
+        os.close(self._duped_fd)
+
+
+def get_sdpa_kernel_from_debug_info(run_func):
+    captured_text = None
+    with scoped_env_var("ORT_ENABLE_ATTENTION_KERNEL_DEBUG_INFO", "1"):
+        with CaptureStdout() as captured:
+            run_func()
+        captured_text = captured.output.decode(errors="replace")
+
+    if captured_text is not None:
+        match = re.search(r"SdpaKernel=(?P<kernel>[A-Z_]+)", captured_text)
+        if match is not None:
+            return match.group("kernel")
+
+        print("Failed to get sdpa kernel from debug info:", captured_text)
+
+    return None
+
+
 @dataclass
 class GQAConfig:
     batch_size: int
@@ -2537,9 +2587,8 @@ def tearDown(self):
     @parameterized.expand(gqa_xqa_sliding_window_test_cases())
     def test_xqa_sliding_window_parity(self, name, config, torch_type, ort_type):
         """Test XQA non-quantized parity with a sliding (local) attention window."""
-        # ORT_ENABLE_XQA=1 forces XQA selection so the sliding-window path is exercised even when
-        # no head_sink input is present (a head_sink input alone would also enable XQA by default).
-        with scoped_env_var("ORT_ENABLE_XQA", "1"):
+
+        def run_parity_check():
             parity_check_gqa_past(
                 config=config,
                 ep="CUDAExecutionProvider",
@@ -2552,6 +2601,11 @@ def test_xqa_sliding_window_parity(self, name, config, torch_type, ort_type):
                 std=0.1,
             )
 
+        # ORT_ENABLE_XQA=1 forces XQA selection so the sliding-window path is exercised even when
+        # no head_sink input is present (a head_sink input alone would also enable XQA by default).
+        with scoped_env_var("ORT_ENABLE_XQA", "1"):
+            self.assertEqual("XQA", get_sdpa_kernel_from_debug_info(run_parity_check))
+
 
 @unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.")
 @unittest.skipIf(not has_quantized_kv_cache(), "Quantized KV Cache is not available, skipping tests.")

From a0c51c8442a8b28bbd61eed9d9b6a2fe95bee00c Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Mon, 22 Jun 2026 06:20:35 +0000
Subject: [PATCH 08/12] update doc: XQA on by default

---
 docs/contrib_ops/cuda/gqa.md | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md
index fcdfaa94a7195..0b91332c926d4 100644
--- a/docs/contrib_ops/cuda/gqa.md
+++ b/docs/contrib_ops/cuda/gqa.md
@@ -269,13 +269,7 @@ XQA (a highly optimized cross/decode attention kernel) is used only when **all**
 (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
 
@@ -322,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. |

From 8af769a12072079cd451e608505cf9cbac6e5f46 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Mon, 22 Jun 2026 06:52:19 +0000
Subject: [PATCH 09/12] reject local_window_size 0

---
 .../cuda/bert/group_query_attention.cc        |  2 ++
 .../test/python/transformers/test_gqa.py      | 25 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
index 961d057057be6..b2f9e8cc0dd5e 100644
--- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
+++ b/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);
diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index e255ff7f293ff..f93724c605d5c 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -2612,6 +2612,31 @@ def run_parity_check():
 class TestGQARegressions(unittest.TestCase):
     """Specific regression tests for historical bugs."""
 
+    def test_gqa_cuda_rejects_zero_local_window_size(self):
+        if not has_cuda_provider():
+            self.skipTest("CUDA required")
+
+        config = GQAConfig(
+            batch_size=1,
+            num_heads=4,
+            kv_num_heads=4,
+            head_size=64,
+            q_sequence_length=1,
+            kv_sequence_length=1,
+            past_kv_sequence_length=0,
+            buffer_sequence_length=1,
+            local_window_size=0,
+            share_buffer=True,
+        )
+        onnx_model_str = create_group_query_attention_graph_prompt(config, TensorProto.FLOAT16)
+
+        with self.assertRaisesRegex(Exception, "local_window_size must be -1 or greater than 0"):
+            InferenceSession(
+                onnx_model_str,
+                SessionOptions(),
+                providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")],
+            )
+
     def test_gqa_rope_separate_qkv_bug(self):
         """
         Regression test for separate QKV + RoPE + FlashAttention bug.

From 4807c9f10aab31cbc23a9c219ebfa9f3b0c31de2 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Mon, 22 Jun 2026 20:56:10 -0700
Subject: [PATCH 10/12] enable for quantized kv cache

---
 .../cuda/bert/group_query_attention.cc        |  14 +-
 .../cuda/bert/xqa/xqa_loader_bf16.cu          |   3 +-
 .../bert/xqa/xqa_loader_bf16_fp8_impl.cuh     |  13 +-
 .../cuda/bert/xqa/xqa_loader_bf16_impl.cuh    |   6 +-
 .../cuda/bert/xqa/xqa_loader_bf16_int8.cu     |  10 +-
 .../bert/xqa/xqa_loader_bf16_int8_impl.cuh    |  13 +-
 .../bert/xqa/xqa_loader_fp16_fp8_impl.cuh     |  13 +-
 .../cuda/bert/xqa/xqa_loader_fp16_impl.cuh    |   6 +-
 .../cuda/bert/xqa/xqa_loader_fp16_int8.cu     |  10 +-
 .../bert/xqa/xqa_loader_fp16_int8_impl.cuh    |  13 +-
 .../test/python/transformers/test_gqa.py      | 173 +++++++++++++-----
 11 files changed, 190 insertions(+), 84 deletions(-)

diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
index 32be810579a04..9fe8460991523 100644
--- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
+++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
@@ -416,8 +416,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. 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).
+  // 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
@@ -437,11 +437,9 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
       is_xqa_smooth_softmax_supported) {
     int group_size = parameters.num_heads / parameters.kv_num_heads;
 
-    // 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 &&
+    // 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 &&
                                         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.
@@ -449,7 +447,7 @@ Status GroupQueryAttention<T, U>::ComputeInternal(OpKernelContext* context) cons
                                         (group_size == 4 || group_size == 8 || group_size == 16 || group_size == 32));
 
 #ifdef USE_FP8_KV_CACHE
-    bool is_fp8_quantized_supported = is_fp8 && is_global_attention &&
+    bool is_fp8_quantized_supported = is_fp8 &&
                                       (k_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                        v_quant_type_ == KVQuantizationType::PER_TENSOR &&
                                        data.k_scale == data.v_scale &&
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu
index 980396be37d11..081f46d8733f5 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16.cu
@@ -98,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,
@@ -133,7 +134,7 @@ 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
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_fp8_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_fp8_impl.cuh
index 773b4810b6b30..8d7f1c3282e84 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_fp8_impl.cuh
+++ b/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);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh
index d39120eff84de..0915c173703cd 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_impl.cuh
@@ -123,6 +123,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,
@@ -144,6 +145,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,
@@ -207,7 +209,7 @@ Status LaunchXQAKernelImpl<__nv_bfloat16>(
     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,
+                                   scale, local_window_size, is_bsnh, past_seq_lens, kv_cache_scale, workspace,
                                    workspace_size);
   }
 
@@ -217,7 +219,7 @@ Status LaunchXQAKernelImpl<__nv_bfloat16>(
     ORT_RETURN_IF(attention_sinks != nullptr, "XQA attention sinks are not supported with FP8 KV cache.");
     return LaunchXQAFp8KernelBF16(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,
+                                  scale, local_window_size, is_bsnh, past_seq_lens, kv_cache_scale, workspace,
                                   workspace_size);
   }
 #endif
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8.cu
index cbca83f27cb87..32d83d0b63103 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8.cu
@@ -23,6 +23,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,
@@ -44,6 +45,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,
@@ -65,6 +67,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,
@@ -86,17 +89,18 @@ 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,
     void* workspace,
     size_t workspace_size) {
   if (head_size == 256) {
-    return H256::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 H256::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);
   } else if (head_size == 128) {
-    return H128::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 H128::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);
   } else if (head_size == 64) {
-    return H64::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 H64::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);
   } else {
     return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA INT8 BF16 only supports head_size=64, 128, or 256. Input has ", head_size);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8_impl.cuh
index 0ad18e99c5841..3e48be0232907 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8_impl.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_bf16_int8_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 INT8 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 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,
@@ -102,13 +107,13 @@ Status LaunchXQAInt8KernelBF16(
   int group_size = num_heads / kv_num_heads;
   switch (group_size) {
     case 4:
-      return grp4_bf16_int8::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_int8::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_int8::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_int8::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_int8::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_int8::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_int8::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_int8::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 INT8 only supports group_size 4, 8, 16, 32. Input has ", group_size);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_fp8_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_fp8_impl.cuh
index 0a613ead5c16e..644c7c7cbc247 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_fp8_impl.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_fp8_impl.cuh
@@ -21,6 +21,10 @@
 #define TOKENS_PER_PAGE 0
 #define INPUT_FP16 1  // Q is FP16
 #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
@@ -93,6 +97,7 @@ Status LaunchXQAFp8Kernel(
     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,
@@ -101,13 +106,13 @@ Status LaunchXQAFp8Kernel(
   int group_size = num_heads / kv_num_heads;
   switch (group_size) {
     case 4:
-      return grp4_fp8::Launch<half>(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_fp8::Launch<half>(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_fp8::Launch<half>(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_fp8::Launch<half>(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_fp8::Launch<half>(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_fp8::Launch<half>(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_fp8::Launch<half>(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_fp8::Launch<half>(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);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh
index 6955261fc983a..e861fb83ef098 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_impl.cuh
@@ -123,6 +123,7 @@ Status LaunchXQAInt8Kernel(
     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,
@@ -144,6 +145,7 @@ Status LaunchXQAFp8Kernel(
     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,
@@ -183,7 +185,7 @@ Status LaunchXQAKernelImpl(
   if (kv_quant_type == XqaQuantType::kInt8) {
     ORT_RETURN_IF(attention_sinks != nullptr, "XQA attention sinks are not supported with INT8 KV cache.");
     if constexpr (std::is_same<T, half>::value) {
-      return LaunchXQAInt8Kernel(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 LaunchXQAInt8Kernel(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);
     } else {
       // BF16 case is handled in xqa_loader_bf16.cu via specialization
       return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA INT8 path mismatch.");
@@ -195,7 +197,7 @@ Status LaunchXQAKernelImpl(
   if (kv_quant_type == XqaQuantType::kFp8) {
     ORT_RETURN_IF(attention_sinks != nullptr, "XQA attention sinks are not supported with FP8 KV cache.");
     if constexpr (std::is_same<T, half>::value) {
-      return LaunchXQAFp8Kernel(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 LaunchXQAFp8Kernel(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);
     } else {
       // BF16 case is handled in xqa_loader_bf16.cu via specialization
       return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA FP8 path mismatch.");
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8.cu b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8.cu
index 4855571c32a57..78b55b7ded8a6 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8.cu
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8.cu
@@ -23,6 +23,7 @@ Status LaunchXQAInt8Kernel(
     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,
@@ -45,6 +46,7 @@ Status LaunchXQAInt8Kernel(
     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,
@@ -67,6 +69,7 @@ Status LaunchXQAInt8Kernel(
     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,
@@ -89,17 +92,18 @@ Status LaunchXQAInt8Kernel(
     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,
     void* workspace,
     size_t workspace_size) {
   if (head_size == 256) {
-    return H256::LaunchXQAInt8Kernel(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 H256::LaunchXQAInt8Kernel(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);
   } else if (head_size == 128) {
-    return H128::LaunchXQAInt8Kernel(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 H128::LaunchXQAInt8Kernel(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);
   } else if (head_size == 64) {
-    return H64::LaunchXQAInt8Kernel(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 H64::LaunchXQAInt8Kernel(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);
   } else {
     return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "XQA INT8 only supports head_size=64, 128, or 256. Input has ", head_size);
   }
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8_impl.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8_impl.cuh
index ebeccfb60c7ba..0a975a9942466 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8_impl.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_loader_fp16_int8_impl.cuh
@@ -21,6 +21,10 @@
 #define TOKENS_PER_PAGE 0
 #define INPUT_FP16 1  // Q is FP16
 #define ALLOW_MULTI_BLOCK_MODE 1
+// Compile the INT8 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
@@ -93,6 +97,7 @@ Status LaunchXQAInt8Kernel(
     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,
@@ -101,13 +106,13 @@ Status LaunchXQAInt8Kernel(
   int group_size = num_heads / kv_num_heads;
   switch (group_size) {
     case 4:
-      return grp4_int8::Launch<half>(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_int8::Launch<half>(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_int8::Launch<half>(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_int8::Launch<half>(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_int8::Launch<half>(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_int8::Launch<half>(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_int8::Launch<half>(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_int8::Launch<half>(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 INT8 only supports group_size 4, 8, 16, 32. Input has ", group_size);
   }
diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index e975818002325..233f7ac102b92 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -2327,17 +2327,16 @@ def test_gqa_qk_norm_past_xqa(self):
             has_qk_norm=True,
         )
 
-        with scoped_env_var("ORT_ENABLE_XQA", "1"):
-            parity_check_gqa_past(
-                config=config,
-                ep="CUDAExecutionProvider",
-                device="cuda",
-                torch_type=torch.float16,
-                ort_type=TensorProto.FLOAT16,
-                causal=True,
-                rtol=rtol["fp16"],
-                atol=atol["fp16"],
-            )
+        parity_check_gqa_past(
+            config=config,
+            ep="CUDAExecutionProvider",
+            device="cuda",
+            torch_type=torch.float16,
+            ort_type=TensorProto.FLOAT16,
+            causal=True,
+            rtol=rtol["fp16"],
+            atol=atol["fp16"],
+        )
 
     def test_gqa_qk_norm_past_shared_kv(self):
         config = GQAConfig(
@@ -2355,17 +2354,16 @@ def test_gqa_qk_norm_past_shared_kv(self):
             has_qk_norm=True,
         )
 
-        with scoped_env_var("ORT_ENABLE_XQA", "1"):
-            parity_check_gqa_past(
-                config=config,
-                ep="CUDAExecutionProvider",
-                device="cuda",
-                torch_type=torch.float16,
-                ort_type=TensorProto.FLOAT16,
-                causal=True,
-                rtol=rtol["fp16"],
-                atol=atol["fp16"],
-            )
+        parity_check_gqa_past(
+            config=config,
+            ep="CUDAExecutionProvider",
+            device="cuda",
+            torch_type=torch.float16,
+            ort_type=TensorProto.FLOAT16,
+            causal=True,
+            rtol=rtol["fp16"],
+            atol=atol["fp16"],
+        )
 
     def test_gqa_qk_norm_past_xqa_bf16(self):
         if not torch.cuda.is_bf16_supported():
@@ -2388,17 +2386,16 @@ def test_gqa_qk_norm_past_xqa_bf16(self):
         )
         config.kv_cache_type = "bfloat16"
 
-        with scoped_env_var("ORT_ENABLE_XQA", "1"):
-            parity_check_gqa_past(
-                config=config,
-                ep="CUDAExecutionProvider",
-                device="cuda",
-                torch_type=torch.bfloat16,
-                ort_type=TensorProto.BFLOAT16,
-                causal=True,
-                rtol=rtol["bf16"],
-                atol=atol["bf16"],
-            )
+        parity_check_gqa_past(
+            config=config,
+            ep="CUDAExecutionProvider",
+            device="cuda",
+            torch_type=torch.bfloat16,
+            ort_type=TensorProto.BFLOAT16,
+            causal=True,
+            rtol=rtol["bf16"],
+            atol=atol["bf16"],
+        )
 
     @parameterized.expand(gqa_qk_norm_test_cases(is_past=True))
     def test_gqa_qk_norm_past_bf16(self, name, config):
@@ -2671,18 +2668,17 @@ def tearDown(self):
     @parameterized.expand(gqa_xqa_test_cases())
     def test_xqa_quantized_parity(self, name, config, torch_type, ort_type):
         """Test XQA per-tensor INT8 quantized parity."""
-        with scoped_env_var("ORT_ENABLE_XQA", "1"):
-            parity_check_gqa_past(
-                config=config,
-                ep="CUDAExecutionProvider",
-                device="cuda",
-                torch_type=torch_type,
-                ort_type=ort_type,
-                causal=True,
-                rtol=rtol["int8_bf16"] if torch_type == torch.bfloat16 else rtol["int8_fp16"],
-                atol=atol["int8_bf16"] if torch_type == torch.bfloat16 else atol["int8_fp16"],
-                std=0.1,
-            )
+        parity_check_gqa_past(
+            config=config,
+            ep="CUDAExecutionProvider",
+            device="cuda",
+            torch_type=torch_type,
+            ort_type=ort_type,
+            causal=True,
+            rtol=rtol["int8_bf16"] if torch_type == torch.bfloat16 else rtol["int8_fp16"],
+            atol=atol["int8_bf16"] if torch_type == torch.bfloat16 else atol["int8_fp16"],
+            std=0.1,
+        )
 
 
 def gqa_xqa_head_sink_test_cases():
@@ -2869,10 +2865,89 @@ def run_parity_check():
                 std=0.1,
             )
 
-        # ORT_ENABLE_XQA=1 forces XQA selection so the sliding-window path is exercised even when
-        # no head_sink input is present (a head_sink input alone would also enable XQA by default).
-        with scoped_env_var("ORT_ENABLE_XQA", "1"):
-            self.assertEqual("XQA", get_sdpa_kernel_from_debug_info(run_parity_check))
+        # XQA is enabled by default for fp16/bf16, so no head_sink input is required to select it.
+        self.assertEqual("XQA", get_sdpa_kernel_from_debug_info(run_parity_check))
+
+
+def gqa_xqa_quantized_sliding_window_test_cases():
+    # Quantized (INT8 / FP8) sliding-window (local attention) decode through the XQA kernel.
+    #
+    # The XQA decode path now supports local_window_size > 0 on the quantized KV-cache paths as
+    # well. Quantized XQA selection requires: decode (q_seq=1), a shared KV buffer, per-tensor
+    # k/v scales that are the same tensor, head_size in {64, 128, 256} and group_size in
+    # {4, 8, 16, 32}. Attention sinks are not supported with quantized KV cache, so no head_sink.
+    #
+    # Two window/past relationships are covered:
+    #   past > window  -> the sliding mask drops the oldest keys (the new code path).
+    #   past <= window -> the window spans the whole cache (parity with global attention).
+    kv_cache_types = ["int8"]
+    if has_fp8_kv_cache:
+        kv_cache_types.append("fp8")
+    for kv_cache_type in kv_cache_types:
+        for torch_type, ort_type in [(torch.float16, TensorProto.FLOAT16), (torch.bfloat16, TensorProto.BFLOAT16)]:
+            for head_size in [64, 128]:
+                for group_size in [4, 8]:
+                    for past_kv_sequence_length, local_window_size in [(512, 128), (64, 128)]:
+                        kv_num_heads = 4
+                        num_heads = kv_num_heads * group_size
+                        config = GQAConfig(
+                            batch_size=2,
+                            q_sequence_length=1,
+                            kv_sequence_length=1,
+                            num_heads=num_heads,
+                            kv_num_heads=kv_num_heads,
+                            head_size=head_size,
+                            past_kv_sequence_length=past_kv_sequence_length,
+                            buffer_sequence_length=past_kv_sequence_length + 128,
+                            local_window_size=local_window_size,
+                            rotary=True,
+                            rotary_interleaved=False,
+                            packed=False,
+                            share_buffer=True,
+                            k_quant_type="PER_TENSOR",
+                            v_quant_type="PER_TENSOR",
+                            kv_cache_type=kv_cache_type,
+                            share_kv_scale=True,
+                        )
+                        type_str = "bf16" if torch_type == torch.bfloat16 else "fp16"
+                        win_str = f"past{past_kv_sequence_length}_win{local_window_size}"
+                        name = f"{kv_cache_type}_{type_str}_g{group_size}_h{head_size}_{win_str}"
+                        yield name, config, torch_type, ort_type
+
+
+@unittest.skipIf(not has_xqa(), "XQA is not available, skipping tests.")
+@unittest.skipIf(not has_quantized_kv_cache(), "Quantized KV Cache is not available, skipping tests.")
+class TestXQAQuantizedSlidingWindowParity(unittest.TestCase):
+    """Verify the quantized (INT8/FP8) XQA sliding-window (local attention) decode path matches the reference."""
+
+    def tearDown(self):
+        """Clear CUDA cache after each test to prevent memory corruption in batch runs."""
+        if torch.cuda.is_available():
+            torch.cuda.synchronize()
+            torch.cuda.empty_cache()
+        gc.collect()
+
+    @parameterized.expand(gqa_xqa_quantized_sliding_window_test_cases())
+    def test_xqa_quantized_sliding_window_parity(self, name, config, torch_type, ort_type):
+        """Test XQA quantized parity with a sliding (local) attention window."""
+        type_str = "bf16" if torch_type == torch.bfloat16 else "fp16"
+        rtol_key = f"{config.kv_cache_type}_{type_str}"
+
+        def run_parity_check():
+            parity_check_gqa_past(
+                config=config,
+                ep="CUDAExecutionProvider",
+                device="cuda",
+                torch_type=torch_type,
+                ort_type=ort_type,
+                causal=True,
+                rtol=rtol[rtol_key],
+                atol=atol[rtol_key],
+                std=0.1,
+            )
+
+        # XQA is enabled by default for fp16/bf16, so the quantized sliding-window path is selected.
+        self.assertEqual("XQA", get_sdpa_kernel_from_debug_info(run_parity_check))
 
 
 @unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.")

From ae18718ea0989180088726a1117a0d4cab53343c Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Mon, 22 Jun 2026 21:08:54 -0700
Subject: [PATCH 11/12] address feedbacks

---
 .../contrib_ops/cuda/bert/group_query_attention.cc | 10 +++++++---
 .../contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh     |  7 ++++++-
 onnxruntime/test/python/transformers/test_gqa.py   | 14 ++++++++++++--
 3 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
index 9fe8460991523..7cfb7c73fd016 100644
--- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc
+++ b/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,9 +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));
-  ORT_ENFORCE(local_window_size_ == -1 || local_window_size_ > 0,
-              "local_window_size must be -1 or greater than 0.");
+  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);
diff --git a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh
index 0f38224d4704f..a759c3d7f9760 100644
--- a/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh
+++ b/onnxruntime/contrib_ops/cuda/bert/xqa/xqa_impl_gen.cuh
@@ -60,7 +60,9 @@ inline Status Launch(
     [[maybe_unused]] const float* kv_cache_scale,
     [[maybe_unused]] void* workspace,
     [[maybe_unused]] size_t workspace_size,
-    [[maybe_unused]] const int local_window_size = -1) {
+    // 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));
@@ -94,6 +96,9 @@ inline Status Launch(
   }
 
 #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
diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index 233f7ac102b92..92a216aea265c 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -74,6 +74,13 @@
 
 
 class CaptureStdout:
+    """Capture output written to OS file descriptor 1 (C++ stdout).
+
+    Uses fd-level dup2 redirection rather than contextlib.redirect_stdout because the kernel
+    debug info is emitted by the native ONNX Runtime library directly to fd 1, which Python's
+    redirect_stdout (which only swaps sys.stdout) cannot intercept.
+    """
+
     def __init__(self):
         self.fd = 1
         self.chunk_size = 1024
@@ -2805,12 +2812,14 @@ def gqa_xqa_sliding_window_test_cases():
     # Two window/past relationships are covered:
     #   past > window  -> the sliding mask drops the oldest keys (the new code path).
     #   past <= window -> the window spans the whole cache (parity with global attention).
+    #   past + 1 == window -> the exact guard boundary (cacheSeqLen == slidingWinSize) that locks
+    #                         down the kernel's `>` vs `>=` window comparison.
     # has_head_sink toggles the GPT-OSS attention-sink input, which composes with the window
     # in-kernel; both with and without a sink are exercised.
     for torch_type, ort_type in [(torch.float16, TensorProto.FLOAT16), (torch.bfloat16, TensorProto.BFLOAT16)]:
         for head_size in [64, 128]:
             for group_size in [4, 8]:
-                for past_kv_sequence_length, local_window_size in [(512, 128), (64, 128)]:
+                for past_kv_sequence_length, local_window_size in [(512, 128), (64, 128), (127, 128)]:
                     for has_head_sink in [False, True]:
                         kv_num_heads = 4
                         num_heads = kv_num_heads * group_size
@@ -2880,6 +2889,7 @@ def gqa_xqa_quantized_sliding_window_test_cases():
     # Two window/past relationships are covered:
     #   past > window  -> the sliding mask drops the oldest keys (the new code path).
     #   past <= window -> the window spans the whole cache (parity with global attention).
+    #   past + 1 == window -> the exact guard boundary (cacheSeqLen == slidingWinSize).
     kv_cache_types = ["int8"]
     if has_fp8_kv_cache:
         kv_cache_types.append("fp8")
@@ -2887,7 +2897,7 @@ def gqa_xqa_quantized_sliding_window_test_cases():
         for torch_type, ort_type in [(torch.float16, TensorProto.FLOAT16), (torch.bfloat16, TensorProto.BFLOAT16)]:
             for head_size in [64, 128]:
                 for group_size in [4, 8]:
-                    for past_kv_sequence_length, local_window_size in [(512, 128), (64, 128)]:
+                    for past_kv_sequence_length, local_window_size in [(512, 128), (64, 128), (127, 128)]:
                         kv_num_heads = 4
                         num_heads = kv_num_heads * group_size
                         config = GQAConfig(

From ba08191267abedb7196dfe4779f7fff934cc5de6 Mon Sep 17 00:00:00 2001
From: Tianlei Wu <tlwu@microsoft.com>
Date: Mon, 22 Jun 2026 22:05:03 -0700
Subject: [PATCH 12/12] address copilot feedback: quantized XQA sliding-window
 doc + hermetic XQA tests

---
 docs/contrib_ops/cuda/gqa.md                  | 23 ++++++++--------
 .../test/python/transformers/test_gqa.py      | 27 +++++++++++++++++++
 2 files changed, 38 insertions(+), 12 deletions(-)

diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md
index 55078a70afea1..b38c6d4fc9c92 100644
--- a/docs/contrib_ops/cuda/gqa.md
+++ b/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 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. |
+| 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,9 +233,9 @@ enabled (see §10).
 
 ### 6.1 XQA
 
-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.
+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
 
@@ -295,8 +295,8 @@ 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. Global attention, **or** local (sliding) window attention (`local_window_size > 0`) on the
-   non-quantized path. The quantized (INT8/FP8) variants remain global-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
@@ -452,12 +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 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.
+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.
diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py
index 92a216aea265c..cd28c25d5c523 100644
--- a/onnxruntime/test/python/transformers/test_gqa.py
+++ b/onnxruntime/test/python/transformers/test_gqa.py
@@ -2665,8 +2665,17 @@ def gqa_xqa_test_cases():
 class TestXQAQuantizedParity(unittest.TestCase):
     """Tests that verify fused kernels produce the same results as unfused kernels."""
 
+    def setUp(self):
+        # Force XQA on so the test is hermetic even if ORT_ENABLE_XQA=0 is set in the environment.
+        self._prev_enable_xqa = os.environ.get("ORT_ENABLE_XQA")
+        os.environ["ORT_ENABLE_XQA"] = "1"
+
     def tearDown(self):
         """Clear CUDA cache after each test to prevent memory corruption in batch runs."""
+        if self._prev_enable_xqa is None:
+            os.environ.pop("ORT_ENABLE_XQA", None)
+        else:
+            os.environ["ORT_ENABLE_XQA"] = self._prev_enable_xqa
         if torch.cuda.is_available():
             torch.cuda.synchronize()
             torch.cuda.empty_cache()
@@ -2850,8 +2859,17 @@ def gqa_xqa_sliding_window_test_cases():
 class TestXQASlidingWindowParity(unittest.TestCase):
     """Verify the non-quantized XQA sliding-window (local attention) decode path matches the reference."""
 
+    def setUp(self):
+        # Force XQA on so the test is hermetic even if ORT_ENABLE_XQA=0 is set in the environment.
+        self._prev_enable_xqa = os.environ.get("ORT_ENABLE_XQA")
+        os.environ["ORT_ENABLE_XQA"] = "1"
+
     def tearDown(self):
         """Clear CUDA cache after each test to prevent memory corruption in batch runs."""
+        if self._prev_enable_xqa is None:
+            os.environ.pop("ORT_ENABLE_XQA", None)
+        else:
+            os.environ["ORT_ENABLE_XQA"] = self._prev_enable_xqa
         if torch.cuda.is_available():
             torch.cuda.synchronize()
             torch.cuda.empty_cache()
@@ -2930,8 +2948,17 @@ def gqa_xqa_quantized_sliding_window_test_cases():
 class TestXQAQuantizedSlidingWindowParity(unittest.TestCase):
     """Verify the quantized (INT8/FP8) XQA sliding-window (local attention) decode path matches the reference."""
 
+    def setUp(self):
+        # Force XQA on so the test is hermetic even if ORT_ENABLE_XQA=0 is set in the environment.
+        self._prev_enable_xqa = os.environ.get("ORT_ENABLE_XQA")
+        os.environ["ORT_ENABLE_XQA"] = "1"
+
     def tearDown(self):
         """Clear CUDA cache after each test to prevent memory corruption in batch runs."""
+        if self._prev_enable_xqa is None:
+            os.environ.pop("ORT_ENABLE_XQA", None)
+        else:
+            os.environ["ORT_ENABLE_XQA"] = self._prev_enable_xqa
         if torch.cuda.is_available():
             torch.cuda.synchronize()
             torch.cuda.empty_cache()