onnxruntime
83e402e0 - [GQA] Make present_key/present_value outputs optional and add Gemma4 support (#28242)

Commit
106 days ago
[GQA] Make present_key/present_value outputs optional and add Gemma4 support (#28242) ## [GQA] Support KV-shared layers with empty K/V inputs (kv_sequence_length=0) ### Summary Enable `GroupQueryAttention` for KV-shared decoder layers (e.g., Gemma4) by allowing `kv_sequence_length=0` when `past_key`/`past_value` contain the borrowed KV cache. No new inputs, no schema changes, no GQA spec changes. ### Motivation Gemma4 has 20 KV-shared layers that borrow K/V from a source layer instead of computing their own. Previously these layers required the standard `Attention` op with Transpose+Reshape to convert the source's BNSH output to BSNH input. This PR enables the optimized GQA kernel for these layers, eliminating the Transpose/Reshape overhead and leveraging flash attention. ### Design KV-shared layers pass empty K/V tensors and wire the source layer's present K/V directly as past: ``` Q: query (RoPE applied by GQA via do_rotary=1) K: empty tensor [B, 0, kv_hidden] (kv_sequence_length = 0) V: empty tensor [B, 0, kv_hidden] (kv_sequence_length = 0) past_key: borrowed KV in BNSH (past_sequence_length > 0) past_value: borrowed KV in BNSH (past_sequence_length > 0) present_key/value: copy of past or aliased via past_present_share_buffer ``` No concatenation is needed since `new_kv_length = 0`. ### Changes | File | Change | |------|--------| | `group_query_attention_helper.h` | Allow `kv_sequence_length=0` when `past_key` is provided (previously required `kv_sequence_length == sequence_length`) | | `group_query_attention_impl.cu` | Add `kv_sequence_length==0` path in `PrepareQKV`: launch `LaunchUnpackRoPEAppend` with `kv_num_heads=0` so only Q head threads are spawned — no K/V memory access | | `group_query_attention.cc` (CPU) | Allow `kv_sequence_length=0` in `do_rotary` path; skip K RoPE when no K tokens exist | | `gqa_attention_base.h` (CPU) | Fix `past_seqlen` for shared KV: when `kv_sequence_length=0` and `past_key` exists, set `past_seqlen=total_seqlen` (all data is from past) instead of 0. Fixes incorrect attention over uninitialized present buffer during prompt phase | ### Why this approach Compared to the `kv_sequence_length != sequence_length` approach (passing full-context K/V as K/V inputs with different Q/K sequence lengths): - **No kernel grid split needed** — `kv_num_heads=0` cleanly eliminates K/V threads - **No Transpose_BSNH_to_BNSH** — shared K/V is already BNSH in the past buffer - **No `past_seqlen` offset issues** — there's nothing to append - **No new schema inputs** — uses existing `past_key`/`past_value` semantics - **Works with `past_present_share_buffer` both true and false** ### Testing - Text-only generation (CPU): 271-token prompt + multiple decode steps ✅ - Multimodal image+text generation (CPU): clean output, no repetition ✅ - Multimodal image+text generation (CUDA): clean output, no crashes ✅ - Model graph verification: 35 GQA nodes (15 source + 20 shared), 0 Attention nodes ✅
Author
Parents
Loading