Fix Gemma4 CUDA EP: GQA head_dim guard, provider_options, opset lowering - #250
Closed
justinchuby wants to merge 10 commits into
Closed
Fix Gemma4 CUDA EP: GQA head_dim guard, provider_options, opset lowering#250justinchuby wants to merge 10 commits into
justinchuby wants to merge 10 commits into
Conversation
ORT's GroupQueryAttention CUDA kernel does not support head_dim > 256, causing illegal memory access for Gemma4 full_attention layers (head_dim=512). Add per-layer head_dim guard in Gemma4TextModel.forward() so that only sliding_attention layers (head_dim=256) use GQA, while full_attention layers (head_dim=512) fall back to standard Attention with manual RoPE. This changes the CUDA EP model from 15 GQA + 20 Attention to 12 GQA + 23 Attention, with layers 4, 9, 14 correctly using Attention. The fallback infrastructure (mask + position_embeddings) is also triggered for these layers. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Two fixes for Gemma4 CUDA EP inference: 1. Remove CUDA provider_options from EP defaults. Explicit CUDA provider_options in genai_config.json conflict with GenAI's internal session setup (ClearOutput, ReuseEmbeddingsBuffer, etc.), causing NaN or crashes for multimodal CUDA models. GenAI's C++ code handles all CUDA EP configuration internally. 2. Use float16 additive masks for all fallback Attention layers (KV-shared and head_dim>256 layers) instead of bool masks. Bool masks triggered NaN in ORT's CUDA ConvertAttnMaskToBias path. Float16 masks match the working default EP model's behavior. Tested: CUDA EP model (12 GQA + 23 Attention) generates at 13.5 tok/s through GenAI with valid output. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
ORT's CUDA EP registers kernels up to opset 23 for standard ops (Reshape, RMSNormalization, etc.). When the model declares opset 24, these ops fall to CPUExecutionProvider, creating ~280 MemcpyFromHost and MemcpyToHost nodes that destroy inference performance. This implements the ort_lower_opset_for_ep flag (which was declared in _flags.py but never wired up). When enabled (default), the opset is lowered from 24 to 23 for all non-default EPs after optimization. Result: 282 memcpy → 4 memcpy for Gemma4 CUDA EP model. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
🏗️ Architecture Diff
gemma4 (gemma4) / decoder — 9 change(s)Op summary: 125 → 127 nodes --- base
+++ head
@@ -47,13 +47,14 @@
MatMul
Constant
Reshape
+Cast
Mul
ReduceMean
Constant
-CastLike
Add
Sqrt
Div
+CastLike
Reshape
Attention
Transpose
@@ -90,13 +91,14 @@
MatMul
Constant
Reshape
+Cast
Mul
ReduceMean
Constant
-CastLike
Add
Sqrt
Div
+CastLike
Reshape
Attention
TransposeAdded nodes:
Removed nodes:
Connectivity changes:
gemma4 (gemma4) / vision_encoder — 41 change(s)Op summary: 202 → 204 nodes --- base
+++ head
@@ -50,14 +50,15 @@
Reshape
RMSNormalization
RMSNormalization
+Cast
Mul
Constant
ReduceMean
Constant
-CastLike
Add
Sqrt
Div
+CastLike
Constant
Gather
Constant
@@ -188,15 +189,16 @@
Constant
CastLike
Mul
+Cast
Mul
Constant
ReduceMean
Constant
-CastLike
Add
Sqrt
Div
-Transpose
-MatMul
-Constant
-Reshape
+CastLike
+Transpose
+MatMul
+Constant
+ReshapeAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
gemma4_text / model — 9 change(s)Op summary: 127 → 129 nodes --- base
+++ head
@@ -49,13 +49,14 @@
MatMul
Constant
Reshape
+Cast
Mul
ReduceMean
Constant
-CastLike
Add
Sqrt
Div
+CastLike
Reshape
Attention
Transpose
@@ -92,13 +93,14 @@
MatMul
Constant
Reshape
+Cast
Mul
ReduceMean
Constant
-CastLike
Add
Sqrt
Div
+CastLike
Reshape
Attention
TransposeAdded nodes:
Removed nodes:
Connectivity changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Performance Comparison
|
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR targets improved Gemma4 performance/correctness on ORT CUDA EP by (1) gating GQA usage per-layer based on head_dim compatibility, (2) removing/avoiding explicit CUDA provider_options in generated GenAI configs, and (3) wiring an opset 24→23 lowering workaround to prevent CUDA kernel fallback/memcpy explosions.
Changes:
- Add per-layer GQA compatibility guard in Gemma4 text decoder and update fallback mask path.
- Adjust ORT GenAI provider-options generation to return
[]when no provider-specific options are needed (notably for CUDA defaults). - Implement
flags.ort_lower_opset_for_epin the build pipeline by lowering the default-domain opset import from 24 to 23 for non-defaultEPs.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| src/mobius/models/gemma4.py | Adds head_dim-based per-layer fallback from GQA to standard Attention and switches fallback masks to additive bias. |
| src/mobius/integrations/ort_genai/ep_config.py | Makes make_provider_options() return an empty list when no options are present to avoid overriding GenAI internal CUDA setup. |
| src/mobius/_execution_providers.py | Removes CUDA default provider_options from the EP registry to prevent GenAI conflicts. |
| src/mobius/_builder.py | Wires opset-lowering flag into build_from_module() by mutating model opset imports post-optimization. |
justinchuby
marked this pull request as draft
May 5, 2026 18:14
Tests now accept empty provider_options for CUDA EP, matching the change where explicit CUDA provider_options were removed to avoid conflicting with GenAI's internal EP configuration. 118 ort_genai tests pass. Signed-off-by: Justin Chu <justinchu@microsoft.com>
ORT now supports head_dim up to 512 in the GQA CUDA kernel. Remove the _MAX_GQA_HEAD_DIM guard that forced full-attention layers (head_dim=512) to fall back to standard Attention. All non-shared layers now use GroupQueryAttention regardless of head_dim. Only KV-shared layers still fall back to Attention (they borrow K,V and have no own KV cache). Signed-off-by: Justin Chu <justinchu@microsoft.com>
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
| assert len(provider_opts) == 1 | ||
| assert "cuda" in provider_opts[0] | ||
| assert isinstance(provider_opts, list) | ||
| pass # CUDA provider_options may be empty |
GenAI derives its providers list from provider_options names in
genai_config.json (config.cpp:1763-1764). Empty provider_options
means GenAI uses CPU-only, even for CUDA models.
Revert the 'empty for CUDA' behavior — always emit [{"cuda": {}}]
so GenAI registers the CUDA execution provider. The empty options
dict is fine — GenAI handles CUDA configuration internally.
118 ort_genai tests pass.
Signed-off-by: Justin Chu <justinchu@microsoft.com>
| assert len(provider_options) == 1, f"{block} missing CUDA provider options" | ||
| assert "cuda" in provider_options[0], f"{block} has wrong EP in provider_options" | ||
| assert isinstance(provider_options, list), f"{block} invalid provider_options" | ||
| pass # CUDA provider_options may be empty |
justinchuby
commented
May 5, 2026
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Member
Author
|
@titaiwangms could you help with reviewing the logic that handles the attention bias? I wasn't sure it this is the right fix. |
Gemma4's parameterless V normalization (v / sqrt(mean(v²) + ε)) squared FP16 values directly via op.Mul(v, v). V projection outputs can reach ~888, and 888² = 788,544 which overflows FP16 max (65504), producing inf → mean(inf) → sqrt(inf) → v/inf = 0. This caused all-zero V outputs on CUDA (CPU uses FP32 accumulation internally). Fix: Cast to FP32 before squaring, compute the full RMSNorm in FP32, then CastLike back to the input dtype. Applied to: - _Gemma4ScaleFreeRMSNorm.forward (vision encoder, projector norms) - Gemma4Attention.forward inline V norm (GQA path) - Gemma4Attention.forward inline V norm (non-GQA path) Result: F16 CUDA inference works — 151.5 tok/s on H200 (was NaN). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
Set stash_type=1 (FLOAT) on all RMSNormalization ops to ensure the variance computation uses FP32 internally when input is FP16/BF16. This prevents potential overflow when squaring large values. Applied to apply_rms_norm() and OffsetRMSNorm (GatedRMSNorm already had stash_type=1). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
justinchuby
force-pushed
the
fix-kv-shared-gqa
branch
from
May 5, 2026 18:48
a7708cc to
ea4e95c
Compare
…lation" This reverts commit ea4e95c.
Member
Author
37 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Fixes Gemma4 CUDA EP inference pipeline:
1. GQA head_dim guard (per-layer)
ORT's GQA CUDA kernel max head_dim is 256. Gemma4 full_attention layers (4, 9, 14) have head_dim=512. Added per-layer gate in
Gemma4TextModel.forward()so only compatible sliding_attention layers use GQA.Model: 12 GQA + 23 Attention (was 15 GQA + 20 Attention)
2. Remove CUDA provider_options from genai_config
Explicit CUDA provider_options in genai_config.json conflicted with GenAI's internal session setup, causing NaN. Empty
provider_options: []lets GenAI handle CUDA configuration internally.3. Float16 masks for fallback layers
Use float16 additive masks instead of bool masks for KV-shared and head_dim>256 layers. More robust on CUDA.
4. Opset 24→23 lowering
Implements the
ort_lower_opset_for_epflag (was declared but never wired up). Lowers opset from 24 to 23 for non-default EPs to avoid missing CUDA kernel registrations. Workaround until ORT adds opset 24 registrations (PR microsoft/onnxruntime#28365).Testing
All 1214 build graph tests pass. GenAI generation: 12.5 tok/s with 4 memcpy.