From 1e9c47b6e6661e2185f2ad1892e6fae884cfc9af Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 27 Jul 2026 08:56:55 +0000 Subject: [PATCH 01/12] GQA support sliding_window_cache --- .../cpu/bert/attention_parameters.h | 15 + .../cpu/bert/group_query_attention.cc | 8 +- .../cpu/bert/group_query_attention_helper.h | 37 +- .../contrib_ops/cuda/bert/attention_data.h | 15 + .../cuda/bert/group_query_attention.cc | 107 +++++- .../cuda/bert/group_query_attention.h | 3 + .../cuda/bert/group_query_attention_impl.cu | 310 +++++++++++++++- .../cuda/bert/group_query_attention_impl.h | 36 ++ .../cuda/bert/group_query_attention_qkv.cuh | 55 +-- .../js/bert/group_query_attention.h | 4 + .../webgpu/bert/group_query_attention.h | 5 + .../core/graph/contrib_ops/bert_defs.cc | 15 +- .../test/python/transformers/test_gqa.py | 333 +++++++++++++++++- 13 files changed, 901 insertions(+), 42 deletions(-) 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..4aa19cc364208 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -43,7 +43,13 @@ REGISTER_KERNEL_TYPED(MLFloat16) template GroupQueryAttention::GroupQueryAttention(const OpKernelInfo& info) - : OpKernel(info), GQAAttentionBase(info, true) {} + : OpKernel(info), GQAAttentionBase(info, true) { + // The windowed KV cache (cache-relative indexing + shift compaction) is implemented only by the + // CUDA kernel. Fail loudly here rather than silently treating the window-sized buffer as a + // full-length cache, which would produce wrong results. + ORT_ENFORCE(info.GetAttrOrDefault("sliding_window_cache", 0) == 0, + "GroupQueryAttention (CPU): sliding_window_cache=1 is not implemented."); +} template Status GroupQueryAttention::Compute(OpKernelContext* context) const { 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..9fe55582ed047 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,38 @@ 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."); + } + 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 +413,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; 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 f5f38a1409930..052103e8cd6f1 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 = windowed_cache_capacity + parameters.sequence_length; + + const size_t staged_bytes = static_cast(parameters.batch_size) * parameters.kv_num_heads * + staged_cache_capacity * dense_head_size * sizeof(U); + staged_key_buffer = GetScratchBuffer(staged_bytes, context->GetComputeStream()); + staged_value_buffer = GetScratchBuffer(staged_bytes, context->GetComputeStream()); + + 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; @@ -617,7 +677,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) @@ -672,18 +734,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(); @@ -692,6 +766,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). + 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 = 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. @@ -843,6 +931,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 dbbe6e3e50db4..d3fc9d88b3f2f 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, 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, @@ -636,9 +646,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 +662,48 @@ __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. + // + // L must stay a pure function of the absolute length T: the op is stateless across steps and + // only ever sees seqlens_k, so it cannot recover a resident count it chose on a previous step. + // That rules out evicting in larger amortized batches. Note that nothing here depends on + // is_first_prompt; P == 0 falls out of the same formulas, which matters because a caller may + // not be able to signal the first prompt (onnxruntime-genai always passes the buffer length as + // total_sequence_length). + const int C = kv_cache_real_capacity; + const int C_used = kv_cache_capacity; + const int Lp = past_len < C ? past_len : C; + const int evicted = Lp + sequence_length - C; + const int shifted = Lp + sequence_length - C_used; + const int E = evicted > 0 ? evicted : 0; + const int D = shifted > 0 ? shifted : 0; + const int resident = Lp + sequence_length; + cache_past_seq_lens[i] = Lp - D; + cache_total_seq_lens[i] = resident < C_used ? resident : C_used; + evict_counts[i] = E; } } @@ -663,17 +712,209 @@ 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. +// ============================================================================ +__global__ void CompactKvCacheRows(uint8_t* __restrict__ dst, + const uint8_t* __restrict__ src, + const int* __restrict__ evict_counts, + const int* __restrict__ cache_past_seq_lens, + const int capacity, + const int row_bytes, + const bool src_is_shifted) { + // grid: (capacity, kv_num_heads, batch_size); block: threads over the row bytes (as uint4 where + // possible, otherwise bytes -- see the launcher, which guarantees row_bytes % 16 == 0). + const int s = blockIdx.x; + const int n = blockIdx.y; + const int b = blockIdx.z; + + const int retained = cache_past_seq_lens[b]; // Lp - D + if (s >= retained) { + return; + } + const int d = evict_counts[b]; + if (d <= 0) { + return; + } + + const int64_t head_stride = static_cast(capacity) * row_bytes; + const int64_t batch_stride = static_cast(gridDim.y) * head_stride; + const int64_t base = static_cast(b) * batch_stride + static_cast(n) * head_stride; + + // Pass 1 (src_is_shifted): scratch[s] <- cache[s + D]. + // Pass 2: cache[s] <- scratch[s]. + const int64_t src_row = base + static_cast(src_is_shifted ? (s + d) : s) * row_bytes; + const int64_t dst_row = base + static_cast(s) * row_bytes; + + const uint4* src_vec = reinterpret_cast(src + src_row); + uint4* dst_vec = reinterpret_cast(dst + dst_row); + const int vec_count = row_bytes / static_cast(sizeof(uint4)); + for (int i = threadIdx.x; i < vec_count; i += blockDim.x) { + dst_vec[i] = src_vec[i]; + } +} + +// 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."); + } + + const int vec_count = row_bytes / static_cast(sizeof(uint4)); + const int threads = vec_count < 256 ? ((vec_count + 31) / 32) * 32 : 256; + const dim3 grid(capacity, kv_num_heads, batch_size); + const dim3 block(threads > 0 ? threads : 32); + + 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."); + uint8_t* scratch_bytes = reinterpret_cast(scratch); + + void* caches[2] = {k_cache, v_cache}; + for (void* cache : caches) { + if (cache == nullptr) { + continue; + } + uint8_t* cache_bytes_ptr = reinterpret_cast(cache); + CompactKvCacheRows<<>>(scratch_bytes, cache_bytes_ptr, evict_counts, + cache_past_seq_lens, capacity, row_bytes, + /*src_is_shifted=*/true); + CompactKvCacheRows<<>>(cache_bytes_ptr, scratch_bytes, evict_counts, + cache_past_seq_lens, capacity, row_bytes, + /*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. +// ============================================================================ +__global__ void CopyKvCacheRows(uint8_t* __restrict__ dst, + const uint8_t* __restrict__ src, + const int* __restrict__ src_offsets, + const int src_capacity, + const int dst_capacity, + const int row_bytes) { + // grid: (rows, kv_num_heads, batch_size) + const int s = blockIdx.x; + const int n = blockIdx.y; + const int b = blockIdx.z; + + const int src_index = s + (src_offsets == nullptr ? 0 : src_offsets[b]); + if (src_index < 0 || src_index >= src_capacity || s >= dst_capacity) { + return; + } + + const int64_t src_head_stride = static_cast(src_capacity) * row_bytes; + const int64_t dst_head_stride = static_cast(dst_capacity) * row_bytes; + const int64_t src_row = (static_cast(b) * gridDim.y + n) * src_head_stride + + static_cast(src_index) * row_bytes; + const int64_t dst_row = (static_cast(b) * gridDim.y + n) * dst_head_stride + + static_cast(s) * row_bytes; + + const uint4* src_vec = reinterpret_cast(src + src_row); + uint4* dst_vec = reinterpret_cast(dst + dst_row); + const int vec_count = row_bytes / static_cast(sizeof(uint4)); + for (int i = threadIdx.x; i < vec_count; i += blockDim.x) { + dst_vec[i] = src_vec[i]; + } +} + +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."); + } + + const int vec_count = row_bytes / static_cast(sizeof(uint4)); + const int threads = vec_count < 256 ? ((vec_count + 31) / 32) * 32 : 256; + const dim3 grid(rows, kv_num_heads, batch_size); + const dim3 block(threads > 0 ? threads : 32); + + const void* srcs[2] = {src_k, src_v}; + void* dsts[2] = {dst_k, dst_v}; + for (int i = 0; i < 2; ++i) { + if (srcs[i] == nullptr || dsts[i] == nullptr) { + continue; + } + CopyKvCacheRows<<>>(reinterpret_cast(dsts[i]), + reinterpret_cast(srcs[i]), + src_offsets, src_capacity, dst_capacity, row_bytes); + } 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, \ @@ -732,8 +973,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, @@ -774,7 +1017,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, data.k_scale, // kv_cache_scale // Map cache type to XqaQuantType: NONE->kNone, Float8E4M3FN->kFp8, int8->kInt8 @@ -902,7 +1146,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; @@ -964,7 +1212,8 @@ Status DequantizeFlashAttentionFallback( q_rot, reinterpret_cast(data.present_key), reinterpret_cast(data.present_value), data.k_scale, data.v_scale, 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), @@ -991,7 +1240,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( @@ -1049,7 +1301,8 @@ Status FlashAttentionAndQuantizeKV( parameters.is_packed_qkv ? nullptr : reinterpret_cast(data.value), q_final, k_final, v_final, nullptr, nullptr, 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 @@ -1090,6 +1343,7 @@ Status FlashAttentionAndQuantizeKV( return Status::OK(); } + #endif #if USE_MEMORY_EFFICIENT_ATTENTION @@ -1152,7 +1406,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; @@ -1263,6 +1519,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, @@ -1375,6 +1640,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 b1f979fef4169..e10a6dba40e9f 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.h @@ -152,12 +152,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_qkv.cuh b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh index 1a5dc6e0fca6b..3e74a42df5117 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh @@ -52,9 +52,13 @@ __global__ void UnpackRoPEAppend( 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 + const int rope_max_pos, // upper bound (exclusive) for absolute RoPE positions (cos/sin cache length) + const int cache_capacity, // C: allocated sequence dim of the KV cache (== rope_max_pos when not windowed) + 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, @@ -201,10 +205,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; @@ -229,11 +233,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; @@ -242,7 +246,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 + @@ -252,7 +256,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 + @@ -322,32 +326,33 @@ Status DispatchUnpackRoPEAppendHeadSize( T* unpacked_q, U* k_cache, U* v_cache, const float* k_scale, const float* v_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, - 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, - 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, - 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, - 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 { @@ -367,8 +372,10 @@ Status LaunchUnpackRoPEAppend( T* unpacked_q, U* k_cache, U* v_cache, const float* k_scale, const float* v_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, @@ -376,6 +383,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) { @@ -418,7 +430,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, num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + k_scale, v_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 @@ -429,7 +442,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, num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + k_scale, v_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 @@ -437,7 +451,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, num_heads, kv_num_heads, head_size, d, max_seqlen, past_seq_lens, + k_scale, v_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..9d1839264e79e 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,14 @@ 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 + sequence_length - 1. 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 3a36c4ec4eea3..b0122925bf8d8 100644 --- a/onnxruntime/test/python/transformers/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_gqa.py @@ -139,6 +139,12 @@ 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 is at least local_window_size + q_sequence_length - 1. + 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 +374,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 +391,7 @@ def create_gqa_node_and_io( qk_output=output_qk, **quantization_attributes, **qk_norm_attributes, + **windowed_cache_attributes, domain="com.microsoft", ) @@ -433,7 +442,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]), @@ -3531,5 +3540,327 @@ 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): + onnx_model_str = create_group_query_attention_graph_past(config, ort_type=ort_type, share_buffer=True) + return InferenceSession( + onnx_model_str, SessionOptions(), providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] + ) + + +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, +): + """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) + 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 + + +@unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.") +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 + + 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): + if not has_cuda_provider(): + self.skipTest("CUDA required") + + device = "cuda" + torch_type = torch.float16 + ort_type = TensorProto.FLOAT16 + 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, + } + 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): + # Chunk size 128 requires capacity >= window_size + 128 - 1. + # 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_capacity_too_small_is_rejected(self): + if not has_cuda_provider(): + self.skipTest("CUDA required") + + device = "cuda" + torch_type = torch.float16 + ort_type = TensorProto.FLOAT16 + # 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, + ) + + def test_requires_local_window_size(self): + if not has_cuda_provider(): + self.skipTest("CUDA required") + + 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, TensorProto.FLOAT16) + + if __name__ == "__main__": unittest.main() From 4ce952731437a337971bba67128a07f5c0a54ba8 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 27 Jul 2026 16:27:22 +0000 Subject: [PATCH 02/12] skip fp8 tests on device < 89 --- onnxruntime/test/python/transformers/test_gqa.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py index b0122925bf8d8..7b326fc2fe649 100644 --- a/onnxruntime/test/python/transformers/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_gqa.py @@ -2208,6 +2208,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, @@ -3059,7 +3064,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)]: From c0360ba25c5808ca60ba0fb474daf49f78af3f8c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 27 Jul 2026 09:31:40 -0700 Subject: [PATCH 03/12] Add checks that sliding_window_cache=1 requires local_window_size > 0 Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../cpu/bert/group_query_attention_helper.h | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) 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 9fe55582ed047..17ec10df2a216 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -340,6 +340,24 @@ Status CheckInputs(const T* query, 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; From f0479487fbf1ad4ff15a4e33a420addd16b7cd2b Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 27 Jul 2026 09:31:59 -0700 Subject: [PATCH 04/12] update op spec Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- onnxruntime/core/graph/contrib_ops/bert_defs.cc | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/onnxruntime/core/graph/contrib_ops/bert_defs.cc b/onnxruntime/core/graph/contrib_ops/bert_defs.cc index 9d1839264e79e..f088fe226cfce 100644 --- a/onnxruntime/core/graph/contrib_ops/bert_defs.cc +++ b/onnxruntime/core/graph/contrib_ops/bert_defs.cc @@ -1240,8 +1240,9 @@ ONNX_MS_OPERATOR_SET_SCHEMA( "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 + sequence_length - 1. Default value is 0 (full-length cache).", + "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", From 591bbe682b590e6dbfa6da34d1d784aaba2fa62c Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 28 Jul 2026 00:20:25 +0000 Subject: [PATCH 05/12] docs: Update GroupQueryAttention GQA documentation for sliding_window_cache support - Add sliding_window_cache attribute to the operator schema table - Document the new sliding window (windowed) KV cache mode and its key behaviors - Include constraints and interactions with other features like Flash Attention --- docs/ContribOperators.md | 65 ++---------------------------------- docs/OperatorKernels.md | 1 - docs/contrib_ops/cuda/gqa.md | 25 ++++++++++++++ 3 files changed, 27 insertions(+), 64 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index 95f0698534a74..e30f096c316eb 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -52,7 +52,6 @@ Do not modify directly.* * com.microsoft.Irfft * com.microsoft.LinearAttention * com.microsoft.LongformerAttention - * com.microsoft.MatMulBlockQuantizedFp4Weight * com.microsoft.MatMulBlockQuantizedFp8Weight * com.microsoft.MatMulBnb4 * com.microsoft.MatMulFpQ4 @@ -2637,6 +2636,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
@@ -2911,68 +2912,6 @@ This version of the operator has been available since version 1 of the 'com.micr -### **com.microsoft.MatMulBlockQuantizedFp4Weight** - - Weight-only NVFP4 (E2M1) matrix multiplication. - - The weight tensor B is stored as packed NVFP4: two E2M1 values per byte (low nibble first). - The dequantized weight value is `e2m1(B) * weight_scale_2 * e4m3(weight_scale[n, k / block_size])`, - where `weight_scale` holds one E4M3 scale per `block_size` (default 16) consecutive K values and - `weight_scale_2` is a single global fp32 scale. The weight is dequantized to the activation type - (FP16/BF16) and multiplied with the FP16/BF16 activation. This path is architecture independent and - runs on Hopper (SM90) as well as Blackwell. - - The output columns `N` and the contraction dimension `K` are derived from the weight shape: - `N = B.shape[0]` and `K = 2 * B.shape[1]`. `K` must therefore be even. - -#### Version - -This version of the operator has been available since version 1 of the 'com.microsoft' operator set. - -#### Attributes - -
-
block_size : int
-
Number of consecutive K values that share one E4M3 weight scale. Default 16.
-
- -#### Inputs (4 - 6) - -
-
A : T
-
Row-major FP16/BF16 activation of shape [..., K].
-
B : T1
-
Packed NVFP4 weight of shape [N, K/2] stored as uint8 (two E2M1 values per byte, low nibble first).
-
weight_scale : T2
-
Per-block E4M3 weight scales of shape [N, ceil(K / block_size)] stored as raw uint8 bytes.
-
weight_scale_2 : T3
-
Global fp32 weight scale (scalar).
-
input_scale (optional) : T3
-
Optional global fp32 activation scale (scalar). Accepted for parity with quantized checkpoints; it is a no-op on the weight-only FP16/BF16 path and is reserved for the native NVFP4 path on Blackwell.
-
bias (optional) : T
-
Optional bias of shape [N].
-
- -#### Outputs - -
-
Y : T
-
Output of shape [..., N] in the activation type.
-
- -#### Type Constraints - -
-
T : tensor(float16), tensor(bfloat16)
-
Constrain activation, bias and output to FP16 or BF16.
-
T1 : tensor(uint8)
-
Constrain packed NVFP4 weight to uint8.
-
T2 : tensor(uint8)
-
Constrain E4M3 weight scales to uint8.
-
T3 : tensor(float)
-
Constrain scalar scales to FP32.
-
- ### **com.microsoft.MatMulBlockQuantizedFp8Weight** Weight-only block-scaled FP8 (E4M3) matrix multiplication. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index 5f2da22aabe65..d31eeec1101c5 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1099,7 +1099,6 @@ The **OpSet Version** column uses the following notation: |Irfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |LinearAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_state:**S**
*in* decay:**T**
*in* beta:**T**
*out* output:**T**
*out* present_state:**S**|1+|**T** = tensor(float), tensor(float16)| |LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| -|MatMulBlockQuantizedFp4Weight|*in* A:**T**
*in* B:**T1**
*in* weight_scale:**T2**
*in* weight_scale_2:**T3**
*in* input_scale:**T3**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(uint8)
**T2** = tensor(uint8)
**T3** = tensor(float)| |MatMulBlockQuantizedFp8Weight|*in* A:**T**
*in* B:**T1**
*in* b_scale:**T2**
*in* a_scale:**T2**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn)
**T2** = tensor(float)| |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)| diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md index 957b0498864b0..5ef91e21d3b26 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,30 @@ 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 should be at least `local_window_size`. +- **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. + ### Quantized KV cache To reduce the KV-cache memory footprint, the cache may be stored quantized while `query` stays From 625dd164fa8ebb582b410223728b5dd71fd3d11d Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Tue, 28 Jul 2026 21:43:39 +0000 Subject: [PATCH 06/12] Fix CUDA plugin EP build and restore generated operator docs - Use CudaKernel::GetComputeStream(context) for the windowed KV staging scratch buffers; OpKernelContext::GetComputeStream() does not exist in the plugin EP adapter, which broke the Linux/Windows CUDA Plugin EP builds. - Restore the MatMulBlockQuantizedFp4Weight sections that were accidentally dropped from ContribOperators.md / OperatorKernels.md, fixing the Windows GPU Kernel Documentation Validation check. --- docs/ContribOperators.md | 64 +++++++++++++++++++ docs/OperatorKernels.md | 1 + .../cuda/bert/group_query_attention.cc | 4 +- 3 files changed, 67 insertions(+), 2 deletions(-) diff --git a/docs/ContribOperators.md b/docs/ContribOperators.md index e30f096c316eb..23cdbac30a042 100644 --- a/docs/ContribOperators.md +++ b/docs/ContribOperators.md @@ -52,6 +52,7 @@ Do not modify directly.* * com.microsoft.Irfft * com.microsoft.LinearAttention * com.microsoft.LongformerAttention + * com.microsoft.MatMulBlockQuantizedFp4Weight * com.microsoft.MatMulBlockQuantizedFp8Weight * com.microsoft.MatMulBnb4 * com.microsoft.MatMulFpQ4 @@ -2912,6 +2913,69 @@ This version of the operator has been available since version 1 of the 'com.micr +### **com.microsoft.MatMulBlockQuantizedFp4Weight** + + Weight-only NVFP4 (E2M1) matrix multiplication. + + The weight tensor B is stored as packed NVFP4: two E2M1 values per byte (low nibble first). + The dequantized weight value is `e2m1(B) * weight_scale_2 * e4m3(weight_scale[n, k / block_size])`, + where `weight_scale` holds one E4M3 scale per `block_size` (default 16) consecutive K values and + `weight_scale_2` is a single global fp32 scale. The weight is dequantized to the activation type + (FP16/BF16) and multiplied with the FP16/BF16 activation. This path is architecture independent and + runs on Hopper (SM90) as well as Blackwell. + + The output columns `N` and the contraction dimension `K` are derived from the weight shape: + `N = B.shape[0]` and `K = 2 * B.shape[1]`. `K` must therefore be even. + +#### Version + +This version of the operator has been available since version 1 of the 'com.microsoft' operator set. + +#### Attributes + +
+
block_size : int
+
Number of consecutive K values that share one E4M3 weight scale. Default 16.
+
+ +#### Inputs (4 - 6) + +
+
A : T
+
Row-major FP16/BF16 activation of shape [..., K].
+
B : T1
+
Packed NVFP4 weight of shape [N, K/2] stored as uint8 (two E2M1 values per byte, low nibble first).
+
weight_scale : T2
+
Per-block E4M3 weight scales of shape [N, ceil(K / block_size)] stored as raw uint8 bytes.
+
weight_scale_2 : T3
+
Global fp32 weight scale (scalar).
+
input_scale (optional) : T3
+
Optional global fp32 activation scale (scalar). Accepted for parity with quantized checkpoints; it is a no-op on the weight-only FP16/BF16 path and is reserved for the native NVFP4 path on Blackwell.
+
bias (optional) : T
+
Optional bias of shape [N].
+
+ +#### Outputs + +
+
Y : T
+
Output of shape [..., N] in the activation type.
+
+ +#### Type Constraints + +
+
T : tensor(float16), tensor(bfloat16)
+
Constrain activation, bias and output to FP16 or BF16.
+
T1 : tensor(uint8)
+
Constrain packed NVFP4 weight to uint8.
+
T2 : tensor(uint8)
+
Constrain E4M3 weight scales to uint8.
+
T3 : tensor(float)
+
Constrain scalar scales to FP32.
+
+ + ### **com.microsoft.MatMulBlockQuantizedFp8Weight** Weight-only block-scaled FP8 (E4M3) matrix multiplication. diff --git a/docs/OperatorKernels.md b/docs/OperatorKernels.md index d31eeec1101c5..5f2da22aabe65 100644 --- a/docs/OperatorKernels.md +++ b/docs/OperatorKernels.md @@ -1099,6 +1099,7 @@ The **OpSet Version** column uses the following notation: |Irfft|*in* X:**T**
*out* Y:**T**|1+|**T** = tensor(double), tensor(float), tensor(float16)| |LinearAttention|*in* query:**T**
*in* key:**T**
*in* value:**T**
*in* past_state:**S**
*in* decay:**T**
*in* beta:**T**
*out* output:**T**
*out* present_state:**S**|1+|**T** = tensor(float), tensor(float16)| |LongformerAttention|*in* input:**T**
*in* weight:**T**
*in* bias:**T**
*in* mask:**T**
*in* global_weight:**T**
*in* global_bias:**T**
*in* global:**G**
*out* output:**T**|1+|**T** = tensor(float), tensor(float16)| +|MatMulBlockQuantizedFp4Weight|*in* A:**T**
*in* B:**T1**
*in* weight_scale:**T2**
*in* weight_scale_2:**T3**
*in* input_scale:**T3**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(uint8)
**T2** = tensor(uint8)
**T3** = tensor(float)| |MatMulBlockQuantizedFp8Weight|*in* A:**T**
*in* B:**T1**
*in* b_scale:**T2**
*in* a_scale:**T2**
*in* bias:**T**
*out* Y:**T**|1+|**T** = tensor(bfloat16), tensor(float16)
**T1** = tensor(float8e4m3fn)
**T2** = tensor(float)| |MatMulBnb4|*in* A:**T1**
*in* B:**T2**
*in* absmax:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)| |MatMulNBits|*in* A:**T1**
*in* B:**T2**
*in* scales:**T1**
*in* zero_points:**T3**
*in* g_idx:**T4**
*in* bias:**T1**
*out* Y:**T1**|1+|**T1** = tensor(bfloat16), tensor(float), tensor(float16)
**T2** = tensor(uint8)
**T3** = tensor(bfloat16), tensor(float), tensor(float16), tensor(uint8)| diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 8d9824b77e212..2b1eb3d4d78d8 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -465,8 +465,8 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons const size_t staged_bytes = static_cast(parameters.batch_size) * parameters.kv_num_heads * staged_cache_capacity * dense_head_size * sizeof(U); - staged_key_buffer = GetScratchBuffer(staged_bytes, context->GetComputeStream()); - staged_value_buffer = GetScratchBuffer(staged_bytes, context->GetComputeStream()); + 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(), From c40bf672dc7d03b87523a132db345c7b6bb3e5d0 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 29 Jul 2026 01:57:51 +0000 Subject: [PATCH 07/12] reject attemtion_bias and windowed kv cache --- .../cpu/bert/group_query_attention_helper.h | 5 ++++ .../cuda/bert/group_query_attention.cc | 4 +-- .../test/python/transformers/test_gqa.py | 29 +++++++++++++++++++ 3 files changed, 36 insertions(+), 2 deletions(-) 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 17ec10df2a216..44e7def5f3bc7 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h @@ -467,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/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 2b1eb3d4d78d8..607f4905f0c23 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -461,9 +461,9 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons windowed_present_key = data.present_key; windowed_present_value = data.present_value; windowed_cache_capacity = parameters.seqlen_present_kv_cache; - staged_cache_capacity = windowed_cache_capacity + parameters.sequence_length; + staged_cache_capacity = SafeInt(windowed_cache_capacity) + parameters.sequence_length; - const size_t staged_bytes = static_cast(parameters.batch_size) * parameters.kv_num_heads * + 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)); diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py index 030ebe5484543..8b832e97a75be 100644 --- a/onnxruntime/test/python/transformers/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_gqa.py @@ -2702,6 +2702,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): From ec2bad9406d99ac0ea9ece92b7405682c0cf5f52 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 29 Jul 2026 05:33:33 +0000 Subject: [PATCH 08/12] Optimize Windowed KV cache compaction kernel --- .../cuda/bert/group_query_attention.cc | 4 +- .../cuda/bert/group_query_attention_impl.cu | 181 +++++++++++------- 2 files changed, 109 insertions(+), 76 deletions(-) diff --git a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc index 607f4905f0c23..897749a3f6b35 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc @@ -779,13 +779,13 @@ Status GroupQueryAttention::ComputeInternal(OpKernelContext* context) cons // 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). + // 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 = static_cast(parameters.batch_size) * parameters.kv_num_heads * + 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(); 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 4212fe6de3cd6..7713ab78b7a7b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -747,42 +747,62 @@ Status LaunchGetSequenceLengths( // row_bytes: number of bytes occupied by one (batch, head, position) entry. // capacity : allocated sequence dimension C of the cache. // ============================================================================ -__global__ void CompactKvCacheRows(uint8_t* __restrict__ dst, - const uint8_t* __restrict__ src, + +// 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: during decoding the cache is +// over its budget on only one step in every cache_slack steps, so nearly every launch does no work +// at all and its cost is dominated by scheduling the grid. +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 row_bytes, + const int vec_count, + const bool copy_v, const bool src_is_shifted) { - // grid: (capacity, kv_num_heads, batch_size); block: threads over the row bytes (as uint4 where - // possible, otherwise bytes -- see the launcher, which guarantees row_bytes % 16 == 0). - const int s = blockIdx.x; - const int n = blockIdx.y; + // grid: (ceil(capacity / blockDim.y), kv_num_heads, batch_size); block: (row chunks, rows). const int b = blockIdx.z; - - const int retained = cache_past_seq_lens[b]; // Lp - D - if (s >= retained) { - return; - } 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) * row_bytes; - const int64_t batch_stride = static_cast(gridDim.y) * head_stride; - const int64_t base = static_cast(b) * batch_stride + static_cast(n) * head_stride; + 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 src_row = base + static_cast(src_is_shifted ? (s + d) : s) * row_bytes; - const int64_t dst_row = base + static_cast(s) * row_bytes; + 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; - const uint4* src_vec = reinterpret_cast(src + src_row); - uint4* dst_vec = reinterpret_cast(dst + dst_row); - const int vec_count = row_bytes / static_cast(sizeof(uint4)); - for (int i = threadIdx.x; i < vec_count; i += blockDim.x) { - dst_vec[i] = src_vec[i]; + 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]; + } } } @@ -807,29 +827,36 @@ Status LaunchCompactKvCache(void* k_cache, " bytes) must be a multiple of 16 for the compaction shift."); } - const int vec_count = row_bytes / static_cast(sizeof(uint4)); - const int threads = vec_count < 256 ? ((vec_count + 31) / 32) * 32 : 256; - const dim3 grid(capacity, kv_num_heads, batch_size); - const dim3 block(threads > 0 ? threads : 32); + // 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."); - uint8_t* scratch_bytes = reinterpret_cast(scratch); + // 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; - void* caches[2] = {k_cache, v_cache}; - for (void* cache : caches) { - if (cache == nullptr) { - continue; - } - uint8_t* cache_bytes_ptr = reinterpret_cast(cache); - CompactKvCacheRows<<>>(scratch_bytes, cache_bytes_ptr, evict_counts, - cache_past_seq_lens, capacity, row_bytes, - /*src_is_shifted=*/true); - CompactKvCacheRows<<>>(cache_bytes_ptr, scratch_bytes, evict_counts, - cache_past_seq_lens, capacity, row_bytes, - /*src_is_shifted=*/false); - } + 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()); } @@ -843,36 +870,38 @@ Status LaunchCompactKvCache(void* k_cache, // 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. +// the cache is written in quantized form. K and V share every index, so they move together. // ============================================================================ -__global__ void CopyKvCacheRows(uint8_t* __restrict__ dst, - const uint8_t* __restrict__ src, +__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 row_bytes) { - // grid: (rows, kv_num_heads, batch_size) - const int s = blockIdx.x; - const int n = blockIdx.y; + 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 || s >= dst_capacity) { + if (src_index < 0 || src_index >= src_capacity) { return; } - const int64_t src_head_stride = static_cast(src_capacity) * row_bytes; - const int64_t dst_head_stride = static_cast(dst_capacity) * row_bytes; - const int64_t src_row = (static_cast(b) * gridDim.y + n) * src_head_stride + - static_cast(src_index) * row_bytes; - const int64_t dst_row = (static_cast(b) * gridDim.y + n) * dst_head_stride + - static_cast(s) * row_bytes; + 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; - const uint4* src_vec = reinterpret_cast(src + src_row); - uint4* dst_vec = reinterpret_cast(dst + dst_row); - const int vec_count = row_bytes / static_cast(sizeof(uint4)); - for (int i = threadIdx.x; i < vec_count; i += blockDim.x) { - dst_vec[i] = src_vec[i]; + 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]; + } } } @@ -897,21 +926,25 @@ Status LaunchCopyKvCacheWindow(void* dst_k, " bytes) must be a multiple of 16 to stage the KV cache."); } - const int vec_count = row_bytes / static_cast(sizeof(uint4)); - const int threads = vec_count < 256 ? ((vec_count + 31) / 32) * 32 : 256; - const dim3 grid(rows, kv_num_heads, batch_size); - const dim3 block(threads > 0 ? threads : 32); - - const void* srcs[2] = {src_k, src_v}; - void* dsts[2] = {dst_k, dst_v}; - for (int i = 0; i < 2; ++i) { - if (srcs[i] == nullptr || dsts[i] == nullptr) { - continue; - } - CopyKvCacheRows<<>>(reinterpret_cast(dsts[i]), - reinterpret_cast(srcs[i]), - src_offsets, src_capacity, dst_capacity, row_bytes); + // 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()); } From f0f834e274e8e2b7502e13d4f1dff3ab8b3988fe Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Wed, 29 Jul 2026 16:25:35 -0700 Subject: [PATCH 09/12] [CPU] Add windowed kv cache for GQA (#31142) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add CPU support for `sliding_window_cache=1` on `com.microsoft::GroupQueryAttention`, matching the existing CUDA/XQA contract: the cache buffer holds only the most recent tokens, indexed in cache-relative coordinates; results are numerically identical to a full-length cache. The second commit introduces a drifting-layout optimisation that amortises the per-step compaction cost, reducing per-token decode latency by 1.35–1.81× depending on buffer size. ## Summary of Changes ### CPU Kernel (`group_query_attention.cc / .h`) | File | Change | |---|---| | `contrib_ops/cpu/bert/group_query_attention.cc` | Three new anonymous-namespace helpers (`WindowedCacheEnd`, `PlanWindowedKvCache`, `MoveKvCacheRows`) plus the windowed block inside `Compute()`. | | `contrib_ops/cpu/bert/group_query_attention.h` | Add private `sliding_window_cache_` field. | **Key design decisions:** - **Single intercept point.** The CPU kernel has four attention paths (naive, quantized, quantized-flash, MLAS-flash). Instead of modifying each, the entire windowed-cache logic lives in `Compute()` and passes a compacted buffer as both `past_key/value` and `present_key/value`. The CPU kernel infers `past_present_share_buffer` from pointer equality, so `ConcatStateChunkGQA` only appends new tokens. - **Drifting layout (commit 2).** Resident entries live at `[0, end)` with the append point tracked via: ``` gap = C - W + 1 end(T) = T while T ≤ C = T − gap × ⌈(T − C) / gap⌉ afterwards (so W ≤ end ≤ C) ``` This is still a pure function of the absolute past length `T`, so the op is stateless. A whole `gap`-sized block is reclaimed at once, so compaction runs once every `gap` decode steps instead of every step. Safety: entries in `[0, end − W)` are older than the window and are already zeroed by the local-window mask — the same property that makes any `C > W` cache correct. - **Skip compaction on drifting steps.** When `seed_offset == 0` (the common decode step) the `MoveKvCacheRows` calls are bypassed entirely; a no-op thread-pool dispatch is not free (~65 µs/token). - **Staging for hazardous steps.** A multi-token step that would compact below the history its own first query reads runs against a temporary staging cache of `max_b(end_before) + S` rows; the surviving tail is written back afterwards. - **Thread-pool parallel `MoveKvCacheRows`.** Compaction is parallelised over `(batch, kv_head)` via `ThreadPool::TryParallelFor` with per-slice `memmove`. At `C = W = 128` (0.50 MiB moved per token) the serial version cost ~0.32 ms/token, making windowed decode 3× slower than the full cache; parallelised over 8 threads it drops to noise. ### Tests (`test_gqa.py`) | Test class / method | What it covers | |---|---| | `TestGQAWindowedKvCacheCpu` | Subclasses the existing CUDA suite; runs the same 11 test scenarios on `CPUExecutionProvider` with `float32` cache. | | `test_small_slack_many_compactions` | `C = W + 8`, 400 decode steps → ~44 compaction events; exercises the drifting layout exhaustively. | | `test_chunked_prefill_small_slack` | 32-token chunks at `C = W + 8`; staging from a drifted append point (the gap is < chunk size). | ### Documentation (`docs/contrib_ops/cuda/gqa.md`) - Corrected the "Cache capacity" bullet: `C ≥ W` is the minimum; modest surplus is recommended for CPU performance. - Added "CPU support" bullet describing the drifting layout and staging contract. - Corrected the PR-level doc note that falsely claimed amortisation was impossible under statelessness. ## Performance Results All measurements on AMD EPYC 7V12, 96 cores, 8 intra-op threads, fp32, 64 Q-heads / 8 KV-heads × 64 (`head_size`), `local_window_size = 128`. ### Before this PR — full-cache is unchanged Windowed decode latency (decode p10 ms) **before** the drifting layout was always-compact-on-every-step. The baseline (no sliding window) did not regress: | Prompt | Full cache (before) | Full cache (after) | |---|---|---| | 512 | 0.129–0.163 | 0.122–0.162 | | 2048 | 0.395–0.427 | 0.356–0.395 | | 8192 | — | within noise | ### Windowed vs full cache — decode throughput (decode p10 ms) `C = W + 256` (slack 256, the old default): | Prompt | Full cache | Windowed | Speedup | |---|---|---|---| | 512 | 0.169 | 0.145 | 1.17× | | 2048 | 0.352 | 0.167 | 2.11× | | 8192 | 1.301 | 0.155 | 8.4× | Windowed decode is **flat in context length** (attention scans only `W` effective keys). The CPU full-cache path does not exploit `local_window_size` and scans the entire cache. KV memory: full cache at 8192 tokens = 33.2 MiB; windowed (`C = 384`) = 1.50 MiB (fixed). ### Drifting layout — capacity sweep (prefill 2048 + 800 decode steps, decode p50 ms, 3 runs) A/B baseline pins `gap = 1` in the planner, which is equivalent to always-compact-on-every-step. | C | slack | gap | before | after | speedup | |---|---|---|---|---|---| | 128 | 0 | 1 | 0.0789 | *(same code, noise)* | — | | 132 | 4 | 5 | 0.0909 | 0.0654 | 1.39× | | **144** | **16** | **17** | 0.0964 | **0.0649** | **1.43×** | | 160 | 32 | 33 | 0.0957 | 0.0675 | 1.42× | | 256 | 128 | 129 | 0.1097 | 0.0773 | 1.42× | | 384 | 256 | 257 | 0.1573 | 0.0972 | 1.62× | | 512 | 384 | 385 | 0.2010 | 0.1109 | 1.81× | **The new decode-latency optimum is `C = W + 16`** (decode p50 0.065 ms, 1.35× better than `C = W` before this change). The curve is now a U-shape with a flat floor over slack 4–128; without drift it was monotonically worse with capacity because every decode step shifted `C − 1` rows. The analytic optimum matches: minimising `attention_slope × slack/2 + shift_cost/gap` over `gap` gives `slack* = sqrt(2 × shift_cost / attention_slope) ≈ 17`. ## Testing ```bash # Full windowed suite (13 CPU + 13 CUDA tests): pytest onnxruntime/test/python/transformers/test_gqa.py -k "TestGQAWindowedKvCache" -q # CPU GQA regression tests: pytest onnxruntime/test/python/transformers/test_gqa_cpu.py -q ``` Expected: 26 windowed tests pass, 4 CPU regression tests pass. ## Limitations (CPU only) The following are rejected with a clear error at session creation or compute time when combined with `sliding_window_cache=1` on CPU: - `attention_bias` input - `qk_output` attribute - Shared QKV layout (`kv_sequence_length == 0`) These limitations do not apply to the CUDA path. ## Checklist - [x] Tests added/updated - [x] No breaking changes - [x] Documentation updated (`docs/contrib_ops/cuda/gqa.md`) --- docs/contrib_ops/cuda/gqa.md | 12 +- .../cpu/bert/group_query_attention.cc | 422 +++++++++++++++--- .../cpu/bert/group_query_attention.h | 5 + .../cuda/bert/group_query_attention_impl.cu | 57 ++- .../test/python/transformers/test_gqa.py | 89 +++- 5 files changed, 480 insertions(+), 105 deletions(-) diff --git a/docs/contrib_ops/cuda/gqa.md b/docs/contrib_ops/cuda/gqa.md index 5ef91e21d3b26..327312e22df0f 100644 --- a/docs/contrib_ops/cuda/gqa.md +++ b/docs/contrib_ops/cuda/gqa.md @@ -166,7 +166,11 @@ 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 should be at least `local_window_size`. + 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) @@ -179,6 +183,12 @@ employ local attention (e.g., GPT-OSS with layer-wise sliding windows). - **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 diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 4aa19cc364208..8c93b00803f88 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -14,12 +14,159 @@ #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. +int WindowedCacheEnd(int past_sequence_length, int capacity, int gap) { + if (past_sequence_length <= capacity) { + return past_sequence_length; // still filling: nothing has been reclaimed yet + } + const int overflow = past_sequence_length - capacity; + const int reclaimed = gap * ((overflow + gap - 1) / gap); + return 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) { + const int past_sequence_length = seqlens_k[b] + 1 - sequence_length; // 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 int 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) \ @@ -44,11 +191,9 @@ REGISTER_KERNEL_TYPED(MLFloat16) template GroupQueryAttention::GroupQueryAttention(const OpKernelInfo& info) : OpKernel(info), GQAAttentionBase(info, true) { - // The windowed KV cache (cache-relative indexing + shift compaction) is implemented only by the - // CUDA kernel. Fail loudly here rather than silently treating the window-sized buffer as a - // full-length cache, which would produce wrong results. - ORT_ENFORCE(info.GetAttrOrDefault("sliding_window_cache", 0) == 0, - "GroupQueryAttention (CPU): sliding_window_cache=1 is not implemented."); + 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 @@ -114,7 +259,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, @@ -149,13 +298,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); @@ -169,7 +323,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) { @@ -201,6 +355,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)); @@ -324,71 +486,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/cuda/bert/group_query_attention_impl.cu b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu index 7713ab78b7a7b..3eb25bc7fff39 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -687,22 +687,47 @@ __global__ void GetSequenceLengths(const int* total_seq_lens_minus_one, // 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. // - // L must stay a pure function of the absolute length T: the op is stateless across steps and - // only ever sees seqlens_k, so it cannot recover a resident count it chose on a previous step. - // That rules out evicting in larger amortized batches. Note that nothing here depends on - // is_first_prompt; P == 0 falls out of the same formulas, which matters because a caller may - // not be able to signal the first prompt (onnxruntime-genai always passes the buffer length as - // total_sequence_length). + // 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 Lp = past_len < C ? past_len : C; - const int evicted = Lp + sequence_length - C; - const int shifted = Lp + sequence_length - C_used; - const int E = evicted > 0 ? evicted : 0; - const int D = shifted > 0 ? shifted : 0; - const int resident = Lp + sequence_length; + 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. + // end(P) = P while P <= C, else P - gap * ceil((P - C) / gap). + // This keeps W <= end <= C once the cache has filled. + auto windowed_end = [&](int P) -> int { + if (P <= C) return P; + const int overflow = P - C; + const int reclaimed = gap * ((overflow + gap - 1) / gap); + return P - reclaimed; + }; + const int end_before = windowed_end(past_len); + const int end_after = windowed_end(past_len + 1); + 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] = resident < C_used ? resident : C_used; + cache_total_seq_lens[i] = (Lp + sequence_length) < C_used ? (Lp + sequence_length) : C_used; evict_counts[i] = E; } } @@ -751,9 +776,9 @@ Status LaunchGetSequenceLengths( // 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: during decoding the cache is -// over its budget on only one step in every cache_slack steps, so nearly every launch does no work -// at all and its cost is dominated by scheduling the grid. +// 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, diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py index 8b832e97a75be..1962f079b0746 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 @@ -140,7 +141,11 @@ class GQAConfig: # 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 is at least local_window_size + q_sequence_length - 1. + # 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. @@ -3871,11 +3876,11 @@ def test_gqa_gemma4_local_window_decode_fp16(self): # ################################################################################################# -def _windowed_make_session(config: GQAConfig, ort_type): +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) - return InferenceSession( - onnx_model_str, SessionOptions(), providers=[resolve_cuda_plugin_ep("CUDAExecutionProvider")] - ) + if providers is None: + providers = [resolve_cuda_plugin_ep("CUDAExecutionProvider")] + return InferenceSession(onnx_model_str, SessionOptions(), providers=providers) def _windowed_run_steps( @@ -3894,6 +3899,7 @@ def _windowed_run_steps( 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. @@ -3924,7 +3930,7 @@ def _windowed_run_steps( config.sliding_window_cache = sliding_window_cache if step_length not in sessions: - sessions[step_length] = _windowed_make_session(config, ort_type) + sessions[step_length] = _windowed_make_session(config, ort_type, providers=providers) session = sessions[step_length] io_binding = session.io_binding() @@ -3980,7 +3986,6 @@ def _windowed_run_steps( return outputs -@unittest.skipIf(not has_flash_attention(), "Flash Attention is not available, skipping tests.") 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 @@ -3990,6 +3995,15 @@ class TestGQAWindowedKvCache(unittest.TestCase): 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, @@ -4007,12 +4021,11 @@ def _base_config(self, **overrides): return config def _check_parity(self, base_config, step_lengths, rtol=1e-3, atol=1e-3, buffer_sequence_length=None): - if not has_cuda_provider(): - self.skipTest("CUDA required") + self._require_ep() - device = "cuda" - torch_type = torch.float16 - ort_type = TensorProto.FLOAT16 + 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) @@ -4054,6 +4067,7 @@ def _check_parity(self, base_config, step_lengths, rtol=1e-3, atol=1e-3, buffer_ "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): @@ -4085,7 +4099,7 @@ def test_prompt_longer_than_capacity_then_decode(self): self._check_parity(self._base_config(), steps) def test_chunked_prefill(self): - # Chunk size 128 requires capacity >= window_size + 128 - 1. + # 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) @@ -4142,13 +4156,29 @@ def test_capacity_equal_to_window_size(self): 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): - if not has_cuda_provider(): - self.skipTest("CUDA required") + self._require_ep() - device = "cuda" - torch_type = torch.float16 - ort_type = TensorProto.FLOAT16 + 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) @@ -4172,11 +4202,11 @@ def test_capacity_too_small_is_rejected(self): device=device, ort_type=ort_type, torch_type=torch_type, + providers=self.providers, ) def test_requires_local_window_size(self): - if not has_cuda_provider(): - self.skipTest("CUDA required") + self._require_ep() config = self._base_config( local_window_size=-1, @@ -4185,7 +4215,24 @@ def test_requires_local_window_size(self): rope_cache_length=self.window_size + self.slack, ) with self.assertRaisesRegex(Exception, "sliding_window_cache"): - _windowed_make_session(config, TensorProto.FLOAT16) + _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__": From aef35fcc27ea0c45a7136476f4cf3320d1e3da8e Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 30 Jul 2026 07:21:02 +0000 Subject: [PATCH 10/12] Clarify comments for windowed KV cache and per-channel scales Comment-only follow-up; no functional change. - UnpackRoPEAppend: expand the rope_max_pos / cache_capacity docs to state explicitly that rope_max_pos is an ABSOLUTE position limit (the cos/sin cache length) and must not be conflated with cache_capacity. They only coincide for a full-length cache; using the capacity as the RoPE bound on a windowed layer would silently skip RoPE for every absolute position >= capacity -- an accuracy bug rather than a crash. - Correct the per-channel scale shape in the q_fold_scale and ScaleHeadsByChannelScaleKernel comments to (kv_num_heads, 1, head_size), matching the operator spec in docs/contrib_ops/cuda/gqa.md, and note that the indexing flattens the singleton middle dimension. - test_gqa.py: drop the now-stale enumeration of quantized XQA selection constraints (the same-tensor k/v scale requirement was lifted and the head_size/group_size sets are configuration dependent). --- .../cuda/bert/group_query_attention_qdq.cuh | 2 +- .../cuda/bert/group_query_attention_qkv.cuh | 19 +++++++++++++------ .../test/python/transformers/test_gqa.py | 6 +++--- 3 files changed, 17 insertions(+), 10 deletions(-) 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 0c3cfa3a07c4b..c026071a3bb2f 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,24 @@ __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 rope_max_pos, // upper bound (exclusive) for absolute RoPE positions (cos/sin cache length) - const int cache_capacity, // C: allocated sequence dim of the KV cache (== rope_max_pos when not windowed) + // 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. diff --git a/onnxruntime/test/python/transformers/test_gqa.py b/onnxruntime/test/python/transformers/test_gqa.py index 1962f079b0746..df0209962cd4a 100644 --- a/onnxruntime/test/python/transformers/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_gqa.py @@ -3089,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: From d3c1a4dd9c750d36c68de2e80b56f889acfc4178 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 30 Jul 2026 07:45:40 +0000 Subject: [PATCH 11/12] Address GQA windowed KV cache review feedback - Widen the windowed-cache end-pointer math to int64_t in both the CPU (WindowedCacheEnd, PlanWindowedKvCache) and CUDA kernels. seqlens_k is caller-supplied, so `seqlens_k[b] + 1` and `overflow + gap - 1` could signed-overflow near INT32_MAX before validation rejects the input. Results stay bounded by capacity, so narrowing back to int is safe. - Clamp the derived past_sequence_length at zero so a bogus seqlens_k cannot produce a negative row count downstream. - Hoist the CUDA windowed_end lambda out of GetSequenceLengths into a file-scope __device__ WindowedCacheEnd() helper that mirrors the CPU implementation, and document that the two must stay in sync. - Fix lintrunner comment alignment in GetSequenceLengths. --- .../cpu/bert/group_query_attention.cc | 23 ++++++++---- .../cuda/bert/group_query_attention_impl.cu | 36 +++++++++++++------ 2 files changed, 41 insertions(+), 18 deletions(-) diff --git a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc index 8c93b00803f88..e6efc4c207300 100644 --- a/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc +++ b/onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc @@ -59,13 +59,18 @@ struct WindowedStep { // 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. -int WindowedCacheEnd(int past_sequence_length, int capacity, int gap) { +// +// 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 past_sequence_length; // still filling: nothing has been reclaimed yet + return static_cast(past_sequence_length); // still filling: nothing has been reclaimed yet } - const int overflow = past_sequence_length - capacity; - const int reclaimed = gap * ((overflow + gap - 1) / gap); - return past_sequence_length - reclaimed; + 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 @@ -88,7 +93,11 @@ void PlanWindowedKvCache(const int32_t* seqlens_k, int max_resident = 0; use_staging = false; for (int b = 0; b < batch_size; ++b) { - const int past_sequence_length = seqlens_k[b] + 1 - sequence_length; // P + // 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 @@ -96,7 +105,7 @@ void PlanWindowedKvCache(const int32_t* seqlens_k, // 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 int required = std::min(past_sequence_length, local_window_size - 1); + 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; 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 3eb25bc7fff39..89144f5e50787 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu @@ -615,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 // ============================================================================ @@ -698,22 +719,15 @@ __global__ void GetSequenceLengths(const int* total_seq_lens_minus_one, 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 + 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. - // end(P) = P while P <= C, else P - gap * ceil((P - C) / gap). + // Drifting-layout end pointer for single-token steps; see WindowedCacheEnd above. // This keeps W <= end <= C once the cache has filled. - auto windowed_end = [&](int P) -> int { - if (P <= C) return P; - const int overflow = P - C; - const int reclaimed = gap * ((overflow + gap - 1) / gap); - return P - reclaimed; - }; - const int end_before = windowed_end(past_len); - const int end_after = windowed_end(past_len + 1); + 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; From 8f3daa0dc58f22bc04c2311b38c3311bb911258e Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Thu, 30 Jul 2026 08:28:21 +0000 Subject: [PATCH 12/12] lintrunner --- onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 c026071a3bb2f..c72d5f9093790 100644 --- a/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh +++ b/onnxruntime/contrib_ops/cuda/bert/group_query_attention_qkv.cuh @@ -59,7 +59,7 @@ __global__ void UnpackRoPEAppend( 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 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