Skip MEA in ONNX Attention for head_dim > 256 (CUTLASS shared memory limit) - #28383
Closed
justinchuby wants to merge 4 commits into
Closed
Skip MEA in ONNX Attention for head_dim > 256 (CUTLASS shared memory limit)#28383justinchuby wants to merge 4 commits into
justinchuby wants to merge 4 commits into
Conversation
ONNX opset 24 bumped Reshape and Cast (added float8e8m0 type support). ORT CUDA EP only had opset 23 registrations, so these ops fell to CPUExecutionProvider on opset 24 models, producing ~280 MemcpyFromHost/MemcpyToHost nodes. Version existing opset 23 registrations to (23, 23) and add new non-versioned opset 24 registrations. Same kernel implementations. Result: 282 memcpy → 4 memcpy for opset 24 models on CUDA EP. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
BiasLoader hardcoded 128-bit (8 fp16 element) vectorized loads via ElementsPerAccess = 128 / sizeof_bits<scalar_t> regardless of the isAligned template parameter. When attention bias stride (total_sequence_length) was not a multiple of 8, the unaligned kernel was selected but BiasLoader still used 128-bit loads on the bias, causing cudaErrorMisalignedAddress. Fix: Use kAlignmentA (kMinimumAlignment=4 for unaligned path, kAlignmentA=8 for aligned path) as BiasLoader's ElementsPerAccess. This allows the unaligned kernel to use 64-bit loads for the bias. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
CUTLASS FMHA shared memory scales with queries_per_block × head_dim. For head_dim > 256 (e.g. Gemma4 full_attention with head_dim=512), the kernel exceeds the SM's shared memory limit, causing cudaErrorInvalidValue on launch. Add head_size <= 256 check to the MEA eligibility in GQA. This lets the unfused fallback (cuBLAS GEMM, any head_dim) handle these layers. Tested: Gemma4 CUDA EP with GQA head_dim=512 layers at 150 tok/s. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
CUTLASS FMHA's custom right-padding kernel variant (used when nonpad_kv_seqlen is provided) has shared memory that scales with queries_per_block × head_dim. For head_dim > 256 (e.g. Gemma4 full-attention with head_dim=512), this exceeds the SM shared memory limit, causing cudaErrorInvalidValue on launch. The standard MEA path (without nonpad_kv_seqlen) works fine for any head_dim, so the guard is scoped to only the nonpad_kv_seqlen case. Fix: skip MEA when nonpad_kv_seqlen is present AND head_size > 256. These layers fall through to unfused attention (cuBLAS GEMM). Tested on H200: - hd=512 without nonpad: OK (MEA still used) - hd=512 with nonpad: OK (unfused fallback, previously crashed) - hd=64/256 with nonpad: OK (MEA still used) Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
fix-cuda-opset24-reshape-cast
branch
from
May 6, 2026 18:03
2601af2 to
5541d36
Compare
Contributor
Author
|
Updated: Narrowed the fix based on feedback. The Refined fix: skip MEA only when Verified:
|
Contributor
Author
|
Outdated |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
CUTLASS FMHA shared memory scales with
queries_per_block × head_dim. Forhead_dim > 256(e.g. Gemma4 full-attention withhead_dim=512), the kernel exceeds the SM's shared memory limit, causingcudaErrorInvalidValueon launch.This is the same fix as c72a5f4 (GQA op) applied to the ONNX Attention op. The GQA contrib op already has
parameters.head_size <= 256in its MEA eligibility check, but the standard ONNX Attention op was missing this guard.Fix
Add
parameters.head_size <= 256to the MEA eligibility check inattention.cc. Layers withhead_dim > 256now fall through to unfused attention (cuBLAS GEMM), which handles any head_dim.Test
Verified with ONNX Attention op (opset 24):
head_dim=64: ✅ (uses Flash/MEA as before)head_dim=256: ✅ (uses Flash/MEA as before)head_dim=512: ✅ (now falls to unfused, previously crashed)head_dim=512: ✅All with
is_causal=0,nonpad_kv_seqlen, fp16.Motivation
This affects Gemma4 models which use dual head dimensions:
head_dim=256for sliding-window attention layers andhead_dim=512for full-attention (global) layers.