From b4d7dbf59bd004b99dd7a96511f5d4e55393001f Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 5 May 2026 19:42:16 +0000 Subject: [PATCH 1/5] Tighten CUDA Attention MEA eligibility on head_size%4 + add GQA test coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Addresses microsoft/onnxruntime#28351 sub-items REG, HS4, 1c, 1e: * HS4 (production): Add (head_size % 4 == 0) clause to the MEA dispatch predicate at core/providers/cuda/llm/attention.cc as forward-looking defense-in-depth. The clause is REDUNDANT TODAY: has_memory_efficient_ attention already enforces (qk_head_size & 7) == 0 (i.e. % 8) at contrib_ops/cuda/bert/cutlass_fmha/memory_efficient_attention.h:71-72, which strictly implies % 4. We are not closing a current bug. The clause is kept as the strictest dtype-agnostic alignment floor that CUTLASS FMHA's BiasLoader actually requires (BiasLoader hardcodes a 128-bit / sizeof_bits-element alignment on Q/bias loads — 4 elements for fp32). Once microsoft/onnxruntime#28365 lands and BiasLoader switches to kAlignmentA / DispatchIsAligned, MEA's own % 8 invariant will be loosened and this clause becomes load-bearing, preventing a correctness regression. The new comment block at the predicate site cites #28365 so the next maintainer can identify the right moment to delete it. * REG (test): TestONNXAttentionGQAAsymmetricHeadSize pins the asymmetric v_head_size != head_size GQA path on fp16 and bf16 to guard against regression of the #28358 fix that removed the LaunchUngroup head_size == v_head_size ENFORCE. * HS4 (test): TestONNXAttentionGQAHeadSizeMod4 sweeps head_size in {6, 10, 12, 16, 24}. Today head_sizes 6/10/12 are filtered upstream by MEA's % 8 gate and take the unfused fall-through path; this test pins that fall-through stays numerically correct. 16/24 satisfy both % 8 and % 4 and exercise the MEA happy path. Once #28365 relaxes MEA's % 8 invariant, head_sizes 6/10 will start exercising the HS4 host-side gate directly. * 1c (test): TestONNXAttentionGQAOutputQK pins the GQA + qk_matmul_output_mode combination (kQK raw scaled QK output) which previously had no test coverage. Threads an optional output_qk parameter through common.py's create_attention_graph_prompt and attention_prompt_func. * 1e (test): TestONNXAttentionGQASoftcapFloat32 pins the fp32 + softcap + GQA combination (symmetric and asymmetric V head). fp32 GQA always falls through to the unfused path on CUDA; existing softcap tests are fp16/bf16, so the fp32 unfused softcap branch had no parity coverage. All 10 new tests pass on H100 (sm_90a). Full test_gqa.py: 91 passed, 3 pre-existing flakes ('Output mismatch between two runs' determinism checks in unrelated decode-flash classes — not regressions from this change). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda-attention-kernel-patterns/SKILL.md | 2 +- .../core/providers/cuda/llm/attention.cc | 16 +- .../test_onnx_attention/common.py | 41 ++- .../test_onnx_attention/test_gqa.py | 316 ++++++++++++++++++ 4 files changed, 367 insertions(+), 8 deletions(-) diff --git a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md b/.agents/skills/cuda-attention-kernel-patterns/SKILL.md index 5325a1bf22bdc..f3e66f6398f3a 100644 --- a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md +++ b/.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) diff --git a/onnxruntime/core/providers/cuda/llm/attention.cc b/onnxruntime/core/providers/cuda/llm/attention.cc index 15f9dcbf8e7f2..909c09b82d473 100644 --- a/onnxruntime/core/providers/cuda/llm/attention.cc +++ b/onnxruntime/core/providers/cuda/llm/attention.cc @@ -1383,7 +1383,21 @@ Status Attention::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::value); + !(is_gqa && std::is_same::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 diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/common.py b/onnxruntime/test/python/transformers/test_onnx_attention/common.py index 1ab38fb1ea0f9..5e90ddf33527f 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/common.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/common.py @@ -100,11 +100,20 @@ def create_attention_node_and_io( 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 @@ def create_attention_node_and_io( "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 @@ def create_attention_node_and_io( 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 @@ def create_attention_node_and_io( 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 @@ def _get_opset_version(config: AttentionConfig): 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)]) @@ -380,6 +390,7 @@ def attention_prompt_func( 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 @@ def attention_prompt_func( 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 @@ def attention_prompt_func( 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 @@ def attention_prompt_func( 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 diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py index 55f07666e8c6f..d3c5ae2f920c8 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py @@ -22,6 +22,7 @@ - Boolean padding mask (converted to seqlens_k internally) """ +import math import os import unittest from unittest.mock import patch @@ -1917,5 +1918,320 @@ def test_flash_gqa_softcap_no_mask_decode_fp16(self): ) +@unittest.skipIf(not has_cuda_device(53), "CUDA EP is not available, skipping tests.") +@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) +class TestONNXAttentionGQAAsymmetricHeadSize(unittest.TestCase): + """ + Regression tests for GQA + asymmetric Q/V head sizes (head_size != v_head_size). + + Guards against the silent-broken-output regression that was fixed by PR #28358 + (microsoft/onnxruntime#28357). Before #28358, the GQA unfused path's + LaunchUngroup ENFORCEd head_size == v_head_size, hard-erroring at runtime, and + the MEA eligibility predicate did not exclude the asymmetric case, leading to + NaN / OOB reads when MEA was attempted on an asymmetric V tile. + + These tests pin down the post-fix behaviour by running an asymmetric-GQA + config (q_num_heads=8, kv_num_heads=1, q/k head_size=32, v_head_size=64, + self-attention seq_len=4) on both fp16 and bf16 and asserting numerical + parity with the reference. + + Asymmetric GQA always falls through to the unfused path on CUDA per the + (!is_gqa || head_size == v_head_size) clause of the MEA eligibility + predicate at core/providers/cuda/llm/attention.cc. + """ + + def _run_asymmetric_gqa_prompt(self, torch_type, ort_type): + config = AttentionConfig( + batch_size=1, + q_sequence_length=4, + kv_sequence_length=4, + q_num_heads=8, + kv_num_heads=1, # MQA: kv_num_heads=1, q_num_heads=8 + head_size=32, # small so the test runs fast on H100 + v_head_size=64, # asymmetric: V head twice as large as Q/K head + is_causal=1, + ) + + torch.manual_seed(0) + device = "cuda" + std = 0.2 + + q = ( + torch.randn( + config.batch_size, + config.q_sequence_length, + config.q_num_heads, + config.head_size, + device=device, + dtype=torch_type, + ) + * std + ) + k = ( + torch.randn( + config.batch_size, + config.kv_sequence_length, + config.kv_num_heads, + config.head_size, + device=device, + dtype=torch_type, + ) + * std + ) + v = ( + torch.randn( + config.batch_size, + config.kv_sequence_length, + config.kv_num_heads, + config.v_head_size, + device=device, + dtype=torch_type, + ) + * std + ) + + out_ref, _ = attention_ref(q=q, k=k, v=v, causal=True) + out_ort, _, _ = attention_prompt_func( + q=q, + k=k, + v=v, + config=config, + attn_mask=None, + ep="CUDAExecutionProvider", + device=device, + ort_type=ort_type, + ) + + out_ort = torch.reshape( + out_ort, + (config.batch_size, config.q_sequence_length, config.q_num_heads, config.v_head_size), + ) + + out_np = out_ort.to(torch.float32).detach().cpu().numpy() + out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy() + # Sanity: no NaN propagation from the previously-broken asymmetric path. + self.assertFalse(numpy.isnan(out_np).any(), "NaN in output — asymmetric GQA path regressed") + # fp16/bf16 attention has wide tolerance bands when reductions are reordered. + atol_key = "fp16" if torch_type == torch.float16 else "bf16" + rtol_key = atol_key + numpy.testing.assert_allclose(out_np, out_ref_np, rtol=rtol[rtol_key], atol=atol[atol_key]) + + def test_gqa_asymmetric_v_head_size_prompt_fp16(self): + self._run_asymmetric_gqa_prompt(torch.float16, TensorProto.FLOAT16) + + def test_gqa_asymmetric_v_head_size_prompt_bf16(self): + if not torch.cuda.is_bf16_supported(): + self.skipTest("BFloat16 not supported on this device") + self._run_asymmetric_gqa_prompt(torch.bfloat16, TensorProto.BFLOAT16) + + +@unittest.skipIf(not has_cuda_device(53), "Memory Efficient Attention is not available, skipping tests.") +@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) +class TestONNXAttentionGQAHeadSizeMod4(unittest.TestCase): + """ + Sweep tests for the host-side head_size % 4 == 0 MEA eligibility guard (HS4). + + The HS4 host-side gate at core/providers/cuda/llm/attention.cc is a + forward-looking alignment floor: Cutlass FMHA's BiasLoader uses a 128-bit / + sizeof_bits-element load for Q/bias tiles (= 4 elements for fp32, + 8 for fp16/bf16), and head_size is the inner stride of those loads. + + Today the HS4 clause is REDUNDANT because has_memory_efficient_attention() + already requires (qk_head_size & 7) == 0, which strictly implies %4. So + head_size values like 6, 10, 12 are filtered by the upstream %8 gate before + HS4 is ever consulted; this test merely verifies the fall-through (unfused) + path stays correct for those values. Once microsoft/onnxruntime#28365 + relaxes the BiasLoader alignment check (and the upstream %8 invariant goes + away), this test will start exercising the HS4 host-side gate directly. + """ + + def _run_with_head_size(self, head_size, torch_type, ort_type): + config = AttentionConfig( + batch_size=1, + q_sequence_length=4, + kv_sequence_length=4, # GQA on CUDA requires self-attention + q_num_heads=4, + kv_num_heads=2, # GQA + head_size=head_size, + is_causal=1, + ) + + torch.manual_seed(0) + device = "cuda" + std = 0.2 + + q = torch.randn(1, 4, 4, head_size, device=device, dtype=torch_type) * std + k = torch.randn(1, 4, 2, head_size, device=device, dtype=torch_type) * std + v = torch.randn(1, 4, 2, head_size, device=device, dtype=torch_type) * std + + out_ref, _ = attention_ref(q=q, k=k, v=v, causal=True) + out_ort, _, _ = attention_prompt_func( + q=q, + k=k, + v=v, + config=config, + attn_mask=None, + ep="CUDAExecutionProvider", + device=device, + ort_type=ort_type, + ) + out_ort = torch.reshape(out_ort, (1, 4, 4, head_size)) + + out_np = out_ort.to(torch.float32).detach().cpu().numpy() + out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy() + self.assertFalse( + numpy.isnan(out_np).any(), + f"NaN in output for head_size={head_size} — fall-through path regressed", + ) + numpy.testing.assert_allclose(out_np, out_ref_np, rtol=rtol["fp16"], atol=atol["fp16"]) + + @parameterized.expand([(hs,) for hs in (6, 10, 12, 16, 24)]) + def test_gqa_head_size_modulo_4_sweep_fp16(self, head_size): + # Today's routing for these head_size values (HS4 is redundant with MEA's %8 gate): + # - 6, 10, 12 are filtered upstream by has_memory_efficient_attention()'s %8 + # check, so they take the unfused fall-through path. This test verifies the + # fall-through stays numerically correct. + # - 16, 24 satisfy MEA's %8 gate AND HS4's %4 gate (no attn_mask keeps the + # kernel selection simple) — exercises the MEA "happy path". + # Once microsoft/onnxruntime#28365 relaxes MEA's %8 invariant, head_size 6/10 + # will start exercising the HS4 host-side gate directly. + self._run_with_head_size(head_size, torch.float16, TensorProto.FLOAT16) + + +@unittest.skipIf(not has_cuda_device(53), "CUDA EP is not available, skipping tests.") +@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) +class TestONNXAttentionGQAOutputQK(unittest.TestCase): + """ + Tests that GQA + qk_matmul_output_mode == 0 (raw QK output) works. + + Issue #28351 sub-item 1c: the output_qk path was implemented in the unfused + kernel but lacked test coverage for the GQA + raw-QK combination. The + output_qk shape is [batch, q_num_heads, q_seq, total_seq]; the unfused + kernel indexes per Q-head, and attention_helper.h infers the shape from + q_num_heads, so this combination should already work — these tests pin it. + + Note: GQA on CUDA requires fp16/bf16 (the MEA fp32-GQA path is excluded + by the predicate in attention.cc; the unfused path requires LaunchUngroup + which only has fp16/bf16 instantiations). + """ + + def test_gqa_output_qk_raw_prompt_fp16(self): + config = AttentionConfig( + batch_size=1, + q_sequence_length=4, + kv_sequence_length=4, + q_num_heads=8, + kv_num_heads=2, + head_size=32, + is_causal=1, + ) + + torch.manual_seed(0) + device = "cuda" + torch_type = torch.float16 + ort_type = TensorProto.FLOAT16 + std = 0.2 + + q = torch.randn(1, 4, 8, 32, device=device, dtype=torch_type) * std + k = torch.randn(1, 4, 2, 32, device=device, dtype=torch_type) * std + v = torch.randn(1, 4, 2, 32, device=device, dtype=torch_type) * std + + # Reference output_qk: raw scaled QK (no mask, no softcap, no softmax). + # Mode kQK == 0 outputs raw Q*K^T / sqrt(d) as the spec defines. + q_f, k_f = q.float(), k.float() + # Repeat K heads for GQA (kv_num_heads=2, q_num_heads=8 -> repeat factor 4) + k_rep = k_f.repeat_interleave(q.shape[2] // k.shape[2], dim=2) + ref_qk = torch.einsum("bthd,bshd->bhts", q_f, k_rep) / math.sqrt(q.shape[-1]) + + # Run ORT with output_qk=0 (kQK in the C++ enum: raw scaled QK). + _out_ort, _, _, qk_ort = attention_prompt_func( + q=q, + k=k, + v=v, + config=config, + attn_mask=None, + ep="CUDAExecutionProvider", + device=device, + ort_type=ort_type, + output_qk=0, # kQK: raw scaled QK + ) + + qk_np = qk_ort.to(torch.float32).detach().cpu().numpy() + ref_qk_np = ref_qk.detach().cpu().numpy() + self.assertFalse(numpy.isnan(qk_np).any(), "NaN in output_qk") + self.assertEqual( + qk_np.shape, + (1, 8, 4, 4), + "output_qk shape must be [batch, q_num_heads, q_seq, total_seq]", + ) + numpy.testing.assert_allclose(qk_np, ref_qk_np, rtol=rtol["fp16"], atol=atol["fp16"]) + + +@unittest.skipIf(not has_cuda_device(53), "CUDA EP is not available, skipping tests.") +@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) +class TestONNXAttentionGQASoftcapFloat32(unittest.TestCase): + """ + Issue #28351 sub-item 1e: softcap coverage for the fp32 path. + + fp32 GQA on CUDA always falls through to the unfused path (the MEA + predicate at attention.cc explicitly excludes is_gqa && std::is_same::value because LaunchUngroup has no fp32 instantiation). Existing + softcap tests are fp16/bf16; this test pins the fp32 + softcap + + asymmetric-or-symmetric-GQA combination so future kernel changes can't + silently break the unfused softcap branch for fp32. + """ + + def _run_softcap_fp32(self, head_size, v_head_size=None): + # v_head_size=None means "same as head_size" (symmetric V). + effective_v_head_size = v_head_size if v_head_size is not None else head_size + config = AttentionConfig( + batch_size=1, + q_sequence_length=4, + kv_sequence_length=4, # GQA on CUDA requires self-attention + q_num_heads=4, + kv_num_heads=2, + head_size=head_size, + # AttentionConfig.v_head_size uses 0 as the "same as head_size" sentinel + # (defined in common.py); translate from the test-local None convention. + v_head_size=v_head_size if v_head_size is not None else 0, + is_causal=1, + softcap=2.0, # small softcap exposes ordering / clipping issues + ) + + torch.manual_seed(0) + device = "cuda" + torch_type = torch.float32 + ort_type = TensorProto.FLOAT + std = 0.5 + + q = torch.randn(1, 4, 4, head_size, device=device, dtype=torch_type) * std + k = torch.randn(1, 4, 2, head_size, device=device, dtype=torch_type) * std + v = torch.randn(1, 4, 2, effective_v_head_size, device=device, dtype=torch_type) * std + + out_ref, _ = attention_ref(q=q, k=k, v=v, causal=True, softcap=2.0) + out_ort, _, _ = attention_prompt_func( + q=q, + k=k, + v=v, + config=config, + attn_mask=None, + ep="CUDAExecutionProvider", + device=device, + ort_type=ort_type, + ) + out_ort = torch.reshape(out_ort, (1, 4, 4, effective_v_head_size)) + + out_np = out_ort.to(torch.float32).detach().cpu().numpy() + out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy() + self.assertFalse(numpy.isnan(out_np).any()) + numpy.testing.assert_allclose(out_np, out_ref_np, rtol=rtol["fp32"], atol=atol["fp32"]) + + def test_gqa_softcap_fp32_symmetric(self): + self._run_softcap_fp32(head_size=16, v_head_size=None) + + def test_gqa_softcap_fp32_asymmetric_v_head(self): + self._run_softcap_fp32(head_size=16, v_head_size=32) + + if __name__ == "__main__": unittest.main() From cb8bd5902a8272272adee0c2253bec1b6bdf1838 Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 5 May 2026 21:40:36 +0000 Subject: [PATCH 2/5] Address copilot-bot review comments on PR #28371 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four targeted follow-ups to the bot review on https://github.com/microsoft/onnxruntime/pull/28371. No production behaviour change beyond the comment text and a Python helper guard. * core/providers/cuda/llm/attention.cc — soften the HS4 deletion criterion. The original 'delete when MEA no longer requires the %8 invariant' is necessary but not sufficient: removing the clause would also need every other host-side gate to keep head_size < 4 out of LaunchUngroup, which still ORT_ENFORCEs head_size %% 4 == 0 internally (see ~line 723-724). Cite microsoft/onnxruntime#28365 and the LaunchUngroup ENFORCE site explicitly. * test/python/transformers/test_onnx_attention/common.py — fix output_qk negative-mode bug. Helper used 'output_qk is not None' to gate the optional 4th output, but a caller mirroring the C++ enum convention (kNone = -1 in attention_parameters.h) would pass -1 and silently get the 4th output bound + the unfused CUDA kernel populating it as raw-QK. Tighten the gate to '>= 0' across all three sites (graph node, output binding, return tuple) and update the prominent NOTE block + docstrings to spell out the convention. * test/python/transformers/test_onnx_attention/test_gqa.py — add test_gqa_softcap_fp32_with_mask_ordering_{symmetric,asymmetric_v_head} to TestONNXAttentionGQASoftcapFloat32. The existing fp32 softcap cases passed attn_mask=None, so they could not detect a wrong softcap-vs-mask order on the unfused fp32 path (without a mask the two orders are arithmetically identical). The new tests use the same poison-V pattern as the fp16/bf16 P1 ordering guards (small softcap, V=1000 in masked slot, attn_mask=-inf there) so wrong ordering produces wild magnitudes / NaN and right ordering yields bounded finite values. Compare against attention_ref(). * test/python/transformers/test_onnx_attention/test_gqa.py — correct the TestONNXAttentionGQAOutputQK docstring. The fp16/bf16 restriction applies only to the MEA LaunchUngroup helper, not the entire GQA-on- CUDA surface; the unfused fall-through DOES support fp32, exercised by TestONNXAttentionGQASoftcapFloat32 in the same file. All four fixes verified locally: - 96/96 in test_onnx_attention/test_gqa.py pass (PR-2 build with HS4 dispatch + GQA-fp32 MEA exclusion). - 4/4 TestONNXAttentionGQASoftcapFloat32 pass (2 existing + 2 new masked ordering tests). - The P1-style ordering guard test_gqa_large_head_unfused_softcap_additive_mask_poison_fp16 still passes; the common.py output_qk change does not affect paths that don't request output_qk. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/cuda/llm/attention.cc | 14 ++- .../test_onnx_attention/common.py | 49 +++++---- .../test_onnx_attention/test_gqa.py | 99 ++++++++++++++++++- 3 files changed, 139 insertions(+), 23 deletions(-) diff --git a/onnxruntime/core/providers/cuda/llm/attention.cc b/onnxruntime/core/providers/cuda/llm/attention.cc index 909c09b82d473..6a379f867d780 100644 --- a/onnxruntime/core/providers/cuda/llm/attention.cc +++ b/onnxruntime/core/providers/cuda/llm/attention.cc @@ -1395,8 +1395,18 @@ Status Attention::ComputeInternal(OpKernelContext* context) const { // 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. + // point MEA's %8 invariant goes away). + // + // Delete this clause ONLY when BOTH conditions hold: + // (a) `has_memory_efficient_attention` no longer enforces the + // (qk_head_size & 7) == 0 invariant (post-microsoft/onnxruntime#28365 + // likely relaxes it to %4 via BiasLoader = kAlignmentA), AND + // (b) no other host-side gate above this point would let head_size < 4 + // (e.g. head_size == 2) reach LaunchUngroup, which still enforces + // head_size % 4 == 0 internally (see ORT_ENFORCE near line 723-724 + // of this file). Removing this clause without (b) would let GQA+MEA + // configurations crash inside LaunchUngroup instead of falling + // through to the unfused path here. (parameters.head_size % 4 == 0); // Cutlass FMHA requires bias strides to satisfy minimum alignment even in the diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/common.py b/onnxruntime/test/python/transformers/test_onnx_attention/common.py index 5e90ddf33527f..c81ef195509a4 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/common.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/common.py @@ -105,14 +105,20 @@ def create_attention_node_and_io( """ 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. + output_qk: when set to a non-negative int, 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) or + when set to any negative value (matching the C++ `kNone = -1` sentinel in + onnxruntime/core/providers/cpu/llm/attention_parameters.h), the 4th output + is not emitted and the attribute defaults to 0. + + NOTE: API contract — `output_qk >= 0` ENABLES the 4th output and selects + the mode; `output_qk=0` enables it in raw-QK mode. `output_qk=None` (the + default) OR any negative int (e.g. `-1`, mirroring the C++ enum's `kNone`) + DISABLES it. Do not pass 0 expecting it to mean "disabled" — pass None + instead. The negative-value branch is defensive: a caller who reads the + C++ enum and passes -1 must not silently bind a 4th output and have the + unfused kernel populate it as raw-QK. ONNX Attention op (opset 23/24) inputs: - 0: Q (query) - required @@ -151,7 +157,11 @@ def create_attention_node_and_io( "present_value", ] - enable_output_qk = output_qk is not None + # Treat `output_qk is None` AND any negative int (e.g. -1, matching the C++ + # `kNone` enum value in attention_parameters.h) as "disabled". Otherwise a + # caller who passes -1 would silently get the 4th output bound and the + # unfused CUDA kernel would populate it as raw-QK regardless. + enable_output_qk = output_qk is not None and output_qk >= 0 if enable_output_qk: outputs.append("output_qk") @@ -405,12 +415,13 @@ def attention_prompt_func( 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. + output_qk: When set to a non-negative int, 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) — or + any negative int (matching the C++ `kNone = -1` enum sentinel) — to + DISABLE; `output_qk=0` enables the 4th output in raw-QK mode. """ if not config.kv_cache_type: config.kv_cache_type = { @@ -495,7 +506,11 @@ def attention_prompt_func( bind_output_tensor(io_binding, "present_value", present_v, device, cache_ort_type) output_qk_torch = None - if output_qk is not None: + # Mirror the `enable_output_qk` convention from create_attention_node_and_io: + # treat negative ints (e.g. C++ `kNone = -1`) as "disabled" so we don't bind + # a 4th output the graph never declared. + output_qk_enabled = output_qk is not None and output_qk >= 0 + 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, @@ -505,7 +520,7 @@ def attention_prompt_func( ort_session.run_with_iobinding(io_binding) - if output_qk is not None: + if output_qk_enabled: return out_torch, present_k, present_v, output_qk_torch return out_torch, present_k, present_v diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py index d3c5ae2f920c8..39dd6be2ea053 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py @@ -2110,9 +2110,10 @@ class TestONNXAttentionGQAOutputQK(unittest.TestCase): kernel indexes per Q-head, and attention_helper.h infers the shape from q_num_heads, so this combination should already work — these tests pin it. - Note: GQA on CUDA requires fp16/bf16 (the MEA fp32-GQA path is excluded - by the predicate in attention.cc; the unfused path requires LaunchUngroup - which only has fp16/bf16 instantiations). + Note: GQA + MEA on CUDA requires fp16/bf16 because the MEA `LaunchUngroup` + helper has no fp32 instantiation; the GQA unfused fall-through DOES + support fp32 (exercised by `TestONNXAttentionGQASoftcapFloat32` below). + These tests pin the fp16 + raw-QK + GQA combination on the unfused path. """ def test_gqa_output_qk_raw_prompt_fp16(self): @@ -2176,9 +2177,14 @@ class TestONNXAttentionGQASoftcapFloat32(unittest.TestCase): fp32 GQA on CUDA always falls through to the unfused path (the MEA predicate at attention.cc explicitly excludes is_gqa && std::is_same::value because LaunchUngroup has no fp32 instantiation). Existing - softcap tests are fp16/bf16; this test pins the fp32 + softcap + + softcap tests are fp16/bf16; this class pins the fp32 + softcap + asymmetric-or-symmetric-GQA combination so future kernel changes can't silently break the unfused softcap branch for fp32. + + The `_with_mask_ordering` variants additionally pin softcap+mask + ORDERING on the fp32 path (poison-V pattern). Without them, the + unmasked softcap tests cannot detect a wrong order — softcap and mask + only diverge when both are present. """ def _run_softcap_fp32(self, head_size, v_head_size=None): @@ -2232,6 +2238,91 @@ def test_gqa_softcap_fp32_symmetric(self): def test_gqa_softcap_fp32_asymmetric_v_head(self): self._run_softcap_fp32(head_size=16, v_head_size=32) + def _run_softcap_fp32_with_mask(self, head_size, v_head_size=None): + """ + Pin softcap+mask ORDERING on the fp32 unfused path. + + The two `_run_softcap_fp32` cases above pass `attn_mask=None`, so they + exercise softcap but cannot detect a wrong ordering of softcap vs + additive mask — without a mask, "softcap then mask" and "mask then + softcap" are arithmetically identical. This test uses the same + poison-V pattern as the fp16/bf16 P1 ordering guards + (test_gqa_large_head_unfused_softcap_additive_mask_poison_fp16): + + - Tiny softcap (2.0) so it would clamp very large logits. + - V values = 1000.0 in the masked KV slot, 0.2 elsewhere. + - attn_mask = -inf for the masked slot, 0 elsewhere. + + Correct order (QK -> softcap -> +mask -> softmax) zeroes out the + masked logit via softmax, so output ~= 0.2. Wrong order (mask before + softcap) would feed -inf through softcap and either clamp it to a + finite value (allowing the poisoned V to leak) or produce NaN. + """ + effective_v_head_size = v_head_size if v_head_size is not None else head_size + config = AttentionConfig( + batch_size=1, + q_sequence_length=1, + kv_sequence_length=3, + q_num_heads=4, + kv_num_heads=2, + head_size=head_size, + v_head_size=v_head_size if v_head_size is not None else 0, + is_causal=0, + softcap=2.0, + has_attn_mask=True, + attn_mask_dims=4, + attn_mask_type="additive", + ) + + device = "cuda" + torch_type = torch.float32 + ort_type = TensorProto.FLOAT + + q = torch.zeros(1, 1, 4, head_size, device=device, dtype=torch_type) + k = torch.zeros(1, 3, 2, head_size, device=device, dtype=torch_type) + v = torch.full((1, 3, 2, effective_v_head_size), 0.2, device=device, dtype=torch_type) + v[:, 1, :, :] = 1000.0 # poison the masked slot + + attn_mask = torch.zeros(1, 4, 1, 3, device=device, dtype=torch_type) + attn_mask[:, :, :, 1] = float("-inf") + + out_ref, _ = attention_ref(q=q, k=k, v=v, attn_bias=attn_mask, softcap=2.0) + out_ort, _, _ = attention_prompt_func( + q=q, + k=k, + v=v, + config=config, + attn_mask=attn_mask, + ep="CUDAExecutionProvider", + device=device, + ort_type=ort_type, + ) + out = out_ort.reshape(1, 1, 4, effective_v_head_size) + + out_np = out.to(torch.float32).detach().cpu().numpy() + out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy() + + self.assertFalse( + numpy.isnan(out_np).any(), + "NaN in fp32 GQA softcap+mask output — wrong softcap/mask ordering on the unfused fp32 path?", + ) + max_abs = numpy.max(numpy.abs(out_np)) + self.assertLess( + max_abs, + 1.0, + f"fp32 GQA softcap+mask leakage: max |output| = {max_abs:.3f}. " + f"Expected ~0.2 (mask zeroes out the poisoned V=1000 slot via softmax). " + f"Wrong ordering (mask before softcap) would let the -inf get clamped " + f"by softcap and the poisoned V to leak through.", + ) + numpy.testing.assert_allclose(out_np, out_ref_np, rtol=rtol["fp32"], atol=atol["fp32"]) + + def test_gqa_softcap_fp32_with_mask_ordering_symmetric(self): + self._run_softcap_fp32_with_mask(head_size=16, v_head_size=None) + + def test_gqa_softcap_fp32_with_mask_ordering_asymmetric_v_head(self): + self._run_softcap_fp32_with_mask(head_size=16, v_head_size=32) + if __name__ == "__main__": unittest.main() From e8a8db723a81d99ecdc263c23eec7ff15f9649df Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 5 May 2026 23:37:50 +0000 Subject: [PATCH 3/5] Address review feedback round 2 on PR #28371: drop HS4 gate, tighten output_qk validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three substantive items + one docstring fix from round-2 reviewer feedback (bot + internal multi-reviewer consolidation). * core/providers/cuda/llm/attention.cc — drop the host-side `head_size %% 4 == 0` HS4 clause from the MEA-eligibility predicate and remove its multi-paragraph comment block. The clause is fully redundant today (`has_memory_efficient_attention` already requires `(qk_head_size & 7) == 0`, which strictly implies %4) and the comment it carried made dtype-aware alignment claims that are wrong for fp16 / bf16 (BiasLoader needs an 8-element stride, not 4, for those dtypes). The dtype-aware alignment floor properly belongs in the BiasLoader fix (microsoft/onnxruntime#28365), not as a vestigial redundant clause here. Predicate is now exactly the upstream/main shape for HS4 purposes. * test/python/transformers/test_onnx_attention/test_gqa.py — delete TestONNXAttentionGQAHeadSizeMod4. With the HS4 clause gone there is no host-side gate left to validate; the parameterized sweep was exercising routing equivalence vs the unfused fall-through, which is already covered by the broader MEA / unfused tests. * test/python/transformers/test_onnx_attention/common.py — tighten the output_qk parameter validation to `output_qk in {0, 1, 2, 3}` or `None`. The previous `is not None and >= 0` guard caught the C++ `kNone = -1` sentinel but still silently accepted invalid modes 4 / 5, which would build an ONNX graph with an out-of-range `qk_matmul_output_mode` attribute and let the unfused CUDA kernel populate the 4th output as raw-QK regardless. Validation now raises immediately with a clear message at the helper boundary; the binding-allocation site downstream is simplified to `is not None` since validation has already happened. NOTE block + both helper docstrings updated to spell out the contract: `None` = disabled; `{0,1,2,3}` = the corresponding QKMatMulOutputMode; anything else raises. * test/python/transformers/test_onnx_attention/test_gqa.py — fix the TestONNXAttentionGQAAsymmetricHeadSize docstring. The pre-#28358 `head_size == v_head_size` ENFORCE in LaunchUngroup is an MEA-path enforcement (LaunchUngroup is the GQA head-expansion helper used by MEA before its FMHA kernel), not an unfused-path one. Docstring now correctly attributes it. Verified locally on the PR-2 build (build_pr2/, sm_90a single-arch): - All targeted PR-2 + ordering-guard tests pass (8/8): the existing OutputQK / SoftcapFloat32 / AsymmetricHeadSize / LargeHeadUnfused poison ordering guard, plus the 2 masked fp32 ordering tests added in the previous fix-up. - test_onnx_attention/test_gqa.py: 89/91 pass on a quiet GPU. The 2 transient failures (FloatMaskDecode, MEAGQASoftcap softcap+mask decode) both pass cleanly when re-run in isolation; they are pre-existing run-to-run flakes (rtol=0/atol=0 strict-equality asserts) under shared-GPU pressure, not caused by this commit. - HS4 sweep class is gone (file count dropped from 96 to 91; the delta is exactly the 5 parameterized HS4 sweep cases, as expected). - Manual negative test of the new validation: output_qk=None -> 3 outputs (disabled, OK) output_qk=2 -> 4 outputs (kQKSoftCap, OK) output_qk=-1 -> AssertionError (OK) output_qk=4 -> AssertionError (OK) output_qk=5 -> AssertionError (OK) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../core/providers/cuda/llm/attention.cc | 26 +----- .../test_onnx_attention/common.py | 65 ++++++++------- .../test_onnx_attention/test_gqa.py | 80 +------------------ 3 files changed, 40 insertions(+), 131 deletions(-) diff --git a/onnxruntime/core/providers/cuda/llm/attention.cc b/onnxruntime/core/providers/cuda/llm/attention.cc index 6a379f867d780..15f9dcbf8e7f2 100644 --- a/onnxruntime/core/providers/cuda/llm/attention.cc +++ b/onnxruntime/core/providers/cuda/llm/attention.cc @@ -1383,31 +1383,7 @@ Status Attention::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::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 ONLY when BOTH conditions hold: - // (a) `has_memory_efficient_attention` no longer enforces the - // (qk_head_size & 7) == 0 invariant (post-microsoft/onnxruntime#28365 - // likely relaxes it to %4 via BiasLoader = kAlignmentA), AND - // (b) no other host-side gate above this point would let head_size < 4 - // (e.g. head_size == 2) reach LaunchUngroup, which still enforces - // head_size % 4 == 0 internally (see ORT_ENFORCE near line 723-724 - // of this file). Removing this clause without (b) would let GQA+MEA - // configurations crash inside LaunchUngroup instead of falling - // through to the unfused path here. - (parameters.head_size % 4 == 0); + !(is_gqa && std::is_same::value); // Cutlass FMHA requires bias strides to satisfy minimum alignment even in the // "unaligned" kernel path. When an attention mask is present (with or without diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/common.py b/onnxruntime/test/python/transformers/test_onnx_attention/common.py index c81ef195509a4..5a06f3ffb596c 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/common.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/common.py @@ -105,20 +105,19 @@ def create_attention_node_and_io( """ Create ONNX Attention op node and I/O definitions for testing. - output_qk: when set to a non-negative int, 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) or - when set to any negative value (matching the C++ `kNone = -1` sentinel in - onnxruntime/core/providers/cpu/llm/attention_parameters.h), the 4th output - is not emitted and the attribute defaults to 0. - - NOTE: API contract — `output_qk >= 0` ENABLES the 4th output and selects - the mode; `output_qk=0` enables it in raw-QK mode. `output_qk=None` (the - default) OR any negative int (e.g. `-1`, mirroring the C++ enum's `kNone`) - DISABLES it. Do not pass 0 expecting it to mean "disabled" — pass None - instead. The negative-value branch is defensive: a caller who reads the - C++ enum and passes -1 must not silently bind a 4th output and have the - unfused kernel populate it as raw-QK. + 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 @@ -157,11 +156,17 @@ def create_attention_node_and_io( "present_value", ] - # Treat `output_qk is None` AND any negative int (e.g. -1, matching the C++ - # `kNone` enum value in attention_parameters.h) as "disabled". Otherwise a - # caller who passes -1 would silently get the 4th output bound and the - # unfused CUDA kernel would populate it as raw-QK regardless. - enable_output_qk = output_qk is not None and 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") @@ -415,13 +420,14 @@ def attention_prompt_func( 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 a non-negative int, 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) — or - any negative int (matching the C++ `kNone = -1` enum sentinel) — to - DISABLE; `output_qk=0` enables the 4th output in raw-QK mode. + 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 = { @@ -506,10 +512,9 @@ def attention_prompt_func( bind_output_tensor(io_binding, "present_value", present_v, device, cache_ort_type) output_qk_torch = None - # Mirror the `enable_output_qk` convention from create_attention_node_and_io: - # treat negative ints (e.g. C++ `kNone = -1`) as "disabled" so we don't bind - # a 4th output the graph never declared. - output_qk_enabled = output_qk is not None and output_qk >= 0 + # `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), diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py index 39dd6be2ea053..04f88de18534d 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py @@ -1925,9 +1925,10 @@ class TestONNXAttentionGQAAsymmetricHeadSize(unittest.TestCase): Regression tests for GQA + asymmetric Q/V head sizes (head_size != v_head_size). Guards against the silent-broken-output regression that was fixed by PR #28358 - (microsoft/onnxruntime#28357). Before #28358, the GQA unfused path's - LaunchUngroup ENFORCEd head_size == v_head_size, hard-erroring at runtime, and - the MEA eligibility predicate did not exclude the asymmetric case, leading to + (microsoft/onnxruntime#28357). Before #28358, the GQA + MEA path's + LaunchUngroup helper (used by MEA to expand K/V heads before the FMHA kernel) + ENFORCEd head_size == v_head_size, hard-erroring at runtime, and the MEA + eligibility predicate did not exclude the asymmetric case, leading to NaN / OOB reads when MEA was attempted on an asymmetric V tile. These tests pin down the post-fix behaviour by running an asymmetric-GQA @@ -2025,79 +2026,6 @@ def test_gqa_asymmetric_v_head_size_prompt_bf16(self): self._run_asymmetric_gqa_prompt(torch.bfloat16, TensorProto.BFLOAT16) -@unittest.skipIf(not has_cuda_device(53), "Memory Efficient Attention is not available, skipping tests.") -@patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) -class TestONNXAttentionGQAHeadSizeMod4(unittest.TestCase): - """ - Sweep tests for the host-side head_size % 4 == 0 MEA eligibility guard (HS4). - - The HS4 host-side gate at core/providers/cuda/llm/attention.cc is a - forward-looking alignment floor: Cutlass FMHA's BiasLoader uses a 128-bit / - sizeof_bits-element load for Q/bias tiles (= 4 elements for fp32, - 8 for fp16/bf16), and head_size is the inner stride of those loads. - - Today the HS4 clause is REDUNDANT because has_memory_efficient_attention() - already requires (qk_head_size & 7) == 0, which strictly implies %4. So - head_size values like 6, 10, 12 are filtered by the upstream %8 gate before - HS4 is ever consulted; this test merely verifies the fall-through (unfused) - path stays correct for those values. Once microsoft/onnxruntime#28365 - relaxes the BiasLoader alignment check (and the upstream %8 invariant goes - away), this test will start exercising the HS4 host-side gate directly. - """ - - def _run_with_head_size(self, head_size, torch_type, ort_type): - config = AttentionConfig( - batch_size=1, - q_sequence_length=4, - kv_sequence_length=4, # GQA on CUDA requires self-attention - q_num_heads=4, - kv_num_heads=2, # GQA - head_size=head_size, - is_causal=1, - ) - - torch.manual_seed(0) - device = "cuda" - std = 0.2 - - q = torch.randn(1, 4, 4, head_size, device=device, dtype=torch_type) * std - k = torch.randn(1, 4, 2, head_size, device=device, dtype=torch_type) * std - v = torch.randn(1, 4, 2, head_size, device=device, dtype=torch_type) * std - - out_ref, _ = attention_ref(q=q, k=k, v=v, causal=True) - out_ort, _, _ = attention_prompt_func( - q=q, - k=k, - v=v, - config=config, - attn_mask=None, - ep="CUDAExecutionProvider", - device=device, - ort_type=ort_type, - ) - out_ort = torch.reshape(out_ort, (1, 4, 4, head_size)) - - out_np = out_ort.to(torch.float32).detach().cpu().numpy() - out_ref_np = out_ref.to(torch.float32).detach().cpu().numpy() - self.assertFalse( - numpy.isnan(out_np).any(), - f"NaN in output for head_size={head_size} — fall-through path regressed", - ) - numpy.testing.assert_allclose(out_np, out_ref_np, rtol=rtol["fp16"], atol=atol["fp16"]) - - @parameterized.expand([(hs,) for hs in (6, 10, 12, 16, 24)]) - def test_gqa_head_size_modulo_4_sweep_fp16(self, head_size): - # Today's routing for these head_size values (HS4 is redundant with MEA's %8 gate): - # - 6, 10, 12 are filtered upstream by has_memory_efficient_attention()'s %8 - # check, so they take the unfused fall-through path. This test verifies the - # fall-through stays numerically correct. - # - 16, 24 satisfy MEA's %8 gate AND HS4's %4 gate (no attn_mask keeps the - # kernel selection simple) — exercises the MEA "happy path". - # Once microsoft/onnxruntime#28365 relaxes MEA's %8 invariant, head_size 6/10 - # will start exercising the HS4 host-side gate directly. - self._run_with_head_size(head_size, torch.float16, TensorProto.FLOAT16) - - @unittest.skipIf(not has_cuda_device(53), "CUDA EP is not available, skipping tests.") @patch.dict(os.environ, {"ORT_DISABLE_FLASH_ATTENTION": "1"}) class TestONNXAttentionGQAOutputQK(unittest.TestCase): From 99a48bc6894830257ae90ac0b7f83a4da96bbb9a Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 5 May 2026 23:39:18 +0000 Subject: [PATCH 4/5] Remove stale SKILL.md HS4 reference after dropping host-side gate MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #28371 commit e8a8db723a dropped the host-side head_size %% 4 == 0 HS4 clause from the MEA-eligibility predicate in core/providers/cuda/llm/attention.cc. This commit removes the now- orphaned reference in the cuda-attention-kernel-patterns skill (§1 MEA eligibility bullet), which previously documented the HS4 floor as a forward-looking alignment guard. The LaunchUngroup, decode, GQA-fp32, and bias-stride clauses are unchanged. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .agents/skills/cuda-attention-kernel-patterns/SKILL.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md b/.agents/skills/cuda-attention-kernel-patterns/SKILL.md index f3e66f6398f3a..4a1547eaf1f9d 100644 --- a/.agents/skills/cuda-attention-kernel-patterns/SKILL.md +++ b/.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 (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). +**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) From 449fa8cfa57f4433b56bfd0572b388d87ce58cdc Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Wed, 6 May 2026 16:36:31 +0000 Subject: [PATCH 5/5] Fix RUFF formatting on output_qk assertion Collapses the multi-line parenthesized assertion message at common.py:163-165 to a single line, per RUFF/Black formatter. Assertion semantics unchanged (verified manually: None and {0,1,2,3} accepted; -1, 4, 5 raise the same AssertionError with the same message). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../test/python/transformers/test_onnx_attention/common.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/onnxruntime/test/python/transformers/test_onnx_attention/common.py b/onnxruntime/test/python/transformers/test_onnx_attention/common.py index 5a06f3ffb596c..a756fa5b2847a 100644 --- a/onnxruntime/test/python/transformers/test_onnx_attention/common.py +++ b/onnxruntime/test/python/transformers/test_onnx_attention/common.py @@ -161,9 +161,7 @@ def create_attention_node_and_io( # 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}" - ) + 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