ONNX Attention MEA errors out with cudaErrorInvalidValue for KV-shared layers with borrowed K/V

[justinchuby]

Problem

ONNX Attention op's CUTLASS MEA (Memory Efficient Attention) aligned kernel crashes with cudaErrorInvalidValue at specific sequence lengths (multiples of 8) for KV-shared layers where K/V tensors are borrowed from another layer's present_key/present_value output.

Root cause

The crash is in the aligned CUTLASS FMHA kernel (selected when total_kv % kAlignmentQ == 0, i.e. total_kv % 8 == 0 for fp16). The unaligned kernel (selected when total_kv % 8 != 0) works correctly for the same inputs.

This affects KV-shared Attention layers where:

• past_key == nullptr (K/V is the full cache from another layer, not incrementally concatenated)
• An additive attention bias (mask) is present
• total_kv_length is a multiple of 8

seq_len	total_kv % 8	Kernel path	Result
1	1	unaligned	✅
4	4	unaligned	✅
7	7	unaligned	✅
8	0	aligned	❌ crash
9	1	unaligned	✅
16	0	aligned	❌ crash
27	3	unaligned	✅
32	0	aligned	❌ crash

Verified workaround

Adding force_unaligned flag to MemoryEfficientAttentionParams and setting it when past_key == nullptr && attn_bias != nullptr prevents the crash at all sequence lengths. However, the unaligned kernel produces slightly different numerics (cos_sim ~0.965 vs CPU), suggesting the MEA kernel may not fully handle the "full cache as K without past" pattern correctly.

Current plan

Default EP (ONNX Attention with unfused fallback) at 157.8 tok/s is the shipping configuration. The ORT team is investigating GQA support for new_kv_length=0 (KV-shared pattern), which would be the proper fix. When GQA supports KV-shared layers natively, the model can switch to CUDA EP with GQA for potentially better performance.

Model architecture (Gemma4 E2B-it)

Gemma4 has 35 decoder layers:

• Layers 0-14: Non-shared (own KV cache) — work fine with both GQA and ONNX Attention
• Layers 15-34: KV-shared — borrow K/V from source layers via Transpose+Reshape of present.N.key. No past_key/past_value inputs.

Performance data

Configuration	tok/s	VRAM	Status
F16 INT4 default EP	176.7	9 GB	✅ Best speed
F16 default EP	157.8	12 GB	✅ Shipping config
F16 CUDA EP (GQA + Attention)	151.4	14 GB	⚠️ Needs head_dim guard
F32 default EP	143.9	22 GB	✅

Environment

• GPU: NVIDIA H200 (SM 9.0)
• ORT: 1.27.0 from main
• CUDA: 13.0, cuDNN: 9
• Model: google/gemma-4-e2b-it (2.6B)

[justinchuby]

Update: Need to confirm whether GQA and ONNX Attention actually use different MEA kernel templates for head_dim=512. A colleague says they use the same CUTLASS kernel. The performance difference (151.4 vs 157.8 tok/s) is NOT from attention kernels — profiling shows CUDA EP raw compute is actually 2x faster (215ms vs 420ms total). The tok/s gap is from GenAI pipeline overhead, not ORT kernel performance.

The GQA MEA crash at seq=8/16/32 with has_custom_right_padding=true still needs investigation — it may be a launch configuration issue specific to the right_padding variant, not shared memory.

[tianleiwu]

If shared K/V in Gemma is in BNSH format, we shall use GQA like the following:

Q: query
K: empty tensor (new_kv_length = 0)
V: empty tensor (new_kv_length = 0)
past_key: shared k in BNSH format (past_sequence_length > 0)
past_value: shared v in BNSH format (past_sequence_length > 0)
present_key: point to same buffer of past_key (No concatenation is needed since new_kv_length = 0)
present_value: point to same buffer of past_value (No concatenation is needed since new_kv_length = 0)

Of course, we need update GQA kernel to support such special case (new_kv_length = 0)

ONNX attention op supports k/v to be 4D (BNSH), we can pass them as K/V directly without kv cache. Normally, it will not have alignment issue since H is usually multiple of 128 even we transpose BNSH to BSNH format. But we shall be careful that usually we allocates max sequence length for kv cache in GenAI, so the S of shared k/v might be max_sequence_length. So we need make sure the length of actual k and v is not derived from tensor shape. If that is the case, we might follow GQA usage to use past_key/past_value to pass shared k/v, which can be consistent between GQA and onnx Attention op usage.

[yuslepukhin]

Is it really a crash? If so, how do you know the error code?