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). 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/test/python/transformers/test_onnx_attention/common.py

@@ -100,11 +100,25 @@
     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 to an int in {0, 1, 2, 3}, enables the optional 4th
+    output `output_qk` and sets the `qk_matmul_output_mode` attribute to
+    that value (0=kQK raw, 1=kQKMask, 2=kQKSoftCap, 3=kQKSoftMax). When None
+    (the default), the 4th output is not emitted and the attribute defaults
+    to 0. Any other int value (negative, or >= 4) raises an AssertionError.
+
+    NOTE: API contract — `output_qk=None` (the default) DISABLES the 4th
+    output. `output_qk=k` for k in {0, 1, 2, 3} ENABLES it and selects the
+    corresponding mode (`output_qk=0` is raw-QK). Passing anything else —
+    including the C++ `kNone = -1` sentinel from attention_parameters.h, or
+    any unknown mode like `4`/`5` — raises immediately rather than silently
+    binding the 4th output (the unfused CUDA kernel would populate it as
+    raw-QK regardless of an out-of-range mode).
+
     ONNX Attention op (opset 23/24) inputs:
     - 0: Q (query) - required
     - 1: K (key) - required
@@ -142,7 +156,16 @@
         "present_value",
     ]
 
-    if output_qk > 0:
+    # Strict validation: only None or one of the known QKMatMulOutputMode values
+    # {0, 1, 2, 3} is accepted. Anything else (including the C++ `kNone = -1`
+    # sentinel, or unknown modes like 4/5) raises immediately, so callers can't
+    # silently bind a 4th output the unfused CUDA kernel would populate as raw-QK.
+    if output_qk is not None:
+        assert output_qk in (0, 1, 2, 3), f"output_qk must be one of {{0, 1, 2, 3}} or None, got {output_qk!r}"
+        enable_output_qk = True
+    else:
+        enable_output_qk = False
+    if enable_output_qk:
         outputs.append("output_qk")
 
     # ONNX Attention op inputs: Q, K, V, attn_mask, past_key, past_value
@@ -171,7 +194,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 +307,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 +324,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)])
@@ -357,18 +380,18 @@
     if ort_type == TensorProto.BFLOAT16:
         return torch.bfloat16
     elif ort_type == TensorProto.FLOAT16:
         return torch.float16
     else:
         return torch.float32
 
 
 def _get_mask_ort_type(config: AttentionConfig, ort_type):
     """Get the ORT type for the attention mask."""
     if config.attn_mask_type == "bool":
         return TensorProto.BOOL
     else:
         return ort_type
 
 
 def attention_prompt_func(
     q,
@@ -380,6 +403,7 @@
     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 +418,14 @@
         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 set to an int in {0, 1, 2, 3}, enables the optional
+            4th output `output_qk` and sets the `qk_matmul_output_mode`
+            attribute to that 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.
+            Pass None (the default) to DISABLE. Any other int value
+            (negative, or >= 4) raises an AssertionError; do NOT pass the
+            C++ `kNone = -1` sentinel — use Python `None`.
     """
     if not config.kv_cache_type:
         config.kv_cache_type = {
@@ -405,6 +437,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 +509,22 @@
     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
+    # `create_attention_node_and_io` (called above via `create_attention_graph_prompt`)
+    # has already validated `output_qk in {0, 1, 2, 3}` or None. Just gate on `is not None`.
+    output_qk_enabled = output_qk is not None
+    if output_qk_enabled:
+        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_enabled:
+        return out_torch, present_k, present_v, output_qk_torch
     return out_torch, present_k, present_v