Fix Local Attention off by 1 bug (#25927)

### Description
Previously, local window size of GQA op excluded the current token. This
does not match standard HuggingFace implementations where tokens are
appended and then local masking occurs; the mismatch can cause the mask
to be off by 1 during generation, leading to accuracy issues. This PR
corrects this mismatch by including the current token. In practice, this
effectively decreases GQA window size by 1.
 


### Motivation and Context
This helps align our models with HuggingFace models.

---------

Co-authored-by: Kunal Vaishnavi <[email protected]>

Author: aciddelgado

onnxruntime/contrib_ops/cpu/bert/attention_parameters.h

@@ -86,7 +86,7 @@ struct GroupQueryAttentionParameters : AttentionParameters {
   int kv_hidden_size;           // hidden size of key or value
   int seqlen_past_kv_cache;     // sequence length of past kv tensor
   int seqlen_present_kv_cache;  // sequence length of present kv tensor
-  int local_window_size;        // The window size excludes current token. It only includes tokens on the left side.
+  int local_window_size;        // Mask out tokens prior to total_sequence_length - local_window_size
   bool kv_share_buffer;
   bool is_subsequent_prompt;  // indicates whether we have past context and seqlen > 1
   bool is_first_prompt;       // indicates whether this is first decoding step
@@ -106,7 +106,7 @@ struct PagedAttentionParameters : AttentionParameters {
   int block_size;              // block size for kv cache
   int max_num_blocks_per_seq;  // max number of blocks per sequence for kv cache
   int num_blocks;              // number of blocks in kv cache
-  int local_window_size;       // The window size excludes current token. It only includes tokens on the left side.
+  int local_window_size;       // The window size includes new token. It only includes tokens on the left side.
   bool rotary_interleaved;
   float softcap;
 };

onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h

@@ -297,16 +297,15 @@ class GQAAttentionBase {
         for (size_t seq = 0; seq < sequence_length; seq++) {
           size_t seq_causal_length = past_seqlen + seq + 1;
 
-          // local_window_size does not include the current query token, while window_size includes it.
           const bool should_apply_local_window = local_window_size_ >= 0 &&
-                                                 seq_causal_length > static_cast<size_t>(local_window_size_) + 1;
+                                                 seq_causal_length > static_cast<size_t>(local_window_size_);
 
-          const size_t start_offset = should_apply_local_window ? seq_causal_length - local_window_size_ - 1 : 0;
-          const size_t window_size = should_apply_local_window ? local_window_size_ + 1 : seq_causal_length;
+          const size_t start_offset = should_apply_local_window ? seq_causal_length - local_window_size_ : 0;
+          const size_t window_size = should_apply_local_window ? local_window_size_ : seq_causal_length;
 
           // Mask everything before local window, if local window should be applied
           if (should_apply_local_window) {
-            for (size_t total_seq_id = 0; total_seq_id < seq_causal_length - local_window_size_ - 1; total_seq_id++) {
+            for (size_t total_seq_id = 0; total_seq_id < seq_causal_length - local_window_size_; total_seq_id++) {
               if constexpr (std::is_same<U, float>::value) {
                 output_softmax[total_seq_id] = 0.f;
               } else {

onnxruntime/contrib_ops/cuda/bert/cutlass_fmha/fmha_launch_template.h

@@ -223,9 +223,7 @@ void LaunchCutlassFmha(const MemoryEfficientAttentionParams& params) {
     }
 
     p.use_smooth_softmax = params.use_smooth_softmax;
-
-    // local_windows_size in GQA does not include current query token, while windows_size in this kernel includes it.
-    p.window_size = params.local_window_size + 1;
+    p.window_size = params.local_window_size;
   }
 
   auto kernel_fn = attention_kernel_batched_impl<Attention>;

onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu

@@ -476,7 +476,7 @@ Status FlashAttention(
       parameters.seqlen_present_kv_cache, kv_sequence_length, parameters.rotary_dim,
       scale, parameters.softcap, is_causal, is_bf16, parameters.use_smooth_softmax, past_bsnh, parameters.num_splits,
       reinterpret_cast<void*>(data.softmax_lse_accum), reinterpret_cast<void*>(data.out_accum),
-      parameters.local_window_size, parameters.rotary_interleaved, parameters.is_packed_qkv));
+      parameters.local_window_size - 1, parameters.rotary_interleaved, parameters.is_packed_qkv));
 
   // if (parameters.left_padding && parameters.is_first_prompt) {
   //   ORT_RETURN_IF_ERROR(LaunchLeftPadLast(parameters, data, stream, device_prop.maxThreadsPerBlock));

onnxruntime/contrib_ops/cuda/bert/paged_attention_impl.cu

@@ -326,7 +326,7 @@ Status FlashAttention(
   ORT_RETURN_IF_ERROR(onnxruntime::flash::mha_varlen_fwd(
       device_prop, stream, q, key_cache, value_cache, output, cumulative_seqlens_q, cumulative_seqlens_kv,
       /*seqused_k*/ nullptr, block_table, softmax_lse, batch_size, num_heads, kv_num_heads, head_size, max_query_len,
-      max_seq_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, local_window_size, max_num_blocks_per_seq,
+      max_seq_len, token_count, scale, softcap, /*is_causal*/ true, is_bf16, local_window_size - 1, max_num_blocks_per_seq,
       block_size));
 
   DUMP_TENSOR_INIT();

onnxruntime/contrib_ops/webgpu/bert/attention.cc

@@ -250,7 +250,7 @@ Status InPlaceSoftmaxProgram::GenerateShaderCode(ShaderHelper& shader) const {
   if (has_sliding_window) {
     // Sliding window
     shader.MainFunctionBody()
-        << "let should_apply_local_window = uniforms.local_window_size >= 0 && seq_causal_length > uniforms.local_window_size + 1;\n"
+        << "let should_apply_local_window = uniforms.local_window_size >= 0 && seq_causal_length > uniforms.local_window_size;\n"
         << "let start_offset = select(0, seq_causal_length - uniforms.local_window_size, should_apply_local_window);\n"
         << "let effective_seq_length = select(seq_causal_length, uniforms.local_window_size, should_apply_local_window);\n";
   } else {

onnxruntime/test/python/transformers/test_gqa.py

@@ -582,7 +582,7 @@ def construct_local_mask(seqlen_q, seqlen_k, window_size, query_padding_mask, ke
         sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
         return torch.logical_or(
             col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk),
-            col_idx < row_idx + sk - sq - window_size[0],
+            col_idx <= row_idx + sk - sq - window_size[0],
         )
 
 

onnxruntime/test/python/transformers/test_gqa_cpu.py

@@ -1122,7 +1122,7 @@ def construct_local_mask(
         sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
         return torch.logical_or(
             col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk),
-            col_idx < row_idx + sk - sq - window_size[0],
+            col_idx <= row_idx + sk - sq - window_size[0],
         )
 
 

onnxruntime/test/python/transformers/test_paged_attention_cuda.py

@@ -331,7 +331,7 @@ def construct_local_mask(
         sk = torch.full_like(col_idx, seqlen_k) if key_padding_mask is None else sk
         return torch.logical_or(
             col_idx > torch.minimum(row_idx + sk - sq + window_size[1], sk),
-            col_idx < row_idx + sk - sq - window_size[0],
+            col_idx <= row_idx + sk - sq - window_size[0],
         )