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

[tianleiwu]

Description

The quantized KV-cache flash-attention path for the CPU GroupQueryAttention contrib op carried two latent batch-stride bugs that produced incorrect results for batch_size > 1. This PR ports the fixes already landed for the FP32 path (PR #28962) to the quantized path, and adds a regression test that exercises the previously-uncovered scenario.

Summary of Changes

Bug fixes

File	Change
onnxruntime/core/mlas/inc/mlas_qkv_quant.h	Add q_batch_stride field to MlasFlashAttentionQuantizedKVArgs so the kernel uses a caller-supplied Q batch stride instead of assuming the unpacked num_heads*S*H layout.
onnxruntime/core/mlas/lib/flashattn_qkv.cpp	Use args->q_batch_stride for the Q pointer in both the main tiled kernel and the flash-decoding kernel; compute the attention-bias batch stride from bias_head_extent = broadcast_head ? 1 : num_heads (two sites) instead of always using num_heads.
onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h	Set q_batch_stride to (num_heads + 2*kv_num_heads)*S*H for packed QKV (else num_heads*S*H) in the unified dispatch; correct the per-batch Q offset and bias slice to use the packed stride and head extent.

Tests

• Add test_int8_bias_broadcast_head_multi_batch to test_gqa_cpu_quantized.py covering [B, 1, S, T] bias with batch_size > 1. The existing test_int8_bias_broadcast_head used batch_size == 1, which masked the head-broadcast batch-stride bug.

Testing

• cd onnxruntime/test/python/transformers && python -m pytest test_gqa_cpu_quantized.py -q → 21 passed, 2 skipped.
• Verified the new test catches the bug: temporarily reverting the bias-stride fix makes test_int8_bias_broadcast_head_multi_batch fail (mismatches starting at batch index 1) while the batch_size == 1 variant still passes; restoring the fix makes all tests pass.
• lintrunner -a → no lint issues.

Motivation and Context

Follow-up to PR #28962, which fixed the identical two bugs (packed-QKV Q batch stride and attention-bias head-broadcast batch stride) in the non-quantized FP32 CPU GQA path. The quantized path was left untouched there as out of scope; existing quantized parity tests did not cover batch > 1 with [B, 1, S, T] bias, so the bugs went undetected.

[Copilot]

Pull request overview

This PR fixes two correctness issues in the CPU quantized KV-cache flash-attention path for the com.microsoft.GroupQueryAttention contrib op that can produce wrong results when batch_size > 1: (1) query (Q) batch stride handling for packed QKV layouts, and (2) attention-bias batch stride computation when the head dimension is broadcast (bias shape [B, 1, S, T]). It also adds a Python regression test to cover the previously-missed multi-batch broadcast-head bias case.

Changes:

• Add q_batch_stride to MlasFlashAttentionQuantizedKVArgs and use it in the quantized flash-attention kernels to support both standard and packed-QKV layouts.
• Fix attention-bias batch-stride calculation in quantized flash attention to account for head broadcasting (bias_head_extent = 1 when broadcast).
• Add a new quantized Python regression test for [B, 1, S, T] bias with batch_size > 1.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
onnxruntime/core/mlas/inc/mlas_qkv_quant.h	Extends quantized flash-attention args with q_batch_stride for correct Q addressing across batches/layouts.
onnxruntime/core/mlas/lib/flashattn_qkv.cpp	Uses caller-provided Q batch stride and fixes bias batch stride when head dim is broadcast (tiled + decoding paths).
onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h	Wires q_batch_stride from the op dispatch and fixes per-batch bias slicing stride for broadcast-head bias.
onnxruntime/test/python/transformers/test_gqa_cpu_quantized.py	Adds a multi-batch regression test for broadcast-head attention bias.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 1 comment.

[yuslepukhin]

Issues Found

Minor: q_batch_stride in unified path not using SafeInt — In gqa_attention_base.h line 834, the packed case uses static_cast<size_t>(packed_batch_stride) without SafeInt, while the unpacked case correctly uses SafeInt<size_t>(num_heads_) * sequence_length * head_size. The packed_batch_stride is already computed with SafeInt elsewhere (so it's fine in practice), but it's inconsistent. Very minor.

Style nit: Line 900 is fairly long (batch_bias += static_cast<size_t>(SafeInt<size_t>(b) * bias_head_extent * sequence_length * attention_bias_seqlen_stride);). Could benefit from breaking across lines, but this is purely cosmetic.

Negative tests / edge cases:

No explicit negative tests (e.g., invalid input shapes). This is acceptable — the PR fixes stride bugs, not input validation.
The existing test_int8_bias_broadcast_head (batch_size=1) continues to pass and serves as the baseline.

Test gap: No test for packed-QKV combined with broadcast-head bias and batch_size > 1. This would exercise both fixes simultaneously. However, the two bugs are independent (different code paths), so separate tests are adequate.

Approve. The PR correctly fixes two real bugs that produce incorrect attention output for batch_size > 1 in the quantized GQA flash-attention path. The fixes are the exact counterpart of what was already landed in PR #28962 for the FP32 path. Arithmetic uses SafeInt where it matters, all review comments have been addressed, and the new regression tests directly target the fixed bugs. No memory safety, OOB, or exception safety concerns.

[titaiwangms]

Review — multi-model review team

I ran a parallel review (readability, code, correctness, deep-semantic, and cross-module integration). Overall: the fix is correct, complete, and well-tested. No blocking issues found. Findings below are all Minor/Nit or optional follow-ups.

Correctness — verified ✅

• Bias head-extent fix is right at both sites (MlasFlashAttentionQuantizedKVThreaded and MlasFlashDecodingQuantizedKVThreaded) and mirrored in the per-batch caller slice in gqa_attention_base.h. For a [B|1, N|1, S, T] row-major bias, the batch stride is head_extent·S·T with head_extent = broadcast_head ? 1 : num_heads; the old num_heads over-strode by a factor of num_heads exactly when broadcast_head && !broadcast_batch && batch>1. No double-counting with the kernel's internal head_idx·S·T term.
• Packed-QKV Q batch stride is correct. Input reaches the kernel as BNSH (MaybeTransposeToBNSH(..., num_heads_ + 2*kv_num_heads_, ...) in group_query_attention.cc), so Q heads are contiguous in the first N·S·H of each (N+2·kvN)·S·H packed batch slice. batch_idx·q_batch_stride + head_idx·S·H matches the established invariant already used by the non-flash path in the same file.
• Parity with Add flash attention for non-quantized CPU GroupQueryAttention #28962 holds; kernel formulas are identical. The only divergence (per-batch loop sets q_batch_stride to the packed-aware value vs. Add flash attention for non-quantized CPU GroupQueryAttention #28962's num_heads*S*H) is harmless — that field is multiplied by batch_idx==0 in the per-batch path.

Cross-module integration — verified ✅

Every consumer of MlasFlashAttentionQuantizedKVArgs (gqa_attention_base.h, bench_qkv_quant.cpp, test_qkv_quant.cpp) was updated to set q_batch_stride. The new default member initializers make a missing assignment default to a deterministic 0 (and bias-broadcast flags to true = safe no-op addressing) rather than UB, and don't break the existing aggregate args{} init in the benchmark.

Minor

1. Unit-test strides hardcode the num_heads == 1 fixture. test_qkv_quant.cpp sets args.q_batch_stride = seq_len * head_size; and = head_size;, omitting num_heads. Correct today only because those fixtures use num_heads == 1/batch_size == 1. If generalized later they'd silently test the wrong stride and could mask the very bug this PR fixes. Suggest the structural form static_cast<size_t>(args.num_heads) * seq_len * head_size (and * head_size for decode), or at least a // num_heads == 1 in this fixture comment.
2. No C++/Python regression for the decode path. The PR also fixes MlasFlashDecodingQuantizedKVThreaded (Q ptr + bias offset), but the new tests are all prompt-style (seq_len > 1). Existing decode unit tests are batch_size==1, num_heads==1, so the multi-batch decode stride fix is unverified. Consider a sequence_length==1, total_seqlen>1, batch_size>1 packed/broadcast-bias case.
3. Duplicated stride logic with divergent spelling in gqa_attention_base.h: the unified-batch path inlines the ternary keyed on packed_qkv, while the per-batch path introduces q_batch_stride_elems keyed on packed_batch_stride > 0. Equivalent but a reader must prove it. Hoisting one named value used by both branches would make them structurally parallel.

Nits

• q_batch_stride name doesn't encode units (elements vs bytes); the // element stride comment carries it. q_batch_stride_elems would be self-documenting.
• The bias_head_extent explanatory comment block is duplicated verbatim in the two sibling kernels — minor drift risk.

Optional follow-up (pre-existing, not introduced here)

packed_batch_stride is built as SafeInt<ptrdiff_t>(num_heads_ + 2 * kv_num_heads_) * ..., where num_heads_ + 2 * kv_num_heads_ is evaluated as plain int before SafeInt wraps it. This PR now also propagates that value into q_batch_stride. It's a pre-existing pattern (lines 311/650/940) and only reachable with adversarial billions-scale head attributes, so it's out of scope here — but worth hardening separately (e.g. SafeInt<ptrdiff_t>(num_heads_) + SafeInt<ptrdiff_t>(2) * kv_num_heads_).

---

Reviewed by a multi-model team (Claude Opus, GPT-5.x, Gemini 3 Pro). All correctness/invariant claims grounded in the source; no test runs were needed for the addressing math.