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.

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

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