diff --git a/docs/ep_quickstart.md b/docs/ep_quickstart.md index ae1fbb076..7f4e645d6 100644 --- a/docs/ep_quickstart.md +++ b/docs/ep_quickstart.md @@ -100,7 +100,9 @@ print(sorted(ep_registry)) # Inspect an EP's capabilities caps = get_ep("cuda") print(caps.gqa_dtypes) # frozenset({FLOAT16, BFLOAT16}) -print(caps.qkv_pack_dtypes) # frozenset({FLOAT, FLOAT16, BFLOAT16}) +print(caps.qkv_pack_dtypes) # frozenset({FLOAT, FLOAT16, BFLOAT16}); legacy GQA-input packing +print(caps.qkv_projection_pack_dtypes) # frozenset(); experimental float projection packing is off +print(caps.matmul_nbits_qkv_pack_dtypes) # frozenset(); experimental INT4 packing is off print(caps.supports_fused_rope) # True print(caps.enable_graph_capture) # True print(caps.provider_options) # {'enable_skip_layer_norm_strict_mode': '1'} @@ -132,6 +134,15 @@ pkg = mobius.build("meta-llama/Llama-3.2-1B", Unrecognised EP names raise `ValueError` during build-time validation, before graph construction or optimization starts — safe-fail by default. +Experimental float `PackQKVProjections` and INT4 projection packing require +independent `qkv_projection_pack_dtypes` and `matmul_nbits_qkv_pack_dtypes` +capability overrides, respectively; both are **default-off on every built-in EP**. +The existing float GQA-input `PackQKV` stage retains its `qkv_pack_dtypes` gate. +Neither INT4 nor CUDA FP16/BF16 float projection-packing numerical accuracy is +qualified by this branch. Read the +[numerical limitations and opt-in example](execution_providers.md#projection-packing-and-numerical-limitations) +before enabling it; activation/component precision flags do not opt in. + --- ## 5. Common EP configurations at a glance @@ -141,7 +152,7 @@ graph construction or optimization starts — safe-fail by default. | Portable ONNX (maximum compatibility) | `"default"` | any | No EP-specific vendor fusions (e.g. no GQA/PackQKV) | | Strict standard ONNX (no custom ops at all) | `"onnx-standard"` | any | All `com.microsoft` ops expanded via InlinePass; safe for non-ORT runtimes | | ORT CPU inference | `"cpu"` | `"f32"` | GQA fusion for FP32 | -| NVIDIA GPU | `"cuda"` | `"f16"` or `"bf16"` | GQA + SkipNorm + PackQKV | +| NVIDIA GPU | `"cuda"` | `"f16"` or `"bf16"` | GQA + SkipNorm + existing float PackQKV plus new default float projection packing (also for masked decoders); experimental INT4 projection packing is off | | Windows GPU (DirectX) | `"dml"` | `"f16"` | RoPE lowered separately | | Browser / WebAssembly | `"webgpu"` | `"f16"` or `"f32"` | Shape ops replaced | | NVIDIA TensorRT-RTX | `"trt-rtx"` | `"f16"` or `"bf16"` | SkipLayerNorm expanded | diff --git a/docs/execution_providers.md b/docs/execution_providers.md index 25297f7ac..313cc8967 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -100,6 +100,73 @@ ops to fused equivalents: These are applied during **Stage 2: Fusion** of the optimization pipeline. +#### Projection packing and numerical limitations + +`PackQKVProjections` combines compatible projections into one GEMM and splits +the result **before normalization, RoPE, attention, or caches**, without changing +their contracts. It supports Q/K/V and Q/K with raw K reused for V. + +| Stage | Roles | Enablement | +|---|---|---| +| Existing float `PackQKV` GQA-input fusion | `decoder` | Unchanged | +| Additional float `PackQKVProjections` | `decoder`, `masked-decoder` | **Experimental; disabled on every built-in EP**; explicit `qkv_projection_pack_dtypes` opt-in | +| INT4 projection packing | `decoder`, `masked-decoder` | **Experimental; disabled on every built-in EP**, including CUDA FP16/BF16 | + +Experimental float and INT4 projection packing use independent process-local +capabilities. Stock `--ep cuda` enables neither new stage; the existing +`qkv_pack_dtypes`-gated GQA-input packing is unchanged. Set only the capabilities +you intend to evaluate (the example opts into both): + +```python +import dataclasses + +import onnx_ir as ir +from mobius import get_ep, register_ep + +register_ep( + dataclasses.replace( + get_ep("cuda"), + qkv_projection_pack_dtypes=frozenset({ir.DataType.FLOAT16}), + matmul_nbits_qkv_pack_dtypes=frozenset({ir.DataType.FLOAT16}), + ), + overwrite=True, +) +``` + +Both opt-ins are independent of legacy float packing, activation/storage dtype, +component precision, and FP8 KV caches. Neither byte-exact payload tests nor +CPU FLOAT execution tests qualify CUDA FP16/BF16 numerical parity. Packing never +dequantizes, requantizes, or promotes weights. +`_qkv_projection_packing.py` implements both float and INT4. + +Packing requires static, non-overridable, compatible parameters and opset >= 13; +ambiguous or overlapping groups and weights with outside consumers are skipped. +INT4 additionally requires a shared rank-3 activation, matching attributes except +output width, canonical UINT8 blocked weights, same-dtype scales, and consistently +absent or packed UINT8 zero points. Flat/row-major sidecars are supported; +block size must be a power of two >= 16 and `accuracy_level` omitted or zero. +Mixed bit widths, extra operands, and caller-owned function bodies are skipped. + +Parameter Concats retain checkpoint names until binding, then initializer +Concats and matrix Transposes fold lazily without changing dtype. Missing leaves +are checked even with `fold_constants=False`. Generic Reshape/Transpose layout +folding is not part of this pipeline; those transformations remain explicit. +Overridable inputs remain untouched. + +**The INT4 numerical gate remains failed and unwaived:** large-dimension CUDA +prefill/cached-decode logits exceeded `rtol=atol=1e-2`. Changing GEMM width can +change floating-point results despite byte-exact weights. MMLU/generation results +do not clear that gate; float projection packing has no separate numerical +qualification here. No production accuracy or speed improvement is claimed. + +For `ORT_FPA_INTB_GEMM=1`, keep exporter or ORT constant folding enabled so +parameter Concats do not block runtime prepacking; behavior is runtime-dependent. +Packed float weights retain adapter-restoration metadata and lazy backing-file +ownership after pruning. `ModelPackage.save()` rejects known same-path, hardlink, +or symlink aliases to backing files across model, shard, policy, adapter, metadata, +and MTP destinations. This is **not full TOCTOU protection**; save to a separate +output directory. + ### Tier 3 — EP-specific constraints (required, correctness) Some EPs cannot execute certain ops. These are handled by **lowering rules** @@ -126,7 +193,8 @@ Every EP is fully described by a single `EpCapabilities` entry in the class EpCapabilities: name: str gqa_dtypes: frozenset[ir.DataType] # dtypes where GQA fusion fires - qkv_pack_dtypes: frozenset[ir.DataType] # dtypes where PackQKV fusion fires + qkv_pack_dtypes: frozenset[ir.DataType] # legacy GQA-input packing only + qkv_projection_pack_dtypes: frozenset[ir.DataType] # experimental float; default empty supports_fused_rope: bool = True # False → SeparateRoPE + UnpackQKV supports_skip_layer_norm: bool = True # False → InlinePass expansion supports_fused_matmul: bool = True # False → Transpose + MatMul via InlinePass diff --git a/docs/gemma4-runtime.md b/docs/gemma4-runtime.md new file mode 100644 index 000000000..6f5762ab1 --- /dev/null +++ b/docs/gemma4-runtime.md @@ -0,0 +1,186 @@ +# 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 +`pkg.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. + +The legacy `mobius.qkv_projection_sources` key is accepted when reading older +exports; new exports write `pkg.mobius.qkv_projection_sources`. + +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 5e9c882a7..f0fefacb7 100644 --- a/docs/model-catalog.md +++ b/docs/model-catalog.md @@ -199,6 +199,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/_component_quantization.py b/src/mobius/_component_quantization.py index b41928916..4bdaa9341 100644 --- a/src/mobius/_component_quantization.py +++ b/src/mobius/_component_quantization.py @@ -754,7 +754,7 @@ def validate_quantized_component_bindings( models: Mapping[str, ir.Model], config: BaseModelConfig, ) -> None: - """Require every affine quantized op input to carry a bound value.""" + """Require bound affine parameters, including deferred axis-0 Concat leaves.""" quantized_input_slots = { "MatMulNBits": (1, 2, 3), "GatherBlockQuantized": (0, 2, 3), @@ -768,13 +768,30 @@ def validate_quantized_component_bindings( if index >= len(node.inputs): continue value = node.inputs[index] - if value is None or value.producer() is not None: + if value is None: continue - if value.const_value is None: - raise ValueError( - f"Quantized component {component!r} has unbound " - f"{node.op_type} parameter {value.name!r}" - ) + pending = [value] + visited: set[ir.Value] = set() + while pending: + parameter = pending.pop() + if parameter in visited: + continue + visited.add(parameter) + producer = parameter.producer() + if producer is not None: + if ( + producer.domain in ("", "ai.onnx") + and producer.op_type == "Concat" + and not producer.overload + and producer.attributes.get_int("axis", None) == 0 + ): + pending.extend(v for v in producer.inputs if v is not None) + continue + if parameter.const_value is None: + raise ValueError( + f"Quantized component {component!r} has unbound " + f"{node.op_type} parameter {parameter.name!r}" + ) def attach_hf_component_sources( 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 c80cc4780..589af1a31 100644 --- a/src/mobius/_constants.py +++ b/src/mobius/_constants.py @@ -9,6 +9,22 @@ # Separated into its own module to avoid circular imports between # tasks, components, and the top-level package. OPSET_VERSION = 24 +QKV_PROJECTION_SOURCES_METADATA_KEY = "pkg.mobius.qkv_projection_sources" +"""Versioned source names and widths retained on projection-packing Split nodes.""" + +LEGACY_QKV_PROJECTION_SOURCES_METADATA_KEY = "mobius.qkv_projection_sources" +"""Read-only key for restoring packed adapters from pre-namespace preview exports.""" + +LAZY_SOURCE_FILES_META_KEY = "mobius.lazy_source_files" +"""Backing-file paths retained by lazy tensors and their model in ``.meta``. + +Folding can remove an ExternalTensor from the initializer registry while a lazy +callback still needs its file. ModelPackage checks these paths before saving so +it cannot overwrite that backing file before the callback reads it. This is +transient Python bookkeeping, not serialized ``metadata_props``. +""" + +GEMMA4_CUSTOM_GEGLU_METADATA_KEY = "mobius.gemma4.required_runtime_feature" # --------------------------------------------------------------------------- # Static KV cache ABI diff --git a/src/mobius/_execution_providers.py b/src/mobius/_execution_providers.py index 21e88b42e..88f9666fe 100644 --- a/src/mobius/_execution_providers.py +++ b/src/mobius/_execution_providers.py @@ -53,11 +53,19 @@ class EpCapabilities: Attributes: name: Canonical EP name (e.g. ``"cuda"``). gqa_dtypes: dtypes for which GroupQueryAttention fusion is supported. - qkv_pack_dtypes: dtypes for which QKV weight packing for - GroupQueryAttention is supported (PackQKV fusion). Set to an - empty frozenset for EPs that do not support packed QKV inputs - (e.g. DML, which always unpacks via UnpackQKV in the lowering - stage). + qkv_pack_dtypes: Dtypes enabling the existing decoder-only GQA-input + packing stage (PackQKV). Does not enable the separate projection-only + float or INT4 stages. Set to an empty frozenset for EPs that do not + support packed QKV inputs (e.g. DML, which always unpacks via + UnpackQKV in the lowering stage). + matmul_nbits_qkv_pack_dtypes: Activation dtypes for projection-only INT4 + packing, which preserves separate normalization/attention branches. + Experimental: empty on every built-in EP, including CUDA. Explicit + opt-in does not waive the unresolved large-model numerical gate. + qkv_projection_pack_dtypes: Dtypes for projection-only float packing + (PackQKVProjections) for decoder and masked-decoder roles, independent + of qkv_pack_dtypes and matmul_nbits_qkv_pack_dtypes. Experimental: + empty on every built-in EP until numerical execution is qualified. supports_fused_rope: ``False`` triggers SeparateRoPE + UnpackQKV lowering (DML). supports_skip_layer_norm: ``False`` expands SkipLayerNormalization / @@ -171,6 +179,13 @@ class EpCapabilities: max_buffer_size: int | None = None layered_per_layer_inputs: bool = False requires_graph_capture_rewrite: bool = False + # Append new capabilities to preserve all existing positional arguments. + matmul_nbits_qkv_pack_dtypes: frozenset[ir.DataType] = dataclasses.field( + default_factory=frozenset + ) + qkv_projection_pack_dtypes: frozenset[ir.DataType] = dataclasses.field( + default_factory=frozenset + ) def __post_init__(self) -> None: if not self.supports_fused_rope and self.qkv_pack_dtypes: 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/_model_package.py b/src/mobius/_model_package.py index 5ecc4f7cf..15a760dcb 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -41,6 +41,7 @@ import torch import tqdm +from mobius._constants import LAZY_SOURCE_FILES_META_KEY from mobius._export_report import ComponentExportReport from mobius._optimizations import fold_initializers_after_weights from mobius._weight_utils import WeightTransform @@ -202,6 +203,8 @@ def _validate_streaming_safetensors_destination( *, external_data: str, component_layout: _ComponentSaveLayout, + include_policy_components: bool = True, + include_adapter_artifacts: bool = True, ) -> None: """Defensively reject known aliases to files read by a streaming package. @@ -263,13 +266,46 @@ def _validate_streaming_safetensors_destination( _QUANTIZATION_REPORT, _EXPORT_REPORT, _WEIGHT_LOADING_REPORT, + "draft_manifest.json", + "gguf-reuse.json", ) ) + if include_policy_components: + output_files.update( + current_directory / "policies" / f"{name}.onnx" + for name in current_package.policy_components + ) + if include_adapter_artifacts: + output_files.update( + _adapter_artifact_destinations(current_package, current_directory) + ) source_directories.update( Path(path).resolve() for path in getattr(current_package, "_native_streaming_source_directories", ()) ) + # Derived lazy slices may own an external source that is no longer an + # initializer after graph cleanup (e.g. restored packed projections). + # Register that ownership before any output can be opened/truncated. + models = chain( + current_package.values(), + ( + component.model + for component in current_package.policy_components.values() + if include_policy_components + ), + ) + for model in models: + current_package.register_lazy_source_artifacts( + files=model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) + for graph in model.graphs(): + for initializer in graph.initializers.values(): + tensor = initializer.const_value + if isinstance(tensor, ir.LazyTensor): + current_package.register_lazy_source_artifacts( + files=tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) source_files.update( Path(path).resolve() for path in getattr(current_package, "_native_streaming_source_files", ()) @@ -309,6 +345,21 @@ def _validate_streaming_safetensors_destination( ) +def _adapter_artifact_destinations(package: ModelPackage, directory: Path) -> Iterator[Path]: + """Inventory the actual files emitted by save_adapter_artifacts, without writes.""" + options = package.adapter_service_options + for alias, artifact in package.adapter_artifacts.items(): + root = directory / "adapters" / alias + if options.portable_fallback: + yield root / "adapter.json" + if options.preserve_source_format: + if artifact.source.format == "onnx_adapter": + yield root / "adapter.onnx_adapter" + elif artifact.source.format == "peft_safetensors": + yield root / "adapter_model.safetensors" + yield root / "adapter_config.json" + + def _validate_output_directory( directory: str, *, kind: str, inspect_contents: bool = True ) -> None: @@ -742,6 +793,8 @@ def save( directory, external_data=external_data, component_layout=component_layout, + include_policy_components=include_policy_components, + include_adapter_artifacts=include_adapter_artifacts, ) _validate_mtp_save_paths( self, diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index 229e29d2f..e21d52bbd 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -21,10 +21,8 @@ Post-weight passes ------------------ :func:`fold_initializers_after_weights` should be called after weights are loaded. -It runs :class:`~mobius._passes.FoldTransposedInitializerPass` and -:class:`~mobius._passes.FoldConcatInitializersPass` to fold runtime Transpose and -Concat nodes over initializers into pre-computed weights, then removes unused -nodes. +It lazily folds initializer matrix Transpose and Concat nodes for projection +packing. Generic Reshape/Transpose layout chains remain explicit. """ from __future__ import annotations @@ -58,6 +56,7 @@ from onnx_ir.passes import common as common_passes from onnxscript.rewriter import rewrite +from mobius._constants import LAZY_SOURCE_FILES_META_KEY from mobius._execution_providers import EpCapabilities, ep_registry from mobius._flags import flags from mobius._passes import ( @@ -75,7 +74,9 @@ gelu_fusion_rules, group_query_attention_rules, htp_rank4_rmsnorm_rules, + pack_matmul_nbits_qkv_pass, pack_qkv_for_gqa_rules, + pack_qkv_projections_pass, separate_rope_rules, skip_layer_norm_rules, skip_norm_rules, @@ -375,6 +376,21 @@ def _apply_stage(model: ir.Model, rules_or_pass: list | ir.passes.InPlacePass) - rewrite(model, pattern_rewrite_rules=rules_or_pass) +def _apply_stages( + model: ir.Model, stages: list[tuple[str, list | ir.passes.InPlacePass]] +) -> None: + """Batch adjacent rule groups without dropping or reordering IR passes.""" + pending: list = [] + for _, rules_or_pass in stages: + if isinstance(rules_or_pass, list): + pending.extend(rules_or_pass) + else: + _apply_stage(model, pending) + pending = [] + _apply_stage(model, rules_or_pass) + _apply_stage(model, pending) + + def _log_trace_entry(entry: _TraceEntry) -> None: if not entry.added and not entry.removed: logger.info("[EP Trace] %-25s: no matches (0 nodes affected)", entry.name) @@ -409,7 +425,10 @@ def _get_optimization_passes( caps: EpCapabilities, dtype: ir.DataType, model_role: str = "decoder", -) -> tuple[list[tuple[str, list]], list[tuple[str, list | ir.passes.InPlacePass]]]: +) -> tuple[ + list[tuple[str, list | ir.passes.InPlacePass]], + list[tuple[str, list | ir.passes.InPlacePass]], +]: """Return ``(fuse_stages, lower_stages)`` for the given capabilities. Queries the :class:`EpCapabilities` object rather than branching on EP @@ -425,7 +444,7 @@ def _get_optimization_passes( ``(fuse_stages, lower_stages)`` — each a list of ``(name, payload)`` tuples where payload is a rule list or IR pass. """ - fuse: list[tuple[str, list]] = [] + fuse: list[tuple[str, list | ir.passes.InPlacePass]] = [] lower: list[tuple[str, list | ir.passes.InPlacePass]] = [] # --- Attention fusion (decoder only) --- @@ -441,6 +460,12 @@ def _get_optimization_passes( if model_role == "decoder" and dtype in caps.qkv_pack_dtypes: fuse.append(("PackQKV", list(pack_qkv_for_gqa_rules()))) + if model_role in {"decoder", "masked-decoder"}: + if dtype in caps.qkv_projection_pack_dtypes: + fuse.append(("PackQKVProjections", pack_qkv_projections_pass())) + if dtype in caps.matmul_nbits_qkv_pack_dtypes: + fuse.append(("PackMatMulNBitsQKV", pack_matmul_nbits_qkv_pass())) + # --- Normalization fusions (all roles, all dtypes) --- # When supports_skip_layer_norm=False (e.g. trt-rtx), fusion is skipped; # InlinePass in optimize_model() expands any pre-existing nodes instead. @@ -666,22 +691,12 @@ def _should_inline(func: ir.Function) -> bool: trace_entries.append(entry) _log_trace_entry(entry) else: - # Batch all rewrite rules for efficiency; apply IR passes separately. - all_fuse_rules = [r for _, rp in fuse_stages if isinstance(rp, list) for r in rp] - all_lower_rules = [r for _, rp in lower_stages if isinstance(rp, list) for r in rp] - lower_ir_passes = [(n, rp) for n, rp in lower_stages if not isinstance(rp, list)] - - if all_fuse_rules: - rewrite(model, pattern_rewrite_rules=all_fuse_rules) + _apply_stages(model, fuse_stages) # Expand unsupported custom ops via registered ir.Function bodies. inline_pass(model) - if all_lower_rules: - rewrite(model, pattern_rewrite_rules=all_lower_rules) - - for _, ir_pass in lower_ir_passes: - ir_pass(model) + _apply_stages(model, lower_stages) # Stage 4: Final dead-node removal, constant folding, and dead input # cleanup after rewrites. @@ -780,20 +795,34 @@ def _should_inline(func: ir.Function) -> bool: def fold_initializers_after_weights(model: ir.Model) -> None: - """Fold weight ``Transpose`` and ``Concat`` nodes after weights are loaded. - - Runs :class:`~onnx_ir.passes.common.LiftConstantsToInitializersPass`, - :class:`~mobius._passes.FoldConcatInitializersPass`, - :class:`~mobius._passes.FoldTransposedInitializerPass`, and - :class:`~onnx_ir.passes.common.RemoveUnusedNodesPass` in order. - FoldConcat must precede FoldTranspose so that packed QKV initializers are - visible before the Transpose fold runs. + """Lazily fold initializer matrix Transpose and Concat nodes. + + Preserve the shipped root Transpose naming convention and fold projection + packing to a fixed point. Generic weight layout chains remain explicit. """ - ir.passes.PassManager( - [ - common_passes.LiftConstantsToInitializersPass(), - FoldConcatInitializersPass(), - FoldTransposedInitializerPass(), - common_passes.RemoveUnusedNodesPass(), - ] - )(model) + common_passes.LiftConstantsToInitializersPass()(model) + for graph in model.graphs(): + for value in graph.initializers.values(): + tensor = value.const_value + if tensor is not None: + source_files = ( + (tensor.path,) + if isinstance(tensor, ir.ExternalTensor) + else tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) + if source_files: + model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( + (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) + ) + common_passes.RemoveUnusedNodesPass()(model) + lazy_passes = ( + FoldTransposedInitializerPass(), + FoldConcatInitializersPass(), + ) + while True: + modified = False + for lazy_pass in lazy_passes: + modified = lazy_pass(model).modified or modified + common_passes.RemoveUnusedNodesPass()(model) + if not modified: + break diff --git a/src/mobius/_passes/__init__.py b/src/mobius/_passes/__init__.py index 7c3a34507..36cfa31f7 100644 --- a/src/mobius/_passes/__init__.py +++ b/src/mobius/_passes/__init__.py @@ -35,6 +35,7 @@ Folds ``Concat(init_0, init_1, ...)`` patterns (where every input is a graph initializer) into a single pre-concatenated initializer. Used to pre-pack QKV weight matrices after weight loading. + """ from __future__ import annotations diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py index 7035f1e98..fd2dfb323 100644 --- a/src/mobius/_passes/_fold_concat.py +++ b/src/mobius/_passes/_fold_concat.py @@ -27,15 +27,65 @@ import numpy as np import onnx_ir as ir +from mobius._constants import LAZY_SOURCE_FILES_META_KEY from mobius._passes._dtype_utils import initializer_dtype logger = logging.getLogger(__name__) +def _retire_sources( + model: ir.Model, + graphs: tuple[ir.Graph, ...], + values: tuple[ir.Value | None, ...], + packed: ir.Value, + *, + prune: bool, +) -> None: + """Track lazy backing files and deregister dead child-Concat inputs. + + A child graph may capture a parent initializer: remove it from its owning + registry, not the consumer's graph. Keep graph I/O and live shared inputs. + Deregistration must not clear const_value: the packed lazy callback still + reads these tensors. Root inputs remain for the usual dead-code pass. + """ + source_files: set[str] = set() + for value in values: + if value is None: + continue + tensor = value.const_value + if tensor is not None: + source_files.update( + (tensor.path,) + if isinstance(tensor, ir.ExternalTensor) + else tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) + if ( + not prune + or value.name is None + or value.uses() + or any(value in scope.inputs or value in scope.outputs for scope in graphs) + ): + continue + owners = [scope for scope in graphs if scope.initializers.get(value.name) is value] + if len(owners) == 1: + del owners[0].initializers[value.name] + if source_files: + if packed.const_value is not None: + packed.const_value.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( + ( + *packed.const_value.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), + *source_files, + ) + ) + model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( + (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) + ) + + class FoldConcatInitializersPass(ir.passes.InPlacePass): """Fold ``Concat(init_0, init_1, ...)`` into a single packed initializer. - For each ``Concat`` node whose **every** input is a graph initializer: + For each ``Concat`` node whose **every** input is a non-overridable initializer: 1. A new initializer ``{name_0}__{name_1}__axis_{axis}__concat`` is registered whose :class:`~onnx_ir.LazyTensor` value lazily concatenates the inputs along @@ -44,20 +94,27 @@ class FoldConcatInitializersPass(ir.passes.InPlacePass): initializer. 3. The ``Concat`` node is removed from the graph. - The original initializers are left in place; a subsequent - :class:`~onnx_ir.passes.common.RemoveUnusedNodesPass` (or the - initializer-dedup pass) will prune them if they have no remaining uses. + For child-scope Concats, dead inputs are removed from their actual owning + registry. Shared inputs and graph I/O are retained. Root Concats preserve + their inputs until subsequent dead-code elimination, as before. + Lazy callbacks keep the source values and tensors alive after deregistration. """ def call(self, model: ir.Model) -> ir.passes.PassResult: modified = False folded_count = 0 + graphs = tuple(model.graphs()) for node in list(model.graph.all_nodes()): if node.op_type != "Concat": continue - inputs = node.inputs + # all_nodes() includes lexical child graphs. Packed names are only + # meaningful within that graph's initializer namespace: sibling + # branches can use the same names for different weight payloads. + graph = node.graph + assert graph is not None + inputs = tuple(node.inputs) if not inputs: continue @@ -65,6 +122,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: if not all(v is not None and v.is_initializer() for v in inputs): continue + # Initializers declared as graph inputs are runtime-overridable, + # including when captured by a lexical child graph. + if any(value in scope.inputs for scope in graphs for value in inputs): + continue + # All inputs must have loaded weights. Folding before weights are # loaded would force the dtype to default to FLOAT and bake a wrong # type into the packed initializer. Mirror FoldTransposedInitializer's @@ -122,15 +184,18 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: # the name collision is not an error — it is the idempotency guard. Reuse the # existing initializer and still remove this Concat node so all eligible Concat # nodes are folded. - if packed_name in model.graph.initializers: - existing_val = model.graph.initializers[packed_name] + if packed_name in graph.initializers: + existing_val = graph.initializers[packed_name] out_val.replace_all_uses_with(existing_val, replace_graph_outputs=True) # safe=True detaches the node from its inputs before removal so # the folded q/k/v initializers' use lists are cleared. Otherwise # `inp.uses()` still points at the removed Concat and the dead # pre-pack weights survive RemoveUnusedNodesPass (~1.8 GB of # orphaned weights serialized into the fp16 GQA model). - model.graph.remove(node, safe=True) + graph.remove(node, safe=True) + _retire_sources( + model, graphs, inputs, existing_val, prune=graph is not model.graph + ) folded_count += 1 modified = True @@ -168,13 +233,14 @@ def _make_packed( _make_packed, dtype=packed_dtype, shape=out_shape ) - model.graph.initializers[new_val.name] = new_val + graph.initializers[new_val.name] = new_val out_val.replace_all_uses_with(new_val, replace_graph_outputs=True) # safe=True detaches the node from its inputs before removal so the # folded q/k/v initializers' use lists are cleared; otherwise the # dead pre-pack weights survive RemoveUnusedNodesPass. - model.graph.remove(node, safe=True) + graph.remove(node, safe=True) + _retire_sources(model, graphs, inputs, new_val, prune=graph is not model.graph) folded_count += 1 modified = True diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py index d8bb04162..9d2361a5b 100644 --- a/src/mobius/_passes/_fold_concat_test.py +++ b/src/mobius/_passes/_fold_concat_test.py @@ -5,9 +5,15 @@ from __future__ import annotations +from pathlib import Path + import numpy as np import onnx_ir as ir +import pytest +from onnx_ir.passes import common +from mobius import ModelPackage +from mobius._constants import LAZY_SOURCE_FILES_META_KEY from mobius._passes._fold_concat import FoldConcatInitializersPass @@ -65,6 +71,164 @@ def _make_concat_model( return model, init_vals +@pytest.mark.parametrize("parent_owned", [False, True]) +@pytest.mark.parametrize("retained", [None, "consumer", "output"]) +def test_nested_concat_prunes_only_dead_owned_sources(parent_owned, retained): + arrays = [ + np.arange(32, dtype=np.float32).reshape(4, 8), + np.full((4, 8), 100, dtype=np.float32), + ] + original, sources = _make_concat_model(arrays) + branch = original.graph + branch.inputs.clear() + branch.outputs[0].name = "packed_output" + if retained == "consumer" or (retained == "output" and parent_owned): + identity = ir.Node("", "Identity", [sources[0]], num_outputs=1) + identity.outputs[0].name = "retained_output" + identity.outputs[0].shape = sources[0].shape + identity.outputs[0].dtype = sources[0].dtype + branch.append(identity) + branch.outputs.append(identity.outputs[0]) + elif retained == "output": + branch.outputs.append(sources[0]) + other_outputs = [] + for index, output in enumerate(branch.outputs): + assert output.shape is not None + other_outputs.append( + ir.Value( + name=f"other_{index}", + shape=output.shape, + type=output.type, + const_value=ir.tensor(np.zeros(output.shape.numpy(), dtype=np.float32)), + ) + ) + other_branch = ir.Graph( + [], other_outputs, nodes=[], initializers=other_outputs, name="other_branch" + ) + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + select = ir.Node( + "", + "If", + [condition], + [ir.AttrGraph("then_branch", branch), ir.AttrGraph("else_branch", other_branch)], + num_outputs=len(branch.outputs), + ) + for index, output in enumerate(select.outputs): + output.name = f"result_{index}" + output.shape = branch.outputs[index].shape + output.type = branch.outputs[index].type + root = ir.Graph( + [condition], + select.outputs, + nodes=[select], + opset_imports={"": 20}, + name="nested_concat_ownership", + ) + model = ir.Model(root, ir_version=10) + if parent_owned: + for source in sources: + assert source.name is not None + del branch.initializers[source.name] + root.register_initializer(source) + if retained == "output": + root.outputs.append(sources[0]) + owner = root if parent_owned else branch + reads = [] + tensors = [] + for source, array in zip(sources, arrays, strict=True): + + def read(data=array): + reads.append(True) + return ir.tensor(data) + + tensor = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) + tensor.meta[LAZY_SOURCE_FILES_META_KEY] = ("checkpoint.bin",) + source.const_value = tensor + tensors.append(tensor) + + assert FoldConcatInitializersPass()(model).modified + packed_name = "init_0__init_1__axis_0__concat" + packed = branch.initializers[packed_name] + expected_keys = {sources[0].name} if retained is not None else set() + if not parent_owned: + expected_keys.add(packed_name) + assert set(owner.initializers) == expected_keys + if parent_owned: + assert set(branch.initializers) == {packed_name} + assert all( + source.const_value is tensor for source, tensor in zip(sources, tensors, strict=True) + ) + assert not reads + assert set(packed.const_value.meta[LAZY_SOURCE_FILES_META_KEY]) == {"checkpoint.bin"} + assert set(model.meta[LAZY_SOURCE_FILES_META_KEY]) == {"checkpoint.bin"} + assert not FoldConcatInitializersPass()(model).modified + np.testing.assert_array_equal(packed.const_value.numpy(), np.concatenate(arrays, axis=0)) + assert len(reads) == 2 + + +def test_concat_preserves_overridable_graph_input_initializer(): + model, sources = _make_concat_model([np.ones((4, 8), dtype=np.float32)] * 2) + model.graph.inputs.append(sources[0]) + assert not FoldConcatInitializersPass()(model).modified + assert [node.op_type for node in model.graph] == ["Concat"] + common.RemoveUnusedNodesPass()(model) + assert model.graph.initializers["init_0"] is sources[0] + assert model.graph.initializers["init_1"] is sources[1] + + +def test_concat_cleanup_preserves_external_source_ownership(tmp_path, monkeypatch): + arrays = [ + np.arange(512, dtype=np.float32).reshape(32, 16), + np.full((32, 16), 3, dtype=np.float32), + ] + model, _ = _make_concat_model(arrays) + model.graph.inputs.clear() + model.graph.outputs[0].name = "output" + base = tmp_path / "base" + ModelPackage({"model": model}).save(base, max_workers=1, progress_bar=False) + before = {path.name: path.read_bytes() for path in base.iterdir()} + model = ModelPackage.load(base)["model"] + sources = list(model.graph.initializers.values()) + assert all(isinstance(value.const_value, ir.ExternalTensor) for value in sources) + source_path = (base / "model.onnx.data").resolve() + + def unexpected_read(*args, **kwargs): + raise AssertionError("folding and alias refusal must not read external payloads") + + packed_name = "init_0__init_1__axis_0__concat" + with monkeypatch.context() as patch: + patch.setattr(ir.ExternalTensor, "numpy", unexpected_read) + assert FoldConcatInitializersPass()(model).modified + assert set(model.graph.initializers) == {"init_0", "init_1", packed_name} + common.RemoveUnusedNodesPass()(model) + assert set(model.graph.initializers) == {packed_name} + packed = model.graph.initializers[packed_name] + assert model.graph.outputs[0] is packed + assert source_path in { + Path(path).resolve() + for path in packed.const_value.meta[LAZY_SOURCE_FILES_META_KEY] + } + assert source_path in { + Path(path).resolve() for path in model.meta[LAZY_SOURCE_FILES_META_KEY] + } + assert all(value.const_value is not None for value in sources) + # A fresh package cannot rely on source-file registration performed + # before pruning: the packed lazy tensor must carry its own backing. + package = ModelPackage({"model": model}) + with pytest.raises(ValueError, match="aliases lazy source file"): + package.save(base, max_workers=1, progress_bar=False) + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + safe = tmp_path / "safe" + package.save(safe, max_workers=1, progress_bar=False) + restored = ModelPackage.load(safe)["model"].graph + assert set(restored.initializers) == {packed_name} + np.testing.assert_array_equal( + restored.initializers[packed_name].const_value.numpy(), np.concatenate(arrays, axis=0) + ) + + class TestFoldConcatInitializersPass: def test_folds_two_initializers_axis0(self): """Concat of two 2-D initializers along axis=0 is folded.""" @@ -336,6 +500,13 @@ def test_duplicate_concat_nodes_reuse_packed_initializer(self): # Only one packed initializer should be registered. packed_keys = [k for k in model.graph.initializers if "__concat" in k] assert len(packed_keys) == 1 + assert {"init_0", "init_1"} <= model.graph.initializers.keys() + common.RemoveUnusedNodesPass()(model) + assert set(model.graph.initializers) == set(packed_keys) + np.testing.assert_array_equal( + model.graph.initializers[packed_keys[0]].const_value.numpy(), + np.concatenate([a, b], axis=0), + ) node_types = [n.op_type for n in model.graph.all_nodes()] assert "Concat" not in node_types diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py index 957d57b86..e14ddc033 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -62,7 +62,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: modified = False for node in list(model.graph.all_nodes()): - if node.op_type != "Transpose": + if node.domain != "" or node.overload or node.op_type != "Transpose": continue if len(node.inputs) != 1: @@ -72,8 +72,16 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: if inp is None or inp.name is None: continue - # Only fold initializer inputs (not graph inputs or computed values). - if not inp.is_initializer(): + # Register legacy ``weight_t`` names only for root-owned weights. + # A Transpose inside If/Loop can capture a root weight and still use + # this pass. Skip child-local weights: registering them at the root + # would collide across sibling scopes. + if ( + not inp.is_initializer() + or inp.is_graph_input() + or inp.is_graph_output() + or model.graph.initializers.get(inp.name) is not inp + ): continue # Only fold when the initializer's sole consumer is this Transpose node. diff --git a/src/mobius/_passes/_fold_transpose_test.py b/src/mobius/_passes/_fold_transpose_test.py index dc64970ef..af7832895 100644 --- a/src/mobius/_passes/_fold_transpose_test.py +++ b/src/mobius/_passes/_fold_transpose_test.py @@ -660,3 +660,58 @@ def test_transpose_inside_loop_body_folds_against_root_initializer(): w_t = model.graph.initializers["weight_t"] np.testing.assert_array_equal(w_t.const_value.numpy(), weight_data.T) + + +def test_post_weight_pipeline_preserves_generic_layout_chain(): + from mobius._optimizations import fold_initializers_after_weights + + data = np.arange(24, dtype=np.uint8).reshape(2, 12) + + def unexpected_read(): + raise AssertionError("generic layout weights must remain untouched") + + weight = ir.Value( + name="checkpoint.weight", + shape=ir.Shape(data.shape), + type=ir.TensorType(ir.DataType.UINT8), + const_value=ir.LazyTensor( + unexpected_read, dtype=ir.DataType.UINT8, shape=ir.Shape(data.shape) + ), + ) + shapes = [ + ir.Value( + name=name, + shape=ir.Shape([len(dimensions)]), + type=ir.TensorType(ir.DataType.INT64), + const_value=ir.tensor(np.array(dimensions, dtype=np.int64)), + ) + for name, dimensions in (("split_shape", (2, 2, 6)), ("final_shape", (2, 12))) + ] + reshape = ir.Node("", "Reshape", [weight, shapes[0]], num_outputs=1) + transpose = ir.Node( + "", + "Transpose", + [reshape.outputs[0]], + attributes=[ir.AttrInt64s("perm", [0, 2, 1])], + num_outputs=1, + ) + result = ir.Node("", "Reshape", [transpose.outputs[0], shapes[1]], num_outputs=1) + result.outputs[0].name = "reordered_weight" + nodes = [reshape, transpose, result] + model = ir.Model( + ir.Graph( + [], + [result.outputs[0]], + nodes=nodes, + initializers=[weight, *shapes], + opset_imports={"": 24}, + ), + ir_version=10, + ) + + fold_initializers_after_weights(model) + + assert list(model.graph) == nodes + assert model.graph.initializers["checkpoint.weight"] is weight + assert reshape.inputs[0] is weight + assert model.graph.outputs[0] is result.outputs[0] diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py index 034840069..14bd199ae 100644 --- a/src/mobius/adapter_io.py +++ b/src/mobius/adapter_io.py @@ -21,8 +21,14 @@ import numpy as np import onnx_ir as ir +from onnx_ir.passes import common as common_passes from safetensors.numpy import load_file +from mobius._constants import ( + LAZY_SOURCE_FILES_META_KEY, + LEGACY_QKV_PROJECTION_SOURCES_METADATA_KEY, + QKV_PROJECTION_SOURCES_METADATA_KEY, +) from mobius.adapters import ( AdapterArtifact, AdapterServiceOptions, @@ -166,11 +172,210 @@ def _peft_module_keys(tensors: Mapping[str, np.ndarray]) -> tuple[str, ...]: ) +def _packed_projection_sources(node: ir.Node, metadata: str) -> list[tuple[str, int]]: + record: object = json.loads(metadata) + if ( + not isinstance(record, dict) + or type(record.get("version")) is not int + or record["version"] != 1 + ): + raise ValueError(f"unsupported packed-projection metadata on {node.name!r}") + names, widths = record.get("weights"), record.get("widths") + if ( + not isinstance(names, list) + or not isinstance(widths, list) + or len(names) != len(widths) + or len(names) not in (2, 3) + ): + raise ValueError(f"invalid packed-projection source list on {node.name!r}") + sources = [] + for name, width in zip(names, widths, strict=True): + if not isinstance(name, str) or not name or type(width) is not int or width <= 0: + raise ValueError(f"invalid packed-projection source on {node.name!r}") + original = name.removesuffix("_t") if name.endswith(".weight_t") else name + sources.append((original, width)) + return sources + + +def _packed_float_projection( + model: ir.Model, split: ir.Node, source_count: int, module_key: str +) -> ir.Node: + """Require the packer's MatMul-to-Split topology before restoring PEFT targets. + + The packed activation must be private to Split; replacing a shared or public + activation could change consumers outside the adapter's projection group. + """ + packed = split.inputs[0] if split.inputs else None + projection = packed.producer() if packed is not None else None + if ( + split.domain != "" + or split.op_type != "Split" + or split.overload + or split.attributes.get_int("axis", 0) not in (-1, 2) + or len(split.outputs) != source_count + or packed is None + or projection is None + or projection.domain != "" + or projection.op_type != "MatMul" + or projection.overload + or projection.attributes + or len(projection.inputs) != 2 + or len(projection.outputs) != 1 + or packed in model.graph.outputs + or len(packed.uses()) != 1 + ): + raise ValueError(f"PEFT module {module_key!r} has an unsupported packed projection") + return projection + + +def _loaded_packed_projection_inputs( + model: ir.Model, projection: ir.Node, total_width: int, module_key: str +) -> tuple[ir.Value, ir.Value]: + """Require a loaded static float weight before lazily slicing its columns. + + Overridable, quantized, or symbolic weights cannot be restored by this + float-only path. Validate without materializing the packed payload. + """ + x, weight = projection.inputs + if ( + x is None + or weight is None + or not weight.is_initializer() + or weight in model.graph.inputs + or weight.const_value is None + or weight.shape is None + or len(weight.shape) != 2 + or weight.dtype not in (ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16) + or not all(isinstance(dim, int) and dim > 0 for dim in weight.shape) + or weight.shape[1] != total_width + ): + raise ValueError( + f"PEFT module {module_key!r} requires a loaded, folded packed projection weight" + ) + return x, weight + + +def _restore_packed_peft_projections( + model: ir.Model, module_key: str, candidates: list[str] +) -> None: + """Restore independent float projections before attaching a PEFT adapter.""" + suffixes = tuple( + suffix + for candidate in candidates + for suffix in (f"{candidate}.weight", f"{candidate}.weight_t") + ) + groups = [] + for node in model.graph: + metadata = node.metadata_props.get(QKV_PROJECTION_SOURCES_METADATA_KEY) + if metadata is None: + # Read previews written before adopting the persistent-key namespace. + metadata = node.metadata_props.get(LEGACY_QKV_PROJECTION_SOURCES_METADATA_KEY) + if metadata is None: + continue + sources = _packed_projection_sources(node, metadata) + if any(name.endswith(suffixes) for name, _ in sources): + groups.append((node, sources)) + if not groups: + return + if ( + len(groups) != 1 + or sum(name.endswith(suffixes) for name, _ in groups[0][1]) != 1 + or any(name.endswith(suffixes) for name in model.graph.initializers) + ): + raise ValueError(f"PEFT module {module_key!r} has ambiguous packed projection targets") + split, sources = groups[0] + projection = _packed_float_projection(model, split, len(sources), module_key) + x, weight = _loaded_packed_projection_inputs( + model, projection, sum(width for _, width in sources), module_key + ) + assert weight.shape is not None and weight.const_value is not None + names = [name for name, _ in sources] + if len(set(names)) != len(names) or any( + name in model.graph.initializers for name in names + ): + raise ValueError(f"PEFT module {module_key!r} has conflicting source weights") + if any( + output.shape is None + or len(output.shape) != 3 + or output.shape[-1] != width + or output.dtype != weight.dtype + for output, (_, width) in zip(split.outputs, sources, strict=True) + ): + raise ValueError(f"PEFT module {module_key!r} has inconsistent packed output metadata") + nodes: list[ir.Node] = [] + outputs = [] + initializers = [] + start = 0 + tensor = weight.const_value + assert weight.dtype is not None + for original, (name, width) in zip(split.outputs, sources, strict=True): + end = start + width + parameter_shape = ir.Shape([width, weight.shape[0]]) + parameter = ir.Value(name=name, type=weight.type, shape=parameter_shape) + + def load_projection( + source: ir.TensorProtocol = tensor, lo: int = start, hi: int = end + ) -> ir.TensorProtocol: + return ir.tensor(np.ascontiguousarray(source.numpy()[:, lo:hi].T)) + + parameter.const_value = ir.LazyTensor( + load_projection, dtype=weight.dtype, shape=parameter_shape, name=name + ) + # RemoveUnusedNodes drops the packed ExternalTensor from the graph. + # Keep its file ownership visible to ModelPackage's existing pre-write + # alias guard without materializing the full packed checkpoint. + source_files = ( + (tensor.path,) + if isinstance(tensor, ir.ExternalTensor) + else tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) + if source_files: + parameter.const_value.meta[LAZY_SOURCE_FILES_META_KEY] = tuple(source_files) + # Folding may replace this tensor; ownership belongs to the model too. + model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( + (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) + ) + initializers.append(parameter) + transposed = ir.Value(type=weight.type, shape=ir.Shape([weight.shape[0], width])) + nodes.append( + ir.Node( + "", + "Transpose", + [parameter], + {"perm": ir.AttrInt64s("perm", [1, 0])}, + outputs=[transposed], + ) + ) + output = ir.Value( + name=original.name, + shape=original.shape, + type=original.type, + doc_string=original.doc_string, + ) + output.metadata_props.update(original.metadata_props) + output.meta.update(original.meta) + outputs.append(output) + nodes.append(ir.Node("", "MatMul", [x, transposed], outputs=[output])) + start = end + for initializer in initializers: + model.graph.initializers.add(initializer) + ir.convenience.replace_nodes_and_values( + model.graph, + insertion_point=projection, + old_nodes=[projection, split], + new_nodes=nodes, + old_values=list(split.outputs), + new_values=outputs, + ) + common_passes.RemoveUnusedNodesPass()(model) + + def _resolve_initializer_name(model: ir.Model, module_key: str) -> str: candidates = [module_key] for prefix in ("base_model.model.", "base_model.", "model."): if module_key.startswith(prefix): candidates.append(module_key[len(prefix) :]) + _restore_packed_peft_projections(model, module_key, candidates) matches = [ name for name in model.graph.initializers 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/functions/__init__.py b/src/mobius/functions/__init__.py index 49d8fc82d..71d4fecfc 100644 --- a/src/mobius/functions/__init__.py +++ b/src/mobius/functions/__init__.py @@ -38,6 +38,8 @@ matmul_block_quantized_fp4_weight, ) from mobius.functions.matmul_nbits import ( + FUNCTION_BODY_METADATA_KEY, + FUNCTION_BODY_METADATA_VALUE, matmul_nbits, ) from mobius.functions.packed_multi_head_attention import ( @@ -49,7 +51,6 @@ ) _DOMAIN = "com.microsoft" -_FUNCTION_BODY_MARKER = "mobius.function_body" _MATMUL_NBITS_ID = (_DOMAIN, "MatMulNBits", "") _MATMUL_NBITS_DEFAULT_ZERO_POINTS_ID = (_DOMAIN, "MatMulNBits", "default_zero_points") @@ -89,13 +90,17 @@ def get_function(op_id: ir.OperatorIdentifier) -> ir.Function | None: if builder is None: return None function = builder() - function.metadata_props[_FUNCTION_BODY_MARKER] = "1" + function.metadata_props[FUNCTION_BODY_METADATA_KEY] = FUNCTION_BODY_METADATA_VALUE return function def _specialize_matmul_nbits_calls(model: ir.Model) -> None: existing = model.functions.get(_MATMUL_NBITS_ID) - if existing is not None and existing.metadata_props.get(_FUNCTION_BODY_MARKER) != "1": + if ( + existing is not None + and existing.metadata_props.get(FUNCTION_BODY_METADATA_KEY) + != FUNCTION_BODY_METADATA_VALUE + ): return graphs = [model.graph, *(function.graph for function in model.functions.values())] diff --git a/src/mobius/functions/matmul_nbits.py b/src/mobius/functions/matmul_nbits.py index f2c5ee43c..2c610d4b6 100644 --- a/src/mobius/functions/matmul_nbits.py +++ b/src/mobius/functions/matmul_nbits.py @@ -48,6 +48,16 @@ from mobius._constants import OPSET_VERSION DOMAIN = "com.microsoft" +FUNCTION_BODY_METADATA_KEY = "pkg.mobius.function_body" +"""Mark registry-generated fallbacks that Mobius may specialize or pack against. + +An unmarked, caller-provided function may implement different semantics despite +using the same operator identifier. The marker distinguishes our known fallback +from those bodies; calling the factory directly does not stamp it. +Legacy-only markers conservatively leave the body caller-owned: restoring an +adapter needs the old source names, but optimizations need not trust an old body. +""" +FUNCTION_BODY_METADATA_VALUE = "1" def matmul_nbits(*, has_zero_points: bool = True) -> ir.Function: @@ -152,4 +162,7 @@ def body(op, a_input, b_packed, scales_input, zero_points_input): ) if not has_zero_points: function.overload = "default_zero_points" + # get_function() attaches automatic-fallback ownership metadata. A direct + # factory result is caller-owned: do not authorize call specialization or + # projection packing merely because the caller used this factory. return function diff --git a/src/mobius/functions/matmul_nbits_test.py b/src/mobius/functions/matmul_nbits_test.py index 454f6795a..e7f19b9c1 100644 --- a/src/mobius/functions/matmul_nbits_test.py +++ b/src/mobius/functions/matmul_nbits_test.py @@ -29,6 +29,10 @@ from onnx_ir.passes.common import InlinePass from mobius.functions import get_function, matmul_nbits, register_function_bodies +from mobius.functions.matmul_nbits import ( + FUNCTION_BODY_METADATA_KEY, + FUNCTION_BODY_METADATA_VALUE, +) def _run(model: ir.Model, feeds: dict) -> np.ndarray: @@ -144,6 +148,49 @@ def test_new_call_is_specialized_after_registration(self): assert node.overload == "default_zero_points" assert node.op_identifier() in model.functions + @pytest.mark.parametrize("has_zero_points", [False, True]) + def test_direct_factory_does_not_claim_registry_ownership(self, has_zero_points): + assert FUNCTION_BODY_METADATA_KEY == "pkg.mobius.function_body" + function = matmul_nbits(has_zero_points=has_zero_points) + assert FUNCTION_BODY_METADATA_KEY not in function.metadata_props + automatic = get_function(function.identifier()) + assert automatic is not None + assert automatic.metadata_props[FUNCTION_BODY_METADATA_KEY] == ( + FUNCTION_BODY_METADATA_VALUE + ) + assert FUNCTION_BODY_METADATA_KEY not in function.metadata_props + + def test_legacy_marker_does_not_authorize_call_specialization(self): + model, *_ = _build_matmulnbits_model(n_out=8, k_in=64, block=32, has_zero_points=False) + function = matmul_nbits() + function.metadata_props["mobius.function_body"] = "1" + model.functions[function.identifier()] = function + + register_function_bodies(model) + + assert model.functions[function.identifier()] is function + assert next(iter(model.graph)).overload == "" + + @pytest.mark.parametrize("has_zero_points", [False, True]) + def test_automatic_registration_retains_provenance_and_specialization( + self, has_zero_points + ): + model, *_ = _build_matmulnbits_model( + n_out=8, k_in=64, block=32, has_zero_points=has_zero_points + ) + register_function_bodies(model) + node = next(iter(model.graph)) + assert node.overload == ("" if has_zero_points else "default_zero_points") + function = model.functions[node.op_identifier()] + assert function.metadata_props[FUNCTION_BODY_METADATA_KEY] == ( + FUNCTION_BODY_METADATA_VALUE + ) + register_function_bodies(model) + assert model.functions[node.op_identifier()] is function + assert function.metadata_props[FUNCTION_BODY_METADATA_KEY] == ( + FUNCTION_BODY_METADATA_VALUE + ) + class TestMatMulNBitsInlineParity: def _inline(self, model: ir.Model) -> None: 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/rewrite_rules/__init__.py b/src/mobius/rewrite_rules/__init__.py index 38d856291..3506eb54b 100644 --- a/src/mobius/rewrite_rules/__init__.py +++ b/src/mobius/rewrite_rules/__init__.py @@ -40,6 +40,8 @@ "htp_rank4_rmsnorm_rules", "layer_norm_fusion_rules", "pack_qkv_for_gqa_rules", + "pack_matmul_nbits_qkv_pass", + "pack_qkv_projections_pass", "packed_attention_rules", "separate_rope_rules", "skip_layer_norm_rules", @@ -64,6 +66,10 @@ layer_norm_fusion_rules, ) from mobius.rewrite_rules._packed_attention import packed_attention_rules +from mobius.rewrite_rules._qkv_projection_packing import ( + pack_matmul_nbits_qkv_pass, + pack_qkv_projections_pass, +) from mobius.rewrite_rules._qmoe_fusion import fuse_dense_moe_to_qmoe from mobius.rewrite_rules._scatternd import tensor_scatter_to_scatternd_rules from mobius.rewrite_rules._separate_rope import separate_rope_rules diff --git a/src/mobius/rewrite_rules/_qkv_projection_packing.py b/src/mobius/rewrite_rules/_qkv_projection_packing.py new file mode 100644 index 000000000..2c029984e --- /dev/null +++ b/src/mobius/rewrite_rules/_qkv_projection_packing.py @@ -0,0 +1,479 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Combine Q/K/V projections into one wider GEMM to reduce kernel launches. + +Split before normalization, RoPE, attention, or caches so their input contracts +remain unchanged. Supports distinct Q/K/V or Q/K with raw K reused for V. +Parameter Concats preserve checkpoint names until weight binding; packing then +joins the loaded weights without dequantizing or requantizing them. +""" + +from __future__ import annotations + +import json +from collections import Counter +from collections.abc import Callable + +import numpy as np +import onnx_ir as ir + +from mobius._constants import QKV_PROJECTION_SOURCES_METADATA_KEY +from mobius.functions.matmul_nbits import ( + FUNCTION_BODY_METADATA_KEY, + FUNCTION_BODY_METADATA_VALUE, +) + +_Group = tuple[ir.Node, ...] +_FLOAT_DTYPES = frozenset({ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16}) +_NBITS_ATTRIBUTES = frozenset({"K", "N", "bits", "block_size", "accuracy_level"}) +_DATA_FIRST = frozenset( + { + "Identity", + "Reshape", + "Transpose", + "Cast", + "CastLike", + "Squeeze", + "Unsqueeze", + "Slice", + "Gather", + "RotaryEmbedding", + "RMSNormalization", + "SimplifiedLayerNormalization", + "ReduceMean", + "Sqrt", + "Reciprocal", + "Pow", + "Neg", + } +) +_DATA_ALL = frozenset({"Add", "Sub", "Mul", "Div", "Concat"}) + + +def _standard(node: ir.Node, op_type: str) -> bool: + return node.domain == "" and node.op_type == op_type and not node.overload + + +def _parameter( + graph: ir.Graph, value: ir.Value | None, dtype: ir.DataType, shape: list[int] +) -> bool: + if ( + value is None + or value.name is None + or value.producer() is not None + or graph.initializers.get(value.name) is not value + or value in graph.inputs + or value.dtype != dtype + or value.shape is None + or list(value.shape) != shape + ): + return False + tensor = value.const_value + return tensor is None or (tensor.dtype == dtype and list(tensor.shape) == shape) + + +def _projection(graph: ir.Graph, value: ir.Value | None) -> ir.Node | None: + pending = [value] + visited: set[ir.Node] = set() + projections: set[ir.Node] = set() + while pending: + current = pending.pop() + if current is None or current.producer() is None: + continue + node = current.producer() + assert node is not None + if node in visited: + continue + if node.graph is not graph or len(visited) >= 128: + return None + visited.add(node) + if _standard(node, "MatMul") or ( + node.domain == "com.microsoft" and node.op_type == "MatMulNBits" + ): + projections.add(node) + elif _standard(node, "Constant"): + continue + elif node.domain in ("", "com.microsoft"): + if node.op_type in _DATA_FIRST: + pending.extend(node.inputs[:1]) + elif node.domain != "com.microsoft" and node.op_type in _DATA_ALL: + pending.extend(node.inputs) + else: + return None + else: + return None + return next(iter(projections)) if len(projections) == 1 else None + + +def _shared_input(group: _Group) -> ir.Value | None: + first = group[0] + if not first.inputs or first.inputs[0] is None: + return None + value = first.inputs[0] + if ( + value.dtype not in _FLOAT_DTYPES + or value.shape is None + or len(value.shape) != 3 + or not isinstance(value.shape[-1], int) + or value.shape[-1] <= 0 + or any(node.inputs[0] is not value or len(node.outputs) != 1 for node in group) + ): + return None + return value + + +def _output_matches(node: ir.Node, value: ir.Value, width: int) -> bool: + assert value.shape is not None + output = node.outputs[0] + if output.dtype not in (None, value.dtype): + return False + if output.shape is None: + return True + expected = [*value.shape[:-1], width] + return len(output.shape) == 3 and all( + (isinstance(dim, ir.SymbolicDim) and dim.value is None) or dim == wanted + for dim, wanted in zip(output.shape, expected, strict=True) + ) + + +def _concat( + nodes: list[ir.Node], values: list[ir.Value], axis: int, shape: list[int | ir.SymbolicDim] +) -> ir.Value: + output = ir.Value(shape=ir.Shape(shape), type=values[0].type) + nodes.append( + ir.Node("", "Concat", values, {"axis": ir.AttrInt64("axis", axis)}, outputs=[output]) + ) + return output + + +def _replace( + graph: ir.Graph, + group: _Group, + value: ir.Value, + widths: list[int], + nodes: list[ir.Node], + packed: ir.Value, +) -> ir.Node: + assert value.shape is not None and value.dtype is not None + tensor = ir.tensor(np.asarray(widths, dtype=np.int64)) + sizes = ir.Value( + shape=ir.Shape([len(widths)]), + type=ir.TensorType(ir.DataType.INT64), + const_value=tensor, + ) + nodes.append( + ir.Node("", "Constant", [], {"value": ir.AttrTensor("value", tensor)}, outputs=[sizes]) + ) + outputs = [] + for node, width in zip(group, widths, strict=True): + original = node.outputs[0] + output = ir.Value( + name=original.name, + doc_string=original.doc_string, + shape=ir.Shape([*value.shape[:-1], width]), + type=ir.TensorType(value.dtype), + ) + output.metadata_props.update(original.metadata_props) + output.meta.update(original.meta) + outputs.append(output) + split = ir.Node( + "", "Split", [packed, sizes], {"axis": ir.AttrInt64("axis", -1)}, outputs=outputs + ) + nodes.append(split) + ir.convenience.replace_nodes_and_values( + graph, + insertion_point=next(node for node in graph if node in group), + old_nodes=list(group), + new_nodes=nodes, + old_values=[node.outputs[0] for node in group], + new_values=outputs, + ) + for output, width in zip(outputs, widths, strict=True): + output.shape = ir.Shape([*value.shape[:-1], width]) + output.type = ir.TensorType(value.dtype) + return split + + +def _pack_float(model: ir.Model, group: _Group) -> bool: + graph = model.graph + if any( + not _standard(node, "MatMul") or node.attributes or len(node.inputs) != 2 + for node in group + ): + return False + value = _shared_input(group) + if value is None: + return False + assert value.shape is not None and value.dtype is not None + k = value.shape[-1] + assert isinstance(k, int) + weights: list[tuple[ir.Value, bool]] = [] + widths = [] + for node in group: + weight = node.inputs[1] + if weight is None: + return False + producer = weight.producer() + transposed = producer is not None and _standard(producer, "Transpose") + if transposed: + assert producer is not None + if ( + producer.graph is not graph + or len(producer.inputs) != 1 + or set(producer.attributes) - {"perm"} + or tuple(producer.attributes.get_ints("perm", [1, 0])) != (1, 0) + ): + return False + weight = producer.inputs[0] + if weight is None or weight.shape is None or len(weight.shape) != 2: + return False + width = weight.shape[0 if transposed else 1] + if not isinstance(width, int) or width <= 0: + return False + if not _parameter( + graph, weight, value.dtype, [width, k] if transposed else [k, width] + ) or not _output_matches(node, value, width): + return False + # Restoring PEFT sources must remain unambiguous. A source initializer + # or its logical Transpose view shared outside this group cannot safely + # be hidden in independently packed groups. Decline before mutating IR. + allowed_consumers = set(group) + if transposed: + assert producer is not None + allowed_consumers.add(producer) + if producer.outputs[0] in graph.outputs or any( + use.node not in group for use in producer.outputs[0].uses() + ): + return False + if weight in graph.outputs or any( + use.node not in allowed_consumers for use in weight.uses() + ): + return False + weights.append((weight, transposed)) + widths.append(width) + # Distinct projection nodes can still read the same source (including + # separate Transpose views). Duplicate names cannot be restored as + # independent PEFT targets. Refuse before creating any replacement nodes. + if len({weight.name for weight, _ in weights}) != len(weights): + return False + nodes: list[ir.Node] = [] + all_transposed = all(transposed for _, transposed in weights) + logical_weights = [] + for weight, transposed in weights: + if transposed and not all_transposed: + assert weight.shape is not None + logical = ir.Value(shape=ir.Shape([k, weight.shape[0]]), type=weight.type) + nodes.append( + ir.Node( + "", + "Transpose", + [weight], + {"perm": ir.AttrInt64s("perm", [1, 0])}, + outputs=[logical], + ) + ) + logical_weights.append(logical) + else: + logical_weights.append(weight) + packed_weight = _concat( + nodes, + logical_weights, + 0 if all_transposed else 1, + [sum(widths), k] if all_transposed else [k, sum(widths)], + ) + if all_transposed: + logical = ir.Value(shape=ir.Shape([k, sum(widths)]), type=packed_weight.type) + nodes.append( + ir.Node( + "", + "Transpose", + [packed_weight], + {"perm": ir.AttrInt64s("perm", [1, 0])}, + outputs=[logical], + ) + ) + packed_weight = logical + packed = ir.Value(shape=ir.Shape([*value.shape[:-1], sum(widths)]), type=value.type) + nodes.append(ir.Node("", "MatMul", [value, packed_weight], outputs=[packed])) + split = _replace(graph, group, value, widths, nodes, packed) + split.metadata_props[QKV_PROJECTION_SOURCES_METADATA_KEY] = json.dumps( + { + "version": 1, + "weights": [weight.name for weight, _ in weights], + "widths": widths, + } + ) + return True + + +def _pack_nbits(model: ir.Model, group: _Group) -> bool: + graph = model.graph + if graph.opset_imports.get("com.microsoft") != 1 or any( + node.domain != "com.microsoft" or node.op_type != "MatMulNBits" or len(node.inputs) < 3 + for node in group + ): + return False + value = _shared_input(group) + if value is None: + return False + assert value.shape is not None and value.dtype is not None + first = group[0] + if any( + attr.type != ir.AttributeType.INT + for node in group + for attr in node.attributes.values() + ): + return False + k = first.attributes.get_int("K", 0) + block = first.attributes.get_int("block_size", 0) + if k != value.shape[-1] or block < 16 or block & (block - 1): + return False + blocks = (k + block - 1) // block + common = {key: attr for key, attr in first.attributes.items() if key != "N"} + has_zp = len(first.inputs) > 3 and first.inputs[3] is not None + layouts: list[tuple[int, ...]] = [] + widths = [] + for node in group: + if node.overload != first.overload or node.overload not in ("", "default_zero_points"): + return False + function = model.functions.get(node.op_identifier()) + if ( + function is not None + and function.metadata_props.get(FUNCTION_BODY_METADATA_KEY) + != FUNCTION_BODY_METADATA_VALUE + ): + return False + if node.overload == "default_zero_points" and has_zp: + return False + if set(node.attributes) - _NBITS_ATTRIBUTES: + return False + if {key: attr for key, attr in node.attributes.items() if key != "N"} != common: + return False + if ( + node.attributes.get_int("bits", 0) != 4 + or node.attributes.get_int("accuracy_level", 0) != 0 + ): + return False + n = node.attributes.get_int("N", 0) + if n <= 0 or any(item is not None for item in node.inputs[4:]): + return False + if not _parameter(graph, node.inputs[1], ir.DataType.UINT8, [n, blocks, block // 2]): + return False + scale = node.inputs[2] + if scale is None or scale.shape is None: + return False + scale_shape = [n, blocks] if len(scale.shape) == 2 else [n * blocks] + if not _parameter(graph, scale, value.dtype, scale_shape): + return False + zp = node.inputs[3] if len(node.inputs) > 3 else None + if (zp is not None) != has_zp: + return False + layout: tuple[int, ...] = (len(scale_shape),) + if zp is not None: + if zp.shape is None: + return False + zp_shape = ( + [n, (blocks + 1) // 2] if len(zp.shape) == 2 else [n * ((blocks + 1) // 2)] + ) + if not _parameter(graph, zp, ir.DataType.UINT8, zp_shape): + return False + layout += (len(zp_shape),) + if layouts and layout != layouts[0]: + return False + if not _output_matches(node, value, n): + return False + layouts.append(layout) + widths.append(n) + parameters = [item for node in group for item in node.inputs[1:4] if item is not None] + if len({item.name for item in parameters}) != len(parameters) or any( + item in graph.outputs or any(use.node not in group for use in item.uses()) + for item in parameters + ): + return False + nodes: list[ir.Node] = [] + inputs: list[ir.Value | None] = [value] + for slot in range(1, 4 if has_zp else 3): + originals = [node.inputs[slot] for node in group] + assert all(item is not None and item.shape is not None for item in originals) + parameters = [item for item in originals if item is not None] + shape = parameters[0].shape + assert shape is not None + extent = 0 + for parameter in parameters: + assert parameter.shape is not None + dimension = parameter.shape[0] + assert isinstance(dimension, int) + extent += dimension + inputs.append(_concat(nodes, parameters, 0, [extent, *shape[1:]])) + packed = ir.Value(shape=ir.Shape([*value.shape[:-1], sum(widths)]), type=value.type) + attributes = dict(first.attributes) + attributes["N"] = ir.AttrInt64("N", sum(widths)) + nodes.append( + ir.Node( + "com.microsoft", + "MatMulNBits", + inputs, + attributes, + overload=first.overload, + outputs=[packed], + ) + ) + _replace(graph, group, value, widths, nodes, packed) + return True + + +def _run(model: ir.Model, pack: Callable[[ir.Model, _Group], bool]) -> ir.passes.PassResult: + graph = model.graph + if graph.opset_imports.get("", 0) < 13: + return ir.passes.PassResult(model, modified=False) + candidates: set[_Group] = set() + for node in graph: + if ( + not ( + _standard(node, "Attention") + or (node.domain == "com.microsoft" and node.op_type == "GroupQueryAttention") + ) + or len(node.inputs) < 3 + or any(item is None for item in node.inputs[:3]) + ): + continue + projections = [_projection(graph, item) for item in node.inputs[:3]] + if any(item is None for item in projections): + continue + # Only K=V reuse is supported. Q=K or Q=V is not an independent + # projection group and must not produce ambiguous source metadata. + if projections[0] in projections[1:]: + continue + group = tuple(dict.fromkeys(item for item in projections if item is not None)) + if len(group) in (2, 3): + candidates.add(group) + counts = Counter(node for group in candidates for node in group) + order = {node: index for index, node in enumerate(graph)} + modified = False + for group in sorted(candidates, key=lambda group: min(order[node] for node in group)): + if all(counts[node] == 1 for node in group): + modified = pack(model, group) or modified + return ir.passes.PassResult(model, modified=modified) + + +class PackQKVProjectionsPass(ir.passes.InPlacePass): + """Pack compatible float projections, retaining all downstream branches.""" + + def call(self, model: ir.Model) -> ir.passes.PassResult: + return _run(model, _pack_float) + + +class PackMatMulNBitsQKVPass(ir.passes.InPlacePass): + """Pack homogeneous INT4 projections without dequantizing or requantizing.""" + + def call(self, model: ir.Model) -> ir.passes.PassResult: + return _run(model, _pack_nbits) + + +def pack_qkv_projections_pass() -> PackQKVProjectionsPass: + return PackQKVProjectionsPass() + + +def pack_matmul_nbits_qkv_pass() -> PackMatMulNBitsQKVPass: + return PackMatMulNBitsQKVPass() diff --git a/src/mobius/rewrite_rules/_qkv_projection_packing_test.py b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py new file mode 100644 index 000000000..0fba9651a --- /dev/null +++ b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py @@ -0,0 +1,1073 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Projection packing payload, cache, rejection, and adapter compatibility tests.""" + +from __future__ import annotations + +import dataclasses +import json +from types import SimpleNamespace + +import ml_dtypes +import numpy as np +import onnx_ir as ir +import pytest +import torch +from safetensors.numpy import save_file + +from mobius import ModelPackage +from mobius._component_quantization import validate_quantized_component_bindings +from mobius._constants import ( + LEGACY_QKV_PROJECTION_SOURCES_METADATA_KEY, + QKV_PROJECTION_SOURCES_METADATA_KEY, +) +from mobius._execution_providers import EpCapabilities, ep_registry, get_ep +from mobius._optimizations import ( + _get_optimization_passes, + fold_initializers_after_weights, + optimize_model, +) +from mobius._pipeline_contract import declare_arbitrary_attention_mask +from mobius._testing.ort_inference import OnnxModelSession +from mobius.adapter_io import _resolve_initializer_name, attach_peft_adapter +from mobius.adapters import AdapterSource +from mobius.functions import register_function_bodies +from mobius.functions.matmul_nbits import matmul_nbits +from mobius.rewrite_rules import pack_matmul_nbits_qkv_pass, pack_qkv_projections_pass + +_DTYPES = { + ir.DataType.FLOAT: np.float32, + ir.DataType.FLOAT16: np.float16, + ir.DataType.BFLOAT16: ml_dtypes.bfloat16, +} + + +def _model( + *, + quantized=False, + shared=False, + dtype=ir.DataType.FLOAT, + deferred=False, + transposed=True, + zero_points=True, + flat=False, + hidden_size=64, + block_size=32, + q_width=32, +): + rng = np.random.default_rng(42) + nodes, initializers, weights, projections = [], [], {}, [] + + def parameter(name, data): + tensor = ir.tensor(data) + value = ir.Value( + name=name, + shape=tensor.shape, + type=ir.TensorType(tensor.dtype), + const_value=None if deferred else tensor, + ) + initializers.append(value) + weights[name] = data + return value + + def op(kind, inputs, name, shape, attrs=None): + output = ir.Value(name=name, shape=ir.Shape(shape), type=ir.TensorType(dtype)) + nodes.append( + ir.Node( + "", + kind, + inputs, + ir.convenience.convert_attributes(attrs or {}), + outputs=[output], + name=name, + ) + ) + return output + + x = ir.Value(name="x", shape=ir.Shape([2, "S", hidden_size]), type=ir.TensorType(dtype)) + blocks = (hidden_size + block_size - 1) // block_size + for index, (name, width) in enumerate( + (("q", q_width), ("k", 16), *(() if shared else (("v", 16),))) + ): + output = ir.Value( + name=name, + shape=ir.Shape([2, "S", width]), + type=ir.TensorType(dtype), + doc_string=f"raw {name} projection", + ) + output.metadata_props["test.projection"] = name + if quantized: + shape = (width, blocks, block_size // 2) + inputs = [ + x, + parameter(f"{name}.weight", rng.integers(0, 256, shape, dtype=np.uint8)), + parameter( + f"{name}.scales", + np.full( + (width * blocks,) if flat else (width, blocks), + 0.01, + dtype=_DTYPES[dtype], + ), + ), + ] + if zero_points: + inputs.append( + parameter( + f"{name}.zero_points", + np.full( + (width * ((blocks + 1) // 2),) + if flat + else (width, (blocks + 1) // 2), + 0x88, + dtype=np.uint8, + ), + ) + ) + node = ir.Node( + "com.microsoft", + "MatMulNBits", + inputs, + ir.convenience.convert_attributes( + {"K": hidden_size, "N": width, "bits": 4, "block_size": block_size} + ), + outputs=[output], + ) + else: + transpose = transposed is True or (transposed == "mixed" and index == 0) + shape = (width, hidden_size) if transpose else (hidden_size, width) + weight = parameter( + f"{name}.weight", (rng.standard_normal(shape) * 0.1).astype(_DTYPES[dtype]) + ) + if transpose: + weight = op( + "Transpose", [weight], name + "_t", [hidden_size, width], {"perm": [1, 0]} + ) + node = ir.Node("", "MatMul", [x, weight], outputs=[output]) + nodes.append(node) + projections.append(node) + q, k = projections[0].outputs[0], projections[1].outputs[0] + v = k if shared else projections[2].outputs[0] + # Normalization branches must remain separate; raw K also feeds V when shared. + q = op( + "RMSNormalization", + [q, parameter("q.norm", np.ones(q_width, dtype=_DTYPES[dtype]))], + "q_norm", + [2, "S", q_width], + {"axis": -1}, + ) + k = op( + "RMSNormalization", + [k, parameter("k.norm", np.ones(16, dtype=_DTYPES[dtype]))], + "k_norm", + [2, "S", 16], + {"axis": -1}, + ) + mask = ir.Value(name="mask", shape=ir.Shape([2, 1, "S", "T"]), type=ir.TensorType(dtype)) + past = [ + ir.Value(name=name, shape=ir.Shape([2, 2, "P", 8]), type=ir.TensorType(dtype)) + for name in ("past_k", "past_v") + ] + outputs = [ + ir.Value(name="y", shape=ir.Shape([2, "S", q_width]), type=ir.TensorType(dtype)), + *[ + ir.Value(name=name, shape=ir.Shape([2, 2, "T", 8]), type=ir.TensorType(dtype)) + for name in ("present_k", "present_v") + ], + ] + attention = ir.Node( + "", + "Attention", + [q, k, v, mask, *past], + ir.convenience.convert_attributes( + {"q_num_heads": q_width // 8, "kv_num_heads": 2, "is_causal": 1} + ), + outputs=outputs, + ) + nodes.append(attention) + graph = ir.Graph( + [x, mask, *past], + [*outputs, *[node.outputs[0] for node in projections]], + nodes=nodes, + initializers=initializers, + opset_imports={"": 24, "com.microsoft": 1}, + ) + return ir.Model(graph, ir_version=10), weights, projections, attention + + +def _pass(quantized): + return pack_matmul_nbits_qkv_pass() if quantized else pack_qkv_projections_pass() + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("shared", [False, True]) +@pytest.mark.parametrize("dtype", list(_DTYPES)) +def test_preserves_branches_metadata_and_is_idempotent(quantized, shared, dtype): + model, _, projections, attention = _model(quantized=quantized, shared=shared, dtype=dtype) + original_inputs = list(attention.inputs) + assert _pass(quantized)(model).modified + assert list(attention.inputs[:2]) == original_inputs[:2] # Norm branches unchanged. + assert list(attention.inputs[3:]) == original_inputs[3:] # Mask/cache unchanged. + split = next(node for node in model.graph if node.op_type == "Split") + assert len(split.outputs) == (2 if shared else 3) + for output in split.outputs: + assert output.dtype == dtype + assert output.metadata_props["test.projection"] == output.name + assert output.doc_string == f"raw {output.name} projection" + assert all(node not in model.graph for node in projections) + assert not _pass(quantized)(model).modified + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("shared", [False, True]) +@pytest.mark.parametrize("transposed", [False, True, "mixed"]) +def test_cpu_prefill_and_reused_cache_parity(quantized, shared, transposed): + original, _, _, _ = _model(quantized=quantized, shared=shared, transposed=transposed) + packed, _, _, _ = _model(quantized=quantized, shared=shared, transposed=transposed) + assert _pass(quantized)(packed).modified + for model in (original, packed): + fold_initializers_after_weights(model) + sessions = [OnnxModelSession(model, device="cpu") for model in (original, packed)] + rng = np.random.default_rng(7) + caches = [[np.zeros((2, 2, 0, 8), dtype=np.float32) for _ in range(2)] for _ in range(2)] + for length in (3, 1, 1): + total = caches[0][0].shape[2] + length + x = rng.standard_normal((2, length, 64)).astype(np.float32) + mask = np.zeros((2, 1, length, total), dtype=np.float32) + mask[1, :, :, -1] = -10000 + results = [ + session.run({"x": x, "mask": mask, "past_k": cache[0], "past_v": cache[1]}) + for session, cache in zip(sessions, caches, strict=True) + ] + for key in results[0]: + np.testing.assert_allclose(results[0][key], results[1][key], atol=1e-5, rtol=1e-5) + caches = [[result["present_k"], result["present_v"]] for result in results] + + +@pytest.mark.parametrize("flat", [False, True]) +@pytest.mark.parametrize("zero_points", [False, True]) +@pytest.mark.parametrize("dtype", list(_DTYPES)) +@pytest.mark.parametrize("hidden_size", [64, 65]) +def test_int4_payload_is_exact_without_requantization(flat, zero_points, dtype, hidden_size): + model, weights, _, _ = _model( + quantized=True, + flat=flat, + zero_points=zero_points, + dtype=dtype, + hidden_size=hidden_size, + ) + assert pack_matmul_nbits_qkv_pass()(model).modified + fold_initializers_after_weights(model) + projection = next(node for node in model.graph if node.op_type == "MatMulNBits") + for index, suffix in enumerate( + ("weight", "scales", *(("zero_points",) if zero_points else ())), 1 + ): + expected = np.concatenate([weights[f"{name}.{suffix}"] for name in "qkv"], axis=0) + np.testing.assert_array_equal(projection.inputs[index].const_value.numpy(), expected) + assert projection.attributes["N"].as_int() == 64 + + +@pytest.mark.parametrize("zero_points", [False, True]) +def test_registry_owned_fallback_remains_packable(zero_points): + model, weights, projections, _ = _model(quantized=True, zero_points=zero_points) + register_function_bodies(model) + identifier = projections[0].op_identifier() + assert identifier[2] == ("" if zero_points else "default_zero_points") + function = model.functions[identifier] + assert pack_matmul_nbits_qkv_pass()(model).modified + assert model.functions[identifier] is function + fold_initializers_after_weights(model) + packed = next(node for node in model.graph if node.op_type == "MatMulNBits") + assert packed.op_identifier() == identifier + np.testing.assert_array_equal( + packed.inputs[1].const_value.numpy(), + np.concatenate([weights[f"{name}.weight"] for name in "qkv"], axis=0), + ) + + +@pytest.mark.parametrize("zero_points", [False, True]) +def test_caller_factory_function_and_calls_are_not_packed_or_specialized(zero_points): + model, _, projections, _ = _model(quantized=True, zero_points=zero_points) + function = matmul_nbits() + model.functions[function.identifier()] = function + register_function_bodies(model) + assert model.functions[function.identifier()] is function + assert all(node.overload == "" for node in projections) + assert not pack_matmul_nbits_qkv_pass()(model).modified + assert model.functions[function.identifier()] is function + assert all(node.overload == "" for node in projections) + + +@pytest.mark.parametrize("transposed", [False, True, "mixed"]) +@pytest.mark.parametrize("quantized", [False, True]) +def test_deferred_binding_and_package_roundtrip(tmp_path, transposed, quantized): + model, weights, _, _ = _model(quantized=quantized, transposed=transposed, deferred=True) + names = set(model.graph.initializers) + assert _pass(quantized)(model).modified + assert names <= set(model.graph.initializers) + package = ModelPackage({"model": model}) + package.apply_weights({name: torch.from_numpy(value) for name, value in weights.items()}) + package.save(tmp_path / "package") + restored = ModelPackage.load(tmp_path / "package") + assert not any(node.op_type == "Concat" for node in restored["model"].graph) + assert any(node.op_type == "Split" for node in restored["model"].graph) + + +@pytest.mark.parametrize("dtype", list(_DTYPES)) +@pytest.mark.parametrize("transposed", [False, True, "mixed"]) +@pytest.mark.parametrize("shared", [False, True]) +def test_float_payload_is_exact_for_all_orientations(dtype, transposed, shared): + model, weights, _, _ = _model(dtype=dtype, transposed=transposed, shared=shared) + assert pack_qkv_projections_pass()(model).modified + fold_initializers_after_weights(model) + packed = next(node for node in model.graph if node.op_type == "MatMul").inputs[1] + logical = [] + for index, name in enumerate(("q", "k") if shared else ("q", "k", "v")): + weight = weights[f"{name}.weight"] + if transposed is True or (transposed == "mixed" and index == 0): + weight = weight.T + logical.append(weight) + assert packed.dtype == dtype + np.testing.assert_array_equal(packed.const_value.numpy(), np.concatenate(logical, axis=1)) + + +@pytest.mark.parametrize( + "case", + [ + "bits", + "accuracy", + "block", + "unknown", + "scale_dtype", + "zp", + "g_idx", + "overridable", + "different_input", + "overlap", + "custom_body", + ], +) +def test_incompatible_int4_is_left_unchanged(case): + model, _, projections, attention = _model(quantized=True) + node = projections[0] + if case in {"bits", "accuracy", "block", "unknown"}: + key, value = { + "bits": ("bits", 8), + "accuracy": ("accuracy_level", 4), + "block": ("block_size", 64), + "unknown": ("weight_prepacked", 1), + }[case] + node.attributes[key] = ir.AttrInt64(key, value) + elif case == "scale_dtype": + node.inputs[2].type = ir.TensorType(ir.DataType.FLOAT16) + elif case == "zp": + node.replace_input_with(3, None) + elif case == "g_idx": + node.resize_inputs(5) + node.replace_input_with(4, model.graph.inputs[0]) + elif case == "overridable": + model.graph.inputs.append(node.inputs[1]) + elif case == "different_input": + node.replace_input_with(0, model.graph.inputs[1]) + elif case == "custom_body": + function = matmul_nbits() + function.metadata_props.clear() + model.functions[function.identifier()] = function + else: + model.graph.append( + ir.Node( + "", + "Attention", + [ + attention.inputs[0], + attention.inputs[1], + attention.inputs[1], + *attention.inputs[3:], + ], + dict(attention.attributes), + outputs=[ir.Value(name="other")], + ) + ) + assert not pack_matmul_nbits_qkv_pass()(model).modified + + +@pytest.mark.parametrize( + "case", + [ + "attribute_type", + "zero_k", + "zero_n", + "block_small", + "block_non_power_two", + "accuracy_presence", + "unknown_overload", + "mixed_overload", + "default_overload_with_zp", + "rank_two", + "unknown_hidden", + "output_shape", + "output_dtype", + "weight_shape", + "weight_dtype", + "tensor_shape", + "tensor_dtype", + "computed_weight", + "foreign_weight", + "fused_bias", + "contrib_opset", + "outside_consumer", + "no_anchor", + ], +) +def test_int4_correctness_guards_are_complete_noops(case): + model, _, projections, attention = _model(quantized=True) + q = projections[0] + if case == "attribute_type": + q.attributes["K"] = ir.AttrFloat32("K", 64.0) + elif case in {"zero_k", "zero_n", "block_small", "block_non_power_two"}: + key, value = { + "zero_k": ("K", 0), + "zero_n": ("N", 0), + "block_small": ("block_size", 8), + "block_non_power_two": ("block_size", 24), + }[case] + q.attributes[key] = ir.AttrInt64(key, value) + elif case == "accuracy_presence": + q.attributes["accuracy_level"] = ir.AttrInt64("accuracy_level", 0) + elif case in {"unknown_overload", "mixed_overload", "default_overload_with_zp"}: + q.overload = "unknown" if case == "unknown_overload" else "default_zero_points" + if case == "default_overload_with_zp": + for node in projections: + node.overload = "default_zero_points" + elif case in {"rank_two", "unknown_hidden"}: + q.inputs[0].shape = ir.Shape([2, 64] if case == "rank_two" else [2, "S", "K"]) + elif case == "output_shape": + q.outputs[0].shape = ir.Shape([2, "S", 31]) + elif case == "output_dtype": + q.outputs[0].dtype = ir.DataType.FLOAT16 + elif case == "weight_shape": + q.inputs[1].shape = ir.Shape([32, 32]) + elif case == "weight_dtype": + q.inputs[1].dtype = ir.DataType.INT8 + elif case == "tensor_shape": + q.inputs[1].const_value = ir.tensor(np.zeros((32, 2, 15), dtype=np.uint8)) + elif case == "tensor_dtype": + q.inputs[2].const_value = ir.tensor(np.ones((32, 2), dtype=np.float16)) + elif case == "computed_weight": + weight = q.inputs[1] + output = ir.Value(name="computed", shape=weight.shape, type=weight.type) + model.graph.insert_before(q, ir.Node("", "Identity", [weight], outputs=[output])) + q.replace_input_with(1, output) + elif case == "foreign_weight": + del model.graph.initializers[q.inputs[1].name] + elif case == "fused_bias": + q.resize_inputs(6) + q.replace_input_with(5, q.inputs[2]) + elif case == "contrib_opset": + model.graph.opset_imports["com.microsoft"] = 2 + elif case == "outside_consumer": + model.graph.append(ir.Node("", "Identity", [q.inputs[1]], num_outputs=1)) + elif case == "no_anchor": + # Detach the public Attention outputs before safely removing their + # producer. Keep raw projection outputs to exercise the no-anchor guard. + for output in attention.outputs: + model.graph.outputs.remove(output) + model.graph.remove(attention, safe=True) + before = tuple(model.graph) + inputs = {node: tuple(node.inputs) for node in before} + initializers = dict(model.graph.initializers) + assert not pack_matmul_nbits_qkv_pass()(model).modified + assert tuple(model.graph) == before + assert {node: tuple(node.inputs) for node in before} == inputs + assert dict(model.graph.initializers) == initializers + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("opset", [None, 12, 13]) +def test_projection_split_opset_boundary(quantized, opset): + model, _, _, _ = _model(quantized=quantized) + if opset is None: + del model.graph.opset_imports[""] + else: + model.graph.opset_imports[""] = opset + before = tuple(model.graph) + assert _pass(quantized)(model).modified == (opset == 13) + if opset != 13: + assert tuple(model.graph) == before + + +@pytest.mark.parametrize("quantized", [False, True]) +def test_projection_matching_does_not_read_lazy_parameters(quantized): + model, _, _, _ = _model(quantized=quantized) + + def unexpected_read(): + raise AssertionError("projection matching must not read weight payloads") + + for value in model.graph.initializers.values(): + value.const_value = ir.LazyTensor( + unexpected_read, dtype=value.dtype, shape=value.shape, name=value.name + ) + assert _pass(quantized)(model).modified + assert not _pass(quantized)(model).modified + + +@pytest.mark.parametrize("transposed", [False, True, "mixed"]) +@pytest.mark.parametrize("shared", [False, True]) +@pytest.mark.parametrize( + "metadata_key", + [QKV_PROJECTION_SOURCES_METADATA_KEY, LEGACY_QKV_PROJECTION_SOURCES_METADATA_KEY], +) +def test_peft_sources_survive_fold_and_serialization( + tmp_path, transposed, shared, metadata_key +): + model, weights, _, _ = _model(transposed=transposed, shared=shared) + assert pack_qkv_projections_pass()(model).modified + split = next(node for node in model.graph if node.op_type == "Split") + assert QKV_PROJECTION_SOURCES_METADATA_KEY == "pkg.mobius.qkv_projection_sources" + metadata = split.metadata_props.pop(QKV_PROJECTION_SOURCES_METADATA_KEY) + split.metadata_props[metadata_key] = metadata + fold_initializers_after_weights(model) + package = ModelPackage({"model": model}) + package.save(tmp_path / "base") + restored = ModelPackage.load(tmp_path / "base")["model"] + assert _resolve_initializer_name(restored, "q") == "q.weight" + for index, name in enumerate(("q", "k") if shared else ("q", "k", "v")): + original = weights[f"{name}.weight"] + if transposed is False or (transposed == "mixed" and index != 0): + original = original.T + np.testing.assert_array_equal( + restored.graph.initializers[f"{name}.weight"].const_value.numpy(), original + ) + assert not any(metadata_key in node.metadata_props for node in restored.graph) + + +def test_invalid_peft_source_metadata_fails_closed(): + model, _, _, _ = _model() + pack_qkv_projections_pass()(model) + split = next(node for node in model.graph if node.op_type == "Split") + split.metadata_props[LEGACY_QKV_PROJECTION_SOURCES_METADATA_KEY] = split.metadata_props[ + QKV_PROJECTION_SOURCES_METADATA_KEY + ] + split.metadata_props[QKV_PROJECTION_SOURCES_METADATA_KEY] = json.dumps({"version": 2}) + with pytest.raises(ValueError, match="unsupported"): + _resolve_initializer_name(model, "q") + + +@pytest.mark.parametrize( + "case", + [ + "split_axis", + "projection_domain", + "projection_attributes", + "public_activation", + "shared_activation", + "missing_input", + "unloaded_weight", + "overridable_weight", + "quantized_weight", + "weight_width", + ], +) +def test_invalid_packed_peft_projection_is_rejected_before_mutation(case): + model, _, _, _ = _model() + assert pack_qkv_projections_pass()(model).modified + fold_initializers_after_weights(model) + split = next(node for node in model.graph if node.op_type == "Split") + packed = split.inputs[0] + projection = packed.producer() + weight = projection.inputs[1] + if case == "split_axis": + split.attributes["axis"] = ir.AttrInt64("axis", 0) + elif case == "projection_domain": + projection.domain = "custom" + elif case == "projection_attributes": + projection.attributes["custom"] = ir.AttrInt64("custom", 1) + elif case == "public_activation": + model.graph.outputs.append(packed) + elif case == "shared_activation": + model.graph.append(ir.Node("", "Identity", [packed], num_outputs=1)) + elif case == "missing_input": + projection.replace_input_with(0, None) + elif case == "unloaded_weight": + weight.const_value = None + elif case == "overridable_weight": + model.graph.inputs.append(weight) + elif case == "quantized_weight": + weight.dtype = ir.DataType.UINT8 + elif case == "weight_width": + weight.shape = ir.Shape([64, 1]) + before = tuple(model.graph) + initializers = dict(model.graph.initializers) + inputs = {node: tuple(node.inputs) for node in before} + message = ( + "unsupported packed projection" + if case + in { + "split_axis", + "projection_domain", + "projection_attributes", + "public_activation", + "shared_activation", + } + else "requires a loaded, folded packed projection weight" + ) + with pytest.raises(ValueError, match=message): + _resolve_initializer_name(model, "q") + assert tuple(model.graph) == before + assert dict(model.graph.initializers) == initializers + assert {node: tuple(node.inputs) for node in before} == inputs + + +def _write_peft_adapter(directory, modules): + directory.mkdir() + (directory / "adapter_config.json").write_text( + json.dumps( + { + "base_model_name_or_path": "offline/tiny", + "r": 2, + "lora_alpha": 4, + "target_modules": list(modules), + } + ) + ) + tensors = {} + for name, (output_size, input_size) in modules.items(): + tensors[f"base_model.model.{name}.lora_A.weight"] = np.ones( + (2, input_size), dtype=np.float32 + ) + tensors[f"base_model.model.{name}.lora_B.weight"] = np.ones( + (output_size, 2), dtype=np.float32 + ) + save_file(tensors, directory / "adapter_model.safetensors") + + +def test_float_initializer_with_outside_consumer_is_not_packed(): + model, _, projections, _ = _model(transposed=False) + external = ir.Value( + name="exposed_weight", + shape=projections[0].inputs[1].shape, + type=projections[0].inputs[1].type, + ) + model.graph.append(ir.Node("", "Identity", [projections[0].inputs[1]], outputs=[external])) + model.graph.outputs.append(external) + assert not pack_qkv_projections_pass()(model).modified + + +@pytest.mark.parametrize("sharing", ["initializer", "logical_view"]) +def test_tied_projection_groups_decline_packing_and_remain_publicly_attachable( + tmp_path, sharing +): + model, weights, projections, attention = _model(transposed=True) + graph = model.graph + first_q_view = projections[0].inputs[1] + source = first_q_view.producer().inputs[0] + second_outputs = [] + for name, width, original in ( + ("second_q", 32, "q"), + ("second_key", 16, "k"), + ("second_value", 16, "v"), + ): + if original == "q": + parameter = source + else: + data = weights[f"{original}.weight"].copy() + parameter = ir.Value( + name=f"{name}.weight", + shape=ir.Shape(data.shape), + type=ir.TensorType(ir.DataType.FLOAT), + const_value=ir.tensor(data), + ) + graph.initializers.add(parameter) + logical = parameter + if original == "q" and sharing == "logical_view": + logical = first_q_view + else: + logical = ir.Value( + shape=ir.Shape([64, width]), type=ir.TensorType(ir.DataType.FLOAT) + ) + graph.append( + ir.Node( + "", + "Transpose", + [parameter], + {"perm": ir.AttrInt64s("perm", [1, 0])}, + outputs=[logical], + ) + ) + output = ir.Value( + name=name, + shape=ir.Shape([2, "S", width]), + type=ir.TensorType(ir.DataType.FLOAT), + ) + graph.append(ir.Node("", "MatMul", [graph.inputs[0], logical], outputs=[output])) + second_outputs.append(output) + outputs = [ + ir.Value(name=f"second_{value.name}", shape=value.shape, type=value.type) + for value in attention.outputs + ] + graph.append( + ir.Node( + "", + "Attention", + [*second_outputs, *attention.inputs[3:]], + dict(attention.attributes), + outputs=outputs, + ) + ) + graph.outputs.extend(outputs) + packing_result = pack_qkv_projections_pass()(model) + fold_initializers_after_weights(model) + package = ModelPackage({"model": model}) + package.save(str(tmp_path / "base"), max_workers=1, progress_bar=False) + restored = ModelPackage.load(str(tmp_path / "base")) + adapter = tmp_path / "adapter" + # Target BOTH groups. Before the outside-consumer guard, packing hid the + # same q.weight in two metadata records: restoring k created q.weight, then + # restoring second_key failed with "conflicting source weights". That + # public attach failure precedes any later save/ownership guard. + _write_peft_adapter(adapter, {"k": (16, 64), "second_key": (16, 64)}) + artifact = attach_peft_adapter(restored, adapter, name="tied") + assert len(artifact.weights) == 2 + assert {weight.target.parameter for weight in artifact.weights} == { + "k.weight", + "second_key.weight", + } + assert not packing_result.modified + restored.adapter_target_manifest.validate(restored) + restored.save(str(tmp_path / "attached"), max_workers=1, progress_bar=False) + roundtrip = ModelPackage.load(str(tmp_path / "attached")) + assert not any( + QKV_PROJECTION_SOURCES_METADATA_KEY in node.metadata_props + for node in roundtrip["model"].graph + ) + + +@pytest.mark.parametrize("refold", [False, True]) +def test_lazy_restoration_rejects_same_path_save_before_read_or_mutation( + tmp_path, monkeypatch, refold +): + model, weights, _, _ = _model(transposed=True) + expected = {name: value.copy() for name, value in weights.items()} + assert pack_qkv_projections_pass()(model).modified + fold_initializers_after_weights(model) + base = tmp_path / "base" + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + before = {path.name: path.read_bytes() for path in base.iterdir()} + restored = ModelPackage.load(str(base)) + assert _resolve_initializer_name(restored["model"], "q") == "q.weight" + assert isinstance( + restored["model"].graph.initializers["q.weight"].const_value, ir.LazyTensor + ) + + def unexpected_read(*args, **kwargs): + raise AssertionError("alias guard must run before reading lazy source tensors") + + with monkeypatch.context() as patch: + patch.setattr(ir.ExternalTensor, "numpy", unexpected_read) + if refold: + fold_initializers_after_weights(restored["model"]) + with pytest.raises(ValueError, match="aliases lazy source file"): + restored.save(str(base), max_workers=1, progress_bar=False) + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + restored.save(str(tmp_path / "safe"), max_workers=1, progress_bar=False) + roundtrip = ModelPackage.load(str(tmp_path / "safe")) + for name in ("q", "k", "v"): + parameter = f"{name}.weight_t" if refold else f"{name}.weight" + actual = roundtrip["model"].graph.initializers[parameter].const_value.numpy() + np.testing.assert_array_equal( + actual.T if refold else actual, + expected[f"{name}.weight"], + ) + + +@pytest.mark.parametrize("alias", ["same_path", "hardlink", "symlink"]) +def test_public_peft_attach_preserves_lazy_source_on_save_alias(tmp_path, alias): + model, weights, _, _ = _model(transposed=True) + expected = {name: value.copy() for name, value in weights.items()} + assert pack_qkv_projections_pass()(model).modified + fold_initializers_after_weights(model) + base = tmp_path / "base" + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + before = {path.name: path.read_bytes() for path in base.iterdir()} + restored = ModelPackage.load(str(base)) + adapter = tmp_path / "adapter" + _write_peft_adapter(adapter, {"q": (32, 64)}) + attach_peft_adapter(restored, adapter, name="packed") + destination = base + if alias != "same_path": + destination = tmp_path / "alias" + destination.mkdir() + target = destination / "model.onnx.data" + source = base / "model.onnx.data" + if alias == "hardlink": + target.hardlink_to(source) + else: + target.symlink_to(source) + with pytest.raises(ValueError, match="aliases lazy source file"): + restored.save(str(destination), max_workers=1, progress_bar=False) + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + if alias != "same_path": + assert list(destination.iterdir()) == [destination / "model.onnx.data"] + restored.save(str(tmp_path / "safe"), max_workers=1, progress_bar=False) + roundtrip = ModelPackage.load(str(tmp_path / "safe")) + for name in ("q", "k", "v"): + np.testing.assert_array_equal( + roundtrip["model"].graph.initializers[f"{name}.weight"].const_value.numpy(), + expected[f"{name}.weight"], + ) + + +@pytest.mark.parametrize( + "artifact_file", + [ + "adapter.json", + "adapter_model.safetensors", + "adapter_config.json", + "adapter.onnx_adapter", + ], +) +@pytest.mark.parametrize("alias", ["hardlink", "symlink"]) +def test_adapter_artifact_alias_rejected_before_any_writes(tmp_path, artifact_file, alias): + model, weights, _, _ = _model() + assert pack_qkv_projections_pass()(model).modified + fold_initializers_after_weights(model) + base = tmp_path / "base" + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + before = {path.name: path.read_bytes() for path in base.iterdir()} + restored = ModelPackage.load(str(base)) + adapter = tmp_path / "adapter" + _write_peft_adapter(adapter, {"q": (32, 64)}) + artifact = attach_peft_adapter(restored, adapter, name="packed") + restored.adapter_service_options = dataclasses.replace( + restored.adapter_service_options, preserve_source_format=True + ) + if artifact_file == "adapter.onnx_adapter": + # Persistence copies this format verbatim; no runtime/parser is needed + # to exercise the exact destination naming and pre-write alias guard. + source = tmp_path / "original.onnx_adapter" + source.write_bytes(b"opaque adapter artifact") + restored.adapter_artifacts["packed"] = dataclasses.replace( + artifact, source=AdapterSource("onnx_adapter", path=str(source)) + ) + destination = tmp_path / "destination" + artifact_directory = destination / "adapters" / "packed" + artifact_directory.mkdir(parents=True) + target = artifact_directory / artifact_file + source = base / "model.onnx.data" + if alias == "hardlink": + target.hardlink_to(source) + else: + target.symlink_to(source) + existing = {path.relative_to(destination) for path in destination.rglob("*")} + with pytest.raises(ValueError, match="aliases lazy source file"): + restored.save(str(destination), max_workers=1, progress_bar=False) + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + assert existing == {path.relative_to(destination) for path in destination.rglob("*")} + assert not (destination / "model.onnx").exists() + safe = tmp_path / "safe" + restored.save(str(safe), max_workers=1, progress_bar=False) + assert (safe / "adapters" / "packed" / artifact_file).is_file() + roundtrip = ModelPackage.load(str(safe)) + for name in ("q", "k", "v"): + np.testing.assert_array_equal( + roundtrip["model"].graph.initializers[f"{name}.weight"].const_value.numpy(), + weights[f"{name}.weight"], + ) + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + + +def test_disabled_adapter_artifacts_do_not_block_safe_model_save(tmp_path): + model, _, _, _ = _model() + assert pack_qkv_projections_pass()(model).modified + fold_initializers_after_weights(model) + base = tmp_path / "base" + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + restored = ModelPackage.load(str(base)) + adapter = tmp_path / "adapter" + _write_peft_adapter(adapter, {"q": (32, 64)}) + attach_peft_adapter(restored, adapter, name="packed") + output = tmp_path / "output" + artifact_directory = output / "adapters" / "packed" + artifact_directory.mkdir(parents=True) + target = artifact_directory / "adapter.json" + target.hardlink_to(base / "model.onnx.data") + before = target.read_bytes() + restored.save( + str(output), include_adapter_artifacts=False, max_workers=1, progress_bar=False + ) + assert target.read_bytes() == before + assert (output / "model.onnx").is_file() + + +def test_ep_dtype_role_gating(): + for ep in ep_registry: + for role in ("decoder", "masked-decoder", "encoder"): + for dtype in _DTYPES: + caps = get_ep(ep) + stages, _ = _get_optimization_passes(caps, dtype, role) + names = {name for name, _ in stages} + assert not caps.matmul_nbits_qkv_pack_dtypes + assert not caps.qkv_projection_pack_dtypes + assert "PackMatMulNBitsQKV" not in names + assert "PackQKVProjections" not in names + assert ("PackQKV" in names) == ( + role == "decoder" and dtype in caps.qkv_pack_dtypes + ) + + +@pytest.mark.parametrize("dtype", list(_DTYPES)) +@pytest.mark.parametrize("role", ["decoder", "masked-decoder", "encoder"]) +def test_int4_capability_requires_explicit_opt_in(dtype, role): + caps = dataclasses.replace(get_ep("cuda"), matmul_nbits_qkv_pack_dtypes=frozenset({dtype})) + stages, _ = _get_optimization_passes(caps, dtype, role) + assert ("PackMatMulNBitsQKV" in {name for name, _ in stages}) == (role != "encoder") + assert "PackQKVProjections" not in {name for name, _ in stages} + + +@pytest.mark.parametrize("dtype", list(_DTYPES)) +@pytest.mark.parametrize("role", ["decoder", "masked-decoder", "encoder"]) +@pytest.mark.parametrize("legacy_packing", [False, True]) +def test_float_capability_requires_independent_opt_in(dtype, role, legacy_packing): + caps = dataclasses.replace( + get_ep("cuda"), + qkv_pack_dtypes=frozenset({dtype}) if legacy_packing else frozenset(), + qkv_projection_pack_dtypes=frozenset({dtype}), + ) + for candidate in _DTYPES: + stages, _ = _get_optimization_passes(caps, candidate, role) + names = {name for name, _ in stages} + assert ("PackQKVProjections" in names) == (role != "encoder" and candidate == dtype) + assert ("PackQKV" in names) == ( + role == "decoder" and legacy_packing and candidate == dtype + ) + assert "PackMatMulNBitsQKV" not in names + + +def test_new_capability_preserves_existing_positional_arguments(): + caps = EpCapabilities("legacy", frozenset(), frozenset(), False) + assert not caps.supports_fused_rope + assert not caps.matmul_nbits_qkv_pack_dtypes + assert not caps.qkv_projection_pack_dtypes + assert [field.name for field in dataclasses.fields(EpCapabilities)[-2:]] == [ + "matmul_nbits_qkv_pack_dtypes", + "qkv_projection_pack_dtypes", + ] + + +@pytest.mark.parametrize("trace", [False, True]) +@pytest.mark.parametrize("opt_in", [False, True]) +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("dtype", list(_DTYPES)) +def test_generic_pipeline_preserves_arbitrary_mask_contract( + monkeypatch, trace, opt_in, quantized, dtype +): + model, _, _, _ = _model(quantized=quantized, deferred=True, dtype=dtype) + declare_arbitrary_attention_mask(model.graph) + caps = dataclasses.replace( + get_ep("cuda"), + name="projection-test", + matmul_nbits_qkv_pack_dtypes=frozenset({dtype}) + if opt_in and quantized + else frozenset(), + qkv_projection_pack_dtypes=frozenset({dtype}) + if opt_in and not quantized + else frozenset(), + ) + monkeypatch.setitem(ep_registry._entries, caps.name, caps) + optimize_model(model, ep=caps.name, dtype=dtype, trace=trace) + attention = next(node for node in model.graph if node.op_type == "Attention") + assert attention.inputs[3].name == "mask" + assert not any(node.op_type == "GroupQueryAttention" for node in model.graph) + projection_type = "MatMulNBits" if quantized else "MatMul" + assert sum(node.op_type == projection_type for node in model.graph) == (1 if opt_in else 3) + + +@pytest.mark.parametrize("quantized", [False, True]) +@pytest.mark.parametrize("transposed", [False, True, "mixed"]) +@pytest.mark.parametrize("sharing", ["source", "projection"]) +def test_repeated_qk_source_declines_before_graph_mutation(quantized, transposed, sharing): + model, _, projections, attention = _model( + quantized=quantized, transposed=transposed, q_width=16 + ) + # Use equal-width K/V sources for Q as well, leaving an independent V. + q, k, _ = projections + if sharing == "projection": + attention.replace_input_with(0, attention.inputs[1]) + else: + q.replace_input_with(1, k.inputs[1]) + if quantized: + q.replace_input_with(2, k.inputs[2]) + q.replace_input_with(3, k.inputs[3]) + q.attributes["N"] = ir.AttrInt64("N", 16) + q.outputs[0].shape = k.outputs[0].shape + before_nodes = tuple(model.graph) + before_inputs = {node: tuple(node.inputs) for node in before_nodes} + before_initializers = dict(model.graph.initializers) + assert not _pass(quantized)(model).modified + assert tuple(model.graph) == before_nodes + assert {node: tuple(node.inputs) for node in model.graph} == before_inputs + assert dict(model.graph.initializers) == before_initializers + assert not any(node.op_type == "Concat" for node in model.graph) + + +@pytest.mark.parametrize("sharing", ["same_transpose", "separate_transposes"]) +def test_repeated_qk_weight_keeps_independent_v_publicly_attachable(tmp_path, sharing): + model, _, projections, _ = _model(transposed=True, q_width=16) + # Separate named Q/K nodes, but both consume k.weight. + q, k, _ = projections + if sharing == "same_transpose": + q.replace_input_with(1, k.inputs[1]) + else: + q.inputs[1].producer().replace_input_with(0, k.inputs[1].producer().inputs[0]) + q.inputs[1].shape = k.inputs[1].shape + q.outputs[0].shape = k.outputs[0].shape + nodes = tuple(model.graph) + assert not pack_qkv_projections_pass()(model).modified + assert tuple(model.graph) == nodes + fold_initializers_after_weights(model) + base = tmp_path / "base" + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + restored = ModelPackage.load(str(base)) + adapter = tmp_path / "adapter" + _write_peft_adapter(adapter, {"v": (16, 64)}) + artifact = attach_peft_adapter(restored, adapter, name="independent_v") + assert {weight.target.parameter for weight in artifact.weights} == {"v.weight"} + restored.adapter_target_manifest.validate(restored) + restored.save(str(tmp_path / "attached"), max_workers=1, progress_bar=False) + + +@pytest.mark.parametrize("transposed", [False, True, "mixed"]) +def test_independent_qk_adapters_restore_publicly_after_packing(tmp_path, transposed): + independent, _, _, _ = _model(transposed=transposed) + assert pack_qkv_projections_pass()(independent).modified + fold_initializers_after_weights(independent) + base = tmp_path / "base" + ModelPackage({"model": independent}).save(str(base), max_workers=1, progress_bar=False) + restored = ModelPackage.load(str(base)) + adapter = tmp_path / "adapter" + _write_peft_adapter(adapter, {"q": (32, 64), "k": (16, 64)}) + artifact = attach_peft_adapter(restored, adapter, name="independent") + assert {weight.target.parameter for weight in artifact.weights} == { + "q.weight", + "k.weight", + } + restored.adapter_target_manifest.validate(restored) + + +@pytest.mark.parametrize("fold", [False, True]) +def test_deferred_quantized_concat_binding_validation(fold): + model, weights, _, _ = _model(quantized=True, deferred=True) + assert pack_matmul_nbits_qkv_pass()(model).modified + config = SimpleNamespace() + with pytest.raises(ValueError, match=r"unbound.*parameter"): + validate_quantized_component_bindings({"model": model}, config) + package = ModelPackage({"model": model}) + package.apply_weights( + {name: torch.from_numpy(data) for name, data in weights.items()}, + fold_constants=fold, + ) + validate_quantized_component_bindings(package, config) + if not fold: + package["model"].graph.initializers["k.scales"].const_value = None + with pytest.raises(ValueError, match=r"unbound MatMulNBits parameter 'k.scales'"): + validate_quantized_component_bindings(package, config) 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(