Fix packed-QKV and broadcast-head bias strides in quantized GQA flash attention

onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h

@@ -831,6 +831,9 @@ class GQAAttentionBase {
       args.buffer = reinterpret_cast<float*>(flash_buffer_alloc);
       args.buffer_size_per_thread = buffer_size_per_thread;
       args.query = Q;
+      args.q_batch_stride = packed_qkv
+                                ? static_cast<size_t>(packed_batch_stride)
+                                : static_cast<size_t>(SafeInt<size_t>(num_heads_) * sequence_length * head_size);
       args.k_cache = present_key_data;
       args.v_cache = present_value_data;
       args.k_scale = k_scale;
@@ -874,7 +877,11 @@ class GQAAttentionBase {
         args.buffer_size_per_thread = buffer_size_per_thread;
 
         // Offset Q and output for this batch
-        args.query = Q + static_cast<size_t>(b) * num_heads_ * sequence_length * head_size;
+        const ptrdiff_t q_batch_stride_elems = packed_batch_stride > 0
+                                                   ? packed_batch_stride
+                                                   : static_cast<ptrdiff_t>(SafeInt<ptrdiff_t>(num_heads_) * sequence_length * head_size);
+        args.query = Q + static_cast<size_t>(b) * static_cast<size_t>(q_batch_stride_elems);
+        args.q_batch_stride = SafeInt<size_t>(num_heads_) * sequence_length * head_size;
         args.k_cache = present_key_data +
                        static_cast<size_t>(b) * kv_num_heads_ * seqlen_present_kv_cache * packed_row_bytes;
         args.v_cache = present_value_data +
@@ -884,10 +891,13 @@ class GQAAttentionBase {
         args.output = output->MutableData<float>() +
                       static_cast<size_t>(b) * sequence_length * hidden_size;
 
-        // Slice attention bias for this batch (the kernel sees batch_size=1, so batch_idx=0 inside)
+        // Slice attention bias for this batch (the kernel sees batch_size=1, so batch_idx=0 inside).
+        // Bias shape is [batch|1, num_heads|1, S, T]; the batch stride uses the actual head
+        // extent (1 when the head dim is broadcast).
         const float* batch_bias = attention_bias_data;
         if (attention_bias_data != nullptr && !attention_bias_broadcast_batch) {
-          batch_bias += static_cast<size_t>(b) * num_heads_ * sequence_length * attention_bias_seqlen_stride;
+          const size_t bias_head_extent = attention_bias_broadcast_head ? 1 : static_cast<size_t>(num_heads_);
+          batch_bias += static_cast<size_t>(b) * bias_head_extent * sequence_length * attention_bias_seqlen_stride;
         }
         args.attention_bias = batch_bias;
         args.attention_bias_seqlen_stride = attention_bias_seqlen_stride;

onnxruntime/core/mlas/inc/mlas_qkv_quant.h

@@ -271,6 +271,8 @@ struct MlasFlashAttentionQuantizedKVArgs {
     size_t buffer_size_per_thread;
 
     const float* query;      // [B, N, S, H] FP32
+    size_t q_batch_stride;   // element stride between consecutive batches in `query`
+                             // (num_heads*S*H for unpacked, (num_heads+2*kv_num_heads)*S*H for packed QKV)
     const uint8_t* k_cache;  // [B, kv_N, seqlen_present, packed_row_bytes] quantized
     const uint8_t* v_cache;  // [B, kv_N, seqlen_present, packed_row_bytes] quantized
     const float* k_scale;    // Scalar or per-channel scales for K

onnxruntime/core/mlas/lib/flashattn_qkv.cpp

@@ -127,10 +127,12 @@ MlasFlashAttentionQuantizedKVThreaded(
             ? args->v_scale + kv_head_idx * static_cast<size_t>(head_size)
             : args->v_scale;
 
-        // Q pointer: layout [batch, num_heads, seq, head_size] or packed
+        // Q pointer: layout [batch, num_heads, seq, head_size]. The batch stride is
+        // supplied separately (args->q_batch_stride) so the kernel works with both the
+        // standard BNSH layout and packed-QKV input where Q/K/V are interleaved per batch.
         const float* q_ptr = args->query +
-            (static_cast<size_t>(batch_idx) * static_cast<size_t>(num_heads) +
-             static_cast<size_t>(head_idx)) * static_cast<size_t>(sequence_length) * static_cast<size_t>(head_size) +
+            static_cast<size_t>(batch_idx) * args->q_batch_stride +
+            static_cast<size_t>(head_idx) * static_cast<size_t>(sequence_length) * static_cast<size_t>(head_size) +
             static_cast<size_t>(q_idx) * static_cast<size_t>(head_size);
 
         // Iterate over KV blocks
@@ -162,10 +164,14 @@ MlasFlashAttentionQuantizedKVThreaded(
                     static_cast<ptrdiff_t>(args->attention_bias_seqlen_stride);
                 const ptrdiff_t bias_matrix_size =
                     static_cast<ptrdiff_t>(sequence_length) * bias_seqlen_stride;
+                // The bias tensor has shape [batch|1, num_heads|1, S, T]; the batch
+                // stride uses the actual head extent (1 when the head dim is broadcast).
+                const ptrdiff_t bias_head_extent =
+                    args->attention_bias_broadcast_head ? 1 : static_cast<ptrdiff_t>(num_heads);
                 ptrdiff_t bias_offset = 0;
                 if (!args->attention_bias_broadcast_batch) {
                     bias_offset += static_cast<ptrdiff_t>(batch_idx) *
-                                   static_cast<ptrdiff_t>(num_heads) * bias_matrix_size;
+                                   bias_head_extent * bias_matrix_size;
                 }
                 if (!args->attention_bias_broadcast_head) {
                     bias_offset += static_cast<ptrdiff_t>(head_idx) * bias_matrix_size;
@@ -378,10 +384,11 @@ MlasFlashDecodingQuantizedKVThreaded(
             ? args->v_scale + kv_head_idx * static_cast<size_t>(head_size)
             : args->v_scale;
 
-        // Q pointer: layout [batch, num_heads, 1, head_size] (sequence_length=1)
+        // Q pointer: layout [batch, num_heads, 1, head_size] (sequence_length=1).
+        // The batch stride is supplied separately to support packed-QKV input.
         const float* q_ptr = args->query +
-            (static_cast<size_t>(batch_idx) * static_cast<size_t>(num_heads) +
-             static_cast<size_t>(head_idx)) * static_cast<size_t>(head_size);
+            static_cast<size_t>(batch_idx) * args->q_batch_stride +
+            static_cast<size_t>(head_idx) * static_cast<size_t>(head_size);
 
         // Step 1: QK^T GEMM for this KV chunk
         const uint8_t* k_block = k_cache_head + static_cast<size_t>(ir) * packed_row_bytes;
@@ -405,10 +412,14 @@ MlasFlashDecodingQuantizedKVThreaded(
             const ptrdiff_t bias_seqlen_stride =
                 static_cast<ptrdiff_t>(args->attention_bias_seqlen_stride);
             const ptrdiff_t bias_matrix_size = bias_seqlen_stride;  // S=1
+            // The bias tensor has shape [batch|1, num_heads|1, S, T]; the batch stride
+            // uses the actual head extent (1 when the head dim is broadcast).
+            const ptrdiff_t bias_head_extent =
+                args->attention_bias_broadcast_head ? 1 : static_cast<ptrdiff_t>(num_heads);
             ptrdiff_t bias_offset = 0;
             if (!args->attention_bias_broadcast_batch) {
                 bias_offset += static_cast<ptrdiff_t>(batch_idx) *
-                               static_cast<ptrdiff_t>(num_heads) * bias_matrix_size;
+                               bias_head_extent * bias_matrix_size;
             }
             if (!args->attention_bias_broadcast_head) {
                 bias_offset += static_cast<ptrdiff_t>(head_idx) * bias_matrix_size;

onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py

@@ -830,6 +830,24 @@ def test_int8_bias_broadcast_head(self):
             bias_broadcast_head=True,
         )
 
+    def test_int8_bias_broadcast_head_multi_batch(self):
+        """Bias shape [B, 1, S, T] with batch_size > 1 and num_heads > 1.
+
+        Regression test: the bias batch stride must use the head extent (1 when the
+        head dimension is broadcast), not num_heads. With batch_size == 1 the bug is
+        masked because batch_idx is always 0.
+        """
+        run_quantized_gqa_bias_test(
+            batch_size=3,
+            seq_len=8,
+            num_heads=4,
+            kv_num_heads=2,
+            head_size=16,
+            quant_type="PER_TENSOR",
+            bit_width=8,
+            bias_broadcast_head=True,
+        )
+
     def test_int8_bias_broadcast_both(self):
         """Bias shape [1, 1, S, T] with batch_size > 1 and num_heads > 1."""
         run_quantized_gqa_bias_test(