Repository navigation
Align shared Gemma4 text and multimodal runtime semantics - #787
Merged
titaiwangms merged 5 commits intoOct 27, 2026
Merged
Conversation
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
🏗️ Architecture Diff
falcon / model — 3 change(s)Op summary: 66 → 68 nodes --- base
+++ head
@@ -1,6 +1,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Connectivity changes:
gemma2 / model — 11 change(s)Op summary: 105 → 109 nodes --- base
+++ head
@@ -2,6 +2,8 @@
Mul
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
Shape
@@ -16,6 +18,8 @@
Where
Cast
Unsqueeze
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
gemma4 (gemma4) / decoder — 11 change(s)Op summary: 125 → 129 nodes --- base
+++ head
@@ -2,6 +2,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
Shape
@@ -16,6 +18,8 @@
Where
Cast
Unsqueeze
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
gemma4 (gemma4) / vision_encoder — 56 change(s)Op summary: 176 → 174 nodes --- base
+++ head
@@ -1,9 +1,7 @@
-Constant
+Cast
+Sub
+Mul
CastLike
-Constant
-CastLike
-Mul
-Sub
Transpose
MatMul
Constant
@@ -133,11 +131,10 @@
Equal
Equal
And
-Not
Cast
-CastLike
Unsqueeze
-Mul
+Constant
+Where
Constant
Max
Constant
@@ -162,12 +159,13 @@
Add
Constant
OneHot
-CastLike
Transpose
MatMul
+CastLike
+Cast
Constant
-CastLike
Mul
+Cast
CastLike
RMSNormalization
TransposeAdded nodes:
Removed nodes:
Modified attributes:
Connectivity changes:
Initializer changes:
gemma4_text / model — 11 change(s)Op summary: 128 → 132 nodes --- base
+++ head
@@ -5,6 +5,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
Shape
@@ -19,6 +21,8 @@
Where
Cast
Unsqueeze
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
qwen2_moe / model — 15 change(s)Op summary: 214 → 216 nodes --- base
+++ head
@@ -1,6 +1,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Connectivity changes:
qwen3_5_moe (hybrid-text-generation) / model — 34 change(s)Op summary: 269 → 271 nodes --- base
+++ head
@@ -1,6 +1,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
qwen3_5_text (hybrid-text-generation) / model — 20 change(s)Op summary: 131 → 133 nodes --- base
+++ head
@@ -1,6 +1,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
qwen3_5_vl (hybrid-qwen-vl) / decoder — 21 change(s)Op summary: 157 → 159 nodes --- base
+++ head
@@ -26,6 +26,8 @@
Where
Where
Where
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
qwen3_moe / model — 19 change(s)Op summary: 194 → 196 nodes --- base
+++ head
@@ -1,6 +1,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Connectivity changes:
qwen3_next (hybrid-text-generation) / model — 56 change(s)Op summary: 565 → 567 nodes --- base
+++ head
@@ -1,6 +1,8 @@
Gather
Gather
Gather
+Shape
+ConstantOfShape
CumSum
Unsqueeze
ShapeAdded nodes:
Modified attributes:
Connectivity changes:
Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Optimized mask handling and packed-adapter restoration have correctness gaps, and artifact hardlinks can bypass checkpoint overwrite protection.
3 open issues
What changed in this PR
This PR aligns Gemma4 multimodal exports with HuggingFace precision and masking behavior, gates experimental CUDA operators, and adds projection packing with adapter-restoration safeguards.
Changes:
- Corrects image visibility, FP32 pixel normalization, pooling, standardization, and router scaling.
- Adds explicit custom GEGLU runtime opt-in and float/INT4 projection packing.
- Expands regression coverage and documents input compatibility and remaining qualification limits.
| File | Description |
|---|---|
| tests/static_cache_bias_parity_test.py | Updates reference mask ordering. |
| tests/integration/gemma4_test.py | Expands comparisons against HuggingFace masks. |
| tests/build_graph/core_test.py | Checks FP32 vision input contracts. |
| src/mobius/tasks/_gemma4.py | Exposes FLOAT vision pixels. |
| src/mobius/rewrite_rules/_qkv_projection_packing.py | Adds float and INT4 projection packing. |
| src/mobius/rewrite_rules/_qkv_projection_packing_test.py | Tests packing, restoration, and save safeguards. |
| src/mobius/rewrite_rules/__init__.py | Exports projection-packing passes. |
| src/mobius/models/gemma4.py | Updates precision, masks, and native expert paths. |
| src/mobius/models/gemma4_test.py | Adds precision and runtime-gating regressions. |
| src/mobius/functions/matmul_nbits.py | Defines ownership metadata constants and clarifies factory ownership. |
| src/mobius/functions/matmul_nbits_test.py | Tests caller-owned and registered functions. |
| src/mobius/components/_moe.py | Makes native expert activation overridable. |
| src/mobius/components/_common.py | Uses physical mask distances and corrected overlay ordering. |
| src/mobius/components/_common_test.py | Covers padding holes and window boundaries. |
| src/mobius/adapter_io.py | Restores packed projections with lazy-source ownership. |
| src/mobius/_optimizations.py | Integrates packing passes and updates folding order. |
| src/mobius/_model_package.py | Registers derived lazy sources before save validation. |
| src/mobius/_flags.py | Adds default-off custom GEGLU opt-in. |
| src/mobius/_flags_test.py | Tests flag defaults and scoped overrides. |
| src/mobius/_execution_providers.py | Declares CUDA INT4 packing eligibility. |
| src/mobius/_constants.py | Adds packing, runtime, and ownership metadata keys. |
| src/mobius/_configs/_vision_defaults.py | Reads vision standardization settings. |
| src/mobius/_configs/_sub_configs.py | Adds vision standardization configuration. |
| src/mobius/_configs/_quantization.py | Adds expert quantization opt-in and legacy alias parsing. |
| examples/gemma4_multimodal.py | Supplies FP32 processor pixels. |
| docs/model-catalog.md | Links Gemma4 runtime qualification guidance. |
| docs/gemma4-runtime.md | Documents contracts, evidence, and unresolved validation. |
Review effort: Balanced
titaiwangms
force-pushed
the
titaiwang/gemma4-vlm-runtime
branch
from
October 24, 2026 00:22
c22ec12 to
0188adc
Compare
titaiwangms
changed the title
titaiwangms
changed the base branch from
main
to
titaiwang/quantized-qkv-packing
October 24, 2026 00:22
3 of 9 tasks
Correct masks, bf16 handling, shared KV behavior, default-off custom GEGLU, causal-global behavior, and preserve arbitrary vision masks. Tests cover 196 cases plus 11 mask-parity cases. This does not claim a strict full numerical pass. Co-authored-by: Copilot <Copilot-Session9a2a5b82-7cd3-4a76-bb94-0998527ea723> Signed-off-by: titaiwang <titaiwang@microsoft.com>
The dense portable fallback must select out unchosen expert outputs before applying routing weights, because zero times Inf or NaN can contaminate the result. Nonfinite outputs from selected experts still propagate as in sparse dispatch. This change makes no tolerance changes and does not qualify CUDA execution. Signed-off-by: titaiwang <titaiwang@microsoft.com> Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723
The test helper accepts the base forward signature, and the restoration assertion is unconditional. 147 tests pass; scoped lint passed. CodeQL is not available locally and is not claimed passed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang <titaiwang@microsoft.com>
titaiwangms
force-pushed
the
titaiwang/gemma4-vlm-runtime
branch
from
October 24, 2026 02:29
0188adc to
fd4513b
Compare
Declare the arbitrary-mask contract conditionally so vision-overlay configurations do not lose information through post-hoc GQA fusion. Add regression coverage for split and text tasks with cached mask holes [1,0,1,1] and [1,1,1,0]; stripped text configurations retain GQA. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang <titaiwang@microsoft.com>
Contributor
There was a problem hiding this comment.
🔵 Needs a closer look
Strict numerical gates remain failed and unwaived, while full-checkpoint validation of this revision and target CUDA/custom-runtime qualification remain incomplete.
No open issues 🎉
1 resolved since last review
Review effort: Balanced
titaiwangms
added a commit
that referenced
this pull request
Oct 26, 2026
) ## Summary Fix VibeVoice CI dependency installation by resolving the pinned Transformers source together with the editable Mobius project, then installing the pinned VibeVoice source with `--no-deps`. ## Failure and fix The failing install sequence first resolved Transformers 5.20.2 with Hub 2.3.0 and `httpx2`, then force-installed Transformers 5.16.0.dev0 from the combined VibeVoice requirements using `--no-deps`. That left Hub 2.3.0 installed despite the pinned Transformers source requiring `huggingface-hub>=1.5.0,<2.0`. Importing `AutoConfig` then failed when Transformers imported the absent `httpx` module. The split install resolves the pinned Transformers source and its dependencies together with the editable project, installs only VibeVoice with `--no-deps`, and immediately exercises `from transformers import AutoConfig`. It neither blindly adds `httpx` nor changes either source revision. Primary evidence: [pinned Transformers dependencies](https://github.com/huggingface/transformers/blob/640a08a597034221ca1c4fc0c129cf0118179225/setup.py), [upstream VibeVoice metadata](https://github.com/microsoft/VibeVoice/blob/79e516a3e20b599f137c9da03410a2a0b473b63b/pyproject.toml), and the [failed synthetic parity job](https://github.com/onnxruntime/mobius/actions/runs/40044682674/job/120723608401). Both 40-character source pins remain unchanged: - Transformers: `640a08a597034221ca1c4fc0c129cf0118179225` - VibeVoice: `79e516a3e20b599f137c9da03410a2a0b473b63b` All five install steps retain actual dependency resolution and AutoConfig import smoke, and the YAML-structure Python tests were removed following review. ## Validation - `ruff format`, `ruff check`, `ruff format --check`, and `git diff --check`: passed. - Validation log: `ci-vibevoice-deps-validation-v2.log`. A local fresh-venv network resolver run was **not** performed; host CI validation remains pending. No installs, full-model, or GPU runs were performed. VibeVoice's upstream metadata still declares Transformers `<5`. This preserves the existing reference integration's deliberate `--no-deps` exception for VibeVoice; it does not resolve that pre-existing metadata mismatch or apply the exception to Transformers itself. This PR is standalone from main `d69efe03a9f4c1ada744bc31dd01077a710b3c39` and independent of #754 and #787. --------- Signed-off-by: titaiwang <titaiwang@microsoft.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723
justinchuby
reviewed
Oct 27, 2026
justinchuby
approved these changes
Oct 27, 2026
titaiwangms
merged commit Oct 27, 2026
6ed5051
into
titaiwang/quantized-qkv-packing
20 of 22 checks passed
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.
Organization
Depends on #754. This PR contains the shared Gemma4 text-only and multimodal
modeling changes. Generic projection packing and weight-layout folding live in
#754 instead of being duplicated here.
Text-only construction already exists on main; this PR corrects and documents
the model semantics used by both text and multimodal exports.
Changes
pooling, and expert contributions.
(causal OR same-image) AND left-window AND padding, using physical KV-slotdistances even when a padding mask contains holes.
optimization does not collapse it to prefix lengths through GQA fusion.
Text-only CUDA exports retain GQA eligibility.
experimental custom-GEGLU opt-in and runtime-requirement metadata.
Review follow-up
Includes cached internal-hole mask regressions using the actual optimized CUDA
export's bias, a text-only GQA eligibility regression, independent HF mask
comparisons, and exception-restoration coverage for the custom-GEGLU flag.
The adapter alias and repeated-source packing findings from this PR are handled
in its #754 dependency.
Qualification limits
Export support is not released-runtime compatibility or numerical qualification.
Custom GEGLU and quantized Gemma4 MoE require the documented experimental
runtime; neither is enabled by ordinary CUDA selection. Previously failed
strict HF numerical gates remain failed and unwaived.
Historical performance and accuracy artifacts are not represented as validation
of this reorganized head. The ongoing paired MMLU-Pro500 evaluation uses the
frozen optimized text artifact and the pinned decoder-only llama.cpp Q4_K_M
artifact, with identical prompts and scoring.
Consolidated-revision validation
selected/unselected overflow with zero/unit expert scale.
quantized-storage regressions and nested-scope persistence coverage.
Exact-base selector runs reproduce 28 failures; the additional GGUF
file-identity failure passes separately on both base and head. This is
not a green full-suite claim.
HF
generateself-type diagnostics remain on both base and head.Two bounded five-role review rounds completed. The review's portable-MoE
question was verified with the actual dispatch graph: an overflowing
unselected expert produced 64 NaNs before the correction, versus a
selected-only zero reference. An independent Boolean selection mask now
selects out unchosen outputs before weighting. The same CPU probe produces
64 finite exact zeros afterward; selected nonfinite outputs still propagate,
including with zero learned expert scale.
The dependency's F1/F2/F3 scope, documentation, and storage findings are fixed
in #754. No Critical/Major reviewer finding was overruled. Remaining exclusions:
full-checkpoint re-export and strict HF logits/KV parity, CUDA qualification of
the fallback safeguard, custom-runtime availability, and performance. The
failed numerical gates remain unwaived.
Complete original captured PR body — historical prior implementation (not validation of the reorganized head)
The historical results in the captured body below describe the prior implementation and are NOT validation of the reorganized head. Older statements about unchanged float behavior are historical and superseded by the new summary above. All original numerical limitations and earlier MMLU/performance evidence are preserved verbatim.
Summary
This draft isolates the Gemma4 runtime work from the unrelated ASR, Qwen4, text-evaluation, and experimental fixed-cache changes.
Compatibility and qualification
The vision input dtype change requires rebuilding exports and passing the original FP32 processor pixels, rather than precasting them to BF16.
Custom native GEGLU is disabled by default.
MOBIUS_GEMMA4_CUSTOM_GEGLU=1explicitly opts into the experimental custom CUDA ABI; it is not implied by selecting CUDA. Ordinary float CUDA uses the portable path, while quantized GEGLU experts without the required opt-in fail clearly. Native operators recordmobius.gemma4.required_runtime_feature="custom-cuda-geglu-v1"; the flag is an assertion about the target runtime, not automatic ABI detection.The native paths require the corresponding ORT runtime/operator support. An exported GQA node does not establish Flash/XQA dispatch, shared KV support, or CUDA Graph capture.
This is not numerical-parity or production approval. Strict full-vocabulary logits/KV checks still fail at the original
atol=0.01, rtol=0.01; the previous packing parity gate is not waived. Video and the experimental TensorScatter/SparsePaged conversions are outside this PR.Recorded pre-port experiment
These results belong to immutable corrected exports built before porting the changes onto the newer main branch, and before the subsequent pool/window, adapter-resave, and runtime-opt-in review fixes. They are not a real-checkpoint re-export/validation of this PR branch. Export graph, source, runtime, and result hashes are retained in the experiment artifacts.
Fixed AI2D-200 quantization-selection subset; not an independent full-AI2D evaluation:
All four runs used matching original FP32 processor tensors and had zero inference failures. Three real-image prompts, each with prefill and two teacher-forced cached decodes, matched all nine first-token argmaxes for each native precision, but failed the strict full-logit/KV gates.
Performance on the same frozen exports, batch 1, three single-image prompts of 333/353/356 input tokens:
Each case used two discarded warmups and five measured trials, generating one first token and 64 greedy decode tokens. TTFT and throughput are medians of the three case medians. Full-vocabulary diagnostic readbacks were excluded from measured trials; EOS suppression was performance-only. HF used eager experts and SDPA without compilation; its Olive INT4 dequantizes on demand, unlike native QMoE. These are different implementations, not an isolated packing speedup.
The selected A100 80GB had no competing compute process, but other GPUs shared the host with text evaluation. Peak memory is a 20ms NVML sample of the whole selected device, not exact model-only allocation. These are short-prompt results, not 2K/8K/16K/32K benchmarks.
The experimental ORT 1.32 runtime was built from revision
942356fce6bdc78c016e4b902e49492fdf695453plus local modifications, not an unmodified upstream release. Its shared-library SHA256 is3f6b7ff482f8fbd4d3c639e807ed2ec35fa240c88c205dbee7822d7ec497e988.Both native decoder graphs retained 30 packed-QKV projection records. Diagnostic profiles executed no Constant nodes, but still averaged 126 CUDA Cast nodes per decoder invocation. Constant folding does not justify deleting dynamic precision boundaries; these counts are not per-op GPU timings.
Isolated-branch validation
Against base
d69efe03a9f4c1ada744bc31dd01077a710b3c39, offline and CPU-only:The four excluded GQA cases were separately reproduced on the exact clean base; unavailable cached remote configurations were not downloaded, and tests were not weakened. The strict typecheck produced 2,798 diagnostics on the head and 2,807 on base with no missing-import diagnostics; this is a baseline-relative result, not clean type approval.
Sensitive regressions include BF16 3x3 pooling against the actual HF implementation, padded NaN/+Inf/-Inf handling, physical-window boundaries with internal padding and cached decode, shared-weight packing refusal followed by public attachment of both independent K adapters, source-preserving same-path/hardlink/symlink save rejection, and safe-destination package round-trips.
Expanded validation caught and fixed a caller-owned
MatMulNBitsfunction regression: raw factory results must not claim registry ownership. The existing preservation test now passes unchanged on both base and head, and new tests retain automatic-fallback packing while preserving caller-owned functions/calls.Review disposition
Two targeted five-reviewer rounds completed. They closed F1-F6 without a remaining grounded Major/Critical in their bounded follow-up. A subsequent adversarial probe promoted the final round's refolding ownership question to F7. The user explicitly approved fixing F7 and running precise validation rather than a third full-team round.
models/gemma4.py: premature low-precision poolingcomponents/_common.py: padding compressed local-window distances_qkv_projection_packing.py/adapter_io.py: shared sources broke independent adapter restorationMatMulNBitsfactory incorrectly claimed automatic function ownershipFor F7, the pre-fix generated fixture's external data file changed from 16,384 bytes to zero before a lazy read failed. With the fix, the identical instrument records zero external reads, zero source write-open requests, and unchanged nonzero sizes/SHA256 digests. The regression also saves elsewhere and checks exact reloaded payloads. Original model artifacts were not touched. F7's final small fix was not independently re-reviewed by the full team; it has the explicitly approved focused execution evidence.
Corrections/minority record: Code withdrew its original claim that directly attaching a shared Q remained supported after folding: the exact base requires one initializer consumer. The real two-independent-K-target regression was retained and fixed. Integration withdrew the proposed fused-RoPE restriction for projection-only packing because
MatMul/MatMulNBits -> Splitdoes not match legacy packed-GQA unpacking. No unresolved reviewer objection to those corrections remains.Attributed open questions: Deep retains current-revision full-checkpoint logits/KV/packing qualification. Integration retains whether released ORT GenAI versions supply the split decoder's required
input_idsalongsideinputs_embeds; the experimental measured runner does not qualify every released runner. Readability's provenance limitation is explicit: the modified runtime patchset and complete experiment records are not included in this repository, so the binary identity is not a reproducible build recipe.Residual risk / exclusions: CPU-only branch validation, two skipped tests, four exact-base offline GQA failures, and non-clean baseline typing remain explicit. No new-branch full-checkpoint export, released-runtime qualification, long-context benchmark, CUDA Graph/shared-KV clearance, arbitrary model-replacement metadata guarantee, or concurrent filesystem-race test is claimed. The unselected-expert nonfinite hypothesis was refuted only for its exact FP16 finite-router probe, not universally. ASR, Qwen4, text benchmark implementation, and experimental static/paged-cache work are excluded.
Critical review threat-model appendix: Inputs include checkpoint/quantization plans, serialized packing metadata and external files, PEFT factors, feature flags, and runtime pixels/positions/IDs/masks. Boundaries are metadata-authorized restoration, lazy backing-file lifetime, and explicit authorization of experimental native operators. Relevant risks are tampering/data integrity and denial of service; no new authentication or elevation path was established. Known source aliases fail before writing; this is not exhaustive transformation or concurrent-filesystem protection.
Remaining validation
Current follow-up validation and CI root cause
Published follow-up
fd4513b5cc2e4692a2cde51fd8fdd57c4bb0878d. The test override accepts base forward parameters, including keyword names and optional causal/static-cache arguments, while retaining the isolated MoE probe; 147 model/flag tests and 225 combined tests passed, with scoped lint passed. There are no production forward changes. The restoration assertion is now infinally, sooriginalis always consumed on either control-flow exit; nested flag contexts/assertions are retained. CodeQL is unavailable locally; no new CodeQL clearance is claimed.The initial CI environment used TensorFlow 5.20.2 / Hugging Face Hub 2.3 / httpx 2, then
--no-depsforced TensorFlow 5.16.dev0 whose metadata requireshuggingface-hub>=1.5,<2; old Hub 1.x supplies the older httpx. This is a dependency-closure violation, not merely missing httpx, and is not changed by this PR. No CI code fix is included and CI green is not claimed. Evidence: pinned TensorFlow setup, Hub v1.5.0 setup.py, and Hub v2.3.0 setup.py. Resolve the pinned TensorFlow dependency normally; use--no-depsonly for VibeVoice plus pip-check/import smoke, with no automatic new PR scope. Strict numerical gates remain failed/unwaived. The full five-role, two-round review is already complete; this follow-up is a direct bounded review only.