[CPU] Add FP32 GEMV decode kernel for GroupQueryAttention

docs/contrib_ops/cpu/gqa.md

@@ -58,7 +58,7 @@ Both the non-quantized and quantized paths have two execution strategies:
 - **Naive (full materialization)**: Computes the full `[S, T]` attention score matrix, applies masking and softmax, then computes the SV product. Simple but memory-intensive for long sequences.
 - **Flash Attention (tiled, online softmax)**: Processes K/V in L2-cache-sized blocks using the online softmax algorithm (Milakov & Gimelshein, 2018). Avoids materializing the full attention matrix, reducing peak memory from O(S×T) to O(S×Bc) per head. Multi-threaded via the MLAS thread pool.
 
-The quantized path uses `MlasFlashAttentionQuantizedKV` (`flashattn_qkv.cpp`); the non-quantized FP32 path uses `MlasFlashAttentionGQA` (`flashattn_gqa.cpp`). Both share the same tiling, masking, and online-softmax structure. The quantized path additionally provides a two-phase flash-decoding strategy for single-token decode; the non-quantized FP32 path is limited to prefill (`sequence_length > 1`) and uses the naive path for decode.
+The quantized path uses `MlasFlashAttentionQuantizedKV` (`flashattn_qkv.cpp`); the non-quantized FP32 path uses `MlasFlashAttentionGQA` (`flashattn_gqa.cpp`). Both share the same tiling, masking, online-softmax, and flash-decoding structure.
 
 The flash path is selected by default when conditions are met (see below). Set `ORT_GQA_DISABLE_FLASH_ATTENTION=1` to force the naive path (applies to both the quantized and non-quantized paths).
 
@@ -213,7 +213,7 @@ The partials buffer is allocated alongside the per-thread scratch in a single al
 
 ## Non-Quantized Flash Attention Path
 
-The non-quantized flash attention path (`MlasFlashAttentionGQA`, in `flashattn_gqa.cpp`) is the FP32-KV-cache counterpart of the quantized path. It is selected for the `float` kernel specialization and reuses the same tiling, online-softmax, and masking structure. Unlike the quantized path, it is limited to prefill / chunked-prefill (`sequence_length > 1`); single-token decode (`sequence_length == 1`) uses the naive path, which is why there is no flash-decoding variant here.
+The non-quantized flash attention path (`MlasFlashAttentionGQA`, in `flashattn_gqa.cpp`) is the FP32-KV-cache counterpart of the quantized path. It is selected for the `float` kernel specialization and reuses the same tiling, online-softmax, masking, and flash-decoding structure.
 
 ### Differences from the Quantized Path
 
@@ -242,26 +242,61 @@ computes a per-q_block bound
 `ir >= kv_causal_limit`, instead of computing and then discarding the masked upper-triangle
 QK/SV GEMMs. This skips roughly half of the QK/SV work for square prefill (S = T) and is the
 main reason the FP32 flash path is faster than naive even at short sequence lengths
-(see the benchmark results below).
+(see the benchmark results below). Decode (q_block of size 1 at the cache tail) attends to all
+KV positions, so the bound equals `total_seqlen` and nothing is skipped.
 
 ### Activation Conditions
 
 The non-quantized flash path is selected when ALL of the following hold:
 
 - The kernel specialization is `float` (FP16 uses the naive path)
 - `ORT_GQA_DISABLE_FLASH_ATTENTION` environment variable is not set (or set to `0`)
-- `sequence_length > 1` (prefill / chunked-prefill; single-token decode uses the naive path)
+- `total_sequence_length > 1`
 - No softcap
 - No smooth softmax
 - No head sink
 - No output QK capture
 - `present_key` and `present_value` are provided
 
-Attention bias, causal masking, local window attention, GQA head grouping (`num_heads != kv_num_heads`), ragged per-batch sequence lengths, and shared past/present buffers are all supported, mirroring the quantized flash path. When any condition is not met, the kernel falls back to the naive full-materialization path.
+Attention bias, causal masking, local window attention, GQA head grouping (`num_heads != kv_num_heads`), ragged per-batch sequence lengths, shared past/present buffers, and flash decoding are all supported, mirroring the quantized flash path. When any condition is not met, the kernel falls back to the naive full-materialization path.
+
+### Block Sizes, Threading, and Flash Decoding
+
+Block-size selection (`kv_block_size`, `q_block_size`), `(batch, head, q_block)` task partitioning, the per-thread working buffer layout (`l`, `m`, `scores`, `temp_output`), and the two-phase flash-decoding strategy for single-token decode are identical to the quantized path described above. The only difference is that the per-thread `temp_output` tile is accumulated directly by the SV SGEMM rather than via a fused dequantization.
+
+#### Decode uses a dedicated GEMV kernel (`sequence_length == 1`)
+
+The tiled online-softmax SGEMM kernel (`MlasFlashAttentionGQAThreaded`) is used **only for
+prefill** (`sequence_length > 1`), where each KV tile is reused across the `q_block_size`
+query rows and tiling delivers real cache-locality and SGEMM packing benefits.
+
+For single-token decode the query tile has `M = 1`, so every K/V element is streamed
+exactly once with no reuse across query rows. Tiling provides **no** cache-locality
+benefit, and routing the `1 × T × H` work through `MlasSgemmOperation` pays the SGEMM
+B-packing/setup cost on every call — which previously made the flash decode path *slower*
+than the naive path (≈0.4–0.6x) for short-to-medium total sequence lengths.
 
-### Block Sizes and Threading
+Decode is therefore handled by a dedicated GEMV kernel (`MlasGQADecodeGQAThreaded`),
+dispatched whenever `sequence_length == 1` and flash decoding is not active. It
+parallelizes over `(batch, head)` and, per head, computes the attention directly with two
+matrix-vector products and a two-pass softmax:
+
+- **QK GEMV** — `scores[t] = scale · dot(q, K[t])` for `t ∈ [0, total_seqlen)`.
+- two-pass softmax over `scores` using the dispatched `ReduceMaximumF32Kernel` /
+  `ComputeSumExpF32Kernel` helpers.
+- **SV GEMV** — `out[h] = Σ_t probs[t] · V[t][h]`, then normalize by `1/Σ probs`.
+
+Both GEMV helpers (`MlasGQADecodeQK`, `MlasGQADecodeSV`) live in the baseline-ISA MLAS
+translation unit, so their inner loops use independent accumulator lanes / map-style
+updates that vectorize under SSE2 without `-ffast-math`. Decode needs no causal mask (the
+single new token is the most recent position and attends to every cached token); only
+optional local-window masking and additive attention bias are applied. The kernel streams
+K and V exactly once each, so it is memory-bandwidth bound.
+
+The two-phase flash-decoding path (active when `batch × heads < threads`, KV partitioned
+across idle threads) now also uses these GEMV helpers for its per-chunk QK and SV products
+instead of `M = 1` SGEMM calls, removing the same packing overhead.
 
-Block-size selection (`kv_block_size`, `q_block_size`), `(batch, head, q_block)` task partitioning, and the per-thread working buffer layout (`l`, `m`, `scores`, `temp_output`) are identical to the quantized path described above. The only difference is that the per-thread `temp_output` tile is accumulated directly by the SV SGEMM rather than via a fused dequantization. Because this path is prefill-only, it does not include the quantized path's two-phase flash-decoding strategy for single-token decode.
 
 ## MLAS Dispatch Paths
 
@@ -513,13 +548,31 @@ offsets the intrinsic per-KV-block online-softmax overhead (running max/exp/outp
 The same advantage holds single-threaded (1.4\u20131.8x at threads=1), confirming the gain is
 algorithmic rather than purely from threading.
 
-#### Decode (S = 1, token generation)
+#### Latency — Decode (S = 1, token generation)
+
+For single-token decode at this head configuration (`batch\u00d7heads = 16 > threads = 8`, so
+flash decoding KV-partitioning is not active), the workload per `Run` is tiny (a `1 × T × H`
+GEMV pair per head) and operator-level latency is dominated by fixed per-`Run` overhead
+(session dispatch, KV-cache concatenation), so operator-level measurements on the EPYC dev
+box are extremely noisy. The numbers below come from a min-of-many-repeats MLAS-path harness
+to suppress that jitter.
 
-Single-token decode (`sequence_length == 1`) is **not** handled by the FP32 flash path; it falls
-back to the naive path. Decode produces only a `[1, total_sequence_length]` score row per head,
-so there is nothing to tile away, and the extra online-softmax bookkeeping made the flash kernel
-slower and noisier in practice. Restricting the flash path to prefill (`sequence_length > 1`) keeps
-the consistent prefill win without regressing decode.
+| Total Seqlen | Naive (ms) | Flash (ms) | Speedup |
+|---:|---:|---:|---:|
+| 513  | 0.50 | 0.42 | ~1.0\u20131.2x (noisy) |
+| 1025 | 0.78 | 0.69 | ~1.0\u20131.1x (noisy) |
+| 2049 | 1.89 | 1.73 | ~1.0\u20131.1x (noisy) |
+| 4097 | 6.1  | 4.5  | 1.35\u20131.5x |
+
+Decode is now handled by the dedicated GEMV kernel (`MlasGQADecodeGQAThreaded`) instead of
+the prefill tiling kernel; see *Decode uses a dedicated GEMV kernel* above. Replacing the
+per-head `M = 1` `MlasSgemmOperation` QK/SV calls with direct GEMVs removes the SGEMM
+B-packing overhead that previously made flash decode noticeably **slower** than naive
+(measured ≈0.4\u20130.6x across all lengths before the change). Flash decode is now at parity
+for short/medium sequences (where the work is memory-bandwidth bound and overhead-dominated)
+and consistently ahead for long contexts (T≥4097, ~1.4\u20131.5x) where the streamed
+single-pass KV access wins. Short decode remains overhead-bound rather than algorithm-bound,
+so it is not the target of the prefill-oriented causal early-termination optimization.
 ## Current CPU Limitations
 
 The current CPU GroupQueryAttention implementation has a few important limitations:

onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h

@@ -1091,16 +1091,49 @@ class GQAAttentionBase {
     int thread_count = concurrency::ThreadPool::DegreeOfParallelism(tp);
     thread_count = std::max(thread_count, 1);
 
-    // Per-thread scratch: l + m + scores[q_block_size * kv_block_size] + temp_output[q_block_size * head_size]
-    const size_t buffer_size_per_thread =
-        (SafeInt<size_t>(q_block_size) * 2 +              // l + m
-         SafeInt<size_t>(q_block_size) * kv_block_size +  // scores
-         SafeInt<size_t>(q_block_size) * head_size) *     // temp_output
-        sizeof(float);
-    size_t total_buffer_bytes = SafeInt<size_t>(buffer_size_per_thread) * thread_count;
+    // Flash decoding: for decode (sequence_length==1), partition KV across threads
+    // to improve parallelism when batch*heads < thread_count. This KV-split is only
+    // wired into the unified kernel (common_past_seqlen >= 0); the ragged/per-batch
+    // fallback runs the single-pass decode kernel instead, which needs a larger
+    // per-thread scratch (scores[total_seqlen] + temp_output[head_size]). Gating on
+    // common_past_seqlen >= 0 keeps the per-thread buffer sizing below consistent
+    // with the kernel that actually runs.
+    const int kv_chunk_count = (max_total_seqlen + kv_block_size - 1) / kv_block_size;
+    const bool use_flash_decoding = (sequence_length == 1 &&
+                                     common_past_seqlen >= 0 &&
+                                     batch_size * num_heads_ < thread_count &&
+                                     kv_chunk_count > 1);
+
+    size_t buffer_size_per_thread;
+    size_t partials_buffer_bytes = 0;
+    if (use_flash_decoding) {
+      // Flash decoding: per-thread scratch only needs scores[kv_block_size]
+      buffer_size_per_thread = SafeInt<size_t>(kv_block_size) * sizeof(float);
+      // Partials: [batch * num_heads * kv_chunk_count * (2 + head_size)] floats
+      partials_buffer_bytes = SafeInt<size_t>(batch_size) * num_heads_ *
+                              kv_chunk_count * (2 + head_size) * sizeof(float);
+    } else if (sequence_length == 1) {
+      // Decode (GEMV kernel, no Q/KV tiling): per-thread scratch holds the full
+      // score row scores[total_seqlen] plus a temp output accumulator[head_size].
+      buffer_size_per_thread =
+          (SafeInt<size_t>(max_total_seqlen) + head_size) * sizeof(float);
+    } else {
+      buffer_size_per_thread =
+          (SafeInt<size_t>(q_block_size) * 2 +              // l + m
+           SafeInt<size_t>(q_block_size) * kv_block_size +  // scores
+           SafeInt<size_t>(q_block_size) * head_size) *     // temp_output
+          sizeof(float);
+    }
+    size_t total_buffer_bytes = SafeInt<size_t>(buffer_size_per_thread) * thread_count + partials_buffer_bytes;
     auto flash_buffer_alloc = allocator->Alloc(total_buffer_bytes);
     BufferUniquePtr flash_buffer(flash_buffer_alloc, BufferDeleter(allocator));
 
+    // Partials buffer is placed after per-thread scratch
+    float* partials_ptr = use_flash_decoding
+                              ? reinterpret_cast<float*>(reinterpret_cast<char*>(flash_buffer_alloc) +
+                                                         buffer_size_per_thread * thread_count)
+                              : nullptr;
+
     const float scale = scale_ == 0.0f ? 1.0f / sqrt(static_cast<float>(head_size)) : scale_;
 
     // If all batch items share the same past_seqlen, use the unified flash kernel.
@@ -1133,6 +1166,8 @@ class GQAAttentionBase {
       args.attention_bias_seqlen_stride = attention_bias_seqlen_stride;
       args.attention_bias_broadcast_batch = attention_bias_broadcast_batch;
       args.attention_bias_broadcast_head = attention_bias_broadcast_head;
+      args.flash_decoding_partials = partials_ptr;
+      args.kv_chunk_count = kv_chunk_count;
 
       MlasFlashAttentionGQA(&args, tp);
     } else {
@@ -1185,6 +1220,8 @@ class GQAAttentionBase {
         args.attention_bias_seqlen_stride = attention_bias_seqlen_stride;
         args.attention_bias_broadcast_batch = true;  // batch offset handled above
         args.attention_bias_broadcast_head = attention_bias_broadcast_head;
+        args.flash_decoding_partials = nullptr;  // per-batch doesn't use flash decoding
+        args.kv_chunk_count = 0;
 
         MlasFlashAttentionGQA(&args, tp);
       }

onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc

@@ -348,13 +348,13 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
   // 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<T, float>) {
-    // Restrict the flash path to prefill / chunked-prefill (query length > 1). Single-token
-    // decode (sequence_length == 1) has no flash benefit: the naive score matrix is only
-    // [1, total_sequence_length] per head, so there is nothing to tile away, and the extra
-    // online-softmax bookkeeping makes it slower in practice.
     const bool use_flash = !disable_gqa_flash_ &&
-                           parameters.sequence_length > 1 &&
+                           parameters.total_sequence_length > 1 &&
                            softcap_ == 0.0f &&
                            !use_smooth_softmax_ &&
                            head_sink_data == nullptr &&

onnxruntime/core/mlas/inc/mlas.h

@@ -2302,9 +2302,9 @@ MlasFlashAttention(
 //
 // Adapts the online-softmax tiled algorithm to operate on an FP32 present
 // K/V cache laid out as BNSH ([batch, kv_num_heads, seqlen_present, head_size]).
-// Supports GQA head grouping, causal masking, local window attention, and
-// additive attention bias. Intended for prefill / chunked-prefill
-// (sequence_length > 1).
+// Supports GQA head grouping, causal masking, local window attention,
+// additive attention bias, and an optional flash-decoding split over the KV
+// sequence dimension for the single-token decode case.
 //
 struct MlasFlashAttentionGQAArgs {
     int batch_size;
@@ -2320,7 +2320,7 @@ struct MlasFlashAttentionGQAArgs {
     int kv_block_size;        // key/value tile size (Bc)
     float scale;              // QK scaling factor
     int thread_count;         // number of partitions / threads
-    float* buffer;            // per-thread scratch
+    float* buffer;            // per-thread scratch (+ optional flash-decoding partials)
     size_t buffer_size_per_thread;
 
     const float* query;       // [batch, num_heads, sequence_length, head_size] BNSH
@@ -2334,6 +2334,12 @@ struct MlasFlashAttentionGQAArgs {
     int attention_bias_seqlen_stride;
     bool attention_bias_broadcast_batch;
     bool attention_bias_broadcast_head;
+
+    // Flash decoding (sequence_length == 1): partition KV across threads.
+    // Set flash_decoding_partials != nullptr to enable; otherwise the standard
+    // per-(batch, head, q_block) partitioning is used.
+    float* flash_decoding_partials;
+    int kv_chunk_count;
 };
 
 /**