diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 95f0698534a74..23cdbac30a042 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -2637,6 +2637,8 @@ This version of the operator has been available since version 1 of the 'com.micr
Rotate using interleaved pattern. Default value is 0 (False).
scale : float
Custom scale will be used if specified. Default value is 1/sqrt(head_size)
+
sliding_window_cache : int
+
Set to 1 when the past/present KV buffers are window-sized instead of holding the whole sequence. The op then keeps only the min(total_sequence_length, cache_capacity) most recent tokens, contiguously, using cache-relative indexing and evicting from the front as needed. Requires local_window_size > 0 and a cache capacity of at least local_window_size. Multi-token steps may use a temporary staging buffer, so the capacity need not cover the entire step. Default value is 0 (full-length cache).
smooth_softmax : int
Use a smooth factor in softmax.
softcap : float
@@ -2973,6 +2975,7 @@ This version of the operator has been available since version 1 of the 'com.micr
Constrain scalar scales to FP32.
+ ### **com.microsoft.MatMulBlockQuantizedFp8Weight** Weight-only block-scaled FP8 (E4M3) matrix multiplication. diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md index 957b0498864b0..327312e22df0f 100644 --- a/docs/contrib_ops/cuda/gqa.md +++ b/docs/contrib_ops/cuda/gqa.md @@ -56,6 +56,7 @@ Selected attributes: | `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. | +| `sliding_window_cache` | Set to `1` when using a windowed (sliding-window) KV cache instead of full-length. When enabled, the operator keeps only the most recent tokens, using cache-relative indexing and evicting from the front as needed. Requires `local_window_size > 0`. Defaults to `0` (full-length cache). | | `do_rotary` / `rotary_interleaved` | Enable RoPE and select interleaved vs. half-rotary layout. | | `smooth_softmax` | Add a smooth factor to the softmax denominator. | | `qk_norm_epsilon` | Epsilon for the fused per-head Q/K RMSNorm (QK-Norm) prologue. Defaults to `1e-6`. | @@ -155,6 +156,40 @@ holds (the past and present tensors alias the same memory), the cache length is 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. +### Sliding Window (Windowed) KV Cache + +When the `sliding_window_cache` attribute is set to `1`, the KV cache operates in a windowed +(sliding-window) mode, keeping only the most recent `local_window_size` tokens instead of growing +to the full sequence length. This significantly reduces memory usage for long-context models that +employ local attention (e.g., GPT-OSS with layer-wise sliding windows). + +**Key behaviors:** + +- **Cache capacity:** The cache buffer size is fixed to `kv_cache_capacity` (the initial allocation + size), which must be at least `local_window_size`. Allocating a little more than the window is + worthwhile on CPU: the surplus is what lets the append point drift, so the cache is compacted once + every `kv_cache_capacity - local_window_size + 1` steps instead of on every step. A modest surplus + (order tens of entries) is the sweet spot; a much larger buffer starts costing more in attention + over out-of-window entries than it saves in compaction. +- **Cache-relative indexing:** New keys/values are appended at cache-relative positions, not global + positions. Once the window is full, old tokens are evicted from the front. +- **RoPE position bounds:** The `rotary_max_position` parameter controls the upper bound (exclusive) + for RoPE position indices, decoupling absolute sequence positions from cache buffer indices. +- **Multi-token staging:** For multi-token (prompt) steps that exceed window size, a temporary + staging buffer is used internally to handle compaction and eviction. +- **Present shape:** When using windowed cache with `past_present_share_buffer`, the `present_key` + and `present_value` shapes remain bounded by `kv_cache_capacity` in the sequence dimension, + rather than growing with `total_sequence_length`. +- **Constraints:** `sliding_window_cache=1` requires `local_window_size > 0`. Windowed caches are + incompatible with the Flash-Attention fast-decode path and instead use the XQA or standard + attention backends. +- **CPU support:** The CPU kernel implements the same contract. Entries live at `[0, end)` and are + appended at `end`, so a step only moves memory when the append would run past the capacity; at + that point the surviving entries are compacted to the front in one go. Multi-token steps that + would drop entries the step itself still reads are staged instead, so results match a full-length + cache exactly. On CPU the `attention_bias` input, the `qk_output` attribute and a shared KV layout + (`key`/`value` folded into `query`) are not supported together with `sliding_window_cache=1`. + ### Quantized KV cache To reduce the KV-cache memory footprint, the cache may be stored quantized while `query` stays diff --git a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h index 11be4be1638df..7b3c38cc883e8 100644 --- a/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h +++ b/onnxruntime/contrib_ops/cpu/bert/attention_parameters.h @@ -104,6 +104,21 @@ struct GroupQueryAttentionParameters : AttentionParameters { KVQuantizationType k_quant_type = KVQuantizationType::NONE; KVQuantizationType v_quant_type = KVQuantizationType::NONE; int kv_cache_bit_width = 0; + + // Windowed (sliding-window) KV cache. Set from the sliding_window_cache attribute. + // When true, past/present KV buffers are allocated with capacity C = kv_cache_capacity, which may + // be much smaller than total_sequence_length. The cache then holds only the min(T, C) most recent + // tokens at cache indices [0, L) and the kernel uses cache-relative indexing plus a shift + // compaction step. Requires local_window_size > 0. See docs: windowed KV cache design. + bool is_windowed_kv_cache = false; + int kv_cache_capacity = 0; // Capacity of the KV buffer used by this step. + int kv_cache_real_capacity = 0; // C: allocated sequence dim of the past/present KV buffers. + // A multi-token step runs against a longer staging buffer (see GroupQueryAttention::ComputeInternal), + // in which case kv_cache_capacity is the staging capacity and differs from kv_cache_real_capacity. + + // Upper bound (exclusive) for absolute RoPE positions: dim 0 of the cos/sin caches. + // 0 when rotary is not configured. + int rotary_max_position = 0; }; // Parameters deduced from node attributes and inputs/outputs. diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index e2d1bebbbae38..e6efc4c207300 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -14,12 +14,168 @@ #include "core/platform/threadpool.h" #include +#include +#include +#include #include using onnxruntime::concurrency::ThreadPool; namespace onnxruntime { namespace contrib { +namespace { + +// --------------------------------------------------------------------------------------------- +// Windowed ("sliding window") KV cache. +// +// The cache holds only the most recent positions of the sequence, at cache indices [0, L) with +// W <= L <= C, where W is the attention window and C the buffer capacity. Both the causal mask and +// the local-window mask depend only on the distance between a query and a key, so translating the +// cache coordinate system leaves them unchanged: running the existing attention paths against +// cache-relative sequence lengths is exact. RoPE is the only position-absolute quantity, and it is +// applied to Q/K from the absolute positions before any of this happens. +// +// Resident entries live at [0, end) and new entries are appended at `end`, so a step only has to +// move memory when the append would run past C. The oldest resident entries are usually already +// older than the window, but the local-window mask hides them, so they can be left alone instead +// of being shifted out one row at a time: compacting reclaims `gap = C - W + 1` rows at once and +// therefore runs once every `gap` decode steps instead of on every step. Slack above the window is +// what buys those free steps; at C == W, gap == 1 and every step compacts. +// +// `end` is a pure function of the absolute past length P, so the op stays stateless across Run() +// calls -- the cache buffer remains the only state carried between steps. +// +// A step whose compaction would drop rows its own earliest query still reads runs against a +// staging cache of `end + S` entries instead, and only the surviving tail is written back. +// --------------------------------------------------------------------------------------------- + +// Per-batch row geometry of one windowed step. +struct WindowedStep { + int seed_offset = 0; // first resident row read out of the real cache (in-place shift distance) + int seed_rows = 0; // resident rows carried into the buffer that attention runs against + int write_offset = 0; // first row of the staging buffer written back into the real cache + int write_rows = 0; // rows written back (0 when the step ran in place) +}; + +// Offset one past the last resident row, i.e. where the next entry is appended. Whole `gap`-sized +// blocks are reclaimed at once, which keeps `end` in [C - gap + 1, C] once the cache has filled. +// +// The intermediate math runs in int64_t: `past_sequence_length` is derived from the caller-supplied +// seqlens_k input, so `overflow + gap - 1` and `gap * blocks` would be able to overflow for values +// near INT32_MAX -- signed overflow is UB, so it has to be avoided even for inputs that later +// validation would reject. The result is bounded by `capacity`, so narrowing back to int is safe. +int WindowedCacheEnd(int64_t past_sequence_length, int64_t capacity, int64_t gap) { + if (past_sequence_length <= capacity) { + return static_cast(past_sequence_length); // still filling: nothing has been reclaimed yet + } + const int64_t overflow = past_sequence_length - capacity; + const int64_t reclaimed = gap * ((overflow + gap - 1) / gap); + return static_cast(past_sequence_length - reclaimed); +} + +// Fills `steps` and the cache-relative seqlens_k, and decides whether the step needs a staging +// cache. `seqlens_k[b]` is the absolute total length minus one, i.e. T - 1 = P + S - 1. +void PlanWindowedKvCache(const int32_t* seqlens_k, + int batch_size, + int sequence_length, + int capacity, + int local_window_size, + std::vector& steps, + std::vector& cache_seqlens_k, + bool& use_staging, + int& staged_capacity) { + steps.assign(batch_size, WindowedStep{}); + cache_seqlens_k.resize(batch_size); + + // CheckInputs validates capacity >= local_window_size, so this is at least 1. + const int gap = capacity - local_window_size + 1; + + int max_resident = 0; + use_staging = false; + for (int b = 0; b < batch_size; ++b) { + // Computed in int64_t for the same reason as in WindowedCacheEnd: seqlens_k is caller-supplied + // and `seqlens_k[b] + 1` would overflow at INT32_MAX. Clamping at zero keeps a bogus (too + // small) seqlens_k from turning into a negative row count downstream. + const int64_t past_sequence_length = + std::max(static_cast(seqlens_k[b]) + 1 - sequence_length, 0); // P + const int end_before = WindowedCacheEnd(past_sequence_length, capacity, gap); + const int end_after = WindowedCacheEnd(past_sequence_length + sequence_length, capacity, gap); + const int kept = end_after - sequence_length; // resident rows that survive this step + + // The first new row of the step is read by a query that also reaches local_window_size - 1 + // rows further back, so compacting down to `kept` rows in place is only valid while that much + // history survives (all of it while the sequence is still shorter than the window). + const int64_t required = std::min(past_sequence_length, local_window_size - 1); + use_staging = use_staging || kept < required; + + steps[b].seed_rows = end_before; + steps[b].seed_offset = end_before - kept; // rows dropped off the front, 0 on a drifting step + max_resident = std::max(max_resident, end_before); + } + + staged_capacity = use_staging ? max_resident + sequence_length : capacity; + + for (int b = 0; b < batch_size; ++b) { + WindowedStep& step = steps[b]; + if (use_staging) { + // Seed the staging cache with everything resident, then write back only the tail that the + // post-step layout keeps. + step.write_offset = step.seed_offset; + step.write_rows = step.seed_rows + sequence_length - step.seed_offset; + step.seed_offset = 0; + } else { + step.seed_rows -= step.seed_offset; // becomes `kept`; 0 shift on a drifting step + } + cache_seqlens_k[b] = static_cast(step.seed_rows + sequence_length - 1); + } +} + +// Moves whole KV rows per (batch, kv_head) between two BNSH caches whose sequence capacities may +// differ. Rows are copied verbatim, so this is independent of the cache element type and of any +// quantization packing. Each (batch, kv_head) slice is disjoint, so the moves are parallelized; +// this is pure memory traffic, which is why the caller skips it on a drifting step. +void MoveKvCacheRows(const uint8_t* src, + uint8_t* dst, + int batch_size, + int kv_num_heads, + int src_capacity, + int dst_capacity, + size_t row_bytes, + const std::vector& steps, + bool write_back, + ThreadPool* thread_pool) { + int max_rows = 0; + for (int b = 0; b < batch_size; ++b) { + max_rows = std::max(max_rows, write_back ? steps[b].write_rows : steps[b].seed_rows); + } + if (max_rows <= 0) { + return; + } + + const double moved_bytes = static_cast(max_rows) * static_cast(row_bytes); + const TensorOpCost cost{moved_bytes, moved_bytes, moved_bytes}; + ThreadPool::TryParallelFor( + thread_pool, static_cast(batch_size) * kv_num_heads, cost, + [&](std::ptrdiff_t begin, std::ptrdiff_t end) { + for (std::ptrdiff_t unit = begin; unit < end; ++unit) { + const int b = static_cast(unit / kv_num_heads); + const int rows = write_back ? steps[b].write_rows : steps[b].seed_rows; + if (rows <= 0) { + continue; + } + const int src_offset = write_back ? steps[b].write_offset : steps[b].seed_offset; + const SafeInt head = SafeInt(unit); + const uint8_t* src_rows = src + static_cast((head * src_capacity + src_offset) * row_bytes); + uint8_t* dst_rows = dst + static_cast(head * dst_capacity * row_bytes); + if (src_rows != dst_rows) { + // memmove, not memcpy: the in-place left shift of the real cache overlaps. + std::memmove(dst_rows, src_rows, static_cast(SafeInt(rows) * row_bytes)); + } + } + }); +} + +} // namespace // These ops are internal-only, so register outside of onnx #define REGISTER_KERNEL_TYPED(T) \ @@ -43,7 +199,11 @@ REGISTER_KERNEL_TYPED(MLFloat16) template GroupQueryAttention::GroupQueryAttention(const OpKernelInfo& info) - : OpKernel(info), GQAAttentionBase(info, true) {} + : OpKernel(info), GQAAttentionBase(info, true) { + sliding_window_cache_ = info.GetAttrOrDefault("sliding_window_cache", 0) == 1; + ORT_ENFORCE(!sliding_window_cache_ || local_window_size_ > 0, + "GroupQueryAttention (CPU): sliding_window_cache=1 requires local_window_size > 0."); +} template Status GroupQueryAttention::Compute(OpKernelContext* context) const { @@ -108,7 +268,11 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { total_seqlen_tensor, scale_, softcap_, - kv_cache_bit_width_)); + kv_cache_bit_width_, + /*max_threads_per_block=*/0, + /*kv_cache_extra_bits=*/0, + sliding_window_cache_, + local_window_size_)); ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckCustomAttentionInputs(position_ids, attention_bias, @@ -143,13 +307,18 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { { const int32_t* seqlens_k_data = seqlens_k->Data(); const int past_kv_seqlen = parameters.seqlen_past_kv_cache; + const bool windowed = parameters.is_windowed_kv_cache; for (int b = 0; b < batch_size; b++) { - if (seqlens_k_data[b] < 0 || seqlens_k_data[b] >= present_kv_seqlen) { + // A windowed cache is deliberately shorter than the sequence it serves, so seqlens_k (which + // stays absolute) is not bounded by the buffer. The resident row count is clamped to the + // capacity instead, when the step is planned. + if (seqlens_k_data[b] < 0 || (!windowed && seqlens_k_data[b] >= present_kv_seqlen)) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "seqlens_k[", b, "] = ", seqlens_k_data[b], " is out of range [0, ", present_kv_seqlen, ")"); } - if (!parameters.is_first_prompt && static_cast(seqlens_k_data[b]) + 1 < sequence_length) { + if ((windowed || !parameters.is_first_prompt) && + static_cast(seqlens_k_data[b]) + 1 < sequence_length) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "seqlens_k[", b, "] = ", seqlens_k_data[b], " is too small for sequence_length ", sequence_length); @@ -163,7 +332,7 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { // Shared KV (kv_sequence_length == 0) appends no new KV and its past read is already // bounded by the present-buffer check together with the total_sequence_length <= // seqlen_past_kv_cache enforcement in the apply-attention paths, so it needs no check here. - if (past_key != nullptr && past_value != nullptr && parameters.kv_sequence_length != 0 && + if (!windowed && past_key != nullptr && past_value != nullptr && parameters.kv_sequence_length != 0 && !parameters.is_first_prompt) { const int64_t past_rows = static_cast(seqlens_k_data[b]) + 1 - sequence_length; if (past_rows > past_kv_seqlen) { @@ -195,6 +364,14 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { ORT_RETURN_IF_ERROR(group_query_attention_helper::CheckOutputs(output_qk, qk_output_)); + // The QK output is laid out over the absolute total_sequence_length, while a windowed cache + // computes scores over cache-relative positions only. Reject the combination rather than + // emitting a buffer whose columns do not mean what the shape says. + if (parameters.is_windowed_kv_cache && output_qk != nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "GroupQueryAttention (CPU): qk_output is not implemented with sliding_window_cache=1."); + } + AllocatorPtr allocator; ORT_RETURN_IF_ERROR(context->GetTempSpaceAllocator(&allocator)); @@ -318,71 +495,197 @@ Status GroupQueryAttention::Compute(OpKernelContext* context) const { const T* head_sink_data = (head_sink != nullptr) ? head_sink->Data() : nullptr; - // Quantized KV cache path: quantize-on-write + direct MLAS QK/SV. - if (kv_quant_enabled_) { - const T* k_data_q = packed_qkv ? nullptr : k_rotary; - const T* v_data_q = packed_qkv ? nullptr : V.Get().Data(); - auto mlas_quant_type = ToMlasKVQuantType(k_quant_type_, kv_cache_bit_width_); - - // Use flash attention path when: - // 1. Total sequence length is long enough to benefit from tiling - // 2. No features that flash path doesn't support (softcap, smooth softmax, output_qk) - const bool use_flash = !disable_gqa_flash_ && - parameters.total_sequence_length > 1 && - softcap_ == 0.0f && - !use_smooth_softmax_ && - head_sink_data == nullptr && - output_qk == nullptr; - - if (use_flash) { - return ApplyAttentionQuantizedFlash( - q_rotary, k_data_q, v_data_q, - attention_bias, - past_key, past_value, - output, present_k, present_v, seqlens_k, - k_scale->Data(), v_scale->Data(), - mlas_quant_type, parameters, allocator, context); + // ---- Windowed KV cache: switch to cache-relative coordinates ---- + // RoPE above is the last consumer of absolute positions, so everything from here on can run + // against a cache that holds only the most recent min(T, C) entries. The attention paths derive + // the causal offset and the mask from seqlens_k alone, so redirecting them at a cache-relative + // seqlens_k (and, when the step would evict entries it still needs, at a staging cache) reuses + // them unchanged. + const Tensor* attention_past_key = past_key; + const Tensor* attention_past_value = past_value; + Tensor* attention_present_key = present_k; + Tensor* attention_present_value = present_v; + const Tensor* attention_seqlens_k = seqlens_k; + + std::vector windowed_steps; + std::vector windowed_cache_seqlens; + std::optional windowed_seqlens_tensor; + std::optional staged_key_tensor; + std::optional staged_value_tensor; + BufferUniquePtr staged_key_buffer; + BufferUniquePtr staged_value_buffer; + bool windowed_use_staging = false; + int windowed_capacity = 0; + int windowed_staged_capacity = 0; + size_t windowed_row_bytes = 0; + + ThreadPool* windowed_thread_pool = context->GetOperatorThreadPool(); + + if (parameters.is_windowed_kv_cache) { + ORT_RETURN_IF(present_k == nullptr || present_v == nullptr, + "sliding_window_cache=1 requires the present_key and present_value outputs."); + windowed_capacity = parameters.kv_cache_capacity; + PlanWindowedKvCache(seqlens_k->Data(), batch_size, sequence_length, windowed_capacity, + local_window_size_, windowed_steps, windowed_cache_seqlens, windowed_use_staging, + windowed_staged_capacity); + + // Rows are moved verbatim, so the row size is taken from the cache tensor itself and covers + // fp32/fp16 as well as the int8 and (nibble-packed) int4 quantized layouts. + windowed_row_bytes = SafeInt(past_key->Shape().GetDims()[3]) * past_key->DataType()->Size(); + const uint8_t* past_key_bytes = static_cast(past_key->DataRaw()); + const uint8_t* past_value_bytes = static_cast(past_value->DataRaw()); + + if (windowed_use_staging) { + const size_t staged_bytes = SafeInt(batch_size) * kv_num_heads_ * + windowed_staged_capacity * windowed_row_bytes; + staged_key_buffer = BufferUniquePtr(allocator->Alloc(staged_bytes), BufferDeleter(allocator)); + staged_value_buffer = BufferUniquePtr(allocator->Alloc(staged_bytes), BufferDeleter(allocator)); + ORT_RETURN_IF(staged_key_buffer.get() == nullptr || staged_value_buffer.get() == nullptr, + "Failed to allocate the sliding_window_cache staging buffers."); + + const TensorShape staged_shape({static_cast(batch_size), static_cast(kv_num_heads_), + static_cast(windowed_staged_capacity), + past_key->Shape().GetDims()[3]}); + staged_key_tensor.emplace(past_key->DataType(), staged_shape, staged_key_buffer.get(), allocator->Info()); + staged_value_tensor.emplace(past_value->DataType(), staged_shape, staged_value_buffer.get(), allocator->Info()); + + MoveKvCacheRows(past_key_bytes, static_cast(staged_key_buffer.get()), batch_size, kv_num_heads_, + windowed_capacity, windowed_staged_capacity, windowed_row_bytes, windowed_steps, + /*write_back=*/false, windowed_thread_pool); + MoveKvCacheRows(past_value_bytes, static_cast(staged_value_buffer.get()), batch_size, kv_num_heads_, + windowed_capacity, windowed_staged_capacity, windowed_row_bytes, windowed_steps, + /*write_back=*/false, windowed_thread_pool); + + attention_past_key = &staged_key_tensor.value(); + attention_past_value = &staged_value_tensor.value(); + attention_present_key = &staged_key_tensor.value(); + attention_present_value = &staged_value_tensor.value(); + } else { + // Compact in place: the surviving window moves to the front of the real cache, which then + // doubles as the past buffer so the append below does not copy it a second time. + uint8_t* present_key_bytes = static_cast(present_k->MutableDataRaw()); + uint8_t* present_value_bytes = static_cast(present_v->MutableDataRaw()); + + // A drifting step appends inside the buffer and reclaims nothing, so it moves no memory at + // all -- this is the common decode step. The copy is still needed when past and present are + // separate buffers, since the append below only writes the new rows. + const bool compacts = std::any_of(windowed_steps.begin(), windowed_steps.end(), + [](const WindowedStep& step) { return step.seed_offset > 0; }); + if (compacts || past_key_bytes != present_key_bytes || past_value_bytes != present_value_bytes) { + MoveKvCacheRows(past_key_bytes, present_key_bytes, batch_size, kv_num_heads_, + windowed_capacity, windowed_capacity, windowed_row_bytes, windowed_steps, + /*write_back=*/false, windowed_thread_pool); + MoveKvCacheRows(past_value_bytes, present_value_bytes, batch_size, kv_num_heads_, + windowed_capacity, windowed_capacity, windowed_row_bytes, windowed_steps, + /*write_back=*/false, windowed_thread_pool); + } + + attention_past_key = present_k; + attention_past_value = present_v; } - return ApplyAttentionQuantized( - q_rotary, k_data_q, v_data_q, head_sink_data, - attention_bias, past_key, past_value, - output, present_k, present_v, output_qk, seqlens_k, - k_scale->Data(), v_scale->Data(), - mlas_quant_type, parameters, allocator, context); + windowed_seqlens_tensor.emplace(DataTypeImpl::GetType(), + TensorShape({static_cast(batch_size)}), + windowed_cache_seqlens.data(), allocator->Info()); + attention_seqlens_k = &windowed_seqlens_tensor.value(); + + // Cache-relative lengths always have a non-empty past (possibly of length zero), so the + // first-prompt shortcut must be off: past_seqlen = total_seqlen - sequence_length already + // yields 0 when nothing is resident. + parameters.is_first_prompt = false; + parameters.is_subsequent_prompt = false; + parameters.seqlen_past_kv_cache = windowed_staged_capacity; + parameters.seqlen_present_kv_cache = windowed_staged_capacity; + parameters.kv_cache_capacity = windowed_staged_capacity; + parameters.total_sequence_length = + *std::max_element(windowed_cache_seqlens.begin(), windowed_cache_seqlens.end()) + 1; } - // Compute the attention score and apply the score to V - const T* k_data = packed_qkv ? nullptr : k_rotary; - const T* v_data = packed_qkv ? nullptr : V.Get().Data(); - - // Non-quantized flash attention path (float only). Uses the tiled online-softmax - // kernel to avoid materializing the full attention score matrix. Falls back to the - // naive path when an unsupported feature is requested (softcap, smooth softmax, - // head sink, or QK output). - // - // Prefill (sequence_length > 1) uses the tiled kernel; single-token decode - // (sequence_length == 1 with total_sequence_length > 1) uses the dedicated GEMV - // decode kernel. Both are reached when total_sequence_length > 1. - if constexpr (std::is_same_v) { - const bool use_flash = !disable_gqa_flash_ && - parameters.total_sequence_length > 1 && - softcap_ == 0.0f && - !use_smooth_softmax_ && - head_sink_data == nullptr && - output_qk == nullptr && - present_k != nullptr && present_v != nullptr; - if (use_flash) { - return ApplyAttentionFlash(q_rotary, k_data, v_data, - attention_bias, past_key, past_value, - output, present_k, present_v, seqlens_k, - parameters, allocator, context); + auto run_attention = [&]() -> Status { + // Quantized KV cache path: quantize-on-write + direct MLAS QK/SV. + if (kv_quant_enabled_) { + const T* k_data_q = packed_qkv ? nullptr : k_rotary; + const T* v_data_q = packed_qkv ? nullptr : V.Get().Data(); + auto mlas_quant_type = ToMlasKVQuantType(k_quant_type_, kv_cache_bit_width_); + + // Use flash attention path when: + // 1. Total sequence length is long enough to benefit from tiling + // 2. No features that flash path doesn't support (softcap, smooth softmax, output_qk) + const bool use_flash = !disable_gqa_flash_ && + parameters.total_sequence_length > 1 && + softcap_ == 0.0f && + !use_smooth_softmax_ && + head_sink_data == nullptr && + output_qk == nullptr; + + if (use_flash) { + return ApplyAttentionQuantizedFlash( + q_rotary, k_data_q, v_data_q, + attention_bias, + attention_past_key, attention_past_value, + output, attention_present_key, attention_present_value, attention_seqlens_k, + k_scale->Data(), v_scale->Data(), + mlas_quant_type, parameters, allocator, context); + } + + return ApplyAttentionQuantized( + q_rotary, k_data_q, v_data_q, head_sink_data, + attention_bias, attention_past_key, attention_past_value, + output, attention_present_key, attention_present_value, output_qk, attention_seqlens_k, + k_scale->Data(), v_scale->Data(), + mlas_quant_type, parameters, allocator, context); } + + // Compute the attention score and apply the score to V + const T* k_data = packed_qkv ? nullptr : k_rotary; + const T* v_data = packed_qkv ? nullptr : V.Get().Data(); + + // Non-quantized flash attention path (float only). Uses the tiled online-softmax + // kernel to avoid materializing the full attention score matrix. Falls back to the + // naive path when an unsupported feature is requested (softcap, smooth softmax, + // head sink, or QK output). + // + // Prefill (sequence_length > 1) uses the tiled kernel; single-token decode + // (sequence_length == 1 with total_sequence_length > 1) uses the dedicated GEMV + // decode kernel. Both are reached when total_sequence_length > 1. + if constexpr (std::is_same_v) { + const bool use_flash = !disable_gqa_flash_ && + parameters.total_sequence_length > 1 && + softcap_ == 0.0f && + !use_smooth_softmax_ && + head_sink_data == nullptr && + output_qk == nullptr && + attention_present_key != nullptr && attention_present_value != nullptr; + if (use_flash) { + return ApplyAttentionFlash(q_rotary, k_data, v_data, + attention_bias, attention_past_key, attention_past_value, + output, attention_present_key, attention_present_value, attention_seqlens_k, + parameters, allocator, context); + } + } + + return ApplyAttention(q_rotary, k_data, v_data, + head_sink_data, attention_bias, attention_past_key, attention_past_value, output, + attention_present_key, attention_present_value, + output_qk, attention_seqlens_k, parameters, allocator, context); + }; + + ORT_RETURN_IF_ERROR(run_attention()); + + if (windowed_use_staging) { + // Attention ran against the staging cache; keep only the entries the sliding window can still + // reach. Everything before write_offset is older than the oldest key any future query can see. + MoveKvCacheRows(static_cast(staged_key_buffer.get()), + static_cast(present_k->MutableDataRaw()), batch_size, kv_num_heads_, + windowed_staged_capacity, windowed_capacity, windowed_row_bytes, windowed_steps, + /*write_back=*/true, windowed_thread_pool); + MoveKvCacheRows(static_cast(staged_value_buffer.get()), + static_cast(present_v->MutableDataRaw()), batch_size, kv_num_heads_, + windowed_staged_capacity, windowed_capacity, windowed_row_bytes, windowed_steps, + /*write_back=*/true, windowed_thread_pool); } - return ApplyAttention(q_rotary, k_data, v_data, - head_sink_data, attention_bias, past_key, past_value, output, present_k, present_v, - output_qk, seqlens_k, parameters, allocator, context); + return Status::OK(); } } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.h index cbe99e0378108..bab368d48d943 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.h @@ -15,6 +15,11 @@ class GroupQueryAttention final : public OpKernel, public GQAAttentionBase { public: GroupQueryAttention(const OpKernelInfo& info); Status Compute(OpKernelContext* context) const override; + + private: + // sliding_window_cache attribute: the past/present KV buffers are window-sized and the kernel + // works in cache-relative coordinates, evicting from the front. Requires local_window_size_ > 0. + bool sliding_window_cache_; }; } // namespace contrib diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h index d7009a2d84b08..44e7def5f3bc7 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -216,7 +216,9 @@ Status CheckInputs(const T* query, float softcap, int kv_cache_bit_width, int max_threads_per_block = 0, - int kv_cache_extra_bits = 0) { + int kv_cache_extra_bits = 0, + bool sliding_window_cache = false, + int local_window_size = -1) { if (max_threads_per_block > 0 && num_heads > max_threads_per_block) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "num_heads should be no larger than ", max_threads_per_block); } @@ -325,9 +327,56 @@ Status CheckInputs(const T* query, int present_sequence_length = std::max(total_sequence_length, past_sequence_length); + // Windowed KV cache: the bound past/present buffer *is* the capacity C, which is intentionally + // smaller than total_sequence_length. The op keeps only the min(T, C) most recent tokens, so the + // present tensor must keep the buffer's own sequence dimension instead of growing to T. + int kv_cache_capacity = 0; + if (sliding_window_cache) { + if (local_window_size <= 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sliding_window_cache=1 requires local_window_size > 0."); + } + if (past_key == nullptr || past_value == nullptr) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sliding_window_cache=1 requires past_key and past_value to be present."); + } + if (kv_sequence_length == 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sliding_window_cache=1 is not supported when kv_sequence_length == 0 (shared KV)."); + } + + // The windowed-cache compaction/staging kernels copy raw KV rows using 16-byte vectors. + // Enforce an aligned row size so sliding_window_cache fails with a clear INVALID_ARGUMENT. + if (kv_cache_bit_width == 8 && (head_size % 16) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sliding_window_cache: head_size (", head_size, + ") must be a multiple of 16 for an 8-bit KV cache."); + } + if (kv_cache_bit_width == 4 && (head_size % 32) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sliding_window_cache: head_size (", head_size, + ") must be a multiple of 32 for a 4-bit KV cache."); + } + + kv_cache_capacity = past_sequence_length; + present_sequence_length = past_sequence_length; + + // The cache must be able to hold the whole window; everything older is unreachable anyway. + // There is deliberately no requirement involving sequence_length: a multi-token step runs + // against a longer staging buffer, so it never has to fit into C. + if (kv_cache_capacity < local_window_size) { + return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, + "sliding_window_cache: KV cache capacity (", kv_cache_capacity, + ") must be at least local_window_size (", local_window_size, ")."); + } + } + + int rotary_max_position = 0; int rotary_dim = 0; if (cos_cache != nullptr && sin_cache != nullptr) { ORT_RETURN_IF_ERROR(CheckRotaryCaches(cos_cache, sin_cache, head_size, total_sequence_length, rotary_dim)); + rotary_max_position = static_cast(std::min(cos_cache->Shape().GetDims()[0], + sin_cache->Shape().GetDims()[0])); // Validate seqlens_k against rotary cache size when rotary embeddings are enabled. // This prevents OOB access when deriving position IDs from seqlens_k during rotary embedding. @@ -382,6 +431,10 @@ Status CheckInputs(const T* query, output_parameters->kv_hidden_size = kv_hidden_size; output_parameters->kv_num_heads = kv_num_heads; output_parameters->rotary_dim = rotary_dim; + output_parameters->rotary_max_position = rotary_max_position; + output_parameters->is_windowed_kv_cache = sliding_window_cache; + output_parameters->kv_cache_capacity = kv_cache_capacity; + output_parameters->kv_cache_real_capacity = kv_cache_capacity; output_parameters->is_packed_qkv = is_packed_qkv; output_parameters->is_unidirectional = true; output_parameters->is_subsequent_prompt = is_subsequent_prompt; @@ -414,6 +467,11 @@ Status CheckCustomAttentionInputs(const T* position_ids, } if (attention_bias != nullptr) { + if (parameters.is_windowed_kv_cache) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "attention_bias with sliding_window_cache is not implemented in GroupQueryAttention."); + } + const auto& attn_bias_shape = attention_bias->Shape(); // TensorShape::operator[] is unchecked — validate the rank before indexing dims. if (attn_bias_shape.NumDimensions() != 4) { diff --git a/onnxruntime/contrib_ops/cuda/bert/attention_data.h b/onnxruntime/contrib_ops/cuda/bert/attention_data.h index 7974c9cd48783..0b87ea831b745 100644 --- a/onnxruntime/contrib_ops/cuda/bert/attention_data.h +++ b/onnxruntime/contrib_ops/cuda/bert/attention_data.h @@ -183,6 +183,21 @@ struct GroupQueryAttentionData { // Only used for first prompt: padded_seq_lens[b] = sequence_length int* padded_seq_lens = nullptr; + // Cache-relative sequence lengths, used when parameters.is_windowed_kv_cache is set. Shape [batch_size]. + // For a full-length (non-windowed) cache these simply alias past_seq_lens / total_seq_lens. + // cache_past_seq_lens[b] : append offset inside the capacity-C buffer, after eviction. May be + // negative on a first prompt longer than the capacity, in which case + // the leading (out-of-window) tokens are skipped by the append kernel. + // cache_total_seq_lens[b] : number of valid cache entries after the append, i.e. min(T, C). + // evict_counts[b] : number of entries D dropped from the front of the cache this step. + int* cache_past_seq_lens = nullptr; + int* cache_total_seq_lens = nullptr; + int* evict_counts = nullptr; + + // Scratch used by the windowed-cache compaction shift. Sized for one KV cache: + // batch_size * kv_num_heads * capacity * head_size elements of the storage type U. + void* compaction_scratch = nullptr; + // Flash buffers T* softmax_lse = nullptr; T* softmax_lse_accum = nullptr; diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 1df4857eaa263..897749a3f6b35 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -108,6 +108,9 @@ GroupQueryAttention::GroupQueryAttention(const OpKernelInfo& info) ORT_ENFORCE(local_window_size_attr == -1 || (local_window_size_attr > 0 && local_window_size_attr <= std::numeric_limits::max()), "local_window_size must be -1 or greater than 0 (and not exceed INT_MAX)."); local_window_size_ = static_cast(local_window_size_attr); + sliding_window_cache_ = info.GetAttrOrDefault("sliding_window_cache", 0) == 1; + ORT_ENFORCE(!sliding_window_cache_ || local_window_size_ > 0, + "GroupQueryAttention (CUDA): sliding_window_cache=1 requires local_window_size > 0."); do_rotary_ = info.GetAttrOrDefault("do_rotary", 0) == 1; rotary_interleaved_ = info.GetAttrOrDefault("rotary_interleaved", 0) == 1; scale_ = info.GetAttrOrDefault("scale", 0.0f); @@ -314,7 +317,10 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons scale_, softcap_, kv_cache_bit_width_, - device_prop.maxThreadsPerBlock)); + device_prop.maxThreadsPerBlock, + /*kv_cache_extra_bits=*/0, + sliding_window_cache_, + local_window_size_)); #ifndef USE_INT4_KV_CACHE if (kv_cache_bit_width_ == 4) { return ORT_MAKE_STATUS(ONNXRUNTIME, INVALID_ARGUMENT, "kv_cache_bit_width==4 is not enabled in this build."); @@ -426,6 +432,60 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons "past_key/present_key and past_value/present_value must be both shared or both separate."); parameters.past_present_share_buffer = past_key_shared; + // Eviction rewrites the cache in place, so past and present must be the same buffer. + ORT_RETURN_IF(parameters.is_windowed_kv_cache && !parameters.past_present_share_buffer, + "sliding_window_cache=1 requires past_key/present_key and past_value/present_value " + "to share the same buffer."); + + // The capacity C of a windowed cache is only guaranteed to cover the attention window, so a step + // that appends S > 1 tokens can transiently need min(P, C) + S entries: the earliest queries of + // the step still have to see keys that the last ones have already pushed out. Redirect such a + // step to a staging cache of that length, seeded with the resident window, and copy the surviving + // tail back afterwards. + // + // This is deliberately not gated on is_first_prompt. The op cannot tell a first prompt from a + // later prefill chunk: total_sequence_length is supplied by the caller and some frameworks + // (onnxruntime-genai) pass the allocated buffer length rather than the running total. Staging + // every multi-token step is correct for both cases and leaves the single-token decode path, which + // never needs more than C entries, on the cheap in-place eviction path. + // + // Doing this here, rather than inside a dedicated kernel path, keeps the staging independent of + // which attention backend is selected and reuses the existing RoPE/QK-Norm/quantization handling. + IAllocatorUniquePtr staged_key_buffer; + IAllocatorUniquePtr staged_value_buffer; + CudaU* windowed_present_key = nullptr; + CudaU* windowed_present_value = nullptr; + int windowed_cache_capacity = 0; + int staged_cache_capacity = 0; + if (parameters.is_windowed_kv_cache && parameters.sequence_length > 1) { + windowed_present_key = data.present_key; + windowed_present_value = data.present_value; + windowed_cache_capacity = parameters.seqlen_present_kv_cache; + staged_cache_capacity = SafeInt(windowed_cache_capacity) + parameters.sequence_length; + + const size_t staged_bytes = SafeInt(parameters.batch_size) * parameters.kv_num_heads * + staged_cache_capacity * dense_head_size * sizeof(U); + staged_key_buffer = GetScratchBuffer(staged_bytes, GetComputeStream(context)); + staged_value_buffer = GetScratchBuffer(staged_bytes, GetComputeStream(context)); + + ORT_RETURN_IF_ERROR(LaunchCopyKvCacheWindow( + staged_key_buffer.get(), staged_value_buffer.get(), + windowed_present_key, windowed_present_value, + parameters.batch_size, parameters.kv_num_heads, + /*src_capacity=*/windowed_cache_capacity, /*dst_capacity=*/staged_cache_capacity, + /*rows=*/windowed_cache_capacity, /*src_offsets=*/nullptr, + static_cast(dense_head_size * sizeof(U)), Stream(context))); + + data.past_key = reinterpret_cast(staged_key_buffer.get()); + data.past_value = reinterpret_cast(staged_value_buffer.get()); + data.present_key = reinterpret_cast(staged_key_buffer.get()); + data.present_value = reinterpret_cast(staged_value_buffer.get()); + + parameters.kv_cache_capacity = staged_cache_capacity; + parameters.seqlen_past_kv_cache = staged_cache_capacity; + parameters.seqlen_present_kv_cache = staged_cache_capacity; + } + bool is_inputs_quantized = (k_quant_type_ != KVQuantizationType::NONE) || (v_quant_type_ != KVQuantizationType::NONE); constexpr bool is_int8 = std::is_same::value; constexpr bool is_fp8 = std::is_same::value; @@ -628,7 +688,9 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons // The fast-decode path lets the flash kernel perform RoPE and KV-append internally, bypassing // PrepareQKV (and therefore the fused QK-Norm prologue). Disable it when q/k norm weights are // present so the regular FlashAttention path (which normalizes via PrepareQKV) is used instead. - data.use_flash_attention_fast_decode = use_flash_attention && !disable_flash_decode_ && !parameters.is_first_prompt && parameters.kv_sequence_length > 0 && parameters.past_present_share_buffer && !is_inputs_quantized && !parameters.use_qk_norm; + // It is also disabled for a windowed KV cache: the kernel derives both the absolute RoPE position + // and the cache append offset from a single seqlens_k value, which those two no longer share. + data.use_flash_attention_fast_decode = use_flash_attention && !disable_flash_decode_ && !parameters.is_first_prompt && parameters.kv_sequence_length > 0 && parameters.past_present_share_buffer && !is_inputs_quantized && !parameters.use_qk_norm && !parameters.is_windowed_kv_cache; if (use_flash_attention) { // Allocate Flash specific buffers (Softmax LSE, Accum) @@ -683,18 +745,30 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons } else { // Compute sequence length buffers (past_seq_lens and total_seq_lens). // Allocate buffer for both: first half is past_seq_lens, second half is total_seq_lens. - seq_lens_buffer = GetScratchBuffer(3 * parameters.batch_size, GetComputeStream(context)); + // A windowed cache needs three more per-batch vectors expressed in cache-relative coordinates. + const int seq_lens_vectors = parameters.is_windowed_kv_cache ? 6 : 3; + seq_lens_buffer = GetScratchBuffer(seq_lens_vectors * parameters.batch_size, GetComputeStream(context)); auto cuda_stream = Stream(context); data.past_seq_lens = seq_lens_buffer.get(); data.total_seq_lens = seq_lens_buffer.get() + parameters.batch_size; data.padded_seq_lens = data.total_seq_lens + parameters.batch_size; + if (parameters.is_windowed_kv_cache) { + data.cache_past_seq_lens = data.padded_seq_lens + parameters.batch_size; + data.cache_total_seq_lens = data.cache_past_seq_lens + parameters.batch_size; + data.evict_counts = data.cache_total_seq_lens + parameters.batch_size; + } ORT_RETURN_IF_ERROR(LaunchGetSequenceLengths(total_seq_lens_minus_one->Data(), data.past_seq_lens, data.total_seq_lens, data.padded_seq_lens, + data.cache_past_seq_lens, + data.cache_total_seq_lens, + data.evict_counts, parameters.batch_size, parameters.sequence_length, parameters.is_first_prompt, + parameters.kv_cache_capacity, + parameters.kv_cache_real_capacity, cuda_stream, device_prop.maxThreadsPerBlock)); DUMP_TENSOR_INIT(); @@ -703,6 +777,20 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons DUMP_TENSOR("padded_seq_lens", data.padded_seq_lens, parameters.batch_size, 1); } + // Scratch used to left-shift the KV cache when the sliding window advances. Sized for the whole + // cache because the number of evicted rows is only known on device (and must stay constant across + // CUDA graph replays). K and V shift in the same kernel, so each needs its own half. + IAllocatorUniquePtr compaction_buffer; + if (parameters.is_windowed_kv_cache && staged_cache_capacity == 0) { + const size_t row_bytes = (parameters.kv_cache_bit_width == 0) + ? static_cast(parameters.head_size) * sizeof(U) + : static_cast(parameters.head_size) * parameters.kv_cache_bit_width / 8; + const size_t compaction_bytes = 2 * static_cast(parameters.batch_size) * parameters.kv_num_heads * + parameters.kv_cache_real_capacity * row_bytes; + compaction_buffer = GetScratchBuffer(compaction_bytes, GetComputeStream(context)); + data.compaction_scratch = compaction_buffer.get(); + } + #if USE_MEMORY_EFFICIENT_ATTENTION if (!data.use_xqa && !data.use_cudnn_sdpa && !data.use_flash_attention) { // Fall back to memory efficient attention. @@ -854,6 +942,19 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons ORT_RETURN_IF_ERROR((QkvToContext( device_prop, cublas, ort_stream.get(), parameters, data))); + + if (windowed_present_key != nullptr) { + // Attention ran against the staging cache; keep only the entries the sliding window can still + // reach. evict_counts[b] = max(0, min(P, C) + S - C) is exactly where that tail starts. The rows + // are already in their final (possibly quantized) form. + const int row_bytes = static_cast(dense_head_size * sizeof(U)); + ORT_RETURN_IF_ERROR(LaunchCopyKvCacheWindow( + windowed_present_key, windowed_present_value, data.present_key, data.present_value, + parameters.batch_size, parameters.kv_num_heads, + /*src_capacity=*/staged_cache_capacity, /*dst_capacity=*/windowed_cache_capacity, + /*rows=*/windowed_cache_capacity, /*src_offsets=*/data.evict_counts, + row_bytes, Stream(context))); + } return Status::OK(); } diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h index dc0d5977c0129..8cdcedc3726f5 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.h @@ -29,6 +29,9 @@ class GroupQueryAttention final : public CudaKernel { int num_heads_; // number of attention heads int kv_num_heads_; // different for k and v for group query attention int local_window_size_; + // sliding_window_cache attribute: past/present KV buffers are window-sized and the kernel uses + // cache-relative indexing with shift compaction. Requires local_window_size_ > 0. + bool sliding_window_cache_; bool is_unidirectional_; bool is_past_bsnh_; bool do_rotary_; 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 59af390035b51..89144f5e50787 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -147,6 +147,15 @@ Status LaunchPerHeadRMSNorm( // Internal helper to get Q, K, V pointers, handling packed input // +// Upper bound (exclusive) for absolute RoPE positions. The cos/sin caches are the real constraint; +// fall back to the KV cache sequence dimension when rotary caches are absent so behaviour is +// unchanged for models that do not use rotary embeddings. With a windowed KV cache these two differ +// by construction: RoPE stays in absolute coordinates while the cache is capacity-relative. +inline int RopeMaxPosition(const GroupQueryAttentionParameters& parameters) { + return parameters.rotary_max_position > 0 ? parameters.rotary_max_position + : parameters.seqlen_present_kv_cache; +} + // This function orchestrates the preparation of Q, K, and V tensors for attention kernels. // It performs: // 1. Handling packed vs. unpacked QKV inputs. @@ -231,7 +240,7 @@ Status PrepareQKV( stream, reinterpret_cast(q_out), reinterpret_cast(q_rope_input), data.position_ids, data.past_seq_lens, reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), - batch_size, sequence_length, num_heads, head_size, parameters.rotary_dim, max_cache_length, + batch_size, sequence_length, num_heads, head_size, parameters.rotary_dim, RopeMaxPosition(parameters), pos_format, parameters.rotary_interleaved, max_threads_per_block, false /* is_input_bnsh_format: Q is BSNH */))); } else if constexpr (std::is_same::value) { @@ -239,7 +248,7 @@ Status PrepareQKV( stream, reinterpret_cast(q_out), reinterpret_cast(q_rope_input), data.position_ids, data.past_seq_lens, reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), - batch_size, sequence_length, num_heads, head_size, parameters.rotary_dim, max_cache_length, + batch_size, sequence_length, num_heads, head_size, parameters.rotary_dim, RopeMaxPosition(parameters), pos_format, parameters.rotary_interleaved, max_threads_per_block, false /* is_input_bnsh_format: Q is BSNH */))); } @@ -254,7 +263,8 @@ Status PrepareQKV( parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.value), q_out, k, v, data.k_scale, data.v_scale, nullptr /* q_fold_scale */, num_heads, kv_num_heads, head_size, sequence_length, batch_size, - max_cache_length, data.past_seq_lens, + RopeMaxPosition(parameters), max_cache_length, + data.past_seq_lens, data.cache_past_seq_lens, reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, is_cache_bnsh, parameters.k_quant_type, @@ -605,6 +615,27 @@ Status LaunchUnpackQKV(const T* packed_qkv, T* unpacked_q, T* unpacked_k, T* unp return CUDA_CALL(cudaGetLastError()); } +// Offset one past the last resident row of a windowed KV cache, i.e. where the next entry is +// appended. Whole `gap`-sized blocks are reclaimed at once, which keeps the result in +// [capacity - gap + 1, capacity] once the cache has filled: +// end(P) = P while P <= capacity, else P - gap * ceil((P - capacity) / gap) +// +// This mirrors WindowedCacheEnd() in the CPU kernel (contrib_ops/cpu/bert/group_query_attention.cc); +// the two layouts must agree or CPU/CUDA parity breaks, so keep them in sync. +// +// The intermediate math runs in int64_t because `past_sequence_length` derives from the +// caller-supplied seqlens_k input: `overflow + gap - 1` and `gap * blocks` could otherwise overflow +// for values near INT32_MAX, and signed overflow is UB. The result is bounded by `capacity`, so +// narrowing back to int is safe. +__device__ __forceinline__ int WindowedCacheEnd(int64_t past_sequence_length, int64_t capacity, int64_t gap) { + if (past_sequence_length <= capacity) { + return static_cast(past_sequence_length); // still filling: nothing has been reclaimed yet + } + const int64_t overflow = past_sequence_length - capacity; + const int64_t reclaimed = gap * ((overflow + gap - 1) / gap); + return static_cast(past_sequence_length - reclaimed); +} + // ============================================================================ // GetSequenceLengths Kernel // ============================================================================ @@ -636,9 +667,14 @@ __global__ void GetSequenceLengths(const int* total_seq_lens_minus_one, int* past_seq_lens, int* total_seq_lens, int* padded_seq_lens, + int* cache_past_seq_lens, + int* cache_total_seq_lens, + int* evict_counts, const int batch_size, const int sequence_length, - const bool is_first_prompt) { + const bool is_first_prompt, + const int kv_cache_capacity, + const int kv_cache_real_capacity) { int i = threadIdx.x + blockIdx.x * blockDim.x; if (i < batch_size) { // total_seq_lens_minus_one is the seqlens_k input and is not range-checked on the device. @@ -647,14 +683,66 @@ __global__ void GetSequenceLengths(const int* total_seq_lens_minus_one, const int seqlens_k = total_seq_lens_minus_one[i]; const int total_len = (seqlens_k > 0 ? seqlens_k : 0) + 1; total_seq_lens[i] = total_len; + int past_len; if (is_first_prompt) { + past_len = 0; past_seq_lens[i] = 0; padded_seq_lens[i] = sequence_length; } else { - const int past_len = total_len - sequence_length; - past_seq_lens[i] = past_len > 0 ? past_len : 0; + past_len = total_len - sequence_length; + past_len = past_len > 0 ? past_len : 0; + past_seq_lens[i] = past_len; padded_seq_lens[i] = 0; } + + if (cache_past_seq_lens == nullptr) { + return; + } + + // Cache-relative coordinates for a windowed KV cache. The cache holds the L = min(T, C) most + // recent tokens contiguously at indices [0, L), where C is the real capacity. + // Lp = min(P, C) tokens currently resident + // E = max(0, Lp + S - C) tokens this append pushes out of the real cache + // D = max(0, Lp + S - C_used) tokens that have to be shifted out of the buffer in use + // A multi-token step writes into a staging buffer long enough to hold Lp + S entries, so there + // D is 0 and E is only the offset of the window inside the staging buffer. For a single-token + // step the staging buffer is the cache itself (C_used == C) and D == E. + // + // For single-token decode steps the drifting-layout optimisation applies: entries live at + // [0, end) and are appended at `end`; compaction reclaims a whole block of `gap = C - W + 1` + // rows at once, so the kernel is load-bearing on only one step in `gap` and on the remaining + // steps the d <= 0 early-exit fires. Note that `end` is still a pure function of the absolute + // past length P, so the op stays stateless across Run() calls. + // For multi-token steps the staging buffer size must be predictable from Lp = min(P, C), and + // the write-back offset is the number of evicted rows under the old formula, so we keep the + // original Lp-based computation for that case. + const int C = kv_cache_real_capacity; + const int C_used = kv_cache_capacity; + const int W = C < C_used ? C : C_used; // == local_window_size when C == C_used + const int gap = C - W + 1; // >= 1; at C == W (no slack) this is 1 and every step compacts + + int E, D; + int Lp; + if (sequence_length == 1) { + // Drifting-layout end pointer for single-token steps; see WindowedCacheEnd above. + // This keeps W <= end <= C once the cache has filled. + const int end_before = WindowedCacheEnd(past_len, C, gap); + const int end_after = WindowedCacheEnd(static_cast(past_len) + 1, C, gap); + const int kept = end_after - 1; + E = end_before - kept > 0 ? end_before - kept : 0; + D = (C_used < C) ? 0 : E; + Lp = end_before; + } else { + // Original formula for multi-token (staging) steps: Lp = min(P, C), E = evicted rows. + Lp = past_len < C ? past_len : C; + const int evicted = Lp + sequence_length - C; + const int shifted = Lp + sequence_length - C_used; + E = evicted > 0 ? evicted : 0; + D = shifted > 0 ? shifted : 0; + } + cache_past_seq_lens[i] = Lp - D; + cache_total_seq_lens[i] = (Lp + sequence_length) < C_used ? (Lp + sequence_length) : C_used; + evict_counts[i] = E; } } @@ -663,17 +751,242 @@ Status LaunchGetSequenceLengths( int* past_seq_lens, int* total_seq_lens, int* padded_seq_lens, + int* cache_past_seq_lens, + int* cache_total_seq_lens, + int* evict_counts, const int batch_size, const int sequence_length, const bool is_first_prompt, + const int kv_cache_capacity, + const int kv_cache_real_capacity, cudaStream_t stream, const int max_threads_per_block) { int blocks = (batch_size + max_threads_per_block - 1) / max_threads_per_block; - GetSequenceLengths<<>>(total_seq_lens_minus_one, past_seq_lens, total_seq_lens, padded_seq_lens, batch_size, sequence_length, is_first_prompt); + GetSequenceLengths<<>>( + total_seq_lens_minus_one, past_seq_lens, total_seq_lens, padded_seq_lens, + cache_past_seq_lens, cache_total_seq_lens, evict_counts, + batch_size, sequence_length, is_first_prompt, kv_cache_capacity, kv_cache_real_capacity); + return CUDA_CALL(cudaGetLastError()); +} + +// ============================================================================ +// Windowed KV cache compaction +// ============================================================================ +// Left-shifts a BNSH KV cache by D = evict_counts[b] positions so that the retained tokens stay +// contiguous at cache indices [0, Lp - D). D is read on the device (it changes every step once the +// cache is full), while the launch geometry is sized for the worst case and threads whose source +// row is out of range exit early. Keeping the grid shape independent of D is what allows the +// enclosing graph to remain capturable. +// +// A direct in-place left shift races (source and destination rows overlap across blocks), so the +// shift runs as copy-to-scratch followed by copy-back. Both kernels operate on whole rows of the +// raw storage type, which makes them quantization agnostic (fp16/bf16/int8/fp8, and int4 as long +// as the row size is expressed in packed bytes). +// +// row_bytes: number of bytes occupied by one (batch, head, position) entry. +// capacity : allocated sequence dimension C of the cache. +// ============================================================================ + +// Shared launch geometry for the two row-copy kernels below. threadIdx.x walks the uint4 chunks of +// a row and threadIdx.y walks rows, so a warp covers several adjacent rows in full and the accesses +// stay coalesced. Handing a block many rows keeps the grid to a few dozen blocks, which matters +// because these copies are launch bound rather than bandwidth bound: with gap = C - W + 1, +// evict_counts[b] is non-zero on only one step in `gap`, so nearly every launch hits the d <= 0 +// early-exit and its cost is dominated by grid scheduling rather than memory traffic. +static void KvRowCopyLaunchConfig(const int vec_count, + const int rows, + const int kv_num_heads, + const int batch_size, + dim3& grid, + dim3& block) { + constexpr int kThreadsPerBlock = 256; + const int threads_x = vec_count < 32 ? vec_count : 32; + const int threads_y = kThreadsPerBlock / threads_x > 0 ? kThreadsPerBlock / threads_x : 1; + block = dim3(threads_x, threads_y); + grid = dim3((rows + threads_y - 1) / threads_y, kv_num_heads, batch_size); +} + +// K and V shift in the same launch: they share every index this kernel computes, so fusing them +// halves the number of launches per layer without adding a single block. +__global__ void CompactKvCacheRows(uint4* __restrict__ dst_k, + const uint4* __restrict__ src_k, + uint4* __restrict__ dst_v, + const uint4* __restrict__ src_v, + const int* __restrict__ evict_counts, + const int* __restrict__ cache_past_seq_lens, + const int capacity, + const int vec_count, + const bool copy_v, + const bool src_is_shifted) { + // grid: (ceil(capacity / blockDim.y), kv_num_heads, batch_size); block: (row chunks, rows). + const int b = blockIdx.z; + const int d = evict_counts[b]; + if (d <= 0) { + return; + } + const int s = blockIdx.x * blockDim.y + threadIdx.y; + if (s >= cache_past_seq_lens[b]) { // cache_past_seq_lens[b] is the retained length Lp - D + return; + } + + const int64_t head_stride = static_cast(capacity) * vec_count; + const int64_t base = (static_cast(b) * gridDim.y + blockIdx.y) * head_stride; + + // Pass 1 (src_is_shifted): scratch[s] <- cache[s + D]. + // Pass 2: cache[s] <- scratch[s]. + const int64_t dst_row = base + static_cast(s) * vec_count; + const int64_t src_row = base + static_cast(src_is_shifted ? (s + d) : s) * vec_count; + + for (int c = threadIdx.x; c < vec_count; c += blockDim.x) { + dst_k[dst_row + c] = src_k[src_row + c]; + if (copy_v) { + dst_v[dst_row + c] = src_v[src_row + c]; + } + } +} + +// Evicts the oldest evict_counts[b] entries from a windowed KV cache by shifting the retained +// entries to the front. Must run before the new tokens are appended. +Status LaunchCompactKvCache(void* k_cache, + void* v_cache, + void* scratch, + const int* evict_counts, + const int* cache_past_seq_lens, + const int batch_size, + const int kv_num_heads, + const int capacity, + const int row_bytes, + cudaStream_t stream) { + if (capacity <= 0 || row_bytes <= 0) { + return Status::OK(); + } + if (row_bytes % static_cast(sizeof(uint4)) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "sliding_window_cache: KV cache row size (", row_bytes, + " bytes) must be a multiple of 16 for the compaction shift."); + } + + // Order the two caches so the fused kernel always has a valid first operand; the second one is + // optional, so a cache pair with only one side present still works. + uint8_t* cache_a = reinterpret_cast(k_cache != nullptr ? k_cache : v_cache); + if (cache_a == nullptr) { + return Status::OK(); + } + uint8_t* cache_b = reinterpret_cast((k_cache != nullptr && v_cache != nullptr) ? v_cache : nullptr); + + const size_t cache_bytes = static_cast(batch_size) * kv_num_heads * capacity * row_bytes; + ORT_RETURN_IF(scratch == nullptr && cache_bytes > 0, + "sliding_window_cache: compaction scratch buffer was not allocated."); + // The two caches cannot share one scratch region: their passes now run concurrently. The buffer + // is allocated with room for both. + uint8_t* scratch_a = reinterpret_cast(scratch); + uint8_t* scratch_b = scratch_a + cache_bytes; + + const int vec_count = row_bytes / static_cast(sizeof(uint4)); + dim3 grid, block; + KvRowCopyLaunchConfig(vec_count, capacity, kv_num_heads, batch_size, grid, block); + + CompactKvCacheRows<<>>( + reinterpret_cast(scratch_a), reinterpret_cast(cache_a), + reinterpret_cast(scratch_b), reinterpret_cast(cache_b), + evict_counts, cache_past_seq_lens, capacity, vec_count, + /*copy_v=*/cache_b != nullptr, /*src_is_shifted=*/true); + CompactKvCacheRows<<>>( + reinterpret_cast(cache_a), reinterpret_cast(scratch_a), + reinterpret_cast(cache_b), reinterpret_cast(scratch_b), + evict_counts, cache_past_seq_lens, capacity, vec_count, + /*copy_v=*/cache_b != nullptr, /*src_is_shifted=*/false); + return CUDA_CALL(cudaGetLastError()); +} + +// ============================================================================ +// Windowed KV cache: copy a window of rows between two BNSH KV caches. +// +// A step with more than one new token runs against a staging cache long enough to hold the +// resident history plus the new tokens, so that queries can still reach every key their window +// covers. This kernel both seeds that staging cache from the real one and copies the surviving +// tail back afterwards. src_offsets is a device vector (one entry per batch) because the offset of +// the surviving tail depends on the past length, which is only known on the device; passing it +// through memory keeps the launch geometry constant and therefore CUDA-graph capturable. +// The copy works on raw rows, which keeps it independent of the cache element type and of whether +// the cache is written in quantized form. K and V share every index, so they move together. +// ============================================================================ +__global__ void CopyKvCacheRows(uint4* __restrict__ dst_k, + const uint4* __restrict__ src_k, + uint4* __restrict__ dst_v, + const uint4* __restrict__ src_v, + const int* __restrict__ src_offsets, + const int src_capacity, + const int dst_capacity, + const int rows, + const int vec_count, + const bool copy_v) { + // grid: (ceil(rows / blockDim.y), kv_num_heads, batch_size); block: (row chunks, rows). + const int b = blockIdx.z; + const int s = blockIdx.x * blockDim.y + threadIdx.y; + if (s >= rows || s >= dst_capacity) { + return; + } + const int src_index = s + (src_offsets == nullptr ? 0 : src_offsets[b]); + if (src_index < 0 || src_index >= src_capacity) { + return; + } + + const int64_t head = static_cast(b) * gridDim.y + blockIdx.y; + const int64_t src_row = (head * src_capacity + src_index) * vec_count; + const int64_t dst_row = (head * dst_capacity + s) * vec_count; + + for (int c = threadIdx.x; c < vec_count; c += blockDim.x) { + dst_k[dst_row + c] = src_k[src_row + c]; + if (copy_v) { + dst_v[dst_row + c] = src_v[src_row + c]; + } + } +} + +Status LaunchCopyKvCacheWindow(void* dst_k, + void* dst_v, + const void* src_k, + const void* src_v, + const int batch_size, + const int kv_num_heads, + const int src_capacity, + const int dst_capacity, + const int rows, + const int* src_offsets, + const int row_bytes, + cudaStream_t stream) { + if (rows <= 0 || row_bytes <= 0) { + return Status::OK(); + } + if (row_bytes % static_cast(sizeof(uint4)) != 0) { + return ORT_MAKE_STATUS(ONNXRUNTIME, NOT_IMPLEMENTED, + "sliding_window_cache: KV cache row size (", row_bytes, + " bytes) must be a multiple of 16 to stage the KV cache."); + } + + // Order the pairs so the first operand is always valid; the second one is optional. + const bool has_k = src_k != nullptr && dst_k != nullptr; + const bool has_v = src_v != nullptr && dst_v != nullptr; + if (!has_k && !has_v) { + return Status::OK(); + } + const void* src_a = has_k ? src_k : src_v; + void* dst_a = has_k ? dst_k : dst_v; + const void* src_b = (has_k && has_v) ? src_v : nullptr; + void* dst_b = (has_k && has_v) ? dst_v : nullptr; + + const int vec_count = row_bytes / static_cast(sizeof(uint4)); + dim3 grid, block; + KvRowCopyLaunchConfig(vec_count, rows, kv_num_heads, batch_size, grid, block); + + CopyKvCacheRows<<>>( + reinterpret_cast(dst_a), reinterpret_cast(src_a), + reinterpret_cast(dst_b), reinterpret_cast(src_b), + src_offsets, src_capacity, dst_capacity, rows, vec_count, /*copy_v=*/src_b != nullptr); return CUDA_CALL(cudaGetLastError()); } -// Trace function for debugging #define ORT_GQA_TRACE(func_name) \ DUMP_PRINTF("[GQA %s] is_packed_qkv: %d, is_first_prompt: %d, is_subsequent_prompt: %d, past_present_share_buffer: %d", \ func_name, \ @@ -746,8 +1059,10 @@ Status ExtremeDecoding( head_size, sequence_length, batch_size, - parameters.seqlen_present_kv_cache, // max_seqlen (capacity) + RopeMaxPosition(parameters), + parameters.seqlen_present_kv_cache, // cache capacity C data.past_seq_lens, + data.cache_past_seq_lens, reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), parameters.do_rotary ? parameters.rotary_dim : 0, @@ -788,7 +1103,8 @@ Status ExtremeDecoding( scale, parameters.local_window_size, // -1 means global attention; >0 enables sliding window past_bsnh, - data.past_seq_lens, + // With a windowed cache the XQA kernel must index the cache in cache-relative coordinates. + parameters.is_windowed_kv_cache ? data.cache_past_seq_lens : data.past_seq_lens, data.xqa_head_sink, fold_k_scale_into_q ? nullptr : data.k_scale, // k_cache_scale (null: already folded into Q) fold_v_scale_into_output ? nullptr : data.v_scale, // v_cache_scale (null: applied to the output below) @@ -925,7 +1241,11 @@ Status FlashAttention( void* kernel_new_v = nullptr; // Use padded seq lens for first prompt since mha_fwd_kvcache assumes uniform seqlen_q. - int* seq_lens = parameters.is_first_prompt ? data.padded_seq_lens : data.total_seq_lens; + // A windowed cache is indexed in cache-relative coordinates, so the attention length is the + // number of resident entries rather than the absolute total sequence length. + int* seq_lens = parameters.is_windowed_kv_cache + ? data.cache_total_seq_lens + : (parameters.is_first_prompt ? data.padded_seq_lens : data.total_seq_lens); // After PrepareQKV, the input for flash attention is no longer packed. constexpr bool is_packed_qkv_for_flash = false; @@ -987,7 +1307,8 @@ Status DequantizeFlashAttentionFallback( q_rot, reinterpret_cast(data.present_key), reinterpret_cast(data.present_value), data.k_scale, data.v_scale, nullptr /* q_fold_scale: this path dequantizes the cache itself */, parameters.num_heads, parameters.kv_num_heads, parameters.head_size, parameters.sequence_length, parameters.batch_size, - parameters.seqlen_present_kv_cache, data.past_seq_lens, + RopeMaxPosition(parameters), parameters.seqlen_present_kv_cache, + data.past_seq_lens, data.cache_past_seq_lens, reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, (parameters.past_kv_format == AttentionQkvFormat::Q_K_V_BNSH), @@ -1019,7 +1340,10 @@ Status DequantizeFlashAttentionFallback( bool is_bf16 = std::is_same::value; // Use the total_seq_lens here since k_dequant/v_dequant has both past and new tokens. - void* seqlens_k_ptr = const_cast(reinterpret_cast(data.total_seq_lens)); + // With a windowed cache the dequantized buffer is capacity-relative, so use the cache-relative + // lengths; the causal / local-window masks only depend on relative distances, so this is exact. + void* seqlens_k_ptr = const_cast(reinterpret_cast( + parameters.is_windowed_kv_cache ? data.cache_total_seq_lens : data.total_seq_lens)); int local_window_size = parameters.local_window_size > 0 ? parameters.local_window_size - 1 : -1; ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_fwd_kvcache( @@ -1077,7 +1401,8 @@ Status FlashAttentionAndQuantizeKV( parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.value), q_final, k_final, v_final, nullptr, nullptr, nullptr /* q_fold_scale */, num_heads, kv_num_heads, head_size, sequence_length, batch_size, - sequence_length, data.past_seq_lens, + RopeMaxPosition(parameters), sequence_length, + data.past_seq_lens, data.past_seq_lens, reinterpret_cast(data.cos_cache), reinterpret_cast(data.sin_cache), parameters.rotary_dim, data.position_ids, parameters.rotary_interleaved, false, // BSNH for scratch @@ -1118,6 +1443,7 @@ Status FlashAttentionAndQuantizeKV( return Status::OK(); } + #endif #if USE_MEMORY_EFFICIENT_ATTENTION @@ -1180,7 +1506,9 @@ Status EfficientAttention( p.causal = true; p.scale = scale; p.softcap = parameters.softcap; - p.seqlen_k_ptr = parameters.is_first_prompt ? data.padded_seq_lens : data.total_seq_lens; + p.seqlen_k_ptr = parameters.is_windowed_kv_cache + ? data.cache_total_seq_lens + : (parameters.is_first_prompt ? data.padded_seq_lens : data.total_seq_lens); p.seqstart_q_ptr = nullptr; p.seqstart_k_ptr = nullptr; p.query = query; @@ -1291,6 +1619,15 @@ Status UnfusedGqaAttention( p.softcap = parameters.softcap; p.seqlens_k = data.total_seq_lens; + if (parameters.is_windowed_kv_cache) { + // The cache holds only the most recent min(T, C) tokens, contiguously from index 0. Causal and + // local-window masks depend solely on query/key distance, so shifting to cache-relative + // coordinates leaves them unchanged. + p.total_kv_length = std::min(parameters.total_sequence_length, max_kv); + p.past_kv_length = p.total_kv_length - parameters.sequence_length; + p.seqlens_k = data.cache_total_seq_lens; + } + ORT_RETURN_IF_ERROR((LaunchUnfusedAttention( device_prop, cublas, stream, p, data.unfused_q_bnsh, @@ -1403,6 +1740,21 @@ Status QkvToContext( GroupQueryAttentionData& data) { auto stream = static_cast(ort_stream->GetHandle()); const float scale = parameters.scale == 0.0f ? 1.f / sqrt(static_cast(parameters.head_size)) : parameters.scale; + + if (parameters.is_windowed_kv_cache && parameters.kv_cache_capacity == parameters.kv_cache_real_capacity) { + // Evict the oldest rows first so the retained window is contiguous at [0, L) and the new + // tokens can simply be appended after it. Skipped when ComputeInternal has redirected this + // step through a staging cache: there the buffer is long enough that nothing is pushed out. + const int row_bytes = (parameters.kv_cache_bit_width == 0) + ? static_cast(parameters.head_size * sizeof(U)) + : parameters.head_size * parameters.kv_cache_bit_width / 8; + ORT_RETURN_IF_ERROR(LaunchCompactKvCache( + data.present_key, data.present_value, data.compaction_scratch, + data.evict_counts, data.cache_past_seq_lens, + parameters.batch_size, parameters.kv_num_heads, parameters.seqlen_present_kv_cache, + row_bytes, stream)); + } + if (data.use_xqa) { return ExtremeDecoding(device_prop, stream, parameters, data, scale); } diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h index f7565a942c85e..59858c83dc171 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h @@ -154,12 +154,48 @@ Status LaunchGetSequenceLengths( int* past_seq_lens, int* total_seq_lens, int* padded_seq_lens, + int* cache_past_seq_lens, + int* cache_total_seq_lens, + int* evict_counts, const int batch_size, const int sequence_length, const bool is_first_prompt, + const int kv_cache_capacity, + const int kv_cache_real_capacity, cudaStream_t stream, const int max_threads_per_block); +// Evicts the oldest evict_counts[b] entries of a windowed (sliding_window_cache) KV cache by +// left-shifting the retained entries to the front, so the cache keeps holding the most recent +// tokens contiguously at indices [0, L). Must run before the new tokens are appended. +Status LaunchCompactKvCache(void* k_cache, + void* v_cache, + void* scratch, + const int* evict_counts, + const int* cache_past_seq_lens, + const int batch_size, + const int kv_num_heads, + const int capacity, + const int row_bytes, + cudaStream_t stream); + +// Copies `rows` entries between two BNSH KV caches, reading from row src_offsets[b] + s (or from +// row s when src_offsets is null) and writing to row s. Used to seed the staging cache that a +// multi-token step of a windowed (sliding_window_cache) KV cache runs against, and to copy the +// surviving window back into the real cache afterwards. +Status LaunchCopyKvCacheWindow(void* dst_k, + void* dst_v, + const void* src_k, + const void* src_v, + const int batch_size, + const int kv_num_heads, + const int src_capacity, + const int dst_capacity, + const int rows, + const int* src_offsets, + const int row_bytes, + cudaStream_t stream); + } // namespace cuda } // namespace contrib } // namespace onnxruntime diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh index 4f9f9aa8e5df8..6655ccd467574 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qdq.cuh @@ -430,7 +430,7 @@ __global__ void ScaleHeadsByChannelScaleKernel(T* __restrict__ dst, } const int h = static_cast(i / head_size) % num_heads; // q head, layout is [..., num_heads, head_size] const int d = static_cast(i % head_size); - const float scale = channel_scale[(h / group_size) * head_size + d]; // scale is [kv_num_heads, head_size] + const float scale = channel_scale[(h / group_size) * head_size + d]; // scale is [kv_num_heads, 1, head_size] dst[i] = static_cast(static_cast(src[i]) * scale); } diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh index 8b0e72f770693..c72d5f9093790 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh @@ -50,17 +50,28 @@ __global__ void UnpackRoPEAppend( U* v_cache, const float* k_scale, const float* v_scale, - // Per-channel K dequantization scale of shape (kv_num_heads, head_size) to fold into the stored - // Q, or nullptr. Used by the XQA per-channel decode path: XQA only accepts a scalar dequant - // scale, and because the channel dim is contracted by Q*K.T, scaling Q by the K scale is exactly - // equivalent to dequantizing K. Folding it here is free -- this kernel already loads and stores Q. + // Per-channel K dequantization scale of shape (kv_num_heads, 1, head_size) to fold into the + // stored Q, or nullptr. The indexing flattens the singleton middle dimension. Used by the XQA + // per-channel decode path: XQA only accepts a scalar dequant scale, and because the channel dim + // is contracted by Q*K.T, scaling Q by the K scale is exactly equivalent to dequantizing K. + // Folding it here is free -- this kernel already loads and stores Q. const float* q_fold_scale, const int num_heads, const int kv_num_heads, const int head_size, - const int d, // packed QKV hidden stride = (num_heads + 2*kv_num_heads) * head_size - const int max_seqlen, // KV cache max sequence length - const int* past_seq_lens, + const int d, // packed QKV hidden stride = (num_heads + 2*kv_num_heads) * head_size + // RoPE position bound: number of valid entries in cos_cache/sin_cache. This is always an + // ABSOLUTE position limit and must not be conflated with the KV cache capacity below. For a + // windowed (shorter-than-total) cache the two differ, and using the capacity here would + // silently skip RoPE for every absolute position >= capacity -- an accuracy bug, not a crash. + const int rope_max_pos, + // KV cache capacity in positions: bounds the cache write index and defines the sequence + // stride. Equal to rope_max_pos for a full-length cache. + const int cache_capacity, + const int* past_seq_lens, // absolute past lengths, used for RoPE positions only + // Cache-relative append offsets. Equals past_seq_lens for a full-length cache; for a windowed + // (sliding_window_cache) layer it is the post-eviction offset within the capacity-C buffer. + const int* cache_past_seq_lens, const T* cos_cache, const T* sin_cache, const int rotary_dim, @@ -207,10 +218,10 @@ __global__ void UnpackRoPEAppend( int64_t pos_val = (position_ids != nullptr) ? position_ids[pos_base + s] : static_cast(past_seq_len + s); #if !defined(NDEBUG) if (tid == 0) { - CUDA_KERNEL_ASSERT(pos_val >= 0 && pos_val < static_cast(max_seqlen)); + CUDA_KERNEL_ASSERT(pos_val >= 0 && pos_val < static_cast(rope_max_pos)); } #endif - if (pos_val >= 0 && pos_val < static_cast(max_seqlen)) { + if (pos_val >= 0 && pos_val < static_cast(rope_max_pos)) { int pos_id = static_cast(pos_val); const int h_idx = h / elements_per_thread; @@ -244,11 +255,11 @@ __global__ void UnpackRoPEAppend( reinterpret_cast(unpacked_q)[q_out_idx / elements_per_thread] = *reinterpret_cast(vals); } } else { - // Store K or V into the KV cache at index (past_seqlen + s) - const int cache_s = past_seq_lens[b] + s; + // Store K or V into the KV cache at the cache-relative index (cache_past_seqlen + s). + const int cache_s = cache_past_seq_lens[b] + s; // Two-sided bound: the lower check mirrors the position guard above and prevents a // negative offset from being sign-extended into the cache index arithmetic below. - if (cache_s >= 0 && cache_s < max_seqlen) { + if (cache_s >= 0 && cache_s < cache_capacity) { void* cache_ptr = (head_type == KEY) ? k_cache : v_cache; if (cache_ptr != nullptr) { int64_t cache_idx; @@ -257,7 +268,7 @@ __global__ void UnpackRoPEAppend( // Note: For BIT_WIDTH=4, head_size refers to the number of UNPACKED elements. // stride_s is the number of bytes occupied by head_size elements. const int64_t stride_s = (BIT_WIDTH == 4) ? (head_size / 2) : head_size; - const int64_t stride_n = max_seqlen * stride_s; + const int64_t stride_n = cache_capacity * stride_s; const int64_t stride_b = kv_num_heads * stride_n; cache_idx = static_cast(b) * stride_b + static_cast(n) * stride_n + @@ -267,7 +278,7 @@ __global__ void UnpackRoPEAppend( // BSNH layout: [Batch, SeqLen, NumHeads, HeadSize] const int64_t stride_n = (BIT_WIDTH == 4) ? (head_size / 2) : head_size; const int64_t stride_s = kv_num_heads * stride_n; - const int64_t stride_b = max_seqlen * stride_s; + const int64_t stride_b = cache_capacity * stride_s; cache_idx = static_cast(b) * stride_b + static_cast(cache_s) * stride_s + static_cast(n) * stride_n + @@ -337,32 +348,33 @@ Status DispatchUnpackRoPEAppendHeadSize( T* unpacked_q, U* k_cache, U* v_cache, const float* k_scale, const float* v_scale, const float* q_fold_scale, const int num_heads, const int kv_num_heads, const int head_size, const int d, - const int max_seqlen, const int* past_seq_lens, + const int rope_max_pos, const int cache_capacity, + const int* past_seq_lens, const int* cache_past_seq_lens, const T* cos_cache, const T* sin_cache, const int rotary_dim, const int64_t* position_ids, const bool interleaved, const bool is_cache_bnsh, const bool per_channel, const T* q_norm_weight, const T* k_norm_weight, const float qk_norm_epsilon) { if (head_size <= 64) { UnpackRoPEAppend<<>>( packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, k_scale, v_scale, q_fold_scale, - num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); } else if (head_size <= 128) { UnpackRoPEAppend<<>>( packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, k_scale, v_scale, q_fold_scale, - num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); } else if (head_size <= 256) { UnpackRoPEAppend<<>>( packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, k_scale, v_scale, q_fold_scale, - num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); } else if (head_size <= 512) { UnpackRoPEAppend<<>>( packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, k_scale, v_scale, q_fold_scale, - num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); } else { @@ -382,8 +394,10 @@ Status LaunchUnpackRoPEAppend( T* unpacked_q, U* k_cache, U* v_cache, const float* k_scale, const float* v_scale, const float* q_fold_scale, const int num_heads, const int kv_num_heads, const int head_size, - const int sequence_length, const int batch_size, const int max_seqlen, - const int* past_seq_lens, const T* cos_cache, const T* sin_cache, + const int sequence_length, const int batch_size, + const int rope_max_pos, const int cache_capacity, + const int* past_seq_lens, const int* cache_past_seq_lens, + const T* cos_cache, const T* sin_cache, const int rotary_dim, const int64_t* position_ids, const bool interleaved, const bool is_cache_bnsh, const KVQuantizationType k_quant_type, const T* q_norm_weight, const T* k_norm_weight, const float qk_norm_epsilon, @@ -391,6 +405,11 @@ Status LaunchUnpackRoPEAppend( static_assert(std::is_same::type>::value); static_assert(std::is_same::type>::value); + // A full-length (non-windowed) cache passes the same buffer for both; keep the call sites simple. + if (cache_past_seq_lens == nullptr) { + cache_past_seq_lens = past_seq_lens; + } + constexpr int elements_per_vector = sizeof(float4) / sizeof(T); if (head_size % elements_per_vector != 0) { @@ -433,7 +452,8 @@ Status LaunchUnpackRoPEAppend( // No quantization: cache type same as input type return DispatchUnpackRoPEAppendHeadSize( grid, block, stream, packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, - k_scale, v_scale, q_fold_scale, num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + k_scale, v_scale, q_fold_scale, num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, + past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); } else if constexpr (std::is_same::value @@ -444,7 +464,8 @@ Status LaunchUnpackRoPEAppend( // INT8 or FP8 quantization (both 8-bit, distinguished inside kernel by type check) return DispatchUnpackRoPEAppendHeadSize( grid, block, stream, packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, - k_scale, v_scale, q_fold_scale, num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + k_scale, v_scale, q_fold_scale, num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, + past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); #ifdef USE_INT4_KV_CACHE @@ -452,7 +473,8 @@ Status LaunchUnpackRoPEAppend( // INT4 quantization (packed 2 elements per byte) return DispatchUnpackRoPEAppendHeadSize( grid, block, stream, packed_qkv, query, key, value, unpacked_q, k_cache, v_cache, - k_scale, v_scale, q_fold_scale, num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + k_scale, v_scale, q_fold_scale, num_heads, kv_num_heads, head_size, d, rope_max_pos, cache_capacity, + past_seq_lens, cache_past_seq_lens, cos_cache, sin_cache, rotary_dim, position_ids, interleaved, is_cache_bnsh, per_channel, q_norm_weight, k_norm_weight, qk_norm_epsilon); #endif diff --git a/onnxruntime/contrib_ops/js/bert/group_query_attention.h b/onnxruntime/contrib_ops/js/bert/group_query_attention.h index dff8663133c31..effe1551dbc87 100644 --- a/onnxruntime/contrib_ops/js/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/js/bert/group_query_attention.h @@ -15,6 +15,10 @@ class GroupQueryAttention : public JsKernel, GQAAttentionBase { public: explicit GroupQueryAttention(const OpKernelInfo& info) : JsKernel(info), GQAAttentionBase(info, false) { + // The windowed KV cache (cache-relative indexing + shift compaction) is implemented only by the + // CUDA kernel. Fail loudly rather than treating the window-sized buffer as a full-length cache. + ORT_ENFORCE(info.GetAttrOrDefault("sliding_window_cache", 0) == 0, + "GroupQueryAttention (JS): sliding_window_cache=1 is not implemented."); JSEP_INIT_KERNEL_ATTRIBUTE(GroupQueryAttention, ({ "numHeads" : $1, "kvNumHeads" : $2, diff --git a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h index cbb5b806eb6ad..8af15937001a2 100644 --- a/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h +++ b/onnxruntime/contrib_ops/webgpu/bert/group_query_attention.h @@ -59,6 +59,11 @@ class GroupQueryAttention final : public WebGpuKernel { local_window_size_ = static_cast(info.GetAttrOrDefault("local_window_size", -1)); + // The windowed KV cache (cache-relative indexing + shift compaction) is implemented only by the + // CUDA kernel. Fail loudly rather than treating the window-sized buffer as a full-length cache. + ORT_ENFORCE(info.GetAttrOrDefault("sliding_window_cache", 0) == 0, + "GroupQueryAttention (WebGPU): sliding_window_cache=1 is not implemented."); + qk_norm_epsilon_ = info.GetAttrOrDefault("qk_norm_epsilon", 1e-6f); } diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 9fdcebf57081d..f088fe226cfce 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -447,7 +447,12 @@ void BaseGroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceConte void GroupQueryAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_key_index, int qk_output_index) { // TODO(aciddelgado): propagate output shapes depending if kv-share buffer is on or not constexpr int use_max_past_present_buffer = -1; - BaseGroupQueryAttentionTypeAndShapeInference(ctx, past_key_index, use_max_past_present_buffer, qk_output_index); + // With a windowed (sliding_window_cache) KV cache, the bound past/present buffer is the cache + // capacity C, which is deliberately smaller than total_sequence_length. present therefore keeps + // the past buffer's own sequence dimension instead of growing with the total sequence length. + const int64_t sliding_window_cache = getAttribute(ctx, "sliding_window_cache", 0); + BaseGroupQueryAttentionTypeAndShapeInference( + ctx, past_key_index, sliding_window_cache == 1 ? 1 : use_max_past_present_buffer, qk_output_index); } void SparseAttentionTypeAndShapeInference(ONNX_NAMESPACE::InferenceContext& ctx, int past_key_index) { @@ -1231,6 +1236,15 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "left_window_size for local attention (like Mistral). Default value is -1 meaning unused.", AttributeProto::INT, static_cast(-1)) + .Attr("sliding_window_cache", + "Set to 1 when the past/present KV buffers are window-sized instead of holding the whole " + "sequence. The op then keeps only the min(total_sequence_length, cache_capacity) most recent " + "tokens, contiguously, using cache-relative indexing and evicting from the front as needed. " + "Requires local_window_size > 0 and a cache capacity of at least local_window_size. " + "Multi-token steps may use a temporary staging buffer, so the capacity need not cover the " + "entire step. Default value is 0 (full-length cache).", + AttributeProto::INT, + static_cast(0)) .Attr("do_rotary", "Whether to use rotary position embedding. Default value is 0.", AttributeProto::INT, diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py index d00f1fc5921c3..df0209962cd4a 100644 --- a/onnxruntime/test/python/transformers/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_gqa.py @@ -17,6 +17,7 @@ import re import sys import threading +import typing import unittest from copy import deepcopy from dataclasses import dataclass @@ -139,6 +140,16 @@ class GQAConfig: buffer_sequence_length: int = 0 # Test-specific parameters local_window_size: int = -1 + # Opt into the cache-relative (windowed) KV cache. Requires local_window_size > 0 and a + # past/present shared buffer whose capacity C is at least local_window_size. There is no + # requirement involving q_sequence_length: a multi-token step that would need more than C + # entries runs against a longer staging buffer, so prefill of any length works at C == W. + # Slack above local_window_size is worth allocating anyway: it lets the append point drift, so + # the CPU kernel compacts once every C - W + 1 steps instead of on every step. + sliding_window_cache: int = 0 + # Length of the cos/sin caches. 0 means "same as buffer_sequence_length". A windowed KV cache is + # shorter than the sequence, so RoPE positions must be driven by a separately sized rotary cache. + rope_cache_length: int = 0 rotary: bool = False rotary_interleaved: bool = False packed: bool = False @@ -368,6 +379,8 @@ def create_gqa_node_and_io( qk_norm_attributes = {"qk_norm_epsilon": config.qk_norm_epsilon} if config.has_qk_norm else {} + windowed_cache_attributes = {"sliding_window_cache": 1} if config.sliding_window_cache else {} + node = helper.make_node( op_type="GroupQueryAttention", inputs=inputs, @@ -383,6 +396,7 @@ def create_gqa_node_and_io( qk_output=output_qk, **quantization_attributes, **qk_norm_attributes, + **windowed_cache_attributes, domain="com.microsoft", ) @@ -433,7 +447,7 @@ def create_gqa_node_and_io( if config.rotary: rotary_dim = (math.floor(config.head_size / 16) * 16) // 2 - cache_seq_len = config.buffer_sequence_length + cache_seq_len = config.rope_cache_length or config.buffer_sequence_length graph_input.extend( [ helper.make_tensor_value_info("cos_cache", ort_type, [cache_seq_len, rotary_dim]), @@ -2199,6 +2213,11 @@ def has_xqa(): return has_cuda_device(80) +def has_fp8_xqa(): + # FP8 XQA kernels require Hopper (SM 8.9) or newer. + return has_cuda_device(89) + + rtol = { "fp16": 5e-3, "bf16": 5e-2, @@ -2688,6 +2707,35 @@ def test_gqa_past_attention_bias(self, name, config): atol=atol["fp16"], ) + def test_windowed_cache_attention_bias_is_rejected(self): + config = GQAConfig( + batch_size=1, + q_sequence_length=1, + kv_sequence_length=1, + past_kv_sequence_length=128, + buffer_sequence_length=129, + num_heads=8, + kv_num_heads=2, + head_size=64, + local_window_size=128, + share_buffer=True, + sliding_window_cache=1, + has_attention_bias=True, + attention_bias_per_head=True, + ) + + with self.assertRaisesRegex(Exception, "attention_bias with sliding_window_cache"): + 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"], + ) + @unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.") class TestFlashGQAPaddingPrompt(unittest.TestCase): @@ -3041,9 +3089,9 @@ 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 covered separately by + # well. Quantized XQA selection requires: decode (q_seq=1), a shared KV buffer, supported + # head_size and group_size combinations for the active kernel configuration. Attention sinks + # are covered separately by # gqa_xqa_quantized_head_sink_test_cases(), so no head_sink here. # # Two window/past relationships are covered: @@ -3051,7 +3099,7 @@ def gqa_xqa_quantized_sliding_window_test_cases(): # 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: + if has_fp8_kv_cache and has_fp8_xqa(): 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)]: @@ -3823,5 +3871,369 @@ def test_gqa_gemma4_local_window_decode_fp16(self): ) +# ################################################################################################# +# Windowed (cache-relative) KV cache: sliding_window_cache=1 +# ################################################################################################# + + +def _windowed_make_session(config: GQAConfig, ort_type, providers=None): + onnx_model_str = create_group_query_attention_graph_past(config, ort_type=ort_type, share_buffer=True) + if providers is None: + providers = [resolve_cuda_plugin_ep("CUDAExecutionProvider")] + return InferenceSession(onnx_model_str, SessionOptions(), providers=providers) + + +def _windowed_run_steps( + base_config: GQAConfig, + buffer_sequence_length: int, + sliding_window_cache: int, + step_lengths, + q_all, + k_all, + v_all, + cos, + sin, + device, + ort_type, + torch_type, + head_sink=None, + k_scale=None, + v_scale=None, + providers=None, +): + """Drives a GroupQueryAttention node token-chunk by token-chunk over a shared past/present buffer. + + `q_all`/`k_all`/`v_all` hold the whole sequence in BSNH layout, so the same inputs can be replayed + against a full-length cache and against a windowed one. Returns the per-step `output` tensors. + """ + batch_size = base_config.batch_size + kv_hidden_size = base_config.kv_num_heads * base_config.head_size + q_hidden_size = base_config.num_heads * base_config.head_size + + cache_torch_type = TORCH_DTYPE_MAP[base_config.kv_cache_type] if base_config.kv_cache_type else torch_type + cache_ort_type = ONNX_TENSOR_TYPE_MAP[base_config.kv_cache_type] if base_config.kv_cache_type else ort_type + dense_head_size = base_config.head_size // 2 if base_config.kv_cache_type == "int4" else base_config.head_size + cache_shape = (batch_size, base_config.kv_num_heads, buffer_sequence_length, dense_head_size) + cache_k = torch.zeros(cache_shape, dtype=cache_torch_type, device=device) + cache_v = torch.zeros(cache_shape, dtype=cache_torch_type, device=device) + + sessions = {} + outputs = [] + past_length = 0 + + for step_length in step_lengths: + config = deepcopy(base_config) + config.q_sequence_length = step_length + config.kv_sequence_length = step_length + config.past_kv_sequence_length = past_length + config.buffer_sequence_length = buffer_sequence_length + config.sliding_window_cache = sliding_window_cache + + if step_length not in sessions: + sessions[step_length] = _windowed_make_session(config, ort_type, providers=providers) + session = sessions[step_length] + + io_binding = session.io_binding() + + q = q_all[:, past_length : past_length + step_length].reshape(batch_size, step_length, q_hidden_size) + k = k_all[:, past_length : past_length + step_length].reshape(batch_size, step_length, kv_hidden_size) + v = v_all[:, past_length : past_length + step_length].reshape(batch_size, step_length, kv_hidden_size) + # Keep references: io_binding only records data pointers, so temporaries would be freed. + q, k, v = q.contiguous(), k.contiguous(), v.contiguous() + bind_tensor(io_binding, "query", q, device, ort_type) + bind_tensor(io_binding, "key", k, device, ort_type) + bind_tensor(io_binding, "value", v, device, ort_type) + bind_tensor(io_binding, "past_key", cache_k, device, cache_ort_type) + bind_tensor(io_binding, "past_value", cache_v, device, cache_ort_type) + + total_length = past_length + step_length + seqlens_k = torch.full((batch_size,), total_length - 1, dtype=torch.int32, device=device) + bind_tensor(io_binding, "seqlens_k", seqlens_k, device, TensorProto.INT32) + cpu_device = torch.device("cpu") + # Keep a reference: io_binding only records the data pointer, so a temporary would be freed. + total_sequence_length = torch.tensor([total_length], dtype=torch.int32, device=cpu_device) + bind_tensor( + io_binding, + "total_sequence_length", + total_sequence_length, + cpu_device, + TensorProto.INT32, + ) + + if cos is not None: + bind_tensor(io_binding, "cos_cache", cos, device, ort_type) + bind_tensor(io_binding, "sin_cache", sin, device, ort_type) + + if head_sink is not None: + bind_tensor(io_binding, "head_sink", head_sink, device, ort_type) + + if k_scale is not None: + bind_tensor(io_binding, "k_scale", k_scale, device, TensorProto.FLOAT) + bind_tensor(io_binding, "v_scale", v_scale, device, TensorProto.FLOAT) + + out = torch.zeros((batch_size, step_length, q_hidden_size), dtype=torch_type, device=device) + bind_output_tensor(io_binding, "output", out, device, ort_type) + io_binding.bind_output("present_key", device, 0, cache_ort_type, tuple(cache_k.shape), cache_k.data_ptr()) + io_binding.bind_output("present_value", device, 0, cache_ort_type, tuple(cache_v.shape), cache_v.data_ptr()) + + io_binding.synchronize_inputs() + session.run_with_iobinding(io_binding) + io_binding.synchronize_outputs() + + outputs.append(out.clone()) + past_length = total_length + + return outputs + + +class TestGQAWindowedKvCache(unittest.TestCase): + """`sliding_window_cache=1` stores only the most recent `capacity` positions of a sliding-window + layer, in cache-relative coordinates. Because RoPE stays absolute and the retained positions are + exactly those the window can attend to, results must match a full-length cache.""" + + max_length = 1024 + window_size = 128 + slack = 256 + + device = "cuda" + torch_type = torch.float16 + ort_type = TensorProto.FLOAT16 + providers = None # None -> resolve the CUDA (plugin) EP. + + def _require_ep(self): + if not has_flash_attention(): + self.skipTest("Flash Attention is not available, skipping test.") + + def _base_config(self, **overrides): + config = GQAConfig( + batch_size=2, + q_sequence_length=1, + kv_sequence_length=1, + num_heads=8, + kv_num_heads=2, + head_size=64, + local_window_size=self.window_size, + share_buffer=True, + rope_cache_length=self.max_length, + ) + for name, value in overrides.items(): + setattr(config, name, value) + return config + + def _check_parity(self, base_config, step_lengths, rtol=1e-3, atol=1e-3, buffer_sequence_length=None): + self._require_ep() + + device = self.device + torch_type = self.torch_type + ort_type = self.ort_type + capacity = buffer_sequence_length or min(self.max_length, self.window_size + self.slack) + total_length = sum(step_lengths) + self.assertLessEqual(total_length, self.max_length) + + torch.manual_seed(0) + shape = (base_config.batch_size, total_length, base_config.num_heads, base_config.head_size) + kv_shape = (base_config.batch_size, total_length, base_config.kv_num_heads, base_config.head_size) + q_all = torch.randn(shape, device=device, dtype=torch_type) * 0.2 + k_all = torch.randn(kv_shape, device=device, dtype=torch_type) * 0.2 + v_all = torch.randn(kv_shape, device=device, dtype=torch_type) * 0.2 + + cos, sin = None, None + if base_config.rotary: + rotary_dim = math.floor(base_config.head_size / 16) * 16 + angle = torch.rand(self.max_length, rotary_dim // 2, device=device) * 2 * math.pi + cos = torch.cos(angle).to(dtype=torch_type) + sin = torch.sin(angle).to(dtype=torch_type) + + head_sink = None + if base_config.has_head_sink: + head_sink = torch.rand(base_config.num_heads, dtype=torch_type, device=device) + + k_scale, v_scale = None, None + if base_config.k_quant_type != "NONE": + k_scale, v_scale = get_static_scale(base_config, device, torch_type, 0.2) + k_scale = k_scale.to(torch.float32).contiguous() + v_scale = v_scale.to(torch.float32).contiguous() + + common = { + "step_lengths": step_lengths, + "q_all": q_all, + "k_all": k_all, + "v_all": v_all, + "cos": cos, + "sin": sin, + "device": device, + "ort_type": ort_type, + "torch_type": torch_type, + "head_sink": head_sink, + "k_scale": k_scale, + "v_scale": v_scale, + "providers": self.providers, + } + reference = _windowed_run_steps(base_config, self.max_length, 0, **common) + if base_config.k_quant_type != "NONE" and any(torch.isnan(step).any() for step in reference): + # A quantized KV cache is only read correctly by the flash-attention prefill kernels. + # Builds without them route the first prompt through memory-efficient attention, which + # reinterprets the quantized cache as unquantized and produces NaN for both the + # windowed and the full-length run, so there is nothing to compare. + self.skipTest("quantized prefill needs a build with flash attention enabled") + windowed = _windowed_run_steps(base_config, capacity, 1, **common) + + for step_index, (expected, actual) in enumerate(zip(reference, windowed, strict=True)): + numpy.testing.assert_allclose( + actual.float().cpu().numpy(), + expected.float().cpu().numpy(), + rtol=rtol, + atol=atol, + err_msg=f"mismatch at step {step_index} (step_lengths={step_lengths})", + ) + + def test_prompt_shorter_than_capacity_then_decode(self): + # Prompt fits in the cache; eviction only starts once the window slides past the capacity. + steps = [64, *([1] * 400)] + self._check_parity(self._base_config(), steps) + + def test_prompt_longer_than_capacity_then_decode(self): + # Prompt exceeds the capacity, so the first call must attend over the full prompt while + # seeding the cache with only the tail. Exercises the windowed first-prompt path. + steps = [512, *([1] * 200)] + self._check_parity(self._base_config(), steps) + + def test_chunked_prefill(self): + # Chunks of 128 fit in the capacity (128 + 256) without staging. + # GroupQueryAttention only allows a subsequent prompt (1 < S < T) at batch_size 1. + steps = [128] * 6 + [1] * 32 + self._check_parity(self._base_config(batch_size=1), steps) + + def test_eviction_boundaries(self): + # Land exactly on capacity - 1, capacity and capacity + 1 with single-token steps. + capacity = self.window_size + self.slack + steps = [capacity - 2, *([1] * 8)] + self._check_parity(self._base_config(), steps) + + def test_with_rotary(self): + # RoPE positions stay absolute even though the cache is indexed relative to its start. + steps = [64, *([1] * 400)] + self._check_parity(self._base_config(rotary=True), steps) + + def test_gpt_oss_shape_with_head_sink(self): + # gpt-oss: attention sinks, RoPE, and alternating 128-token sliding-window layers. + steps = [512, *([1] * 200)] + self._check_parity( + self._base_config(num_heads=8, kv_num_heads=2, head_size=64, rotary=True, has_head_sink=True), steps + ) + + @unittest.skipIf(not has_quantized_kv_cache(), "Quantized KV cache is not available, skipping test.") + def test_int8_quantized_cache(self): + # gpt-oss stores the cache as PER_CHANNEL int8; eviction and staging must move the + # quantized rows around without dequantizing them. + steps = [512, *([1] * 200)] + self._check_parity( + self._base_config( + num_heads=8, + kv_num_heads=2, + head_size=64, + rotary=True, + has_head_sink=True, + kv_cache_type="int8", + k_quant_type="PER_CHANNEL", + v_quant_type="PER_CHANNEL", + kv_cache_bit_width=8, + ), + steps, + rtol=5e-2, + atol=1e-1, + ) + + def test_batch_one_packed_head_layout(self): + steps = [200, *([1] * 300)] + self._check_parity(self._base_config(batch_size=1, num_heads=4, kv_num_heads=4, head_size=128), steps) + + def test_capacity_equal_to_window_size(self): + # C == W leaves no slack, so a multi-token step needs more entries than the cache can hold. + # The op stages such a step, so this must still match a full cache exactly. + capacity = self.window_size + self._check_parity( + self._base_config(batch_size=1), step_lengths=[32, 16, 1, 1], buffer_sequence_length=capacity + ) + + def test_small_slack_many_compactions(self): + # C == W + 8 reclaims only 9 rows per compaction, so a long decode run crosses the + # compaction boundary dozens of times instead of once or twice. + self._check_parity( + self._base_config(), step_lengths=[64, *([1] * 400)], buffer_sequence_length=self.window_size + 8 + ) + + def test_chunked_prefill_small_slack(self): + # 32-token chunks onto a cache with 9 free rows: once the cache has filled, every chunk + # drops rows its own first query still reads, so it must stage from a drifted append point. + # GroupQueryAttention only allows a subsequent prompt (1 < S < T) at batch_size 1. + self._check_parity( + self._base_config(batch_size=1), + step_lengths=[32] * 10 + [1] * 20, + buffer_sequence_length=self.window_size + 8, + ) + + def test_capacity_too_small_is_rejected(self): + self._require_ep() + + device = self.device + torch_type = self.torch_type + ort_type = self.ort_type + # A cache shorter than the attention window can never hold everything the window covers. + capacity = self.window_size - 1 + config = self._base_config(rope_cache_length=capacity) + with self.assertRaisesRegex(Exception, "sliding_window_cache"): + _windowed_run_steps( + config, + capacity, + 1, + step_lengths=[32, 16], + q_all=torch.zeros( + (config.batch_size, 48, config.num_heads, config.head_size), device=device, dtype=torch_type + ), + k_all=torch.zeros( + (config.batch_size, 48, config.kv_num_heads, config.head_size), device=device, dtype=torch_type + ), + v_all=torch.zeros( + (config.batch_size, 48, config.kv_num_heads, config.head_size), device=device, dtype=torch_type + ), + cos=None, + sin=None, + device=device, + ort_type=ort_type, + torch_type=torch_type, + providers=self.providers, + ) + + def test_requires_local_window_size(self): + self._require_ep() + + config = self._base_config( + local_window_size=-1, + sliding_window_cache=1, + buffer_sequence_length=self.window_size + self.slack, + rope_cache_length=self.window_size + self.slack, + ) + with self.assertRaisesRegex(Exception, "sliding_window_cache"): + _windowed_make_session(config, self.ort_type, providers=self.providers) + + +class TestGQAWindowedKvCacheCpu(TestGQAWindowedKvCache): + """Same parity contract as the CUDA suite, run against the CPU GroupQueryAttention kernel.""" + + device = "cpu" + torch_type = torch.float32 + ort_type = TensorProto.FLOAT + providers: typing.ClassVar[list[str]] = ["CPUExecutionProvider"] + + def _require_ep(self): + pass + + def _base_config(self, **overrides): + # The float32 CPU kernel keeps an unquantized float32 cache by default. + overrides.setdefault("kv_cache_type", "float32") + return super()._base_config(**overrides) + + if __name__ == "__main__": unittest.main()