From da598502d6ebaef963a85ffb35fac24525806704 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Fri, 23 Oct 2026 23:39:44 +0000 Subject: [PATCH 1/4] Gemma4: align shared text and multimodal runtime semantics 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 Signed-off-by: titaiwang --- docs/gemma4-runtime.md | 177 +++++++ docs/model-catalog.md | 4 + examples/gemma4_multimodal.py | 5 +- src/mobius/_configs/_quantization.py | 3 + src/mobius/_configs/_sub_configs.py | 2 + src/mobius/_configs/_vision_defaults.py | 1 + src/mobius/_constants.py | 2 + src/mobius/_flags.py | 16 + src/mobius/_flags_test.py | 23 + src/mobius/components/_common.py | 61 +-- src/mobius/components/_common_test.py | 78 ++- src/mobius/components/_moe.py | 5 +- src/mobius/models/gemma4.py | 431 +++++++++++----- src/mobius/models/gemma4_test.py | 649 +++++++++++++++++++++++- src/mobius/tasks/_gemma4.py | 12 +- tests/build_graph/core_test.py | 8 +- tests/integration/gemma4_test.py | 120 +++-- tests/static_cache_bias_parity_test.py | 8 +- 18 files changed, 1374 insertions(+), 231 deletions(-) create mode 100644 docs/gemma4-runtime.md diff --git a/docs/gemma4-runtime.md b/docs/gemma4-runtime.md new file mode 100644 index 000000000..4465b27a6 --- /dev/null +++ b/docs/gemma4-runtime.md @@ -0,0 +1,177 @@ +# Gemma4 text and multimodal runtime qualification + +This runtime work is draft, not production or equivalence approved. It does not +qualify optimized text-only performance, ASR, or experimental fixed/paged-cache +execution. + +## HuggingFace alignment and input contract + +- Global decoder layers use causal attention only; image-block IDs must not + unmask future keys in global layers. +- Local decoder layers use + `valid_key AND (causal OR same_image) AND (key > query - window)`. + Same-image future keys can be visible, but an image block cannot bypass the + strict left window boundary. Distances use physical query/KV slots, including + past offsets; padding holes do not compress the window. Only current-chunk + image-block IDs are available; + these IDs are not persisted alongside the KV cache. +- The separate vision encoder's `pixel_values` input deliberately uses ONNX + `FLOAT`, including with BF16/FP16 projection weights. Supply FP32 processor + patches in `[0, 1]`, not BF16 patches or pre-normalized `[-1, 1]` values. + Normalization computes `2 * (pixels - 0.5)` in FP32 before casting to the + projection's working dtype. Existing consumers must update their input dtype; + `examples/gemma4_multimodal.py` demonstrates the new ABI. This does not change + the encoder-free `gemma4_unified` vision ABI. +- Checkpoint `router.scale` is preserved. The router first multiplies the + normalized working-dtype activation by the learned scale, then applies the + separate `hidden_size ** -0.5` scalar with FP32 opmath and rounds back to the + working dtype. Folding that scalar into a low-precision learned scale changes + rounding and is not equivalent. +- Checkpoint-enabled vision standardization retains `std_bias`/`std_scale` + directly on the vision submodel. Pooling coefficients and feature accumulation + are FP32, including the `1/9` coefficients for a 3x3 kernel. The pooled mean + then rounds to working dtype before FP32 scaling/standardization and the + final cast back for projection. + +## Native runtime requirements and limits + +Ordinary `execution_provider="cuda"` does **not** assert custom GEGLU support. +The default-off `gemma4_custom_geglu` flag explicitly opts into the experimental +custom CUDA ABI: set `MOBIUS_GEMMA4_CUSTOM_GEGLU=1` before importing Mobius, or +use `override_flags(gemma4_custom_geglu=True)` around construction, export and +weight preprocessing. Environment values are captured when the module-level +flags instance is created on first import; changing the environment afterward +does not update that instance. Without opt-in, float tanh-GELU experts use the portable +exact-activation fallback; quantized GEGLU experts fail at construction with an +explicit requirement message. Opt-in does not enable native GEGLU on non-CUDA +providers. Tanh-GELU must not be replaced by a sigmoid-based SwiGLU approximation. + +Each emitted native `MoE`/`QMoE` operator persists +`mobius.gemma4.required_runtime_feature="custom-cuda-geglu-v1"` in its ONNX +metadata. This declares the experimental CUDA `activation_type="geglu"`, +`swiglu_fusion=1` requirement; it is not a claim of released ORT support or a +claim that the exporting machine used a particular runtime binary. + +The frozen pre-port experiments used an **experimental modified ORT 1.32 build**: +source `942356fce6bdc78c016e4b902e49492fdf695453` with seven dirty paths; binary +SHA-256 `3f6b7ff482f8fbd4d3c639e807ed2ec35fa240c88c205dbee7822d7ec497e988`. +This is not an unmodified upstream release. Its local patchset and full +experiment records are not included in this repository; these identity fields +do not provide a reproducible runtime build. Released ORT support is unverified. + +Dynamic vision-overlay decoders declare the arbitrary-attention-mask contract. +Both local and global layers retain ONNX `Attention`: reducing a mask to a valid +token count would discard internal padding holes even in causal global layers. +This preserves the full additive bias through CUDA optimization. The default +VLM qualification path uses dynamic KV cache, graph capture off, and past/present +buffer sharing off. + +Text-only export uses the existing `text_only=True` builder option (CLI +`--features text-only`) to remove the vision overlay. Its causal/windowed +attention can use the existing GQA path when supported by the execution provider. +GQA sequence lengths describe prefix validity, not arbitrary interior mask holes; +use a standard-ONNX attention path when those holes must be preserved. +That eligibility is separate from projection packing, custom GEGLU opt-in, and +runtime KV-sharing/graph-capture settings. No experimental static-cache or +paged-attention changes are part of this modeling work. + +## Recorded pre-port experiments, not validation of this branch + +The corrected real-weight native BF16 and INT4 VLM runs each scored +**166/200**; HF BF16 scored **167/200** and HF INT4 **162/200**. These results use +a fixed quantization-selection AI2D subset, not an independent full-AI2D +evaluation, and must not be described as such. + +Those runs used frozen corrected exports built from the original dirty source +worktree, with source hashes recorded in the originating experiment records. +They were **not** exports of this isolated origin/main-based branch. Newer main +has different leaf-layout/configuration features; these recorded scores (and +any separately recorded pre-port performance results) do not validate this +branch. The frozen exports also predate this branch's corrected FP32 pooling +and physical-slot window fixes. Publication requires the original source hashes, +checkpoint/quantization configuration, export identity, and runtime version to +remain attached to the experiment provenance. + +Strict full-vocabulary logit parity **failed 9/9**, and KV parity **failed +540/540**, at `atol=0.01, rtol=0.01`. Task-level scores do not establish numerical +equivalence or production readiness. The isolated +patch's tiny-graph and weight-layout tests are separate evidence. Full-checkpoint +re-export and revalidation on the ported branch remain **unverified**. No model +rebuild, download, GPU run, or performance measurement is implied by this note. +Lead review is required before publication. + +## Prior combined-branch evidence + +Offline, CPU-only validation against base +`d69efe03a9f4c1ada744bc31dd01077a710b3c39`: + +- Related runtime, package, expert-layout, adapter, and packing tests: + **656 passed, 2 skipped**. These include `models/gemma4_test.py`, + `components/_common_test.py`, `_model_package_test.py`, + `functions/matmul_nbits_test.py`, and + `rewrite_rules/_qkv_projection_packing_test.py` under `src/mobius/`. +- GQA rewrite tests: **32 passed**. Four remote-config cases were excluded from + that offline invocation after reproducing their failures on the exact clean + base; their tests were not changed. +- `tests/build_graph/core_test.py -k gemma4`: **10 passed**. + `tests/integration/gemma4_test.py::test_gemma4_bidirectional_mask_parity`: + **11 passed** against the actual pinned HF mask entry points. + +Changed-file lint, formatting, compilation, and whitespace checks passed. +Strict mypy with resolved editable ONNX Script imports has zero new normalized +diagnostic instances relative to base, but neither revision typechecks cleanly. +These checks do not replace the unresolved full-checkpoint numerical gates. +The modeling/pass split and subsequent review fixes require their own validation; +the historical counts above are not results for the reorganized revision. + +## Packed-checkpoint contract + +Native Gemma4 experts require an explicit decoder quantization config: +`quant_method="olive"`, integer-affine `bits=4`, a power-of-two +`group_size >= 16`, `sym`, and canonical `quantize_moe=true`. +Metadata also accepts the legacy `moe` alias; `quantize_moe` takes precedence +when both are supplied. Expert quantization opt-in and the custom-runtime flag +are separate requirements. +Compute dtype must be FP16/BF16 and expert activation must be tanh-GELU. +Both expert input dimensions must be divisible by the group size. Preserve the +authoritative HF source-module names and any component/per-projection override +or exclusion plan; unsupported expert layouts fail closed, while ordinary +projection plans remain handled by the shared component loader. + +For each layer with `E` experts, hidden size `H`, intermediate size `I`, and +group size `G`, Olive sidecars under `...layers.L.experts` are: + +| Stem | `_qweight` (UINT8) | `_scales` (floating point) | `_qzeros` (UINT8, asymmetric only) | +|---|---|---|---| +| `gate_up_proj` | `[E, 2*I, H/2]` | `[E, 2*I, H/G]` | `[E, 2*I, ceil((H/G)/2)]` | +| `down_proj` | `[E, H, I/2]` | `[E, H, I/G]` | `[E, H, ceil((I/G)/2)]` | + +Weights use contiguous gate-then-up rows before graph-level interleaving. +The adapter preserves packed codes and folds `router.per_expert_scale` into +down-projection block scales; low-precision folding changes rounding order +and is not claimed bitwise exact. The learned `router.scale` remains unchanged. +Incomplete, stale, or shape-incompatible expert sidecars are rejected. + +Float projection packing retains version-1 +`mobius.qkv_projection_sources` metadata on the `Split`: original initializer +names and output widths. Keep it through serialization so PEFT attachment can +restore independent projections after weight folding. INT4 projection packing +requires homogeneous bits/block size/accuracy level, scale dtype, zero-point +presence/layout, and a known shape-generic fallback body when one is registered. +Float groups whose source initializer or logical weight view has outside +consumers are conservatively left unpacked, preserving existing tied-weight +adapter attachment rather than inventing multiple-group restoration semantics. + +Restored projections remain lazy slices, not eager copies of a large packed +checkpoint. Their backing-file ownership is registered with the package save +alias guard. Saving over a backing file (same path, hard link or symlink) is +rejected before any output is written, with the source unchanged. Use a distinct, +non-aliasing output location; ordinary packages without such derived lazy +dependencies retain their existing save behavior. + +Backing-file ownership is also retained on the in-memory model, so initializer +folding cannot hide it by replacing the restored tensor wrappers. A regression +refolds the lazy projections before saving, checks refusal before any external +read or source mutation, then verifies an exact round-trip at a safe destination. +This metadata is transient, not serialized provenance or a guarantee for arbitrary +model-replacement transformations. diff --git a/docs/model-catalog.md b/docs/model-catalog.md index c42960214..7659e1d31 100644 --- a/docs/model-catalog.md +++ b/docs/model-catalog.md @@ -112,6 +112,10 @@ Models combining Mamba/SSM layers with transformer attention layers. Models that process both images and text. +For Gemma4 26B-A4B, see the [draft runtime qualification and FP32 pixel input +contract](gemma4-runtime.md), including the failed strict parity gate and limits +of the prior evaluation evidence. + | Model Type | Module Class | Task | Example HuggingFace Model | |---|---|---|---| | `llava` | `LLaVAModel` | `vision-language` | `llava-hf/llava-1.5-7b-hf` | diff --git a/examples/gemma4_multimodal.py b/examples/gemma4_multimodal.py index e55104121..df3544be5 100644 --- a/examples/gemma4_multimodal.py +++ b/examples/gemma4_multimodal.py @@ -132,7 +132,7 @@ def prepare_vision_feeds( Args: processor: ``AutoProcessor`` loaded for the Gemma 4 model. image_path: Path to a local image file (JPEG, PNG, etc.). - np_dtype: Numpy dtype for pixel values (default float32). + np_dtype: Retained for caller compatibility; pixels always use float32. Returns: ``{"pixel_values": [B, N, 3*P^2], "pixel_position_ids": int64[B, N, 2]}`` @@ -143,7 +143,8 @@ def prepare_vision_feeds( # Gemma 4 processor requires a text argument alongside images. # Use a placeholder image token so the processor computes correct dims. processed = processor(images=img, text="", return_tensors="np") - pixel_values = processed["pixel_values"].astype(np_dtype) + # The vision ABI is FLOAT even when decoder/projection weights use BF16. + pixel_values = processed["pixel_values"].astype(np.float32) # The HF processor returns "image_position_ids"; our ONNX vision model input is "pixel_position_ids" pixel_position_ids = processed["image_position_ids"].astype(np.int64) return {"pixel_values": pixel_values, "pixel_position_ids": pixel_position_ids} diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index defb49ffc..8c124bb33 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -108,6 +108,8 @@ class QuantizationConfig: # positional construction, so inserting a field earlier would silently # change the meaning of existing callers' arguments. weight_format: QuantizedWeightFormat = QuantizedWeightFormat.INTEGER_AFFINE + # Olive explicitly opts fused expert tensors into native QMoE packing. + quantize_moe: bool = False def __post_init__(self) -> None: """Normalize serialized enum values without inferring storage semantics.""" @@ -249,6 +251,7 @@ def from_value( quantize_embeddings=bool(qc.get("embeds")), quantize_lm_head=bool(qc.get("lm_head")), quantize_vision=bool(qc.get("quantize_vision")), + quantize_moe=bool(qc.get("quantize_moe", qc.get("moe", False))), tie_word_embeddings=bool(qc.get("tie_word_embeddings")), modules_to_not_convert=exclusions, overrides=overrides, diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index a4bca1539..0b5381ef1 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/src/mobius/_configs/_sub_configs.py @@ -108,6 +108,8 @@ class VisionConfig: window_kernel_size: tuple[int, int] = (2, 2) merge_kernel_size: tuple[int, int] = (2, 2) merger_times: int = 1 + # Gemma4 post-pooling checkpoint affine transform, computed in FP32. + standardize: bool = False @dataclasses.dataclass diff --git a/src/mobius/_configs/_vision_defaults.py b/src/mobius/_configs/_vision_defaults.py index eabb9af36..63c4b8adb 100644 --- a/src/mobius/_configs/_vision_defaults.py +++ b/src/mobius/_configs/_vision_defaults.py @@ -97,6 +97,7 @@ def apply_vision_defaults(config, parent_config, model_type: str, fields: dict) hidden_act=getattr(vc, "hidden_act", None), use_clipped_linears=getattr(vc, "use_clipped_linears", False), position_embedding_size=getattr(vc, "position_embedding_size", None), + standardize=getattr(vc, "standardize", False), ) # This runs before per-model hooks, so plain assignment is safe — any # per-model hook may overwrite later via fields.update(...). diff --git a/src/mobius/_constants.py b/src/mobius/_constants.py index 45ff4e9ac..8653d4f04 100644 --- a/src/mobius/_constants.py +++ b/src/mobius/_constants.py @@ -15,6 +15,8 @@ # Transient tensor.meta ownership, never serialized as ONNX metadata. LAZY_SOURCE_FILES_META_KEY = "mobius.lazy_source_files" +GEMMA4_CUSTOM_GEGLU_METADATA_KEY = "mobius.gemma4.required_runtime_feature" + # --------------------------------------------------------------------------- # Static KV cache ABI # --------------------------------------------------------------------------- diff --git a/src/mobius/_flags.py b/src/mobius/_flags.py index e17fe621c..9dfa16659 100644 --- a/src/mobius/_flags.py +++ b/src/mobius/_flags.py @@ -97,6 +97,11 @@ class _Flags: block overlay + padding) on the external-KV static-cache ``Attention`` path (``is_causal=0``) for float-bias decoders, instead of the maskless ``is_causal=1`` default. + * - ``gemma4_custom_geglu`` + - ``MOBIUS_GEMMA4_CUSTOM_GEGLU`` + - ``False`` + - Opt in to the experimental custom CUDA GEGLU MoE/QMoE ABI. + Ordinary CUDA selection does not imply this runtime support. """ suppress_dedup_warning: bool = dataclasses.field( @@ -178,6 +183,17 @@ class _Flags: unchanged, so no shipped model's graph changes unless this flag is set. """ + gemma4_custom_geglu: bool = dataclasses.field( + default_factory=lambda: _env_bool("MOBIUS_GEMMA4_CUSTOM_GEGLU", False) + ) + """Opt in to Gemma4's experimental custom CUDA tanh-GEGLU MoE/QMoE ABI. + + Requires a separately qualified custom runtime, not an unmodified released + ORT wheel. Default float exports use the portable exact-activation fallback; + quantized experts fail closed without this opt-in. The flag does not enable + native dispatch for non-CUDA providers. + """ + allow_dense_moe_experts: bool = dataclasses.field( default_factory=lambda: _env_bool("MOBIUS_ALLOW_DENSE_MOE_EXPERTS", False) ) diff --git a/src/mobius/_flags_test.py b/src/mobius/_flags_test.py index 412aa9953..38ae2fc47 100644 --- a/src/mobius/_flags_test.py +++ b/src/mobius/_flags_test.py @@ -18,6 +18,10 @@ def test_suppress_dedup_warning_default_on(self, monkeypatch): f = _flags._Flags() assert f.suppress_dedup_warning is True + def test_gemma4_custom_geglu_default_off(self, monkeypatch): + monkeypatch.delenv("MOBIUS_GEMMA4_CUSTOM_GEGLU", raising=False) + assert _flags._Flags().gemma4_custom_geglu is False + class TestEnvVarOverride: """Flags are read from environment variables at construction time.""" @@ -39,6 +43,13 @@ def test_unknown_value_falls_back_to_default(self, monkeypatch): f = _flags._Flags() assert f.suppress_dedup_warning is True # default + @pytest.mark.parametrize( + ("value", "expected"), [("1", True), ("true", True), ("0", False), ("maybe", False)] + ) + def test_gemma4_custom_geglu_opt_in(self, monkeypatch, value, expected): + monkeypatch.setenv("MOBIUS_GEMMA4_CUSTOM_GEGLU", value) + assert _flags._Flags().gemma4_custom_geglu is expected + class TestProgrammaticOverride: """Flags can be assigned directly on the singleton.""" @@ -80,6 +91,18 @@ def test_nested_overrides(self): assert _flags.flags.suppress_dedup_warning is not original assert _flags.flags.suppress_dedup_warning is original + def test_custom_geglu_override_is_scoped_even_on_exception(self): + original = _flags.flags.gemma4_custom_geglu + with _flags.override_flags(gemma4_custom_geglu=False): + assert not _flags.flags.gemma4_custom_geglu + with pytest.raises(RuntimeError, match="test restoration"): + assert not _flags.flags.gemma4_custom_geglu + with _flags.override_flags(gemma4_custom_geglu=True): + assert _flags.flags.gemma4_custom_geglu + raise RuntimeError("test restoration") + assert not _flags.flags.gemma4_custom_geglu + assert _flags.flags.gemma4_custom_geglu is original + def test_unknown_flag_raises_value_error(self): with ( pytest.raises(ValueError, match="Unknown flag"), diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py index 3254a25e8..6019cde81 100644 --- a/src/mobius/components/_common.py +++ b/src/mobius/components/_common.py @@ -179,9 +179,10 @@ def create_attention_bias( (batch_size, query_length) giving a contiguous vision-block id per current-sequence position (``>= 0`` for vision tokens in the same block, ``-1`` for text). When provided, a bidirectional - "blockwise overlay" is OR-ed onto the causal (and sliding) mask: + "blockwise overlay" is OR-ed onto the causal mask: two positions in the same block (same id ``>= 0``) may attend to - each other regardless of causal order. This mirrors HuggingFace + each other regardless of causal order, subject to the sliding + window's left boundary. This mirrors HuggingFace ``blockwise_overlay`` for Gemma4 ``use_bidirectional_attention``. The returned bias bakes in causal + sliding + padding + blockwise, so the consuming ``Attention`` op MUST be called with @@ -198,8 +199,13 @@ def create_attention_bias( def as_numeric_mask(value): return op.Cast(value, to=ir.DataType.INT64) - # cumsum on attention_mask gives indices - all_indices = op.CumSum(attention_mask, 1) # axis=1 + # Distances are physical KV slots, not counts of valid tokens. An internal + # padding hole must not pull an old key back inside the sliding window. + # CumSum of ones avoids Range on providers (e.g. QNN) without a Range kernel. + slots = op.ConstantOfShape( + op.Shape(attention_mask), value=ir.tensor(np.array([1], dtype=np.int64)) + ) + all_indices = op.CumSum(slots, 1) # 1-based physical slots; differences are unchanged # kv_indices: (batch_size, 1, total_length) kv_indices = op.Unsqueeze(all_indices, [1]) @@ -228,27 +234,15 @@ def as_numeric_mask(value): if openvino_numeric_mask: full_mask = as_numeric_mask(full_mask) - if sliding_window is not None: - # Also mask out positions too far away - dist = op.Sub(q_indices, kv_indices) - within_window = op.Less(dist, sliding_window) - full_mask = ( - op.Mul(full_mask, as_numeric_mask(within_window)) - if openvino_numeric_mask - else op.And(full_mask, within_window) - ) - if block_sequence_ids is not None: - # Bidirectional vision-block overlay (OR-ed onto causal/sliding mask, - # BEFORE the padding AND, matching HF blockwise_overlay ordering). + # Apply the image overlay before the window and padding constraints. # # q_group: block id per query position -> (batch, query_length, 1). # block_sequence_ids covers the current input (== the query), so it # aligns 1:1 with the query positions. q_group = op.Unsqueeze(block_sequence_ids, [2]) # (B, q_len, 1) # kv_group: block id per kv position -> (batch, 1, total_length). - # The kv axis spans past + current; past positions are text in the - # cache, so left-pad with -1 to width total_length. + # Only current-chunk IDs are supplied; left-pad past slots with -1. pad_width = op.Sub(total_length, query_length) # [1], == past length zero_1d = op.Constant(value_ints=[0]) # Pad spec for a 2-D tensor [B, q_len]: [b_begin, s_begin, b_end, s_end]. @@ -272,6 +266,16 @@ def as_numeric_mask(value): same_block = op.And(same_group, valid_group) full_mask = op.Or(full_mask, same_block) + if sliding_window is not None: + # Same-image future keys are visible, but old keys cannot escape the window. + dist = op.Sub(q_indices, kv_indices) + within_window = op.Less(dist, sliding_window) + full_mask = ( + op.Mul(full_mask, as_numeric_mask(within_window)) + if openvino_numeric_mask + else op.And(full_mask, within_window) + ) + # Combine with attention_mask mask_value = float(dtype.min) if openvino_numeric_mask: @@ -317,7 +321,7 @@ def create_static_cache_attention_bias( ``total_length``. * **KV positions** are dense slot ids ``0 .. max_seq_len - 1`` (the absolute cache position each slot holds), so the causal/sliding comparisons use - absolute positions directly instead of cumsum indices. + absolute positions directly, as does the dynamic helper. * **Q positions** are absolute: ``write_indices[b] + arange(S_q)`` — the cache slots the current chunk is being scattered into. * **KV validity** is ``slot < nonpad_kv_seqlen[b]`` (replaces the cumsum @@ -349,8 +353,9 @@ def create_static_cache_attention_bias( block_sequence_ids: Optional ``[B, S_q]`` int64 block id per *current* query position (``>= 0`` for vision tokens sharing a block, ``-1`` for text). When set, a bidirectional overlay is OR-ed onto the - causal/sliding mask: two positions in the same block may attend to - each other regardless of causal order — mirroring HuggingFace + causal mask: two positions in the same block may attend to + each other regardless of causal order, subject to the window's + left boundary — mirroring HuggingFace ``blockwise_overlay`` for Gemma4. The per-slot KV block ids are built by scattering ``block_sequence_ids`` into a ``-1``-filled ``[B, max_seq_len]`` buffer at ``write_indices`` (the same scatter @@ -406,15 +411,8 @@ def create_static_cache_attention_bias( # Broadcasting (B, S_q, 1) >= (max_seq_len,) -> (B, S_q, max_seq_len). full_mask = op.GreaterOrEqual(q_abs, kv_slots) - if sliding_window is not None: - # Local window: keep slots within `sliding_window` of the query. - dist = op.Sub(q_abs, kv_slots) # (B, S_q, max_seq_len) - within_window = op.Less(dist, op.Constant(value_int=sliding_window)) - full_mask = op.And(full_mask, within_window) - if block_sequence_ids is not None: - # Bidirectional vision-block overlay (OR-ed BEFORE the padding AND, - # matching create_attention_bias / HF blockwise_overlay ordering). + # Apply the overlay before the left-window and padding constraints. # # q_group: block id per current query position -> (B, S_q, 1). q_group = op.Unsqueeze(block_sequence_ids, [2]) # (B, S_q, 1) @@ -438,6 +436,11 @@ def create_static_cache_attention_bias( ) full_mask = op.Or(full_mask, same_block) + if sliding_window is not None: + dist = op.Sub(q_abs, kv_slots) # (B, S_q, max_seq_len) + within_window = op.Less(dist, op.Constant(value_int=sliding_window)) + full_mask = op.And(full_mask, within_window) + # KV validity (padding): a slot is real iff it lies inside the valid # prefix, slot < nonpad_kv_seqlen[b]. Broadcasting (max_seq_len,) < # (B, 1) -> (B, max_seq_len) -> (B, 1, max_seq_len). diff --git a/src/mobius/components/_common_test.py b/src/mobius/components/_common_test.py index fad559d5c..05787a74b 100644 --- a/src/mobius/components/_common_test.py +++ b/src/mobius/components/_common_test.py @@ -7,6 +7,7 @@ import numpy as np import onnx_ir as ir +import pytest from mobius._testing import count_op_type, create_test_builder, create_test_input from mobius._testing.ort_inference import OnnxModelSession @@ -140,7 +141,7 @@ def build(execution_provider): op, input_ids, attention_mask, - sliding_window=4, + sliding_window=2, block_sequence_ids=block_sequence_ids, ) bias.name = "bias" @@ -165,12 +166,61 @@ def build(execution_provider): qnn_bias = OnnxModelSession(qnn_model, device="cpu").run(feeds)["bias"] np.testing.assert_array_equal(openvino_bias, qnn_bias) + # Exercise a same-image block wider than the window on both code paths. + feeds["attention_mask"] = np.ones((1, 8), dtype=np.int64) + feeds["block_sequence_ids"] = np.zeros((1, 4), dtype=np.int64) + for model in (openvino_model, qnn_model): + bias = OnnxModelSession(model, device="cpu").run(feeds)["bias"] + assert bias[0, 0, 0, 7] == 0 # Same-image future key. + assert bias[0, 0, 3, 4] < -1 # Old same-image key. + assert bias[0, 0, 3, 5] < -1 # Strict window boundary. + assert bias[0, 0, 3, 6] == 0 + + @pytest.mark.parametrize("execution_provider", ["default", "openvino"]) + @pytest.mark.parametrize("query_length", [1, 2, 1025]) + def test_internal_padding_hole_does_not_shrink_physical_window( + self, execution_provider, query_length + ): + from mobius._build_context import build_context + from mobius._execution_providers import get_ep + + builder, op, graph = create_test_builder() + ii = create_test_input( + builder, "input_ids", [1, query_length], dtype=ir.DataType.INT64 + ) + am = create_test_input(builder, "attention_mask", [1, 1025], dtype=ir.DataType.INT64) + bk = create_test_input( + builder, "block_sequence_ids", [1, query_length], dtype=ir.DataType.INT64 + ) + with build_context(get_ep(execution_provider)): + bias = create_attention_bias( + op, ii, am, sliding_window=1024, block_sequence_ids=bk + ) + bias.name = "bias" + graph.outputs.append(bias) + mask = np.ones((1, 1025), dtype=np.int64) + mask[0, 1] = 0 + actual = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu").run( + { + "input_ids": np.zeros((1, query_length), dtype=np.int64), + "attention_mask": mask, + "block_sequence_ids": np.zeros((1, query_length), dtype=np.int64), + } + )["bias"] + assert actual[0, 0, -1, 0] < -1 # q=1024, k=0: exact left boundary. + assert actual[0, 0, -1, 1] < -1 # Invalid key stays masked. + assert actual[0, 0, -1, 2] == 0 + if query_length > 1: + assert actual[0, 0, -2, 0] == 0 # q=1023: boundary not yet reached. + if query_length == 1025: + assert actual[0, 0, 0, 1024] == 0 # Same-image future key. + class TestBlockwiseAttentionBias: """Numerically verify the Gemma4 vision-block bidirectional overlay. ``create_attention_bias(block_sequence_ids=...)`` must bake the FULL mask - (causal [+ sliding] OR same-block, AND padding) so the Attention op can be + ((causal OR same-block) AND sliding AND padding) so the Attention op can be called with ``is_causal=0``. We build the graph, run it via ORT, and compare the attended pattern (bias == 0) against a numpy reference. """ @@ -195,15 +245,15 @@ def _build(sliding): @staticmethod def _ref(block_ids, attn, sliding): - cumsum = np.cumsum(attn) - qi = cumsum[:, None] - ki = cumsum[None, :] + slots = np.arange(len(attn)) + qi = slots[:, None] + ki = slots[None, :] m = qi >= ki - if sliding is not None: - m = m & ((qi - ki) < sliding) qg = np.array(block_ids)[:, None] kg = np.array(block_ids)[None, :] m = m | ((qg == kg) & (qg >= 0)) + if sliding is not None: + m = m & ((qi - ki) < sliding) return m & (np.array(attn)[None, :].astype(bool)) def _run(self, sliding, block_ids, attn): @@ -231,17 +281,17 @@ def test_multi_block_full_attention(self): assert not attended[3, 4] def test_block_wider_than_sliding_window(self): - # Single block spanning positions 1..4 with a window of 2: the block - # must escape the sliding window (same-block OR overrides the window). + # Same-image future keys are visible, but the strict left boundary applies. attended, expected = self._run(2, [-1, 0, 0, 0, 0, -1, -1, -1], [1] * 8) assert np.array_equal(attended, expected) - assert attended[4, 1] # distance 3 >= window, allowed via same block + assert attended[1, 4] + assert not attended[4, 1] + assert not attended[4, 2] # key == query - window is masked. + assert attended[4, 3] def test_padding_still_masked(self): - # Last two positions are padding (attention_mask == 0). - attended, expected = self._run( - 2, [-1, 0, 0, -1, -1, -1, -1, -1], [1, 1, 1, 1, 1, 1, 0, 0] - ) + # Padding must stay masked even when it carries the same image-block ID. + attended, expected = self._run(2, [-1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 0, 0]) assert np.array_equal(attended, expected) assert not attended[:, 6:].any() # nothing attends to padding diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 563dbd1e4..63bc6083d 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -8,6 +8,7 @@ import dataclasses import math from collections.abc import Callable +from typing import ClassVar import onnx_ir as ir from onnxscript import OpBuilder, nn @@ -446,6 +447,8 @@ class MoELayer(nn.Module): and compatibility path, not the grouped-expert performance representation. """ + _qmoe_activation_type: ClassVar[str] = "swiglu" + def __init__( self, config: ArchitectureConfig, @@ -637,7 +640,7 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value, *gate_args): self.fc2_experts_zero_points, None, router_weights, - activation_type="swiglu", + activation_type=self._qmoe_activation_type, normalize_routing_weights=int(normalize), k=self.top_k, # A uniform override can use the global attribute directly. diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 484957e89..f41acb0cf 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -35,9 +35,13 @@ from mobius._build_context import ep_capabilities, is_prefill_prefix_pruning_enabled from mobius._configs import ArchitectureConfig, Gemma4Config, QuantizationConfig +from mobius._constants import GEMMA4_CUSTOM_GEGLU_METADATA_KEY +from mobius._flags import flags from mobius._weight_utils import ( + QMoEQuantizationLayout, is_packed_quant_key, preprocess_quantized_weights, + resolve_qmoe_quantization, vlm_decoder_weights, vlm_embedding_weights, ) @@ -47,6 +51,7 @@ Embedding, LayerNorm, Linear, + MoELayer, RMSNorm, ScaledQuantizedEmbedding, ScaleFreeRMSNorm, @@ -91,6 +96,146 @@ ), } +_TANH_GELU_ACTIVATIONS = frozenset({"gelu_pytorch_tanh", "gelu_new", "gelu_fast"}) + + +def _mark_custom_geglu_runtime(result: ir.Value) -> ir.Value: + """Persist the experimental operator ABI requirement, not a release claim.""" + node = result.producer() + assert node is not None + node.metadata_props[GEMMA4_CUSTOM_GEGLU_METADATA_KEY] = "custom-cuda-geglu-v1" + return result + + +def _expert_source_paths(layer_idx: int) -> tuple[str, ...]: + return ( + f"model.language_model.layers.{layer_idx}", + f"model.layers.{layer_idx}", + ) + + +def _expert_quantization( + config: Gemma4Config, layer_idx: int +) -> QMoEQuantizationLayout | None: + """Resolve native experts without bypassing main's per-projection module plan.""" + if not config.enable_moe_block: + return None + quantization = _active_quantization(config.quantization_for("decoder")) + if quantization is None or not quantization.quantize_moe: + return None + if not flags.gemma4_custom_geglu: + raise NotImplementedError( + "Quantized Gemma4 GEGLU experts require explicit custom-runtime opt-in: " + "MOBIUS_GEMMA4_CUSTOM_GEGLU=1 (or override_flags(gemma4_custom_geglu=True)). " + "Unmodified released ORT support is unverified." + ) + if config.dtype not in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: + raise NotImplementedError("Quantized Gemma4 MoE requires FLOAT16 or BFLOAT16.") + layout = resolve_qmoe_quantization(quantization, _expert_source_paths(layer_idx)) + if ( + config.disable_qmoe + or config.hidden_act not in _TANH_GELU_ACTIVATIONS + or layout is None + or any( + projection.bits != 4 or projection.quant_method != "olive" + for projection in (layout.fc1, layout.fc2) + ) + ): + raise NotImplementedError( + "Quantized Gemma4 MoE experts require Olive integer-affine INT4, " + "tanh-GELU, and native QMoE enabled." + ) + return layout + + +def _extract_native_expert_weights( + state_dict: dict[str, torch.Tensor], config: Gemma4Config, *, prefix: str +) -> dict[str, torch.Tensor]: + """Consume only Gemma4 expert sidecars, leaving ordinary module plans untouched. + + Return canonical QMoE groups separately so the legacy/root Olive converter + cannot interpret ``fc1_scales`` as another raw ``*_scales`` sidecar. The shared + component loader already recognizes complete canonical MoELayer groups. + """ + result: dict[str, torch.Tensor] = {} + for index in range(config.num_hidden_layers): + layout = _expert_quantization(config, index) + if layout is None: + continue + assert config.moe_intermediate_size is not None + assert config.num_local_experts is not None + root = f"{prefix}layers.{index}" + for projection, quantization, output_size, input_size, target in ( + ( + "gate_up_proj", + layout.fc1, + 2 * config.moe_intermediate_size, + config.hidden_size, + "fc1", + ), + ("down_proj", layout.fc2, config.hidden_size, config.moe_intermediate_size, "fc2"), + ): + stem = f"{root}.experts.{projection}" + keys = [key for key in state_dict if key.startswith(stem)] + if not keys: + continue + group = quantization.group_size + if input_size % group: + raise ValueError(f"Gemma4 QMoE {stem}: input size must divide group_size.") + blocks = input_size // group + shapes = { + "_qweight": (config.num_local_experts, output_size, input_size // 2), + "_scales": (config.num_local_experts, output_size, blocks), + } + if not quantization.sym: + shapes["_qzeros"] = (config.num_local_experts, output_size, (blocks + 1) // 2) + expected = {stem + suffix for suffix in shapes} + if set(keys) != expected: + raise ValueError( + f"Gemma4 QMoE {stem}: expected sidecars {sorted(expected)}, " + f"got {sorted(keys)}." + ) + suffixes = { + "_qweight": "experts_weights", + "_scales": "scales", + "_qzeros": "experts_zero_points", + } + for suffix, shape in shapes.items(): + value = state_dict[stem + suffix] + if tuple(value.shape) != shape: + raise ValueError( + f"Gemma4 QMoE {stem + suffix}: expected shape {shape}, " + f"got {tuple(value.shape)}." + ) + if suffix == "_scales": + if ( + not value.is_floating_point() + or not torch.isfinite(value).all() + or torch.any(value < 0) + ): + raise ValueError(f"Gemma4 QMoE {stem}: invalid quantization scales.") + elif value.dtype != torch.uint8: + raise ValueError(f"Gemma4 QMoE {stem + suffix} must be uint8.") + result[f"{root}.moe.{target}_{suffixes[suffix]}"] = state_dict.pop( + stem + suffix + ) + # Fold expert output scales into block scales, never the packed INT4 codes. + scale_key = f"{root}.router.per_expert_scale" + fc2_key = f"{root}.moe.fc2_scales" + if fc2_key in result and scale_key in state_dict: + scale = state_dict[scale_key] + if ( + tuple(scale.shape) != (config.num_local_experts,) + or not torch.isfinite(scale).all() + ): + raise ValueError(f"Gemma4 QMoE {scale_key}: invalid expert output scales.") + result[fc2_key] = result[fc2_key] * scale.to(result[fc2_key].dtype)[:, None, None] + state_dict[scale_key] = torch.ones_like(scale) + for key in list(state_dict): + if key.startswith(f"{root}.router."): + result[key.replace(".router.", ".moe.gate.")] = state_dict.pop(key) + return result + def _split_per_layer_projection_weight( state_dict: dict[str, torch.Tensor], @@ -440,7 +585,7 @@ def _remap_moe_expert_weights( state_dict: dict[str, torch.Tensor], config: Gemma4Config, ) -> None: - """Rename HF MoE expert weights and fold router scale in-place. + """Rename HF MoE expert weights and fold expert output scales in-place. Shared by ``Gemma4CausalLMModel`` and ``Gemma4Model`` to avoid duplicating the rename/fold logic. @@ -479,13 +624,9 @@ def _remap_moe_expert_weights( new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") state_dict[new_key] = state_dict.pop(key) - # Fold hidden_size^-0.5 into router.scale if config.enable_moe_block: - scale_factor = float(config.hidden_size**-0.5) pes_by_prefix: dict[str, torch.Tensor] = {} for key in list(state_dict.keys()): - if ".router.scale" in key and ".per_expert_scale" not in key: - state_dict[key] = state_dict[key] * scale_factor if key.endswith(".router.per_expert_scale"): pes_by_prefix[key[: -len(".router.per_expert_scale")]] = state_dict[key] # Fold per_expert_scale into the expert down projection (fc2). HF's router @@ -807,12 +948,14 @@ def forward( y_pos = op.Gather(pixel_position_ids, op.Constant(value_int=1), axis=2) # [B, T] is_padding = op.And(op.Equal(x_pos, neg_one), op.Equal(y_pos, neg_one)) # [B, T] bool - # zero out padding hidden states: features * (1 - is_padding) - not_padding_f = op.CastLike( - op.Cast(op.Not(is_padding), to=ir.DataType.FLOAT), vision_features - ) - not_padding_f = op.Unsqueeze(not_padding_f, [2]) # [B, T, 1] - vision_features = op.Mul(vision_features, not_padding_f) # [B, T, D] + # Match HF masked_fill, including padded NaN/Inf, in the FP32 pooling + # path. Keep the original features as the mean-rounding dtype reference. + features_f32 = op.Cast(vision_features, to=ir.DataType.FLOAT) + masked_features_f32 = op.Where( + op.Unsqueeze(is_padding, [2]), + op.Constant(value_float=0.0), + features_f32, + ) # [B, T, D] # --- 2. Clamp positions to [0, ∞) so padding doesn't break Div ------ clamped_x = op.Max(x_pos, op.Constant(value_int=0)) # [B, T] @@ -854,23 +997,25 @@ def forward( on_val = 1.0 / float(k2) one_hot_vals = op.Constant(value_floats=[0.0, on_val]) weights = op.OneHot(kernel_idxs, valid_depth, one_hot_vals) # [B, T, valid_depth] f32 - # Cast to model dtype for the subsequent MatMul - weights = op.CastLike(weights, vision_features) - # --- 6. Weighted sum: [B, valid_depth, T] @ [B, T, D] --------------- + # HF accumulates FP32 coefficients and FP32 features, then rounds the + # pooled mean once to working dtype. BF16 coefficients (notably 1/9) + # change that mean even if the MatMul itself uses FP32 accumulation. weights_t = op.Transpose(weights, perm=[0, 2, 1]) # [B, valid_depth, T] - output = op.MatMul(weights_t, vision_features) # [B, valid_depth, D] + output = op.MatMul(weights_t, masked_features_f32) # [B, valid_depth, D] + output = op.CastLike(output, vision_features) # --- 7. Scale by sqrt(hidden_size) matching HF VisionPooler ---------- - scale = op.CastLike(op.Constant(value_float=self._pooler_scale), output) - return op.Mul(output, scale) + # HF rounds the pooled mean to working dtype before FP32 scaling/standardization. + output = op.Cast(output, to=ir.DataType.FLOAT) + return op.Mul(output, op.Constant(value_float=self._pooler_scale)) class _Gemma4VisionPatchEmbedder(nn.Module): """Gemma4 patch embedder: linear projection + 2D position lookup. Inputs: - - ``pixel_values [B, N, 3*patch_size^2]``: pre-patchified, normalized to ``[-1, 1]`` + - ``pixel_values [B, N, 3*patch_size^2]``: FP32 processor patches in ``[0, 1]`` - ``pixel_position_ids [B, N, 2]``: (x, y) coordinates for each patch Output: ``[B, N, hidden_size]`` @@ -905,10 +1050,11 @@ def forward( position embeddings are zeroed exactly as HF does in ``Gemma4VisionPatchEmbedder._position_embeddings``. """ - # pixel_values in [0,1] -> normalize to [-1, 1]: 2*(v - 0.5) = 2v - 1 - two = op.CastLike(op.Constant(value_float=2.0), pixel_values) - one = op.CastLike(op.Constant(value_float=1.0), pixel_values) - pixel_values = op.Sub(op.Mul(pixel_values, two), one) + # HF normalizes processor pixels in FP32, then rounds to the projection dtype. + # Casting processor pixels first loses information, especially for BF16. + pixel_values = op.Cast(pixel_values, to=ir.DataType.FLOAT) + pixel_values = op.Mul(op.Sub(pixel_values, 0.5), 2.0) + pixel_values = op.CastLike(pixel_values, self.position_embedding_table) hidden_states = self.input_proj(op, pixel_values) # [B, N, hidden] # Detect padding patches: x-coord == -1 means the patch is padding. @@ -949,6 +1095,7 @@ class _Gemma4VisionEncoderCore(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() vc = config.vision # VisionConfig for the SigLIP encoder + linear_class, clippable_linear_class = _vision_linear_classes(config) self.patch_embedder = _Gemma4VisionPatchEmbedder( patch_size=vc.patch_size or 16, @@ -1326,7 +1473,7 @@ def forward( past_key = past_key_value[0] if past_key_value is not None else None past_value = past_key_value[1] if past_key_value is not None else None - gqa_attrs: dict = { + gqa_attrs = { "num_heads": self.num_attention_heads, "kv_num_heads": self.num_key_value_heads, "scale": self.scaling, @@ -1442,44 +1589,45 @@ def forward( class _Gemma4MoeRouter(nn.Module): - """Gemma4 MoE router: scale-free RMSNorm → learned scale → linear → softmax. + """Gemma4 MoE router: scale-free RMSNorm → learned scale → raw logits. Matches ``Gemma4TextRouter`` in HuggingFace. - Parameters (aligned with HF state_dict, with one weight-fold): - - ``scale`` [hidden_size]: ``router.scale * hidden_size^-0.5`` (folded during - ``preprocess_weights`` — see ``Gemma4CausalLMModel.preprocess_weights``) + Parameters (aligned with HF state_dict): + - ``scale`` [hidden_size]: unmodified learned router scale - ``proj.weight`` [num_experts, hidden_size]: router logits (no bias) - ``per_expert_scale`` [num_experts]: per-expert output weight scaling ``router.norm`` is scale-free (no learnable parameter, omitted from state_dict). - The ``hidden_size^-0.5`` scale factor is absorbed into ``self.scale`` at weight-load - time to avoid graph-level float-constant collisions across multiple decoder layers. - Args: hidden_size: Model hidden dimension. num_experts: Total number of experts. rms_norm_eps: Epsilon for the scale-free RMSNorm. """ + # Olive expert quantization leaves the learned router projection floating point. + # Use the shared component loader's exclusion hook rather than rewriting its plan. + component_quantization_excluded_methods: ClassVar[frozenset[str]] = frozenset( + {"olive", "gptq", "awq"} + ) + def __init__(self, hidden_size: int, num_experts: int, rms_norm_eps: float = 1e-6): super().__init__() self._eps = rms_norm_eps self._hidden_size = hidden_size - # ``scale`` stores router.scale * hidden_size^-0.5 (folded in preprocess_weights) self.scale = nn.Parameter([hidden_size]) self.proj = Linear(hidden_size, num_experts, bias=False) self.per_expert_scale = nn.Parameter([num_experts]) def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: - """Compute router probabilities over all experts. + """Compute raw logits; native MoE/QMoE performs softmax internally. Args: hidden_states: [num_tokens, hidden_size] (batch-sequence flattened) Returns: - router_probs: [num_tokens, num_experts] full softmax probabilities + Router logits: [num_tokens, num_experts]. """ # Scale-free RMSNorm using RMSNormalization op (epsilon as attribute avoids # graph-level float constant name collisions across multiple decoder layers). @@ -1489,11 +1637,29 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: hidden_states, ) x_normed = op.RMSNormalization(hidden_states, ones, epsilon=self._eps, axis=-1) - # Scale: x_normed * self.scale (hidden_size^-0.5 already folded into scale) x_scaled = op.Mul(x_normed, self.scale) + # Keep the learned-scale rounding, then use FP32 opmath for HF's Python scalar. + x_scaled = op.CastLike( + op.Mul( + op.Cast(x_scaled, to=ir.DataType.FLOAT), + op.Constant(value_float=float(self._hidden_size**-0.5)), + ), + x_scaled, + ) # Linear projection then softmax -> router_probs [num_tokens, num_experts] - expert_scores = self.proj(op, x_scaled) - return op.Softmax(expert_scores, axis=-1) + return self.proj(op, x_scaled) + + def qmoe_routing( + self, op: OpBuilder, hidden_states: ir.Value, router_input: ir.Value + ) -> tuple[ir.Value, None, bool, float]: + # Router sees the residual; experts see a separately normalized activation. + return self.forward(op, router_input), None, True, 1.0 + + +class _Gemma4QMoELayer(MoELayer): + """Native INT4 experts with exact tanh-GEGLU instead of approximate SwiGLU.""" + + _qmoe_activation_type = "geglu" class Gemma4DecoderLayer(nn.Module): @@ -1591,6 +1757,7 @@ def __init__(self, config: Gemma4Config, layer_idx: int): self.act_fn = get_activation(config.hidden_act) self._enable_moe_block = config.enable_moe_block + self.moe: MoELayer | None = None if config.enable_moe_block: assert config.num_local_experts is not None, "num_local_experts required for MoE" assert config.num_experts_per_tok is not None, ( @@ -1604,18 +1771,17 @@ def __init__(self, config: Gemma4Config, layer_idx: int): self._num_experts = config.num_local_experts self._moe_intermediate_size = config.moe_intermediate_size self._hidden_size = config.hidden_size - # The fused com.microsoft::MoE op's only GATED activation is - # ``swiglu`` = ``x * sigmoid(activation_alpha * x) * up``. Gemma 4 - # experts use ``act_fn(gate) * up`` where ``act_fn`` is the config's - # ``hidden_act`` (``gelu_pytorch_tanh`` for the 26B-A4B / 31B MoE - # checkpoints), NOT silu. Since ``x * sigmoid(1.702 * x)`` closely - # approximates gelu, gelu-family activations map to - # ``activation_alpha=1.702``; a genuine silu expert maps to 1.0. _moe_act = (config.hidden_act or "").lower() - self._moe_swiglu_alpha = 1.702 if "gelu" in _moe_act else 1.0 + self._fused_moe_activation = ( + "geglu" + if _moe_act in _TANH_GELU_ACTIVATIONS + else "swiglu" + if _moe_act in {"silu", "swish"} + else None + ) # The unfused fallback (EPs without a fused MoE op) uses explicit # ONNX ops, so it applies the EXACT expert activation via the shared - # ACT2FN mapping instead of the swiglu approximation. This keeps + # ACT2FN mapping. This keeps # ``gelu`` (erf) distinct from the tanh-approx variants # (``gelu_pytorch_tanh``/``gelu_new``/``gelu_fast``) and ``silu`` -> # Swish, faithfully matching HuggingFace's ``act_fn(gate) * up`` for @@ -1623,19 +1789,32 @@ def __init__(self, config: Gemma4Config, layer_idx: int): self._moe_act_fn = get_activation(config.hidden_act) moe_inter = config.moe_intermediate_size - self.router = _Gemma4MoeRouter( + router = _Gemma4MoeRouter( config.hidden_size, config.num_local_experts, config.rms_norm_eps ) + if _expert_quantization(config, layer_idx) is not None: + self.moe = _Gemma4QMoELayer( + dataclasses.replace( + config, + quantization=config.quantization_for("decoder"), + component_quantization=None, + qmoe_source_paths=_expert_source_paths(layer_idx), + ), + gate=router, + ) + else: + self.router = router # Expert weights stored as 3D tensors (all experts stacked). # fc1_experts_weights: gate+up combined [E, 2*moe_inter, hidden]. # fc2_experts_weights: down projection [E, hidden, moe_inter]. # These map to HF experts.gate_up_proj and experts.down_proj. - self.fc1_experts_weights = nn.Parameter( - [config.num_local_experts, 2 * moe_inter, config.hidden_size] - ) - self.fc2_experts_weights = nn.Parameter( - [config.num_local_experts, config.hidden_size, moe_inter] - ) + if self.moe is None: + self.fc1_experts_weights = nn.Parameter( + [config.num_local_experts, 2 * moe_inter, config.hidden_size] + ) + self.fc2_experts_weights = nn.Parameter( + [config.num_local_experts, config.hidden_size, moe_inter] + ) # MoE-specific norms (in addition to the shared pre/post_feedforward norms) self.pre_feedforward_layernorm_2 = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -1699,7 +1878,7 @@ def forward( # combined: post_ff_norm(dense + moe) + residual dense_out = self.post_feedforward_layernorm_1(op, hidden_states) - # MoE input is the pre-attention residual (flattened to 2D for routing). + # Route the feed-forward residual, flattened to [B*S, H]. batch_size = op.Shape(residual, start=0, end=1) # [1] scalar seq_len = op.Shape(residual, start=1, end=2) # [1] scalar hidden_size = op.Shape(residual, start=2, end=3) # [1] scalar @@ -1707,25 +1886,33 @@ def forward( flat_shape = op.Concat(num_tokens, hidden_size, axis=0) # [2] residual_flat = op.Reshape(residual, flat_shape) # [B*S, H] - # router_probs: [B*S, E] full softmax over all experts - router_probs = self.router(op, residual_flat) # [B*S, E] - # Norm residual before experts normed_flat = self.pre_feedforward_layernorm_2(op, residual_flat) # [B*S, H] caps = ep_capabilities() - if caps.supports_fused_moe: + custom_geglu = caps.name == "cuda" and flags.gemma4_custom_geglu + if self.moe is not None: + if not (caps.supports_fused_moe and custom_geglu): + raise NotImplementedError( + "Quantized Gemma4 MoE requires a GEGLU-capable custom CUDA runtime " + "and MOBIUS_GEMMA4_CUSTOM_GEGLU=1." + ) + moe_out_flat = op.CastLike( + _mark_custom_geglu_runtime(self.moe(op, normed_flat, residual_flat)), + normed_flat, + ) + elif ( + caps.supports_fused_moe + and self._fused_moe_activation is not None + and (self._fused_moe_activation != "geglu" or custom_geglu) + ): # Fused com.microsoft::MoE op handles top-k selection + # expert dispatch internally. Requires ORT main (post # microsoft/onnxruntime#28467, MoE GEMM Refactor) which # plumbs the SwiGLU schema attributes to the kernel. # - # Gemma 4 gated-FFN semantics via the fused SwiGLU schema: - # activation_type=swiglu — x * sigmoid(activation_alpha*x) * up - # activation_alpha=1.702 — approximates gelu_pytorch_tanh, the - # activation Gemma 4 experts actually - # use (act_fn(gate)*up); a silu expert - # would use 1.0. (See _moe_swiglu_alpha.) + # GEGLU is tanh-GELU(gate)*up, not sigmoid(1.702*gate)*gate*up. + # SwiGLU uses alpha=1 for genuine SiLU; GEGLU ignores alpha. # activation_beta=0.0 — linear * gate (no GPT-OSS "+1" bias) # swiglu_limit=inf — no clipping (≤0 disables the clamp) # swiglu_fusion=1 — interleaved layout @@ -1763,28 +1950,33 @@ def forward( ), op.Constant(value_ints=[e_dim, 2 * inter, hidden]), ) # [E, 2*inter, H] interleaved - moe_out_flat = op.CastLike( - op.MoE( # type: ignore[attr-defined] - normed_flat, - router_probs, - fc1_interleaved, - None, # fc1_experts_bias (slot 3, optional) - self.fc2_experts_weights, - activation_type="swiglu", - k=self._top_k, - normalize_routing_weights=1, - activation_alpha=self._moe_swiglu_alpha, - activation_beta=0.0, - swiglu_limit=float("inf"), - swiglu_fusion=1, - _domain="com.microsoft", - ), - normed_flat, # match input dtype (bf16/fp16/fp32) - ) # [B*S, H] + moe_result = op.MoE( # type: ignore[attr-defined] + normed_flat, + self.router(op, residual_flat), + fc1_interleaved, + None, # fc1_experts_bias (slot 3, optional) + self.fc2_experts_weights, + activation_type=self._fused_moe_activation, + k=self._top_k, + normalize_routing_weights=1, + activation_alpha=1.0, + activation_beta=0.0, + swiglu_limit=float("inf"), + swiglu_fusion=1, + _domain="com.microsoft", + ) + if self._fused_moe_activation == "geglu": + moe_result = _mark_custom_geglu_runtime(moe_result) + moe_out_flat = op.CastLike(moe_result, normed_flat) # [B*S, H] else: # EPs without fused MoE support fall back to a static # per-expert unroll. - moe_out_flat = self._dispatch_moe_fallback(op, normed_flat, router_probs) + router_logits = self.router(op, residual_flat) + moe_out_flat = self._dispatch_moe_fallback( + op, + normed_flat, + op.Softmax(op.Cast(router_logits, to=ir.DataType.FLOAT), axis=-1), + ) moe_out = op.Reshape(moe_out_flat, op.Shape(residual)) # [B, S, H] moe_out = self.post_feedforward_layernorm_2(op, moe_out) @@ -1874,7 +2066,7 @@ def _dispatch_moe_fallback( up = op.Einsum(normed_flat, up_w, equation="th,eih->tei") # Exact expert activation via ACT2FN: ``gelu_pytorch_tanh`` -> tanh-approx # GELU, bare ``gelu`` -> erf GELU, ``silu``/``swish`` -> Swish, etc. This - # avoids the swiglu approximation used on the fused path. + # preserves the distinction between erf GELU and tanh-GEGLU. activated = self._moe_act_fn(op, gate) inter = op.Mul(activated, up) # [T, E, moe_inter] @@ -1882,7 +2074,11 @@ def _dispatch_moe_fallback( expert_out = op.Einsum(inter, self.fc2_experts_weights, equation="tei,ehi->teh") # Weight by the dense routing vector and sum over experts. - weighted = op.Mul(expert_out, op.Unsqueeze(dense_w, [2])) # [T, E, H] * [T, E, 1] + # HF rounds each FP32-weighted contribution back to the expert dtype. + weighted = op.CastLike( + op.Mul(op.Cast(expert_out, to=ir.DataType.FLOAT), op.Unsqueeze(dense_w, [2])), + expert_out, + ) # [T, E, H] * [T, E, 1] return op.ReduceSum(weighted, [1], keepdims=0) # [T, H] @@ -1984,9 +2180,9 @@ def __init__(self, config: Gemma4Config): self.layer_types = layer_types self.sliding_window = config.sliding_window # Bidirectional attention mode (None | "vision"). When "vision", - # contiguous image-token blocks attend bidirectionally; the overlay is - # derived from input_ids at runtime via ``block_sequence_ids`` and forces - # the float-bias attention path (``is_causal=0``). HF also defines an + # contiguous image-token blocks attend bidirectionally in local layers. + # The overlay is derived from input_ids via ``block_sequence_ids`` and + # uses the float-bias attention path (``is_causal=0``) locally. HF defines an # "all" mode (every token bidirectional, no causal mask); it is not used # by any supported Gemma4 checkpoint and not implemented here, so reject # it explicitly rather than silently falling back to causal attention. @@ -2267,12 +2463,10 @@ def forward( from mobius.components._attention import GQAContext, StaticCacheState dtype = get_build_dtype() - # Bidirectional vision-block overlay (Gemma4 larger models). When - # active, contiguous vision-token blocks attend bidirectionally on - # BOTH full and sliding layers. This cannot be expressed by the - # GroupQueryAttention op (causal / local-window only), so we force - # the float-bias Attention path with ``is_causal=0`` and bake the - # full mask (causal + sliding + padding + blockwise OR) into the bias. + # HF applies the image overlay only to local layers: + # valid_key AND (causal OR same-image) AND (key > query - window). + # Global layers stay causal. Retain the conservative all-Attention + # construction here; the normal optimizer can fuse causal globals to GQA. # # ``block_sequence_ids`` [B, S] identifies contiguous image token # blocks. It is derived from ``input_ids`` (image token spans only; @@ -2283,16 +2477,8 @@ def forward( # to the decoder but cannot forward an arbitrary int tensor). When a # caller supplies ``block_sequence_ids`` directly it is used as-is. # - # Trade-off: the overlay is a *static* graph choice — once a model is - # built with ``use_bidirectional_attention == "vision"`` the decoder - # always takes the float-bias (``is_causal=0``) path and forgoes GQA, - # for text-only prompts and every decode step too. This is unavoidable - # in a single static graph: image-token presence is data-dependent at - # runtime, and the same graph serves both image prefill and text decode. - # It stays numerically correct everywhere — when there are no image - # tokens the block ids are all ``-1`` and the ``q_group >= 0`` guard in - # ``create_attention_bias`` makes the overlay a pure no-op (plain causal, - # matching HuggingFace). Only the fused/GQA fast path is given up. + # The overlay is a static graph choice even for text-only decode. + # All -1 block IDs make it a no-op at runtime. bidirectional = self._use_bidirectional_attention == "vision" and self._has_image_token if bidirectional and block_sequence_ids is None and input_ids is not None: block_sequence_ids = _compute_block_sequence_ids( @@ -2311,11 +2497,6 @@ def forward( and caps.supports_fused_rope and not use_block_overlay ) - # When the blockwise overlay is active the Attention op must NOT - # re-apply its built-in causal mask (it would cancel the - # future-position unmasking baked into the float bias). - attn_is_causal = 0 if use_block_overlay else 1 - if use_gqa: # Calling forward() on the RoPE modules materializes their # cos_cache / sin_cache nn.Parameters as ONNX graph initializers. @@ -2447,7 +2628,6 @@ def forward( input_ids=query_input, attention_mask=attention_mask, dtype=self._dtype, - block_sequence_ids=block_sequence_ids if use_block_overlay else None, ), } fallback_pos_dict = position_embeddings_dict @@ -2529,7 +2709,8 @@ def forward( shared_kv_states=shared_kv_states, per_layer_input=per_layer_input, past_key_value=past_kv, - is_causal=attn_is_causal, + # Only local layers may unmask future keys within an image block. + is_causal=0 if use_block_overlay and layer_type == "sliding_attention" else 1, static_kv_seqlen=static_kv_seqlen, ) # KV-shared layers borrow K,V from source layers — exclude from @@ -2684,10 +2865,14 @@ def preprocess_weights( # HF's model.embed_tokens_per_layer.weight [V, L*D] maps directly # to our fused embedding table — no splitting needed. # (For WebGPU, splitting is handled by _Gemma4DecoderModel.preprocess_weights.) - # Map HF expert weight names and fold router scale + config = self.config + if not isinstance(config, Gemma4Config): + raise TypeError("Gemma4 expert preprocessing requires a Gemma4Config.") + native = _extract_native_expert_weights(state_dict, config, prefix="model.") + # Map HF expert weight names and fold expert output scales. _remap_moe_expert_weights(state_dict, self.config) _split_per_layer_projection_weight(state_dict, "model.", self.config) - return super().preprocess_weights(state_dict) + return super().preprocess_weights(state_dict) | native def static_kv_cache_specs(self) -> list[tuple[int, int]]: """Per-cache-layer ``(num_kv_heads, head_dim)`` for static-cache mode.""" @@ -2793,11 +2978,16 @@ class _Gemma4VisionEncoderModel(nn.Module): def __init__(self, config: Gemma4Config): super().__init__() vc = config.vision # VisionConfig for the SigLIP encoder + if vc is None or vc.hidden_size is None: + raise ValueError("Gemma4 vision encoder requires vision.hidden_size.") self._text_hidden_size = config.hidden_size + self._dtype = config.dtype self.encoder = _Gemma4VisionEncoderCore(config) # Gemma4VisionPooler: 3x3 spatial average pooling + sqrt(hidden) scaling. # Reduces N patches to N/9 before projection. self.pooler = Gemma4VisionPooler(vc.hidden_size, vc.pooling_kernel_size or 3) + self.std_bias = nn.Parameter([vc.hidden_size]) if vc.standardize else None + self.std_scale = nn.Parameter([vc.hidden_size]) if vc.standardize else None self.projector_norm = ScaleFreeRMSNorm(vc.hidden_size, eps=vc.norm_eps) linear_class, _ = _vision_linear_classes(config) self.projector = linear_class(vc.hidden_size, config.hidden_size, bias=False) @@ -2815,6 +3005,13 @@ def forward( # Passes pixel_position_ids so the pooler can assign patches to spatial buckets. vision_features = self.pooler(op, vision_features, pixel_position_ids) + if self.std_bias is not None and self.std_scale is not None: + vision_features = op.Mul( + op.Sub(vision_features, op.Cast(self.std_bias, to=ir.DataType.FLOAT)), + op.Cast(self.std_scale, to=ir.DataType.FLOAT), + ) + vision_features = op.Cast(vision_features, to=self._dtype) + # Scale-free norm + linear projection -> [B, N/9, text_hidden] vision_features = self.projector_norm(op, vision_features) vision_features = self.projector(op, vision_features) @@ -2832,7 +3029,9 @@ def preprocess_weights( ) -> dict[str, torch.Tensor]: renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - if key.startswith("vision_tower."): + if key in ("vision_tower.std_bias", "vision_tower.std_scale"): + renamed[key.removeprefix("vision_tower.")] = value + elif key.startswith("vision_tower."): new_key = "encoder." + key[len("vision_tower.") :] # HF Gemma4VisionModel wraps its layer list in a Gemma4VisionEncoder # submodule, adding an extra "encoder." level. Our _Gemma4VisionEncoderCore @@ -3543,6 +3742,8 @@ def preprocess_weights( elif key.startswith("lm_head."): renamed["decoder." + key] = value + elif key in ("vision_tower.std_bias", "vision_tower.std_scale"): + renamed["vision_encoder." + key.removeprefix("vision_tower.")] = value elif key.startswith("vision_tower."): new_key = "vision_encoder.encoder." + key[len("vision_tower.") :] # HF wraps encoder layers under an extra "encoder." sub-module; strip it @@ -3599,10 +3800,12 @@ def preprocess_weights( else: renamed[key] = value - # Map HF expert weight names and fold router scale + native = _extract_native_expert_weights(renamed, self.config, prefix="decoder.model.") + # Map HF expert weight names and fold expert output scales. _remap_moe_expert_weights(renamed, self.config) renamed = _preprocess_component_quantized_weights(renamed, self.config) + renamed.update(native) # For WebGPU: the fused [V, L*D] embed_tokens_per_layer exceeds the 256 MiB # per-buffer limit. Split it into L separate [V, D] tables in the decoder. diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index aa8abbcbb..a5709e2ee 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -58,6 +58,133 @@ def _tiny_gemma4_config(**overrides) -> Gemma4Config: return Gemma4Config(**defaults) +class TestGemma4HfPrecisionAlignment: + @pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] + ) + @pytest.mark.parametrize("quantized", [False, True]) + @pytest.mark.parametrize("optimized", [False, True]) + def test_vision_accepts_float32_processor_pixels(self, dtype, quantized, optimized): + from mobius._builder import _cast_module_dtype, build_from_module + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_config( + dtype=dtype, + enable_moe_block=False, + quantization=( + QuantizationConfig( + quant_method="olive", bits=4, group_size=16, quantize_vision=True + ) + if quantized + else None + ), + ) + module = Gemma4Model(config) + _cast_module_dtype(module, dtype) + package = ( + build_from_module(module, config, task="gemma4") + if optimized + else Gemma4Task().build(module, config) + ) + graph = package["vision_encoder"].graph + assert graph.inputs[0].name == "pixel_values" + assert graph.inputs[0].dtype == ir.DataType.FLOAT + projection = next(node for node in graph if node.op_type in {"MatMul", "MatMulNBits"}) + assert projection.inputs[0].dtype == dtype + + @pytest.mark.parametrize( + ("dtype", "torch_dtype"), + [ + (ir.DataType.FLOAT, torch.float32), + (ir.DataType.FLOAT16, torch.float16), + (ir.DataType.BFLOAT16, torch.bfloat16), + ], + ) + def test_pixel_normalization_before_projection_cast(self, dtype, torch_dtype): + import numpy as np + from onnx_ir import tensor_adapters + + from mobius._builder import _cast_module_dtype + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _Gemma4VisionPatchEmbedder + from mobius.tasks._base import _make_graph + + module = _Gemma4VisionPatchEmbedder(1, 4, 1) + _cast_module_dtype(module, dtype) + graph, builder = _make_graph() + pixels = builder.input("pixel_values", ir.DataType.FLOAT, [1, 256, 3]) + positions = builder.input("positions", ir.DataType.INT64, [1, 256, 2]) + module(builder.op, pixels, positions) + projection = next(node for node in graph if node.op_type == "MatMul") + normalized = projection.inputs[0] + # Execute only normalization: no CPU BF16 MatMul kernel is needed. + normalized_f32 = builder.op.Cast(normalized, to=ir.DataType.FLOAT) + builder.add_output(normalized_f32, "normalized") + module.position_embedding_table.const_value = tensor_adapters.TorchTensor( + torch.zeros(2, 1, 4, dtype=torch_dtype) + ) + prefix = ir.convenience.extract(graph, [pixels], [normalized_f32]) + values = (torch.arange(256).float() / 255).reshape(1, 256, 1).expand(1, 256, 3) + expected = (2 * (values - 0.5)).to(torch_dtype).float().numpy() + actual = OnnxModelSession(ir.Model(prefix, ir_version=10), device="cpu").run( + {"pixel_values": values.numpy().copy()} + )["normalized"] + np.testing.assert_array_equal(actual, expected) + if dtype == ir.DataType.BFLOAT16: + old = (values.to(torch_dtype) * 2 - 1).float().numpy() + assert np.count_nonzero(old != expected) == 128 * 3 + + @pytest.mark.parametrize( + ("dtype", "torch_dtype"), + [ + (ir.DataType.FLOAT, torch.float32), + (ir.DataType.FLOAT16, torch.float16), + (ir.DataType.BFLOAT16, torch.bfloat16), + ], + ) + def test_router_scalar_after_learned_scale(self, dtype, torch_dtype): + import ml_dtypes + import numpy as np + + from mobius._builder import _cast_module_dtype + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _Gemma4MoeRouter + from mobius.tasks._base import _make_graph + + # A non-power-of-two hidden size exposes BF16 scalar-rounding differences. + generator = torch.Generator().manual_seed(7) + normed = torch.randn(4, 2816, generator=generator).to(torch_dtype) + scale = (0.2 + torch.rand(2816, generator=generator)).to(torch_dtype) + learned = normed * scale + expected = (learned * (2816**-0.5)).float().numpy() + module = _Gemma4MoeRouter(2816, 4) + _cast_module_dtype(module, dtype) + graph, builder = _make_graph() + hidden = builder.input("hidden", dtype, [4, 2816]) + module(builder.op, hidden) + projection = next(node for node in graph if node.op_type == "MatMul") + scaled = projection.inputs[0] + assert scaled.producer().op_type == "CastLike" + learned_value = scaled.producer().inputs[1] + assert learned_value.producer().op_type == "Mul" + assert learned_value.producer().inputs[0].producer().op_type == "RMSNormalization" + output = builder.op.Cast(scaled, to=ir.DataType.FLOAT) + builder.add_output(output, "scaled") + prefix = ir.convenience.extract(graph, [learned_value], [output]) + np_dtype = { + ir.DataType.FLOAT: np.float32, + ir.DataType.FLOAT16: np.float16, + ir.DataType.BFLOAT16: ml_dtypes.bfloat16, + }[dtype] + actual = OnnxModelSession(ir.Model(prefix, ir_version=10), device="cpu").run( + {learned_value.name: learned.float().numpy().astype(np_dtype)} + )["scaled"] + np.testing.assert_array_equal(actual, expected) + if dtype == ir.DataType.BFLOAT16: + folded = normed * (scale * (2816**-0.5)) + assert np.count_nonzero(folded.float().numpy() != expected) > 0 + + class TestGemma4CausalLMPreprocessWeights: """Test Gemma4CausalLMModel.preprocess_weights.""" @@ -75,16 +202,16 @@ def test_expert_weight_rename(self): assert "model.layers.0.fc2_experts_weights" in result assert "model.layers.0.experts.gate_up_proj" not in result - def test_router_scale_folding(self): + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_router_scale_preserved(self, dtype): config = _tiny_gemma4_config() model = Gemma4CausalLMModel(config) - scale_val = torch.tensor([2.0]) + scale_val = torch.tensor([2.0], dtype=dtype) fake_sd = {"model.layers.0.router.scale": scale_val.clone()} result = model.preprocess_weights(fake_sd) - expected = 2.0 * (64**-0.5) # hidden_size=64 - assert abs(result["model.layers.0.router.scale"].item() - expected) < 1e-6 + torch.testing.assert_close(result["model.layers.0.router.scale"], scale_val) def test_quantized_moe_experts_fail_closed(self): config = _tiny_gemma4_config( @@ -121,18 +248,18 @@ def test_expert_weight_rename(self): assert "decoder.model.layers.0.fc1_experts_weights" in result assert "decoder.model.layers.0.fc2_experts_weights" in result - def test_router_scale_folding(self): + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_router_scale_preserved(self, dtype): config = _tiny_gemma4_config() model = Gemma4Model(config) - scale_val = torch.tensor([2.0]) + scale_val = torch.tensor([2.0], dtype=dtype) fake_sd = {"model.language_model.layers.0.router.scale": scale_val.clone()} result = model.preprocess_weights(fake_sd) - expected = 2.0 * (64**-0.5) key = "decoder.model.layers.0.router.scale" assert key in result - assert abs(result[key].item() - expected) < 1e-6 + torch.testing.assert_close(result[key], scale_val) def test_quantized_moe_experts_fail_closed(self): config = _tiny_gemma4_config( @@ -1103,14 +1230,28 @@ def test_package_wires_block_sequence_ids_end_to_end(self): assert "block_sequence_ids" not in dec_inputs assert "input_ids" in dec_inputs - # Decoder attention must drop GQA and disable the op's built-in causal - # mask (is_causal=0) so the baked blockwise bias is honored. + # Only local layers disable the built-in causal mask for the image overlay. dec = pkg["decoder"].graph assert not any(n.op_type == "GroupQueryAttention" for n in dec) attn_nodes = [n for n in dec if n.op_type == "Attention"] - assert attn_nodes - for n in attn_nodes: - assert n.attributes["is_causal"].as_int() == 0 + assert len(attn_nodes) == len(config.layer_types) + for node, layer_type in zip(attn_nodes, config.layer_types, strict=True): + expected = 0 if layer_type == "sliding_attention" else 1 + assert node.attributes["is_causal"].as_int() == expected + # A global mask must have no data dependency on the image block overlay. + global_bias = attn_nodes[1].inputs[3] + pending = [global_bias] + visited = set() + while pending: + value = pending.pop() + if value is None or value.producer() is None: + continue + producer = value.producer() + if producer in visited: + continue + visited.add(producer) + assert producer.op_type != "Or" + pending.extend(producer.inputs) def test_no_block_sequence_ids_when_causal(self): """Without bidirectional attention, no input_ids/overlay is wired.""" @@ -1126,6 +1267,81 @@ def test_no_block_sequence_ids_when_causal(self): assert "block_sequence_ids" not in dec_inputs assert "input_ids" not in dec_inputs + @pytest.mark.parametrize("mask", [[1, 0, 1, 1], [1, 1, 1, 0]]) + def test_cuda_export_preserves_global_cached_mask_holes(self, mask): + import numpy as np + + from mobius import build_from_module + from mobius._pipeline_contract import requires_arbitrary_attention_mask + from mobius._testing.ort_inference import OnnxModelSession + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_config( + dtype=ir.DataType.FLOAT16, + enable_moe_block=False, + use_bidirectional_attention="vision", + image_token_id=255, + ) + package = build_from_module( + Gemma4Model(config), config, task=Gemma4Task(), execution_provider="cuda" + ) + decoder = package["decoder"] + assert requires_arbitrary_attention_mask(decoder.graph) + assert not any(node.op_type == "GroupQueryAttention" for node in decoder.graph) + global_attention = next( + node + for node in decoder.graph + if node.op_type == "Attention" + and node.attributes["kv_num_heads"].as_int() == config.num_global_key_value_heads + ) + assert global_attention.attributes["is_causal"].as_int() == 1 + bias = global_attention.inputs[3] + assert bias is not None and bias.name is not None + # Execute the actual optimized bias with one query and three past slots. + # Equal token counts must not collapse these distinct visible-key sets. + inputs = [ + value + for value in decoder.graph.inputs + if value.name in {"input_ids", "attention_mask"} + ] + graph = ir.convenience.extract(decoder.graph, inputs, [bias]) + session = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu") + try: + actual = session.run( + { + "input_ids": np.array([[5]], dtype=np.int64), + "attention_mask": np.array([mask], dtype=np.int64), + } + )[bias.name] + finally: + session.close() + np.testing.assert_array_equal(actual[0, 0, 0] > -1, np.array(mask, dtype=bool)) + + def test_text_only_cuda_export_keeps_gqa_eligibility(self): + from mobius import build_from_module + from mobius._pipeline_contract import requires_arbitrary_attention_mask + + config = _tiny_gemma4_config( + model_type="gemma4_text", + dtype=ir.DataType.FLOAT16, + enable_moe_block=False, + vision=None, + image_token_id=None, + use_bidirectional_attention=None, + ) + package = build_from_module( + Gemma4CausalLMModel(config), + config, + task="gemma4-text-generation", + execution_provider="cuda", + ) + graph = package["model"].graph + assert not requires_arbitrary_attention_mask(graph) + assert ( + sum(node.op_type == "GroupQueryAttention" for node in graph) + == config.num_hidden_layers + ) + def _tiny_gemma4_unified_config(**overrides) -> Gemma4Config: """Minimal gemma4_unified config (dense decoder + encoder-free embedders).""" @@ -1360,6 +1576,413 @@ def _build_gemma4_text(config: Gemma4Config) -> ir.Model: return task.build(module, config)["model"] +class TestGemma4NativeMoe: + @pytest.fixture(autouse=True) + def custom_runtime_opt_in(self): + from mobius._flags import override_flags + + with override_flags(gemma4_custom_geglu=True): + yield + + @staticmethod + def _quantization(**overrides): + return QuantizationConfig( + quant_method="olive", bits=4, group_size=16, quantize_moe=True, **overrides + ) + + @pytest.mark.parametrize("quantized", [False, True]) + @pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] + ) + def test_native_geglu_uses_raw_logits_and_separate_expert_norm( + self, quantized, dtype, tmp_path + ): + import numpy as np + + from mobius._build_context import build_context + from mobius._builder import _cast_module_dtype + from mobius._constants import GEMMA4_CUSTOM_GEGLU_METADATA_KEY + from mobius._execution_providers import get_ep + from mobius.tasks._gemma4 import Gemma4TextCausalLMTask + + config = _tiny_gemma4_text_moe_config( + dtype=dtype, quantization=self._quantization() if quantized else None + ) + if quantized and dtype == ir.DataType.FLOAT: + with pytest.raises(NotImplementedError, match="FLOAT16 or BFLOAT16"): + Gemma4CausalLMModel(config) + return + module = Gemma4CausalLMModel(config) + _cast_module_dtype(module, dtype) + with build_context(get_ep("cuda"), dtype): + graph = Gemma4TextCausalLMTask().build(module, config)["model"].graph + nodes = [node for node in graph if node.op_type == ("QMoE" if quantized else "MoE")] + assert len(nodes) == 2 + for node in nodes: + assert node.metadata_props[GEMMA4_CUSTOM_GEGLU_METADATA_KEY] == ( + "custom-cuda-geglu-v1" + ) + assert node.attributes["activation_type"].value == "geglu" + assert node.attributes["swiglu_fusion"].value == 1 + router = node.inputs[1] + while router.producer().op_type == "Reshape": + router = router.producer().inputs[0] + assert router.producer().op_type == "MatMul" + scaled = router.producer().inputs[0] + assert scaled.producer().op_type == "CastLike" + learned = scaled.producer().inputs[1] + norm = learned.producer().inputs[0].producer() + assert norm.op_type == "RMSNormalization" + expert_norm = node.inputs[0].producer() + assert expert_norm.op_type == "RMSNormalization" + assert norm.inputs[0] is expert_norm.inputs[0] + assert norm is not expert_norm + if quantized: + assert graph.initializers["model.layers.0.moe.fc1_experts_weights"].dtype == ( + ir.DataType.UINT8 + ) + assert graph.initializers["model.layers.0.moe.gate.proj.weight"].dtype == dtype + # The requirement belongs to the exported operator, not transient flags. + for initializer in graph.initializers.values(): + if initializer.const_value is None: + assert initializer.shape is not None and initializer.dtype is not None + initializer.const_value = ir.tensor( + np.zeros(list(initializer.shape), dtype=initializer.dtype.numpy()) + ) + ir.save(ir.Model(graph, ir_version=10), tmp_path / "native.onnx") + restored = ir.load(tmp_path / "native.onnx") + restored_nodes = [ + node for node in restored.graph if node.op_type == ("QMoE" if quantized else "MoE") + ] + assert len(restored_nodes) == 2 + assert all( + node.metadata_props[GEMMA4_CUSTOM_GEGLU_METADATA_KEY] == "custom-cuda-geglu-v1" + for node in restored_nodes + ) + + def test_default_cuda_is_portable_and_quantized_experts_require_opt_in(self): + from mobius._builder import build_from_module + from mobius._flags import override_flags + + with override_flags(gemma4_custom_geglu=False): + config = _tiny_gemma4_text_moe_config(dtype=ir.DataType.BFLOAT16) + package = build_from_module( + Gemma4CausalLMModel(config), config, execution_provider="cuda" + ) + assert not any(node.op_type in {"MoE", "QMoE"} for node in package["model"].graph) + assert any(node.op_type == "Softmax" for node in package["model"].graph) + with pytest.raises(NotImplementedError, match="MOBIUS_GEMMA4_CUSTOM_GEGLU"): + Gemma4CausalLMModel( + dataclasses.replace(config, quantization=self._quantization()) + ) + + def test_native_experts_fail_closed_on_unsupported_ep_and_layout(self): + config = _tiny_gemma4_text_moe_config( + dtype=ir.DataType.BFLOAT16, quantization=self._quantization() + ) + with pytest.raises(NotImplementedError, match="GEGLU-capable"): + _build_gemma4_text(config) + for hidden_act in ("gelu", "silu"): + with pytest.raises(NotImplementedError, match="tanh-GELU"): + Gemma4CausalLMModel(dataclasses.replace(config, hidden_act=hidden_act)) + quantization = dataclasses.replace( + config.quantization, + overrides={ + "model.language_model.layers.0.experts.down_proj": QuantizationOverride(bits=8) + }, + ) + with pytest.raises(NotImplementedError, match="INT4"): + Gemma4CausalLMModel(dataclasses.replace(config, quantization=quantization)) + + @pytest.mark.parametrize("multimodal", [False, True]) + @pytest.mark.parametrize("symmetric", [False, True]) + @pytest.mark.parametrize("component_plan", [False, True]) + def test_expert_payloads_and_module_layouts(self, multimodal, symmetric, component_plan): + from mobius._component_quantization import configure_component_quantization + from mobius.tasks._gemma4 import Gemma4Task + + quantization = self._quantization( + sym=symmetric, + overrides=( + { + "model.language_model.layers.0.self_attn.q_proj": QuantizationOverride( + bits=8 + ) + } + if component_plan + else {} + ), + ) + config = _tiny_gemma4_config( + dtype=ir.DataType.BFLOAT16, + hidden_act="gelu_pytorch_tanh", + quantization=quantization, + component_quantization={"decoder": quantization} if component_plan else None, + ) + prefix = "model.language_model.layers.0" if multimodal else "model.layers.0" + target = "decoder.model.layers.0.moe" if multimodal else "model.layers.0.moe" + gate = torch.arange(4 * 64 * 32).to(torch.uint8).reshape(4, 64, 32) + down = torch.arange(4 * 64 * 16).to(torch.uint8).reshape(4, 64, 16) + output_scale = torch.tensor([-2.0, 0.0, 0.5, 1.5], dtype=torch.bfloat16) + state = { + f"{prefix}.experts.gate_up_proj_qweight": gate, + f"{prefix}.experts.gate_up_proj_scales": torch.ones( + 4, 64, 4, dtype=torch.bfloat16 + ), + f"{prefix}.experts.down_proj_qweight": down, + f"{prefix}.experts.down_proj_scales": torch.ones(4, 64, 2, dtype=torch.bfloat16), + f"{prefix}.router.scale": torch.ones(64, dtype=torch.bfloat16), + f"{prefix}.router.proj.weight": torch.ones(4, 64, dtype=torch.bfloat16), + f"{prefix}.router.per_expert_scale": output_scale, + } + if not symmetric: + state[f"{prefix}.experts.gate_up_proj_qzeros"] = torch.full( + (4, 64, 2), 0x93, dtype=torch.uint8 + ) + state[f"{prefix}.experts.down_proj_qzeros"] = torch.full( + (4, 64, 1), 0x57, dtype=torch.uint8 + ) + module = (Gemma4Model if multimodal else Gemma4CausalLMModel)(config) + result = module.preprocess_weights(dict(state)) + torch.testing.assert_close(result[f"{target}.fc1_experts_weights"], gate) + torch.testing.assert_close(result[f"{target}.fc2_experts_weights"], down) + torch.testing.assert_close( + result[f"{target}.fc2_scales"], + torch.ones(4, 64, 2, dtype=torch.bfloat16) * output_scale[:, None, None], + ) + torch.testing.assert_close( + result[f"{target}.gate.scale"], torch.ones(64, dtype=torch.bfloat16) + ) + assert (f"{target}.fc1_experts_zero_points" in result) is not symmetric + if not symmetric: + assert torch.all(result[f"{target}.fc1_experts_zero_points"] == 0x93) + if component_plan and multimodal: + from mobius.components import Linear + + task = Gemma4Task() + manifest = configure_component_quantization(module, config, task) + normalized = normalize_component_quantized_weights( + result, + module, + config, + ("decoder", "vision_encoder", "embedding"), + manifest=manifest, + task=task, + ) + assert normalized.keys() == result.keys() + torch.testing.assert_close(normalized[f"{target}.fc1_experts_weights"], gate) + assert module.decoder.model.layers[0].self_attn.q_proj._bits == 8 + assert isinstance(module.decoder.model.layers[0].moe.gate.proj, Linear) + invalid = dict(state) + invalid[f"{prefix}.experts.down_proj_qweight"] = down[..., :-1] + with pytest.raises(ValueError, match="expected shape"): + module.preprocess_weights(invalid) + invalid = dict(state) + invalid.pop(f"{prefix}.experts.down_proj_scales") + with pytest.raises(ValueError, match="expected sidecars"): + module.preprocess_weights(invalid) + + @pytest.mark.parametrize("key", ["moe", "quantize_moe"]) + def test_quantization_metadata_opt_in(self, key): + quantization = QuantizationConfig.from_value( + {"quant_method": "olive", "bits": 4, "group_size": 16, key: True} + ) + assert quantization.quantize_moe + + def test_canonical_quantize_moe_takes_precedence_over_legacy_alias(self): + quantization = QuantizationConfig.from_value( + {"quant_method": "olive", "quantize_moe": False, "moe": True} + ) + assert not quantization.quantize_moe + + +class TestGemma4VisionStandardization: + def test_bf16_kernel_three_pool_matches_actual_hf_rounding(self): + import numpy as np + from transformers import Gemma4VisionConfig as HFVisionConfig + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4VisionPooler as HFVisionPooler, + ) + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import Gemma4VisionPooler + + hidden = torch.ones(1, 9, 4, dtype=torch.bfloat16) + hidden[:, -1] = 1.03125 + positions = torch.tensor([[(x, y) for y in range(3) for x in range(3)]]) + hf_pooler = HFVisionPooler(HFVisionConfig(hidden_size=4)) + expected, _ = hf_pooler( + hidden, positions, torch.zeros(1, 9, dtype=torch.bool), output_length=1 + ) + # Actual pinned HF rounds the FP32 mean to 1, then scales by sqrt(4). + torch.testing.assert_close(expected, torch.full((1, 1, 4), 2.0), rtol=0, atol=0) + rounded_coefficients = torch.full((1, 1, 9), 1 / 9, dtype=torch.bfloat16) + old_mean = (rounded_coefficients.float() @ hidden.float()).to(torch.bfloat16) + # Sensitive to premature coefficient cast; compare exactly in BF16. + torch.testing.assert_close( + old_mean, torch.full_like(old_mean, 1.0078125), rtol=0, atol=0 + ) + + builder, op, graph = create_test_builder() + features = create_test_input(builder, "features", [1, 9, 4], dtype=ir.DataType.FLOAT) + pos = create_test_input(builder, "positions", [1, 9, 2], dtype=ir.DataType.INT64) + output = Gemma4VisionPooler(4, kernel_size=3)( + op, op.Cast(features, to=ir.DataType.BFLOAT16), pos + ) + output.name = "pooled" + graph.outputs.append(output) + matmul = next(node for node in graph if node.op_type == "MatMul") + assert all(value.dtype == ir.DataType.FLOAT for value in matmul.inputs) + actual = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu").run( + {"features": hidden.float().numpy(), "positions": positions.numpy()} + )["pooled"] + np.testing.assert_array_equal(actual, expected.numpy()) + + @pytest.mark.parametrize("nonfinite", [float("nan"), float("inf"), float("-inf")]) + def test_bf16_pool_zeroes_padded_nonfinite_like_actual_hf(self, nonfinite): + import numpy as np + from transformers import Gemma4VisionConfig as HFVisionConfig + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4VisionPooler as HFVisionPooler, + ) + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import Gemma4VisionPooler + + hidden = torch.ones(1, 9, 4, dtype=torch.bfloat16) + hidden[:, -1] = nonfinite + positions = torch.tensor([[(x, y) for y in range(3) for x in range(3)]]) + positions[:, -1] = -1 + padding = (positions == -1).all(dim=-1) + expected, _ = HFVisionPooler(HFVisionConfig(hidden_size=4))( + hidden, positions, padding, output_length=1 + ) + assert torch.isfinite(expected).all() + rounded_mean = torch.tensor(8 / 9, dtype=torch.bfloat16).float() + torch.testing.assert_close( + expected, (rounded_mean * 2).expand(1, 1, 4), rtol=0, atol=0 + ) + + builder, op, graph = create_test_builder() + features = create_test_input(builder, "features", [1, 9, 4], dtype=ir.DataType.FLOAT) + pos = create_test_input(builder, "positions", [1, 9, 2], dtype=ir.DataType.INT64) + output = Gemma4VisionPooler(4, kernel_size=3)( + op, op.Cast(features, to=ir.DataType.BFLOAT16), pos + ) + output.name = "pooled" + graph.outputs.append(output) + actual = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu").run( + {"features": hidden.float().numpy(), "positions": positions.numpy()} + )["pooled"] + np.testing.assert_array_equal(actual, expected.numpy()) + + @pytest.mark.parametrize("missing", ["vision", "hidden_size"]) + def test_vision_encoder_requires_concrete_hidden_size(self, missing): + from mobius._configs import VisionConfig + from mobius.models.gemma4 import _Gemma4VisionEncoderModel + + with pytest.raises(ValueError, match=r"vision\.hidden_size"): + _Gemma4VisionEncoderModel( + _tiny_gemma4_config( + vision=None if missing == "vision" else VisionConfig(hidden_size=None) + ) + ) + + @pytest.mark.parametrize("standardize", [False, True]) + def test_actual_hf_vision_entrypoint_with_padding(self, standardize): + import numpy as np + from transformers import Gemma4VisionConfig as HFVisionConfig + from transformers.models.gemma4.modeling_gemma4 import Gemma4VisionModel + + from mobius import ModelPackage + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _Gemma4VisionEncoderModel + from mobius.tasks._gemma4 import Gemma4Task + + torch.manual_seed(42) + hf_config = HFVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=16, + patch_size=4, + pooling_kernel_size=2, + position_embedding_size=8, + standardize=standardize, + use_clipped_linears=False, + ) + hf_config._attn_implementation = "eager" + reference = Gemma4VisionModel(hf_config).eval() + if standardize: + reference.std_bias.copy_(torch.linspace(-5.0, 5.0, 32)) + reference.std_scale.copy_(torch.linspace(0.1, 2.0, 32)) + vision = dataclasses.replace( + _tiny_gemma4_config().vision, + pooling_kernel_size=2, + position_embedding_size=8, + standardize=standardize, + ) + config = _tiny_gemma4_config(vision=vision) + module = _Gemma4VisionEncoderModel(config) + package = ModelPackage({"vision_encoder": Gemma4Task()._build_vision(module, config)}) + projection = torch.randn(64, 32) * 0.1 + weights = { + "vision_tower." + name: value for name, value in reference.state_dict().items() + } + weights["embed_vision.embedding_projection.weight"] = projection + package.apply_weights(module.preprocess_weights(weights)) + positions = torch.tensor( + [[*[[x, y] for y in range(2) for x in range(4)], *[[-1, -1]] * 4]] + ) + pixels = torch.rand(1, 12, 48) + with torch.inference_mode(): + features = reference( + pixel_values=pixels, pixel_position_ids=positions + ).last_hidden_state + expected = torch.nn.functional.linear( + torch.nn.functional.rms_norm(features, (32,), eps=vision.norm_eps), projection + ) + actual = OnnxModelSession(package["vision_encoder"], device="cpu").run( + {"pixel_values": pixels.numpy(), "pixel_position_ids": positions.numpy()} + )["image_features"] + assert actual.shape == (2, 64) + np.testing.assert_allclose(actual, expected.numpy(), atol=1e-4, rtol=1e-4) + + @pytest.mark.parametrize("standardize", [False, True]) + def test_config_and_weight_mapping(self, standardize): + from types import SimpleNamespace + + from mobius._configs._vision_defaults import apply_vision_defaults + from mobius.models.gemma4 import _Gemma4VisionEncoderModel + + fields = {} + apply_vision_defaults( + SimpleNamespace(vision_config={"standardize": standardize}), + None, + "gemma4", + fields, + ) + assert fields["standardize"] is standardize + config = _tiny_gemma4_config( + vision=dataclasses.replace(_tiny_gemma4_config().vision, standardize=standardize) + ) + module = _Gemma4VisionEncoderModel(config) + assert (module.std_bias is not None) is standardize + if standardize: + bias, scale = torch.arange(32).float(), torch.ones(32) + state = {"vision_tower.std_bias": bias, "vision_tower.std_scale": scale} + assert module.preprocess_weights(state) == {"std_bias": bias, "std_scale": scale} + result = Gemma4Model(config).preprocess_weights( + {"model." + key: value for key, value in state.items()} + ) + assert result["vision_encoder.std_bias"] is bias + assert result["vision_encoder.std_scale"] is scale + + class TestGemma4MoeFallbackActivation: """The unfused MoE fallback must apply the EXACT ``hidden_act`` (ACT2FN). diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index 0beb78c1d..de5f69b2e 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -34,6 +34,7 @@ ) from mobius._model_package import ModelPackage from mobius._pipeline_contract import ( + declare_arbitrary_attention_mask, declare_component_presence, declare_optional_input, ) @@ -537,6 +538,13 @@ def _build_decoder( seq_len = ir.SymbolicDim("sequence_len") graph, builder = _make_graph(name="decoder") + if ( + config.use_bidirectional_attention == "vision" + and config.image_token_id is not None + ): + # A causal global bias still carries arbitrary padding positions. + # Prefix-length GQA fusion cannot preserve internal mask holes. + declare_arbitrary_attention_mask(graph) op = builder.op inputs_embeds = builder.input( @@ -643,7 +651,7 @@ def _build_vision( """Build vision encoder: pixel_values + pixel_position_ids -> image_features. Inputs: - - ``pixel_values [1, N, 3*P^2]``: one pre-patchified image with ``N`` + - ``pixel_values [1, N, 3*P^2]``: FP32 patches in ``[0, 1]``, with ``N`` patches (padded to ``max_soft_tokens``, typically 280). - ``pixel_position_ids [1, N, 2]``: (x, y) patch coordinates. Unused patch slots (from images smaller than the maximum) use @@ -666,7 +674,7 @@ def _build_vision( pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[1, num_patches, pixel_dim], ) pixel_position_ids = builder.input( diff --git a/tests/build_graph/core_test.py b/tests/build_graph/core_test.py index a1fa025b3..a7ea1b5f2 100644 --- a/tests/build_graph/core_test.py +++ b/tests/build_graph/core_test.py @@ -797,8 +797,8 @@ def test_multimodal_encoder_inputs_match_model_dtype(self, dtype_str): "dtype_str", ["f16", "bf16"], ) - def test_gemma4_encoder_inputs_match_model_dtype(self, dtype_str): - """Gemma4 vision/audio encoder inputs match the model's compute dtype.""" + def test_gemma4_encoder_input_dtypes(self, dtype_str): + """Vision takes FP32 processor pixels; audio keeps the compute dtype.""" from mobius._configs import Gemma4AudioConfig, Gemma4Config config = Gemma4Config( @@ -846,11 +846,11 @@ def test_gemma4_encoder_inputs_match_model_dtype(self, dtype_str): task = get_task("gemma4") pkg = task.build(module, config) - # Vision encoder pixel_values must match config.dtype + # Normalize processor pixels in FP32 before casting to the projection dtype. vision_model = pkg["vision_encoder"] pv_input = vision_model.graph.inputs[0] assert pv_input.name == "pixel_values" - assert pv_input.dtype == config.dtype + assert pv_input.dtype == ir.DataType.FLOAT # Audio encoder input_features must match config.dtype audio_model = pkg["audio_encoder"] diff --git a/tests/integration/gemma4_test.py b/tests/integration/gemma4_test.py index f97c3be61..107cd7caf 100644 --- a/tests/integration/gemma4_test.py +++ b/tests/integration/gemma4_test.py @@ -558,47 +558,68 @@ def test_gemma4_unified_12b_multimodal_prefill(): @pytest.mark.integration -def test_gemma4_bidirectional_mask_parity(): +@pytest.mark.parametrize( + ("static", "case"), + [ + (static, case) + for static in (False, True) + for case in ("text", "image", "two_images", "long_image", "padding") + ] + + [(False, "internal_hole")], +) +def test_gemma4_bidirectional_mask_parity(static, case): """Mobius's vision-block attention bias matches HuggingFace exactly. - Larger Gemma 4 models use ``use_bidirectional_attention="vision"``: a - contiguous run of image/audio placeholder tokens attends bidirectionally - within its block (on BOTH full and sliding layers), while text stays - causal. This test compares mobius's ``create_attention_bias`` output - (with ``block_sequence_ids`` from ``_compute_block_sequence_ids``) against - HuggingFace's real ``create_causal_mask`` / ``create_sliding_window_causal_mask`` - for the actual ``gemma-4-26b-a4b-it`` config. It needs only the config - (no weights), so it is cheap and deterministic. + Global layers remain causal; local layers use + ``(causal OR same-image) AND left-window AND padding``. Compare the actual + HF Gemma4 entry point with a local config, without downloads or weights. """ import onnx_ir as ir import torch - from transformers.masking_utils import ( - create_causal_mask, - create_sliding_window_causal_mask, + from transformers import Gemma4TextConfig + from transformers.models.gemma4.modeling_gemma4 import ( + create_masks_for_vision_model, + get_block_sequence_ids_for_mask, ) from mobius._testing import create_test_builder, create_test_input - from mobius.components._common import create_attention_bias + from mobius.components import create_attention_bias, create_static_cache_attention_bias from mobius.models.gemma4 import _compute_block_sequence_ids from mobius.tasks._base import _make_graph, _make_model - model_id = "google/gemma-4-26b-a4b-it" - if not _model_accessible(model_id): - pytest.skip(f"{model_id} not accessible (requires HuggingFace authentication)") - - hf_config = transformers.AutoConfig.from_pretrained(model_id) - text_cfg = hf_config.text_config + text_cfg = Gemma4TextConfig( + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + sliding_window=1024, + layer_types=["sliding_attention", "full_attention"], + use_bidirectional_attention="vision", + ) text_cfg._attn_implementation = "eager" # always returns a dense float mask assert text_cfg.use_bidirectional_attention == "vision" - image_token_id = hf_config.image_token_id - - # Synthetic layout: text, image block, text, image block, text. - seq_len = 12 - input_ids = np.full((1, seq_len), 5, dtype=np.int64) - input_ids[0, 2:7] = image_token_id - input_ids[0, 9:11] = image_token_id + image_token_id = 255 + types = { + "text": [0] * 7, + "image": [0, 1, 1, 1, 0], + "two_images": [0, 1, 1, 0, 1, 1, 0], + "long_image": [0] + [1] * 1120 + [0], + "padding": [0, 1, 1, 1, 0, 0], + "internal_hole": [1] * 1025, + }[case] + token_types = torch.tensor([types]) + seq_len = len(types) + input_ids = np.where(token_types.numpy() == 1, image_token_id, 5).astype(np.int64) attention_mask = np.ones((1, seq_len), dtype=np.int64) position_ids = np.arange(seq_len, dtype=np.int64)[None, :] + if case == "padding": + attention_mask[0, -2:] = 0 + elif case == "internal_hole": + # Physical q=1024,k=0 is outside W=1024 despite an invalid slot at k=1. + # The static primitive has prefix validity, not arbitrary hole validity; + # retain its ten existing comparisons and test this case dynamically. + attention_mask[0, 1] = 0 # Mobius block_sequence_ids from input_ids. graph, builder = _make_graph() @@ -609,10 +630,9 @@ def test_gemma4_bidirectional_mask_parity(): block_ids = OnnxModelSession(_make_model(graph), device="cpu").run( {"input_ids": input_ids} )["bsid"] - # Two contiguous image runs separated by text -> groups 0 and 1. np.testing.assert_array_equal( - block_ids[0], - np.array([-1, -1, 0, 0, 0, 0, 0, -1, -1, 1, 1, -1], dtype=np.int64), + block_ids, + get_block_sequence_ids_for_mask(token_types, torch.device("cpu")).numpy(), ) def mobius_attended(sliding_window): @@ -620,14 +640,22 @@ def mobius_attended(sliding_window): ii = create_test_input(b, "input_ids", [1, seq_len], dtype=ir.DataType.INT64) am = create_test_input(b, "attention_mask", [1, seq_len], dtype=ir.DataType.INT64) bk = create_test_input(b, "block_sequence_ids", [1, seq_len], dtype=ir.DataType.INT64) - bias = create_attention_bias( - bop, - ii, - am, + arguments = dict( sliding_window=sliding_window, dtype=ir.DataType.FLOAT, - block_sequence_ids=bk, + block_sequence_ids=bk if sliding_window is not None else None, ) + if static: + bias = create_static_cache_attention_bias( + bop, + write_indices=bop.Constant(value_ints=[0]), + seq_len=bop.Shape(ii, start=1, end=2), + nonpad_kv_seqlen=bop.ReduceSum(am, [1], keepdims=0), + max_seq_len=seq_len, + **arguments, + ) + else: + bias = create_attention_bias(bop, ii, am, **arguments) bias.name = "bias" g.outputs.append(bias) out = OnnxModelSession(ir.Model(g, ir_version=10), device="cpu").run( @@ -641,30 +669,26 @@ def mobius_attended(sliding_window): inputs_embeds = torch.zeros(1, seq_len, 8) blk = torch.from_numpy(block_ids) - hf_full = create_causal_mask( - text_cfg, - inputs_embeds, - torch.from_numpy(attention_mask), - None, - torch.from_numpy(position_ids), - block_sequence_ids=blk, - ) - hf_sliding = create_sliding_window_causal_mask( + hf_masks = create_masks_for_vision_model( text_cfg, inputs_embeds, torch.from_numpy(attention_mask), None, torch.from_numpy(position_ids), - block_sequence_ids=blk, + blk, ) def hf_attended(mask): m = mask[0, 0].numpy() return m if m.dtype == bool else (m > -1e30) - # Full-attention layers: causal OR same-block, AND padding. - np.testing.assert_array_equal(mobius_attended(None), hf_attended(hf_full)) - # Sliding layers: (causal AND window) OR same-block, AND padding. np.testing.assert_array_equal( - mobius_attended(text_cfg.sliding_window), hf_attended(hf_sliding) + mobius_attended(None), hf_attended(hf_masks["full_attention"]) ) + actual_local = mobius_attended(text_cfg.sliding_window) + hf_local = hf_attended(hf_masks["sliding_attention"]) + np.testing.assert_array_equal(actual_local, hf_local) + if case == "internal_hole": + assert not hf_local[1024, 0] + assert not actual_local[1024, 0] + assert actual_local[1024, 2] diff --git a/tests/static_cache_bias_parity_test.py b/tests/static_cache_bias_parity_test.py index ffef477db..861fc52e2 100644 --- a/tests/static_cache_bias_parity_test.py +++ b/tests/static_cache_bias_parity_test.py @@ -187,7 +187,7 @@ def _dense_reference( ``full_*_cache`` are the *post-scatter* caches (current chunk already written). For each (batch, query, head) we build the boolean mask in the exact rule order ``create_static_cache_attention_bias`` uses — causal, - AND sliding window, OR block overlay, AND padding validity — gather the + OR block overlay, AND sliding window, AND padding validity — gather the valid cache slots (with GQA head sharing) and softmax over them. Every query row keeps its own diagonal slot valid (``q_abs[t] == write+t`` is causal, within any window, and ``< nonpad``), so a fully-masked row is @@ -218,10 +218,10 @@ def _dense_reference( for t in range(query_len): mask = q_abs[t] >= kv_slots # causal - if sliding_window is not None: - mask &= (q_abs[t] - kv_slots) < sliding_window # local window if block_sequence_ids is not None and block_sequence_ids[b, t] >= 0: mask |= kv_group == block_sequence_ids[b, t] # bidirectional block + if sliding_window is not None: + mask &= (q_abs[t] - kv_slots) < sliding_window # local window mask &= kv_slots < npad # padding validity idx = np.nonzero(mask)[0] # The bias-MEA contract keeps every query row's own diagonal slot @@ -371,7 +371,7 @@ def test_prefill_chunk_bias_matches_dense_reference(): """Prefill chunk (write=0, S_q=N, nonpad=N) matches the dense reference. The whole chunk is the valid region; query row ``i`` attends slots within - ``[max(0, i - sliding_window + 1), i]``, OR-ed with same-block tokens. + ``causal OR same-block``, constrained by the strict left window boundary. """ query_len = 6 onnx_attn, ref = _run_and_compare( From 086aeb5ea2b470ee0e482cd6f5ae1e5a3347cce1 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Sat, 24 Oct 2026 00:03:15 +0000 Subject: [PATCH 2/4] Gemma4: mask unselected expert overflow before weighting 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 Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 --- docs/gemma4-runtime.md | 6 ++++ src/mobius/models/gemma4.py | 19 +++++++--- src/mobius/models/gemma4_test.py | 59 ++++++++++++++++++++++++++++++++ 3 files changed, 80 insertions(+), 4 deletions(-) diff --git a/docs/gemma4-runtime.md b/docs/gemma4-runtime.md index 4465b27a6..0acd1f66c 100644 --- a/docs/gemma4-runtime.md +++ b/docs/gemma4-runtime.md @@ -46,6 +46,12 @@ exact-activation fallback; quantized GEGLU experts fail at construction with an explicit requirement message. Opt-in does not enable native GEGLU on non-CUDA providers. Tanh-GELU must not be replaced by a sigmoid-based SwiGLU approximation. +The portable fallback evaluates experts densely, but selects out unchosen expert +outputs before weighting. This prevents an overflowing unselected expert from +contaminating the result through `0 * Inf` or `0 * NaN`; nonfinite selected +experts still propagate. This safeguard is not full-checkpoint numerical or +performance qualification. + Each emitted native `MoE`/`QMoE` operator persists `mobius.gemma4.required_runtime_feature="custom-cuda-geglu-v1"` in its ONNX metadata. This declares the experimental CUDA `activation_type="geglu"`, diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index f41acb0cf..1b42629ee 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -2016,10 +2016,8 @@ def _dispatch_moe_fallback( count), unlike a per-expert static unroll which produces O(E x K) nodes (~6k per layer at E=128/K=8 — impractical for large MoE checkpoints). - Standard ``MatMul``/``Einsum`` accumulate in fp32 on CUDA even for fp16 - inputs, so this path is numerically faithful to HuggingFace — unlike the - released-ORT fused ``com.microsoft::MoE`` kernel, which fp16-accumulates - the expert GEMMs and loses accuracy at large hidden size. + This preserves the reference activation and routing boundaries, but + kernel-dependent accumulation still requires numerical qualification. Args: normed_flat: [T, H] — pre-normed input for the experts (T = B*S). @@ -2053,6 +2051,14 @@ def _dispatch_moe_fallback( op.ConstantOfShape(dense_shape, value=ir.tensor(np.zeros(1, dtype=np.float32))), top_weights, ) + selected = op.ScatterElements( + op.Cast(dense_w, to=ir.DataType.BOOL), + top_indices, + op.ConstantOfShape( + op.Shape(top_indices), value=ir.tensor(np.ones(1, dtype=np.bool_)) + ), + axis=1, + ) # [T, E], independent of routing weights or learned expert scales dense_w = op.ScatterElements(dense_w, top_indices, top_weights, axis=1) # [T, E] # fc1_experts_weights [E, 2*moe_inter, H] stores gate rows then up rows. @@ -2073,6 +2079,11 @@ def _dispatch_moe_fallback( # Down projection: [T, E, moe_inter] x [E, H, moe_inter] -> [T, E, H] expert_out = op.Einsum(inter, self.fc2_experts_weights, equation="tei,ehi->teh") + # Sparse HF dispatch never evaluates unselected experts. Multiplying + # their nonfinite outputs by zero would contaminate selected results. + expert_out = op.Where( + op.Unsqueeze(selected, [2]), expert_out, op.CastLike(0.0, expert_out) + ) # Weight by the dense routing vector and sum over experts. # HF rounds each FP32-weighted contribution back to the expert dtype. weighted = op.CastLike( diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index a5709e2ee..8a3dcb6ad 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -2021,6 +2021,65 @@ def test_gelu_pytorch_tanh_uses_tanh_approx(self): assert acts, "MoE fallback must emit an activation op" assert all(op == "Gelu" and approx == "tanh" for op, approx in acts), acts + @pytest.mark.parametrize("selected_expert", [0, 1]) + @pytest.mark.parametrize("expert_scale", [0.0, 1.0]) + def test_nonfinite_expert_is_masked_by_selection_not_weight( + self, selected_expert, expert_scale + ): + import numpy as np + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import Gemma4DecoderLayer + + class FallbackOnly(Gemma4DecoderLayer): + def forward(self, op, normed_flat, router_probs): + self.router(op, normed_flat) + return self._dispatch_moe_fallback(op, normed_flat, router_probs) + + config = _tiny_gemma4_text_moe_config( + dtype=ir.DataType.FLOAT, + num_local_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=16, + ) + module = FallbackOnly(config, 0) + fc1 = np.zeros((2, 32, 64), dtype=np.float32) + fc1[1] = np.finfo(np.float32).max + fc2 = np.zeros((2, 64, 16), dtype=np.float32) + fc2[1] = 1 + module.fc1_experts_weights.const_value = ir.tensor(fc1) + module.fc2_experts_weights.const_value = ir.tensor(fc2) + module.router.per_expert_scale.const_value = ir.tensor( + np.full(2, expert_scale, dtype=np.float32) + ) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "normed_flat", [1, 64]) + probs = create_test_input(builder, "router_probs", [1, 2]) + output = module(op, x, probs) + output.name = "fallback" + output.shape = ir.Shape([1, 64]) + output.dtype = ir.DataType.FLOAT + graph.outputs.append(output) + graph = ir.convenience.extract(graph, [x, probs], [output]) + session = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu") + try: + actual = session.run( + { + "normed_flat": np.ones((1, 64), dtype=np.float32), + "router_probs": np.eye(2, dtype=np.float32)[ + selected_expert : selected_expert + 1 + ], + } + )["fallback"] + finally: + session.close() + if selected_expert == 0: + np.testing.assert_array_equal(actual, np.zeros((1, 64), dtype=np.float32)) + else: + # Even a zero learned scale must not hide a selected expert's NaNs. + assert not np.isfinite(actual).any() + class TestGemma4OutputLayerIndices: """``output_layer_indices`` taps the post-final-norm hidden as ``hidden_states.{idx}``. From fd4513b5cc2e4692a2cde51fd8fdd57c4bb0878d Mon Sep 17 00:00:00 2001 From: titaiwang Date: Sat, 24 Oct 2026 02:26:11 +0000 Subject: [PATCH 3/4] Clarify Gemma4 regression-test override and flag restoration 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 --- src/mobius/_flags_test.py | 18 ++++++++++-------- src/mobius/models/gemma4_test.py | 17 ++++++++++++++--- 2 files changed, 24 insertions(+), 11 deletions(-) diff --git a/src/mobius/_flags_test.py b/src/mobius/_flags_test.py index 38ae2fc47..66a68a0f6 100644 --- a/src/mobius/_flags_test.py +++ b/src/mobius/_flags_test.py @@ -93,15 +93,17 @@ def test_nested_overrides(self): def test_custom_geglu_override_is_scoped_even_on_exception(self): original = _flags.flags.gemma4_custom_geglu - with _flags.override_flags(gemma4_custom_geglu=False): - assert not _flags.flags.gemma4_custom_geglu - with pytest.raises(RuntimeError, match="test restoration"): + try: + with _flags.override_flags(gemma4_custom_geglu=False): + assert not _flags.flags.gemma4_custom_geglu + with pytest.raises(RuntimeError, match="test restoration"): + assert not _flags.flags.gemma4_custom_geglu + with _flags.override_flags(gemma4_custom_geglu=True): + assert _flags.flags.gemma4_custom_geglu + raise RuntimeError("test restoration") assert not _flags.flags.gemma4_custom_geglu - with _flags.override_flags(gemma4_custom_geglu=True): - assert _flags.flags.gemma4_custom_geglu - raise RuntimeError("test restoration") - assert not _flags.flags.gemma4_custom_geglu - assert _flags.flags.gemma4_custom_geglu is original + finally: + assert _flags.flags.gemma4_custom_geglu is original def test_unknown_flag_raises_value_error(self): with ( diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index 8a3dcb6ad..661ea8c4a 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -2033,9 +2033,20 @@ def test_nonfinite_expert_is_masked_by_selection_not_weight( from mobius.models.gemma4 import Gemma4DecoderLayer class FallbackOnly(Gemma4DecoderLayer): - def forward(self, op, normed_flat, router_probs): - self.router(op, normed_flat) - return self._dispatch_moe_fallback(op, normed_flat, router_probs) + def forward( + self, + op, + hidden_states, + attention_bias, + position_embeddings=None, + shared_kv_states=None, + per_layer_input=None, + past_key_value=None, + is_causal=1, + static_kv_seqlen=None, + ): + self.router(op, hidden_states) + return self._dispatch_moe_fallback(op, hidden_states, attention_bias) config = _tiny_gemma4_text_moe_config( dtype=ir.DataType.FLOAT, From c3b2ec76131c1db166bab8cbd506f033d375448b Mon Sep 17 00:00:00 2001 From: titaiwang Date: Sun, 25 Oct 2026 01:19:53 +0000 Subject: [PATCH 4/4] Gemma4: preserve arbitrary masks in text-task exports 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 --- src/mobius/models/gemma4_test.py | 13 +++++++------ src/mobius/tasks/_gemma4.py | 9 +++++++++ 2 files changed, 16 insertions(+), 6 deletions(-) diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index 661ea8c4a..809585867 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -1268,13 +1268,14 @@ def test_no_block_sequence_ids_when_causal(self): assert "input_ids" not in dec_inputs @pytest.mark.parametrize("mask", [[1, 0, 1, 1], [1, 1, 1, 0]]) - def test_cuda_export_preserves_global_cached_mask_holes(self, mask): + @pytest.mark.parametrize("text_task", [False, True], ids=["split", "text"]) + def test_cuda_export_preserves_global_cached_mask_holes(self, mask, text_task): import numpy as np from mobius import build_from_module from mobius._pipeline_contract import requires_arbitrary_attention_mask from mobius._testing.ort_inference import OnnxModelSession - from mobius.tasks._gemma4 import Gemma4Task + from mobius.tasks._gemma4 import Gemma4Task, Gemma4TextCausalLMTask config = _tiny_gemma4_config( dtype=ir.DataType.FLOAT16, @@ -1282,10 +1283,10 @@ def test_cuda_export_preserves_global_cached_mask_holes(self, mask): use_bidirectional_attention="vision", image_token_id=255, ) - package = build_from_module( - Gemma4Model(config), config, task=Gemma4Task(), execution_provider="cuda" - ) - decoder = package["decoder"] + module = Gemma4CausalLMModel(config) if text_task else Gemma4Model(config) + task = Gemma4TextCausalLMTask() if text_task else Gemma4Task() + package = build_from_module(module, config, task=task, execution_provider="cuda") + decoder = package["model" if text_task else "decoder"] assert requires_arbitrary_attention_mask(decoder.graph) assert not any(node.op_type == "GroupQueryAttention" for node in decoder.graph) global_attention = next( diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index de5f69b2e..dbd1f7267 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -273,6 +273,9 @@ class Gemma4TextCausalLMTask(ModelTask): Supports ``static_cache=True`` for pre-allocated TensorScatter-based KV cache (requires ORT ≥ 1.25.0). + Configs retaining vision-overlay attention preserve explicit masks during + optimization; stripped text-only configs remain eligible for GQA fusion. + Inputs (dynamic cache): - input_ids: [batch, sequence_len] INT64 - attention_mask: [batch, past_seq_len + seq_len] INT64 @@ -327,6 +330,12 @@ def build( seq_len = ir.SymbolicDim("sequence_len") graph, builder = _make_graph() + if ( + config.use_bidirectional_attention == "vision" + and config.image_token_id is not None + ): + # Prefix-length GQA fusion cannot preserve internal mask holes. + declare_arbitrary_attention_mask(graph) op = builder.op input_ids = builder.input(