Enable GQA for KV-shared layers and fix dtype issues in vision and audio by apsonawane · Pull Request #279 · onnxruntime/mobius

apsonawane · 2026-05-06T21:30:32Z

[Gemma4] Enable GQA for KV-shared layers, fix vision encoder dtype, fix audio SkipNorm fusion

Summary

Four changes to enable optimized Gemma4 multimodal inference:

GQA for KV-shared layers — use GroupQueryAttention instead of standard Attention for the 20 KV-shared decoder layers
Vision encoder float32 input — accept float32 pixel_values and cast inside the model to prevent NaN
Audio encoder SkipNorm fix — prevent incorrect fusion of Add(MatMul, 1D_bias) + RMSNorm
Enable past_present_share_buffer — avoid per-step KV cache allocation

1. GQA for KV-shared layers (`gemma4.py`)

KV-shared layers borrow K/V from a source layer's present_key/present_value (already in BNSH format). Instead of Transpose+Reshape into BSNH for the standard Attention op, we now pass:

Empty K/V tensors [B, 0, kv_hidden] (kv_sequence_length=0)
Source layer's present K/V directly as past_key/past_value

This eliminates Transpose+Reshape ops from the graph for all 20 shared layers and uses the optimized GQA flash attention kernel. Requires corresponding ORT changes to support kv_sequence_length=0 (separate ORT PR).

Decoder loop change: All layers (shared and non-shared) now receive GQA context when the EP supports it. The previous is_shared exclusion and fallback attention bias/position embedding generation for shared layers are removed.

2. Vision encoder float32 input (`_gemma4.py`)

Problem: The image processor in ORT-GenAI outputs float32 pixel_values, but the vision encoder declared its input as config.dtype (float16). The ORT-GenAI trimming path did a raw memcpy from float32 source into an fp16 tensor using fp16 strides — reinterpreting float32 bit patterns as fp16, producing NaN.

Fix: Declare pixel_values as ir.DataType.FLOAT and add op.Cast(pixel_values, to=config.dtype) inside the ONNX model when needed. This keeps the host-side pipeline in float32 and lets the model handle the precision conversion safely.

3. Audio encoder SkipNorm fix (`_skip_norm.py`)

Problem: The AddRMSNormToSkipNorm rewrite rule fused Add(MatMul_result, 1D_bias) + RMSNorm into SkipSimplifiedLayerNormalization. The 1D bias [hidden] became the skip input, but the op requires both inputs to have the same shape ([B, S, hidden]). This caused:

CPU: skip is expected to have 3 or 2 dimensions, got 1
CUDA: input and skip shall have same shape when enable_skip_layer_norm_strict_mode is True

Fix: Add a rank equality check in the check() method. Fusion is skipped when the two Add inputs have different ranks (e.g., 3D tensor + 1D bias via broadcasting).

4. Enable `past_present_share_buffer` (`genai_config.py`)

Set past_present_share_buffer to true by default for all EPs. This pre-allocates KV cache buffers and aliases past/present pointers, avoiding per-step allocation and past→present copies.

Testing

Scenario	Status
Text-only generation (CPU)	✅
Text-only generation (CUDA)	✅
Multimodal image+text (CPU)	✅ Clean output, no repetition
Multimodal image+text (CUDA)	✅ Clean output, no crashes
Model graph	✅ 35 GQA nodes (15 source + 20 shared), 0 Attention nodes

…ix audio SkipNorm fusion

Copilot

Pull request overview

This PR updates Mobius’s Gemma4 multimodal export path to improve inference performance and correctness across text/vision/audio by expanding GroupQueryAttention usage, correcting vision input dtype handling, tightening a SkipNorm fusion rule, and changing ORT-GenAI config defaults.

Changes:

Gemma4 KV-shared decoder layers can now use com.microsoft.GroupQueryAttention by wiring shared KV as past_key/past_value and passing empty K/V.
Gemma4 vision encoder now declares pixel_values as FLOAT (fp32) and casts internally to the model dtype to avoid NaNs from host-side dtype reinterpretation.
SkipNorm fusion rule now rejects Add+broadcast patterns by checking Add input ranks; ORT-GenAI config default enables past_present_share_buffer.

Reviewed changes

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

File	Description
src/mobius/tasks/_gemma4.py	Makes Gemma4 vision encoder input `pixel_values` fp32 and casts to model dtype inside the graph.
src/mobius/rewrite_rules/_skip_norm.py	Adds a rank-equality guard to prevent invalid SkipNorm fusion when Add uses broadcasting (e.g., 1D bias).
src/mobius/models/gemma4.py	Enables GQA path for KV-shared layers by passing empty K/V and using shared present KV as past; simplifies fallback bias/pos-emb logic.
src/mobius/integrations/ort_genai/genai_config.py	Changes genai_config generation to always set `past_present_share_buffer` true (and apply max_length capping logic).

Co-authored-by: Copilot <copilot@github.com>

github-actions · 2026-05-06T21:39:45Z

Performance Comparison

Comparing fdf2c2e → 2dd2e58

Model	Metric	Baseline	Current	Delta
bert (feature-extraction)	model_size_bytes	359 KB	359 KB	+0.0%	⚪
bert (feature-extraction)	num_nodes	60	60	+0.0%	⚪
falcon	model_size_bytes	364 KB	364 KB	+0.0%	⚪
falcon	num_nodes	66	66	+0.0%	⚪
gemma2	model_size_bytes	428 KB	428 KB	+0.0%	⚪
gemma2	num_nodes	107	107	+0.0%	⚪
gpt2	model_size_bytes	388 KB	388 KB	+0.0%	⚪
gpt2	num_nodes	53	53	+0.0%	⚪
llama	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama	num_nodes	61	61	+0.0%	⚪
llama (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
llama (static-cache)	num_nodes	58	58	+0.0%	⚪
mamba (ssm-text-generation)	model_size_bytes	296 KB	296 KB	+0.0%	⚪
mamba (ssm-text-generation)	num_nodes	98	98	+0.0%	⚪
phi3	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3	num_nodes	59	59	+0.0%	⚪
phi3 (static-cache)	model_size_bytes	421 KB	421 KB	+0.0%	⚪
phi3 (static-cache)	num_nodes	56	56	+0.0%	⚪
qwen2	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2	num_nodes	61	61	+0.0%	⚪
qwen2 (static-cache)	model_size_bytes	425 KB	425 KB	+0.0%	⚪
qwen2 (static-cache)	num_nodes	58	58	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	model_size_bytes	506 KB	506 KB	+0.0%	⚪
qwen3_5_moe (hybrid-text-generation)	num_nodes	275	275	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	model_size_bytes	458 KB	458 KB	+0.0%	⚪
qwen3_5_text (hybrid-text-generation)	num_nodes	129	129	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	model_size_bytes	977 KB	977 KB	+0.0%	⚪
qwen3_5_vl (hybrid-qwen-vl)	num_nodes	408	408	+0.0%	⚪
t5 (seq2seq)	model_size_bytes	836 KB	836 KB	+0.0%	⚪
t5 (seq2seq)	num_nodes	166	166	+0.0%	⚪
whisper (speech-to-text)	model_size_bytes	1008 KB	1008 KB	+0.0%	⚪
whisper (speech-to-text)	num_nodes	128	128	+0.0%	⚪

No performance regressions.

codecov · 2026-05-06T21:40:22Z

The author of this PR, apsonawane, is not an activated member of this organization on Codecov.
Please activate this user on Codecov to display this PR comment.
Coverage data is still being uploaded to Codecov.io for purposes of overall coverage calculations.
Please don't hesitate to email us at support@codecov.io with any questions.

github-actions · 2026-05-06T21:42:09Z

🏗️ Architecture Diff

Comparing fdf2c2e → 2dd2e58

Model	Sub-model	Status
gemma4 (gemma4)	decoder	⚪
gemma4 (gemma4)	embedding	⚪
gemma4 (gemma4)	vision_encoder	⚪
gemma4_text	model	⚪

No architecture changes detected. ✅

Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed)

PR #279 sets do_rotary=1 for KV-shared GQA layers so RoPE is applied to queries. The source key already has RoPE baked in. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>

[Gemma4] Enable GQA for KV-shared layers, fix vision encoder dtype, f…

0068b23

…ix audio SkipNorm fusion

apsonawane requested review from a team and Copilot May 6, 2026 21:30

Copilot started reviewing on behalf of apsonawane May 6, 2026 21:31 View session

Copilot AI reviewed May 6, 2026

View reviewed changes

Comment thread src/mobius/models/gemma4.py

Comment thread src/mobius/models/gemma4.py Outdated

Comment thread src/mobius/rewrite_rules/_skip_norm.py Outdated

Comment thread src/mobius/integrations/ort_genai/genai_config.py Outdated

Merge branch 'main' into asonawane/gemma4

32fd042

Co-authored-by: Copilot <copilot@github.com>

apsonawane added 2 commits May 6, 2026 21:59

Address copilot comments

c2f984e

Fix lint

206ae99

justinchuby self-requested a review May 6, 2026 22:38

justinchuby reviewed May 6, 2026

View reviewed changes

Comment thread src/mobius/models/gemma4.py

justinchuby reviewed May 6, 2026

View reviewed changes

Comment thread src/mobius/tasks/_gemma4.py Outdated

justinchuby reviewed May 6, 2026

View reviewed changes

Comment thread src/mobius/models/gemma4.py Outdated

justinchuby approved these changes May 6, 2026

View reviewed changes

justinchuby mentioned this pull request May 6, 2026

Gemma4 ONNX + CUDA support: all PRs tracker #257

Open

37 tasks

apsonawane and others added 3 commits May 7, 2026 00:29

Address comments

552a729

Merge branch 'main' into asonawane/gemma4

ac0798a

Address comments

56f3570

apsonawane enabled auto-merge (squash) May 7, 2026 00:56

Merge branch 'main' into asonawane/gemma4

2dd2e58

apsonawane merged commit 41d2641 into main May 7, 2026
18 of 21 checks passed

apsonawane deleted the asonawane/gemma4 branch May 7, 2026 01:03

titaiwangms mentioned this pull request May 8, 2026

docs: refresh skills for latest Gemma4 state and references #283

Merged

Navigation Menu

Uh oh!

Enable GQA for KV-shared layers and fix dtype issues in vision and audio - #279

Enable GQA for KV-shared layers and fix dtype issues in vision and audio#279
apsonawane merged 8 commits into
mainfrom
asonawane/gemma4

apsonawane commented May 6, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented May 6, 2026 •

edited

Loading

Uh oh!

codecov Bot commented May 6, 2026

Uh oh!

github-actions Bot commented May 6, 2026 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Uh oh!

Conversation

apsonawane commented May 6, 2026

[Gemma4] Enable GQA for KV-shared layers, fix vision encoder dtype, fix audio SkipNorm fusion

Summary

1. GQA for KV-shared layers (gemma4.py)

2. Vision encoder float32 input (_gemma4.py)

3. Audio encoder SkipNorm fix (_skip_norm.py)

4. Enable past_present_share_buffer (genai_config.py)

Testing

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Performance Comparison

Uh oh!

codecov Bot commented May 6, 2026

Uh oh!

github-actions Bot commented May 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🏗️ Architecture Diff

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

1. GQA for KV-shared layers (`gemma4.py`)

2. Vision encoder float32 input (`_gemma4.py`)

3. Audio encoder SkipNorm fix (`_skip_norm.py`)

4. Enable `past_present_share_buffer` (`genai_config.py`)

github-actions Bot commented May 6, 2026 •

edited

Loading

github-actions Bot commented May 6, 2026 •

edited

Loading