Speed up CPU GroupQueryAttention decode with a grouped flash kernel

onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc

@@ -797,17 +797,24 @@ Status GroupQueryAttention<T>::Compute(OpKernelContext* context) const {
     const T* k_data = packed_qkv ? nullptr : k_rotary;
     const T* v_data = packed_qkv ? nullptr : V.Get<Tensor>().Data<T>();
 
-    // 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.
     const bool read_only_cache = kv_sequence_length == 0 && attention_past_key != nullptr &&
                                  attention_past_value != nullptr && attention_present_key == nullptr &&
                                  attention_present_value == nullptr;
+
+    // fp32 and fp16 caches: all query heads of a KV head share each cache tile, the softmax stays in fp32, and decode
+    // splits the cache across threads when there are fewer (batch, KV head) pairs than threads. Steps it doesn't take
+    // go to the fp32 flash path or the per-head path below.
+    if (CanApplyAttentionGrouped<T>(output_qk != nullptr,
+                                    attention_present_key != nullptr && attention_present_value != nullptr,
+                                    read_only_cache, parameters.sequence_length)) {
+      return ApplyAttentionGrouped(q_rotary, k_data, v_data, head_sink_data,
+                                   attention_bias, attention_bias_offsets, attention_past_key, attention_past_value,
+                                   output, attention_present_key, attention_present_value, attention_seqlens_k,
+                                   parameters, allocator, context, attention_past_seqlens, query_seqlens_data);
+    }
+
+    // fp32 flash path: MLAS's tiled online-softmax kernel for prefill and its GEMV kernel for single-token decode, which
+    // off x64 is still faster than the grouped kernel when every query head has its own KV head.
     if constexpr (std::is_same_v<T, float>) {
       const bool use_flash = !disable_gqa_flash_ &&
                              is_unidirectional_ &&

onnxruntime/contrib_ops/cpu/bert/grouped_flash_attention.h

@@ -0,0 +1,87 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <cstddef>
+
+#include "core/common/status.h"
+#include "core/framework/allocator.h"
+#include "core/mlas/inc/mlas.h"
+#include "core/platform/threadpool.h"
+
+namespace onnxruntime {
+namespace contrib {
+
+// Attention of query heads that share key/value heads (grouped-query attention), over a key/value cache in fp32 or
+// fp16. Strides are in elements. The kernel only reads the cache: appending new rows is the caller's job.
+template <typename T>
+struct GroupedFlashAttentionArgs {
+  int batch_size = 0;
+  int sequence_length = 0;  // query tokens per batch entry
+  int num_heads = 0;        // query heads, a multiple of kv_num_heads
+  int kv_num_heads = 0;
+  int head_size = 0;    // width of Q and K
+  int v_head_size = 0;  // width of V and of the output
+
+  // Query (b, head h, token s) starts at query + b * q_batch_stride + h * q_head_stride + s * q_seq_stride.
+  const T* query = nullptr;
+  size_t q_batch_stride = 0;
+  size_t q_head_stride = 0;
+  size_t q_seq_stride = 0;
+
+  // Cached row t of (b, kv head k) starts at key + b * k_batch_stride + k * k_head_stride + t * k_row_stride, and
+  // likewise for value. Query head h reads kv head h / (num_heads / kv_num_heads).
+  const T* key = nullptr;
+  size_t k_batch_stride = 0;
+  size_t k_head_stride = 0;
+  size_t k_row_stride = 0;
+  const T* value = nullptr;
+  size_t v_batch_stride = 0;
+  size_t v_head_stride = 0;
+  size_t v_row_stride = 0;
+
+  // Output (b, token s, head h) starts at output + b * o_batch_stride + s * o_seq_stride + h * o_head_stride.
+  T* output = nullptr;
+  size_t o_batch_stride = 0;
+  size_t o_seq_stride = 0;
+  size_t o_head_stride = 0;
+
+  // Per batch entry. Query s attends to cached rows [0, min(causal_offsets[b] + s + 1, total_lengths[b])) when causal,
+  // else to [0, total_lengths[b]). Only the first query_lengths[b] tokens are computed when query_lengths is given;
+  // the output rows of the others are left untouched.
+  const int* total_lengths = nullptr;
+  const int* causal_offsets = nullptr;
+  const int* query_lengths = nullptr;
+  bool causal = true;
+  // When positive, a query sees only the last local_window_size of the rows it would otherwise see.
+  int local_window_size = -1;
+
+  float scale = 1.0f;
+  float softcap = 0.0f;  // when positive, scores become softcap * tanh(scores / softcap) before the bias
+
+  // Optional additive bias for (b, h, s, cached row t) at
+  // attention_bias + b * bias_batch_stride + h * bias_head_stride + s * bias_seq_stride + bias_offsets[b] + t.
+  const T* attention_bias = nullptr;
+  size_t bias_batch_stride = 0;
+  size_t bias_head_stride = 0;
+  size_t bias_seq_stride = 0;
+  const int* bias_offsets = nullptr;
+
+  // Smooth softmax: one extra logit per head (head_sink[h], or 0 without head_sink) joins the softmax denominator.
+  bool use_sink = false;
+  const T* head_sink = nullptr;
+
+  const MLAS_BACKEND_KERNEL_SELECTOR_CONFIG* mlas_config = nullptr;
+};
+
+// Runs the attention with all query heads of a kv head stacked as the rows of one GEMM against each cache tile, so a
+// tile is read (and, for fp16, converted to fp32) once per group instead of once per query head. The softmax runs
+// online in fp32 tile by tile, and when there are fewer row blocks than threads (decode), the cache is also split
+// into chunks whose partial softmax states are merged with a log-sum-exp reduction.
+template <typename T>
+Status RunGroupedFlashAttention(const GroupedFlashAttentionArgs<T>& args, AllocatorPtr allocator,
+                                concurrency::ThreadPool* thread_pool);
+
+}  // namespace contrib
+}  // namespace onnxruntime