Data-dependent bidirectional mask (Gemma-4 vision-block) forces standard Attention over GQA — split prefill/decode decoder graphs?

[justinchuby]

Summary

For models with a data-dependent bidirectional attention mask (e.g. Gemma‑4 multimodal's vision‑block bidirectional attention), the exported decoder is forced to use the standard ONNX Attention op with a float attention bias instead of com.microsoft.GroupQueryAttention. Because GQA is given up, the decoder also gives up its decode‑time fast paths — most importantly past_present_share_buffer (in‑place KV cache), plus XQA / Flash FastDecode.

The key observation is that this penalty is only fundamentally required during prefill. During token generation the bias degenerates to plain causal, so a decode‑only graph could legitimately use GQA. This looks like a strong fit for the existing decoder pipeline (run_on_prompt / run_on_token_gen), and I'd like guidance on whether a prefill/decode split is the recommended approach (and how to handle the KV‑cache handoff).

Background / concrete model

• Model: Gemma‑4‑12B (gemma4_unified), exported as a 4‑model ORT‑GenAI multimodal package (vision / audio / embedding / decoder) via mobius tooling.
• Decoder: 48 layers, mixed attention — sliding‑window local layers (head_dim=256) interleaved with global full‑attention layers (head_dim=512).
• HuggingFace semantics: contiguous image‑token blocks attend bidirectionally to each other (audio/text stay causal). This is the "vision‑block overlay".

Why GQA is dropped

The bidirectional image‑block mask cannot be expressed by GroupQueryAttention (causal / local‑window only — it rejects an arbitrary attention_bias). So mobius bakes the full mask (causal + sliding window + padding + the block‑wise bidirectional OR) into a float additive bias and runs the standard ONNX Attention op with is_causal=0. Result: the decoder graph is 48× Attention, 0× GQA, and genai_config.json therefore has past_present_share_buffer effectively off (no in‑place KV cache).

This is numerically correct everywhere (verified against HF — text/image/audio all match). It's purely a performance trade‑off.

The prefill vs decode asymmetry (the actual point)

The bidirectional bias only matters when image tokens are queries, which only happens during prefill (the image placeholders are in the current input sequence). During decode:

• the single query is a freshly generated text token,
• image tokens are only keys (already in past KV), never queries,
• so the new token just attends causally to all past KV → the overlay bias is a no‑op.

⇒ A decode step is plain causal attention and is fully GQA‑eligible. But because mobius emits a single static graph that must also serve prefill, every decode step pays the no‑GQA cost: no past_present_share_buffer, KV cache reallocated/copied as the sequence grows, no XQA/FastDecode. For long generations this is the dominant cost.

Proposal / question: split prefill and decode decoder graphs?

ORT‑GenAI already appears to support exactly this via the decoder pipeline (Config::Model::Decoder::PipelineModel with run_on_prompt / run_on_token_gen, dispatched in decoder_only_pipeline.cpp). Concretely the idea would be:

Stage	Graph	Mask	Op	Flags
Prefill	decoder_prefill.onnx	causal + sliding + padding + block OR (float bias)	Attention, is_causal=0	run_on_prompt=true, run_on_token_gen=false
Decode	decoder_decode.onnx	plain causal	GroupQueryAttention + past_present_share_buffer	run_on_prompt=false, run_on_token_gen=true

This would keep prefill correct while recovering GQA's in‑place KV cache and fast decode kernels.

Questions for maintainers:

1. Is the run_on_prompt / run_on_token_gen decoder pipeline the intended/supported mechanism for a heterogeneous prefill (standard Attention, dynamic present KV) + decode (GroupQueryAttention, pre‑allocated shared‑buffer KV) split?
2. KV‑cache handoff is the crux: prefill (Attention) emits a dynamic present.* (separate past/present), while the decode (GQA) graph expects a pre‑allocated past_present_share_buffer. Is there an existing path for the runtime to materialize the prefill KV into the GQA shared buffer (copy on first decode step), or must both stages agree on the same KV layout/allocation? Any precedent (e.g. NPU decoder‑pipeline models)?
3. Are there known examples of two pipeline stages sharing KV cache state, or is KV cache assumed private per pipeline model today?

Alternative (ORT core, not genai)

Alternatively, if com.microsoft.GroupQueryAttention accepted an additive attention_bias (or a block/custom mask input), a single GQA graph could serve both prefill and decode and this whole split would be unnecessary. Filing here first because the pipeline mechanism seems to already exist; happy to mirror to onnxruntime core if that's the preferred fix.

Environment

• onnxruntime‑genai built from source (CUDA), ORT 1.27.
• Model exported via mobius; package verified correct (text/image/audio match HF reference). This issue is strictly about decode‑time performance, not correctness.

[justinchuby]

Refinement: the real goal is single-graph dispatch to a fast kernel, and a key constraint on Flash

A reviewer raised a good point: rather than (or in addition to) splitting models, what ultimately matters is that the single exported graph can still dispatch to a fast attention kernel. Two important clarifications on what's actually reachable:

1. An explicit attention_bias fundamentally precludes Flash — and dispatch keys on presence, not content.

Flash requires attn_mask == nullptr. So any explicit additive bias disqualifies Flash, even a numerically no-op bias (e.g. all-causal with no block tokens during decode). The kernel selection looks at whether a bias tensor is wired into the op, not at its values. This is the core reason a static graph that carries the overlay bias node can never reach Flash at decode time — even though decode is semantically plain causal. The com.microsoft.GroupQueryAttention contrib op goes further and rejects a bias outright.

2. For Gemma-4 specifically, Flash is unreachable for every layer regardless of the bias.

• Local layers: sliding-window ⇒ require an explicit mask ⇒ no Flash.
• Global layers: head_dim = 512 ⇒ exceeds Flash's 256 limit ⇒ no Flash.

So "dispatch to Flash" is not the realistic target here. The realistic fast path for this shape is:

• Decode: GroupQueryAttention + past_present_share_buffer (in-place KV), plus FastDecode (and XQA for the head_dim=256 local layers). The in-place KV cache is the dominant win.
• Prefill: MEA / CUTLASS memory-efficient attention (bias-aligned).

Implication for the ask. For a single graph to stay on the fast path at decode, the runtime/op would need to treat a decode step (seq_len==1, no image-token queries) as plain causal and route it to GQA — which today is blocked because the bias node is statically present and GQA rejects bias. Concretely the options are:

1. Split prefill/decode via the existing run_on_prompt / run_on_token_gen pipeline (issue body) — decode graph uses bias-free GQA + shared buffer. (KV-cache handoff question still stands.)
2. A bias-aware efficient attention in ORT core: let GroupQueryAttention (or an MEA-class GQA path) accept an optional additive bias / block-mask and keep past_present_share_buffer. This would not be Flash (a bias always precludes Flash), but MEA-class + in-place KV on a single graph would recover most of the loss for shapes like Gemma-4's.

Would appreciate maintainers' view on (1) vs (2), and on which fast kernels are reachable for a head_dim=512 + sliding-window GQA decode in the current CUDA cascade.

[justinchuby]

Update: the ONNX Attention op already supports bias + GQA + fast external KV cache — the gap is runtime KV management, and a prefill/decode split may be unnecessary

Digging into ORT's unified ONNX-domain Attention CUDA impl (onnxruntime/core/providers/cuda/llm/attention.cc) changes the conclusion above. That op already supports the combination this model needs, on the MEA (CUTLASS) path:

• nonpad_kv_seqlen (opset 24) external KV cache + attn_mask/additive bias + GQA are supported together. The code explicitly converts the mask to an additive bias and runs MEA "for both MHA and GQA when nonpad_kv_seqlen is provided" (attention.cc ~L778-781, L834).
• The in-source perf note (L205-231) states the external-cache path (TensorScatter + nonpad_kv_seqlen) "achieves near-parity with contrib GQA performance"; the internal past_key/present_key path is the ~15-30% slower one.

So for Gemma-4's shape the realistic fast target is MEA + external KV cache with the bias, on a single standard-Attention graph — which means the prefill/decode split proposed above may not be necessary. (Flash stays out regardless: any bias precludes it, and the head_dim=512 global layers exceed Flash's 256 limit. MEA-external ≈ GQA parity is the ceiling here, and it accepts the bias.)

Where the fast path is actually lost today:

1. genai doesn't drive the opset-24 external-cache path. The KV-cache manager only knows past_present_share_buffer (contrib GQA in-place) vs the default dynamic past/present; there's no reference to nonpad_kv_seqlen / TensorScatter external cache in src/models/kv_cache.cpp. So even an exported graph using nonpad_kv_seqlen wouldn't be driven on the fast path.
2. Exporter emits this decoder with the internal past/present cache (the slower concat path), not the external-cache form.

Revised ask (supersedes the split proposal): would it be feasible for genai to support the opset-24 external KV cache (nonpad_kv_seqlen + TensorScatter) in the decoder KV-cache manager? That lets a single standard-Attention graph carry the bidirectional-block bias in prefill and hit the near-GQA-parity MEA external-cache path at decode — no model split, no contrib GQA, and no reliance on Flash. The exporter would then emit the nonpad_kv_seqlen form for models that need a bias (Gemma-4 vision-block, sliding-window models, etc.).

Does that match the intended long-term direction for KV-cache management in genai? Happy to prototype the exporter side once the runtime path is confirmed.