From a8265a7873588fd60bf9815c991401d8cef70757 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Fri, 23 Oct 2026 23:33:42 +0000 Subject: [PATCH 1/7] Pass consolidation: fold layouts and pack QKV projections Consolidate generic QKV packing, lazy constant layout folding, adapter aliases, and default-off INT4 behavior with focused tests. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/ep_quickstart.md | 9 +- docs/execution_providers.md | 74 ++ src/mobius/_component_quantization.py | 31 +- src/mobius/_constants.py | 5 + src/mobius/_execution_providers.py | 8 + src/mobius/_model_package.py | 53 + src/mobius/_optimizations.py | 121 ++- src/mobius/_passes/__init__.py | 6 + src/mobius/_passes/_fold_transpose.py | 16 +- src/mobius/_passes/_fold_weight_layout.py | 217 ++++ .../_passes/_fold_weight_layout_test.py | 539 ++++++++++ src/mobius/adapter_io.py | 171 ++++ src/mobius/functions/matmul_nbits.py | 5 + src/mobius/functions/matmul_nbits_test.py | 35 + src/mobius/rewrite_rules/__init__.py | 6 + .../rewrite_rules/_qkv_projection_packing.py | 477 +++++++++ .../_qkv_projection_packing_test.py | 963 ++++++++++++++++++ 17 files changed, 2692 insertions(+), 44 deletions(-) create mode 100644 src/mobius/_passes/_fold_weight_layout.py create mode 100644 src/mobius/_passes/_fold_weight_layout_test.py create mode 100644 src/mobius/rewrite_rules/_qkv_projection_packing.py create mode 100644 src/mobius/rewrite_rules/_qkv_projection_packing_test.py diff --git a/docs/ep_quickstart.md b/docs/ep_quickstart.md index ae1fbb076..6a805f64d 100644 --- a/docs/ep_quickstart.md +++ b/docs/ep_quickstart.md @@ -101,6 +101,7 @@ print(sorted(ep_registry)) caps = get_ep("cuda") print(caps.gqa_dtypes) # frozenset({FLOAT16, BFLOAT16}) print(caps.qkv_pack_dtypes) # frozenset({FLOAT, FLOAT16, BFLOAT16}) +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 +133,12 @@ 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 INT4 projection packing requires an explicit +`matmul_nbits_qkv_pack_dtypes` capability override; no built-in EP enables it. +Existing float `PackQKV` behavior is unchanged. 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 +148,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 + float PackQKV; 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..6a56a564e 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -100,6 +100,80 @@ ops to fused equivalents: These are applied during **Stage 2: Fusion** of the optimization pipeline. +#### Projection packing and numerical limitations + +The existing float `PackQKV` GQA-input fusion is unchanged. A generic +projection-only pass also handles compatible float projections that remain +separate after that fusion, including masked decoders and raw K reused for V. +It splits the packed result **before** normalization, RoPE, attention, or cache +consumers. Sources must be static, non-overridable initializer parameters with +compatible shapes and dtypes; ambiguous repeated sources, shared weights with +outside consumers, and overlapping groups are declined before graph mutation. +No attention-mask contract is relaxed to make packing possible. + +**INT4 projection packing is experimental and default-disabled on every +built-in EP, including CUDA, for FP16 and BF16 alike.** Opt in only through +the existing capability API: + +```python +import dataclasses + +import onnx_ir as ir +from mobius import get_ep, register_ep + +register_ep( + dataclasses.replace( + get_ep("cuda"), + matmul_nbits_qkv_pack_dtypes=frozenset({ir.DataType.FLOAT16}), + ), + overwrite=True, +) +``` + +This overrides the process-local profile for subsequent builds. Stock +`--ep cuda` does not opt in. The capability is independent of float +`qkv_pack_dtypes`; activation dtype (`--dtype`), weight quantization/storage, +component precision overrides, and FP8 KV-cache flags remain independent +choices. Packing does not promote precision, dequantize, or requantize weights. + +Compatible INT4 groups share a rank-3 floating-point activation, attributes +other than output width, and canonical UINT8 blocked weights. Scales match the +activation dtype; zero points must be absent throughout or uniformly packed +UINT8. Flat and row-major scale/zero-point storage are supported. Block size +must be a power of two of at least 16; `accuracy_level` must be omitted or zero. +Caller-owned function bodies, mixed bit widths, extra nonempty inputs, and +incompatible layouts are not packed. Standard opset >= 13 is required. +The canonical implementation is `_qkv_projection_packing.py` for both float +and INT4, including the supported Q/K pair with K reused for V. + +Parameter Concats retain checkpoint names until binding. Missing parameter +leaves are validated even with `apply_weights(..., fold_constants=False)`. +Normally the post-weight pipeline folds Concats and static Reshape/Transpose +chains lazily, including layouts exposed by small integer shape-expression +folding. That generic folding is always independent of experimental packing; +it preserves dtype and does not eagerly load large weight tensors. + +**The PR754 numerical gate remains failed and unwaived.** Byte-exact packing +does not imply identical floating-point outputs: changing GEMM output width can +change CUDA kernel behavior. Representative large-dimension prefill and +cached-decode logits exceeded the existing `rtol=atol=1e-2` criterion. Bounded +MMLU/generation evidence does not resolve pointwise parity; neither tolerances +nor automatic BF16/FP16 enablement have changed. No production numerical +accuracy or speed improvement is claimed. + +For `ORT_FPA_INTB_GEMM=1`, keep exporter folding or ORT constant folding enabled: +unfolded parameter Concats can prevent runtime prepacking when both are off. +This is not a guarantee for every runtime version/configuration. + +Packed float sources retain adapter restoration metadata. Lazy restored/folded +weights keep their backing-file ownership after source initializers are pruned. +`ModelPackage.save()` checks model/shard, policy, adapter, package metadata, and +MTP destinations before writing, rejecting known same-path, hardlink, or symlink +aliases to lazy sources. Adapter files use `adapters//adapter.json`, +`adapter.onnx_adapter`, or `adapter_model.safetensors` plus `adapter_config.json` +according to persistence options. These checks are **not** full TOCTOU +protection for non-transactional saves; use a separate output directory. + ### Tier 3 — EP-specific constraints (required, correctness) Some EPs cannot execute certain ops. These are handled by **lowering rules** 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/_constants.py b/src/mobius/_constants.py index c80cc4780..45ff4e9ac 100644 --- a/src/mobius/_constants.py +++ b/src/mobius/_constants.py @@ -9,6 +9,11 @@ # 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 = "mobius.qkv_projection_sources" +"""Versioned source names and widths retained on projection-packing Split nodes.""" + +# Transient tensor.meta ownership, never serialized as ONNX metadata. +LAZY_SOURCE_FILES_META_KEY = "mobius.lazy_source_files" # --------------------------------------------------------------------------- # Static KV cache ABI diff --git a/src/mobius/_execution_providers.py b/src/mobius/_execution_providers.py index 21e88b42e..e062c5452 100644 --- a/src/mobius/_execution_providers.py +++ b/src/mobius/_execution_providers.py @@ -58,6 +58,10 @@ class EpCapabilities: 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. supports_fused_rope: ``False`` triggers SeparateRoPE + UnpackQKV lowering (DML). supports_skip_layer_norm: ``False`` expands SkipLayerNormalization / @@ -171,6 +175,10 @@ 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 + ) def __post_init__(self) -> None: if not self.supports_fused_rope and self.qkv_pack_dtypes: diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index d25c5d8b8..f3fb54897 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..e597fb217 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 Transpose, Concat, and static Reshape/Transpose +chains, independently of experimental projection packing. """ from __future__ import annotations @@ -49,6 +47,7 @@ import contextlib import dataclasses import logging +import math import warnings import onnx @@ -58,11 +57,13 @@ 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 ( FoldConcatInitializersPass, FoldTransposedInitializerPass, + FoldWeightLayoutPass, Fp8KvCachePass, RemoveDeadGraphInputsPass, ) @@ -75,7 +76,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 +378,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 +427,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 +446,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 +462,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_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 +693,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 +797,56 @@ 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 weight Transpose, Concat, and static layout chains. + + Fold small integer shape expressions first without exposing weight payloads + to the general constant folder. Preserve the shipped root Transpose naming + convention, then fold layouts and packing to a fixed point. """ - ir.passes.PassManager( - [ - common_passes.LiftConstantsToInitializersPass(), - FoldConcatInitializersPass(), - FoldTransposedInitializerPass(), - common_passes.RemoveUnusedNodesPass(), - ] - )(model) + common_passes.LiftConstantsToInitializersPass()(model) + hidden: dict[ir.Value, ir.TensorProtocol] = {} + 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) + ) + if tensor is not None and not ( + tensor.dtype == ir.DataType.INT64 + and all(isinstance(dim, int) for dim in tensor.shape) + and math.prod(tensor.shape.numpy()) <= 8192 + ): + hidden[value] = tensor + value.const_value = None + try: + # Partial evaluators also simplify Identity without reading its input. + # In particular, an output Identity transfers its public output name + # to the producer value before lazy layout folding. Size limits and + # should_fold do not gate those upstream symbolic simplifications. + onnxscript.optimizer._constant_folding.FoldConstantsPass( + shape_inference=False, input_size_limit=8192, output_size_limit=8192 + )(model) + finally: + for value, tensor in hidden.items(): + value.const_value = tensor + common_passes.LiftConstantsToInitializersPass()(model) + common_passes.RemoveUnusedNodesPass()(model) + lazy_passes = ( + FoldTransposedInitializerPass(), + FoldWeightLayoutPass(), + 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..499473ec4 100644 --- a/src/mobius/_passes/__init__.py +++ b/src/mobius/_passes/__init__.py @@ -35,6 +35,10 @@ 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. + +:class:`FoldWeightLayoutPass` + Folds static Reshape/Transpose chains into lazy initializers, including + higher-rank layout conversions of quantized weights, scales, and zero points. """ from __future__ import annotations @@ -42,11 +46,13 @@ __all__ = [ "FoldTransposedInitializerPass", "FoldConcatInitializersPass", + "FoldWeightLayoutPass", "Fp8KvCachePass", "RemoveDeadGraphInputsPass", ] from mobius._passes._fold_concat import FoldConcatInitializersPass from mobius._passes._fold_transpose import FoldTransposedInitializerPass +from mobius._passes._fold_weight_layout import FoldWeightLayoutPass from mobius._passes._fp8_kv_cache import Fp8KvCachePass from mobius._passes._remove_dead_inputs import RemoveDeadGraphInputsPass diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py index 957d57b86..c0ed635d6 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -62,7 +62,11 @@ 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 not in ("", "ai.onnx") + or node.overload + or node.op_type != "Transpose" + ): continue if len(node.inputs) != 1: @@ -73,7 +77,15 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: continue # Only fold initializer inputs (not graph inputs or computed values). - if not inp.is_initializer(): + if ( + not inp.is_initializer() + or inp.is_graph_input() + or inp.is_graph_output() + # Preserve legacy root ownership/naming even when the + # consumer is in a Loop/If body. Only child-local initializers + # need the scoped FoldWeightLayoutPass. + 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_weight_layout.py b/src/mobius/_passes/_fold_weight_layout.py new file mode 100644 index 000000000..ddefb9f5a --- /dev/null +++ b/src/mobius/_passes/_fold_weight_layout.py @@ -0,0 +1,217 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Lazily fold static weight Reshape/Transpose chains after checkpoint loading.""" + +from __future__ import annotations + +import logging +import math + +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 _reshape_shape(node: ir.Node, source_shape: tuple[int, ...]) -> tuple[int, ...] | None: + """Resolve a static shape, skip dynamic shapes, and reject invalid reshapes.""" + shape_value = node.inputs[1] + if shape_value is None or shape_value.is_graph_input(): + return None + if shape_value.is_initializer(): + shape_tensor = shape_value.const_value + else: + constant = shape_value.producer() + if ( + constant is None + or constant.domain not in ("", "ai.onnx") + or constant.op_type != "Constant" + or len(constant.attributes) != 1 + ): + return None + if "value" in constant.attributes: + shape_tensor = constant.attributes["value"].as_tensor() + elif "value_ints" in constant.attributes: + shape_tensor = ir.tensor( + constant.attributes["value_ints"].as_ints(), dtype=ir.DataType.INT64 + ) + else: + return None + if shape_tensor is None: + return None + if shape_tensor.dtype != ir.DataType.INT64 or len(shape_tensor.shape) != 1: + raise ValueError(f"Reshape {node.name!r} requires a one-dimensional INT64 shape") + requested = tuple(int(dim) for dim in shape_tensor.numpy()) + allowzero = node.attributes.get_int("allowzero", 0) + if allowzero not in (0, 1) or requested.count(-1) > 1 or any(d < -1 for d in requested): + raise ValueError(f"Invalid Reshape attributes or shape at {node.name!r}: {requested}") + resolved = list(requested) + if not allowzero: + for index, dim in enumerate(resolved): + if dim == 0: + if index >= len(source_shape): + raise ValueError(f"Reshape {node.name!r} copies a nonexistent dimension") + resolved[index] = source_shape[index] + size = math.prod(source_shape) + known_size = math.prod(dim for dim in resolved if dim != -1) + if -1 in resolved: + if known_size == 0 or size % known_size: + raise ValueError(f"Reshape {node.name!r} cannot infer dimension in {requested}") + resolved[resolved.index(-1)] = size // known_size + elif known_size != size: + raise ValueError( + f"Reshape {node.name!r} changes element count: {source_shape} to {requested}" + ) + return tuple(resolved) + + +class FoldWeightLayoutPass(ir.passes.InPlacePass): + """Fold named static Reshape/Transpose outputs into lazy initializers. + + Graphs must be topologically ordered. A single traversal folds layout chains + without reading weight payloads: only constant Reshape shape vectors are read. + Materialization preserves the source dtype and uses no arithmetic or casts. + Overridable initializers, shared sources, and dynamic shapes remain untouched. + Dead folded inputs are removed in their owning graph; run dead-code elimination + afterward to remove unused Constant shape nodes. + Malformed static layout nodes raise ValueError rather than being skipped. + """ + + def call(self, model: ir.Model) -> ir.passes.PassResult: + folded = 0 + graphs = tuple(model.graphs()) + + # Registry identity is authoritative, including parent-owned captures. + # Do not infer ownership from the graph containing the consumer node. + def owners(value: ir.Value) -> list[ir.Graph]: + return [ + scope + for scope in graphs + if value.name is not None and scope.initializers.get(value.name) is value + ] + + for graph in graphs: + for node in list(graph): + if ( + node.overload + or node.domain not in ("", "ai.onnx") + or node.op_type + not in ( + "Reshape", + "Transpose", + ) + ): + continue + expected_inputs = 2 if node.op_type == "Reshape" else 1 + if len(node.inputs) != expected_inputs or len(node.outputs) != 1: + raise ValueError(f"Malformed weight layout node {node.name!r}") + source = node.inputs[0] + if source is None or not source.is_initializer() or source.is_graph_input(): + continue + # Keep tied tables shared instead of serializing a second full layout. + if ( + len(source.uses()) != 1 + or len(owners(source)) != 1 + or any( + source in scope.inputs or source in scope.outputs for scope in graphs + ) + ): + continue + tensor = source.const_value + if tensor is None: + logger.warning( + "FoldWeightLayoutPass: skipping %r; initializer %r has no data", + node.name, + source.name, + ) + continue + output = node.outputs[0] + if not output.name or any(not isinstance(dim, int) for dim in tensor.shape): + logger.debug( + "FoldWeightLayoutPass: skipping unnamed or symbolic %r", node.name + ) + continue + source_shape = tuple(tensor.shape.numpy()) + dtype = initializer_dtype(source) + if dtype is None: + raise ValueError(f"Weight layout source {source.name!r} has no dtype") + is_transpose = node.op_type == "Transpose" + if is_transpose: + layout = tuple( + node.attributes.get_ints( + "perm", tuple(reversed(range(len(source_shape)))) + ) + ) + if sorted(layout) != list(range(len(source_shape))): + raise ValueError( + f"Invalid Transpose permutation at {node.name!r}: {layout}" + ) + shape = tuple(source_shape[index] for index in layout) + else: + resolved = _reshape_shape(node, source_shape) + if resolved is None: + logger.debug( + "FoldWeightLayoutPass: skipping dynamic shape at %r", node.name + ) + continue + layout = shape = resolved + + # Bind this iteration's values before the lazy callback runs later. + def materialize( + data: ir.TensorProtocol = tensor, + dimensions: tuple[int, ...] = layout, + transpose: bool = is_transpose, + data_type: ir.DataType = dtype, + ) -> ir.TensorProtocol: + array = data.numpy() + result = ( + array.transpose(dimensions) if transpose else array.reshape(dimensions) + ) + return ir.tensor(result, dtype=data_type) + + output_shape = ir.Shape(shape) + replacement = ir.Value( + name=output.name, + shape=output_shape, + type=ir.TensorType(dtype), + doc_string=output.doc_string, + metadata_props=dict(output.metadata_props), + ) + replacement.const_value = ir.LazyTensor( + materialize, dtype=dtype, shape=output_shape, name=output.name + ) + source_files = ( + (tensor.path,) + if isinstance(tensor, ir.ExternalTensor) + else tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) + ) + if source_files: + replacement.const_value.meta[LAZY_SOURCE_FILES_META_KEY] = tuple( + source_files + ) + model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( + (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) + ) + output.replace_all_uses_with(replacement, replace_graph_outputs=True) + inputs = tuple(node.inputs) + graph.remove(node, safe=True) + # Downstream layout nodes can now capture this lazy initializer. + graph.register_initializer(replacement) + for value in inputs: + if ( + value is not None + and value.name is not None + and not value.uses() + and not any( + value in scope.inputs or value in scope.outputs for scope in graphs + ) + ): + registries = owners(value) + if len(registries) == 1: + del registries[0].initializers[value.name] + folded += 1 + logger.debug("FoldWeightLayoutPass: folded %d weight layout nodes", folded) + return ir.passes.PassResult(model, modified=folded > 0) diff --git a/src/mobius/_passes/_fold_weight_layout_test.py b/src/mobius/_passes/_fold_weight_layout_test.py new file mode 100644 index 000000000..f6ac88cd0 --- /dev/null +++ b/src/mobius/_passes/_fold_weight_layout_test.py @@ -0,0 +1,539 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Value, laziness, and persistence coverage for static weight layout folding.""" + +from __future__ import annotations + +from pathlib import Path + +import ml_dtypes +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._optimizations import fold_initializers_after_weights +from mobius._passes import FoldWeightLayoutPass + + +def _initializer(name, data): + tensor = ir.tensor(data, name=name) + return ir.Value( + name=name, shape=tensor.shape, type=ir.TensorType(tensor.dtype), const_value=tensor + ) + + +def _model(data, *, first=(2, 2, 4, 3), perm=(0, 2, 1, 3), last=(2, 8, 3)): + weight = _initializer("weight", data) + shape1 = _initializer("shape1", np.array(first, dtype=np.int64)) + shape2 = _initializer("shape2", np.array(last, dtype=np.int64)) + reshape1 = ir.Node("", "Reshape", [weight, shape1], num_outputs=1) + transpose = ir.Node( + "", + "Transpose", + [reshape1.outputs[0]], + attributes=[ir.AttrInt64s("perm", perm)], + num_outputs=1, + ) + reshape2 = ir.Node("", "Reshape", [transpose.outputs[0], shape2], num_outputs=1) + identity = ir.Node("", "Identity", [reshape2.outputs[0]], num_outputs=1) + for node, name in zip( + (reshape1, transpose, reshape2, identity), + ("reshaped", "transposed", "layout", "output"), + strict=True, + ): + node.outputs[0].name = name + node.outputs[0].dtype = weight.dtype + graph = ir.Graph( + [], + [identity.outputs[0]], + nodes=[reshape1, transpose, reshape2, identity], + initializers=[weight, shape1, shape2], + opset_imports={"": 24}, + name="layout_test", + ) + return ir.Model(graph, ir_version=10) + + +@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, ml_dtypes.bfloat16]) +def test_quantized_layout_values_and_package_roundtrip(dtype, tmp_path): + data = np.arange(48).astype(dtype).reshape(2, 8, 3) + model = _model(data) + fold_initializers_after_weights(model) + # The shape-expression constant folder also resolves Identity symbolically, + # preserving the public output name on the eventual lazy initializer. + assert not list(model.graph) + assert set(model.graph.initializers) == {"output"} + folded = model.graph.initializers["output"] + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + expected = data.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3) + assert folded.dtype == ir.tensor(data).dtype + assert folded.const_value.numpy().tobytes() == expected.tobytes() + package = ModelPackage({"model": model}) + package.save(tmp_path, progress_bar=False) + loaded = ModelPackage.load(tmp_path) + restored = loaded["model"].graph.initializers["output"] + assert loaded["model"].graph.outputs[0] is restored + assert restored.dtype == folded.dtype + assert restored.shape == ir.Shape(expected.shape) + assert restored.const_value.numpy().tobytes() == expected.tobytes() + + +def test_bfloat16_special_value_bits_preserved(tmp_path): + bits = np.resize( + np.array([0, 0x8000, 0x7F80, 0xFF80, 0x7FC1, 0x7FFF, 1, 0x8001], dtype=np.uint16), + 48, + ).reshape(2, 8, 3) + model = _model(bits.view(ml_dtypes.bfloat16)) + fold_initializers_after_weights(model) + folded = model.graph.initializers["output"] + assert set(model.graph.initializers) == {"output"} + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + actual = folded.const_value.numpy() + expected = bits.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3) + assert actual.tobytes() == expected.tobytes() + ModelPackage({"model": model}).save(tmp_path, progress_bar=False) + restored_graph = ModelPackage.load(tmp_path)["model"].graph + restored = restored_graph.initializers["output"] + assert restored_graph.outputs[0] is restored + assert restored.dtype == ir.DataType.BFLOAT16 + assert restored.const_value.numpy().tobytes() == expected.tobytes() + + +def test_weight_payload_remains_lazy_and_pass_is_idempotent(): + data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) + model = _model(data) + source = model.graph.initializers["weight"] + calls = [] + + def read(): + calls.append(True) + return ir.tensor(data) + + source.const_value = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) + assert FoldWeightLayoutPass()(model).modified + assert not calls + common.RemoveUnusedNodesPass()(model) + assert set(model.graph.initializers) == {"layout"} + assert not FoldWeightLayoutPass()(model).modified + assert not calls + actual = model.graph.initializers["layout"].const_value.numpy() + assert len(calls) == 1 + assert ( + actual.tobytes() + == data.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3).tobytes() + ) + + +@pytest.mark.parametrize("dynamic", ["weight", "shape1"]) +@pytest.mark.parametrize("overridable", [False, True]) +def test_dynamic_or_overridable_inputs_are_not_folded(dynamic, overridable): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + value = model.graph.initializers[dynamic] + model.graph.inputs.append(value) + if not overridable: + del model.graph.initializers[dynamic] + value.const_value = None + assert not FoldWeightLayoutPass()(model).modified + assert len(model.graph) == 4 + + +def test_shared_sources_are_not_duplicated(): + data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) + model = _model(data) + source = model.graph.initializers["weight"] + extra = ir.Node("", "Identity", [source], num_outputs=1) + extra.outputs[0].name = "original_weight" + extra.outputs[0].dtype = source.dtype + model.graph.append(extra) + model.graph.outputs.append(extra.outputs[0]) + layout = next(node for node in model.graph if node.outputs[0].name == "layout") + model.graph.outputs.append(layout.outputs[0]) + assert not FoldWeightLayoutPass()(model).modified + common.RemoveUnusedNodesPass()(model) + assert [output.name for output in model.graph.outputs] == [ + "output", + "original_weight", + "layout", + ] + assert extra.inputs[0] is source + assert source.name in model.graph.initializers + assert model.graph.outputs[2] is layout.outputs[0] + assert set(model.graph.initializers) == {"weight", "shape1", "shape2"} + assert source.const_value.numpy().tobytes() == data.tobytes() + + +def test_unshared_graph_output_layout_name_survives(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + layout = next(node for node in model.graph if node.outputs[0].name == "layout") + model.graph.outputs.append(layout.outputs[0]) + assert FoldWeightLayoutPass()(model).modified + assert model.graph.outputs[1] is model.graph.initializers["layout"] + + +@pytest.mark.parametrize("op_type", ["Reshape", "Transpose"]) +def test_graph_output_source_is_not_duplicated(op_type): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + if op_type == "Transpose": + first = next(iter(model.graph)) + first.op_type = "Transpose" + first.replace_input_with(1, None) + first.resize_inputs(1) + first.attributes["perm"] = ir.AttrInt64s("perm", [2, 1, 0]) + model.graph.outputs.append(model.graph.initializers["weight"]) + assert not FoldWeightLayoutPass()(model).modified + assert "layout" not in model.graph.initializers + + +@pytest.mark.parametrize("style", ["value", "value_ints"]) +@pytest.mark.parametrize("shared", [False, True]) +def test_post_weight_pipeline_folds_constant_shape_nodes(style, shared): + data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) + model = _model(data) + if shared: + extra = ir.Node("", "Identity", [model.graph.initializers["weight"]], num_outputs=1) + extra.outputs[0].name = "shared" + model.graph.append(extra) + model.graph.outputs.append(extra.outputs[0]) + for name in ("shape1", "shape2"): + shape = model.graph.initializers[name] + attribute = ( + ir.AttrTensor("value", shape.const_value) + if style == "value" + else ir.AttrInt64s("value_ints", shape.const_value.numpy().tolist()) + ) + constant = ir.Node("", "Constant", [], attributes=[attribute], num_outputs=1) + constant.outputs[0].name = name + model.graph.insert_before(next(iter(model.graph)), constant) + shape.replace_all_uses_with(constant.outputs[0]) + del model.graph.initializers[name] + fold_initializers_after_weights(model) + if shared: + assert set(model.graph.initializers) == {"weight"} + assert sum(node.op_type == "Reshape" for node in model.graph) == 2 + assert sum(node.op_type == "Transpose" for node in model.graph) == 1 + return + assert not list(model.graph) + assert set(model.graph.initializers) == {"output"} + folded = model.graph.initializers["output"] + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + expected = data.reshape(2, 2, 4, 3).swapaxes(1, 2).reshape(2, 8, 3) + assert folded.const_value.numpy().tobytes() == expected.tobytes() + + +def test_tied_embedding_pipeline_keeps_one_table(tmp_path): + weight = _initializer("weight", np.arange(24, dtype=np.float16).reshape(8, 3)) + indices = ir.Value( + name="indices", shape=ir.Shape([2]), type=ir.TensorType(ir.DataType.INT64) + ) + gather = ir.Node("", "Gather", [weight, indices], num_outputs=1) + transpose = ir.Node( + "", "Transpose", [weight], attributes=[ir.AttrInt64s("perm", [1, 0])], num_outputs=1 + ) + matmul = ir.Node("", "MatMul", [gather.outputs[0], transpose.outputs[0]], num_outputs=1) + for node, name in ((gather, "embedding"), (transpose, "head"), (matmul, "output")): + node.outputs[0].name = name + node.outputs[0].dtype = ir.DataType.FLOAT16 + graph = ir.Graph( + [indices], + [matmul.outputs[0]], + nodes=[gather, transpose, matmul], + initializers=[weight], + opset_imports={"": 24}, + name="tied_table", + ) + model = ir.Model(graph, ir_version=10) + fold_initializers_after_weights(model) + assert set(graph.initializers) == {"weight"} + assert transpose.inputs[0] is gather.inputs[0] + ModelPackage({"model": model}).save(tmp_path, progress_bar=False) + loaded = ModelPackage.load(tmp_path)["model"].graph + assert set(loaded.initializers) == {"weight"} + assert [node.op_type for node in loaded] == ["Gather", "Transpose", "MatMul"] + + +@pytest.mark.parametrize( + ("source", "requested", "allowzero", "expected"), + [ + ((2, 3), (0, -1), 0, (2, 3)), + ((2, 3), (-1, 2), 0, (3, 2)), + ((0, 3), (0, 3), 1, (0, 3)), + ((0, 3), (-1, 3), 0, (0, 3)), + ((1,), (), 0, ()), + ], +) +def test_reshape_zero_inferred_and_scalar_shapes(source, requested, allowzero, expected): + weight = _initializer("weight", np.zeros(source, dtype=np.float16)) + shape = _initializer("shape", np.array(requested, dtype=np.int64)) + node = ir.Node( + "", + "Reshape", + [weight, shape], + attributes=[ir.AttrInt64("allowzero", allowzero)], + num_outputs=1, + ) + node.outputs[0].name = "output" + graph = ir.Graph( + [], + [node.outputs[0]], + nodes=[node], + initializers=[weight, shape], + opset_imports={"": 24}, + ) + model = ir.Model(graph, ir_version=10) + assert FoldWeightLayoutPass()(model).modified + assert model.graph.outputs[0].shape == ir.Shape(expected) + assert model.graph.outputs[0].const_value.numpy().shape == expected + + +@pytest.mark.parametrize("requested", [(0, -1), (-1, -1), (4,), (-2,)]) +def test_invalid_static_reshape_fails_explicitly(requested): + model = _model(np.zeros((0, 3), dtype=np.float16), first=requested) + with pytest.raises(ValueError, match="Reshape"): + FoldWeightLayoutPass()(model) + + +def test_default_transpose_and_missing_declared_dtype(): + data = np.arange(24, dtype=np.float16).reshape(2, 3, 4) + weight = _initializer("weight", data) + weight.type = None + node = ir.Node("", "Transpose", [weight], num_outputs=1) + node.outputs[0].name = "output" + graph = ir.Graph([], [node.outputs[0]], nodes=[node], initializers=[weight]) + model = ir.Model(graph, ir_version=10) + assert FoldWeightLayoutPass()(model).modified + output = graph.outputs[0] + assert output.dtype == ir.DataType.FLOAT16 + assert output.const_value.numpy().tobytes() == data.transpose().tobytes() + + +def test_invalid_permutation_and_conflicting_dtype_fail(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3), perm=(0, 0, 1, 2)) + with pytest.raises(ValueError, match="permutation"): + FoldWeightLayoutPass()(model) + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + model.graph.initializers["weight"].dtype = ir.DataType.FLOAT + with pytest.raises(ValueError, match="declares dtype"): + FoldWeightLayoutPass()(model) + + +def test_custom_domain_layout_op_is_not_folded(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + transpose = next(node for node in model.graph if node.op_type == "Transpose") + transpose.domain = "custom" + FoldWeightLayoutPass()(model) + assert transpose in list(model.graph) + assert transpose.outputs[0].producer() is transpose + + +def test_missing_weights_are_warned_and_left_unmodified(caplog): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + model.graph.initializers["weight"].const_value = None + assert not FoldWeightLayoutPass()(model).modified + assert "has no data" in caplog.text + + +def test_invalid_shape_tensor_dtype_fails_explicitly(): + model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) + shape = model.graph.initializers["shape1"] + shape.dtype = ir.DataType.INT32 + shape.const_value = ir.tensor(np.array([2, 2, 4, 3], dtype=np.int32)) + with pytest.raises(ValueError, match="INT64 shape"): + FoldWeightLayoutPass()(model) + + +def test_nested_graph_initializers_are_registered_in_their_scope(): + first = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)).graph + second = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)).graph + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + node = ir.Node( + "", + "If", + [condition], + attributes=[ir.AttrGraph("then_branch", first), ir.AttrGraph("else_branch", second)], + num_outputs=1, + ) + node.outputs[0].name = "result" + graph = ir.Graph([condition], [node.outputs[0]], nodes=[node], opset_imports={"": 24}) + model = ir.Model(graph, ir_version=10) + assert FoldWeightLayoutPass()(model).modified + assert not graph.initializers + for branch in (first, second): + assert [item.op_type for item in branch] == ["Identity"] + assert set(branch.initializers) == {"layout"} + + +@pytest.mark.parametrize("retained", [None, "shared", "graph_output"]) +def test_parent_capture_initializer_ownership_and_persisted_payload(tmp_path, retained): + data = np.arange(6, dtype=np.float32).reshape(2, 3) + source = _initializer("parent_weight", data) + transpose = ir.Node( + "", "Transpose", [source], [ir.AttrInt64s("perm", [1, 0])], num_outputs=1 + ) + transpose.outputs[0].name = "captured_layout" + transpose.outputs[0].shape = ir.Shape([3, 2]) + transpose.outputs[0].dtype = ir.DataType.FLOAT + identity = ir.Node("", "Identity", [transpose.outputs[0]], num_outputs=1) + identity.outputs[0].name = "branch_output" + identity.outputs[0].shape = ir.Shape([3, 2]) + identity.outputs[0].dtype = ir.DataType.FLOAT + branch = ir.Graph([], [identity.outputs[0]], nodes=[transpose, identity]) + other = _initializer("other_weight", np.zeros((3, 2), dtype=np.float32)) + other_branch = ir.Graph([], [other], nodes=[], initializers=[other]) + 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=1, + ) + select.outputs[0].name = "output" + select.outputs[0].shape = ir.Shape([3, 2]) + select.outputs[0].dtype = ir.DataType.FLOAT + graph = ir.Graph( + [condition], + [select.outputs[0]], + nodes=[select], + initializers=[source], + opset_imports={"": 24}, + ) + if retained == "shared": + extra = ir.Node("", "Identity", [source], num_outputs=1) + extra.outputs[0].name = "shared_source" + extra.outputs[0].shape = source.shape + extra.outputs[0].dtype = source.dtype + graph.append(extra) + graph.outputs.append(extra.outputs[0]) + elif retained == "graph_output": + graph.outputs.append(source) + model = ir.Model(graph, ir_version=10) + result = FoldWeightLayoutPass()(model) + assert result.modified == (retained is None) + if retained is None: + assert "parent_weight" not in graph.initializers + assert set(branch.initializers) == {"captured_layout"} + assert not FoldWeightLayoutPass()(model).modified + np.testing.assert_array_equal( + branch.initializers["captured_layout"].const_value.numpy(), data.T + ) + else: + assert graph.initializers["parent_weight"] is source + assert not branch.initializers + assert transpose.inputs[0] is source + # Exercise the public serializer, not a protobuf encoding. The retained + # parent case must stay shared; the unshared case must serialize only the + # folded child weight, not both it and the dead parent payload. + fold_initializers_after_weights(model) + ModelPackage({"model": model}).save(str(tmp_path), max_workers=1, progress_bar=False) + restored = ModelPackage.load(str(tmp_path))["model"] + persisted = [value for scope in restored.graphs() for value in scope.initializers.values()] + assert sum(value.const_value.numpy().nbytes for value in persisted) == data.nbytes * 2 + if retained is None: + assert "parent_weight" not in restored.graph.initializers + restored_branch = next(iter(restored.graph)).attributes["then_branch"].as_graph() + assert set(restored_branch.initializers) == {"captured_layout"} + np.testing.assert_array_equal( + restored_branch.initializers["captured_layout"].const_value.numpy(), data.T + ) + else: + np.testing.assert_array_equal( + restored.graph.initializers["parent_weight"].const_value.numpy(), data + ) + + +@pytest.mark.parametrize("width", [3, 3 * 8192]) +def test_shape_constant_folding_exposes_layout_without_reading_large_weights(width): + data = np.arange(16 * width, dtype=np.float32).reshape(2, 8, width) + model = _model(data, first=(2, 2, 4, width), last=(2, 8, width)) + source = model.graph.initializers["weight"] + reads = [] + + def read(): + reads.append(True) + return ir.tensor(data) + + source.const_value = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) + shape = model.graph.initializers["shape1"] + offset = _initializer("shape_offset", np.zeros(4, dtype=np.int64)) + model.graph.register_initializer(offset) + expression = ir.Node("", "Add", [shape, offset], num_outputs=1) + expression.outputs[0].name = "computed_shape" + expression.outputs[0].shape = shape.shape + expression.outputs[0].dtype = ir.DataType.INT64 + first = next(iter(model.graph)) + model.graph.insert_before(first, expression) + first.replace_input_with(1, expression.outputs[0]) + fold_initializers_after_weights(model) + assert not reads + assert set(model.graph.initializers) == {"output"} + assert not list(model.graph) + folded = model.graph.initializers["output"] + assert model.graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + fold_initializers_after_weights(model) + assert not reads + assert model.graph.initializers["output"] is folded + assert model.graph.outputs[0] is folded + expected = data.reshape(2, 2, 4, width).transpose(0, 2, 1, 3).reshape(2, 8, width) + np.testing.assert_array_equal(folded.const_value.numpy(), expected) + assert len(reads) == 1 + + +def test_folded_layout_keeps_pruned_external_source_ownership(tmp_path, monkeypatch): + # The 48-element fixture used elsewhere is only 192 bytes: ir.save keeps + # tensors below its 256-byte default threshold inline, so it has no file + # backing to protect. This fixture must exercise a real ExternalTensor. + width = 24 + data = np.arange(16 * width, dtype=np.float32).reshape(2, 8, width) + base = tmp_path / "base" + model = _model(data, first=(2, 2, 4, width), last=(2, 8, width)) + ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) + before = {path.name: path.read_bytes() for path in base.iterdir()} + package = ModelPackage.load(str(base)) + model = package["model"] + source = model.graph.initializers["weight"].const_value + assert isinstance(source, ir.ExternalTensor) + source_path = Path(source.path).resolve() + assert source_path == (base / "model.onnx.data").resolve() + + def unexpected_read(*args, **kwargs): + raise AssertionError("source-alias refusal must precede lazy payload reads") + + with monkeypatch.context() as patch: + patch.setattr(ir.ExternalTensor, "numpy", unexpected_read) + fold_initializers_after_weights(model) + folded_graph = model.graph + assert set(folded_graph.initializers) == {"output"} + folded = folded_graph.initializers["output"] + assert folded_graph.outputs[0] is folded + assert isinstance(folded.const_value, ir.LazyTensor) + assert source_path in { + Path(path).resolve() for path in model.meta[LAZY_SOURCE_FILES_META_KEY] + } + assert source_path in { + Path(path).resolve() + for path in folded.const_value.meta[LAZY_SOURCE_FILES_META_KEY] + } + with pytest.raises(ValueError, match="aliases lazy source file"): + package.save(str(base), max_workers=1, progress_bar=False) + assert source_path in package._native_streaming_source_files + assert before == {path.name: path.read_bytes() for path in base.iterdir()} + safe = tmp_path / "safe" + package.save(str(safe), max_workers=1, progress_bar=False) + restored_graph = ModelPackage.load(str(safe))["model"].graph + assert set(restored_graph.initializers) == {"output"} + restored = restored_graph.initializers["output"] + assert restored_graph.outputs[0] is restored + expected = data.reshape(2, 2, 4, width).transpose(0, 2, 1, 3).reshape(2, 8, width) + np.testing.assert_array_equal(restored.const_value.numpy(), expected) diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py index 034840069..4ec370e95 100644 --- a/src/mobius/adapter_io.py +++ b/src/mobius/adapter_io.py @@ -21,8 +21,10 @@ 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, QKV_PROJECTION_SOURCES_METADATA_KEY from mobius.adapters import ( AdapterArtifact, AdapterServiceOptions, @@ -166,11 +168,180 @@ 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 _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: + 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] + packed = split.inputs[0] if split.inputs else None + projection = packed.producer() if packed is not None else None + if ( + split.domain not in ("", "ai.onnx") + or split.op_type != "Split" + or split.overload + or split.attributes.get_int("axis", 0) not in (-1, 2) + or len(split.outputs) != len(sources) + or packed is None + or projection is None + or projection.domain not in ("", "ai.onnx") + 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") + 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] != sum(width for _, width in sources) + ): + raise ValueError( + f"PEFT module {module_key!r} requires a loaded, folded packed projection weight" + ) + 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/functions/matmul_nbits.py b/src/mobius/functions/matmul_nbits.py index f2c5ee43c..59dd3d5ef 100644 --- a/src/mobius/functions/matmul_nbits.py +++ b/src/mobius/functions/matmul_nbits.py @@ -48,6 +48,8 @@ from mobius._constants import OPSET_VERSION DOMAIN = "com.microsoft" +FUNCTION_BODY_METADATA_KEY = "mobius.function_body" +FUNCTION_BODY_METADATA_VALUE = "1" def matmul_nbits(*, has_zero_points: bool = True) -> ir.Function: @@ -152,4 +154,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..cd0920e2e 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,37 @@ 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): + 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 + + @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/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..60672209e --- /dev/null +++ b/src/mobius/rewrite_rules/_qkv_projection_packing.py @@ -0,0 +1,477 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Pack projections, then split before normalization, RoPE, attention, or caches. + +Supports distinct Q/K/V or Q/K with raw K reused for V. Parameter Concats defer +packing until checkpoint binding. No weights are dequantized or requantized. +""" + +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 in ("", "ai.onnx") 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 ("", "ai.onnx", "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..19474fda4 --- /dev/null +++ b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py @@ -0,0 +1,963 @@ +# 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 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]) +def test_peft_sources_survive_fold_and_serialization(tmp_path, transposed, shared): + model, weights, _, _ = _model(transposed=transposed, shared=shared) + assert pack_qkv_projections_pass()(model).modified + 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( + QKV_PROJECTION_SOURCES_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[QKV_PROJECTION_SOURCES_METADATA_KEY] = json.dumps({"version": 2}) + with pytest.raises(ValueError, match="unsupported"): + _resolve_initializer_name(model, "q") + + +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 "PackMatMulNBitsQKV" not in names + assert ("PackQKV" in names) == ( + role == "decoder" and dtype in caps.qkv_pack_dtypes + ) + assert ("PackQKVProjections" in names) == ( + role != "encoder" 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") + + +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 dataclasses.fields(EpCapabilities)[-1].name == "matmul_nbits_qkv_pack_dtypes" + + +@pytest.mark.parametrize("trace", [False, True]) +@pytest.mark.parametrize("opt_in", [False, True]) +def test_generic_pipeline_preserves_arbitrary_mask_contract(monkeypatch, trace, opt_in): + model, _, _, _ = _model(quantized=True, deferred=True, dtype=ir.DataType.FLOAT16) + declare_arbitrary_attention_mask(model.graph) + caps = dataclasses.replace( + get_ep("cuda"), + name="projection-test", + matmul_nbits_qkv_pack_dtypes=( + frozenset({ir.DataType.FLOAT16}) if opt_in else frozenset() + ), + ) + monkeypatch.setitem(ep_registry._entries, caps.name, caps) + optimize_model(model, ep=caps.name, dtype=ir.DataType.FLOAT16, 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) + assert sum(node.op_type == "MatMulNBits" 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) From e01ac8b83290090fa93f4f67344fc3f0bff732d7 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Sat, 24 Oct 2026 00:09:09 +0000 Subject: [PATCH 2/7] Consolidate nested Concat cleanup and quantized layout folding Scope nested Concat folding with child cleanup, document quantized storage-contract eligibility, and record float capability documentation. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/ep_quickstart.md | 11 +- docs/execution_providers.md | 26 +- src/mobius/_execution_providers.py | 13 +- src/mobius/_passes/_fold_concat.py | 68 ++++- src/mobius/_passes/_fold_concat_test.py | 171 ++++++++++++ src/mobius/_passes/_fold_weight_layout.py | 46 ++++ .../_passes/_fold_weight_layout_test.py | 244 ++++++++++++++++++ 7 files changed, 555 insertions(+), 24 deletions(-) diff --git a/docs/ep_quickstart.md b/docs/ep_quickstart.md index 6a805f64d..5414225fc 100644 --- a/docs/ep_quickstart.md +++ b/docs/ep_quickstart.md @@ -100,7 +100,7 @@ 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}); both float packing stages 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 @@ -135,7 +135,12 @@ graph construction or optimization starts — safe-fail by default. Experimental INT4 projection packing requires an explicit `matmul_nbits_qkv_pack_dtypes` capability override; no built-in EP enables it. -Existing float `PackQKV` behavior is unchanged. Read the +The existing float GQA-input `PackQKV` stage is unchanged, but the additional +`PackQKVProjections` float stage is **new and default-enabled** for decoder +and masked-decoder roles whenever the dtype is in `qkv_pack_dtypes`, including +CUDA's existing FLOAT/FP16/BF16 profile. Only INT4 projection packing is +default-off. Its failed/unwaived numerical gate does not qualify the new float +stage; float numerical accuracy is not separately 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. @@ -148,7 +153,7 @@ before enabling it; activation/component precision flags do not opt in. | 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 + float PackQKV; experimental INT4 projection packing is off | +| 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 6a56a564e..3cc8004b7 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -102,9 +102,14 @@ These are applied during **Stage 2: Fusion** of the optimization pipeline. #### Projection packing and numerical limitations -The existing float `PackQKV` GQA-input fusion is unchanged. A generic -projection-only pass also handles compatible float projections that remain -separate after that fusion, including masked decoders and raw K reused for V. +The existing float `PackQKV` GQA-input fusion stage is unchanged, but +**`PackQKVProjections` is a new, additional float projection-packing stage**. +It runs by default for both `decoder` and `masked-decoder` roles whenever the +activation dtype is in the EP's existing `qkv_pack_dtypes` capability. Thus +nonempty built-in float profiles, including CUDA's FLOAT/FP16/BF16 profile, +enable this new behavior without a new opt-in. It handles compatible float +projections that remain separate after GQA-input packing, including masked +decoders and raw K reused for V. Only INT4 projection packing is default-off. It splits the packed result **before** normalization, RoPE, attention, or cache consumers. Sources must be static, non-overridable initializer parameters with compatible shapes and dtypes; ambiguous repeated sources, shared weights with @@ -152,14 +157,21 @@ Normally the post-weight pipeline folds Concats and static Reshape/Transpose chains lazily, including layouts exposed by small integer shape-expression folding. That generic folding is always independent of experimental packing; it preserves dtype and does not eagerly load large weight tensors. +Checkpoint-storage layout paths feeding MoE/QMoE parameters or +`DequantizeLinear` are conservatively retained: their named expert/code/scale +initializers and original byte layouts remain available to loaders and +serialized-storage contracts. Ordinary layout folding and the existing +2-D transpose optimization remain enabled. -**The PR754 numerical gate remains failed and unwaived.** Byte-exact packing +**The PR754 INT4-packing numerical gate remains failed and unwaived.** Byte-exact packing does not imply identical floating-point outputs: changing GEMM output width can change CUDA kernel behavior. Representative large-dimension prefill and cached-decode logits exceeded the existing `rtol=atol=1e-2` criterion. Bounded MMLU/generation evidence does not resolve pointwise parity; neither tolerances -nor automatic BF16/FP16 enablement have changed. No production numerical -accuracy or speed improvement is claimed. +nor automatic BF16/FP16 INT4-packing enablement have changed. This branch does +not separately qualify numerical accuracy for the new float projection stage; +do not infer a float numerical pass from the INT4 gate's scope. No production +numerical accuracy or speed improvement is claimed. For `ORT_FPA_INTB_GEMM=1`, keep exporter folding or ORT constant folding enabled: unfolded parameter Concats can prevent runtime prepacking when both are off. @@ -200,7 +212,7 @@ 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] # both float packing stages 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/src/mobius/_execution_providers.py b/src/mobius/_execution_providers.py index e062c5452..51cb1c224 100644 --- a/src/mobius/_execution_providers.py +++ b/src/mobius/_execution_providers.py @@ -53,11 +53,14 @@ 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 both the existing decoder-only + GroupQueryAttention input packing (PackQKV) and the new additional + float projection packing (PackQKVProjections) for decoder and + masked-decoder roles. Nonempty built-in profiles, including CUDA, + enable the new float stage by default; INT4 uses its separate + capability below. 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 diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py index 7035f1e98..00ef2f118 100644 --- a/src/mobius/_passes/_fold_concat.py +++ b/src/mobius/_passes/_fold_concat.py @@ -27,6 +27,7 @@ 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__) @@ -44,20 +45,67 @@ 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()) + + def retire_sources( + values: tuple[ir.Value | None, ...], packed: ir.Value, *, prune: bool + ) -> None: + 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: + # Deregister only: lazy callbacks still hold these values + # and need their const_value tensors to materialize later. + 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) + ) 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 @@ -122,15 +170,16 @@ 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(inputs, existing_val, prune=graph is not model.graph) folded_count += 1 modified = True @@ -168,13 +217,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(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..f13173944 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_cleanup_retains_graph_input_initializer(): + model, sources = _make_concat_model([np.ones((4, 8), dtype=np.float32)] * 2) + model.graph.inputs.append(sources[0]) + assert FoldConcatInitializersPass()(model).modified + assert model.graph.initializers["init_1"] is sources[1] + common.RemoveUnusedNodesPass()(model) + assert model.graph.initializers["init_0"] is sources[0] + assert "init_1" not in model.graph.initializers + + +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_weight_layout.py b/src/mobius/_passes/_fold_weight_layout.py index ddefb9f5a..b7bd9ca5f 100644 --- a/src/mobius/_passes/_fold_weight_layout.py +++ b/src/mobius/_passes/_fold_weight_layout.py @@ -68,6 +68,47 @@ def _reshape_shape(node: ir.Node, source_shape: tuple[int, ...]) -> tuple[int, . return tuple(resolved) +def _storage_layout_values(model: ir.Model) -> set[ir.Value]: + """Find parameter views whose checkpoint storage ABI must remain explicit.""" + pending: list[ir.Value] = [] + for node in model.graph.all_nodes(): + if node.domain in ("", "ai.onnx") and node.op_type == "DequantizeLinear": + parameters = node.inputs[:3] + elif node.domain == "com.microsoft" and node.op_type in ("MoE", "QMoE"): + # Skip activations/routing. Expert weights, biases, scales, and + # zero points retain their checkpoint-format names and row order; + # the explicit layout views adapt that storage to the runtime ABI. + parameters = node.inputs[2 : 14 if node.op_type == "QMoE" else 8] + else: + continue + pending.extend(value for value in parameters if value is not None) + protected: set[ir.Value] = set() + while pending: + value = pending.pop() + if value in protected: + continue + protected.add(value) + producer = value.producer() + if producer is None or producer.domain not in ("", "ai.onnx") or producer.overload: + continue + if producer.op_type in ( + "Identity", + "Reshape", + "Transpose", + "Cast", + "CastLike", + "Pad", + "Slice", + "Expand", + ): + pending.extend(item for item in producer.inputs[:1] if item is not None) + elif producer.op_type in ("Add", "Sub", "Mul", "Div", "Concat"): + # FP8 scales can carry an exact-zero execution dependency before + # DQ. Follow the scale path without evaluating its arithmetic. + pending.extend(item for item in producer.inputs if item is not None) + return protected + + class FoldWeightLayoutPass(ir.passes.InPlacePass): """Fold named static Reshape/Transpose outputs into lazy initializers. @@ -75,6 +116,8 @@ class FoldWeightLayoutPass(ir.passes.InPlacePass): without reading weight payloads: only constant Reshape shape vectors are read. Materialization preserves the source dtype and uses no arithmetic or casts. Overridable initializers, shared sources, and dynamic shapes remain untouched. + MoE/QMoE and QDQ parameter-layout paths retain checkpoint storage names and + byte layouts; their explicit runtime format conversions are not folded. Dead folded inputs are removed in their owning graph; run dead-code elimination afterward to remove unused Constant shape nodes. Malformed static layout nodes raise ValueError rather than being skipped. @@ -83,6 +126,7 @@ class FoldWeightLayoutPass(ir.passes.InPlacePass): def call(self, model: ir.Model) -> ir.passes.PassResult: folded = 0 graphs = tuple(model.graphs()) + storage_layouts = _storage_layout_values(model) # Registry identity is authoritative, including parent-owned captures. # Do not infer ownership from the graph containing the consumer node. @@ -108,6 +152,8 @@ def owners(value: ir.Value) -> list[ir.Graph]: expected_inputs = 2 if node.op_type == "Reshape" else 1 if len(node.inputs) != expected_inputs or len(node.outputs) != 1: raise ValueError(f"Malformed weight layout node {node.name!r}") + if node.outputs[0] in storage_layouts: + continue source = node.inputs[0] if source is None or not source.is_initializer() or source.is_graph_input(): continue diff --git a/src/mobius/_passes/_fold_weight_layout_test.py b/src/mobius/_passes/_fold_weight_layout_test.py index f6ac88cd0..626c18ebf 100644 --- a/src/mobius/_passes/_fold_weight_layout_test.py +++ b/src/mobius/_passes/_fold_weight_layout_test.py @@ -17,6 +17,7 @@ from mobius._constants import LAZY_SOURCE_FILES_META_KEY from mobius._optimizations import fold_initializers_after_weights from mobius._passes import FoldWeightLayoutPass +from mobius._testing.ort_inference import OnnxModelSession def _initializer(name, data): @@ -168,6 +169,139 @@ def test_shared_sources_are_not_duplicated(): assert source.const_value.numpy().tobytes() == data.tobytes() +@pytest.mark.parametrize( + ("domain", "op_type", "parameter_slot", "protected"), + [ + ("com.microsoft", "MoE", 2, True), + ("com.microsoft", "QMoE", 2, True), + ("com.microsoft", "QMoE", 3, True), + ("", "DequantizeLinear", 0, True), + ("", "DequantizeLinear", 1, True), + ("com.microsoft", "MatMulNBits", 1, False), + ("", "MatMul", 1, False), + ], +) +def test_post_weight_storage_layout_paths_preserve_named_sources( + domain, op_type, parameter_slot, protected +): + model = _model(np.arange(48, dtype=np.float32).reshape(2, 8, 3)) + source = model.graph.initializers["weight"] + + def unexpected_read(): + raise AssertionError("eligibility and layout folding must not read weight payloads") + + tensor = ir.LazyTensor(unexpected_read, dtype=source.dtype, shape=source.shape) + source.const_value = tensor + # Only the consumer's parameter position matters to this structural test; + # complete model/schema and public save/load contracts are covered by the + # existing QMoE/FP8/streaming model tests, unchanged by this regression fix. + inputs = [] + for index in range(parameter_slot): + argument = ir.Value( + name=f"argument_{index}", + type=ir.TensorType(ir.DataType.FLOAT), + shape=ir.Shape([1]), + ) + model.graph.inputs.append(argument) + inputs.append(argument) + inputs.append(model.graph.outputs[0]) + consumer = ir.Node(domain, op_type, inputs, num_outputs=1) + consumer.outputs[0].name = "consumer_output" + model.graph.append(consumer) + model.graph.outputs.append(consumer.outputs[0]) + fold_initializers_after_weights(model) + if protected: + assert model.graph.initializers["weight"] is source + assert source.const_value is tensor + assert sum(node.op_type == "Reshape" for node in model.graph) == 2 + assert sum(node.op_type == "Transpose" for node in model.graph) == 1 + else: + assert "weight" not in model.graph.initializers + assert not any(node.op_type in ("Reshape", "Transpose") for node in model.graph) + + +@pytest.mark.parametrize("with_dependency", [False, True]) +def test_qdq_scalar_scale_storage_survives_pipeline_and_package_roundtrip( + tmp_path, with_dependency +): + data = np.array([[0.5]], dtype=np.float32) + scale = _initializer("checkpoint_scale", data) + reads = [] + + def read_scale(): + reads.append(True) + return ir.tensor(data) + + tensor = ir.LazyTensor(read_scale, dtype=scale.dtype, shape=scale.shape) + scale.const_value = tensor + scalar_shape = _initializer("scalar_shape", np.array([], dtype=np.int64)) + reshape = ir.Node("", "Reshape", [scale, scalar_shape], num_outputs=1) + reshape.outputs[0].name = "runtime_scalar_scale" + reshape.outputs[0].dtype = ir.DataType.FLOAT + reshape.outputs[0].shape = ir.Shape([]) + codes = ir.Value(name="codes", type=ir.TensorType(ir.DataType.UINT8), shape=ir.Shape([4])) + inputs = [codes] + nodes = [reshape] + runtime_scale = reshape.outputs[0] + if with_dependency: + dependency = ir.Value( + name="dependency", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([]) + ) + inputs.append(dependency) + add = ir.Node("", "Add", [runtime_scale, dependency], num_outputs=1) + add.outputs[0].name = "dependent_scale" + add.outputs[0].dtype = ir.DataType.FLOAT + add.outputs[0].shape = ir.Shape([]) + nodes.append(add) + runtime_scale = add.outputs[0] + dequantize = ir.Node("", "DequantizeLinear", [codes, runtime_scale], num_outputs=1) + dequantize.outputs[0].name = "output" + dequantize.outputs[0].dtype = ir.DataType.FLOAT + dequantize.outputs[0].shape = ir.Shape([4]) + nodes.append(dequantize) + model = ir.Model( + ir.Graph( + inputs, + [dequantize.outputs[0]], + nodes=nodes, + initializers=[scale, scalar_shape], + opset_imports={"": 21}, + ), + ir_version=10, + ) + fold_initializers_after_weights(model) + assert model.graph.initializers["checkpoint_scale"] is scale + assert scale.const_value is tensor + assert not reads + assert reshape.graph is model.graph + assert reshape.inputs[0] is scale + # The protected source stays in checkpoint orientation, not the scalar + # runtime format. This is a storage contract, not just numerical parity. + assert scale.shape == ir.Shape([1, 1]) + assert scale.dtype == ir.DataType.FLOAT + fold_initializers_after_weights(model) + assert model.graph.initializers["checkpoint_scale"] is scale + assert not reads + + ModelPackage({"model": model}).save(tmp_path, progress_bar=False) + restored = ModelPackage.load(tmp_path)["model"] + restored_scale = restored.graph.initializers["checkpoint_scale"] + assert restored_scale.shape == scale.shape + assert restored_scale.dtype == scale.dtype + assert restored_scale.const_value.numpy().tobytes() == data.tobytes() + for candidate in (model, restored): + session = OnnxModelSession(candidate, device="cpu") + for dependency_value in (0.0, 0.25) if with_dependency else (0.0,): + code_data = np.array([0, 1, 127, 255], dtype=np.uint8) + feeds: dict[str, np.ndarray] = {"codes": code_data} + if with_dependency: + feeds["dependency"] = np.asarray(dependency_value, dtype=np.float32) + actual = session.run(feeds)["output"] + np.testing.assert_array_equal( + actual, code_data.astype(np.float32) * (0.5 + dependency_value) + ) + + def test_unshared_graph_output_layout_name_survives(): model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) layout = next(node for node in model.graph if node.outputs[0].name == "layout") @@ -371,6 +505,116 @@ def test_nested_graph_initializers_are_registered_in_their_scope(): assert set(branch.initializers) == {"layout"} +def test_post_weight_pipeline_concat_layouts_are_isolated_in_if_branches(tmp_path): + branches = [] + expected_outputs = [] + for offset in (0, 100): + data = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + offset + zeros = np.zeros((1, 4, 3), dtype=np.float32) + # Each lexical scope deliberately uses the same names for different + # payloads. The rank-3 Transpose is folded by FoldWeightLayoutPass, + # which exposes an all-initializer Concat in each branch. + weight = _initializer("W", data) + padding = _initializer("Z", zeros) + transpose = ir.Node( + "", "Transpose", [weight], [ir.AttrInt64s("perm", [0, 2, 1])], num_outputs=1 + ) + transpose.outputs[0].name = "layout" + transpose.outputs[0].shape = ir.Shape([2, 4, 3]) + transpose.outputs[0].dtype = ir.DataType.FLOAT + concat = ir.Node( + "", + "Concat", + [transpose.outputs[0], padding], + [ir.AttrInt64("axis", 0)], + num_outputs=1, + ) + concat.outputs[0].name = "branch_output" + concat.outputs[0].shape = ir.Shape([3, 4, 3]) + concat.outputs[0].dtype = ir.DataType.FLOAT + branches.append( + ir.Graph( + [], + [concat.outputs[0]], + nodes=[transpose, concat], + initializers=[weight, padding], + name=f"branch_{offset}", + ) + ) + expected_outputs.append(np.concatenate([data.transpose(0, 2, 1), zeros], axis=0)) + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + select = ir.Node( + "", + "If", + [condition], + [ + ir.AttrGraph("then_branch", branches[0]), + ir.AttrGraph("else_branch", branches[1]), + ], + num_outputs=1, + ) + select.outputs[0].name = "result" + select.outputs[0].shape = ir.Shape([3, 4, 3]) + select.outputs[0].dtype = ir.DataType.FLOAT + model = ir.Model( + ir.Graph( + [condition], + [select.outputs[0]], + nodes=[select], + opset_imports={"": 24}, + name="layout_concat_if", + ), + ir_version=10, + ) + + fold_initializers_after_weights(model) + assert not model.graph.initializers + packed_name = "layout__Z__axis_0__concat" + packed_values = [] + for branch, expected in zip(branches, expected_outputs, strict=True): + assert not list(branch) + assert set(branch.initializers) == {packed_name} + packed = branch.initializers[packed_name] + assert branch.outputs[0] is packed + assert isinstance(packed.const_value, ir.LazyTensor) + np.testing.assert_array_equal(packed.const_value.numpy(), expected) + packed_values.append(packed) + assert packed_values[0] is not packed_values[1] + fold_initializers_after_weights(model) + assert not model.graph.initializers + assert all( + branch.initializers[packed_name] is packed + for branch, packed in zip(branches, packed_values, strict=True) + ) + + output_directory = tmp_path / "package" + ModelPackage({"model": model}).save( + str(output_directory), max_workers=1, progress_bar=False + ) + restored = ModelPackage.load(str(output_directory))["model"] + assert not restored.graph.initializers + restored_select = next(iter(restored.graph)) + for attribute, expected in zip( + ("then_branch", "else_branch"), expected_outputs, strict=True + ): + branch = restored_select.attributes[attribute].as_graph() + assert not list(branch) + assert set(branch.initializers) == {packed_name} + packed = branch.initializers[packed_name] + assert branch.outputs[0] is packed + np.testing.assert_array_equal(packed.const_value.numpy(), expected) + + # Both branches must execute their own payload, before and after the public + # package roundtrip, rather than reusing a same-named sibling initializer. + for candidate in (model, restored): + session = OnnxModelSession(candidate, device="cpu") + for take_then, expected in zip((True, False), expected_outputs, strict=True): + result = session.run({"condition": np.asarray(take_then, dtype=np.bool_)}) + np.testing.assert_array_equal(result["result"], expected) + + @pytest.mark.parametrize("retained", [None, "shared", "graph_output"]) def test_parent_capture_initializer_ownership_and_persisted_payload(tmp_path, retained): data = np.arange(6, dtype=np.float32).reshape(2, 3) From c914cc42094baca183dce244e42ed6c11e3c620d Mon Sep 17 00:00:00 2001 From: titaiwang Date: Sat, 24 Oct 2026 02:26:11 +0000 Subject: [PATCH 3/7] Preserve runtime-overridable inputs during Concat folding Captured parent initializers cannot be packed when they are runtime-overridable. Added root and captured runtime/default-override coverage, including public save/load behavior. State: 78 pass, 2 before-fix expected fail, scoped lint passed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/execution_providers.md | 2 + src/mobius/_passes/_fold_concat.py | 7 +- src/mobius/_passes/_fold_concat_test.py | 8 +- .../_passes/_fold_weight_layout_test.py | 89 +++++++++++++++++++ 4 files changed, 101 insertions(+), 5 deletions(-) diff --git a/docs/execution_providers.md b/docs/execution_providers.md index 3cc8004b7..358e9ba13 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -157,6 +157,8 @@ Normally the post-weight pipeline folds Concats and static Reshape/Transpose chains lazily, including layouts exposed by small integer shape-expression folding. That generic folding is always independent of experimental packing; it preserves dtype and does not eagerly load large weight tensors. +Concats with runtime-overridable initializer inputs remain unfused, including +initializers captured from a parent graph's inputs. Checkpoint-storage layout paths feeding MoE/QMoE parameters or `DequantizeLinear` are conservatively retained: their named expert/code/scale initializers and original byte layouts remain available to loaders and diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py index 00ef2f118..dbbbf5673 100644 --- a/src/mobius/_passes/_fold_concat.py +++ b/src/mobius/_passes/_fold_concat.py @@ -36,7 +36,7 @@ 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 @@ -113,6 +113,11 @@ def retire_sources( 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 diff --git a/src/mobius/_passes/_fold_concat_test.py b/src/mobius/_passes/_fold_concat_test.py index f13173944..9d2361a5b 100644 --- a/src/mobius/_passes/_fold_concat_test.py +++ b/src/mobius/_passes/_fold_concat_test.py @@ -168,14 +168,14 @@ def read(data=array): assert len(reads) == 2 -def test_concat_cleanup_retains_graph_input_initializer(): +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 FoldConcatInitializersPass()(model).modified - assert model.graph.initializers["init_1"] is sources[1] + 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 "init_1" not in model.graph.initializers + assert model.graph.initializers["init_1"] is sources[1] def test_concat_cleanup_preserves_external_source_ownership(tmp_path, monkeypatch): diff --git a/src/mobius/_passes/_fold_weight_layout_test.py b/src/mobius/_passes/_fold_weight_layout_test.py index 626c18ebf..8bae303e0 100644 --- a/src/mobius/_passes/_fold_weight_layout_test.py +++ b/src/mobius/_passes/_fold_weight_layout_test.py @@ -505,6 +505,95 @@ def test_nested_graph_initializers_are_registered_in_their_scope(): assert set(branch.initializers) == {"layout"} +@pytest.mark.parametrize("parent_capture", [False, True]) +def test_post_weight_pipeline_preserves_overridable_concat_input(tmp_path, parent_capture): + import onnxruntime as ort + + data = np.arange(8, dtype=np.float32) + weight = _initializer("W", data) + shape = _initializer("shape", np.array([2, 4], dtype=np.int64)) + bias = _initializer("B", np.zeros((1, 4), dtype=np.float32)) + reshape = ir.Node("", "Reshape", [weight, shape], num_outputs=1) + reshape.outputs[0].name = "layout" + reshape.outputs[0].shape = ir.Shape([2, 4]) + reshape.outputs[0].dtype = ir.DataType.FLOAT + concat = ir.Node( + "", + "Concat", + [reshape.outputs[0], bias], + [ir.AttrInt64("axis", 0)], + num_outputs=1, + ) + concat.outputs[0].name = "output" + concat.outputs[0].shape = ir.Shape([3, 4]) + concat.outputs[0].dtype = ir.DataType.FLOAT + graph = ir.Graph( + [bias], + concat.outputs, + nodes=[reshape, concat], + initializers=[weight, shape, bias], + name="overridable_concat", + opset_imports={"": 24}, + ) + if parent_capture: + graph.inputs.clear() + del graph.initializers["B"] + condition = ir.Value( + name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) + ) + fallback = ir.Node("", "Identity", [bias], num_outputs=1) + fallback.outputs[0].name = "fallback" + fallback.outputs[0].shape = bias.shape + fallback.outputs[0].dtype = bias.dtype + other = ir.Graph([], fallback.outputs, nodes=[fallback], name="other_branch") + select = ir.Node( + "", + "If", + [condition], + [ir.AttrGraph("then_branch", graph), ir.AttrGraph("else_branch", other)], + num_outputs=1, + ) + select.outputs[0].name = "result" + select.outputs[0].shape = ir.Shape([None, 4]) + select.outputs[0].dtype = ir.DataType.FLOAT + graph = ir.Graph( + [condition, bias], + select.outputs, + nodes=[select], + initializers=[bias], + name="captured_overridable_concat", + opset_imports={"": 24}, + ) + model = ir.Model(graph, ir_version=10) + fold_initializers_after_weights(model) + assert all(node.op_type != "Reshape" for node in model.graph.all_nodes()) + assert sum(node.op_type == "Concat" for node in model.graph.all_nodes()) == 1 + assert model.graph.initializers["B"] is bias + assert bias in model.graph.inputs + + ModelPackage({"model": model}).save(tmp_path, progress_bar=False, max_workers=1) + restored = ModelPackage.load(tmp_path)["model"] + for index, candidate in enumerate((model, restored)): + model_path = tmp_path / f"runtime_{index}.onnx" + ir.save(candidate, model_path) + # The test wrapper filters out overridable initializers via get_inputs(); + # use ORT directly so B is actually supplied to the runtime. + session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) + assert [value.name for value in session.get_overridable_initializers()] == ["B"] + for override in (None, np.full((1, 4), 17, dtype=np.float32)): + feeds = {"condition": np.asarray(True)} if parent_capture else {} + if override is not None: + feeds["B"] = override + actual = session.run(None, feeds)[0] + expected = np.concatenate( + [ + data.reshape(2, 4), + np.zeros((1, 4), dtype=np.float32) if override is None else override, + ] + ) + np.testing.assert_array_equal(actual, expected) + + def test_post_weight_pipeline_concat_layouts_are_isolated_in_if_branches(tmp_path): branches = [] expected_outputs = [] From aee5dcca9cb64cddaa84a3b554f7478e6ac79840 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Mon, 26 Oct 2026 22:50:44 +0000 Subject: [PATCH 4/7] Clarify projection packing helpers and metadata conventions Address Justin's review by moving source retirement and initializer ownership lookup into documented module-level helpers, and naming the topology and loaded-weight validation stages used during packed PEFT restoration. Explain parent-captured initializer ownership, lazy backing-file lifetime, registry-generated function provenance, and why projection packing defers checkpoint binding without changing downstream attention contracts. Use pkg.mobius. for the persistent projection and function-body metadata keys, sharing the function marker between producers and consumers. Keep the old projection metadata readable for previously exported preview packages. Legacy-only function markers conservatively leave the body caller-owned; document and test that specialization is not authorized by those markers. Keep transient lazy-file bookkeeping in .meta rather than serialized metadata_props. Remove redundant ai.onnx checks because IR normalizes it. Shorten the execution-provider documentation while retaining automatic float-stage behavior, experimental INT4 default-off status, the failed numerical gate, and non-transactional save limitations. Validation: 463 focused CPU tests passed, 2 skipped. Scoped Ruff, lintrunner, and whitespace checks passed. MyPy has the same 26 diagnostics as the exact parent revision, with no new diagnostics. No CUDA or full-model numerical requalification is claimed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/execution_providers.md | 112 +++++++----------- src/mobius/_constants.py | 13 +- src/mobius/_passes/_fold_concat.py | 95 ++++++++------- src/mobius/_passes/_fold_transpose.py | 14 +-- src/mobius/_passes/_fold_weight_layout.py | 35 +++--- src/mobius/adapter_io.py | 96 ++++++++++----- src/mobius/functions/__init__.py | 11 +- src/mobius/functions/matmul_nbits.py | 10 +- src/mobius/functions/matmul_nbits_test.py | 12 ++ .../rewrite_rules/_qkv_projection_packing.py | 12 +- .../_qkv_projection_packing_test.py | 89 +++++++++++++- 11 files changed, 317 insertions(+), 182 deletions(-) diff --git a/docs/execution_providers.md b/docs/execution_providers.md index 358e9ba13..8b0701748 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -102,23 +102,18 @@ These are applied during **Stage 2: Fusion** of the optimization pipeline. #### Projection packing and numerical limitations -The existing float `PackQKV` GQA-input fusion stage is unchanged, but -**`PackQKVProjections` is a new, additional float projection-packing stage**. -It runs by default for both `decoder` and `masked-decoder` roles whenever the -activation dtype is in the EP's existing `qkv_pack_dtypes` capability. Thus -nonempty built-in float profiles, including CUDA's FLOAT/FP16/BF16 profile, -enable this new behavior without a new opt-in. It handles compatible float -projections that remain separate after GQA-input packing, including masked -decoders and raw K reused for V. Only INT4 projection packing is default-off. -It splits the packed result **before** normalization, RoPE, attention, or cache -consumers. Sources must be static, non-overridable initializer parameters with -compatible shapes and dtypes; ambiguous repeated sources, shared weights with -outside consumers, and overlapping groups are declined before graph mutation. -No attention-mask contract is relaxed to make packing possible. - -**INT4 projection packing is experimental and default-disabled on every -built-in EP, including CUDA, for FP16 and BF16 alike.** Opt in only through -the existing capability API: +`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` | Automatic for `qkv_pack_dtypes`, including CUDA FLOAT/FP16/BF16 | +| INT4 projection packing | `decoder`, `masked-decoder` | **Experimental; disabled on every built-in EP**, including CUDA FP16/BF16 | + +Experimental INT4 opt-in uses the process-local capability API; stock `--ep cuda` +does not enable it: ```python import dataclasses @@ -135,58 +130,37 @@ register_ep( ) ``` -This overrides the process-local profile for subsequent builds. Stock -`--ep cuda` does not opt in. The capability is independent of float -`qkv_pack_dtypes`; activation dtype (`--dtype`), weight quantization/storage, -component precision overrides, and FP8 KV-cache flags remain independent -choices. Packing does not promote precision, dequantize, or requantize weights. - -Compatible INT4 groups share a rank-3 floating-point activation, attributes -other than output width, and canonical UINT8 blocked weights. Scales match the -activation dtype; zero points must be absent throughout or uniformly packed -UINT8. Flat and row-major scale/zero-point storage are supported. Block size -must be a power of two of at least 16; `accuracy_level` must be omitted or zero. -Caller-owned function bodies, mixed bit widths, extra nonempty inputs, and -incompatible layouts are not packed. Standard opset >= 13 is required. -The canonical implementation is `_qkv_projection_packing.py` for both float -and INT4, including the supported Q/K pair with K reused for V. - -Parameter Concats retain checkpoint names until binding. Missing parameter -leaves are validated even with `apply_weights(..., fold_constants=False)`. -Normally the post-weight pipeline folds Concats and static Reshape/Transpose -chains lazily, including layouts exposed by small integer shape-expression -folding. That generic folding is always independent of experimental packing; -it preserves dtype and does not eagerly load large weight tensors. -Concats with runtime-overridable initializer inputs remain unfused, including -initializers captured from a parent graph's inputs. -Checkpoint-storage layout paths feeding MoE/QMoE parameters or -`DequantizeLinear` are conservatively retained: their named expert/code/scale -initializers and original byte layouts remain available to loaders and -serialized-storage contracts. Ordinary layout folding and the existing -2-D transpose optimization remain enabled. - -**The PR754 INT4-packing numerical gate remains failed and unwaived.** Byte-exact packing -does not imply identical floating-point outputs: changing GEMM output width can -change CUDA kernel behavior. Representative large-dimension prefill and -cached-decode logits exceeded the existing `rtol=atol=1e-2` criterion. Bounded -MMLU/generation evidence does not resolve pointwise parity; neither tolerances -nor automatic BF16/FP16 INT4-packing enablement have changed. This branch does -not separately qualify numerical accuracy for the new float projection stage; -do not infer a float numerical pass from the INT4 gate's scope. No production -numerical accuracy or speed improvement is claimed. - -For `ORT_FPA_INTB_GEMM=1`, keep exporter folding or ORT constant folding enabled: -unfolded parameter Concats can prevent runtime prepacking when both are off. -This is not a guarantee for every runtime version/configuration. - -Packed float sources retain adapter restoration metadata. Lazy restored/folded -weights keep their backing-file ownership after source initializers are pruned. -`ModelPackage.save()` checks model/shard, policy, adapter, package metadata, and -MTP destinations before writing, rejecting known same-path, hardlink, or symlink -aliases to lazy sources. Adapter files use `adapters//adapter.json`, -`adapter.onnx_adapter`, or `adapter_model.safetensors` plus `adapter_config.json` -according to persistence options. These checks are **not** full TOCTOU -protection for non-transactional saves; use a separate output directory. +The INT4 capability is independent of float packing, activation/storage dtype, +component precision, and FP8 KV caches. 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. + +Post-weight folding is independent of experimental packing: parameter Concats +retain checkpoint names until binding, then Concat/Reshape/Transpose chains fold +lazily without changing dtype. Missing leaves are checked even with +`fold_constants=False`. Overridable inputs and MoE/QMoE/QDQ storage-layout paths +remain explicit to preserve loader names and byte layouts. + +**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) diff --git a/src/mobius/_constants.py b/src/mobius/_constants.py index 45ff4e9ac..7781a3f27 100644 --- a/src/mobius/_constants.py +++ b/src/mobius/_constants.py @@ -9,11 +9,20 @@ # 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 = "mobius.qkv_projection_sources" +QKV_PROJECTION_SOURCES_METADATA_KEY = "pkg.mobius.qkv_projection_sources" """Versioned source names and widths retained on projection-packing Split nodes.""" -# Transient tensor.meta ownership, never serialized as ONNX metadata. +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``. +""" # --------------------------------------------------------------------------- # Static KV cache ABI diff --git a/src/mobius/_passes/_fold_concat.py b/src/mobius/_passes/_fold_concat.py index dbbbf5673..fd2dfb323 100644 --- a/src/mobius/_passes/_fold_concat.py +++ b/src/mobius/_passes/_fold_concat.py @@ -33,6 +33,55 @@ 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. @@ -56,46 +105,6 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: folded_count = 0 graphs = tuple(model.graphs()) - def retire_sources( - values: tuple[ir.Value | None, ...], packed: ir.Value, *, prune: bool - ) -> None: - 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: - # Deregister only: lazy callbacks still hold these values - # and need their const_value tensors to materialize later. - 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) - ) - for node in list(model.graph.all_nodes()): if node.op_type != "Concat": continue @@ -184,7 +193,9 @@ def retire_sources( # pre-pack weights survive RemoveUnusedNodesPass (~1.8 GB of # orphaned weights serialized into the fp16 GQA model). graph.remove(node, safe=True) - retire_sources(inputs, existing_val, prune=graph is not model.graph) + _retire_sources( + model, graphs, inputs, existing_val, prune=graph is not model.graph + ) folded_count += 1 modified = True @@ -229,7 +240,7 @@ def _make_packed( # folded q/k/v initializers' use lists are cleared; otherwise the # dead pre-pack weights survive RemoveUnusedNodesPass. graph.remove(node, safe=True) - retire_sources(inputs, new_val, prune=graph is not model.graph) + _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_transpose.py b/src/mobius/_passes/_fold_transpose.py index c0ed635d6..52fb232dd 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -62,11 +62,7 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: modified = False for node in list(model.graph.all_nodes()): - if ( - node.domain not in ("", "ai.onnx") - or node.overload - or node.op_type != "Transpose" - ): + if node.domain != "" or node.overload or node.op_type != "Transpose": continue if len(node.inputs) != 1: @@ -76,14 +72,14 @@ 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). + # 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. Child-local weights belong to FoldWeightLayoutPass: + # 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() - # Preserve legacy root ownership/naming even when the - # consumer is in a Loop/If body. Only child-local initializers - # need the scoped FoldWeightLayoutPass. or model.graph.initializers.get(inp.name) is not inp ): continue diff --git a/src/mobius/_passes/_fold_weight_layout.py b/src/mobius/_passes/_fold_weight_layout.py index b7bd9ca5f..a208bd1a3 100644 --- a/src/mobius/_passes/_fold_weight_layout.py +++ b/src/mobius/_passes/_fold_weight_layout.py @@ -27,7 +27,7 @@ def _reshape_shape(node: ir.Node, source_shape: tuple[int, ...]) -> tuple[int, . constant = shape_value.producer() if ( constant is None - or constant.domain not in ("", "ai.onnx") + or constant.domain != "" or constant.op_type != "Constant" or len(constant.attributes) != 1 ): @@ -72,7 +72,7 @@ def _storage_layout_values(model: ir.Model) -> set[ir.Value]: """Find parameter views whose checkpoint storage ABI must remain explicit.""" pending: list[ir.Value] = [] for node in model.graph.all_nodes(): - if node.domain in ("", "ai.onnx") and node.op_type == "DequantizeLinear": + if node.domain == "" and node.op_type == "DequantizeLinear": parameters = node.inputs[:3] elif node.domain == "com.microsoft" and node.op_type in ("MoE", "QMoE"): # Skip activations/routing. Expert weights, biases, scales, and @@ -89,7 +89,7 @@ def _storage_layout_values(model: ir.Model) -> set[ir.Value]: continue protected.add(value) producer = value.producer() - if producer is None or producer.domain not in ("", "ai.onnx") or producer.overload: + if producer is None or producer.domain != "" or producer.overload: continue if producer.op_type in ( "Identity", @@ -109,6 +109,20 @@ def _storage_layout_values(model: ir.Model) -> set[ir.Value]: return protected +def _initializer_owners(value: ir.Value, graphs: tuple[ir.Graph, ...]) -> list[ir.Graph]: + """Find registries containing this exact value, not just its name. + + An If branch can capture a root initializer without owning it, and sibling + branches can own different values with the same name. Registry identity + determines where it is safe to remove a folded source. + """ + return [ + graph + for graph in graphs + if value.name is not None and graph.initializers.get(value.name) is value + ] + + class FoldWeightLayoutPass(ir.passes.InPlacePass): """Fold named static Reshape/Transpose outputs into lazy initializers. @@ -128,20 +142,11 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: graphs = tuple(model.graphs()) storage_layouts = _storage_layout_values(model) - # Registry identity is authoritative, including parent-owned captures. - # Do not infer ownership from the graph containing the consumer node. - def owners(value: ir.Value) -> list[ir.Graph]: - return [ - scope - for scope in graphs - if value.name is not None and scope.initializers.get(value.name) is value - ] - for graph in graphs: for node in list(graph): if ( node.overload - or node.domain not in ("", "ai.onnx") + or node.domain != "" or node.op_type not in ( "Reshape", @@ -160,7 +165,7 @@ def owners(value: ir.Value) -> list[ir.Graph]: # Keep tied tables shared instead of serializing a second full layout. if ( len(source.uses()) != 1 - or len(owners(source)) != 1 + or len(_initializer_owners(source, graphs)) != 1 or any( source in scope.inputs or source in scope.outputs for scope in graphs ) @@ -255,7 +260,7 @@ def materialize( value in scope.inputs or value in scope.outputs for scope in graphs ) ): - registries = owners(value) + registries = _initializer_owners(value, graphs) if len(registries) == 1: del registries[0].initializers[value.name] folded += 1 diff --git a/src/mobius/adapter_io.py b/src/mobius/adapter_io.py index 4ec370e95..14bd199ae 100644 --- a/src/mobius/adapter_io.py +++ b/src/mobius/adapter_io.py @@ -24,7 +24,11 @@ from onnx_ir.passes import common as common_passes from safetensors.numpy import load_file -from mobius._constants import LAZY_SOURCE_FILES_META_KEY, QKV_PROJECTION_SOURCES_METADATA_KEY +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, @@ -193,43 +197,25 @@ def _packed_projection_sources(node: ir.Node, metadata: str) -> list[tuple[str, return sources -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: - 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] +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 not in ("", "ai.onnx") + 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) != len(sources) + or len(split.outputs) != source_count or packed is None or projection is None - or projection.domain not in ("", "ai.onnx") + or projection.domain != "" or projection.op_type != "MatMul" or projection.overload or projection.attributes @@ -239,6 +225,17 @@ def _restore_packed_peft_projections( 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 @@ -250,11 +247,48 @@ def _restore_packed_peft_projections( 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] != sum(width for _, width in sources) + 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 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 59dd3d5ef..2c610d4b6 100644 --- a/src/mobius/functions/matmul_nbits.py +++ b/src/mobius/functions/matmul_nbits.py @@ -48,7 +48,15 @@ from mobius._constants import OPSET_VERSION DOMAIN = "com.microsoft" -FUNCTION_BODY_METADATA_KEY = "mobius.function_body" +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" diff --git a/src/mobius/functions/matmul_nbits_test.py b/src/mobius/functions/matmul_nbits_test.py index cd0920e2e..e7f19b9c1 100644 --- a/src/mobius/functions/matmul_nbits_test.py +++ b/src/mobius/functions/matmul_nbits_test.py @@ -150,6 +150,7 @@ def test_new_call_is_specialized_after_registration(self): @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()) @@ -159,6 +160,17 @@ def test_direct_factory_does_not_claim_registry_ownership(self, has_zero_points) ) 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 diff --git a/src/mobius/rewrite_rules/_qkv_projection_packing.py b/src/mobius/rewrite_rules/_qkv_projection_packing.py index 60672209e..2c029984e 100644 --- a/src/mobius/rewrite_rules/_qkv_projection_packing.py +++ b/src/mobius/rewrite_rules/_qkv_projection_packing.py @@ -1,10 +1,12 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""Pack projections, then split before normalization, RoPE, attention, or caches. +"""Combine Q/K/V projections into one wider GEMM to reduce kernel launches. -Supports distinct Q/K/V or Q/K with raw K reused for V. Parameter Concats defer -packing until checkpoint binding. No weights are dequantized or requantized. +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 @@ -50,7 +52,7 @@ def _standard(node: ir.Node, op_type: str) -> bool: - return node.domain in ("", "ai.onnx") and node.op_type == op_type and not node.overload + return node.domain == "" and node.op_type == op_type and not node.overload def _parameter( @@ -92,7 +94,7 @@ def _projection(graph: ir.Graph, value: ir.Value | None) -> ir.Node | None: projections.add(node) elif _standard(node, "Constant"): continue - elif node.domain in ("", "ai.onnx", "com.microsoft"): + 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: diff --git a/src/mobius/rewrite_rules/_qkv_projection_packing_test.py b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py index 19474fda4..f65a48983 100644 --- a/src/mobius/rewrite_rules/_qkv_projection_packing_test.py +++ b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py @@ -18,7 +18,10 @@ from mobius import ModelPackage from mobius._component_quantization import validate_quantized_component_bindings -from mobius._constants import QKV_PROJECTION_SOURCES_METADATA_KEY +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, @@ -510,9 +513,19 @@ def unexpected_read(): @pytest.mark.parametrize("transposed", [False, True, "mixed"]) @pytest.mark.parametrize("shared", [False, True]) -def test_peft_sources_survive_fold_and_serialization(tmp_path, transposed, shared): +@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") @@ -525,20 +538,86 @@ def test_peft_sources_survive_fold_and_serialization(tmp_path, transposed, share np.testing.assert_array_equal( restored.graph.initializers[f"{name}.weight"].const_value.numpy(), original ) - assert not any( - QKV_PROJECTION_SOURCES_METADATA_KEY in node.metadata_props for node in restored.graph - ) + 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( From 6ed50516a7c15af00a1018418cd982a05c5310cd Mon Sep 17 00:00:00 2001 From: Ti-Tai Wang Date: Tue, 27 Oct 2026 09:48:13 -0700 Subject: [PATCH 5/7] Align shared Gemma4 text and multimodal runtime semantics (#787) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Organization **Depends on #754.** This PR contains the shared Gemma4 text-only and multimodal modeling changes. Generic projection packing and weight-layout folding live in #754 instead of being duplicated here. Text-only construction already exists on main; this PR corrects and documents the model semantics used by both text and multimodal exports. ## Changes - Preserve HuggingFace arithmetic boundaries in BF16 norms, routing, vision pooling, and expert contributions. - Keep global attention causal. Compose local vision visibility as `(causal OR same-image) AND left-window AND padding`, using physical KV-slot distances even when a padding mask contains holes. - Declare the richer arbitrary-mask contract for vision-overlay exports, so optimization does not collapse it to prefix lengths through GQA fusion. Text-only CUDA exports retain GQA eligibility. - Correct RoPE/shared-KV and quantized expert handling, with explicit experimental custom-GEGLU opt-in and runtime-requirement metadata. - Keep processor-facing vision inputs float32, with internal dtype conversion. ## Review follow-up Includes cached internal-hole mask regressions using the actual optimized CUDA export's bias, a text-only GQA eligibility regression, independent HF mask comparisons, and exception-restoration coverage for the custom-GEGLU flag. The adapter alias and repeated-source packing findings from this PR are handled in its #754 dependency. ## Qualification limits Export support is not released-runtime compatibility or numerical qualification. Custom GEGLU and quantized Gemma4 MoE require the documented experimental runtime; neither is enabled by ordinary CUDA selection. Previously failed strict HF numerical gates remain failed and unwaived. Historical performance and accuracy artifacts are not represented as validation of this reorganized head. The ongoing paired MMLU-Pro500 evaluation uses the frozen optimized text artifact and the pinned decoder-only llama.cpp Q4_K_M artifact, with identical prompts and scoring. ## Consolidated-revision validation - Gemma4 model suite: **118 passed**, including four adversarial cases crossing selected/unselected overflow with zero/unit expert scale. - Independent offline HF attention-mask comparisons: **11 passed**. - Its #754 dependency: **516 passed, 3 skipped**, including all 22 unchanged quantized-storage regressions and nested-scope persistence coverage. - Combined non-integration suite: **11,311 passed, 29 failed, 138 skipped**. Exact-base selector runs reproduce 28 failures; the additional GGUF file-identity failure passes separately on both base and head. This is **not a green full-suite claim**. - Scoped Ruff checks passed. Scoped MyPy remains baseline-unclean; existing HF `generate` self-type diagnostics remain on both base and head. Two bounded five-role review rounds completed. The review's portable-MoE question was verified with the actual dispatch graph: an overflowing **unselected** expert produced 64 NaNs before the correction, versus a selected-only zero reference. An independent Boolean selection mask now selects out unchosen outputs before weighting. The same CPU probe produces 64 finite exact zeros afterward; selected nonfinite outputs still propagate, including with zero learned expert scale. | ID | Severity | Raised by | Finding | Disposition | |---|---|---|---|---| | F4 | Major | Deep / critical / CPU QA | Dense fallback allowed unselected nonfinite experts to contaminate the result | Fixed and CPU-verified; CUDA equivalence not inferred | The dependency's F1/F2/F3 scope, documentation, and storage findings are fixed in #754. No Critical/Major reviewer finding was overruled. Remaining exclusions: full-checkpoint re-export and strict HF logits/KV parity, CUDA qualification of the fallback safeguard, custom-runtime availability, and performance. The failed numerical gates remain unwaived.
Complete original captured PR body — historical prior implementation (not validation of the reorganized head) The historical results in the captured body below describe the prior implementation and are **NOT validation of the reorganized head**. Older statements about unchanged float behavior are historical and superseded by the new summary above. All original numerical limitations and earlier MMLU/performance evidence are preserved verbatim. ## Summary This draft isolates the Gemma4 runtime work from the unrelated ASR, Qwen4, text-evaluation, and experimental fixed-cache changes. - Align global/local image visibility with the pinned HuggingFace reference: global layers stay causal; local same-image visibility remains subject to the sliding-window left boundary. - Normalize processor pixels in FP32 before casting for the vision projection. The separate vision encoder now accepts FLOAT pixels. - Accumulate vision pooling in FP32, round the pooled mean to the working dtype once, then scale/standardize in FP32. Local window distances use physical slots even with internal padding. - Preserve the learned router scale and apply the hidden-size scalar separately with FP32 opmath and working-dtype rounding. - Adapt the necessary CUDA native MoE/QMoE and projection-packing support to current main without reverting its weight-layout normalization. - Decline float packing when projection weights have outside consumers, and protect lazily restored PEFT weights against destructive same-path/alias saves. ## Compatibility and qualification The vision input dtype change requires rebuilding exports and passing the original FP32 processor pixels, rather than precasting them to BF16. Custom native GEGLU is disabled by default. `MOBIUS_GEMMA4_CUSTOM_GEGLU=1` explicitly opts into the experimental custom CUDA ABI; it is not implied by selecting CUDA. Ordinary float CUDA uses the portable path, while quantized GEGLU experts without the required opt-in fail clearly. Native operators record `mobius.gemma4.required_runtime_feature="custom-cuda-geglu-v1"`; the flag is an assertion about the target runtime, not automatic ABI detection. The native paths require the corresponding ORT runtime/operator support. An exported GQA node does not establish Flash/XQA dispatch, shared KV support, or CUDA Graph capture. **This is not numerical-parity or production approval.** Strict full-vocabulary logits/KV checks still fail at the original `atol=0.01, rtol=0.01`; the previous packing parity gate is not waived. Video and the experimental TensorScatter/SparsePaged conversions are outside this PR. ## Recorded pre-port experiment These results belong to immutable corrected exports built before porting the changes onto the newer main branch, and before the subsequent pool/window, adapter-resave, and runtime-opt-in review fixes. They are not a real-checkpoint re-export/validation of this PR branch. Export graph, source, runtime, and result hashes are retained in the experiment artifacts. Fixed AI2D-200 quantization-selection subset; not an independent full-AI2D evaluation: | Precision | HF reference | Native GenAI | |---|---:|---:| | BF16 | 167/200 | 166/200 | | INT4 | 162/200 | 166/200 | All four runs used matching original FP32 processor tensors and had zero inference failures. Three real-image prompts, each with prefill and two teacher-forced cached decodes, matched all nine first-token argmaxes for each native precision, but failed the strict full-logit/KV gates. Performance on the same frozen exports, batch 1, three single-image prompts of 333/353/356 input tokens: | Variant | TTFT (ms) | Decode (tokens/s) | Sampled device peak (GiB) | |---|---:|---:|---:| | Native BF16 | 127.44 | 78.99 | 56.51 | | Native INT4 | 128.59 | 87.77 | 34.98 | | HF BF16 | 1181.61 | 6.28 | 49.50 | | HF INT4 | 3351.30 | 2.43 | 17.28 | Each case used two discarded warmups and five measured trials, generating one first token and 64 greedy decode tokens. TTFT and throughput are medians of the three case medians. Full-vocabulary diagnostic readbacks were excluded from measured trials; EOS suppression was performance-only. HF used eager experts and SDPA without compilation; its Olive INT4 dequantizes on demand, unlike native QMoE. These are different implementations, not an isolated packing speedup. The selected A100 80GB had no competing compute process, but other GPUs shared the host with text evaluation. Peak memory is a 20ms NVML sample of the whole selected device, not exact model-only allocation. These are short-prompt results, not 2K/8K/16K/32K benchmarks. The experimental ORT 1.32 runtime was built from revision `942356fce6bdc78c016e4b902e49492fdf695453` plus local modifications, not an unmodified upstream release. Its shared-library SHA256 is `3f6b7ff482f8fbd4d3c639e807ed2ec35fa240c88c205dbee7822d7ec497e988`. Both native decoder graphs retained 30 packed-QKV projection records. Diagnostic profiles executed no Constant nodes, but still averaged 126 CUDA Cast nodes per decoder invocation. Constant folding does not justify deleting dynamic precision boundaries; these counts are not per-op GPU timings. ## Isolated-branch validation Against base `d69efe03a9f4c1ada744bc31dd01077a710b3c39`, offline and CPU-only: | Check | Result | |---|---| | Related runtime, package, expert-layout, adapter, and packing regressions | 656 passed, 2 skipped | | GQA rewrite regressions | 32 passed; 4 known offline baseline failures excluded from this invocation | | Public Gemma4 builder ABI | 10 passed | | Actual pinned-HF dynamic/static mask comparisons | 11 passed | | Changed-file Ruff, format, compilation, whitespace | Passed | | Strict mypy with resolved editable ONNX Script imports | Not clean on either revision; zero new normalized diagnostic instances | The four excluded GQA cases were separately reproduced on the exact clean base; unavailable cached remote configurations were not downloaded, and tests were not weakened. The strict typecheck produced 2,798 diagnostics on the head and 2,807 on base with no missing-import diagnostics; this is a baseline-relative result, not clean type approval. Sensitive regressions include BF16 3x3 pooling against the actual HF implementation, padded NaN/+Inf/-Inf handling, physical-window boundaries with internal padding and cached decode, shared-weight packing refusal followed by public attachment of both independent K adapters, source-preserving same-path/hardlink/symlink save rejection, and safe-destination package round-trips. Expanded validation caught and fixed a caller-owned `MatMulNBits` function regression: raw factory results must not claim registry ownership. The existing preservation test now passes unchanged on both base and head, and new tests retain automatic-fallback packing while preserving caller-owned functions/calls. ## Review disposition Two targeted five-reviewer rounds completed. They closed F1-F6 without a remaining grounded Major/Critical in their bounded follow-up. A subsequent adversarial probe promoted the final round's refolding ownership question to F7. The user explicitly approved fixing F7 and running precise validation rather than a third full-team round. | ID | Severity | Raised by | Finding / affected location | Disposition | |---|---|---|---|---| | F1 | Major | Deep | `models/gemma4.py`: premature low-precision pooling | Confirmed; fixed and actual-HF BF16/nonfinite regressions pass; Deep closed | | F2 | Major | Deep | `components/_common.py`: padding compressed local-window distances | Confirmed; fixed; physical-slot prefill/decode boundaries and actual HF masks pass; Deep closed | | F3 | Major | Code | `_qkv_projection_packing.py` / `adapter_io.py`: shared sources broke independent adapter restoration | Confirmed with corrected two-K-target public regression; conservative packing refusal fixes it; Code closed | | F4 | Major | Critical | CUDA implicitly authorized experimental GEGLU | Confirmed; default-off custom-runtime opt-in and persistent requirement metadata; Critical/Integration closed | | F5 | Major | Critical question + QA | Restored lazy slices hid backing files during alias saves | Confirmed; source ownership retained, aliases rejected before writes, safe-output round-trips pass; bounded lifecycle closed | | F6 | Major | Expanded validation | Raw `MatMulNBits` factory incorrectly claimed automatic function ownership | Confirmed by unchanged head/base test; caller-owned factory left unmarked, automatic registry provenance preserved; Code closed | | F7 | Major | Critical question + QA | Refolding replaced tensor wrappers and lost source ownership before first save | Confirmed; lead retains ownership on the in-memory model too; focused regression and unchanged adversarial probe now pass | For F7, the pre-fix generated fixture's external data file changed from 16,384 bytes to zero before a lazy read failed. With the fix, the identical instrument records **zero external reads, zero source write-open requests, and unchanged nonzero sizes/SHA256 digests**. The regression also saves elsewhere and checks exact reloaded payloads. Original model artifacts were not touched. F7's final small fix was not independently re-reviewed by the full team; it has the explicitly approved focused execution evidence. **Corrections/minority record:** Code withdrew its original claim that directly attaching a shared Q remained supported after folding: the exact base requires one initializer consumer. The real two-independent-K-target regression was retained and fixed. Integration withdrew the proposed fused-RoPE restriction for projection-only packing because `MatMul`/`MatMulNBits -> Split` does not match legacy packed-GQA unpacking. No unresolved reviewer objection to those corrections remains. **Attributed open questions:** Deep retains current-revision full-checkpoint logits/KV/packing qualification. Integration retains whether released ORT GenAI versions supply the split decoder's required `input_ids` alongside `inputs_embeds`; the experimental measured runner does not qualify every released runner. Readability's provenance limitation is explicit: the modified runtime patchset and complete experiment records are not included in this repository, so the binary identity is not a reproducible build recipe. **Residual risk / exclusions:** CPU-only branch validation, two skipped tests, four exact-base offline GQA failures, and non-clean baseline typing remain explicit. No new-branch full-checkpoint export, released-runtime qualification, long-context benchmark, CUDA Graph/shared-KV clearance, arbitrary model-replacement metadata guarantee, or concurrent filesystem-race test is claimed. The unselected-expert nonfinite hypothesis was refuted only for its exact FP16 finite-router probe, not universally. ASR, Qwen4, text benchmark implementation, and experimental static/paged-cache work are excluded. **Critical review threat-model appendix:** Inputs include checkpoint/quantization plans, serialized packing metadata and external files, PEFT factors, feature flags, and runtime pixels/positions/IDs/masks. Boundaries are metadata-authorized restoration, lazy backing-file lifetime, and explicit authorization of experimental native operators. Relevant risks are tampering/data integrity and denial of service; no new authentication or elevation path was established. Known source aliases fail before writing; this is not exhaustive transformation or concurrent-filesystem protection. ## Remaining validation - Strict BF16/INT4 full-logit/KV equivalence and packed-projection equivalence. - Real-checkpoint rebuild and revalidation on the main-based PR branch. - Long-context 2K/8K/16K/32K performance; shared/static KV and CUDA Graph evaluation. - Full AI2D and completion of the independent text-only task-quality pilot.
## Current follow-up validation and CI root cause Published follow-up `fd4513b5cc2e4692a2cde51fd8fdd57c4bb0878d`. The test override accepts base forward parameters, including keyword names and optional causal/static-cache arguments, while retaining the isolated MoE probe; 147 model/flag tests and 225 combined tests passed, with scoped lint passed. There are no production forward changes. The restoration assertion is now in `finally`, so `original` is always consumed on either control-flow exit; nested flag contexts/assertions are retained. CodeQL is unavailable locally; no new CodeQL clearance is claimed. The initial CI environment used TensorFlow 5.20.2 / Hugging Face Hub 2.3 / httpx 2, then `--no-deps` forced TensorFlow 5.16.dev0 whose metadata requires `huggingface-hub>=1.5,<2`; old Hub 1.x supplies the older httpx. This is a dependency-closure violation, not merely missing httpx, and is not changed by this PR. No CI code fix is included and CI green is not claimed. Evidence: [pinned TensorFlow setup](https://github.com/onnxruntime/mobius/blob/fd4513b5cc2e4692a2cde51fd8fdd57c4bb0878d/.github/workflows/ci.yml#L80), [Hub v1.5.0 setup.py](https://github.com/huggingface/huggingface_hub/blob/v1.5.0/setup.py), and [Hub v2.3.0 setup.py](https://github.com/huggingface/huggingface_hub/blob/v2.3.0/setup.py). Resolve the pinned TensorFlow dependency normally; use `--no-deps` only for VibeVoice plus pip-check/import smoke, with no automatic new PR scope. Strict numerical gates remain failed/unwaived. The full five-role, two-round review is already complete; this follow-up is a direct bounded review only. --------- Signed-off-by: titaiwang Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 --- docs/gemma4-runtime.md | 183 ++++++ docs/model-catalog.md | 4 + examples/gemma4_multimodal.py | 5 +- src/mobius/_configs/_quantization.py | 3 + src/mobius/_configs/_sub_configs.py | 2 + src/mobius/_configs/_vision_defaults.py | 1 + src/mobius/_constants.py | 2 + src/mobius/_flags.py | 16 + src/mobius/_flags_test.py | 25 + src/mobius/components/_common.py | 61 +- src/mobius/components/_common_test.py | 78 ++- src/mobius/components/_moe.py | 5 +- src/mobius/models/gemma4.py | 450 +++++++++++---- src/mobius/models/gemma4_test.py | 720 +++++++++++++++++++++++- src/mobius/tasks/_gemma4.py | 21 +- tests/build_graph/core_test.py | 8 +- tests/integration/gemma4_test.py | 120 ++-- tests/static_cache_bias_parity_test.py | 8 +- 18 files changed, 1477 insertions(+), 235 deletions(-) create mode 100644 docs/gemma4-runtime.md diff --git a/docs/gemma4-runtime.md b/docs/gemma4-runtime.md new file mode 100644 index 000000000..0acd1f66c --- /dev/null +++ b/docs/gemma4-runtime.md @@ -0,0 +1,183 @@ +# Gemma4 text and multimodal runtime qualification + +This runtime work is draft, not production or equivalence approved. It does not +qualify optimized text-only performance, ASR, or experimental fixed/paged-cache +execution. + +## HuggingFace alignment and input contract + +- Global decoder layers use causal attention only; image-block IDs must not + unmask future keys in global layers. +- Local decoder layers use + `valid_key AND (causal OR same_image) AND (key > query - window)`. + Same-image future keys can be visible, but an image block cannot bypass the + strict left window boundary. Distances use physical query/KV slots, including + past offsets; padding holes do not compress the window. Only current-chunk + image-block IDs are available; + these IDs are not persisted alongside the KV cache. +- The separate vision encoder's `pixel_values` input deliberately uses ONNX + `FLOAT`, including with BF16/FP16 projection weights. Supply FP32 processor + patches in `[0, 1]`, not BF16 patches or pre-normalized `[-1, 1]` values. + Normalization computes `2 * (pixels - 0.5)` in FP32 before casting to the + projection's working dtype. Existing consumers must update their input dtype; + `examples/gemma4_multimodal.py` demonstrates the new ABI. This does not change + the encoder-free `gemma4_unified` vision ABI. +- Checkpoint `router.scale` is preserved. The router first multiplies the + normalized working-dtype activation by the learned scale, then applies the + separate `hidden_size ** -0.5` scalar with FP32 opmath and rounds back to the + working dtype. Folding that scalar into a low-precision learned scale changes + rounding and is not equivalent. +- Checkpoint-enabled vision standardization retains `std_bias`/`std_scale` + directly on the vision submodel. Pooling coefficients and feature accumulation + are FP32, including the `1/9` coefficients for a 3x3 kernel. The pooled mean + then rounds to working dtype before FP32 scaling/standardization and the + final cast back for projection. + +## Native runtime requirements and limits + +Ordinary `execution_provider="cuda"` does **not** assert custom GEGLU support. +The default-off `gemma4_custom_geglu` flag explicitly opts into the experimental +custom CUDA ABI: set `MOBIUS_GEMMA4_CUSTOM_GEGLU=1` before importing Mobius, or +use `override_flags(gemma4_custom_geglu=True)` around construction, export and +weight preprocessing. Environment values are captured when the module-level +flags instance is created on first import; changing the environment afterward +does not update that instance. Without opt-in, float tanh-GELU experts use the portable +exact-activation fallback; quantized GEGLU experts fail at construction with an +explicit requirement message. Opt-in does not enable native GEGLU on non-CUDA +providers. Tanh-GELU must not be replaced by a sigmoid-based SwiGLU approximation. + +The portable fallback evaluates experts densely, but selects out unchosen expert +outputs before weighting. This prevents an overflowing unselected expert from +contaminating the result through `0 * Inf` or `0 * NaN`; nonfinite selected +experts still propagate. This safeguard is not full-checkpoint numerical or +performance qualification. + +Each emitted native `MoE`/`QMoE` operator persists +`mobius.gemma4.required_runtime_feature="custom-cuda-geglu-v1"` in its ONNX +metadata. This declares the experimental CUDA `activation_type="geglu"`, +`swiglu_fusion=1` requirement; it is not a claim of released ORT support or a +claim that the exporting machine used a particular runtime binary. + +The frozen pre-port experiments used an **experimental modified ORT 1.32 build**: +source `942356fce6bdc78c016e4b902e49492fdf695453` with seven dirty paths; binary +SHA-256 `3f6b7ff482f8fbd4d3c639e807ed2ec35fa240c88c205dbee7822d7ec497e988`. +This is not an unmodified upstream release. Its local patchset and full +experiment records are not included in this repository; these identity fields +do not provide a reproducible runtime build. Released ORT support is unverified. + +Dynamic vision-overlay decoders declare the arbitrary-attention-mask contract. +Both local and global layers retain ONNX `Attention`: reducing a mask to a valid +token count would discard internal padding holes even in causal global layers. +This preserves the full additive bias through CUDA optimization. The default +VLM qualification path uses dynamic KV cache, graph capture off, and past/present +buffer sharing off. + +Text-only export uses the existing `text_only=True` builder option (CLI +`--features text-only`) to remove the vision overlay. Its causal/windowed +attention can use the existing GQA path when supported by the execution provider. +GQA sequence lengths describe prefix validity, not arbitrary interior mask holes; +use a standard-ONNX attention path when those holes must be preserved. +That eligibility is separate from projection packing, custom GEGLU opt-in, and +runtime KV-sharing/graph-capture settings. No experimental static-cache or +paged-attention changes are part of this modeling work. + +## Recorded pre-port experiments, not validation of this branch + +The corrected real-weight native BF16 and INT4 VLM runs each scored +**166/200**; HF BF16 scored **167/200** and HF INT4 **162/200**. These results use +a fixed quantization-selection AI2D subset, not an independent full-AI2D +evaluation, and must not be described as such. + +Those runs used frozen corrected exports built from the original dirty source +worktree, with source hashes recorded in the originating experiment records. +They were **not** exports of this isolated origin/main-based branch. Newer main +has different leaf-layout/configuration features; these recorded scores (and +any separately recorded pre-port performance results) do not validate this +branch. The frozen exports also predate this branch's corrected FP32 pooling +and physical-slot window fixes. Publication requires the original source hashes, +checkpoint/quantization configuration, export identity, and runtime version to +remain attached to the experiment provenance. + +Strict full-vocabulary logit parity **failed 9/9**, and KV parity **failed +540/540**, at `atol=0.01, rtol=0.01`. Task-level scores do not establish numerical +equivalence or production readiness. The isolated +patch's tiny-graph and weight-layout tests are separate evidence. Full-checkpoint +re-export and revalidation on the ported branch remain **unverified**. No model +rebuild, download, GPU run, or performance measurement is implied by this note. +Lead review is required before publication. + +## Prior combined-branch evidence + +Offline, CPU-only validation against base +`d69efe03a9f4c1ada744bc31dd01077a710b3c39`: + +- Related runtime, package, expert-layout, adapter, and packing tests: + **656 passed, 2 skipped**. These include `models/gemma4_test.py`, + `components/_common_test.py`, `_model_package_test.py`, + `functions/matmul_nbits_test.py`, and + `rewrite_rules/_qkv_projection_packing_test.py` under `src/mobius/`. +- GQA rewrite tests: **32 passed**. Four remote-config cases were excluded from + that offline invocation after reproducing their failures on the exact clean + base; their tests were not changed. +- `tests/build_graph/core_test.py -k gemma4`: **10 passed**. + `tests/integration/gemma4_test.py::test_gemma4_bidirectional_mask_parity`: + **11 passed** against the actual pinned HF mask entry points. + +Changed-file lint, formatting, compilation, and whitespace checks passed. +Strict mypy with resolved editable ONNX Script imports has zero new normalized +diagnostic instances relative to base, but neither revision typechecks cleanly. +These checks do not replace the unresolved full-checkpoint numerical gates. +The modeling/pass split and subsequent review fixes require their own validation; +the historical counts above are not results for the reorganized revision. + +## Packed-checkpoint contract + +Native Gemma4 experts require an explicit decoder quantization config: +`quant_method="olive"`, integer-affine `bits=4`, a power-of-two +`group_size >= 16`, `sym`, and canonical `quantize_moe=true`. +Metadata also accepts the legacy `moe` alias; `quantize_moe` takes precedence +when both are supplied. Expert quantization opt-in and the custom-runtime flag +are separate requirements. +Compute dtype must be FP16/BF16 and expert activation must be tanh-GELU. +Both expert input dimensions must be divisible by the group size. Preserve the +authoritative HF source-module names and any component/per-projection override +or exclusion plan; unsupported expert layouts fail closed, while ordinary +projection plans remain handled by the shared component loader. + +For each layer with `E` experts, hidden size `H`, intermediate size `I`, and +group size `G`, Olive sidecars under `...layers.L.experts` are: + +| Stem | `_qweight` (UINT8) | `_scales` (floating point) | `_qzeros` (UINT8, asymmetric only) | +|---|---|---|---| +| `gate_up_proj` | `[E, 2*I, H/2]` | `[E, 2*I, H/G]` | `[E, 2*I, ceil((H/G)/2)]` | +| `down_proj` | `[E, H, I/2]` | `[E, H, I/G]` | `[E, H, ceil((I/G)/2)]` | + +Weights use contiguous gate-then-up rows before graph-level interleaving. +The adapter preserves packed codes and folds `router.per_expert_scale` into +down-projection block scales; low-precision folding changes rounding order +and is not claimed bitwise exact. The learned `router.scale` remains unchanged. +Incomplete, stale, or shape-incompatible expert sidecars are rejected. + +Float projection packing retains version-1 +`mobius.qkv_projection_sources` metadata on the `Split`: original initializer +names and output widths. Keep it through serialization so PEFT attachment can +restore independent projections after weight folding. INT4 projection packing +requires homogeneous bits/block size/accuracy level, scale dtype, zero-point +presence/layout, and a known shape-generic fallback body when one is registered. +Float groups whose source initializer or logical weight view has outside +consumers are conservatively left unpacked, preserving existing tied-weight +adapter attachment rather than inventing multiple-group restoration semantics. + +Restored projections remain lazy slices, not eager copies of a large packed +checkpoint. Their backing-file ownership is registered with the package save +alias guard. Saving over a backing file (same path, hard link or symlink) is +rejected before any output is written, with the source unchanged. Use a distinct, +non-aliasing output location; ordinary packages without such derived lazy +dependencies retain their existing save behavior. + +Backing-file ownership is also retained on the in-memory model, so initializer +folding cannot hide it by replacing the restored tensor wrappers. A regression +refolds the lazy projections before saving, checks refusal before any external +read or source mutation, then verifies an exact round-trip at a safe destination. +This metadata is transient, not serialized provenance or a guarantee for arbitrary +model-replacement transformations. diff --git a/docs/model-catalog.md b/docs/model-catalog.md index c42960214..7659e1d31 100644 --- a/docs/model-catalog.md +++ b/docs/model-catalog.md @@ -112,6 +112,10 @@ Models combining Mamba/SSM layers with transformer attention layers. Models that process both images and text. +For Gemma4 26B-A4B, see the [draft runtime qualification and FP32 pixel input +contract](gemma4-runtime.md), including the failed strict parity gate and limits +of the prior evaluation evidence. + | Model Type | Module Class | Task | Example HuggingFace Model | |---|---|---|---| | `llava` | `LLaVAModel` | `vision-language` | `llava-hf/llava-1.5-7b-hf` | diff --git a/examples/gemma4_multimodal.py b/examples/gemma4_multimodal.py index e55104121..df3544be5 100644 --- a/examples/gemma4_multimodal.py +++ b/examples/gemma4_multimodal.py @@ -132,7 +132,7 @@ def prepare_vision_feeds( Args: processor: ``AutoProcessor`` loaded for the Gemma 4 model. image_path: Path to a local image file (JPEG, PNG, etc.). - np_dtype: Numpy dtype for pixel values (default float32). + np_dtype: Retained for caller compatibility; pixels always use float32. Returns: ``{"pixel_values": [B, N, 3*P^2], "pixel_position_ids": int64[B, N, 2]}`` @@ -143,7 +143,8 @@ def prepare_vision_feeds( # Gemma 4 processor requires a text argument alongside images. # Use a placeholder image token so the processor computes correct dims. processed = processor(images=img, text="", return_tensors="np") - pixel_values = processed["pixel_values"].astype(np_dtype) + # The vision ABI is FLOAT even when decoder/projection weights use BF16. + pixel_values = processed["pixel_values"].astype(np.float32) # The HF processor returns "image_position_ids"; our ONNX vision model input is "pixel_position_ids" pixel_position_ids = processed["image_position_ids"].astype(np.int64) return {"pixel_values": pixel_values, "pixel_position_ids": pixel_position_ids} diff --git a/src/mobius/_configs/_quantization.py b/src/mobius/_configs/_quantization.py index defb49ffc..8c124bb33 100644 --- a/src/mobius/_configs/_quantization.py +++ b/src/mobius/_configs/_quantization.py @@ -108,6 +108,8 @@ class QuantizationConfig: # positional construction, so inserting a field earlier would silently # change the meaning of existing callers' arguments. weight_format: QuantizedWeightFormat = QuantizedWeightFormat.INTEGER_AFFINE + # Olive explicitly opts fused expert tensors into native QMoE packing. + quantize_moe: bool = False def __post_init__(self) -> None: """Normalize serialized enum values without inferring storage semantics.""" @@ -249,6 +251,7 @@ def from_value( quantize_embeddings=bool(qc.get("embeds")), quantize_lm_head=bool(qc.get("lm_head")), quantize_vision=bool(qc.get("quantize_vision")), + quantize_moe=bool(qc.get("quantize_moe", qc.get("moe", False))), tie_word_embeddings=bool(qc.get("tie_word_embeddings")), modules_to_not_convert=exclusions, overrides=overrides, diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index a4bca1539..0b5381ef1 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/src/mobius/_configs/_sub_configs.py @@ -108,6 +108,8 @@ class VisionConfig: window_kernel_size: tuple[int, int] = (2, 2) merge_kernel_size: tuple[int, int] = (2, 2) merger_times: int = 1 + # Gemma4 post-pooling checkpoint affine transform, computed in FP32. + standardize: bool = False @dataclasses.dataclass diff --git a/src/mobius/_configs/_vision_defaults.py b/src/mobius/_configs/_vision_defaults.py index eabb9af36..63c4b8adb 100644 --- a/src/mobius/_configs/_vision_defaults.py +++ b/src/mobius/_configs/_vision_defaults.py @@ -97,6 +97,7 @@ def apply_vision_defaults(config, parent_config, model_type: str, fields: dict) hidden_act=getattr(vc, "hidden_act", None), use_clipped_linears=getattr(vc, "use_clipped_linears", False), position_embedding_size=getattr(vc, "position_embedding_size", None), + standardize=getattr(vc, "standardize", False), ) # This runs before per-model hooks, so plain assignment is safe — any # per-model hook may overwrite later via fields.update(...). diff --git a/src/mobius/_constants.py b/src/mobius/_constants.py index 7781a3f27..589af1a31 100644 --- a/src/mobius/_constants.py +++ b/src/mobius/_constants.py @@ -24,6 +24,8 @@ 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/_flags.py b/src/mobius/_flags.py index e17fe621c..9dfa16659 100644 --- a/src/mobius/_flags.py +++ b/src/mobius/_flags.py @@ -97,6 +97,11 @@ class _Flags: block overlay + padding) on the external-KV static-cache ``Attention`` path (``is_causal=0``) for float-bias decoders, instead of the maskless ``is_causal=1`` default. + * - ``gemma4_custom_geglu`` + - ``MOBIUS_GEMMA4_CUSTOM_GEGLU`` + - ``False`` + - Opt in to the experimental custom CUDA GEGLU MoE/QMoE ABI. + Ordinary CUDA selection does not imply this runtime support. """ suppress_dedup_warning: bool = dataclasses.field( @@ -178,6 +183,17 @@ class _Flags: unchanged, so no shipped model's graph changes unless this flag is set. """ + gemma4_custom_geglu: bool = dataclasses.field( + default_factory=lambda: _env_bool("MOBIUS_GEMMA4_CUSTOM_GEGLU", False) + ) + """Opt in to Gemma4's experimental custom CUDA tanh-GEGLU MoE/QMoE ABI. + + Requires a separately qualified custom runtime, not an unmodified released + ORT wheel. Default float exports use the portable exact-activation fallback; + quantized experts fail closed without this opt-in. The flag does not enable + native dispatch for non-CUDA providers. + """ + allow_dense_moe_experts: bool = dataclasses.field( default_factory=lambda: _env_bool("MOBIUS_ALLOW_DENSE_MOE_EXPERTS", False) ) diff --git a/src/mobius/_flags_test.py b/src/mobius/_flags_test.py index 412aa9953..66a68a0f6 100644 --- a/src/mobius/_flags_test.py +++ b/src/mobius/_flags_test.py @@ -18,6 +18,10 @@ def test_suppress_dedup_warning_default_on(self, monkeypatch): f = _flags._Flags() assert f.suppress_dedup_warning is True + def test_gemma4_custom_geglu_default_off(self, monkeypatch): + monkeypatch.delenv("MOBIUS_GEMMA4_CUSTOM_GEGLU", raising=False) + assert _flags._Flags().gemma4_custom_geglu is False + class TestEnvVarOverride: """Flags are read from environment variables at construction time.""" @@ -39,6 +43,13 @@ def test_unknown_value_falls_back_to_default(self, monkeypatch): f = _flags._Flags() assert f.suppress_dedup_warning is True # default + @pytest.mark.parametrize( + ("value", "expected"), [("1", True), ("true", True), ("0", False), ("maybe", False)] + ) + def test_gemma4_custom_geglu_opt_in(self, monkeypatch, value, expected): + monkeypatch.setenv("MOBIUS_GEMMA4_CUSTOM_GEGLU", value) + assert _flags._Flags().gemma4_custom_geglu is expected + class TestProgrammaticOverride: """Flags can be assigned directly on the singleton.""" @@ -80,6 +91,20 @@ def test_nested_overrides(self): assert _flags.flags.suppress_dedup_warning is not original assert _flags.flags.suppress_dedup_warning is original + def test_custom_geglu_override_is_scoped_even_on_exception(self): + original = _flags.flags.gemma4_custom_geglu + try: + with _flags.override_flags(gemma4_custom_geglu=False): + assert not _flags.flags.gemma4_custom_geglu + with pytest.raises(RuntimeError, match="test restoration"): + assert not _flags.flags.gemma4_custom_geglu + with _flags.override_flags(gemma4_custom_geglu=True): + assert _flags.flags.gemma4_custom_geglu + raise RuntimeError("test restoration") + assert not _flags.flags.gemma4_custom_geglu + finally: + assert _flags.flags.gemma4_custom_geglu is original + def test_unknown_flag_raises_value_error(self): with ( pytest.raises(ValueError, match="Unknown flag"), diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py index 3254a25e8..6019cde81 100644 --- a/src/mobius/components/_common.py +++ b/src/mobius/components/_common.py @@ -179,9 +179,10 @@ def create_attention_bias( (batch_size, query_length) giving a contiguous vision-block id per current-sequence position (``>= 0`` for vision tokens in the same block, ``-1`` for text). When provided, a bidirectional - "blockwise overlay" is OR-ed onto the causal (and sliding) mask: + "blockwise overlay" is OR-ed onto the causal mask: two positions in the same block (same id ``>= 0``) may attend to - each other regardless of causal order. This mirrors HuggingFace + each other regardless of causal order, subject to the sliding + window's left boundary. This mirrors HuggingFace ``blockwise_overlay`` for Gemma4 ``use_bidirectional_attention``. The returned bias bakes in causal + sliding + padding + blockwise, so the consuming ``Attention`` op MUST be called with @@ -198,8 +199,13 @@ def create_attention_bias( def as_numeric_mask(value): return op.Cast(value, to=ir.DataType.INT64) - # cumsum on attention_mask gives indices - all_indices = op.CumSum(attention_mask, 1) # axis=1 + # Distances are physical KV slots, not counts of valid tokens. An internal + # padding hole must not pull an old key back inside the sliding window. + # CumSum of ones avoids Range on providers (e.g. QNN) without a Range kernel. + slots = op.ConstantOfShape( + op.Shape(attention_mask), value=ir.tensor(np.array([1], dtype=np.int64)) + ) + all_indices = op.CumSum(slots, 1) # 1-based physical slots; differences are unchanged # kv_indices: (batch_size, 1, total_length) kv_indices = op.Unsqueeze(all_indices, [1]) @@ -228,27 +234,15 @@ def as_numeric_mask(value): if openvino_numeric_mask: full_mask = as_numeric_mask(full_mask) - if sliding_window is not None: - # Also mask out positions too far away - dist = op.Sub(q_indices, kv_indices) - within_window = op.Less(dist, sliding_window) - full_mask = ( - op.Mul(full_mask, as_numeric_mask(within_window)) - if openvino_numeric_mask - else op.And(full_mask, within_window) - ) - if block_sequence_ids is not None: - # Bidirectional vision-block overlay (OR-ed onto causal/sliding mask, - # BEFORE the padding AND, matching HF blockwise_overlay ordering). + # Apply the image overlay before the window and padding constraints. # # q_group: block id per query position -> (batch, query_length, 1). # block_sequence_ids covers the current input (== the query), so it # aligns 1:1 with the query positions. q_group = op.Unsqueeze(block_sequence_ids, [2]) # (B, q_len, 1) # kv_group: block id per kv position -> (batch, 1, total_length). - # The kv axis spans past + current; past positions are text in the - # cache, so left-pad with -1 to width total_length. + # Only current-chunk IDs are supplied; left-pad past slots with -1. pad_width = op.Sub(total_length, query_length) # [1], == past length zero_1d = op.Constant(value_ints=[0]) # Pad spec for a 2-D tensor [B, q_len]: [b_begin, s_begin, b_end, s_end]. @@ -272,6 +266,16 @@ def as_numeric_mask(value): same_block = op.And(same_group, valid_group) full_mask = op.Or(full_mask, same_block) + if sliding_window is not None: + # Same-image future keys are visible, but old keys cannot escape the window. + dist = op.Sub(q_indices, kv_indices) + within_window = op.Less(dist, sliding_window) + full_mask = ( + op.Mul(full_mask, as_numeric_mask(within_window)) + if openvino_numeric_mask + else op.And(full_mask, within_window) + ) + # Combine with attention_mask mask_value = float(dtype.min) if openvino_numeric_mask: @@ -317,7 +321,7 @@ def create_static_cache_attention_bias( ``total_length``. * **KV positions** are dense slot ids ``0 .. max_seq_len - 1`` (the absolute cache position each slot holds), so the causal/sliding comparisons use - absolute positions directly instead of cumsum indices. + absolute positions directly, as does the dynamic helper. * **Q positions** are absolute: ``write_indices[b] + arange(S_q)`` — the cache slots the current chunk is being scattered into. * **KV validity** is ``slot < nonpad_kv_seqlen[b]`` (replaces the cumsum @@ -349,8 +353,9 @@ def create_static_cache_attention_bias( block_sequence_ids: Optional ``[B, S_q]`` int64 block id per *current* query position (``>= 0`` for vision tokens sharing a block, ``-1`` for text). When set, a bidirectional overlay is OR-ed onto the - causal/sliding mask: two positions in the same block may attend to - each other regardless of causal order — mirroring HuggingFace + causal mask: two positions in the same block may attend to + each other regardless of causal order, subject to the window's + left boundary — mirroring HuggingFace ``blockwise_overlay`` for Gemma4. The per-slot KV block ids are built by scattering ``block_sequence_ids`` into a ``-1``-filled ``[B, max_seq_len]`` buffer at ``write_indices`` (the same scatter @@ -406,15 +411,8 @@ def create_static_cache_attention_bias( # Broadcasting (B, S_q, 1) >= (max_seq_len,) -> (B, S_q, max_seq_len). full_mask = op.GreaterOrEqual(q_abs, kv_slots) - if sliding_window is not None: - # Local window: keep slots within `sliding_window` of the query. - dist = op.Sub(q_abs, kv_slots) # (B, S_q, max_seq_len) - within_window = op.Less(dist, op.Constant(value_int=sliding_window)) - full_mask = op.And(full_mask, within_window) - if block_sequence_ids is not None: - # Bidirectional vision-block overlay (OR-ed BEFORE the padding AND, - # matching create_attention_bias / HF blockwise_overlay ordering). + # Apply the overlay before the left-window and padding constraints. # # q_group: block id per current query position -> (B, S_q, 1). q_group = op.Unsqueeze(block_sequence_ids, [2]) # (B, S_q, 1) @@ -438,6 +436,11 @@ def create_static_cache_attention_bias( ) full_mask = op.Or(full_mask, same_block) + if sliding_window is not None: + dist = op.Sub(q_abs, kv_slots) # (B, S_q, max_seq_len) + within_window = op.Less(dist, op.Constant(value_int=sliding_window)) + full_mask = op.And(full_mask, within_window) + # KV validity (padding): a slot is real iff it lies inside the valid # prefix, slot < nonpad_kv_seqlen[b]. Broadcasting (max_seq_len,) < # (B, 1) -> (B, max_seq_len) -> (B, 1, max_seq_len). diff --git a/src/mobius/components/_common_test.py b/src/mobius/components/_common_test.py index fad559d5c..05787a74b 100644 --- a/src/mobius/components/_common_test.py +++ b/src/mobius/components/_common_test.py @@ -7,6 +7,7 @@ import numpy as np import onnx_ir as ir +import pytest from mobius._testing import count_op_type, create_test_builder, create_test_input from mobius._testing.ort_inference import OnnxModelSession @@ -140,7 +141,7 @@ def build(execution_provider): op, input_ids, attention_mask, - sliding_window=4, + sliding_window=2, block_sequence_ids=block_sequence_ids, ) bias.name = "bias" @@ -165,12 +166,61 @@ def build(execution_provider): qnn_bias = OnnxModelSession(qnn_model, device="cpu").run(feeds)["bias"] np.testing.assert_array_equal(openvino_bias, qnn_bias) + # Exercise a same-image block wider than the window on both code paths. + feeds["attention_mask"] = np.ones((1, 8), dtype=np.int64) + feeds["block_sequence_ids"] = np.zeros((1, 4), dtype=np.int64) + for model in (openvino_model, qnn_model): + bias = OnnxModelSession(model, device="cpu").run(feeds)["bias"] + assert bias[0, 0, 0, 7] == 0 # Same-image future key. + assert bias[0, 0, 3, 4] < -1 # Old same-image key. + assert bias[0, 0, 3, 5] < -1 # Strict window boundary. + assert bias[0, 0, 3, 6] == 0 + + @pytest.mark.parametrize("execution_provider", ["default", "openvino"]) + @pytest.mark.parametrize("query_length", [1, 2, 1025]) + def test_internal_padding_hole_does_not_shrink_physical_window( + self, execution_provider, query_length + ): + from mobius._build_context import build_context + from mobius._execution_providers import get_ep + + builder, op, graph = create_test_builder() + ii = create_test_input( + builder, "input_ids", [1, query_length], dtype=ir.DataType.INT64 + ) + am = create_test_input(builder, "attention_mask", [1, 1025], dtype=ir.DataType.INT64) + bk = create_test_input( + builder, "block_sequence_ids", [1, query_length], dtype=ir.DataType.INT64 + ) + with build_context(get_ep(execution_provider)): + bias = create_attention_bias( + op, ii, am, sliding_window=1024, block_sequence_ids=bk + ) + bias.name = "bias" + graph.outputs.append(bias) + mask = np.ones((1, 1025), dtype=np.int64) + mask[0, 1] = 0 + actual = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu").run( + { + "input_ids": np.zeros((1, query_length), dtype=np.int64), + "attention_mask": mask, + "block_sequence_ids": np.zeros((1, query_length), dtype=np.int64), + } + )["bias"] + assert actual[0, 0, -1, 0] < -1 # q=1024, k=0: exact left boundary. + assert actual[0, 0, -1, 1] < -1 # Invalid key stays masked. + assert actual[0, 0, -1, 2] == 0 + if query_length > 1: + assert actual[0, 0, -2, 0] == 0 # q=1023: boundary not yet reached. + if query_length == 1025: + assert actual[0, 0, 0, 1024] == 0 # Same-image future key. + class TestBlockwiseAttentionBias: """Numerically verify the Gemma4 vision-block bidirectional overlay. ``create_attention_bias(block_sequence_ids=...)`` must bake the FULL mask - (causal [+ sliding] OR same-block, AND padding) so the Attention op can be + ((causal OR same-block) AND sliding AND padding) so the Attention op can be called with ``is_causal=0``. We build the graph, run it via ORT, and compare the attended pattern (bias == 0) against a numpy reference. """ @@ -195,15 +245,15 @@ def _build(sliding): @staticmethod def _ref(block_ids, attn, sliding): - cumsum = np.cumsum(attn) - qi = cumsum[:, None] - ki = cumsum[None, :] + slots = np.arange(len(attn)) + qi = slots[:, None] + ki = slots[None, :] m = qi >= ki - if sliding is not None: - m = m & ((qi - ki) < sliding) qg = np.array(block_ids)[:, None] kg = np.array(block_ids)[None, :] m = m | ((qg == kg) & (qg >= 0)) + if sliding is not None: + m = m & ((qi - ki) < sliding) return m & (np.array(attn)[None, :].astype(bool)) def _run(self, sliding, block_ids, attn): @@ -231,17 +281,17 @@ def test_multi_block_full_attention(self): assert not attended[3, 4] def test_block_wider_than_sliding_window(self): - # Single block spanning positions 1..4 with a window of 2: the block - # must escape the sliding window (same-block OR overrides the window). + # Same-image future keys are visible, but the strict left boundary applies. attended, expected = self._run(2, [-1, 0, 0, 0, 0, -1, -1, -1], [1] * 8) assert np.array_equal(attended, expected) - assert attended[4, 1] # distance 3 >= window, allowed via same block + assert attended[1, 4] + assert not attended[4, 1] + assert not attended[4, 2] # key == query - window is masked. + assert attended[4, 3] def test_padding_still_masked(self): - # Last two positions are padding (attention_mask == 0). - attended, expected = self._run( - 2, [-1, 0, 0, -1, -1, -1, -1, -1], [1, 1, 1, 1, 1, 1, 0, 0] - ) + # Padding must stay masked even when it carries the same image-block ID. + attended, expected = self._run(2, [-1, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 0, 0]) assert np.array_equal(attended, expected) assert not attended[:, 6:].any() # nothing attends to padding diff --git a/src/mobius/components/_moe.py b/src/mobius/components/_moe.py index 563dbd1e4..63bc6083d 100644 --- a/src/mobius/components/_moe.py +++ b/src/mobius/components/_moe.py @@ -8,6 +8,7 @@ import dataclasses import math from collections.abc import Callable +from typing import ClassVar import onnx_ir as ir from onnxscript import OpBuilder, nn @@ -446,6 +447,8 @@ class MoELayer(nn.Module): and compatibility path, not the grouped-expert performance representation. """ + _qmoe_activation_type: ClassVar[str] = "swiglu" + def __init__( self, config: ArchitectureConfig, @@ -637,7 +640,7 @@ def _qmoe_forward(self, op: OpBuilder, hidden_states: ir.Value, *gate_args): self.fc2_experts_zero_points, None, router_weights, - activation_type="swiglu", + activation_type=self._qmoe_activation_type, normalize_routing_weights=int(normalize), k=self.top_k, # A uniform override can use the global attribute directly. diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 484957e89..1b42629ee 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -35,9 +35,13 @@ from mobius._build_context import ep_capabilities, is_prefill_prefix_pruning_enabled from mobius._configs import ArchitectureConfig, Gemma4Config, QuantizationConfig +from mobius._constants import GEMMA4_CUSTOM_GEGLU_METADATA_KEY +from mobius._flags import flags from mobius._weight_utils import ( + QMoEQuantizationLayout, is_packed_quant_key, preprocess_quantized_weights, + resolve_qmoe_quantization, vlm_decoder_weights, vlm_embedding_weights, ) @@ -47,6 +51,7 @@ Embedding, LayerNorm, Linear, + MoELayer, RMSNorm, ScaledQuantizedEmbedding, ScaleFreeRMSNorm, @@ -91,6 +96,146 @@ ), } +_TANH_GELU_ACTIVATIONS = frozenset({"gelu_pytorch_tanh", "gelu_new", "gelu_fast"}) + + +def _mark_custom_geglu_runtime(result: ir.Value) -> ir.Value: + """Persist the experimental operator ABI requirement, not a release claim.""" + node = result.producer() + assert node is not None + node.metadata_props[GEMMA4_CUSTOM_GEGLU_METADATA_KEY] = "custom-cuda-geglu-v1" + return result + + +def _expert_source_paths(layer_idx: int) -> tuple[str, ...]: + return ( + f"model.language_model.layers.{layer_idx}", + f"model.layers.{layer_idx}", + ) + + +def _expert_quantization( + config: Gemma4Config, layer_idx: int +) -> QMoEQuantizationLayout | None: + """Resolve native experts without bypassing main's per-projection module plan.""" + if not config.enable_moe_block: + return None + quantization = _active_quantization(config.quantization_for("decoder")) + if quantization is None or not quantization.quantize_moe: + return None + if not flags.gemma4_custom_geglu: + raise NotImplementedError( + "Quantized Gemma4 GEGLU experts require explicit custom-runtime opt-in: " + "MOBIUS_GEMMA4_CUSTOM_GEGLU=1 (or override_flags(gemma4_custom_geglu=True)). " + "Unmodified released ORT support is unverified." + ) + if config.dtype not in {ir.DataType.FLOAT16, ir.DataType.BFLOAT16}: + raise NotImplementedError("Quantized Gemma4 MoE requires FLOAT16 or BFLOAT16.") + layout = resolve_qmoe_quantization(quantization, _expert_source_paths(layer_idx)) + if ( + config.disable_qmoe + or config.hidden_act not in _TANH_GELU_ACTIVATIONS + or layout is None + or any( + projection.bits != 4 or projection.quant_method != "olive" + for projection in (layout.fc1, layout.fc2) + ) + ): + raise NotImplementedError( + "Quantized Gemma4 MoE experts require Olive integer-affine INT4, " + "tanh-GELU, and native QMoE enabled." + ) + return layout + + +def _extract_native_expert_weights( + state_dict: dict[str, torch.Tensor], config: Gemma4Config, *, prefix: str +) -> dict[str, torch.Tensor]: + """Consume only Gemma4 expert sidecars, leaving ordinary module plans untouched. + + Return canonical QMoE groups separately so the legacy/root Olive converter + cannot interpret ``fc1_scales`` as another raw ``*_scales`` sidecar. The shared + component loader already recognizes complete canonical MoELayer groups. + """ + result: dict[str, torch.Tensor] = {} + for index in range(config.num_hidden_layers): + layout = _expert_quantization(config, index) + if layout is None: + continue + assert config.moe_intermediate_size is not None + assert config.num_local_experts is not None + root = f"{prefix}layers.{index}" + for projection, quantization, output_size, input_size, target in ( + ( + "gate_up_proj", + layout.fc1, + 2 * config.moe_intermediate_size, + config.hidden_size, + "fc1", + ), + ("down_proj", layout.fc2, config.hidden_size, config.moe_intermediate_size, "fc2"), + ): + stem = f"{root}.experts.{projection}" + keys = [key for key in state_dict if key.startswith(stem)] + if not keys: + continue + group = quantization.group_size + if input_size % group: + raise ValueError(f"Gemma4 QMoE {stem}: input size must divide group_size.") + blocks = input_size // group + shapes = { + "_qweight": (config.num_local_experts, output_size, input_size // 2), + "_scales": (config.num_local_experts, output_size, blocks), + } + if not quantization.sym: + shapes["_qzeros"] = (config.num_local_experts, output_size, (blocks + 1) // 2) + expected = {stem + suffix for suffix in shapes} + if set(keys) != expected: + raise ValueError( + f"Gemma4 QMoE {stem}: expected sidecars {sorted(expected)}, " + f"got {sorted(keys)}." + ) + suffixes = { + "_qweight": "experts_weights", + "_scales": "scales", + "_qzeros": "experts_zero_points", + } + for suffix, shape in shapes.items(): + value = state_dict[stem + suffix] + if tuple(value.shape) != shape: + raise ValueError( + f"Gemma4 QMoE {stem + suffix}: expected shape {shape}, " + f"got {tuple(value.shape)}." + ) + if suffix == "_scales": + if ( + not value.is_floating_point() + or not torch.isfinite(value).all() + or torch.any(value < 0) + ): + raise ValueError(f"Gemma4 QMoE {stem}: invalid quantization scales.") + elif value.dtype != torch.uint8: + raise ValueError(f"Gemma4 QMoE {stem + suffix} must be uint8.") + result[f"{root}.moe.{target}_{suffixes[suffix]}"] = state_dict.pop( + stem + suffix + ) + # Fold expert output scales into block scales, never the packed INT4 codes. + scale_key = f"{root}.router.per_expert_scale" + fc2_key = f"{root}.moe.fc2_scales" + if fc2_key in result and scale_key in state_dict: + scale = state_dict[scale_key] + if ( + tuple(scale.shape) != (config.num_local_experts,) + or not torch.isfinite(scale).all() + ): + raise ValueError(f"Gemma4 QMoE {scale_key}: invalid expert output scales.") + result[fc2_key] = result[fc2_key] * scale.to(result[fc2_key].dtype)[:, None, None] + state_dict[scale_key] = torch.ones_like(scale) + for key in list(state_dict): + if key.startswith(f"{root}.router."): + result[key.replace(".router.", ".moe.gate.")] = state_dict.pop(key) + return result + def _split_per_layer_projection_weight( state_dict: dict[str, torch.Tensor], @@ -440,7 +585,7 @@ def _remap_moe_expert_weights( state_dict: dict[str, torch.Tensor], config: Gemma4Config, ) -> None: - """Rename HF MoE expert weights and fold router scale in-place. + """Rename HF MoE expert weights and fold expert output scales in-place. Shared by ``Gemma4CausalLMModel`` and ``Gemma4Model`` to avoid duplicating the rename/fold logic. @@ -479,13 +624,9 @@ def _remap_moe_expert_weights( new_key = key.replace(".experts.down_proj", ".fc2_experts_weights") state_dict[new_key] = state_dict.pop(key) - # Fold hidden_size^-0.5 into router.scale if config.enable_moe_block: - scale_factor = float(config.hidden_size**-0.5) pes_by_prefix: dict[str, torch.Tensor] = {} for key in list(state_dict.keys()): - if ".router.scale" in key and ".per_expert_scale" not in key: - state_dict[key] = state_dict[key] * scale_factor if key.endswith(".router.per_expert_scale"): pes_by_prefix[key[: -len(".router.per_expert_scale")]] = state_dict[key] # Fold per_expert_scale into the expert down projection (fc2). HF's router @@ -807,12 +948,14 @@ def forward( y_pos = op.Gather(pixel_position_ids, op.Constant(value_int=1), axis=2) # [B, T] is_padding = op.And(op.Equal(x_pos, neg_one), op.Equal(y_pos, neg_one)) # [B, T] bool - # zero out padding hidden states: features * (1 - is_padding) - not_padding_f = op.CastLike( - op.Cast(op.Not(is_padding), to=ir.DataType.FLOAT), vision_features - ) - not_padding_f = op.Unsqueeze(not_padding_f, [2]) # [B, T, 1] - vision_features = op.Mul(vision_features, not_padding_f) # [B, T, D] + # Match HF masked_fill, including padded NaN/Inf, in the FP32 pooling + # path. Keep the original features as the mean-rounding dtype reference. + features_f32 = op.Cast(vision_features, to=ir.DataType.FLOAT) + masked_features_f32 = op.Where( + op.Unsqueeze(is_padding, [2]), + op.Constant(value_float=0.0), + features_f32, + ) # [B, T, D] # --- 2. Clamp positions to [0, ∞) so padding doesn't break Div ------ clamped_x = op.Max(x_pos, op.Constant(value_int=0)) # [B, T] @@ -854,23 +997,25 @@ def forward( on_val = 1.0 / float(k2) one_hot_vals = op.Constant(value_floats=[0.0, on_val]) weights = op.OneHot(kernel_idxs, valid_depth, one_hot_vals) # [B, T, valid_depth] f32 - # Cast to model dtype for the subsequent MatMul - weights = op.CastLike(weights, vision_features) - # --- 6. Weighted sum: [B, valid_depth, T] @ [B, T, D] --------------- + # HF accumulates FP32 coefficients and FP32 features, then rounds the + # pooled mean once to working dtype. BF16 coefficients (notably 1/9) + # change that mean even if the MatMul itself uses FP32 accumulation. weights_t = op.Transpose(weights, perm=[0, 2, 1]) # [B, valid_depth, T] - output = op.MatMul(weights_t, vision_features) # [B, valid_depth, D] + output = op.MatMul(weights_t, masked_features_f32) # [B, valid_depth, D] + output = op.CastLike(output, vision_features) # --- 7. Scale by sqrt(hidden_size) matching HF VisionPooler ---------- - scale = op.CastLike(op.Constant(value_float=self._pooler_scale), output) - return op.Mul(output, scale) + # HF rounds the pooled mean to working dtype before FP32 scaling/standardization. + output = op.Cast(output, to=ir.DataType.FLOAT) + return op.Mul(output, op.Constant(value_float=self._pooler_scale)) class _Gemma4VisionPatchEmbedder(nn.Module): """Gemma4 patch embedder: linear projection + 2D position lookup. Inputs: - - ``pixel_values [B, N, 3*patch_size^2]``: pre-patchified, normalized to ``[-1, 1]`` + - ``pixel_values [B, N, 3*patch_size^2]``: FP32 processor patches in ``[0, 1]`` - ``pixel_position_ids [B, N, 2]``: (x, y) coordinates for each patch Output: ``[B, N, hidden_size]`` @@ -905,10 +1050,11 @@ def forward( position embeddings are zeroed exactly as HF does in ``Gemma4VisionPatchEmbedder._position_embeddings``. """ - # pixel_values in [0,1] -> normalize to [-1, 1]: 2*(v - 0.5) = 2v - 1 - two = op.CastLike(op.Constant(value_float=2.0), pixel_values) - one = op.CastLike(op.Constant(value_float=1.0), pixel_values) - pixel_values = op.Sub(op.Mul(pixel_values, two), one) + # HF normalizes processor pixels in FP32, then rounds to the projection dtype. + # Casting processor pixels first loses information, especially for BF16. + pixel_values = op.Cast(pixel_values, to=ir.DataType.FLOAT) + pixel_values = op.Mul(op.Sub(pixel_values, 0.5), 2.0) + pixel_values = op.CastLike(pixel_values, self.position_embedding_table) hidden_states = self.input_proj(op, pixel_values) # [B, N, hidden] # Detect padding patches: x-coord == -1 means the patch is padding. @@ -949,6 +1095,7 @@ class _Gemma4VisionEncoderCore(nn.Module): def __init__(self, config: ArchitectureConfig): super().__init__() vc = config.vision # VisionConfig for the SigLIP encoder + linear_class, clippable_linear_class = _vision_linear_classes(config) self.patch_embedder = _Gemma4VisionPatchEmbedder( patch_size=vc.patch_size or 16, @@ -1326,7 +1473,7 @@ def forward( past_key = past_key_value[0] if past_key_value is not None else None past_value = past_key_value[1] if past_key_value is not None else None - gqa_attrs: dict = { + gqa_attrs = { "num_heads": self.num_attention_heads, "kv_num_heads": self.num_key_value_heads, "scale": self.scaling, @@ -1442,44 +1589,45 @@ def forward( class _Gemma4MoeRouter(nn.Module): - """Gemma4 MoE router: scale-free RMSNorm → learned scale → linear → softmax. + """Gemma4 MoE router: scale-free RMSNorm → learned scale → raw logits. Matches ``Gemma4TextRouter`` in HuggingFace. - Parameters (aligned with HF state_dict, with one weight-fold): - - ``scale`` [hidden_size]: ``router.scale * hidden_size^-0.5`` (folded during - ``preprocess_weights`` — see ``Gemma4CausalLMModel.preprocess_weights``) + Parameters (aligned with HF state_dict): + - ``scale`` [hidden_size]: unmodified learned router scale - ``proj.weight`` [num_experts, hidden_size]: router logits (no bias) - ``per_expert_scale`` [num_experts]: per-expert output weight scaling ``router.norm`` is scale-free (no learnable parameter, omitted from state_dict). - The ``hidden_size^-0.5`` scale factor is absorbed into ``self.scale`` at weight-load - time to avoid graph-level float-constant collisions across multiple decoder layers. - Args: hidden_size: Model hidden dimension. num_experts: Total number of experts. rms_norm_eps: Epsilon for the scale-free RMSNorm. """ + # Olive expert quantization leaves the learned router projection floating point. + # Use the shared component loader's exclusion hook rather than rewriting its plan. + component_quantization_excluded_methods: ClassVar[frozenset[str]] = frozenset( + {"olive", "gptq", "awq"} + ) + def __init__(self, hidden_size: int, num_experts: int, rms_norm_eps: float = 1e-6): super().__init__() self._eps = rms_norm_eps self._hidden_size = hidden_size - # ``scale`` stores router.scale * hidden_size^-0.5 (folded in preprocess_weights) self.scale = nn.Parameter([hidden_size]) self.proj = Linear(hidden_size, num_experts, bias=False) self.per_expert_scale = nn.Parameter([num_experts]) def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: - """Compute router probabilities over all experts. + """Compute raw logits; native MoE/QMoE performs softmax internally. Args: hidden_states: [num_tokens, hidden_size] (batch-sequence flattened) Returns: - router_probs: [num_tokens, num_experts] full softmax probabilities + Router logits: [num_tokens, num_experts]. """ # Scale-free RMSNorm using RMSNormalization op (epsilon as attribute avoids # graph-level float constant name collisions across multiple decoder layers). @@ -1489,11 +1637,29 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: hidden_states, ) x_normed = op.RMSNormalization(hidden_states, ones, epsilon=self._eps, axis=-1) - # Scale: x_normed * self.scale (hidden_size^-0.5 already folded into scale) x_scaled = op.Mul(x_normed, self.scale) + # Keep the learned-scale rounding, then use FP32 opmath for HF's Python scalar. + x_scaled = op.CastLike( + op.Mul( + op.Cast(x_scaled, to=ir.DataType.FLOAT), + op.Constant(value_float=float(self._hidden_size**-0.5)), + ), + x_scaled, + ) # Linear projection then softmax -> router_probs [num_tokens, num_experts] - expert_scores = self.proj(op, x_scaled) - return op.Softmax(expert_scores, axis=-1) + return self.proj(op, x_scaled) + + def qmoe_routing( + self, op: OpBuilder, hidden_states: ir.Value, router_input: ir.Value + ) -> tuple[ir.Value, None, bool, float]: + # Router sees the residual; experts see a separately normalized activation. + return self.forward(op, router_input), None, True, 1.0 + + +class _Gemma4QMoELayer(MoELayer): + """Native INT4 experts with exact tanh-GEGLU instead of approximate SwiGLU.""" + + _qmoe_activation_type = "geglu" class Gemma4DecoderLayer(nn.Module): @@ -1591,6 +1757,7 @@ def __init__(self, config: Gemma4Config, layer_idx: int): self.act_fn = get_activation(config.hidden_act) self._enable_moe_block = config.enable_moe_block + self.moe: MoELayer | None = None if config.enable_moe_block: assert config.num_local_experts is not None, "num_local_experts required for MoE" assert config.num_experts_per_tok is not None, ( @@ -1604,18 +1771,17 @@ def __init__(self, config: Gemma4Config, layer_idx: int): self._num_experts = config.num_local_experts self._moe_intermediate_size = config.moe_intermediate_size self._hidden_size = config.hidden_size - # The fused com.microsoft::MoE op's only GATED activation is - # ``swiglu`` = ``x * sigmoid(activation_alpha * x) * up``. Gemma 4 - # experts use ``act_fn(gate) * up`` where ``act_fn`` is the config's - # ``hidden_act`` (``gelu_pytorch_tanh`` for the 26B-A4B / 31B MoE - # checkpoints), NOT silu. Since ``x * sigmoid(1.702 * x)`` closely - # approximates gelu, gelu-family activations map to - # ``activation_alpha=1.702``; a genuine silu expert maps to 1.0. _moe_act = (config.hidden_act or "").lower() - self._moe_swiglu_alpha = 1.702 if "gelu" in _moe_act else 1.0 + self._fused_moe_activation = ( + "geglu" + if _moe_act in _TANH_GELU_ACTIVATIONS + else "swiglu" + if _moe_act in {"silu", "swish"} + else None + ) # The unfused fallback (EPs without a fused MoE op) uses explicit # ONNX ops, so it applies the EXACT expert activation via the shared - # ACT2FN mapping instead of the swiglu approximation. This keeps + # ACT2FN mapping. This keeps # ``gelu`` (erf) distinct from the tanh-approx variants # (``gelu_pytorch_tanh``/``gelu_new``/``gelu_fast``) and ``silu`` -> # Swish, faithfully matching HuggingFace's ``act_fn(gate) * up`` for @@ -1623,19 +1789,32 @@ def __init__(self, config: Gemma4Config, layer_idx: int): self._moe_act_fn = get_activation(config.hidden_act) moe_inter = config.moe_intermediate_size - self.router = _Gemma4MoeRouter( + router = _Gemma4MoeRouter( config.hidden_size, config.num_local_experts, config.rms_norm_eps ) + if _expert_quantization(config, layer_idx) is not None: + self.moe = _Gemma4QMoELayer( + dataclasses.replace( + config, + quantization=config.quantization_for("decoder"), + component_quantization=None, + qmoe_source_paths=_expert_source_paths(layer_idx), + ), + gate=router, + ) + else: + self.router = router # Expert weights stored as 3D tensors (all experts stacked). # fc1_experts_weights: gate+up combined [E, 2*moe_inter, hidden]. # fc2_experts_weights: down projection [E, hidden, moe_inter]. # These map to HF experts.gate_up_proj and experts.down_proj. - self.fc1_experts_weights = nn.Parameter( - [config.num_local_experts, 2 * moe_inter, config.hidden_size] - ) - self.fc2_experts_weights = nn.Parameter( - [config.num_local_experts, config.hidden_size, moe_inter] - ) + if self.moe is None: + self.fc1_experts_weights = nn.Parameter( + [config.num_local_experts, 2 * moe_inter, config.hidden_size] + ) + self.fc2_experts_weights = nn.Parameter( + [config.num_local_experts, config.hidden_size, moe_inter] + ) # MoE-specific norms (in addition to the shared pre/post_feedforward norms) self.pre_feedforward_layernorm_2 = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -1699,7 +1878,7 @@ def forward( # combined: post_ff_norm(dense + moe) + residual dense_out = self.post_feedforward_layernorm_1(op, hidden_states) - # MoE input is the pre-attention residual (flattened to 2D for routing). + # Route the feed-forward residual, flattened to [B*S, H]. batch_size = op.Shape(residual, start=0, end=1) # [1] scalar seq_len = op.Shape(residual, start=1, end=2) # [1] scalar hidden_size = op.Shape(residual, start=2, end=3) # [1] scalar @@ -1707,25 +1886,33 @@ def forward( flat_shape = op.Concat(num_tokens, hidden_size, axis=0) # [2] residual_flat = op.Reshape(residual, flat_shape) # [B*S, H] - # router_probs: [B*S, E] full softmax over all experts - router_probs = self.router(op, residual_flat) # [B*S, E] - # Norm residual before experts normed_flat = self.pre_feedforward_layernorm_2(op, residual_flat) # [B*S, H] caps = ep_capabilities() - if caps.supports_fused_moe: + custom_geglu = caps.name == "cuda" and flags.gemma4_custom_geglu + if self.moe is not None: + if not (caps.supports_fused_moe and custom_geglu): + raise NotImplementedError( + "Quantized Gemma4 MoE requires a GEGLU-capable custom CUDA runtime " + "and MOBIUS_GEMMA4_CUSTOM_GEGLU=1." + ) + moe_out_flat = op.CastLike( + _mark_custom_geglu_runtime(self.moe(op, normed_flat, residual_flat)), + normed_flat, + ) + elif ( + caps.supports_fused_moe + and self._fused_moe_activation is not None + and (self._fused_moe_activation != "geglu" or custom_geglu) + ): # Fused com.microsoft::MoE op handles top-k selection + # expert dispatch internally. Requires ORT main (post # microsoft/onnxruntime#28467, MoE GEMM Refactor) which # plumbs the SwiGLU schema attributes to the kernel. # - # Gemma 4 gated-FFN semantics via the fused SwiGLU schema: - # activation_type=swiglu — x * sigmoid(activation_alpha*x) * up - # activation_alpha=1.702 — approximates gelu_pytorch_tanh, the - # activation Gemma 4 experts actually - # use (act_fn(gate)*up); a silu expert - # would use 1.0. (See _moe_swiglu_alpha.) + # GEGLU is tanh-GELU(gate)*up, not sigmoid(1.702*gate)*gate*up. + # SwiGLU uses alpha=1 for genuine SiLU; GEGLU ignores alpha. # activation_beta=0.0 — linear * gate (no GPT-OSS "+1" bias) # swiglu_limit=inf — no clipping (≤0 disables the clamp) # swiglu_fusion=1 — interleaved layout @@ -1763,28 +1950,33 @@ def forward( ), op.Constant(value_ints=[e_dim, 2 * inter, hidden]), ) # [E, 2*inter, H] interleaved - moe_out_flat = op.CastLike( - op.MoE( # type: ignore[attr-defined] - normed_flat, - router_probs, - fc1_interleaved, - None, # fc1_experts_bias (slot 3, optional) - self.fc2_experts_weights, - activation_type="swiglu", - k=self._top_k, - normalize_routing_weights=1, - activation_alpha=self._moe_swiglu_alpha, - activation_beta=0.0, - swiglu_limit=float("inf"), - swiglu_fusion=1, - _domain="com.microsoft", - ), - normed_flat, # match input dtype (bf16/fp16/fp32) - ) # [B*S, H] + moe_result = op.MoE( # type: ignore[attr-defined] + normed_flat, + self.router(op, residual_flat), + fc1_interleaved, + None, # fc1_experts_bias (slot 3, optional) + self.fc2_experts_weights, + activation_type=self._fused_moe_activation, + k=self._top_k, + normalize_routing_weights=1, + activation_alpha=1.0, + activation_beta=0.0, + swiglu_limit=float("inf"), + swiglu_fusion=1, + _domain="com.microsoft", + ) + if self._fused_moe_activation == "geglu": + moe_result = _mark_custom_geglu_runtime(moe_result) + moe_out_flat = op.CastLike(moe_result, normed_flat) # [B*S, H] else: # EPs without fused MoE support fall back to a static # per-expert unroll. - moe_out_flat = self._dispatch_moe_fallback(op, normed_flat, router_probs) + router_logits = self.router(op, residual_flat) + moe_out_flat = self._dispatch_moe_fallback( + op, + normed_flat, + op.Softmax(op.Cast(router_logits, to=ir.DataType.FLOAT), axis=-1), + ) moe_out = op.Reshape(moe_out_flat, op.Shape(residual)) # [B, S, H] moe_out = self.post_feedforward_layernorm_2(op, moe_out) @@ -1824,10 +2016,8 @@ def _dispatch_moe_fallback( count), unlike a per-expert static unroll which produces O(E x K) nodes (~6k per layer at E=128/K=8 — impractical for large MoE checkpoints). - Standard ``MatMul``/``Einsum`` accumulate in fp32 on CUDA even for fp16 - inputs, so this path is numerically faithful to HuggingFace — unlike the - released-ORT fused ``com.microsoft::MoE`` kernel, which fp16-accumulates - the expert GEMMs and loses accuracy at large hidden size. + This preserves the reference activation and routing boundaries, but + kernel-dependent accumulation still requires numerical qualification. Args: normed_flat: [T, H] — pre-normed input for the experts (T = B*S). @@ -1861,6 +2051,14 @@ def _dispatch_moe_fallback( op.ConstantOfShape(dense_shape, value=ir.tensor(np.zeros(1, dtype=np.float32))), top_weights, ) + selected = op.ScatterElements( + op.Cast(dense_w, to=ir.DataType.BOOL), + top_indices, + op.ConstantOfShape( + op.Shape(top_indices), value=ir.tensor(np.ones(1, dtype=np.bool_)) + ), + axis=1, + ) # [T, E], independent of routing weights or learned expert scales dense_w = op.ScatterElements(dense_w, top_indices, top_weights, axis=1) # [T, E] # fc1_experts_weights [E, 2*moe_inter, H] stores gate rows then up rows. @@ -1874,15 +2072,24 @@ def _dispatch_moe_fallback( up = op.Einsum(normed_flat, up_w, equation="th,eih->tei") # Exact expert activation via ACT2FN: ``gelu_pytorch_tanh`` -> tanh-approx # GELU, bare ``gelu`` -> erf GELU, ``silu``/``swish`` -> Swish, etc. This - # avoids the swiglu approximation used on the fused path. + # preserves the distinction between erf GELU and tanh-GEGLU. activated = self._moe_act_fn(op, gate) inter = op.Mul(activated, up) # [T, E, moe_inter] # Down projection: [T, E, moe_inter] x [E, H, moe_inter] -> [T, E, H] expert_out = op.Einsum(inter, self.fc2_experts_weights, equation="tei,ehi->teh") + # Sparse HF dispatch never evaluates unselected experts. Multiplying + # their nonfinite outputs by zero would contaminate selected results. + expert_out = op.Where( + op.Unsqueeze(selected, [2]), expert_out, op.CastLike(0.0, expert_out) + ) # Weight by the dense routing vector and sum over experts. - weighted = op.Mul(expert_out, op.Unsqueeze(dense_w, [2])) # [T, E, H] * [T, E, 1] + # HF rounds each FP32-weighted contribution back to the expert dtype. + weighted = op.CastLike( + op.Mul(op.Cast(expert_out, to=ir.DataType.FLOAT), op.Unsqueeze(dense_w, [2])), + expert_out, + ) # [T, E, H] * [T, E, 1] return op.ReduceSum(weighted, [1], keepdims=0) # [T, H] @@ -1984,9 +2191,9 @@ def __init__(self, config: Gemma4Config): self.layer_types = layer_types self.sliding_window = config.sliding_window # Bidirectional attention mode (None | "vision"). When "vision", - # contiguous image-token blocks attend bidirectionally; the overlay is - # derived from input_ids at runtime via ``block_sequence_ids`` and forces - # the float-bias attention path (``is_causal=0``). HF also defines an + # contiguous image-token blocks attend bidirectionally in local layers. + # The overlay is derived from input_ids via ``block_sequence_ids`` and + # uses the float-bias attention path (``is_causal=0``) locally. HF defines an # "all" mode (every token bidirectional, no causal mask); it is not used # by any supported Gemma4 checkpoint and not implemented here, so reject # it explicitly rather than silently falling back to causal attention. @@ -2267,12 +2474,10 @@ def forward( from mobius.components._attention import GQAContext, StaticCacheState dtype = get_build_dtype() - # Bidirectional vision-block overlay (Gemma4 larger models). When - # active, contiguous vision-token blocks attend bidirectionally on - # BOTH full and sliding layers. This cannot be expressed by the - # GroupQueryAttention op (causal / local-window only), so we force - # the float-bias Attention path with ``is_causal=0`` and bake the - # full mask (causal + sliding + padding + blockwise OR) into the bias. + # HF applies the image overlay only to local layers: + # valid_key AND (causal OR same-image) AND (key > query - window). + # Global layers stay causal. Retain the conservative all-Attention + # construction here; the normal optimizer can fuse causal globals to GQA. # # ``block_sequence_ids`` [B, S] identifies contiguous image token # blocks. It is derived from ``input_ids`` (image token spans only; @@ -2283,16 +2488,8 @@ def forward( # to the decoder but cannot forward an arbitrary int tensor). When a # caller supplies ``block_sequence_ids`` directly it is used as-is. # - # Trade-off: the overlay is a *static* graph choice — once a model is - # built with ``use_bidirectional_attention == "vision"`` the decoder - # always takes the float-bias (``is_causal=0``) path and forgoes GQA, - # for text-only prompts and every decode step too. This is unavoidable - # in a single static graph: image-token presence is data-dependent at - # runtime, and the same graph serves both image prefill and text decode. - # It stays numerically correct everywhere — when there are no image - # tokens the block ids are all ``-1`` and the ``q_group >= 0`` guard in - # ``create_attention_bias`` makes the overlay a pure no-op (plain causal, - # matching HuggingFace). Only the fused/GQA fast path is given up. + # The overlay is a static graph choice even for text-only decode. + # All -1 block IDs make it a no-op at runtime. bidirectional = self._use_bidirectional_attention == "vision" and self._has_image_token if bidirectional and block_sequence_ids is None and input_ids is not None: block_sequence_ids = _compute_block_sequence_ids( @@ -2311,11 +2508,6 @@ def forward( and caps.supports_fused_rope and not use_block_overlay ) - # When the blockwise overlay is active the Attention op must NOT - # re-apply its built-in causal mask (it would cancel the - # future-position unmasking baked into the float bias). - attn_is_causal = 0 if use_block_overlay else 1 - if use_gqa: # Calling forward() on the RoPE modules materializes their # cos_cache / sin_cache nn.Parameters as ONNX graph initializers. @@ -2447,7 +2639,6 @@ def forward( input_ids=query_input, attention_mask=attention_mask, dtype=self._dtype, - block_sequence_ids=block_sequence_ids if use_block_overlay else None, ), } fallback_pos_dict = position_embeddings_dict @@ -2529,7 +2720,8 @@ def forward( shared_kv_states=shared_kv_states, per_layer_input=per_layer_input, past_key_value=past_kv, - is_causal=attn_is_causal, + # Only local layers may unmask future keys within an image block. + is_causal=0 if use_block_overlay and layer_type == "sliding_attention" else 1, static_kv_seqlen=static_kv_seqlen, ) # KV-shared layers borrow K,V from source layers — exclude from @@ -2684,10 +2876,14 @@ def preprocess_weights( # HF's model.embed_tokens_per_layer.weight [V, L*D] maps directly # to our fused embedding table — no splitting needed. # (For WebGPU, splitting is handled by _Gemma4DecoderModel.preprocess_weights.) - # Map HF expert weight names and fold router scale + config = self.config + if not isinstance(config, Gemma4Config): + raise TypeError("Gemma4 expert preprocessing requires a Gemma4Config.") + native = _extract_native_expert_weights(state_dict, config, prefix="model.") + # Map HF expert weight names and fold expert output scales. _remap_moe_expert_weights(state_dict, self.config) _split_per_layer_projection_weight(state_dict, "model.", self.config) - return super().preprocess_weights(state_dict) + return super().preprocess_weights(state_dict) | native def static_kv_cache_specs(self) -> list[tuple[int, int]]: """Per-cache-layer ``(num_kv_heads, head_dim)`` for static-cache mode.""" @@ -2793,11 +2989,16 @@ class _Gemma4VisionEncoderModel(nn.Module): def __init__(self, config: Gemma4Config): super().__init__() vc = config.vision # VisionConfig for the SigLIP encoder + if vc is None or vc.hidden_size is None: + raise ValueError("Gemma4 vision encoder requires vision.hidden_size.") self._text_hidden_size = config.hidden_size + self._dtype = config.dtype self.encoder = _Gemma4VisionEncoderCore(config) # Gemma4VisionPooler: 3x3 spatial average pooling + sqrt(hidden) scaling. # Reduces N patches to N/9 before projection. self.pooler = Gemma4VisionPooler(vc.hidden_size, vc.pooling_kernel_size or 3) + self.std_bias = nn.Parameter([vc.hidden_size]) if vc.standardize else None + self.std_scale = nn.Parameter([vc.hidden_size]) if vc.standardize else None self.projector_norm = ScaleFreeRMSNorm(vc.hidden_size, eps=vc.norm_eps) linear_class, _ = _vision_linear_classes(config) self.projector = linear_class(vc.hidden_size, config.hidden_size, bias=False) @@ -2815,6 +3016,13 @@ def forward( # Passes pixel_position_ids so the pooler can assign patches to spatial buckets. vision_features = self.pooler(op, vision_features, pixel_position_ids) + if self.std_bias is not None and self.std_scale is not None: + vision_features = op.Mul( + op.Sub(vision_features, op.Cast(self.std_bias, to=ir.DataType.FLOAT)), + op.Cast(self.std_scale, to=ir.DataType.FLOAT), + ) + vision_features = op.Cast(vision_features, to=self._dtype) + # Scale-free norm + linear projection -> [B, N/9, text_hidden] vision_features = self.projector_norm(op, vision_features) vision_features = self.projector(op, vision_features) @@ -2832,7 +3040,9 @@ def preprocess_weights( ) -> dict[str, torch.Tensor]: renamed: dict[str, torch.Tensor] = {} for key, value in state_dict.items(): - if key.startswith("vision_tower."): + if key in ("vision_tower.std_bias", "vision_tower.std_scale"): + renamed[key.removeprefix("vision_tower.")] = value + elif key.startswith("vision_tower."): new_key = "encoder." + key[len("vision_tower.") :] # HF Gemma4VisionModel wraps its layer list in a Gemma4VisionEncoder # submodule, adding an extra "encoder." level. Our _Gemma4VisionEncoderCore @@ -3543,6 +3753,8 @@ def preprocess_weights( elif key.startswith("lm_head."): renamed["decoder." + key] = value + elif key in ("vision_tower.std_bias", "vision_tower.std_scale"): + renamed["vision_encoder." + key.removeprefix("vision_tower.")] = value elif key.startswith("vision_tower."): new_key = "vision_encoder.encoder." + key[len("vision_tower.") :] # HF wraps encoder layers under an extra "encoder." sub-module; strip it @@ -3599,10 +3811,12 @@ def preprocess_weights( else: renamed[key] = value - # Map HF expert weight names and fold router scale + native = _extract_native_expert_weights(renamed, self.config, prefix="decoder.model.") + # Map HF expert weight names and fold expert output scales. _remap_moe_expert_weights(renamed, self.config) renamed = _preprocess_component_quantized_weights(renamed, self.config) + renamed.update(native) # For WebGPU: the fused [V, L*D] embed_tokens_per_layer exceeds the 256 MiB # per-buffer limit. Split it into L separate [V, D] tables in the decoder. diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index aa8abbcbb..809585867 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -58,6 +58,133 @@ def _tiny_gemma4_config(**overrides) -> Gemma4Config: return Gemma4Config(**defaults) +class TestGemma4HfPrecisionAlignment: + @pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] + ) + @pytest.mark.parametrize("quantized", [False, True]) + @pytest.mark.parametrize("optimized", [False, True]) + def test_vision_accepts_float32_processor_pixels(self, dtype, quantized, optimized): + from mobius._builder import _cast_module_dtype, build_from_module + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_config( + dtype=dtype, + enable_moe_block=False, + quantization=( + QuantizationConfig( + quant_method="olive", bits=4, group_size=16, quantize_vision=True + ) + if quantized + else None + ), + ) + module = Gemma4Model(config) + _cast_module_dtype(module, dtype) + package = ( + build_from_module(module, config, task="gemma4") + if optimized + else Gemma4Task().build(module, config) + ) + graph = package["vision_encoder"].graph + assert graph.inputs[0].name == "pixel_values" + assert graph.inputs[0].dtype == ir.DataType.FLOAT + projection = next(node for node in graph if node.op_type in {"MatMul", "MatMulNBits"}) + assert projection.inputs[0].dtype == dtype + + @pytest.mark.parametrize( + ("dtype", "torch_dtype"), + [ + (ir.DataType.FLOAT, torch.float32), + (ir.DataType.FLOAT16, torch.float16), + (ir.DataType.BFLOAT16, torch.bfloat16), + ], + ) + def test_pixel_normalization_before_projection_cast(self, dtype, torch_dtype): + import numpy as np + from onnx_ir import tensor_adapters + + from mobius._builder import _cast_module_dtype + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _Gemma4VisionPatchEmbedder + from mobius.tasks._base import _make_graph + + module = _Gemma4VisionPatchEmbedder(1, 4, 1) + _cast_module_dtype(module, dtype) + graph, builder = _make_graph() + pixels = builder.input("pixel_values", ir.DataType.FLOAT, [1, 256, 3]) + positions = builder.input("positions", ir.DataType.INT64, [1, 256, 2]) + module(builder.op, pixels, positions) + projection = next(node for node in graph if node.op_type == "MatMul") + normalized = projection.inputs[0] + # Execute only normalization: no CPU BF16 MatMul kernel is needed. + normalized_f32 = builder.op.Cast(normalized, to=ir.DataType.FLOAT) + builder.add_output(normalized_f32, "normalized") + module.position_embedding_table.const_value = tensor_adapters.TorchTensor( + torch.zeros(2, 1, 4, dtype=torch_dtype) + ) + prefix = ir.convenience.extract(graph, [pixels], [normalized_f32]) + values = (torch.arange(256).float() / 255).reshape(1, 256, 1).expand(1, 256, 3) + expected = (2 * (values - 0.5)).to(torch_dtype).float().numpy() + actual = OnnxModelSession(ir.Model(prefix, ir_version=10), device="cpu").run( + {"pixel_values": values.numpy().copy()} + )["normalized"] + np.testing.assert_array_equal(actual, expected) + if dtype == ir.DataType.BFLOAT16: + old = (values.to(torch_dtype) * 2 - 1).float().numpy() + assert np.count_nonzero(old != expected) == 128 * 3 + + @pytest.mark.parametrize( + ("dtype", "torch_dtype"), + [ + (ir.DataType.FLOAT, torch.float32), + (ir.DataType.FLOAT16, torch.float16), + (ir.DataType.BFLOAT16, torch.bfloat16), + ], + ) + def test_router_scalar_after_learned_scale(self, dtype, torch_dtype): + import ml_dtypes + import numpy as np + + from mobius._builder import _cast_module_dtype + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _Gemma4MoeRouter + from mobius.tasks._base import _make_graph + + # A non-power-of-two hidden size exposes BF16 scalar-rounding differences. + generator = torch.Generator().manual_seed(7) + normed = torch.randn(4, 2816, generator=generator).to(torch_dtype) + scale = (0.2 + torch.rand(2816, generator=generator)).to(torch_dtype) + learned = normed * scale + expected = (learned * (2816**-0.5)).float().numpy() + module = _Gemma4MoeRouter(2816, 4) + _cast_module_dtype(module, dtype) + graph, builder = _make_graph() + hidden = builder.input("hidden", dtype, [4, 2816]) + module(builder.op, hidden) + projection = next(node for node in graph if node.op_type == "MatMul") + scaled = projection.inputs[0] + assert scaled.producer().op_type == "CastLike" + learned_value = scaled.producer().inputs[1] + assert learned_value.producer().op_type == "Mul" + assert learned_value.producer().inputs[0].producer().op_type == "RMSNormalization" + output = builder.op.Cast(scaled, to=ir.DataType.FLOAT) + builder.add_output(output, "scaled") + prefix = ir.convenience.extract(graph, [learned_value], [output]) + np_dtype = { + ir.DataType.FLOAT: np.float32, + ir.DataType.FLOAT16: np.float16, + ir.DataType.BFLOAT16: ml_dtypes.bfloat16, + }[dtype] + actual = OnnxModelSession(ir.Model(prefix, ir_version=10), device="cpu").run( + {learned_value.name: learned.float().numpy().astype(np_dtype)} + )["scaled"] + np.testing.assert_array_equal(actual, expected) + if dtype == ir.DataType.BFLOAT16: + folded = normed * (scale * (2816**-0.5)) + assert np.count_nonzero(folded.float().numpy() != expected) > 0 + + class TestGemma4CausalLMPreprocessWeights: """Test Gemma4CausalLMModel.preprocess_weights.""" @@ -75,16 +202,16 @@ def test_expert_weight_rename(self): assert "model.layers.0.fc2_experts_weights" in result assert "model.layers.0.experts.gate_up_proj" not in result - def test_router_scale_folding(self): + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_router_scale_preserved(self, dtype): config = _tiny_gemma4_config() model = Gemma4CausalLMModel(config) - scale_val = torch.tensor([2.0]) + scale_val = torch.tensor([2.0], dtype=dtype) fake_sd = {"model.layers.0.router.scale": scale_val.clone()} result = model.preprocess_weights(fake_sd) - expected = 2.0 * (64**-0.5) # hidden_size=64 - assert abs(result["model.layers.0.router.scale"].item() - expected) < 1e-6 + torch.testing.assert_close(result["model.layers.0.router.scale"], scale_val) def test_quantized_moe_experts_fail_closed(self): config = _tiny_gemma4_config( @@ -121,18 +248,18 @@ def test_expert_weight_rename(self): assert "decoder.model.layers.0.fc1_experts_weights" in result assert "decoder.model.layers.0.fc2_experts_weights" in result - def test_router_scale_folding(self): + @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16]) + def test_router_scale_preserved(self, dtype): config = _tiny_gemma4_config() model = Gemma4Model(config) - scale_val = torch.tensor([2.0]) + scale_val = torch.tensor([2.0], dtype=dtype) fake_sd = {"model.language_model.layers.0.router.scale": scale_val.clone()} result = model.preprocess_weights(fake_sd) - expected = 2.0 * (64**-0.5) key = "decoder.model.layers.0.router.scale" assert key in result - assert abs(result[key].item() - expected) < 1e-6 + torch.testing.assert_close(result[key], scale_val) def test_quantized_moe_experts_fail_closed(self): config = _tiny_gemma4_config( @@ -1103,14 +1230,28 @@ def test_package_wires_block_sequence_ids_end_to_end(self): assert "block_sequence_ids" not in dec_inputs assert "input_ids" in dec_inputs - # Decoder attention must drop GQA and disable the op's built-in causal - # mask (is_causal=0) so the baked blockwise bias is honored. + # Only local layers disable the built-in causal mask for the image overlay. dec = pkg["decoder"].graph assert not any(n.op_type == "GroupQueryAttention" for n in dec) attn_nodes = [n for n in dec if n.op_type == "Attention"] - assert attn_nodes - for n in attn_nodes: - assert n.attributes["is_causal"].as_int() == 0 + assert len(attn_nodes) == len(config.layer_types) + for node, layer_type in zip(attn_nodes, config.layer_types, strict=True): + expected = 0 if layer_type == "sliding_attention" else 1 + assert node.attributes["is_causal"].as_int() == expected + # A global mask must have no data dependency on the image block overlay. + global_bias = attn_nodes[1].inputs[3] + pending = [global_bias] + visited = set() + while pending: + value = pending.pop() + if value is None or value.producer() is None: + continue + producer = value.producer() + if producer in visited: + continue + visited.add(producer) + assert producer.op_type != "Or" + pending.extend(producer.inputs) def test_no_block_sequence_ids_when_causal(self): """Without bidirectional attention, no input_ids/overlay is wired.""" @@ -1126,6 +1267,82 @@ def test_no_block_sequence_ids_when_causal(self): assert "block_sequence_ids" not in dec_inputs assert "input_ids" not in dec_inputs + @pytest.mark.parametrize("mask", [[1, 0, 1, 1], [1, 1, 1, 0]]) + @pytest.mark.parametrize("text_task", [False, True], ids=["split", "text"]) + def test_cuda_export_preserves_global_cached_mask_holes(self, mask, text_task): + import numpy as np + + from mobius import build_from_module + from mobius._pipeline_contract import requires_arbitrary_attention_mask + from mobius._testing.ort_inference import OnnxModelSession + from mobius.tasks._gemma4 import Gemma4Task, Gemma4TextCausalLMTask + + config = _tiny_gemma4_config( + dtype=ir.DataType.FLOAT16, + enable_moe_block=False, + use_bidirectional_attention="vision", + image_token_id=255, + ) + module = Gemma4CausalLMModel(config) if text_task else Gemma4Model(config) + task = Gemma4TextCausalLMTask() if text_task else Gemma4Task() + package = build_from_module(module, config, task=task, execution_provider="cuda") + decoder = package["model" if text_task else "decoder"] + assert requires_arbitrary_attention_mask(decoder.graph) + assert not any(node.op_type == "GroupQueryAttention" for node in decoder.graph) + global_attention = next( + node + for node in decoder.graph + if node.op_type == "Attention" + and node.attributes["kv_num_heads"].as_int() == config.num_global_key_value_heads + ) + assert global_attention.attributes["is_causal"].as_int() == 1 + bias = global_attention.inputs[3] + assert bias is not None and bias.name is not None + # Execute the actual optimized bias with one query and three past slots. + # Equal token counts must not collapse these distinct visible-key sets. + inputs = [ + value + for value in decoder.graph.inputs + if value.name in {"input_ids", "attention_mask"} + ] + graph = ir.convenience.extract(decoder.graph, inputs, [bias]) + session = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu") + try: + actual = session.run( + { + "input_ids": np.array([[5]], dtype=np.int64), + "attention_mask": np.array([mask], dtype=np.int64), + } + )[bias.name] + finally: + session.close() + np.testing.assert_array_equal(actual[0, 0, 0] > -1, np.array(mask, dtype=bool)) + + def test_text_only_cuda_export_keeps_gqa_eligibility(self): + from mobius import build_from_module + from mobius._pipeline_contract import requires_arbitrary_attention_mask + + config = _tiny_gemma4_config( + model_type="gemma4_text", + dtype=ir.DataType.FLOAT16, + enable_moe_block=False, + vision=None, + image_token_id=None, + use_bidirectional_attention=None, + ) + package = build_from_module( + Gemma4CausalLMModel(config), + config, + task="gemma4-text-generation", + execution_provider="cuda", + ) + graph = package["model"].graph + assert not requires_arbitrary_attention_mask(graph) + assert ( + sum(node.op_type == "GroupQueryAttention" for node in graph) + == config.num_hidden_layers + ) + def _tiny_gemma4_unified_config(**overrides) -> Gemma4Config: """Minimal gemma4_unified config (dense decoder + encoder-free embedders).""" @@ -1360,6 +1577,413 @@ def _build_gemma4_text(config: Gemma4Config) -> ir.Model: return task.build(module, config)["model"] +class TestGemma4NativeMoe: + @pytest.fixture(autouse=True) + def custom_runtime_opt_in(self): + from mobius._flags import override_flags + + with override_flags(gemma4_custom_geglu=True): + yield + + @staticmethod + def _quantization(**overrides): + return QuantizationConfig( + quant_method="olive", bits=4, group_size=16, quantize_moe=True, **overrides + ) + + @pytest.mark.parametrize("quantized", [False, True]) + @pytest.mark.parametrize( + "dtype", [ir.DataType.FLOAT, ir.DataType.FLOAT16, ir.DataType.BFLOAT16] + ) + def test_native_geglu_uses_raw_logits_and_separate_expert_norm( + self, quantized, dtype, tmp_path + ): + import numpy as np + + from mobius._build_context import build_context + from mobius._builder import _cast_module_dtype + from mobius._constants import GEMMA4_CUSTOM_GEGLU_METADATA_KEY + from mobius._execution_providers import get_ep + from mobius.tasks._gemma4 import Gemma4TextCausalLMTask + + config = _tiny_gemma4_text_moe_config( + dtype=dtype, quantization=self._quantization() if quantized else None + ) + if quantized and dtype == ir.DataType.FLOAT: + with pytest.raises(NotImplementedError, match="FLOAT16 or BFLOAT16"): + Gemma4CausalLMModel(config) + return + module = Gemma4CausalLMModel(config) + _cast_module_dtype(module, dtype) + with build_context(get_ep("cuda"), dtype): + graph = Gemma4TextCausalLMTask().build(module, config)["model"].graph + nodes = [node for node in graph if node.op_type == ("QMoE" if quantized else "MoE")] + assert len(nodes) == 2 + for node in nodes: + assert node.metadata_props[GEMMA4_CUSTOM_GEGLU_METADATA_KEY] == ( + "custom-cuda-geglu-v1" + ) + assert node.attributes["activation_type"].value == "geglu" + assert node.attributes["swiglu_fusion"].value == 1 + router = node.inputs[1] + while router.producer().op_type == "Reshape": + router = router.producer().inputs[0] + assert router.producer().op_type == "MatMul" + scaled = router.producer().inputs[0] + assert scaled.producer().op_type == "CastLike" + learned = scaled.producer().inputs[1] + norm = learned.producer().inputs[0].producer() + assert norm.op_type == "RMSNormalization" + expert_norm = node.inputs[0].producer() + assert expert_norm.op_type == "RMSNormalization" + assert norm.inputs[0] is expert_norm.inputs[0] + assert norm is not expert_norm + if quantized: + assert graph.initializers["model.layers.0.moe.fc1_experts_weights"].dtype == ( + ir.DataType.UINT8 + ) + assert graph.initializers["model.layers.0.moe.gate.proj.weight"].dtype == dtype + # The requirement belongs to the exported operator, not transient flags. + for initializer in graph.initializers.values(): + if initializer.const_value is None: + assert initializer.shape is not None and initializer.dtype is not None + initializer.const_value = ir.tensor( + np.zeros(list(initializer.shape), dtype=initializer.dtype.numpy()) + ) + ir.save(ir.Model(graph, ir_version=10), tmp_path / "native.onnx") + restored = ir.load(tmp_path / "native.onnx") + restored_nodes = [ + node for node in restored.graph if node.op_type == ("QMoE" if quantized else "MoE") + ] + assert len(restored_nodes) == 2 + assert all( + node.metadata_props[GEMMA4_CUSTOM_GEGLU_METADATA_KEY] == "custom-cuda-geglu-v1" + for node in restored_nodes + ) + + def test_default_cuda_is_portable_and_quantized_experts_require_opt_in(self): + from mobius._builder import build_from_module + from mobius._flags import override_flags + + with override_flags(gemma4_custom_geglu=False): + config = _tiny_gemma4_text_moe_config(dtype=ir.DataType.BFLOAT16) + package = build_from_module( + Gemma4CausalLMModel(config), config, execution_provider="cuda" + ) + assert not any(node.op_type in {"MoE", "QMoE"} for node in package["model"].graph) + assert any(node.op_type == "Softmax" for node in package["model"].graph) + with pytest.raises(NotImplementedError, match="MOBIUS_GEMMA4_CUSTOM_GEGLU"): + Gemma4CausalLMModel( + dataclasses.replace(config, quantization=self._quantization()) + ) + + def test_native_experts_fail_closed_on_unsupported_ep_and_layout(self): + config = _tiny_gemma4_text_moe_config( + dtype=ir.DataType.BFLOAT16, quantization=self._quantization() + ) + with pytest.raises(NotImplementedError, match="GEGLU-capable"): + _build_gemma4_text(config) + for hidden_act in ("gelu", "silu"): + with pytest.raises(NotImplementedError, match="tanh-GELU"): + Gemma4CausalLMModel(dataclasses.replace(config, hidden_act=hidden_act)) + quantization = dataclasses.replace( + config.quantization, + overrides={ + "model.language_model.layers.0.experts.down_proj": QuantizationOverride(bits=8) + }, + ) + with pytest.raises(NotImplementedError, match="INT4"): + Gemma4CausalLMModel(dataclasses.replace(config, quantization=quantization)) + + @pytest.mark.parametrize("multimodal", [False, True]) + @pytest.mark.parametrize("symmetric", [False, True]) + @pytest.mark.parametrize("component_plan", [False, True]) + def test_expert_payloads_and_module_layouts(self, multimodal, symmetric, component_plan): + from mobius._component_quantization import configure_component_quantization + from mobius.tasks._gemma4 import Gemma4Task + + quantization = self._quantization( + sym=symmetric, + overrides=( + { + "model.language_model.layers.0.self_attn.q_proj": QuantizationOverride( + bits=8 + ) + } + if component_plan + else {} + ), + ) + config = _tiny_gemma4_config( + dtype=ir.DataType.BFLOAT16, + hidden_act="gelu_pytorch_tanh", + quantization=quantization, + component_quantization={"decoder": quantization} if component_plan else None, + ) + prefix = "model.language_model.layers.0" if multimodal else "model.layers.0" + target = "decoder.model.layers.0.moe" if multimodal else "model.layers.0.moe" + gate = torch.arange(4 * 64 * 32).to(torch.uint8).reshape(4, 64, 32) + down = torch.arange(4 * 64 * 16).to(torch.uint8).reshape(4, 64, 16) + output_scale = torch.tensor([-2.0, 0.0, 0.5, 1.5], dtype=torch.bfloat16) + state = { + f"{prefix}.experts.gate_up_proj_qweight": gate, + f"{prefix}.experts.gate_up_proj_scales": torch.ones( + 4, 64, 4, dtype=torch.bfloat16 + ), + f"{prefix}.experts.down_proj_qweight": down, + f"{prefix}.experts.down_proj_scales": torch.ones(4, 64, 2, dtype=torch.bfloat16), + f"{prefix}.router.scale": torch.ones(64, dtype=torch.bfloat16), + f"{prefix}.router.proj.weight": torch.ones(4, 64, dtype=torch.bfloat16), + f"{prefix}.router.per_expert_scale": output_scale, + } + if not symmetric: + state[f"{prefix}.experts.gate_up_proj_qzeros"] = torch.full( + (4, 64, 2), 0x93, dtype=torch.uint8 + ) + state[f"{prefix}.experts.down_proj_qzeros"] = torch.full( + (4, 64, 1), 0x57, dtype=torch.uint8 + ) + module = (Gemma4Model if multimodal else Gemma4CausalLMModel)(config) + result = module.preprocess_weights(dict(state)) + torch.testing.assert_close(result[f"{target}.fc1_experts_weights"], gate) + torch.testing.assert_close(result[f"{target}.fc2_experts_weights"], down) + torch.testing.assert_close( + result[f"{target}.fc2_scales"], + torch.ones(4, 64, 2, dtype=torch.bfloat16) * output_scale[:, None, None], + ) + torch.testing.assert_close( + result[f"{target}.gate.scale"], torch.ones(64, dtype=torch.bfloat16) + ) + assert (f"{target}.fc1_experts_zero_points" in result) is not symmetric + if not symmetric: + assert torch.all(result[f"{target}.fc1_experts_zero_points"] == 0x93) + if component_plan and multimodal: + from mobius.components import Linear + + task = Gemma4Task() + manifest = configure_component_quantization(module, config, task) + normalized = normalize_component_quantized_weights( + result, + module, + config, + ("decoder", "vision_encoder", "embedding"), + manifest=manifest, + task=task, + ) + assert normalized.keys() == result.keys() + torch.testing.assert_close(normalized[f"{target}.fc1_experts_weights"], gate) + assert module.decoder.model.layers[0].self_attn.q_proj._bits == 8 + assert isinstance(module.decoder.model.layers[0].moe.gate.proj, Linear) + invalid = dict(state) + invalid[f"{prefix}.experts.down_proj_qweight"] = down[..., :-1] + with pytest.raises(ValueError, match="expected shape"): + module.preprocess_weights(invalid) + invalid = dict(state) + invalid.pop(f"{prefix}.experts.down_proj_scales") + with pytest.raises(ValueError, match="expected sidecars"): + module.preprocess_weights(invalid) + + @pytest.mark.parametrize("key", ["moe", "quantize_moe"]) + def test_quantization_metadata_opt_in(self, key): + quantization = QuantizationConfig.from_value( + {"quant_method": "olive", "bits": 4, "group_size": 16, key: True} + ) + assert quantization.quantize_moe + + def test_canonical_quantize_moe_takes_precedence_over_legacy_alias(self): + quantization = QuantizationConfig.from_value( + {"quant_method": "olive", "quantize_moe": False, "moe": True} + ) + assert not quantization.quantize_moe + + +class TestGemma4VisionStandardization: + def test_bf16_kernel_three_pool_matches_actual_hf_rounding(self): + import numpy as np + from transformers import Gemma4VisionConfig as HFVisionConfig + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4VisionPooler as HFVisionPooler, + ) + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import Gemma4VisionPooler + + hidden = torch.ones(1, 9, 4, dtype=torch.bfloat16) + hidden[:, -1] = 1.03125 + positions = torch.tensor([[(x, y) for y in range(3) for x in range(3)]]) + hf_pooler = HFVisionPooler(HFVisionConfig(hidden_size=4)) + expected, _ = hf_pooler( + hidden, positions, torch.zeros(1, 9, dtype=torch.bool), output_length=1 + ) + # Actual pinned HF rounds the FP32 mean to 1, then scales by sqrt(4). + torch.testing.assert_close(expected, torch.full((1, 1, 4), 2.0), rtol=0, atol=0) + rounded_coefficients = torch.full((1, 1, 9), 1 / 9, dtype=torch.bfloat16) + old_mean = (rounded_coefficients.float() @ hidden.float()).to(torch.bfloat16) + # Sensitive to premature coefficient cast; compare exactly in BF16. + torch.testing.assert_close( + old_mean, torch.full_like(old_mean, 1.0078125), rtol=0, atol=0 + ) + + builder, op, graph = create_test_builder() + features = create_test_input(builder, "features", [1, 9, 4], dtype=ir.DataType.FLOAT) + pos = create_test_input(builder, "positions", [1, 9, 2], dtype=ir.DataType.INT64) + output = Gemma4VisionPooler(4, kernel_size=3)( + op, op.Cast(features, to=ir.DataType.BFLOAT16), pos + ) + output.name = "pooled" + graph.outputs.append(output) + matmul = next(node for node in graph if node.op_type == "MatMul") + assert all(value.dtype == ir.DataType.FLOAT for value in matmul.inputs) + actual = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu").run( + {"features": hidden.float().numpy(), "positions": positions.numpy()} + )["pooled"] + np.testing.assert_array_equal(actual, expected.numpy()) + + @pytest.mark.parametrize("nonfinite", [float("nan"), float("inf"), float("-inf")]) + def test_bf16_pool_zeroes_padded_nonfinite_like_actual_hf(self, nonfinite): + import numpy as np + from transformers import Gemma4VisionConfig as HFVisionConfig + from transformers.models.gemma4.modeling_gemma4 import ( + Gemma4VisionPooler as HFVisionPooler, + ) + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import Gemma4VisionPooler + + hidden = torch.ones(1, 9, 4, dtype=torch.bfloat16) + hidden[:, -1] = nonfinite + positions = torch.tensor([[(x, y) for y in range(3) for x in range(3)]]) + positions[:, -1] = -1 + padding = (positions == -1).all(dim=-1) + expected, _ = HFVisionPooler(HFVisionConfig(hidden_size=4))( + hidden, positions, padding, output_length=1 + ) + assert torch.isfinite(expected).all() + rounded_mean = torch.tensor(8 / 9, dtype=torch.bfloat16).float() + torch.testing.assert_close( + expected, (rounded_mean * 2).expand(1, 1, 4), rtol=0, atol=0 + ) + + builder, op, graph = create_test_builder() + features = create_test_input(builder, "features", [1, 9, 4], dtype=ir.DataType.FLOAT) + pos = create_test_input(builder, "positions", [1, 9, 2], dtype=ir.DataType.INT64) + output = Gemma4VisionPooler(4, kernel_size=3)( + op, op.Cast(features, to=ir.DataType.BFLOAT16), pos + ) + output.name = "pooled" + graph.outputs.append(output) + actual = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu").run( + {"features": hidden.float().numpy(), "positions": positions.numpy()} + )["pooled"] + np.testing.assert_array_equal(actual, expected.numpy()) + + @pytest.mark.parametrize("missing", ["vision", "hidden_size"]) + def test_vision_encoder_requires_concrete_hidden_size(self, missing): + from mobius._configs import VisionConfig + from mobius.models.gemma4 import _Gemma4VisionEncoderModel + + with pytest.raises(ValueError, match=r"vision\.hidden_size"): + _Gemma4VisionEncoderModel( + _tiny_gemma4_config( + vision=None if missing == "vision" else VisionConfig(hidden_size=None) + ) + ) + + @pytest.mark.parametrize("standardize", [False, True]) + def test_actual_hf_vision_entrypoint_with_padding(self, standardize): + import numpy as np + from transformers import Gemma4VisionConfig as HFVisionConfig + from transformers.models.gemma4.modeling_gemma4 import Gemma4VisionModel + + from mobius import ModelPackage + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _Gemma4VisionEncoderModel + from mobius.tasks._gemma4 import Gemma4Task + + torch.manual_seed(42) + hf_config = HFVisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + head_dim=16, + patch_size=4, + pooling_kernel_size=2, + position_embedding_size=8, + standardize=standardize, + use_clipped_linears=False, + ) + hf_config._attn_implementation = "eager" + reference = Gemma4VisionModel(hf_config).eval() + if standardize: + reference.std_bias.copy_(torch.linspace(-5.0, 5.0, 32)) + reference.std_scale.copy_(torch.linspace(0.1, 2.0, 32)) + vision = dataclasses.replace( + _tiny_gemma4_config().vision, + pooling_kernel_size=2, + position_embedding_size=8, + standardize=standardize, + ) + config = _tiny_gemma4_config(vision=vision) + module = _Gemma4VisionEncoderModel(config) + package = ModelPackage({"vision_encoder": Gemma4Task()._build_vision(module, config)}) + projection = torch.randn(64, 32) * 0.1 + weights = { + "vision_tower." + name: value for name, value in reference.state_dict().items() + } + weights["embed_vision.embedding_projection.weight"] = projection + package.apply_weights(module.preprocess_weights(weights)) + positions = torch.tensor( + [[*[[x, y] for y in range(2) for x in range(4)], *[[-1, -1]] * 4]] + ) + pixels = torch.rand(1, 12, 48) + with torch.inference_mode(): + features = reference( + pixel_values=pixels, pixel_position_ids=positions + ).last_hidden_state + expected = torch.nn.functional.linear( + torch.nn.functional.rms_norm(features, (32,), eps=vision.norm_eps), projection + ) + actual = OnnxModelSession(package["vision_encoder"], device="cpu").run( + {"pixel_values": pixels.numpy(), "pixel_position_ids": positions.numpy()} + )["image_features"] + assert actual.shape == (2, 64) + np.testing.assert_allclose(actual, expected.numpy(), atol=1e-4, rtol=1e-4) + + @pytest.mark.parametrize("standardize", [False, True]) + def test_config_and_weight_mapping(self, standardize): + from types import SimpleNamespace + + from mobius._configs._vision_defaults import apply_vision_defaults + from mobius.models.gemma4 import _Gemma4VisionEncoderModel + + fields = {} + apply_vision_defaults( + SimpleNamespace(vision_config={"standardize": standardize}), + None, + "gemma4", + fields, + ) + assert fields["standardize"] is standardize + config = _tiny_gemma4_config( + vision=dataclasses.replace(_tiny_gemma4_config().vision, standardize=standardize) + ) + module = _Gemma4VisionEncoderModel(config) + assert (module.std_bias is not None) is standardize + if standardize: + bias, scale = torch.arange(32).float(), torch.ones(32) + state = {"vision_tower.std_bias": bias, "vision_tower.std_scale": scale} + assert module.preprocess_weights(state) == {"std_bias": bias, "std_scale": scale} + result = Gemma4Model(config).preprocess_weights( + {"model." + key: value for key, value in state.items()} + ) + assert result["vision_encoder.std_bias"] is bias + assert result["vision_encoder.std_scale"] is scale + + class TestGemma4MoeFallbackActivation: """The unfused MoE fallback must apply the EXACT ``hidden_act`` (ACT2FN). @@ -1398,6 +2022,76 @@ def test_gelu_pytorch_tanh_uses_tanh_approx(self): assert acts, "MoE fallback must emit an activation op" assert all(op == "Gelu" and approx == "tanh" for op, approx in acts), acts + @pytest.mark.parametrize("selected_expert", [0, 1]) + @pytest.mark.parametrize("expert_scale", [0.0, 1.0]) + def test_nonfinite_expert_is_masked_by_selection_not_weight( + self, selected_expert, expert_scale + ): + import numpy as np + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import Gemma4DecoderLayer + + class FallbackOnly(Gemma4DecoderLayer): + def forward( + self, + op, + hidden_states, + attention_bias, + position_embeddings=None, + shared_kv_states=None, + per_layer_input=None, + past_key_value=None, + is_causal=1, + static_kv_seqlen=None, + ): + self.router(op, hidden_states) + return self._dispatch_moe_fallback(op, hidden_states, attention_bias) + + config = _tiny_gemma4_text_moe_config( + dtype=ir.DataType.FLOAT, + num_local_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=16, + ) + module = FallbackOnly(config, 0) + fc1 = np.zeros((2, 32, 64), dtype=np.float32) + fc1[1] = np.finfo(np.float32).max + fc2 = np.zeros((2, 64, 16), dtype=np.float32) + fc2[1] = 1 + module.fc1_experts_weights.const_value = ir.tensor(fc1) + module.fc2_experts_weights.const_value = ir.tensor(fc2) + module.router.per_expert_scale.const_value = ir.tensor( + np.full(2, expert_scale, dtype=np.float32) + ) + builder, op, graph = create_test_builder() + x = create_test_input(builder, "normed_flat", [1, 64]) + probs = create_test_input(builder, "router_probs", [1, 2]) + output = module(op, x, probs) + output.name = "fallback" + output.shape = ir.Shape([1, 64]) + output.dtype = ir.DataType.FLOAT + graph.outputs.append(output) + graph = ir.convenience.extract(graph, [x, probs], [output]) + session = OnnxModelSession(ir.Model(graph, ir_version=10), device="cpu") + try: + actual = session.run( + { + "normed_flat": np.ones((1, 64), dtype=np.float32), + "router_probs": np.eye(2, dtype=np.float32)[ + selected_expert : selected_expert + 1 + ], + } + )["fallback"] + finally: + session.close() + if selected_expert == 0: + np.testing.assert_array_equal(actual, np.zeros((1, 64), dtype=np.float32)) + else: + # Even a zero learned scale must not hide a selected expert's NaNs. + assert not np.isfinite(actual).any() + class TestGemma4OutputLayerIndices: """``output_layer_indices`` taps the post-final-norm hidden as ``hidden_states.{idx}``. diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index 0beb78c1d..dbd1f7267 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -34,6 +34,7 @@ ) from mobius._model_package import ModelPackage from mobius._pipeline_contract import ( + declare_arbitrary_attention_mask, declare_component_presence, declare_optional_input, ) @@ -272,6 +273,9 @@ class Gemma4TextCausalLMTask(ModelTask): Supports ``static_cache=True`` for pre-allocated TensorScatter-based KV cache (requires ORT ≥ 1.25.0). + Configs retaining vision-overlay attention preserve explicit masks during + optimization; stripped text-only configs remain eligible for GQA fusion. + Inputs (dynamic cache): - input_ids: [batch, sequence_len] INT64 - attention_mask: [batch, past_seq_len + seq_len] INT64 @@ -326,6 +330,12 @@ def build( seq_len = ir.SymbolicDim("sequence_len") graph, builder = _make_graph() + if ( + config.use_bidirectional_attention == "vision" + and config.image_token_id is not None + ): + # Prefix-length GQA fusion cannot preserve internal mask holes. + declare_arbitrary_attention_mask(graph) op = builder.op input_ids = builder.input( @@ -537,6 +547,13 @@ def _build_decoder( seq_len = ir.SymbolicDim("sequence_len") graph, builder = _make_graph(name="decoder") + if ( + config.use_bidirectional_attention == "vision" + and config.image_token_id is not None + ): + # A causal global bias still carries arbitrary padding positions. + # Prefix-length GQA fusion cannot preserve internal mask holes. + declare_arbitrary_attention_mask(graph) op = builder.op inputs_embeds = builder.input( @@ -643,7 +660,7 @@ def _build_vision( """Build vision encoder: pixel_values + pixel_position_ids -> image_features. Inputs: - - ``pixel_values [1, N, 3*P^2]``: one pre-patchified image with ``N`` + - ``pixel_values [1, N, 3*P^2]``: FP32 patches in ``[0, 1]``, with ``N`` patches (padded to ``max_soft_tokens``, typically 280). - ``pixel_position_ids [1, N, 2]``: (x, y) patch coordinates. Unused patch slots (from images smaller than the maximum) use @@ -666,7 +683,7 @@ def _build_vision( pixel_values = builder.input( "pixel_values", - dtype=config.dtype, + dtype=ir.DataType.FLOAT, shape=[1, num_patches, pixel_dim], ) pixel_position_ids = builder.input( diff --git a/tests/build_graph/core_test.py b/tests/build_graph/core_test.py index a1fa025b3..a7ea1b5f2 100644 --- a/tests/build_graph/core_test.py +++ b/tests/build_graph/core_test.py @@ -797,8 +797,8 @@ def test_multimodal_encoder_inputs_match_model_dtype(self, dtype_str): "dtype_str", ["f16", "bf16"], ) - def test_gemma4_encoder_inputs_match_model_dtype(self, dtype_str): - """Gemma4 vision/audio encoder inputs match the model's compute dtype.""" + def test_gemma4_encoder_input_dtypes(self, dtype_str): + """Vision takes FP32 processor pixels; audio keeps the compute dtype.""" from mobius._configs import Gemma4AudioConfig, Gemma4Config config = Gemma4Config( @@ -846,11 +846,11 @@ def test_gemma4_encoder_inputs_match_model_dtype(self, dtype_str): task = get_task("gemma4") pkg = task.build(module, config) - # Vision encoder pixel_values must match config.dtype + # Normalize processor pixels in FP32 before casting to the projection dtype. vision_model = pkg["vision_encoder"] pv_input = vision_model.graph.inputs[0] assert pv_input.name == "pixel_values" - assert pv_input.dtype == config.dtype + assert pv_input.dtype == ir.DataType.FLOAT # Audio encoder input_features must match config.dtype audio_model = pkg["audio_encoder"] diff --git a/tests/integration/gemma4_test.py b/tests/integration/gemma4_test.py index f97c3be61..107cd7caf 100644 --- a/tests/integration/gemma4_test.py +++ b/tests/integration/gemma4_test.py @@ -558,47 +558,68 @@ def test_gemma4_unified_12b_multimodal_prefill(): @pytest.mark.integration -def test_gemma4_bidirectional_mask_parity(): +@pytest.mark.parametrize( + ("static", "case"), + [ + (static, case) + for static in (False, True) + for case in ("text", "image", "two_images", "long_image", "padding") + ] + + [(False, "internal_hole")], +) +def test_gemma4_bidirectional_mask_parity(static, case): """Mobius's vision-block attention bias matches HuggingFace exactly. - Larger Gemma 4 models use ``use_bidirectional_attention="vision"``: a - contiguous run of image/audio placeholder tokens attends bidirectionally - within its block (on BOTH full and sliding layers), while text stays - causal. This test compares mobius's ``create_attention_bias`` output - (with ``block_sequence_ids`` from ``_compute_block_sequence_ids``) against - HuggingFace's real ``create_causal_mask`` / ``create_sliding_window_causal_mask`` - for the actual ``gemma-4-26b-a4b-it`` config. It needs only the config - (no weights), so it is cheap and deterministic. + Global layers remain causal; local layers use + ``(causal OR same-image) AND left-window AND padding``. Compare the actual + HF Gemma4 entry point with a local config, without downloads or weights. """ import onnx_ir as ir import torch - from transformers.masking_utils import ( - create_causal_mask, - create_sliding_window_causal_mask, + from transformers import Gemma4TextConfig + from transformers.models.gemma4.modeling_gemma4 import ( + create_masks_for_vision_model, + get_block_sequence_ids_for_mask, ) from mobius._testing import create_test_builder, create_test_input - from mobius.components._common import create_attention_bias + from mobius.components import create_attention_bias, create_static_cache_attention_bias from mobius.models.gemma4 import _compute_block_sequence_ids from mobius.tasks._base import _make_graph, _make_model - model_id = "google/gemma-4-26b-a4b-it" - if not _model_accessible(model_id): - pytest.skip(f"{model_id} not accessible (requires HuggingFace authentication)") - - hf_config = transformers.AutoConfig.from_pretrained(model_id) - text_cfg = hf_config.text_config + text_cfg = Gemma4TextConfig( + hidden_size=64, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + sliding_window=1024, + layer_types=["sliding_attention", "full_attention"], + use_bidirectional_attention="vision", + ) text_cfg._attn_implementation = "eager" # always returns a dense float mask assert text_cfg.use_bidirectional_attention == "vision" - image_token_id = hf_config.image_token_id - - # Synthetic layout: text, image block, text, image block, text. - seq_len = 12 - input_ids = np.full((1, seq_len), 5, dtype=np.int64) - input_ids[0, 2:7] = image_token_id - input_ids[0, 9:11] = image_token_id + image_token_id = 255 + types = { + "text": [0] * 7, + "image": [0, 1, 1, 1, 0], + "two_images": [0, 1, 1, 0, 1, 1, 0], + "long_image": [0] + [1] * 1120 + [0], + "padding": [0, 1, 1, 1, 0, 0], + "internal_hole": [1] * 1025, + }[case] + token_types = torch.tensor([types]) + seq_len = len(types) + input_ids = np.where(token_types.numpy() == 1, image_token_id, 5).astype(np.int64) attention_mask = np.ones((1, seq_len), dtype=np.int64) position_ids = np.arange(seq_len, dtype=np.int64)[None, :] + if case == "padding": + attention_mask[0, -2:] = 0 + elif case == "internal_hole": + # Physical q=1024,k=0 is outside W=1024 despite an invalid slot at k=1. + # The static primitive has prefix validity, not arbitrary hole validity; + # retain its ten existing comparisons and test this case dynamically. + attention_mask[0, 1] = 0 # Mobius block_sequence_ids from input_ids. graph, builder = _make_graph() @@ -609,10 +630,9 @@ def test_gemma4_bidirectional_mask_parity(): block_ids = OnnxModelSession(_make_model(graph), device="cpu").run( {"input_ids": input_ids} )["bsid"] - # Two contiguous image runs separated by text -> groups 0 and 1. np.testing.assert_array_equal( - block_ids[0], - np.array([-1, -1, 0, 0, 0, 0, 0, -1, -1, 1, 1, -1], dtype=np.int64), + block_ids, + get_block_sequence_ids_for_mask(token_types, torch.device("cpu")).numpy(), ) def mobius_attended(sliding_window): @@ -620,14 +640,22 @@ def mobius_attended(sliding_window): ii = create_test_input(b, "input_ids", [1, seq_len], dtype=ir.DataType.INT64) am = create_test_input(b, "attention_mask", [1, seq_len], dtype=ir.DataType.INT64) bk = create_test_input(b, "block_sequence_ids", [1, seq_len], dtype=ir.DataType.INT64) - bias = create_attention_bias( - bop, - ii, - am, + arguments = dict( sliding_window=sliding_window, dtype=ir.DataType.FLOAT, - block_sequence_ids=bk, + block_sequence_ids=bk if sliding_window is not None else None, ) + if static: + bias = create_static_cache_attention_bias( + bop, + write_indices=bop.Constant(value_ints=[0]), + seq_len=bop.Shape(ii, start=1, end=2), + nonpad_kv_seqlen=bop.ReduceSum(am, [1], keepdims=0), + max_seq_len=seq_len, + **arguments, + ) + else: + bias = create_attention_bias(bop, ii, am, **arguments) bias.name = "bias" g.outputs.append(bias) out = OnnxModelSession(ir.Model(g, ir_version=10), device="cpu").run( @@ -641,30 +669,26 @@ def mobius_attended(sliding_window): inputs_embeds = torch.zeros(1, seq_len, 8) blk = torch.from_numpy(block_ids) - hf_full = create_causal_mask( - text_cfg, - inputs_embeds, - torch.from_numpy(attention_mask), - None, - torch.from_numpy(position_ids), - block_sequence_ids=blk, - ) - hf_sliding = create_sliding_window_causal_mask( + hf_masks = create_masks_for_vision_model( text_cfg, inputs_embeds, torch.from_numpy(attention_mask), None, torch.from_numpy(position_ids), - block_sequence_ids=blk, + blk, ) def hf_attended(mask): m = mask[0, 0].numpy() return m if m.dtype == bool else (m > -1e30) - # Full-attention layers: causal OR same-block, AND padding. - np.testing.assert_array_equal(mobius_attended(None), hf_attended(hf_full)) - # Sliding layers: (causal AND window) OR same-block, AND padding. np.testing.assert_array_equal( - mobius_attended(text_cfg.sliding_window), hf_attended(hf_sliding) + mobius_attended(None), hf_attended(hf_masks["full_attention"]) ) + actual_local = mobius_attended(text_cfg.sliding_window) + hf_local = hf_attended(hf_masks["sliding_attention"]) + np.testing.assert_array_equal(actual_local, hf_local) + if case == "internal_hole": + assert not hf_local[1024, 0] + assert not actual_local[1024, 0] + assert actual_local[1024, 2] diff --git a/tests/static_cache_bias_parity_test.py b/tests/static_cache_bias_parity_test.py index ffef477db..861fc52e2 100644 --- a/tests/static_cache_bias_parity_test.py +++ b/tests/static_cache_bias_parity_test.py @@ -187,7 +187,7 @@ def _dense_reference( ``full_*_cache`` are the *post-scatter* caches (current chunk already written). For each (batch, query, head) we build the boolean mask in the exact rule order ``create_static_cache_attention_bias`` uses — causal, - AND sliding window, OR block overlay, AND padding validity — gather the + OR block overlay, AND sliding window, AND padding validity — gather the valid cache slots (with GQA head sharing) and softmax over them. Every query row keeps its own diagonal slot valid (``q_abs[t] == write+t`` is causal, within any window, and ``< nonpad``), so a fully-masked row is @@ -218,10 +218,10 @@ def _dense_reference( for t in range(query_len): mask = q_abs[t] >= kv_slots # causal - if sliding_window is not None: - mask &= (q_abs[t] - kv_slots) < sliding_window # local window if block_sequence_ids is not None and block_sequence_ids[b, t] >= 0: mask |= kv_group == block_sequence_ids[b, t] # bidirectional block + if sliding_window is not None: + mask &= (q_abs[t] - kv_slots) < sliding_window # local window mask &= kv_slots < npad # padding validity idx = np.nonzero(mask)[0] # The bias-MEA contract keeps every query row's own diagonal slot @@ -371,7 +371,7 @@ def test_prefill_chunk_bias_matches_dense_reference(): """Prefill chunk (write=0, S_q=N, nonpad=N) matches the dense reference. The whole chunk is the valid region; query row ``i`` attends slots within - ``[max(0, i - sliding_window + 1), i]``, OR-ed with same-block tokens. + ``causal OR same-block``, constrained by the strict left window boundary. """ query_len = 6 onnx_attn, ref = _run_and_compare( From 3dc6ecddf637071671d47c85bff27488da2e322c Mon Sep 17 00:00:00 2001 From: titaiwang Date: Wed, 28 Oct 2026 21:18:31 +0000 Subject: [PATCH 6/7] Separate generic weight-layout folding from projection packing Move generic static Reshape/Transpose folding and integer shape folding to a separate draft, while retaining legacy matrix Transpose+Concat packing, dead-node cleanup, lazy source ownership, adapter behavior, and the new no-fold regression. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/execution_providers.md | 10 +- src/mobius/_optimizations.py | 34 +- src/mobius/_passes/__init__.py | 5 - src/mobius/_passes/_fold_transpose.py | 4 +- src/mobius/_passes/_fold_transpose_test.py | 55 ++ src/mobius/_passes/_fold_weight_layout.py | 268 ------ .../_passes/_fold_weight_layout_test.py | 872 ------------------ 7 files changed, 67 insertions(+), 1181 deletions(-) delete mode 100644 src/mobius/_passes/_fold_weight_layout.py delete mode 100644 src/mobius/_passes/_fold_weight_layout_test.py diff --git a/docs/execution_providers.md b/docs/execution_providers.md index 8b0701748..dcc35e39d 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -142,11 +142,11 @@ 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. -Post-weight folding is independent of experimental packing: parameter Concats -retain checkpoint names until binding, then Concat/Reshape/Transpose chains fold -lazily without changing dtype. Missing leaves are checked even with -`fold_constants=False`. Overridable inputs and MoE/QMoE/QDQ storage-layout paths -remain explicit to preserve loader names and byte layouts. +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 diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index e597fb217..f46efdbf3 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -21,8 +21,8 @@ Post-weight passes ------------------ :func:`fold_initializers_after_weights` should be called after weights are loaded. -It lazily folds initializer Transpose, Concat, and static Reshape/Transpose -chains, independently of experimental projection packing. +It lazily folds initializer matrix Transpose and Concat nodes for projection +packing. Generic Reshape/Transpose layout chains remain explicit. """ from __future__ import annotations @@ -47,7 +47,6 @@ import contextlib import dataclasses import logging -import math import warnings import onnx @@ -63,7 +62,6 @@ from mobius._passes import ( FoldConcatInitializersPass, FoldTransposedInitializerPass, - FoldWeightLayoutPass, Fp8KvCachePass, RemoveDeadGraphInputsPass, ) @@ -797,14 +795,12 @@ def _should_inline(func: ir.Function) -> bool: def fold_initializers_after_weights(model: ir.Model) -> None: - """Lazily fold weight Transpose, Concat, and static layout chains. + """Lazily fold initializer matrix Transpose and Concat nodes. - Fold small integer shape expressions first without exposing weight payloads - to the general constant folder. Preserve the shipped root Transpose naming - convention, then fold layouts and packing to a fixed point. + Preserve the shipped root Transpose naming convention and fold projection + packing to a fixed point. Generic weight layout chains remain explicit. """ common_passes.LiftConstantsToInitializersPass()(model) - hidden: dict[ir.Value, ir.TensorProtocol] = {} for graph in model.graphs(): for value in graph.initializers.values(): tensor = value.const_value @@ -818,29 +814,9 @@ def fold_initializers_after_weights(model: ir.Model) -> None: model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) ) - if tensor is not None and not ( - tensor.dtype == ir.DataType.INT64 - and all(isinstance(dim, int) for dim in tensor.shape) - and math.prod(tensor.shape.numpy()) <= 8192 - ): - hidden[value] = tensor - value.const_value = None - try: - # Partial evaluators also simplify Identity without reading its input. - # In particular, an output Identity transfers its public output name - # to the producer value before lazy layout folding. Size limits and - # should_fold do not gate those upstream symbolic simplifications. - onnxscript.optimizer._constant_folding.FoldConstantsPass( - shape_inference=False, input_size_limit=8192, output_size_limit=8192 - )(model) - finally: - for value, tensor in hidden.items(): - value.const_value = tensor - common_passes.LiftConstantsToInitializersPass()(model) common_passes.RemoveUnusedNodesPass()(model) lazy_passes = ( FoldTransposedInitializerPass(), - FoldWeightLayoutPass(), FoldConcatInitializersPass(), ) while True: diff --git a/src/mobius/_passes/__init__.py b/src/mobius/_passes/__init__.py index 499473ec4..36cfa31f7 100644 --- a/src/mobius/_passes/__init__.py +++ b/src/mobius/_passes/__init__.py @@ -36,9 +36,6 @@ graph initializer) into a single pre-concatenated initializer. Used to pre-pack QKV weight matrices after weight loading. -:class:`FoldWeightLayoutPass` - Folds static Reshape/Transpose chains into lazy initializers, including - higher-rank layout conversions of quantized weights, scales, and zero points. """ from __future__ import annotations @@ -46,13 +43,11 @@ __all__ = [ "FoldTransposedInitializerPass", "FoldConcatInitializersPass", - "FoldWeightLayoutPass", "Fp8KvCachePass", "RemoveDeadGraphInputsPass", ] from mobius._passes._fold_concat import FoldConcatInitializersPass from mobius._passes._fold_transpose import FoldTransposedInitializerPass -from mobius._passes._fold_weight_layout import FoldWeightLayoutPass from mobius._passes._fp8_kv_cache import Fp8KvCachePass from mobius._passes._remove_dead_inputs import RemoveDeadGraphInputsPass diff --git a/src/mobius/_passes/_fold_transpose.py b/src/mobius/_passes/_fold_transpose.py index 52fb232dd..e14ddc033 100644 --- a/src/mobius/_passes/_fold_transpose.py +++ b/src/mobius/_passes/_fold_transpose.py @@ -74,8 +74,8 @@ def call(self, model: ir.Model) -> ir.passes.PassResult: # 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. Child-local weights belong to FoldWeightLayoutPass: - # registering them at the root would collide across sibling scopes. + # 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() 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/_passes/_fold_weight_layout.py b/src/mobius/_passes/_fold_weight_layout.py deleted file mode 100644 index a208bd1a3..000000000 --- a/src/mobius/_passes/_fold_weight_layout.py +++ /dev/null @@ -1,268 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Lazily fold static weight Reshape/Transpose chains after checkpoint loading.""" - -from __future__ import annotations - -import logging -import math - -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 _reshape_shape(node: ir.Node, source_shape: tuple[int, ...]) -> tuple[int, ...] | None: - """Resolve a static shape, skip dynamic shapes, and reject invalid reshapes.""" - shape_value = node.inputs[1] - if shape_value is None or shape_value.is_graph_input(): - return None - if shape_value.is_initializer(): - shape_tensor = shape_value.const_value - else: - constant = shape_value.producer() - if ( - constant is None - or constant.domain != "" - or constant.op_type != "Constant" - or len(constant.attributes) != 1 - ): - return None - if "value" in constant.attributes: - shape_tensor = constant.attributes["value"].as_tensor() - elif "value_ints" in constant.attributes: - shape_tensor = ir.tensor( - constant.attributes["value_ints"].as_ints(), dtype=ir.DataType.INT64 - ) - else: - return None - if shape_tensor is None: - return None - if shape_tensor.dtype != ir.DataType.INT64 or len(shape_tensor.shape) != 1: - raise ValueError(f"Reshape {node.name!r} requires a one-dimensional INT64 shape") - requested = tuple(int(dim) for dim in shape_tensor.numpy()) - allowzero = node.attributes.get_int("allowzero", 0) - if allowzero not in (0, 1) or requested.count(-1) > 1 or any(d < -1 for d in requested): - raise ValueError(f"Invalid Reshape attributes or shape at {node.name!r}: {requested}") - resolved = list(requested) - if not allowzero: - for index, dim in enumerate(resolved): - if dim == 0: - if index >= len(source_shape): - raise ValueError(f"Reshape {node.name!r} copies a nonexistent dimension") - resolved[index] = source_shape[index] - size = math.prod(source_shape) - known_size = math.prod(dim for dim in resolved if dim != -1) - if -1 in resolved: - if known_size == 0 or size % known_size: - raise ValueError(f"Reshape {node.name!r} cannot infer dimension in {requested}") - resolved[resolved.index(-1)] = size // known_size - elif known_size != size: - raise ValueError( - f"Reshape {node.name!r} changes element count: {source_shape} to {requested}" - ) - return tuple(resolved) - - -def _storage_layout_values(model: ir.Model) -> set[ir.Value]: - """Find parameter views whose checkpoint storage ABI must remain explicit.""" - pending: list[ir.Value] = [] - for node in model.graph.all_nodes(): - if node.domain == "" and node.op_type == "DequantizeLinear": - parameters = node.inputs[:3] - elif node.domain == "com.microsoft" and node.op_type in ("MoE", "QMoE"): - # Skip activations/routing. Expert weights, biases, scales, and - # zero points retain their checkpoint-format names and row order; - # the explicit layout views adapt that storage to the runtime ABI. - parameters = node.inputs[2 : 14 if node.op_type == "QMoE" else 8] - else: - continue - pending.extend(value for value in parameters if value is not None) - protected: set[ir.Value] = set() - while pending: - value = pending.pop() - if value in protected: - continue - protected.add(value) - producer = value.producer() - if producer is None or producer.domain != "" or producer.overload: - continue - if producer.op_type in ( - "Identity", - "Reshape", - "Transpose", - "Cast", - "CastLike", - "Pad", - "Slice", - "Expand", - ): - pending.extend(item for item in producer.inputs[:1] if item is not None) - elif producer.op_type in ("Add", "Sub", "Mul", "Div", "Concat"): - # FP8 scales can carry an exact-zero execution dependency before - # DQ. Follow the scale path without evaluating its arithmetic. - pending.extend(item for item in producer.inputs if item is not None) - return protected - - -def _initializer_owners(value: ir.Value, graphs: tuple[ir.Graph, ...]) -> list[ir.Graph]: - """Find registries containing this exact value, not just its name. - - An If branch can capture a root initializer without owning it, and sibling - branches can own different values with the same name. Registry identity - determines where it is safe to remove a folded source. - """ - return [ - graph - for graph in graphs - if value.name is not None and graph.initializers.get(value.name) is value - ] - - -class FoldWeightLayoutPass(ir.passes.InPlacePass): - """Fold named static Reshape/Transpose outputs into lazy initializers. - - Graphs must be topologically ordered. A single traversal folds layout chains - without reading weight payloads: only constant Reshape shape vectors are read. - Materialization preserves the source dtype and uses no arithmetic or casts. - Overridable initializers, shared sources, and dynamic shapes remain untouched. - MoE/QMoE and QDQ parameter-layout paths retain checkpoint storage names and - byte layouts; their explicit runtime format conversions are not folded. - Dead folded inputs are removed in their owning graph; run dead-code elimination - afterward to remove unused Constant shape nodes. - Malformed static layout nodes raise ValueError rather than being skipped. - """ - - def call(self, model: ir.Model) -> ir.passes.PassResult: - folded = 0 - graphs = tuple(model.graphs()) - storage_layouts = _storage_layout_values(model) - - for graph in graphs: - for node in list(graph): - if ( - node.overload - or node.domain != "" - or node.op_type - not in ( - "Reshape", - "Transpose", - ) - ): - continue - expected_inputs = 2 if node.op_type == "Reshape" else 1 - if len(node.inputs) != expected_inputs or len(node.outputs) != 1: - raise ValueError(f"Malformed weight layout node {node.name!r}") - if node.outputs[0] in storage_layouts: - continue - source = node.inputs[0] - if source is None or not source.is_initializer() or source.is_graph_input(): - continue - # Keep tied tables shared instead of serializing a second full layout. - if ( - len(source.uses()) != 1 - or len(_initializer_owners(source, graphs)) != 1 - or any( - source in scope.inputs or source in scope.outputs for scope in graphs - ) - ): - continue - tensor = source.const_value - if tensor is None: - logger.warning( - "FoldWeightLayoutPass: skipping %r; initializer %r has no data", - node.name, - source.name, - ) - continue - output = node.outputs[0] - if not output.name or any(not isinstance(dim, int) for dim in tensor.shape): - logger.debug( - "FoldWeightLayoutPass: skipping unnamed or symbolic %r", node.name - ) - continue - source_shape = tuple(tensor.shape.numpy()) - dtype = initializer_dtype(source) - if dtype is None: - raise ValueError(f"Weight layout source {source.name!r} has no dtype") - is_transpose = node.op_type == "Transpose" - if is_transpose: - layout = tuple( - node.attributes.get_ints( - "perm", tuple(reversed(range(len(source_shape)))) - ) - ) - if sorted(layout) != list(range(len(source_shape))): - raise ValueError( - f"Invalid Transpose permutation at {node.name!r}: {layout}" - ) - shape = tuple(source_shape[index] for index in layout) - else: - resolved = _reshape_shape(node, source_shape) - if resolved is None: - logger.debug( - "FoldWeightLayoutPass: skipping dynamic shape at %r", node.name - ) - continue - layout = shape = resolved - - # Bind this iteration's values before the lazy callback runs later. - def materialize( - data: ir.TensorProtocol = tensor, - dimensions: tuple[int, ...] = layout, - transpose: bool = is_transpose, - data_type: ir.DataType = dtype, - ) -> ir.TensorProtocol: - array = data.numpy() - result = ( - array.transpose(dimensions) if transpose else array.reshape(dimensions) - ) - return ir.tensor(result, dtype=data_type) - - output_shape = ir.Shape(shape) - replacement = ir.Value( - name=output.name, - shape=output_shape, - type=ir.TensorType(dtype), - doc_string=output.doc_string, - metadata_props=dict(output.metadata_props), - ) - replacement.const_value = ir.LazyTensor( - materialize, dtype=dtype, shape=output_shape, name=output.name - ) - source_files = ( - (tensor.path,) - if isinstance(tensor, ir.ExternalTensor) - else tensor.meta.get(LAZY_SOURCE_FILES_META_KEY, ()) - ) - if source_files: - replacement.const_value.meta[LAZY_SOURCE_FILES_META_KEY] = tuple( - source_files - ) - model.meta[LAZY_SOURCE_FILES_META_KEY] = frozenset( - (*model.meta.get(LAZY_SOURCE_FILES_META_KEY, ()), *source_files) - ) - output.replace_all_uses_with(replacement, replace_graph_outputs=True) - inputs = tuple(node.inputs) - graph.remove(node, safe=True) - # Downstream layout nodes can now capture this lazy initializer. - graph.register_initializer(replacement) - for value in inputs: - if ( - value is not None - and value.name is not None - and not value.uses() - and not any( - value in scope.inputs or value in scope.outputs for scope in graphs - ) - ): - registries = _initializer_owners(value, graphs) - if len(registries) == 1: - del registries[0].initializers[value.name] - folded += 1 - logger.debug("FoldWeightLayoutPass: folded %d weight layout nodes", folded) - return ir.passes.PassResult(model, modified=folded > 0) diff --git a/src/mobius/_passes/_fold_weight_layout_test.py b/src/mobius/_passes/_fold_weight_layout_test.py deleted file mode 100644 index 8bae303e0..000000000 --- a/src/mobius/_passes/_fold_weight_layout_test.py +++ /dev/null @@ -1,872 +0,0 @@ -# Copyright (c) Microsoft Corporation. -# Licensed under the MIT License. - -"""Value, laziness, and persistence coverage for static weight layout folding.""" - -from __future__ import annotations - -from pathlib import Path - -import ml_dtypes -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._optimizations import fold_initializers_after_weights -from mobius._passes import FoldWeightLayoutPass -from mobius._testing.ort_inference import OnnxModelSession - - -def _initializer(name, data): - tensor = ir.tensor(data, name=name) - return ir.Value( - name=name, shape=tensor.shape, type=ir.TensorType(tensor.dtype), const_value=tensor - ) - - -def _model(data, *, first=(2, 2, 4, 3), perm=(0, 2, 1, 3), last=(2, 8, 3)): - weight = _initializer("weight", data) - shape1 = _initializer("shape1", np.array(first, dtype=np.int64)) - shape2 = _initializer("shape2", np.array(last, dtype=np.int64)) - reshape1 = ir.Node("", "Reshape", [weight, shape1], num_outputs=1) - transpose = ir.Node( - "", - "Transpose", - [reshape1.outputs[0]], - attributes=[ir.AttrInt64s("perm", perm)], - num_outputs=1, - ) - reshape2 = ir.Node("", "Reshape", [transpose.outputs[0], shape2], num_outputs=1) - identity = ir.Node("", "Identity", [reshape2.outputs[0]], num_outputs=1) - for node, name in zip( - (reshape1, transpose, reshape2, identity), - ("reshaped", "transposed", "layout", "output"), - strict=True, - ): - node.outputs[0].name = name - node.outputs[0].dtype = weight.dtype - graph = ir.Graph( - [], - [identity.outputs[0]], - nodes=[reshape1, transpose, reshape2, identity], - initializers=[weight, shape1, shape2], - opset_imports={"": 24}, - name="layout_test", - ) - return ir.Model(graph, ir_version=10) - - -@pytest.mark.parametrize("dtype", [np.uint8, np.float16, np.float32, ml_dtypes.bfloat16]) -def test_quantized_layout_values_and_package_roundtrip(dtype, tmp_path): - data = np.arange(48).astype(dtype).reshape(2, 8, 3) - model = _model(data) - fold_initializers_after_weights(model) - # The shape-expression constant folder also resolves Identity symbolically, - # preserving the public output name on the eventual lazy initializer. - assert not list(model.graph) - assert set(model.graph.initializers) == {"output"} - folded = model.graph.initializers["output"] - assert model.graph.outputs[0] is folded - assert isinstance(folded.const_value, ir.LazyTensor) - expected = data.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3) - assert folded.dtype == ir.tensor(data).dtype - assert folded.const_value.numpy().tobytes() == expected.tobytes() - package = ModelPackage({"model": model}) - package.save(tmp_path, progress_bar=False) - loaded = ModelPackage.load(tmp_path) - restored = loaded["model"].graph.initializers["output"] - assert loaded["model"].graph.outputs[0] is restored - assert restored.dtype == folded.dtype - assert restored.shape == ir.Shape(expected.shape) - assert restored.const_value.numpy().tobytes() == expected.tobytes() - - -def test_bfloat16_special_value_bits_preserved(tmp_path): - bits = np.resize( - np.array([0, 0x8000, 0x7F80, 0xFF80, 0x7FC1, 0x7FFF, 1, 0x8001], dtype=np.uint16), - 48, - ).reshape(2, 8, 3) - model = _model(bits.view(ml_dtypes.bfloat16)) - fold_initializers_after_weights(model) - folded = model.graph.initializers["output"] - assert set(model.graph.initializers) == {"output"} - assert model.graph.outputs[0] is folded - assert isinstance(folded.const_value, ir.LazyTensor) - actual = folded.const_value.numpy() - expected = bits.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3) - assert actual.tobytes() == expected.tobytes() - ModelPackage({"model": model}).save(tmp_path, progress_bar=False) - restored_graph = ModelPackage.load(tmp_path)["model"].graph - restored = restored_graph.initializers["output"] - assert restored_graph.outputs[0] is restored - assert restored.dtype == ir.DataType.BFLOAT16 - assert restored.const_value.numpy().tobytes() == expected.tobytes() - - -def test_weight_payload_remains_lazy_and_pass_is_idempotent(): - data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) - model = _model(data) - source = model.graph.initializers["weight"] - calls = [] - - def read(): - calls.append(True) - return ir.tensor(data) - - source.const_value = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) - assert FoldWeightLayoutPass()(model).modified - assert not calls - common.RemoveUnusedNodesPass()(model) - assert set(model.graph.initializers) == {"layout"} - assert not FoldWeightLayoutPass()(model).modified - assert not calls - actual = model.graph.initializers["layout"].const_value.numpy() - assert len(calls) == 1 - assert ( - actual.tobytes() - == data.reshape(2, 2, 4, 3).transpose(0, 2, 1, 3).reshape(2, 8, 3).tobytes() - ) - - -@pytest.mark.parametrize("dynamic", ["weight", "shape1"]) -@pytest.mark.parametrize("overridable", [False, True]) -def test_dynamic_or_overridable_inputs_are_not_folded(dynamic, overridable): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - value = model.graph.initializers[dynamic] - model.graph.inputs.append(value) - if not overridable: - del model.graph.initializers[dynamic] - value.const_value = None - assert not FoldWeightLayoutPass()(model).modified - assert len(model.graph) == 4 - - -def test_shared_sources_are_not_duplicated(): - data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) - model = _model(data) - source = model.graph.initializers["weight"] - extra = ir.Node("", "Identity", [source], num_outputs=1) - extra.outputs[0].name = "original_weight" - extra.outputs[0].dtype = source.dtype - model.graph.append(extra) - model.graph.outputs.append(extra.outputs[0]) - layout = next(node for node in model.graph if node.outputs[0].name == "layout") - model.graph.outputs.append(layout.outputs[0]) - assert not FoldWeightLayoutPass()(model).modified - common.RemoveUnusedNodesPass()(model) - assert [output.name for output in model.graph.outputs] == [ - "output", - "original_weight", - "layout", - ] - assert extra.inputs[0] is source - assert source.name in model.graph.initializers - assert model.graph.outputs[2] is layout.outputs[0] - assert set(model.graph.initializers) == {"weight", "shape1", "shape2"} - assert source.const_value.numpy().tobytes() == data.tobytes() - - -@pytest.mark.parametrize( - ("domain", "op_type", "parameter_slot", "protected"), - [ - ("com.microsoft", "MoE", 2, True), - ("com.microsoft", "QMoE", 2, True), - ("com.microsoft", "QMoE", 3, True), - ("", "DequantizeLinear", 0, True), - ("", "DequantizeLinear", 1, True), - ("com.microsoft", "MatMulNBits", 1, False), - ("", "MatMul", 1, False), - ], -) -def test_post_weight_storage_layout_paths_preserve_named_sources( - domain, op_type, parameter_slot, protected -): - model = _model(np.arange(48, dtype=np.float32).reshape(2, 8, 3)) - source = model.graph.initializers["weight"] - - def unexpected_read(): - raise AssertionError("eligibility and layout folding must not read weight payloads") - - tensor = ir.LazyTensor(unexpected_read, dtype=source.dtype, shape=source.shape) - source.const_value = tensor - # Only the consumer's parameter position matters to this structural test; - # complete model/schema and public save/load contracts are covered by the - # existing QMoE/FP8/streaming model tests, unchanged by this regression fix. - inputs = [] - for index in range(parameter_slot): - argument = ir.Value( - name=f"argument_{index}", - type=ir.TensorType(ir.DataType.FLOAT), - shape=ir.Shape([1]), - ) - model.graph.inputs.append(argument) - inputs.append(argument) - inputs.append(model.graph.outputs[0]) - consumer = ir.Node(domain, op_type, inputs, num_outputs=1) - consumer.outputs[0].name = "consumer_output" - model.graph.append(consumer) - model.graph.outputs.append(consumer.outputs[0]) - fold_initializers_after_weights(model) - if protected: - assert model.graph.initializers["weight"] is source - assert source.const_value is tensor - assert sum(node.op_type == "Reshape" for node in model.graph) == 2 - assert sum(node.op_type == "Transpose" for node in model.graph) == 1 - else: - assert "weight" not in model.graph.initializers - assert not any(node.op_type in ("Reshape", "Transpose") for node in model.graph) - - -@pytest.mark.parametrize("with_dependency", [False, True]) -def test_qdq_scalar_scale_storage_survives_pipeline_and_package_roundtrip( - tmp_path, with_dependency -): - data = np.array([[0.5]], dtype=np.float32) - scale = _initializer("checkpoint_scale", data) - reads = [] - - def read_scale(): - reads.append(True) - return ir.tensor(data) - - tensor = ir.LazyTensor(read_scale, dtype=scale.dtype, shape=scale.shape) - scale.const_value = tensor - scalar_shape = _initializer("scalar_shape", np.array([], dtype=np.int64)) - reshape = ir.Node("", "Reshape", [scale, scalar_shape], num_outputs=1) - reshape.outputs[0].name = "runtime_scalar_scale" - reshape.outputs[0].dtype = ir.DataType.FLOAT - reshape.outputs[0].shape = ir.Shape([]) - codes = ir.Value(name="codes", type=ir.TensorType(ir.DataType.UINT8), shape=ir.Shape([4])) - inputs = [codes] - nodes = [reshape] - runtime_scale = reshape.outputs[0] - if with_dependency: - dependency = ir.Value( - name="dependency", type=ir.TensorType(ir.DataType.FLOAT), shape=ir.Shape([]) - ) - inputs.append(dependency) - add = ir.Node("", "Add", [runtime_scale, dependency], num_outputs=1) - add.outputs[0].name = "dependent_scale" - add.outputs[0].dtype = ir.DataType.FLOAT - add.outputs[0].shape = ir.Shape([]) - nodes.append(add) - runtime_scale = add.outputs[0] - dequantize = ir.Node("", "DequantizeLinear", [codes, runtime_scale], num_outputs=1) - dequantize.outputs[0].name = "output" - dequantize.outputs[0].dtype = ir.DataType.FLOAT - dequantize.outputs[0].shape = ir.Shape([4]) - nodes.append(dequantize) - model = ir.Model( - ir.Graph( - inputs, - [dequantize.outputs[0]], - nodes=nodes, - initializers=[scale, scalar_shape], - opset_imports={"": 21}, - ), - ir_version=10, - ) - fold_initializers_after_weights(model) - assert model.graph.initializers["checkpoint_scale"] is scale - assert scale.const_value is tensor - assert not reads - assert reshape.graph is model.graph - assert reshape.inputs[0] is scale - # The protected source stays in checkpoint orientation, not the scalar - # runtime format. This is a storage contract, not just numerical parity. - assert scale.shape == ir.Shape([1, 1]) - assert scale.dtype == ir.DataType.FLOAT - fold_initializers_after_weights(model) - assert model.graph.initializers["checkpoint_scale"] is scale - assert not reads - - ModelPackage({"model": model}).save(tmp_path, progress_bar=False) - restored = ModelPackage.load(tmp_path)["model"] - restored_scale = restored.graph.initializers["checkpoint_scale"] - assert restored_scale.shape == scale.shape - assert restored_scale.dtype == scale.dtype - assert restored_scale.const_value.numpy().tobytes() == data.tobytes() - for candidate in (model, restored): - session = OnnxModelSession(candidate, device="cpu") - for dependency_value in (0.0, 0.25) if with_dependency else (0.0,): - code_data = np.array([0, 1, 127, 255], dtype=np.uint8) - feeds: dict[str, np.ndarray] = {"codes": code_data} - if with_dependency: - feeds["dependency"] = np.asarray(dependency_value, dtype=np.float32) - actual = session.run(feeds)["output"] - np.testing.assert_array_equal( - actual, code_data.astype(np.float32) * (0.5 + dependency_value) - ) - - -def test_unshared_graph_output_layout_name_survives(): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - layout = next(node for node in model.graph if node.outputs[0].name == "layout") - model.graph.outputs.append(layout.outputs[0]) - assert FoldWeightLayoutPass()(model).modified - assert model.graph.outputs[1] is model.graph.initializers["layout"] - - -@pytest.mark.parametrize("op_type", ["Reshape", "Transpose"]) -def test_graph_output_source_is_not_duplicated(op_type): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - if op_type == "Transpose": - first = next(iter(model.graph)) - first.op_type = "Transpose" - first.replace_input_with(1, None) - first.resize_inputs(1) - first.attributes["perm"] = ir.AttrInt64s("perm", [2, 1, 0]) - model.graph.outputs.append(model.graph.initializers["weight"]) - assert not FoldWeightLayoutPass()(model).modified - assert "layout" not in model.graph.initializers - - -@pytest.mark.parametrize("style", ["value", "value_ints"]) -@pytest.mark.parametrize("shared", [False, True]) -def test_post_weight_pipeline_folds_constant_shape_nodes(style, shared): - data = np.arange(48, dtype=np.uint8).reshape(2, 8, 3) - model = _model(data) - if shared: - extra = ir.Node("", "Identity", [model.graph.initializers["weight"]], num_outputs=1) - extra.outputs[0].name = "shared" - model.graph.append(extra) - model.graph.outputs.append(extra.outputs[0]) - for name in ("shape1", "shape2"): - shape = model.graph.initializers[name] - attribute = ( - ir.AttrTensor("value", shape.const_value) - if style == "value" - else ir.AttrInt64s("value_ints", shape.const_value.numpy().tolist()) - ) - constant = ir.Node("", "Constant", [], attributes=[attribute], num_outputs=1) - constant.outputs[0].name = name - model.graph.insert_before(next(iter(model.graph)), constant) - shape.replace_all_uses_with(constant.outputs[0]) - del model.graph.initializers[name] - fold_initializers_after_weights(model) - if shared: - assert set(model.graph.initializers) == {"weight"} - assert sum(node.op_type == "Reshape" for node in model.graph) == 2 - assert sum(node.op_type == "Transpose" for node in model.graph) == 1 - return - assert not list(model.graph) - assert set(model.graph.initializers) == {"output"} - folded = model.graph.initializers["output"] - assert model.graph.outputs[0] is folded - assert isinstance(folded.const_value, ir.LazyTensor) - expected = data.reshape(2, 2, 4, 3).swapaxes(1, 2).reshape(2, 8, 3) - assert folded.const_value.numpy().tobytes() == expected.tobytes() - - -def test_tied_embedding_pipeline_keeps_one_table(tmp_path): - weight = _initializer("weight", np.arange(24, dtype=np.float16).reshape(8, 3)) - indices = ir.Value( - name="indices", shape=ir.Shape([2]), type=ir.TensorType(ir.DataType.INT64) - ) - gather = ir.Node("", "Gather", [weight, indices], num_outputs=1) - transpose = ir.Node( - "", "Transpose", [weight], attributes=[ir.AttrInt64s("perm", [1, 0])], num_outputs=1 - ) - matmul = ir.Node("", "MatMul", [gather.outputs[0], transpose.outputs[0]], num_outputs=1) - for node, name in ((gather, "embedding"), (transpose, "head"), (matmul, "output")): - node.outputs[0].name = name - node.outputs[0].dtype = ir.DataType.FLOAT16 - graph = ir.Graph( - [indices], - [matmul.outputs[0]], - nodes=[gather, transpose, matmul], - initializers=[weight], - opset_imports={"": 24}, - name="tied_table", - ) - model = ir.Model(graph, ir_version=10) - fold_initializers_after_weights(model) - assert set(graph.initializers) == {"weight"} - assert transpose.inputs[0] is gather.inputs[0] - ModelPackage({"model": model}).save(tmp_path, progress_bar=False) - loaded = ModelPackage.load(tmp_path)["model"].graph - assert set(loaded.initializers) == {"weight"} - assert [node.op_type for node in loaded] == ["Gather", "Transpose", "MatMul"] - - -@pytest.mark.parametrize( - ("source", "requested", "allowzero", "expected"), - [ - ((2, 3), (0, -1), 0, (2, 3)), - ((2, 3), (-1, 2), 0, (3, 2)), - ((0, 3), (0, 3), 1, (0, 3)), - ((0, 3), (-1, 3), 0, (0, 3)), - ((1,), (), 0, ()), - ], -) -def test_reshape_zero_inferred_and_scalar_shapes(source, requested, allowzero, expected): - weight = _initializer("weight", np.zeros(source, dtype=np.float16)) - shape = _initializer("shape", np.array(requested, dtype=np.int64)) - node = ir.Node( - "", - "Reshape", - [weight, shape], - attributes=[ir.AttrInt64("allowzero", allowzero)], - num_outputs=1, - ) - node.outputs[0].name = "output" - graph = ir.Graph( - [], - [node.outputs[0]], - nodes=[node], - initializers=[weight, shape], - opset_imports={"": 24}, - ) - model = ir.Model(graph, ir_version=10) - assert FoldWeightLayoutPass()(model).modified - assert model.graph.outputs[0].shape == ir.Shape(expected) - assert model.graph.outputs[0].const_value.numpy().shape == expected - - -@pytest.mark.parametrize("requested", [(0, -1), (-1, -1), (4,), (-2,)]) -def test_invalid_static_reshape_fails_explicitly(requested): - model = _model(np.zeros((0, 3), dtype=np.float16), first=requested) - with pytest.raises(ValueError, match="Reshape"): - FoldWeightLayoutPass()(model) - - -def test_default_transpose_and_missing_declared_dtype(): - data = np.arange(24, dtype=np.float16).reshape(2, 3, 4) - weight = _initializer("weight", data) - weight.type = None - node = ir.Node("", "Transpose", [weight], num_outputs=1) - node.outputs[0].name = "output" - graph = ir.Graph([], [node.outputs[0]], nodes=[node], initializers=[weight]) - model = ir.Model(graph, ir_version=10) - assert FoldWeightLayoutPass()(model).modified - output = graph.outputs[0] - assert output.dtype == ir.DataType.FLOAT16 - assert output.const_value.numpy().tobytes() == data.transpose().tobytes() - - -def test_invalid_permutation_and_conflicting_dtype_fail(): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3), perm=(0, 0, 1, 2)) - with pytest.raises(ValueError, match="permutation"): - FoldWeightLayoutPass()(model) - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - model.graph.initializers["weight"].dtype = ir.DataType.FLOAT - with pytest.raises(ValueError, match="declares dtype"): - FoldWeightLayoutPass()(model) - - -def test_custom_domain_layout_op_is_not_folded(): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - transpose = next(node for node in model.graph if node.op_type == "Transpose") - transpose.domain = "custom" - FoldWeightLayoutPass()(model) - assert transpose in list(model.graph) - assert transpose.outputs[0].producer() is transpose - - -def test_missing_weights_are_warned_and_left_unmodified(caplog): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - model.graph.initializers["weight"].const_value = None - assert not FoldWeightLayoutPass()(model).modified - assert "has no data" in caplog.text - - -def test_invalid_shape_tensor_dtype_fails_explicitly(): - model = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)) - shape = model.graph.initializers["shape1"] - shape.dtype = ir.DataType.INT32 - shape.const_value = ir.tensor(np.array([2, 2, 4, 3], dtype=np.int32)) - with pytest.raises(ValueError, match="INT64 shape"): - FoldWeightLayoutPass()(model) - - -def test_nested_graph_initializers_are_registered_in_their_scope(): - first = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)).graph - second = _model(np.arange(48, dtype=np.uint8).reshape(2, 8, 3)).graph - condition = ir.Value( - name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) - ) - node = ir.Node( - "", - "If", - [condition], - attributes=[ir.AttrGraph("then_branch", first), ir.AttrGraph("else_branch", second)], - num_outputs=1, - ) - node.outputs[0].name = "result" - graph = ir.Graph([condition], [node.outputs[0]], nodes=[node], opset_imports={"": 24}) - model = ir.Model(graph, ir_version=10) - assert FoldWeightLayoutPass()(model).modified - assert not graph.initializers - for branch in (first, second): - assert [item.op_type for item in branch] == ["Identity"] - assert set(branch.initializers) == {"layout"} - - -@pytest.mark.parametrize("parent_capture", [False, True]) -def test_post_weight_pipeline_preserves_overridable_concat_input(tmp_path, parent_capture): - import onnxruntime as ort - - data = np.arange(8, dtype=np.float32) - weight = _initializer("W", data) - shape = _initializer("shape", np.array([2, 4], dtype=np.int64)) - bias = _initializer("B", np.zeros((1, 4), dtype=np.float32)) - reshape = ir.Node("", "Reshape", [weight, shape], num_outputs=1) - reshape.outputs[0].name = "layout" - reshape.outputs[0].shape = ir.Shape([2, 4]) - reshape.outputs[0].dtype = ir.DataType.FLOAT - concat = ir.Node( - "", - "Concat", - [reshape.outputs[0], bias], - [ir.AttrInt64("axis", 0)], - num_outputs=1, - ) - concat.outputs[0].name = "output" - concat.outputs[0].shape = ir.Shape([3, 4]) - concat.outputs[0].dtype = ir.DataType.FLOAT - graph = ir.Graph( - [bias], - concat.outputs, - nodes=[reshape, concat], - initializers=[weight, shape, bias], - name="overridable_concat", - opset_imports={"": 24}, - ) - if parent_capture: - graph.inputs.clear() - del graph.initializers["B"] - condition = ir.Value( - name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) - ) - fallback = ir.Node("", "Identity", [bias], num_outputs=1) - fallback.outputs[0].name = "fallback" - fallback.outputs[0].shape = bias.shape - fallback.outputs[0].dtype = bias.dtype - other = ir.Graph([], fallback.outputs, nodes=[fallback], name="other_branch") - select = ir.Node( - "", - "If", - [condition], - [ir.AttrGraph("then_branch", graph), ir.AttrGraph("else_branch", other)], - num_outputs=1, - ) - select.outputs[0].name = "result" - select.outputs[0].shape = ir.Shape([None, 4]) - select.outputs[0].dtype = ir.DataType.FLOAT - graph = ir.Graph( - [condition, bias], - select.outputs, - nodes=[select], - initializers=[bias], - name="captured_overridable_concat", - opset_imports={"": 24}, - ) - model = ir.Model(graph, ir_version=10) - fold_initializers_after_weights(model) - assert all(node.op_type != "Reshape" for node in model.graph.all_nodes()) - assert sum(node.op_type == "Concat" for node in model.graph.all_nodes()) == 1 - assert model.graph.initializers["B"] is bias - assert bias in model.graph.inputs - - ModelPackage({"model": model}).save(tmp_path, progress_bar=False, max_workers=1) - restored = ModelPackage.load(tmp_path)["model"] - for index, candidate in enumerate((model, restored)): - model_path = tmp_path / f"runtime_{index}.onnx" - ir.save(candidate, model_path) - # The test wrapper filters out overridable initializers via get_inputs(); - # use ORT directly so B is actually supplied to the runtime. - session = ort.InferenceSession(str(model_path), providers=["CPUExecutionProvider"]) - assert [value.name for value in session.get_overridable_initializers()] == ["B"] - for override in (None, np.full((1, 4), 17, dtype=np.float32)): - feeds = {"condition": np.asarray(True)} if parent_capture else {} - if override is not None: - feeds["B"] = override - actual = session.run(None, feeds)[0] - expected = np.concatenate( - [ - data.reshape(2, 4), - np.zeros((1, 4), dtype=np.float32) if override is None else override, - ] - ) - np.testing.assert_array_equal(actual, expected) - - -def test_post_weight_pipeline_concat_layouts_are_isolated_in_if_branches(tmp_path): - branches = [] - expected_outputs = [] - for offset in (0, 100): - data = np.arange(24, dtype=np.float32).reshape(2, 3, 4) + offset - zeros = np.zeros((1, 4, 3), dtype=np.float32) - # Each lexical scope deliberately uses the same names for different - # payloads. The rank-3 Transpose is folded by FoldWeightLayoutPass, - # which exposes an all-initializer Concat in each branch. - weight = _initializer("W", data) - padding = _initializer("Z", zeros) - transpose = ir.Node( - "", "Transpose", [weight], [ir.AttrInt64s("perm", [0, 2, 1])], num_outputs=1 - ) - transpose.outputs[0].name = "layout" - transpose.outputs[0].shape = ir.Shape([2, 4, 3]) - transpose.outputs[0].dtype = ir.DataType.FLOAT - concat = ir.Node( - "", - "Concat", - [transpose.outputs[0], padding], - [ir.AttrInt64("axis", 0)], - num_outputs=1, - ) - concat.outputs[0].name = "branch_output" - concat.outputs[0].shape = ir.Shape([3, 4, 3]) - concat.outputs[0].dtype = ir.DataType.FLOAT - branches.append( - ir.Graph( - [], - [concat.outputs[0]], - nodes=[transpose, concat], - initializers=[weight, padding], - name=f"branch_{offset}", - ) - ) - expected_outputs.append(np.concatenate([data.transpose(0, 2, 1), zeros], axis=0)) - condition = ir.Value( - name="condition", shape=ir.Shape([]), type=ir.TensorType(ir.DataType.BOOL) - ) - select = ir.Node( - "", - "If", - [condition], - [ - ir.AttrGraph("then_branch", branches[0]), - ir.AttrGraph("else_branch", branches[1]), - ], - num_outputs=1, - ) - select.outputs[0].name = "result" - select.outputs[0].shape = ir.Shape([3, 4, 3]) - select.outputs[0].dtype = ir.DataType.FLOAT - model = ir.Model( - ir.Graph( - [condition], - [select.outputs[0]], - nodes=[select], - opset_imports={"": 24}, - name="layout_concat_if", - ), - ir_version=10, - ) - - fold_initializers_after_weights(model) - assert not model.graph.initializers - packed_name = "layout__Z__axis_0__concat" - packed_values = [] - for branch, expected in zip(branches, expected_outputs, strict=True): - assert not list(branch) - assert set(branch.initializers) == {packed_name} - packed = branch.initializers[packed_name] - assert branch.outputs[0] is packed - assert isinstance(packed.const_value, ir.LazyTensor) - np.testing.assert_array_equal(packed.const_value.numpy(), expected) - packed_values.append(packed) - assert packed_values[0] is not packed_values[1] - fold_initializers_after_weights(model) - assert not model.graph.initializers - assert all( - branch.initializers[packed_name] is packed - for branch, packed in zip(branches, packed_values, strict=True) - ) - - output_directory = tmp_path / "package" - ModelPackage({"model": model}).save( - str(output_directory), max_workers=1, progress_bar=False - ) - restored = ModelPackage.load(str(output_directory))["model"] - assert not restored.graph.initializers - restored_select = next(iter(restored.graph)) - for attribute, expected in zip( - ("then_branch", "else_branch"), expected_outputs, strict=True - ): - branch = restored_select.attributes[attribute].as_graph() - assert not list(branch) - assert set(branch.initializers) == {packed_name} - packed = branch.initializers[packed_name] - assert branch.outputs[0] is packed - np.testing.assert_array_equal(packed.const_value.numpy(), expected) - - # Both branches must execute their own payload, before and after the public - # package roundtrip, rather than reusing a same-named sibling initializer. - for candidate in (model, restored): - session = OnnxModelSession(candidate, device="cpu") - for take_then, expected in zip((True, False), expected_outputs, strict=True): - result = session.run({"condition": np.asarray(take_then, dtype=np.bool_)}) - np.testing.assert_array_equal(result["result"], expected) - - -@pytest.mark.parametrize("retained", [None, "shared", "graph_output"]) -def test_parent_capture_initializer_ownership_and_persisted_payload(tmp_path, retained): - data = np.arange(6, dtype=np.float32).reshape(2, 3) - source = _initializer("parent_weight", data) - transpose = ir.Node( - "", "Transpose", [source], [ir.AttrInt64s("perm", [1, 0])], num_outputs=1 - ) - transpose.outputs[0].name = "captured_layout" - transpose.outputs[0].shape = ir.Shape([3, 2]) - transpose.outputs[0].dtype = ir.DataType.FLOAT - identity = ir.Node("", "Identity", [transpose.outputs[0]], num_outputs=1) - identity.outputs[0].name = "branch_output" - identity.outputs[0].shape = ir.Shape([3, 2]) - identity.outputs[0].dtype = ir.DataType.FLOAT - branch = ir.Graph([], [identity.outputs[0]], nodes=[transpose, identity]) - other = _initializer("other_weight", np.zeros((3, 2), dtype=np.float32)) - other_branch = ir.Graph([], [other], nodes=[], initializers=[other]) - 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=1, - ) - select.outputs[0].name = "output" - select.outputs[0].shape = ir.Shape([3, 2]) - select.outputs[0].dtype = ir.DataType.FLOAT - graph = ir.Graph( - [condition], - [select.outputs[0]], - nodes=[select], - initializers=[source], - opset_imports={"": 24}, - ) - if retained == "shared": - extra = ir.Node("", "Identity", [source], num_outputs=1) - extra.outputs[0].name = "shared_source" - extra.outputs[0].shape = source.shape - extra.outputs[0].dtype = source.dtype - graph.append(extra) - graph.outputs.append(extra.outputs[0]) - elif retained == "graph_output": - graph.outputs.append(source) - model = ir.Model(graph, ir_version=10) - result = FoldWeightLayoutPass()(model) - assert result.modified == (retained is None) - if retained is None: - assert "parent_weight" not in graph.initializers - assert set(branch.initializers) == {"captured_layout"} - assert not FoldWeightLayoutPass()(model).modified - np.testing.assert_array_equal( - branch.initializers["captured_layout"].const_value.numpy(), data.T - ) - else: - assert graph.initializers["parent_weight"] is source - assert not branch.initializers - assert transpose.inputs[0] is source - # Exercise the public serializer, not a protobuf encoding. The retained - # parent case must stay shared; the unshared case must serialize only the - # folded child weight, not both it and the dead parent payload. - fold_initializers_after_weights(model) - ModelPackage({"model": model}).save(str(tmp_path), max_workers=1, progress_bar=False) - restored = ModelPackage.load(str(tmp_path))["model"] - persisted = [value for scope in restored.graphs() for value in scope.initializers.values()] - assert sum(value.const_value.numpy().nbytes for value in persisted) == data.nbytes * 2 - if retained is None: - assert "parent_weight" not in restored.graph.initializers - restored_branch = next(iter(restored.graph)).attributes["then_branch"].as_graph() - assert set(restored_branch.initializers) == {"captured_layout"} - np.testing.assert_array_equal( - restored_branch.initializers["captured_layout"].const_value.numpy(), data.T - ) - else: - np.testing.assert_array_equal( - restored.graph.initializers["parent_weight"].const_value.numpy(), data - ) - - -@pytest.mark.parametrize("width", [3, 3 * 8192]) -def test_shape_constant_folding_exposes_layout_without_reading_large_weights(width): - data = np.arange(16 * width, dtype=np.float32).reshape(2, 8, width) - model = _model(data, first=(2, 2, 4, width), last=(2, 8, width)) - source = model.graph.initializers["weight"] - reads = [] - - def read(): - reads.append(True) - return ir.tensor(data) - - source.const_value = ir.LazyTensor(read, dtype=source.dtype, shape=source.shape) - shape = model.graph.initializers["shape1"] - offset = _initializer("shape_offset", np.zeros(4, dtype=np.int64)) - model.graph.register_initializer(offset) - expression = ir.Node("", "Add", [shape, offset], num_outputs=1) - expression.outputs[0].name = "computed_shape" - expression.outputs[0].shape = shape.shape - expression.outputs[0].dtype = ir.DataType.INT64 - first = next(iter(model.graph)) - model.graph.insert_before(first, expression) - first.replace_input_with(1, expression.outputs[0]) - fold_initializers_after_weights(model) - assert not reads - assert set(model.graph.initializers) == {"output"} - assert not list(model.graph) - folded = model.graph.initializers["output"] - assert model.graph.outputs[0] is folded - assert isinstance(folded.const_value, ir.LazyTensor) - fold_initializers_after_weights(model) - assert not reads - assert model.graph.initializers["output"] is folded - assert model.graph.outputs[0] is folded - expected = data.reshape(2, 2, 4, width).transpose(0, 2, 1, 3).reshape(2, 8, width) - np.testing.assert_array_equal(folded.const_value.numpy(), expected) - assert len(reads) == 1 - - -def test_folded_layout_keeps_pruned_external_source_ownership(tmp_path, monkeypatch): - # The 48-element fixture used elsewhere is only 192 bytes: ir.save keeps - # tensors below its 256-byte default threshold inline, so it has no file - # backing to protect. This fixture must exercise a real ExternalTensor. - width = 24 - data = np.arange(16 * width, dtype=np.float32).reshape(2, 8, width) - base = tmp_path / "base" - model = _model(data, first=(2, 2, 4, width), last=(2, 8, width)) - ModelPackage({"model": model}).save(str(base), max_workers=1, progress_bar=False) - before = {path.name: path.read_bytes() for path in base.iterdir()} - package = ModelPackage.load(str(base)) - model = package["model"] - source = model.graph.initializers["weight"].const_value - assert isinstance(source, ir.ExternalTensor) - source_path = Path(source.path).resolve() - assert source_path == (base / "model.onnx.data").resolve() - - def unexpected_read(*args, **kwargs): - raise AssertionError("source-alias refusal must precede lazy payload reads") - - with monkeypatch.context() as patch: - patch.setattr(ir.ExternalTensor, "numpy", unexpected_read) - fold_initializers_after_weights(model) - folded_graph = model.graph - assert set(folded_graph.initializers) == {"output"} - folded = folded_graph.initializers["output"] - assert folded_graph.outputs[0] is folded - assert isinstance(folded.const_value, ir.LazyTensor) - assert source_path in { - Path(path).resolve() for path in model.meta[LAZY_SOURCE_FILES_META_KEY] - } - assert source_path in { - Path(path).resolve() - for path in folded.const_value.meta[LAZY_SOURCE_FILES_META_KEY] - } - with pytest.raises(ValueError, match="aliases lazy source file"): - package.save(str(base), max_workers=1, progress_bar=False) - assert source_path in package._native_streaming_source_files - assert before == {path.name: path.read_bytes() for path in base.iterdir()} - safe = tmp_path / "safe" - package.save(str(safe), max_workers=1, progress_bar=False) - restored_graph = ModelPackage.load(str(safe))["model"].graph - assert set(restored_graph.initializers) == {"output"} - restored = restored_graph.initializers["output"] - assert restored_graph.outputs[0] is restored - expected = data.reshape(2, 2, 4, width).transpose(0, 2, 1, 3).reshape(2, 8, width) - np.testing.assert_array_equal(restored.const_value.numpy(), expected) From 54327d3433f69456de4d34eb9c44d1515b457c55 Mon Sep 17 00:00:00 2001 From: titaiwang Date: Wed, 28 Oct 2026 22:24:56 +0000 Subject: [PATCH 7/7] Require independent opt-in for float projection packing Separate the new float projection stage from the unchanged legacy GQA PackQKV capability. Append the default-empty capability without shifting existing positional fields. Expand CPU gating and arbitrary-mask tests, and document the current metadata namespace. Validation542 passed CPU/offline; no CUDA numerical qualification claimed. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Copilot-Session: 9a2a5b82-7cd3-4a76-bb94-0998527ea723 Signed-off-by: titaiwang --- docs/ep_quickstart.md | 17 +++--- docs/execution_providers.md | 20 ++++--- docs/gemma4-runtime.md | 5 +- src/mobius/_execution_providers.py | 16 +++--- src/mobius/_optimizations.py | 2 +- .../_qkv_projection_packing_test.py | 53 +++++++++++++++---- 6 files changed, 78 insertions(+), 35 deletions(-) diff --git a/docs/ep_quickstart.md b/docs/ep_quickstart.md index 5414225fc..7f4e645d6 100644 --- a/docs/ep_quickstart.md +++ b/docs/ep_quickstart.md @@ -100,7 +100,8 @@ 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}); both float packing stages +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 @@ -133,14 +134,12 @@ 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 INT4 projection packing requires an explicit -`matmul_nbits_qkv_pack_dtypes` capability override; no built-in EP enables it. -The existing float GQA-input `PackQKV` stage is unchanged, but the additional -`PackQKVProjections` float stage is **new and default-enabled** for decoder -and masked-decoder roles whenever the dtype is in `qkv_pack_dtypes`, including -CUDA's existing FLOAT/FP16/BF16 profile. Only INT4 projection packing is -default-off. Its failed/unwaived numerical gate does not qualify the new float -stage; float numerical accuracy is not separately qualified by this branch. Read the +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. diff --git a/docs/execution_providers.md b/docs/execution_providers.md index dcc35e39d..313cc8967 100644 --- a/docs/execution_providers.md +++ b/docs/execution_providers.md @@ -109,11 +109,13 @@ 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` | Automatic for `qkv_pack_dtypes`, including CUDA FLOAT/FP16/BF16 | +| 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 INT4 opt-in uses the process-local capability API; stock `--ep cuda` -does not enable it: +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 @@ -124,15 +126,18 @@ 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, ) ``` -The INT4 capability is independent of float packing, activation/storage dtype, -component precision, and FP8 KV caches. Packing never dequantizes, requantizes, or -promotes weights. `_qkv_projection_packing.py` implements both float and INT4. +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. @@ -188,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] # both float packing stages + 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 index 0acd1f66c..6f5762ab1 100644 --- a/docs/gemma4-runtime.md +++ b/docs/gemma4-runtime.md @@ -159,7 +159,7 @@ and is not claimed bitwise exact. The learned `router.scale` remains unchanged. Incomplete, stale, or shape-incompatible expert sidecars are rejected. Float projection packing retains version-1 -`mobius.qkv_projection_sources` metadata on the `Split`: original initializer +`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 @@ -168,6 +168,9 @@ 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 diff --git a/src/mobius/_execution_providers.py b/src/mobius/_execution_providers.py index 51cb1c224..88f9666fe 100644 --- a/src/mobius/_execution_providers.py +++ b/src/mobius/_execution_providers.py @@ -53,18 +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 enabling both the existing decoder-only - GroupQueryAttention input packing (PackQKV) and the new additional - float projection packing (PackQKVProjections) for decoder and - masked-decoder roles. Nonempty built-in profiles, including CUDA, - enable the new float stage by default; INT4 uses its separate - capability below. Set to an empty frozenset for EPs that do not + 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 / @@ -182,6 +183,9 @@ class EpCapabilities: 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/_optimizations.py b/src/mobius/_optimizations.py index f46efdbf3..e21d52bbd 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -461,7 +461,7 @@ def _get_optimization_passes( fuse.append(("PackQKV", list(pack_qkv_for_gqa_rules()))) if model_role in {"decoder", "masked-decoder"}: - if dtype in caps.qkv_pack_dtypes: + 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())) diff --git a/src/mobius/rewrite_rules/_qkv_projection_packing_test.py b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py index f65a48983..0fba9651a 100644 --- a/src/mobius/rewrite_rules/_qkv_projection_packing_test.py +++ b/src/mobius/rewrite_rules/_qkv_projection_packing_test.py @@ -907,13 +907,12 @@ def test_ep_dtype_role_gating(): 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 ) - assert ("PackQKVProjections" in names) == ( - role != "encoder" and dtype in caps.qkv_pack_dtypes - ) @pytest.mark.parametrize("dtype", list(_DTYPES)) @@ -922,33 +921,65 @@ 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 dataclasses.fields(EpCapabilities)[-1].name == "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]) -def test_generic_pipeline_preserves_arbitrary_mask_contract(monkeypatch, trace, opt_in): - model, _, _, _ = _model(quantized=True, deferred=True, dtype=ir.DataType.FLOAT16) +@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({ir.DataType.FLOAT16}) if opt_in else frozenset() - ), + 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=ir.DataType.FLOAT16, trace=trace) + 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) - assert sum(node.op_type == "MatMulNBits" for node in model.graph) == (1 if opt_in else 3) + 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])