Gemma4 CUDA EP: float additive bias + hybrid GQA/Attention dispatch - #269
Merged
Conversation
Three changes for Gemma4 CUDA EP support: 1. Use float additive bias for all fallback Attention layers (both sliding-window and full-attention). Encodes causal + sliding window + padding constraints. Required for batch>1 correctness. 2. Add feature flag MOBIUS_USE_GQA_FOR_KV_SHARED (default False). When False: normal layers use GQA, KV-shared layers use Attention. When True: all layers use GQA (requires ORT new_kv_length=0). 3. Remove unused create_padding_mask/create_sliding_window_mask imports. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
37 tasks
🏗️ Architecture Diff
No architecture changes detected. ✅ 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! |
rui-ren
approved these changes
May 6, 2026
ORT fuses Add(output_proj.bias) + RMSNormalization into SkipSimplifiedLayerNormalization, placing the 1D bias as the 'skip' input. ORT's CUDA kernel rejects 1D skip (requires 2D+), while the CPU kernel accepts it. This fusion was enabled by PR #253 which changed _Gemma4ScaleFreeRMSNorm from manual primitive ops to op.RMSNormalization(stash_type=1). The RMSNormalization op is recognized by ORT's SkipLayerNorm fusion pattern. Fix: inline manual RMSNorm ops in _Gemma4AudioEncoderModel.forward() for the pre_projection_norm, preventing ORT from recognizing the fusion pattern. This preserves the FP32 accumulation for numerical stability. Tested: audio encoder runs on CUDA with correct output (standalone ORT and GenAI at 117 tok/s). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchu@microsoft.com>
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
2 participants
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
Two changes for Gemma4 CUDA EP support:
1. Float additive bias for all fallback Attention layers
All KV-shared layers (and all layers on default EP) now use
create_attention_bias()which encodes causal + sliding window + padding constraints as a float additive bias. This:Removes unused
create_padding_maskandcreate_sliding_window_maskimports.2. Hybrid GQA/Attention dispatch with feature flag
Adds
MOBIUS_USE_GQA_FOR_KV_SHAREDflag (_flags.py, defaultFalse):False(default)Truenew_kv_length=0support)KV-shared layers cannot use GQA today because they pass the full KV cache as K input with no
past_key— the CUTLASS MEA aligned kernel crashes (tracked in microsoft/onnxruntime#28376). SetMOBIUS_USE_GQA_FOR_KV_SHARED=1when ORT GQA supports this pattern.Testing