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

### 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](https://github.com/microsoft/onnxruntime) 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.