Initialize large-head FMHA shared memory per kernel variant

[yinli-systems]

Description

Fixes #28388.

Large-head Memory Efficient Attention kernels require an explicit cudaFuncAttributeMaxDynamicSharedMemorySize opt-in when their dynamic shared-memory requirement exceeds 0xc000. The default batched kernel and the custom right-padding kernel are distinct device functions, but the previous function-local once configured only whichever function pointer was selected on the first invocation. If the other variant ran later in the same process, it launched without the required opt-in and failed with excessive shared-memory usage.

This change:

• keeps separate thread-safe initialization state for the default and right-padding kernel variants;
• applies cudaFuncSetAttribute to the exact kernel that will launch;
• checks the CUDA API result with CUDA_CALL_THROW; and
• launches the matching kernel directly after its own shared-memory initialization.

Tests

• Added TestONNXAttentionGQALargeHeadNonpadMEA, which runs the default large-head MEA kernel first and then the right-padding/nonpad variant in the same process. This ordering reproduces the original process-global initialization bug and verifies that both distinct kernels opt in independently.
• python3 -m py_compile onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py
• git diff --check

The targeted CUDA test requires an ONNX Runtime CUDA build and a compatible GPU, so execution is delegated to CI.

[yinli-systems]

@microsoft-github-policy-service agree

[Copilot]

Pull request overview

This PR updates the CUDA dispatch logic for the ONNX-domain Attention op to avoid launching CUTLASS Memory Efficient Attention (MEA) in a known-crashing configuration (external-cache nonpad_kv_seqlen + head_size > 256). Instead, it forces that specific case to fall through to the unified unfused attention implementation, which supports large head sizes safely.

Changes:

• Add an MEA eligibility guard: when nonpad_kv_seqlen != nullptr and head_size > 256, skip MEA and use the unified unfused fallback.
• Add a Python parity test that exercises head_size=512 with nonpad_kv_seqlen on CUDA to ensure the large-head external-cache scenario works (via fallback).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
onnxruntime/core/providers/cuda/llm/attention.cc	Adds a targeted MEA dispatch guard for the nonpad (custom right-padding) path when head_size > 256, forcing fallback to unified unfused attention.
onnxruntime/test/python/transformers/test_onnx_attention/test_gqa.py	Adds a CUDA test covering head_size=512 + nonpad_kv_seqlen to validate the new fallback behavior doesn’t crash and remains numerically correct.

[titaiwangms]

Review summary

The fallback itself is semantically safe and non-regressing — the unified unfused path genuinely handles nonpad_kv_seqlen + GQA + softcap, so forcing nonpad large-head to unfused only changes performance, not results. The concerns below are about whether the guard condition is the right one.

Major

1. The 256 cutoff doesn't match the kernel — 256 and 512 use the identical instantiation and shared-memory footprint.
DispatchBlockSize (cutlass_fmha/fmha_launch_template.h:306-312) buckets on v_head_size: ≤64, ≤128, else → <32,128,1024>. So v_head_size 256 and 512 both instantiate the same <...,1024> kernel, and SharedStorage is a compile-time, template-only size (:242) — byte-identical for 256 and 512. The comment's "large heads exceed the per-SM dynamic shared-memory opt-in limit" therefore cannot explain a 256/512 split. The real dispatch seam is v_head_size > 128. The 256 looks inherited from a stale comment in test_gqa.py:1352,1430 (claims MEA caps at 256; the actual cap is 1024 — memory_efficient_attention.h:73). Suggest: rewrite the comment to state it's an empirical safety margin, or align the boundary to the real seam, and make it a named constant.

2. The guard checks head_size, but kernel selection keys on v_head_size.
DispatchBlockSize buckets on v_head_size, while the new clause checks parameters.head_size. For nonpad MHA with no past and non-GQA, head_size != v_head_size is allowed (equality is only enforced for GQA/past at attention.cc:1382-1383), so e.g. head_size=128, v_head_size=512 bypasses the guard into the heavy kernel. Suggest gating on max(head_size, v_head_size).

3. Likely true root cause the PR doesn't address: the smem opt-in is registered for only one of two kernel variants.
In fmha_launch_template.h:236-249, kernel_fn is reassigned to attention_kernel_batched_impl_right_padding (:239), but static bool once (:245) calls cudaFuncSetAttribute exactly once per template instantiation — for whichever kernel_fn ran first. cudaFuncSetAttribute is per-function, so a prior normal MEA call sets the attribute on the normal kernel, and a later nonpad call uses the right-padding kernel that never received the opt-in → an order-dependent, nonpad-only launch fault (and the return value is discarded). This matches the nonpad-specific signature of the bug; the head-size guard masks the symptom for one band rather than fixing it. A proper fix would set the attribute per function pointer (drop static) and check the return.

4. The test doesn't prove the fallback occurred.
test_gqa.py does a parity check only; it never asserts MEA was skipped, so it can pass with or without the guard on a GPU where the kernel doesn't crash. Combined with #3, an isolated run may set the opt-in on the right-padding kernel first and not crash even pre-fix. Suggest asserting dispatch (e.g. capture VERBOSE logs and assert the "using Memory Efficient Attention" line is absent / unfused is used).

Out-of-scope follow-up

• The contrib sibling GroupQueryAttention sets has_custom_right_padding = true unconditionally (group_query_attention_impl.cu:1063) gated only at 1024 (group_query_attention.cc:524) — same latent crash, worth a follow-up PR.

Minor / Nit

• 256 is a magic number — name it (the adjacent min_bias_align = 4 constant is the model to follow).
• Comment wording "smaller architectures" → more precise as "pre-Ampere / Turing (SM75)".
• Test name ..._falls_back_from_mea_fp16 breaks the sibling _unfused_fp16 suffix convention.
• Stale head_size=256 cap comments at test_gqa.py:1352,1430 should be corrected.

Open questions (need pre-Ampere/SM75 hardware to settle)

• Does head_size=256 itself also exceed the opt-in on the smallest arch (cutoff too loose)?
• Does the normal non-nonpad head_size>256 path also fault, or does the static once ordering spare it (favoring update HighLevelDesign.md #3 over the asymmetry concern)?
• Performance cost of forcing large-head nonpad to unfused on high-smem Ampere/Hopper.

---

Reviewed by a multi-model agent team (readability / correctness / adversarial / spec-depth / integration). Structural claims above were verified against the cited source; smem/crash claims that require SM75 hardware are flagged as open questions.

[tianleiwu]

Summary

This fix correctly addresses the root cause of #28388 in fmha_launch_template.h. The previous code opted a single kernel_fn into the large dynamic-shared-memory limit via one static bool once guard. Because the default kernel (attention_kernel_batched_impl) and the right-padding/nonpad kernel (attention_kernel_batched_impl_right_padding) are distinct device functions that both need the cudaFuncAttributeMaxDynamicSharedMemorySize opt-in when smem_bytes > 0xc000 (e.g. head_size=512), once only configured whichever variant ran first. The second-seen variant then launched without the opt-in and crashed (too much shared memory). Splitting into per-variant statics (right_padding_once / default_once) and launching the matching kernel directly is the right fix. Upgrading the ignored cudaFuncSetAttribute return to CUDA_CALL_THROW is a nice diagnosability improvement.

Verified: queries_per_block is a compile-time template parameter of LaunchCutlassFmha, so both statics and the direct launches are well-formed per instantiation; function-local statics keep thread-safe init; CUDA_CALL_THROW is available transitively via memory_efficient_attention.h -> cuda_common.h -> cuda_call.h.

High priority: PR description/title no longer match the implementation

The PR description (and the auto-generated reviewer overview) describe a different approach: adding an MEA eligibility guard in core/providers/cuda/llm/attention.cc so that nonpad_kv_seqlen != nullptr && head_size > 256 falls through to the unified unfused path. The current diff touches no such file and instead keeps MEA enabled for the large-head nonpad path while fixing the shared-memory opt-in. The Tests section also references a *_falls_back_from_mea_* test that is not present in the diff. Please update the title, description, and Tests section to match the shipped kernel-launch fix so reviewers and the #28388 resolution rationale are not misled.

The code change itself looks correct and well-scoped.

[tianleiwu]

The latest head addresses my previous feedback. The PR title/body now match the actual per-kernel shared-memory opt-in fix, and the regression test documents why the default-then-nonpad ordering is intentional.

I re-reviewed the current diff and the CUDA launch change still looks correct: the default and custom right-padding kernels are distinct device functions, each now gets its own cudaFuncAttributeMaxDynamicSharedMemorySize initialization, and the CUDA API result is checked with CUDA_CALL_THROW. The author response on the optional reverse-order test is reasonable because that coverage would require process isolation once the first sequence has initialized both variants.

No remaining findings from me.