[GQA] Make present_key/present_value outputs optional and add Gemma4 support

[Akshay Sonawane (apsonawane)]

[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 ✅

[github-actions]

You can commit the suggested changes from lintrunner.

[Copilot]

Pull request overview

Adds support for KV-shared decoder layers by allowing com.microsoft.GroupQueryAttention to optionally consume pre-computed external K/V tensors (instead of maintaining/updating its own KV cache), enabling architectures like Gemma4-style KV sharing.

Changes:

• Extended the GroupQueryAttention schema with optional inputs external_key / external_value (indices 14/15).
• Added new parameters + validation helpers to detect/configure “external KV” mode and enforce do_rotary=0.
• Updated CPU and CUDA kernels to source KV from external tensors and bypass KV-cache update / RoPE-on-KV paths.

Reviewed changes

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

Show a summary per file

File	Description
onnxruntime/core/graph/contrib_ops/bert_defs.cc	Adds schema inputs for external KV (but type/shape inference also needs external-KV awareness).
onnxruntime/contrib_ops/cpu/bert/attention_parameters.h	Adds use_external_kv and external_kv_sequence_length to GQA parameters.
onnxruntime/contrib_ops/cpu/bert/group_query_attention_helper.h	Adds Q-only checks and external-KV shape validation/configuration helpers; updates CheckInputs to distinguish packed-QKV vs Q-only.
onnxruntime/contrib_ops/cpu/bert/group_query_attention.cc	Plumbs external KV inputs into CPU kernel and skips K/V transpose + rotary when external KV is used.
onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h	Updates CPU attention core to skip KV concatenation and copy external KV to present outputs once per KV head.
onnxruntime/contrib_ops/cuda/bert/attention_data.h	Adds external KV pointers to CUDA attention data struct.
onnxruntime/contrib_ops/cuda/bert/group_query_attention.cc	Plumbs external KV inputs into CUDA kernel and enforces do_rotary=0 for external KV mode.
onnxruntime/contrib_ops/cuda/bert/group_query_attention_impl.cu	Adds external-KV path in PrepareQKV (copy external KV to present and skip append/RoPE).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Tianlei Wu (tianleiwu)]

There is no need to add extra inputs, you can use key/value for that, and make past_key/past_value/present_key/present_value as optional.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated 3 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Tianlei Wu (tianleiwu)]

I found one additional correctness issue on the current head. There are also already-open current-head threads covering the optional-present tests and the CUDA/documentation mismatch, so I am not duplicating those here.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated 7 comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Tianlei Wu (tianleiwu)]

Most of my concerns on this head are already covered by existing review threads (CPU GEMM/concat invariant, CUDA hard-error vs. optional schema, mixed-output configurations, scratch buffer sizing on CUDA, seqlen_present_kv_cache initialization, test tolerance). Two items I did not see covered:

1. PR description does not describe this PR. The description discusses adding external_key/external_value inputs (inputs 14/15), Check_Q_Only, CheckExternalKV, use_external_kv, CUDA attention_data.h fields, etc. None of that appears in this diff — the actual change makes existing present_key/present_value outputs optional. Please rewrite the description to match the implementation; downstream tooling and release notes rely on it.

2. CUDA test coverage is missing. All new tests use DefaultCpuExecutionProvider() only. The CUDA path has its own contracts added in this PR — the early guard in group_query_attention.cc (claims first-prompt is supported) and the unconditional rejection in PrepareQKV (group_query_attention_impl.cu). These two messages contradict each other, but no CUDA test exercises omit_present=true, so the contradiction is invisible to CI. Please add at least one CUDA-gated negative test that asserts the kernel rejects omitted present outputs with a stable error message (or, if the intent is to support it on CUDA, a positive equivalence test). Otherwise a future refactor of PrepareQKV can silently change the user-facing error path.

No new inline comments — the existing threads already pinpoint the code locations.

[Tianlei Wu (tianleiwu)]

Review of head 6cbe62c

The two concerns I raised on 9a14803a have been addressed:

• PR description body now accurately describes the optional-present-output change.
• CUDA test (OptionalPresent_CudaOmitMatchesConnected) was added.

The is_first_prompt validation in both CPU and CUDA kernels properly constrains the omitted-present case, and the scratch-buffer approach on CUDA correctly keeps data.present_key/data.present_value non-null for downstream kernels. I resolved my earlier thread on gqa_attention_base.h line 177.

One remaining concern:

WebGPU EP still has no guard for omitted present outputs. The schema relaxation (OpSchema::Optional on outputs 1/2) is global across all EPs. The WebGPU kernel (onnxruntime/contrib_ops/webgpu/bert/group_query_attention.cc ~line 240) calls context.Output(1, present_kv_shape) and then passes the result to ApplyFlashAttention/ApplyAttention without a nullptr check. A model that omits these outputs will crash at runtime on WebGPU. Please add a validation that rejects present_key == nullptr || present_value == nullptr with a clear message (e.g., "WebGPU GroupQueryAttention requires present_key and present_value outputs").

Nitpick: The PR title still says "Add external_key/external_value inputs" — please update to match the description body.

[Tianlei Wu (tianleiwu)]

Two suggestions below — neither is blocking.

[Tianlei Wu (tianleiwu)]

I found two CUDA shared-KV correctness issues that look worth fixing before this merges. Both are in the decode/cache-handling path, so I am marking this as request changes.

[Copilot]

Pull request overview

Copilot reviewed 7 out of 7 changed files in this pull request and generated 3 comments.

Comments suppressed due to low confidence (1)

onnxruntime/contrib_ops/cpu/bert/gqa_attention_base.h:495

• In ComputeVxAttentionScore, when present_value is nullptr (omitted output), the code never concatenates or otherwise incorporates past_value (the ConcatStateChunkGQA call is guarded by if (nullptr != present_value)). If past_value is provided, this will compute the output using only the (possibly empty) V input and ignore cached history. Consider the same fix as for keys: require present outputs when past is provided, or allocate a temporary concatenated buffer / handle kv_sequence_length==0 by reading directly from past.

        const size_t batch_index = i / num_heads_;
        const size_t head_index = i % num_heads_;
        const size_t total_seqlen = SafeInt<size_t>(seqlens_k[batch_index]) + 1;
        size_t past_seqlen;
        if (past_value == nullptr) {
          past_seqlen = 0;
        } else if (kv_sequence_length == 0) {
          past_seqlen = total_seqlen;
        } else if (is_prompt) {
          past_seqlen = 0;
        } else {
          past_seqlen = total_seqlen - sequence_length;
        }
        const size_t past_chunk_length = SafeInt<size_t>(past_seqlen) * head_size;

        const T* v;
        if (packed_qkv) {
          v = V + packed_batch_stride * batch_index + kv_input_chunk_length * (head_index / kv_num_heads_factor);
        } else {
          v = V + kv_input_chunk_length * (i / kv_num_heads_factor);
        }
        if (nullptr != present_value) {
          v = ConcatStateChunkGQA(past_value, v, present_value, present_buff_chunk_length, past_buff_chunk_length,
                                  past_chunk_length, kv_input_chunk_length, past_present_share_buffer,
                                  i / kv_num_heads_factor);
        }

[Tianlei Wu (tianleiwu)]

Thanks for the updates. The earlier CUDA alias consistency and shared-KV fast-decode routing concerns look addressed in this head, but I found one remaining test issue that leaves the new CUDA shared-KV path effectively unverified.

[Tianlei Wu (tianleiwu)]

Thanks for the update. The CUDA shared-KV tests now use an fp16 dtype that is registered for the CUDA GroupQueryAttention kernel, and the previous test coverage concern is addressed. I do not see remaining issues in the latest diff.