diff --git a/docs/gemma4-runtime.md b/docs/gemma4-runtime.md new file mode 100644 index 000000000..0acd1f66c --- /dev/null +++ b/docs/gemma4-runtime.md @@ -0,0 +1,183 @@ +# 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. + +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"`, +`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..66a68a0f6 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,20 @@ 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 + 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 + finally: + 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..1b42629ee 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) @@ -1824,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). @@ -1861,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. @@ -1874,15 +2072,24 @@ 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] # 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. - 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 +2191,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 +2474,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 +2488,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 +2508,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 +2639,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 +2720,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 +2876,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 +2989,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 +3016,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 +3040,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 +3753,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 +3811,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..809585867 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,82 @@ 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]]) + @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, Gemma4TextCausalLMTask + + config = _tiny_gemma4_config( + dtype=ir.DataType.FLOAT16, + enable_moe_block=False, + use_bidirectional_attention="vision", + image_token_id=255, + ) + 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( + 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 +1577,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). @@ -1398,6 +2022,76 @@ 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, + 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, + 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}``. diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index 0beb78c1d..dbd1f7267 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, ) @@ -272,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 @@ -326,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( @@ -537,6 +547,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 +660,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 +683,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(