Add CUDA Attention spec-coverage tests (GQA asymmetric head-size, output_qk, fp32 softcap+mask ordering)

.agents/skills/cuda-attention-kernel-patterns/SKILL.md

@@ -36,7 +36,7 @@ Unified Unfused      → RunUnfusedAttention()
 
 **Flash eligibility**: fp16/bf16 only, SM≥8.0 (Ampere+), `head_size == v_head_size`, `head_size <= 256`, no `output_qk`, `attn_mask == nullptr`. Uses `mha_fwd` / `mha_fwd_kvcache`.
 
-**MEA eligibility**: SM50+/53+/80+ by dtype, `head_size <= 1024` and divisible by 8, no `output_qk`. Decode requires `head_size == v_head_size` (for `LaunchConcatNewToPastKV`). Bias stride must satisfy `total_sequence_length % 4 == 0`. GQA with FP32 is excluded (LaunchUngroup only has fp16/bf16 instantiations). Supports `softcap + attn_mask` — CUTLASS applies softcap before bias in kernel tiles, matching ONNX spec ordering (onnx/onnx#7865).
+**MEA eligibility**: SM50+/53+/80+ by dtype, `head_size <= 1024` and divisible by 8 (enforced by `has_memory_efficient_attention`), no `output_qk`. GQA additionally requires `head_size == v_head_size` (for `LaunchUngroup`); decode also requires it (for `LaunchConcatNewToPastKV`). Bias stride must satisfy `total_sequence_length % 4 == 0`. GQA with FP32 is excluded (LaunchUngroup only has fp16/bf16 instantiations). The host-side dispatch in `core/providers/cuda/llm/attention.cc` additionally enforces `head_size % 4 == 0` as a forward-looking alignment floor (Cutlass FMHA's BiasLoader uses 128-bit / sizeof_bits-element loads = 4 elements for fp32, 8 for fp16/bf16); this is redundant with the `% 8` check today but becomes load-bearing once microsoft/onnxruntime#28365 relaxes BiasLoader to use `kAlignmentA` / `DispatchIsAligned`. Supports `softcap + attn_mask` — CUTLASS applies softcap before bias in kernel tiles, matching ONNX spec ordering (onnx/onnx#7865).
 
 **Unified Unfused Attention**: Always available as the final fallback. Handles both MHA (`num_heads == kv_num_heads`, group=1) and GQA (`num_heads != kv_num_heads`, group>1) via a reshape-Q trick with stride-based cuBLAS batched GEMM (no K/V head replication). Uses FP32 QK scratch for precision. Supports all features:
 - softcap + attn_mask (spec-correct ordering)

onnxruntime/core/providers/cuda/llm/attention.cc

@@ -1383,7 +1383,21 @@ Status Attention<T>::ComputeInternal(OpKernelContext* context) const {
         (past_key == nullptr || parameters.head_size == parameters.v_head_size) &&
         // GQA+MEA requires LaunchUngroup which only has fp16/bf16 instantiations.
         // FP32 GQA must fall through to the unfused path.
-        !(is_gqa && std::is_same<T, float>::value);
+        !(is_gqa && std::is_same<T, float>::value) &&
+        // Forward-looking alignment floor: head_size must be a multiple of 4
+        // for Cutlass FMHA's BiasLoader, which uses a 128-bit / sizeof_bits
+        // element-aligned load for the Q and additive-bias tiles (= 4 elements
+        // for fp32, 8 for fp16/bf16). head_size is the inner stride of those
+        // loads.
+        //
+        // Redundant today: has_memory_efficient_attention() already requires
+        // (qk_head_size & 7) == 0 (memory_efficient_attention.h), which strictly
+        // implies %4. Kept as an explicit, dtype-agnostic alignment floor so
+        // it remains correct once microsoft/onnxruntime#28365 lands and
+        // relaxes BiasLoader to use kAlignmentA / DispatchIsAligned (at which
+        // point MEA's %8 invariant goes away). Delete this clause when MEA
+        // no longer requires the %8 invariant upstream.
+        (parameters.head_size % 4 == 0);
 
     // Cutlass FMHA requires bias strides to satisfy minimum alignment even in the
     // "unaligned" kernel path. When an attention mask is present (with or without

onnxruntime/test/python/transformers/test_onnx_attention/common.py

@@ -100,11 +100,20 @@
     config: AttentionConfig,
     ort_type,
     is_past=False,
-    output_qk: int = 0,  # CUDA does not support output_qk for GQA path
+    output_qk: int | None = None,
 ):
     """
     Create ONNX Attention op node and I/O definitions for testing.
 
+    output_qk: when set, enables the optional 4th output `output_qk` and sets
+    the `qk_matmul_output_mode` attribute to this value (0=kQK raw, 1=kQKMask,
+    2=kQKSoftCap, 3=kQKSoftMax). When None (default), the 4th output is not
+    emitted and the attribute defaults to 0.
+
+    NOTE: API contract — `output_qk=0` ENABLES the 4th output in raw-QK mode;
+    `output_qk=None` (the default) DISABLES it. Do not pass 0 expecting it to
+    mean "disabled" — pass None instead.
+
     ONNX Attention op (opset 23/24) inputs:
     - 0: Q (query) - required
     - 1: K (key) - required
@@ -142,7 +151,8 @@
         "present_value",
     ]
 
-    if output_qk > 0:
+    enable_output_qk = output_qk is not None
+    if enable_output_qk:
         outputs.append("output_qk")
 
     # ONNX Attention op inputs: Q, K, V, attn_mask, past_key, past_value
@@ -171,7 +181,7 @@
         kv_num_heads=config.kv_num_heads,
         q_num_heads=config.q_num_heads,
         softcap=config.softcap,
-        qk_matmul_output_mode=output_qk,
+        qk_matmul_output_mode=output_qk if enable_output_qk else 0,
         domain="",  # ai.onnx domain
     )
 
@@ -284,7 +294,7 @@
         helper.make_tensor_value_info("present_value", cache_ort_type, output_v_shape),
     ]
 
-    if output_qk > 0:
+    if enable_output_qk:
         graph_output.append(
             helper.make_tensor_value_info(
                 "output_qk",
@@ -301,9 +311,9 @@
     return 24 if config.has_nonpad_kv_seqlen else 23
 
 
-def create_attention_graph_prompt(config: AttentionConfig, ort_type):
+def create_attention_graph_prompt(config: AttentionConfig, ort_type, output_qk: int | None = None):
     """Create ONNX graph for prompt phase (no past KV cache)."""
-    node, graph_input, graph_output = create_attention_node_and_io(config, ort_type, is_past=False)
+    node, graph_input, graph_output = create_attention_node_and_io(config, ort_type, is_past=False, output_qk=output_qk)
     graph = helper.make_graph([node], "Attention_Graph", graph_input, graph_output)
     opset = _get_opset_version(config)
     model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", opset)])
@@ -370,17 +380,18 @@
         return ort_type
 
 
 def attention_prompt_func(
     q,
     k,
     v,
     config: AttentionConfig,
     attn_mask,
     ep,
     device,
     ort_type=TensorProto.FLOAT16,
     nonpad_kv_seqlen=None,
+    output_qk: int | None = None,
 ):
     """
     Run ONNX Attention op for prompt phase (no past KV cache).
 
@@ -394,6 +405,12 @@
         device: Device string (e.g., "cuda")
         ort_type: ONNX tensor type
         nonpad_kv_seqlen: Optional int64 tensor [batch_size] for opset 24
+        output_qk: When not None, enables the optional 4th output `output_qk`
+            and sets the `qk_matmul_output_mode` attribute to this value
+            (0=kQK raw, 1=kQKMask, 2=kQKSoftCap, 3=kQKSoftMax). The function
+            then returns a 4-tuple (out, present_k, present_v, qk) instead
+            of the usual 3-tuple. NOTE: pass None (the default) to DISABLE;
+            `output_qk=0` enables the 4th output in raw-QK mode.
     """
     if not config.kv_cache_type:
         config.kv_cache_type = {
@@ -405,6 +422,7 @@
     onnx_model_str = create_attention_graph_prompt(
         config=config,
         ort_type=ort_type,
+        output_qk=output_qk,
     )
 
     # Reshape inputs for ONNX graph
@@ -476,8 +494,19 @@
     bind_output_tensor(io_binding, "present_key", present_k, device, cache_ort_type)
     bind_output_tensor(io_binding, "present_value", present_v, device, cache_ort_type)
 
+    output_qk_torch = None
+    if output_qk is not None:
+        output_qk_torch = torch.zeros(
+            (config.batch_size, config.q_num_heads, config.q_sequence_length, present_seqlen),
+            dtype=out_dtype,
+            device=device,
+        )
+        bind_output_tensor(io_binding, "output_qk", output_qk_torch, device, ort_type)
+
     ort_session.run_with_iobinding(io_binding)
 
+    if output_qk is not None:
+        return out_torch, present_k, present_v, output_qk_torch
     return out_torch, present_k, present_v