Skip to content

Emit opset-24 external KV cache (nonpad_kv_seqlen) for bias decoders to recover near-GQA decode (Gemma-4, sliding-window) #349

Description

@justinchuby

Summary

For models that require a float attention bias on the decoder (e.g. Gemma‑4 multimodal's bidirectional vision‑block overlay, or any sliding‑window model), mobius currently emits the decoder with an internal past/present KV cache. With a bias present this routes ORT's ONNX‑domain Attention op (core/providers/cuda/llm/attention.cc) to the internal‑cache MEA path, which the op's own in‑source perf note (L205‑231) calls ~15‑30% slower than contrib GQA's in‑place decode.

ORT already supports a faster alternative on the same standard Attention op: the opset‑24 external KV cache (nonpad_kv_seqlen + TensorScatter). On the MEA path it combines attn_mask/additive bias + GQA + external cache and "achieves near‑parity with contrib GQA performance" (same note, plus ~L778‑781, L834). This would let a single standard‑Attention graph carry the bidirectional‑block bias during prefill and hit the near‑GQA decode path — no contrib GQA, no prefill/decode model split, and correctly (Flash is precluded by any bias and by head_dim=512 global layers, so MEA‑external ≈ GQA is the realistic ceiling here).

Task

Add an exporter path that emits the decoder using the opset‑24 external KV cache form (nonpad_kv_seqlen + TensorScatter) for models that need a decoder attention bias, instead of internal past/present.

Applies to Gemma‑4 (gemma4_unified, use_bidirectional_attention == "vision") and is generally relevant to sliding‑window / float‑bias decoders.

Blocked on / coordinate with runtime

This is only end‑to‑end useful once onnxruntime‑genai drives the external‑cache path. Today genai's KV‑cache manager (src/models/kv_cache.cpp) only knows past_present_share_buffer (contrib GQA in‑place) vs the default dynamic past/present; it has no nonpad_kv_seqlen / TensorScatter support. Tracked upstream in microsoft/onnxruntime-genai#2204.

Acceptance / verification

  • Decoder graph emits nonpad_kv_seqlen (opset 24) external cache; genai_config.json wired accordingly.
  • Numerical parity vs HuggingFace preserved for text / image / audio (currently verified on the internal‑cache build at justinchuby/gemma-4-12b-onnx).
  • Decode‑time profiling shows the MEA external‑cache path (near contrib‑GQA), not the internal‑cache concat path.

Background / current state

  • Root cause + the bias→no‑GQA trade‑off is documented in PR Fix Gemma4 bidirectional attention + add gemma-4-12B (gemma4_unified) #338 and the genai issue above.
  • Gemma‑4‑12B currently ships (correct, internal‑cache) at justinchuby/gemma-4-12b-onnx across f16/bf16/Q4_K_M × cuda/default/cpu. The decoder is 48× standard Attention (0 GQA) because the vision‑block overlay needs a float bias.

Metadata

Metadata

Assignees

Labels

No labels
No labels

Type

No type

Projects

No projects

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions