diff --git a/examples/gemma4_unified_ort_genai.py b/examples/gemma4_unified_ort_genai.py new file mode 100644 index 000000000..0aa5bda7f --- /dev/null +++ b/examples/gemma4_unified_ort_genai.py @@ -0,0 +1,491 @@ +#!/usr/bin/env python +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +r"""Gemma-4-12B (``gemma4_unified``) multimodal generation with onnxruntime-genai. + +``google/gemma-4-12B`` is an **encoder-free unified multimodal** checkpoint +(HuggingFace ``model_type == "gemma4_unified"``). ``mobius.build`` auto-detects +it and produces a **4-model package**: ``decoder`` + ``embedding`` + +``vision_encoder`` + ``audio_encoder``. Unlike the released ``gemma4`` models +there is no SigLIP / Conformer tower — raw merged pixel patches (48x48, 6912-dim) +and raw waveform-frame features (640-dim) are projected directly into language +space by ``vision_encoder`` / ``audio_encoder``, then scattered into +``inputs_embeds`` by ``embedding`` at the image / audio placeholder positions. + +This script builds the full package via the real +``mobius.integrations.ort_genai`` export path and runs **text**, **image+text**, +and **audio+text** generation through onnxruntime-genai on the same package. + +Two runtime caveats are handled below: + +1. **Patched onnxruntime-genai required.** gemma-4-12B has a *mixed* KV cache: + most layers use 8 heads x head_dim 256, but the global-attention layers use a + single 1 head x head_dim 512 KV. A stock genai build assumes uniform KV + shapes and fails to load the decoder. Use a build that reads per-layer + ``num_heads`` / ``head_dim`` from each ``past_key_values.*`` input shape. + +2. **HuggingFace preprocessing for image / audio.** genai's built-in + ``Gemma4ImageTransform`` targets the SigLIP ``gemma4`` patch contract + (16px, 768-dim), which does **not** match this encoder-free unified model + (48px merged patches, 6912-dim). Until a genai-native unified transform + exists, this example preprocesses image / audio with the HuggingFace + ``AutoProcessor`` and feeds the resulting tensors to genai via + ``Generator.set_inputs(NamedTensors)`` (which bypasses genai's own + transform). Text generation uses the native genai path. + +3. **Structural-token suppression.** HF's ``generation_config`` suppresses the + ```` / ```` tokens (``suppress_tokens``). genai + has no native equivalent, so without it this base checkpoint degenerates into + repeating ```` after an image instead of describing it. The decode + loop masks those token ids to ``-inf`` before sampling (see SUPPRESS_TOKEN_IDS + and ``_decode_loop``). + +``google/gemma-4-12B`` is a **base** (non-instruction-tuned) checkpoint, so +completion-style leads work far better than instructions (the defaults use +"This image shows" / "The audio says", not "Describe ..."). Verified outputs on +GPU (f16, greedy, matching HuggingFace ``model.generate``): + +- text "The capital of Japan is" -> " Tokyo." +- image (Sydney Chinatown photo) -> " the Chinese Arch in the Chinatown of + Sydney, Australia." +- audio (LibriSpeech clip) -> " He hoped there would be stew for + dinner, turnips and carrots and bruised potatoes ..." + +Numerical correctness of the image / audio pipeline is also verified by the +integration tests (``tests/integration_test.py:: +test_gemma4_unified_12b_multimodal_prefill``, vision cosine 1.0 vs HuggingFace). + +Optional ``--quantize Q4_K_M`` INT4-quantizes the decoder with Olive (~23GB -> +~6.8GB, 3.4x smaller). Spot-checked quality on this base model: coherent text +/ image / audio generation, ~0.986 last-token logit cosine and ~75% greedy +top-1 agreement vs f16 on short factual prompts (the base model itself emits +some off-distribution tokens, so a few disagreements are not quantization +artifacts). + +Requirements:: + + pip install mobius-ai[ort-genai] transformers pillow librosa + # plus a per-layer-KV-aware onnxruntime-genai build (see caveat 1) + +Usage:: + + # Build the package once and reuse it across modes: + python examples/gemma4_unified_ort_genai.py --mode text --save-to out/gemma4_12b/ + + # Image + text (reuse the built dir): + python examples/gemma4_unified_ort_genai.py --mode image \ + --model-dir out/gemma4_12b/ --image path/to/photo.jpg + + # Audio + text: + python examples/gemma4_unified_ort_genai.py --mode audio \ + --model-dir out/gemma4_12b/ --audio path/to/clip.flac + + # INT4-quantize the decoder with Olive (Q4_K_M, ~3.4x smaller) and run it: + python examples/gemma4_unified_ort_genai.py --mode image \ + --model-dir out/gemma4_12b/ --image path/to/photo.jpg \ + --quantize Q4_K_M --quantized-out out/gemma4_12b-Q4_K_M/ +""" + +from __future__ import annotations + +import argparse +import os +import sys +import tempfile + +import numpy as np +import onnxruntime_genai as og + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +MODEL_ID = "google/gemma-4-12B" +# gemma uses BOS id 2. genai's tokenizer defaults to add_special_tokens=false +# and this base checkpoint has no chat template, so BOS must be added manually. +BOS_TOKEN_ID = 2 +IMAGE_TOKEN_ID = 258880 +AUDIO_TOKEN_ID = 258881 +# Structural multimodal tokens that the base model tends to emit verbatim during +# generation (```` / ````). HF's generation_config +# lists them in ``suppress_tokens``; genai has no native suppression, so the +# decode loop forces their logits to -inf (see _decode_loop). Without this the +# base checkpoint degenerates into repeating ```` instead of captioning. +SUPPRESS_TOKEN_IDS = (258882, 258883) +MAX_NEW_TOKENS = 40 + + +# --------------------------------------------------------------------------- +# Export +# --------------------------------------------------------------------------- + + +def build_and_export(model_id: str, output_dir: str, dtype: str, ep: str) -> None: + """Build the full unified multimodal package and write ORT GenAI artifacts. + + Produces ``decoder`` + ``embedding`` + ``vision_encoder`` + ``audio_encoder`` + ONNX models plus ``genai_config.json``, tokenizer files, ``image_processor`` + and audio feature-extraction configs, via the real + ``mobius.integrations.ort_genai.auto_export`` path. + + Args: + model_id: HuggingFace model ID (``google/gemma-4-12B``). + output_dir: Directory to write all outputs. + dtype: Model dtype (``"f16"`` recommended for CUDA). + ep: Execution provider for ``genai_config.json`` session options. + """ + from mobius.integrations.ort_genai import auto_export + + print( + f"Building unified multimodal package for {model_id!r} " + f"(dtype={dtype}, ep={ep}) — this downloads ~25GB of weights ..." + ) + manifest = auto_export(model_id, output_dir, dtype=dtype, ep=ep) + print(f"Export complete -> {output_dir}") + for name, path in sorted(manifest.items()): + print(f" {name}: {path}") + + +# --------------------------------------------------------------------------- +# Olive INT4 quantization (decoder only) +# --------------------------------------------------------------------------- + + +def quantize_decoder( + src_dir: str, + dst_dir: str, + *, + precision: str = "Q4_K_M", + block_size: int = 32, +) -> None: + """INT4-quantize the decoder sub-model with Olive (k-quant or NF4). + + Only the decoder is quantized — it holds ~23GB of the package's weights + (>95%). The embedding / vision_encoder / audio_encoder sub-models, the + tokenizer, ``genai_config.json``, and the processor configs are copied + over unchanged, so the result is a drop-in ORT GenAI package. + + Args: + src_dir: full-precision package (output of :func:`build_and_export`). + dst_dir: destination directory for the quantized package. + precision: ``"Q4_K_M"`` (k-quant; install ``cupy-cuda12x`` for the + 19-51x GPU speedup) or ``"NF4"`` (4-bit NormalFloat, native C++). + block_size: k-quant block size; ignored for NF4. + + Requires ``olive-ai`` (``pip install olive-ai``). + """ + import shutil + + from olive.workflows import run as olive_run + + if precision == "Q4_K_M": + pass_cfg = {"type": "OnnxKQuantQuantization", "bits": 4, "block_size": block_size} + elif precision == "NF4": + pass_cfg = {"type": "OnnxBnb4Quantization", "precision": "nf4"} + else: + raise ValueError(f"Unsupported precision: {precision!r}") + + decoder_dst = os.path.join(dst_dir, "decoder") + os.makedirs(decoder_dst, exist_ok=True) + config = { + "input_model": { + "type": "OnnxModel", + "model_path": os.path.join(src_dir, "decoder", "model.onnx"), + }, + "passes": {precision.lower(): pass_cfg}, + "output_dir": decoder_dst, + } + print(f"Quantizing decoder ({precision}) -> {decoder_dst} ...") + olive_run(config) + + # Olive writes bookkeeping files that ORT GenAI must not see in the + # decoder/ folder; keep only the ONNX model + its external data. + for f in os.listdir(decoder_dst): + if not (f == "model.onnx" or f.startswith("model.onnx.data")): + path = os.path.join(decoder_dst, f) + shutil.rmtree(path) if os.path.isdir(path) else os.remove(path) + + _ensure_logits_output(os.path.join(decoder_dst, "model.onnx")) + + # Copy the non-decoder sub-models + config + tokenizer unchanged. + for child in os.listdir(src_dir): + path = os.path.join(src_dir, child) + target = os.path.join(dst_dir, child) + if child == "decoder" or os.path.exists(target): + continue + shutil.copytree(path, target) if os.path.isdir(path) else shutil.copy2(path, target) + print(f"Quantized package ready -> {dst_dir}") + + +def _ensure_logits_output(decoder_path: str) -> None: + """Rename a quantized decoder's ``logits_Q4`` output back to ``logits``. + + Some Olive versions rename the decoder's ``logits`` output to ``logits_Q4`` + during k-quant. ORT GenAI maps outputs by the names in the (copied) + ``genai_config.json`` (which says ``logits``), so we rename the graph + output back when needed to keep the quantized decoder a drop-in. + """ + import onnx_ir as ir + + model = ir.load(decoder_path) + renamed = False + for value in model.graph.outputs: + if value.name == "logits_Q4": + value.name = "logits" + renamed = True + if renamed: + ir.save(model, decoder_path, external_data="model.onnx.data") + print(" Renamed decoder output logits_Q4 -> logits") + + +# --------------------------------------------------------------------------- +# Generation helpers +# --------------------------------------------------------------------------- + + +def _load_model(model_dir: str, ep: str) -> tuple[og.Model, og.Tokenizer]: + config = og.Config(model_dir) + config.clear_providers() + if ep != "cpu": + config.append_provider(ep) + model = og.Model(config) + return model, og.Tokenizer(model) + + +def _decode_loop( + model: og.Model, generator: og.Generator, tokenizer: og.Tokenizer, max_new: int +) -> str: + stream = tokenizer.create_stream() + tokens: list[int] = [] + for _ in range(max_new): + if generator.is_done(): + break + # Suppress the structural multimodal tokens before sampling, mirroring + # HF generation_config's ``suppress_tokens`` (genai has no native + # equivalent). get_logits -> mask -> set_logits -> sample. + logits = generator.get_logits() + logits[..., list(SUPPRESS_TOKEN_IDS)] = float("-inf") + generator.set_logits(logits) + generator.generate_next_token() + tok = int(generator.get_next_tokens()[0]) + tokens.append(tok) + print(stream.decode(tok), end="", flush=True) + print() + return tokenizer.decode(tokens) + + +def generate_text(model_dir: str, prompt: str, ep: str, max_new: int) -> str: + """Greedy text generation through the native genai path (BOS prepended).""" + model, tokenizer = _load_model(model_dir, ep) + # Prepend BOS manually: base checkpoint, no chat template (see BOS_TOKEN_ID). + input_ids = [BOS_TOKEN_ID, *tokenizer.encode(prompt)] + + params = og.GeneratorParams(model) + params.set_search_options(max_length=len(input_ids) + max_new, do_sample=False) + generator = og.Generator(model, params) + generator.append_tokens(input_ids) + + print(f"\nPrompt: {prompt}\n" + "-" * 40) + text = _decode_loop(model, generator, tokenizer, max_new) + print("-" * 40) + del generator + return text + + +def generate_image( + model_dir: str, model_id: str, image_path: str, prompt: str, ep: str, max_new: int +) -> str: + """Image+text generation: HF unified processor -> genai ``set_inputs``. + + genai's built-in image transform targets the SigLIP ``gemma4`` contract, so + we preprocess with the HuggingFace processor (48px merged patches, 6912-dim) + and inject the tensors directly. ``set_inputs`` bypasses genai's transform; + genai then runs ``vision_encoder -> embedding -> decoder``. + """ + from PIL import Image + from transformers import AutoProcessor + + model, tokenizer = _load_model(model_dir, ep) + processor = AutoProcessor.from_pretrained(model_id) + image = Image.open(image_path).convert("RGB") + + # The HF processor inserts IMAGE_TOKEN_ID placeholders and (for gemma) BOS. + proc = processor( + text=[f"{processor.image_token}{prompt}"], images=[image], return_tensors="pt" + ) + input_ids = proc["input_ids"].numpy().astype(np.int32) + n_image_tokens = int((input_ids == IMAGE_TOKEN_ID).sum()) + + nt = og.NamedTensors() + nt["input_ids"] = input_ids + # Graph input names: pixel_values, pixel_position_ids (HF names them + # pixel_values / image_position_ids). + nt["pixel_values"] = proc["pixel_values"].numpy().astype(np.float16) + nt["pixel_position_ids"] = proc["image_position_ids"].numpy().astype(np.int64) + nt["num_image_tokens"] = np.array([n_image_tokens], dtype=np.int64) + + params = og.GeneratorParams(model) + params.set_search_options(max_length=input_ids.shape[1] + max_new, do_sample=False) + generator = og.Generator(model, params) + generator.set_inputs(nt) + + print(f"\nImage: {image_path}\nPrompt: {prompt}\n" + "-" * 40) + text = _decode_loop(model, generator, tokenizer, max_new) + print("-" * 40) + del generator + return text + + +def generate_audio( + model_dir: str, model_id: str, audio_path: str, prompt: str, ep: str, max_new: int +) -> str: + """Audio+text generation: HF unified processor -> genai ``set_inputs``. + + Mirrors :func:`generate_image` for the audio branch. genai derives the + audio-token count from the summed ``audio_sizes`` input, then runs + ``audio_encoder -> embedding -> decoder``. + """ + import librosa + from transformers import AutoProcessor + + model, tokenizer = _load_model(model_dir, ep) + processor = AutoProcessor.from_pretrained(model_id) + waveform, _ = librosa.load(audio_path, sr=16000) + + proc = processor( + text=[f"{processor.audio_token}{prompt}"], + audio=[waveform], + return_tensors="pt", + ) + input_ids = proc["input_ids"].numpy().astype(np.int32) + n_audio_tokens = int((input_ids == AUDIO_TOKEN_ID).sum()) + + nt = og.NamedTensors() + nt["input_ids"] = input_ids + nt["input_features"] = proc["input_features"].numpy().astype(np.float16) + nt["input_features_mask"] = proc["input_features_mask"].numpy().astype(bool) + # genai sums audio_sizes to get the audio-token count for the speech branch. + nt["audio_sizes"] = np.array([n_audio_tokens], dtype=np.int64) + + params = og.GeneratorParams(model) + params.set_search_options(max_length=input_ids.shape[1] + max_new, do_sample=False) + generator = og.Generator(model, params) + generator.set_inputs(nt) + + print(f"\nAudio: {audio_path}\nPrompt: {prompt}\n" + "-" * 40) + text = _decode_loop(model, generator, tokenizer, max_new) + print("-" * 40) + del generator + return text + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__.split("\n", 1)[0]) + parser.add_argument( + "--mode", + default="text", + choices=["text", "image", "audio"], + help="Which modality to demonstrate.", + ) + parser.add_argument("--model-id", default=MODEL_ID, help="HuggingFace model ID.") + parser.add_argument("--prompt", default=None, help="Generation prompt.") + parser.add_argument("--image", default=None, help="Image path (image mode).") + parser.add_argument("--audio", default=None, help="Audio path (audio mode).") + parser.add_argument( + "--max-new-tokens", type=int, default=MAX_NEW_TOKENS, help="Max new tokens." + ) + parser.add_argument( + "--dtype", default="f16", choices=["f32", "f16", "bf16"], help="Model dtype." + ) + parser.add_argument("--ep", default="cuda", help="Execution provider.") + parser.add_argument( + "--model-dir", default=None, help="Reuse a pre-built export directory." + ) + parser.add_argument( + "--save-to", default=None, help="Build + save to this directory (skip cleanup)." + ) + parser.add_argument( + "--quantize", + default=None, + choices=["Q4_K_M", "NF4"], + help="INT4-quantize the decoder with Olive and run against the result.", + ) + parser.add_argument( + "--quantized-out", + default=None, + help="Output dir for --quantize (default: -).", + ) + args = parser.parse_args() + + # Default prompts per mode. This is a *base* (non-instruction-tuned) + # checkpoint, so completion-style leads work far better than instructions. + prompt = ( + args.prompt + or { + "text": "The capital of France is", + "image": "This image shows", + "audio": "The audio says", + }[args.mode] + ) + + if args.mode == "image" and not args.image: + parser.error("--image is required for --mode image") + if args.mode == "audio" and not args.audio: + parser.error("--audio is required for --mode audio") + + # Determine the export directory. + tmp_dir: str | None = None + if args.model_dir is not None: + model_dir = args.model_dir + else: + model_dir = args.save_to + if model_dir is None: + tmp_dir = tempfile.mkdtemp(prefix="gemma4_12b_mm_") + model_dir = tmp_dir + build_and_export(args.model_id, model_dir, args.dtype, args.ep) + + # Optionally INT4-quantize the decoder and run against the quantized package. + if args.quantize: + quant_dir = args.quantized_out or f"{model_dir.rstrip('/')}-{args.quantize}" + quantize_decoder(model_dir, quant_dir, precision=args.quantize) + model_dir = quant_dir + + try: + if args.mode == "text": + out = generate_text(model_dir, prompt, args.ep, args.max_new_tokens) + elif args.mode == "image": + out = generate_image( + model_dir, + args.model_id, + args.image, + prompt, + args.ep, + args.max_new_tokens, + ) + else: + out = generate_audio( + model_dir, + args.model_id, + args.audio, + prompt, + args.ep, + args.max_new_tokens, + ) + print(f"\nORT GenAI output:\n{out}") + finally: + if tmp_dir is not None and args.save_to is None: + import shutil + + shutil.rmtree(tmp_dir, ignore_errors=True) + + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index 1d8451715..ad377e3d3 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -403,10 +403,14 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N # Load images from testdata/ images = [Image.open(Path("testdata") / img_path) for img_path in case.images] - # Build chat-formatted prompt with image placeholders if the - # processor supports apply_chat_template (Qwen-VL, Gemma-3, etc.) + # Build chat-formatted prompt with image placeholders if the processor has a + # usable chat template (Qwen-VL, Gemma-3, etc.). Base checkpoints (e.g. + # google/gemma-4-12B) ship no chat template, so fall back to manually + # prepending one image placeholder token per image — the processor then + # expands each into the correct number of soft tokens (mirrors how + # examples/gemma4_unified_ort_genai.py formats image prompts). prompt_text = case.prompts[0] - if hasattr(processor, "apply_chat_template"): + if getattr(processor, "chat_template", None): content: list[dict[str, str]] = [] for img_path in case.images: content.append({"type": "image", "image": str(Path("testdata") / img_path)}) @@ -415,6 +419,8 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N prompt_text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) + elif getattr(processor, "image_token", None): + prompt_text = processor.image_token * len(case.images) + prompt_text # Process multimodal inputs through the HF processor processed = processor( @@ -739,7 +745,7 @@ def _prepare_speech_language_inputs( else: # Gemma4-style: text prompt + audio prompt_text = case.prompts[0] - if hasattr(processor, "apply_chat_template"): + if getattr(processor, "chat_template", None): content: list[dict[str, str]] = [ {"type": "audio", "audio": str(audio_path)}, {"type": "text", "text": prompt_text}, @@ -748,6 +754,10 @@ def _prepare_speech_language_inputs( prompt_text = processor.apply_chat_template( messages, tokenize=False, add_generation_prompt=True ) + elif getattr(processor, "audio_token", None): + # Base checkpoint (no chat template): manually prepend the audio + # placeholder; the processor expands it to the right token count. + prompt_text = processor.audio_token + prompt_text model_device = _get_model_device(model, device) processed = processor( text=prompt_text, diff --git a/src/mobius/_configs/_base.py b/src/mobius/_configs/_base.py index 17592fc88..7d00d6cbc 100644 --- a/src/mobius/_configs/_base.py +++ b/src/mobius/_configs/_base.py @@ -1267,6 +1267,22 @@ class Gemma4Config(VisionLanguageConfig): enable_moe_block: bool = False attention_k_eq_v: bool = False boa_token_id: int | None = None + use_bidirectional_attention: str | None = None + """Bidirectional attention mode for the text decoder. + + Mirrors HF ``Gemma4TextConfig.use_bidirectional_attention``: + - ``None``: fully causal (smaller Gemma4 models, e.g. E2B). + - ``"vision"``: text stays causal, but contiguous image-token blocks + attend bidirectionally within each block (larger models, e.g. + 12B/26B/32B). Implemented via a per-position ``block_sequence_ids`` + overlay added onto the causal mask. Audio placeholders are *not* + included (HF marks audio as token-type 3, excluded from the vision + block mask), so audio tokens keep causal attention. + - ``"all"``: HF mode where every token attends bidirectionally. Not used + by any currently supported Gemma4 model and not implemented here; the + decoder raises ``NotImplementedError`` rather than silently degrading to + causal attention (only ``None`` and ``"vision"`` are accepted). + """ @classmethod def from_transformers(cls, config, parent_config=None) -> Gemma4Config: @@ -1340,6 +1356,7 @@ def from_transformers(cls, config, parent_config=None) -> Gemma4Config: enable_moe_block=getattr(config, "enable_moe_block", False), attention_k_eq_v=getattr(config, "attention_k_eq_v", False), boa_token_id=getattr(parent_config, "boa_token_id", None), + use_bidirectional_attention=getattr(config, "use_bidirectional_attention", None), ) diff --git a/src/mobius/_configs/_sub_configs.py b/src/mobius/_configs/_sub_configs.py index bb2cbc5e3..f872d9a7f 100644 --- a/src/mobius/_configs/_sub_configs.py +++ b/src/mobius/_configs/_sub_configs.py @@ -201,6 +201,9 @@ class AudioConfig: audio_start_token_id: int | None = None audio_end_token_id: int | None = None classify_num: int | None = None + # RMSNorm epsilon for the audio encoder/embedder (may differ from the text + # decoder's rms_norm_eps). Falls back to the text value when unset. + rms_norm_eps: float | None = None # Qwen3-ASR chunked conv parameters. ``n_window`` is half the # number of mel frames per conv chunk (so chunk_size = 2 * # n_window). ``n_window_infer`` is the attention window in mel diff --git a/src/mobius/_configs/per_model/__init__.py b/src/mobius/_configs/per_model/__init__.py index 9c509ae50..fa21685ab 100644 --- a/src/mobius/_configs/per_model/__init__.py +++ b/src/mobius/_configs/per_model/__init__.py @@ -24,6 +24,8 @@ # may freely re-sort this block. from mobius._configs.per_model import ( # noqa: F401 _gemma4_audio, + _gemma4_unified_audio, + _gemma4_unified_vision, _hunyuan_vl_mot_vision, _internvl_vision, _phi4mm_audio, diff --git a/src/mobius/_configs/per_model/_gemma4_unified_audio.py b/src/mobius/_configs/per_model/_gemma4_unified_audio.py new file mode 100644 index 000000000..8cd97c2af --- /dev/null +++ b/src/mobius/_configs/per_model/_gemma4_unified_audio.py @@ -0,0 +1,38 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Gemma4-unified (gemma-4-12B) audio extractor hook. + +The ``gemma4_unified`` audio config describes an *encoder-free* embedder (no +Conformer tower). It exposes only ``audio_embed_dim`` (input feature size for +the projection) and ``rms_norm_eps``. This hook maps those onto +:class:`Gemma4AudioConfig` so +:class:`~mobius.models.gemma4._Gemma4UnifiedAudioEmbedderModel` can read them. +""" + +from __future__ import annotations + +from mobius._configs._extractors import register_audio_hook +from mobius._configs._sub_configs import Gemma4AudioConfig + +_UNIFIED_TYPES = ("gemma4_unified", "gemma4_unified_text", "gemma4_unified_audio") + + +@register_audio_hook +def _gemma4_unified_audio(config, parent_config, model_type: str, fields: dict): + composite = parent_config or config + parent_model_type = getattr(composite, "model_type", "") + if model_type not in _UNIFIED_TYPES and parent_model_type != "gemma4_unified": + return None + hf_audio = getattr(composite, "audio_config", None) + if hf_audio is None: + return None + audio_embed_dim = getattr(hf_audio, "audio_embed_dim", 640) + return { + "audio": Gemma4AudioConfig( + hidden_size=audio_embed_dim, + output_proj_dims=audio_embed_dim, + audio_token_id=getattr(composite, "audio_token_id", None), + rms_norm_eps=getattr(hf_audio, "rms_norm_eps", None), + ) + } diff --git a/src/mobius/_configs/per_model/_gemma4_unified_vision.py b/src/mobius/_configs/per_model/_gemma4_unified_vision.py new file mode 100644 index 000000000..bb336f71b --- /dev/null +++ b/src/mobius/_configs/per_model/_gemma4_unified_vision.py @@ -0,0 +1,49 @@ +# Copyright (c) Microsoft Corporation. +# Licensed under the MIT License. + +"""Gemma4-unified (gemma-4-12B) vision extractor hook. + +The ``gemma4_unified`` vision config describes an *encoder-free* embedder, not +a SigLIP tower. Its fields differ from the generic ``vision_config``: + +- ``patch_size`` / ``pooling_kernel_size`` → merged ``model_patch_size`` +- ``mm_embed_dim`` → embedder hidden size (``VisionConfig.hidden_size``) +- ``mm_posemb_size`` → factorized positional-embedding table size + (``VisionConfig.position_embedding_size``) +- ``output_proj_dims`` → projection input dim (``VisionConfig.out_hidden_size``) + +This hook maps those onto :class:`VisionConfig` so +:class:`~mobius.models.gemma4._Gemma4UnifiedVisionEmbedderModel` can read them. +""" + +from __future__ import annotations + +from mobius._configs._extractors import register_vision_hook + +_UNIFIED_TYPES = ("gemma4_unified", "gemma4_unified_text", "gemma4_unified_vision") + + +@register_vision_hook +def _gemma4_unified_vision(config, parent_config, model_type: str, fields: dict): + composite = parent_config or config + parent_model_type = getattr(composite, "model_type", "") + if model_type not in _UNIFIED_TYPES and parent_model_type != "gemma4_unified": + return None + hf_vision = getattr(composite, "vision_config", None) + if hf_vision is None: + return None + + def _get(name, default=None): + return getattr(hf_vision, name, default) + + fields.update( + model_type="gemma4_unified_vision", + hidden_size=_get("mm_embed_dim", 3840), + patch_size=_get("patch_size", 16), + pooling_kernel_size=_get("pooling_kernel_size", 3), + position_embedding_size=_get("mm_posemb_size", 1120), + out_hidden_size=_get("output_proj_dims", _get("mm_embed_dim", 3840)), + norm_eps=_get("rms_norm_eps", 1e-6), + ) + fields["image_token_id"] = getattr(composite, "image_token_id", None) + return None diff --git a/src/mobius/_optimizations.py b/src/mobius/_optimizations.py index fd9f972e8..1cfa65901 100644 --- a/src/mobius/_optimizations.py +++ b/src/mobius/_optimizations.py @@ -284,7 +284,12 @@ def _get_optimization_passes( # --- Attention fusion (decoder only) --- if model_role == "decoder" and dtype in caps.gqa_dtypes: - fuse.append(("GQAFusion", list(group_query_attention_rules()))) + fuse.append( + ( + "GQAFusion", + list(group_query_attention_rules()), + ) + ) # --- QKV packing (decoder only, gated by qkv_pack_dtypes) --- if model_role == "decoder" and dtype in caps.qkv_pack_dtypes: diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index 071a338b1..1278f97f5 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -45,6 +45,7 @@ Gemma3MultiModalModel, Gemma4CausalLMModel, Gemma4Model, + Gemma4UnifiedModel, GemmaCausalLMModel, Glm4CausalLMModel, Glm4MoECausalLMModel, @@ -398,6 +399,7 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "gemma3n": ModelRegistration(Gemma3nCausalLMModel), "gemma3n_text": ModelRegistration(Gemma3nCausalLMModel), "gemma4_text": ModelRegistration(Gemma4CausalLMModel, config_class=Gemma4Config), + "gemma4_unified_text": ModelRegistration(Gemma4CausalLMModel, config_class=Gemma4Config), "glm": ModelRegistration(GlmCausalLMModel), "glm4": ModelRegistration(Glm4CausalLMModel), "gpt_neox": ModelRegistration(GPTNeoXCausalLMModel), @@ -472,6 +474,9 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "florence2": ModelRegistration(LLaVAModel, task="vision-language"), "fuyu": ModelRegistration(LLaVAModel, task="vision-language"), "gemma4": ModelRegistration(Gemma4Model, task="gemma4", config_class=Gemma4Config), + "gemma4_unified": ModelRegistration( + Gemma4UnifiedModel, task="gemma4-unified", config_class=Gemma4Config + ), "glm4v": ModelRegistration(LLaVAModel, task="vision-language"), "glm4v_moe": ModelRegistration(LLaVAModel, task="vision-language"), "glm4v_moe_text": ModelRegistration(Glm4MoECausalLMModel), @@ -819,6 +824,8 @@ def _create_default_registry() -> ModelRegistry: "llava_next": "llava-hf/llava-v1.6-mistral-7b-hf", "mllama": "meta-llama/Llama-3.2-11B-Vision-Instruct", "gemma4": "google/gemma-4-E2B-it", + "gemma4_unified": "google/gemma-4-12B", + "gemma4_unified_text": "google/gemma-4-12B", "internvl2": "OpenGVLab/InternVL2-1B", "phi4mm": "microsoft/Phi-4-multimodal-instruct", "phi4_multimodal": "microsoft/Phi-4-multimodal-instruct", diff --git a/src/mobius/components/_attention.py b/src/mobius/components/_attention.py index 9c076d2ec..2ff62d77d 100644 --- a/src/mobius/components/_attention.py +++ b/src/mobius/components/_attention.py @@ -88,6 +88,7 @@ def _apply_attention( scale: float, softcap: float = 0.0, static_cache: StaticCacheState | None = None, + is_causal: int = 1, ) -> tuple[ir.Value, ir.Value, ir.Value]: """Apply the ONNX Attention op with internal or static KV cache. @@ -103,10 +104,20 @@ def _apply_attention( Also uses ``is_causal=1``. Returns ``(attn_output, updated_key_cache, updated_value_cache)``. + Args: + is_causal: Whether the Attention op applies its built-in causal + mask (default ``1``). Set to ``0`` when ``attn_mask`` already + bakes the FULL mask (causal + sliding + padding, and any + bidirectional unmasking such as Gemma4's vision-block overlay) + into a float additive bias. Leaving ``is_causal=1`` in that + case would re-apply causality and cancel any future-position + unmasking encoded in the bias. + Note: - Both paths set ``is_causal=1`` on the Attention op, which enables - built-in causal masking. This means ``attn_mask`` should encode - only padding information (as a bool mask), not causality. + Both paths default to ``is_causal=1`` on the Attention op, which + enables built-in causal masking. This means ``attn_mask`` should + encode only padding information (as a bool mask), not causality, + unless ``is_causal=0`` is passed explicitly. Note: ``nonpad_kv_seqlen`` (input #6) is only valid in static cache mode @@ -167,7 +178,7 @@ def _apply_attention( kv_num_heads=num_key_value_heads, scale=scale, softcap=softcap, - is_causal=1, + is_causal=is_causal, _outputs=3, ) return attn_output, updated_k, updated_v @@ -191,7 +202,7 @@ def _apply_attention( kv_num_heads=num_key_value_heads, scale=scale, softcap=softcap, - is_causal=1, + is_causal=is_causal, _outputs=3, ) return attn_output, present_key, present_value diff --git a/src/mobius/components/_common.py b/src/mobius/components/_common.py index 2d4836f56..4867df77c 100644 --- a/src/mobius/components/_common.py +++ b/src/mobius/components/_common.py @@ -161,6 +161,7 @@ def create_attention_bias( attention_mask, sliding_window: int | None = None, dtype: ir.DataType = ir.DataType.FLOAT, + block_sequence_ids=None, ): """Create causal attention bias for use in attention mechanisms. @@ -172,6 +173,18 @@ def create_attention_bias( dtype: Data type for the bias tensor. The masked value uses the minimum representable value for this dtype (e.g. -65504 for float16, -3.4e38 for float32). + block_sequence_ids: Optional INT tensor of shape + (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: + two positions in the same block (same id ``>= 0``) may attend to + each other regardless of causal order. 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 + ``is_causal=0`` to avoid re-applying the causal constraint and + cancelling the bidirectional unmasking. Returns: Attention bias tensor of shape (batch_size, 1, query_length, total_length). @@ -210,6 +223,34 @@ def create_attention_bias( within_window = op.Less(dist, sliding_window) full_mask = 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). + # + # 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. + 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]. + pads = op.Concat(zero_1d, pad_width, zero_1d, zero_1d, axis=0) + kv_group_2d = op.Pad( + block_sequence_ids, + pads, + op.Constant(value_int=-1), + ) # (B, total_length) + kv_group = op.Unsqueeze(kv_group_2d, [1]) # (B, 1, total_length) + # same_block = (q_group == kv_group) AND (q_group >= 0) + same_block = op.And( + op.Equal(q_group, kv_group), + op.GreaterOrEqual(q_group, op.Constant(value_int=0)), + ) + full_mask = op.Or(full_mask, same_block) + # Combine with attention_mask attn_mask_bool = op.Cast(op.Unsqueeze(attention_mask, [1]), to=ir.DataType.BOOL) full_mask = op.And(attn_mask_bool, full_mask) diff --git a/src/mobius/components/_common_test.py b/src/mobius/components/_common_test.py index 3ebc1bf92..09e7f17b1 100644 --- a/src/mobius/components/_common_test.py +++ b/src/mobius/components/_common_test.py @@ -5,9 +5,11 @@ from __future__ import annotations +import numpy as np import onnx_ir as ir from mobius._testing import count_op_type, create_test_builder, create_test_input +from mobius._testing.ort_inference import OnnxModelSession from mobius.components._common import ( Embedding, Linear, @@ -117,6 +119,100 @@ def test_query_length_from_input_ids_not_attention_mask(self): ) +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 + called with ``is_causal=0``. We build the graph, run it via ORT, and + compare the attended pattern (bias == 0) against a numpy reference. + """ + + @staticmethod + def _build(sliding): + b, op, g = create_test_builder() + input_ids = create_test_input(b, "input_ids", [1, "S"], dtype=ir.DataType.INT64) + attn = create_test_input(b, "attention_mask", [1, "T"], dtype=ir.DataType.INT64) + bsid = create_test_input(b, "block_sequence_ids", [1, "S"], dtype=ir.DataType.INT64) + bias = create_attention_bias( + op, + input_ids, + attn, + sliding_window=sliding, + dtype=ir.DataType.FLOAT, + block_sequence_ids=bsid, + ) + bias.name = "bias" + g.outputs.append(bias) + return ir.Model(g, ir_version=10) + + @staticmethod + def _ref(block_ids, attn, sliding): + cumsum = np.cumsum(attn) + qi = cumsum[:, None] + ki = cumsum[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)) + return m & (np.array(attn)[None, :].astype(bool)) + + def _run(self, sliding, block_ids, attn): + sess = OnnxModelSession(self._build(sliding), device="cpu") + block_ids = np.array([block_ids], dtype=np.int64) + attn = np.array([attn], dtype=np.int64) + out = sess.run( + { + "input_ids": np.zeros_like(block_ids), + "attention_mask": attn, + "block_sequence_ids": block_ids, + } + )["bias"] + attended = out[0, 0] > -1.0 + expected = self._ref(block_ids[0], attn[0], sliding) + return attended, expected + + def test_multi_block_full_attention(self): + # Two vision blocks (pos 1-2 and 4-5) separated by text. + attended, expected = self._run(None, [-1, 0, 0, -1, 1, 1, -1, -1], [1] * 8) + assert np.array_equal(attended, expected) + # A vision token attends to a LATER token in the same block (bidirectional). + assert attended[1, 2] + # But text stays causal: position 3 cannot see position 4. + 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). + 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 + + 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] + ) + assert np.array_equal(attended, expected) + assert not attended[:, 6:].any() # nothing attends to padding + + def test_decode_single_query_is_causal(self): + # Decode step: q_len=1 (new text token, group -1), kv total length 8. + sess = OnnxModelSession(self._build(None), device="cpu") + out = sess.run( + { + "input_ids": np.zeros((1, 1), dtype=np.int64), + "attention_mask": np.ones((1, 8), dtype=np.int64), + "block_sequence_ids": np.array([[-1]], dtype=np.int64), + } + )["bias"] + assert out.shape == (1, 1, 1, 8) + # Text decode token attends to all past positions (pure causal row). + assert bool((out[0, 0, 0] > -1.0).all()) + + class TestCreatePaddingMask: def test_creates_bool_mask_with_2d_input_ids(self): """Standard path: input_ids is 2D [batch, q_len].""" diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 9d7dbeae5..5b6f49329 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -75,6 +75,12 @@ "gemma2": "gemma", "gemma4": "gemma4", "gemma4_text": "gemma4_text", + # gemma-4-12B "unified" (encoder-free) variant reuses the gemma4 ORT GenAI + # pipelines: the multimodal package (decoder taking inputs_embeds + vision + # embedder + embedding fusion) maps to "gemma4"; the standalone text + # backbone maps to "gemma4_text". + "gemma4_unified": "gemma4", + "gemma4_unified_text": "gemma4_text", "mistral": "mistral", "mistral3": "mistral3", # HunYuan-V1 dense / Hy-MT1.5 — generic decoder LLM type accepted by @@ -88,7 +94,18 @@ "qwen3_5_vl": "qwen2_5_vl", } -_GEMMA4_MODEL_TYPES = frozenset({"gemma4", "gemma4_text"}) +_GEMMA4_MODEL_TYPES = frozenset( + {"gemma4", "gemma4_text", "gemma4_unified", "gemma4_unified_text"} +) +# Encoder-free gemma-4-12B "unified" variants. Their image/audio inputs are raw +# merged pixel patches (48px, 6912-dim) / raw waveform frames (640-dim), NOT the +# SigLIP 16px / 128-dim log-mel contract that the ort-extensions +# ``Gemma4ImageTransform`` / ``Gemma4LogMel`` ops implement. There is no +# genai-native transform for the unified contract, so we deliberately do NOT +# emit image_processor.json / audio_processor.json for these models — callers +# must preprocess with the HuggingFace processor and feed tensors via +# ``Generator.set_inputs`` (see examples/gemma4_unified_ort_genai.py). +_GEMMA4_UNIFIED_MODEL_TYPES = frozenset({"gemma4_unified", "gemma4_unified_text"}) _PIXTRAL_MODEL_TYPES = frozenset({"mistral3"}) _QWEN_VL_MODEL_TYPES = frozenset( { @@ -384,6 +401,9 @@ def _write_vision_processor_config( - **Gemma4** (``gemma4``, ``gemma4_text``): Writes ``image_processor.json`` with a ``DecodeImage → Gemma4ImageTransform`` pipeline. + - **Gemma4 unified** (``gemma4_unified*``): Returns ``None`` — the + encoder-free model has no matching ort-extensions transform; callers feed + HF-preprocessed pixel_values via ``Generator.set_inputs``. - **Pixtral / Mistral3**: Writes ``processor_config.json`` with a 7-step pipeline (DecodeImage → ConvertRGB → Resize → Rescale → Normalize → Permute3D → PixtralImageSizes). @@ -397,6 +417,17 @@ def _write_vision_processor_config( return None model_type = getattr(config, "model_type", "") + if model_type in _GEMMA4_UNIFIED_MODEL_TYPES: + # Encoder-free unified model: no ort-extensions transform matches its + # raw merged-patch contract. Emit no image_processor.json; callers feed + # HF-preprocessed pixel_values via Generator.set_inputs. + logger.info( + "Skipping image_processor.json for encoder-free %s " + "(no native ort-extensions transform; use HF processor + set_inputs)", + model_type, + ) + return None + vision_model_type = getattr(vision, "model_type", None) is_pixtral = vision_model_type == "pixtral" or model_type in _PIXTRAL_MODEL_TYPES @@ -573,6 +604,17 @@ def _write_audio_processor_config( model_type = getattr(config, "model_type", "") + if model_type in _GEMMA4_UNIFIED_MODEL_TYPES: + # Encoder-free unified model: raw 640-dim waveform frames, not the + # 128-dim log-mel Gemma4LogMel contract. Emit no audio_processor.json; + # callers feed HF-preprocessed input_features via Generator.set_inputs. + logger.info( + "Skipping audio_processor.json for encoder-free %s " + "(no native ort-extensions transform; use HF processor + set_inputs)", + model_type, + ) + return None + if model_type in _GEMMA4_MODEL_TYPES: # Gemma4 USM-style 128-dim log-mel spectrogram. # OrtxCreateSpeechFeatureExtractor requires the feature_extraction.sequence format. diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 7a7966f1e..345003351 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -85,6 +85,16 @@ def test_phi4mm_model_types(self): assert _resolve_ort_genai_model_type("phi4_multimodal") == "phi4mm" assert _resolve_ort_genai_model_type("phi") == "phi" + def test_gemma4_unified_model_types(self): + # The gemma-4-12B unified checkpoint (model_type "gemma4_unified") + # reuses the multimodal "gemma4" ORT GenAI pipeline; its standalone + # text decoder ("gemma4_unified_text") maps to "gemma4_text". + assert _resolve_ort_genai_model_type("gemma4_unified") == "gemma4" + assert _resolve_ort_genai_model_type("gemma4_unified_text") == "gemma4_text" + # Released gemma4 mappings remain unchanged. + assert _resolve_ort_genai_model_type("gemma4") == "gemma4" + assert _resolve_ort_genai_model_type("gemma4_text") == "gemma4_text" + class TestWriteProcessorConfig: def test_no_vision_returns_none(self, tmp_path): @@ -128,6 +138,15 @@ def test_writes_transform_pipeline(self, tmp_path): assert len(norm_attrs["mean"]) == 3 assert len(norm_attrs["std"]) == 3 + def test_gemma4_unified_skips_image_processor(self, tmp_path): + """Encoder-free gemma4_unified has no native transform: no image_processor.json.""" + vision = mock.MagicMock() + vision.model_type = None + config = mock.MagicMock() + config.vision = vision + config.model_type = "gemma4_unified" + assert _write_vision_processor_config(config, str(tmp_path)) is None + def test_pixtral_vision_config(self, tmp_path): """Generates pixtral-specific processor config with 7 transforms.""" vision = mock.MagicMock() @@ -268,6 +287,13 @@ def test_audio_non_gemma4_returns_none(self, tmp_path): config.model_type = "whisper" assert _write_audio_processor_config(config, str(tmp_path)) is None + def test_audio_gemma4_unified_skips_audio_processor(self, tmp_path): + """Encoder-free gemma4_unified has no native transform: no audio_processor.json.""" + config = mock.MagicMock() + config.audio = mock.MagicMock() + config.model_type = "gemma4_unified" + assert _write_audio_processor_config(config, str(tmp_path)) is None + def test_audio_gemma4_writes_feature_extraction_json(self, tmp_path): config = mock.MagicMock() config.audio = mock.MagicMock() diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index 75efcee39..b091e7a2e 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -46,6 +46,7 @@ "Gemma3nCausalLMModel", "Gemma4CausalLMModel", "Gemma4Model", + "Gemma4UnifiedModel", "GemmaCausalLMModel", "Glm4CausalLMModel", "Glm4MoECausalLMModel", @@ -165,6 +166,7 @@ from mobius.models.gemma4 import ( Gemma4CausalLMModel, Gemma4Model, + Gemma4UnifiedModel, ) from mobius.models.glm import Glm4CausalLMModel, GlmCausalLMModel from mobius.models.gpt2 import GPT2CausalLMModel diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py index 86a954d7d..7bc7d70f6 100644 --- a/src/mobius/models/gemma4.py +++ b/src/mobius/models/gemma4.py @@ -38,6 +38,8 @@ from mobius.components import ( MLP, ClippableLinear, + Embedding, + LayerNorm, Linear, RMSNorm, create_attention_bias, @@ -53,6 +55,21 @@ from mobius.components._attention import GQAContext +def _dtype_safe_compress( + op: OpBuilder, data: ir.Value, condition: ir.Value, *, axis: int +) -> ir.Value: + """Row-select ``data`` by ``condition`` in a dtype that ORT supports. + + ORT does not register a ``Compress`` kernel for ``bfloat16``, so a bf16 + package would fail to load. Run the selection in float32 and cast the + result back to ``data``'s dtype. The float16/bfloat16 round-trip through + float32 is lossless, so this is exact for every supported build dtype. + """ + data_f32 = op.Cast(data, to=ir.DataType.FLOAT) + selected = op.Compress(data_f32, condition, axis=axis) + return op.CastLike(selected, data) + + # --------------------------------------------------------------------------- # Shared weight preprocessing helpers # --------------------------------------------------------------------------- @@ -707,6 +724,7 @@ def forward( position_embeddings: tuple | None = None, shared_kv_states: dict | None = None, past_key_value: tuple | None = None, + is_causal: int = 1, ): from mobius.components._attention import ( GQAContext, @@ -806,6 +824,7 @@ def forward( num_key_value_heads=self.num_key_value_heads, scale=self.scaling, softcap=self.softcap, + is_causal=is_causal, ) elif use_gqa: # GQA path: emit com.microsoft.GroupQueryAttention directly. @@ -928,6 +947,7 @@ def forward( num_key_value_heads=self.num_key_value_heads, scale=self.scaling, softcap=self.softcap, + is_causal=is_causal, ) # Source layers store K,V for downstream KV-shared layers. @@ -1136,6 +1156,7 @@ def forward( shared_kv_states: dict, per_layer_input: ir.Value | None, past_key_value: tuple | None, + is_causal: int = 1, ): # Attention block: pre-norm -> attn -> post-norm -> residual residual = hidden_states @@ -1147,6 +1168,7 @@ def forward( position_embeddings=position_embeddings, shared_kv_states=shared_kv_states, past_key_value=past_key_value, + is_causal=is_causal, ) hidden_states = self.post_attention_layernorm(op, attn_output) hidden_states = op.Add(residual, hidden_states) @@ -1367,6 +1389,62 @@ def _dispatch_moe_fallback( # --------------------------------------------------------------------------- +def _compute_block_sequence_ids( + op: OpBuilder, + input_ids: ir.Value, + *, + image_token_id: int, +) -> ir.Value: + """Compute Gemma4 ``block_sequence_ids`` [B, S] from ``input_ids``. + + Mirrors HuggingFace ``get_block_sequence_ids_for_mask``: each contiguous + run of image placeholder tokens gets a unique, monotonically increasing + block id (``>= 0``); every other position (text **and audio**) gets ``-1``. + Tokens within the same block may attend to each other bidirectionally. + + Only image tokens form blocks. HF derives ``is_vision`` from + ``mm_token_type_ids`` as ``(== 1) | (== 2)`` (image or video); audio is + token-type ``3`` and is deliberately excluded, so audio placeholders keep + plain causal attention. gemma4_unified has no video modality, so this + reduces to image tokens alone. + + A new block starts only on a non-image -> image transition. + + Returns an INT64 tensor of shape ``[B, S]``. + """ + # is_vision [B, S] BOOL: token is an image placeholder. Audio tokens are + # intentionally NOT included (HF block mask covers image/video only). + is_vision = op.Equal(input_ids, op.Constant(value_int=image_token_id)) + + # is_prev_vision: is_vision shifted right by one along the sequence axis, + # with position 0 forced to False. Implemented without ConstantOfShape + # (which blocks ONNX shape inference): left-pad the int mask with one + # zero column, then drop the last column. + is_vision_int = op.Cast(is_vision, to=ir.DataType.INT64) + padded = op.Pad( + is_vision_int, + op.Constant(value_ints=[0, 1, 0, 0]), # prepend 1 col on axis 1 + op.Constant(value_int=0), + ) # [B, S + 1] + prev_int = op.Slice( + padded, + op.Constant(value_ints=[0]), + op.Constant(value_ints=[-1]), + op.Constant(value_ints=[1]), + ) # [B, S] + is_prev_vision = op.Cast(prev_int, to=ir.DataType.BOOL) + + # new_vision_starts = is_vision AND NOT is_prev_vision + new_starts = op.And(is_vision, op.Not(is_prev_vision)) + # vision_group_ids = cumsum(new_starts) - 1 (along sequence axis) + group_ids = op.Sub( + op.CumSum(op.Cast(new_starts, to=ir.DataType.INT64), op.Constant(value_int=1)), + op.Constant(value_int=1), + ) + # block_sequence_ids = where(is_vision, group_ids, -1) + return op.Where(is_vision, group_ids, op.Constant(value_int=-1)) + + class Gemma4TextModel(nn.Module): """Gemma4 text transformer with hybrid local/global attention. @@ -1402,6 +1480,21 @@ 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 + # "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. + if config.use_bidirectional_attention not in (None, "vision"): + raise NotImplementedError( + "Gemma4 use_bidirectional_attention=" + f"{config.use_bidirectional_attention!r} is not supported; only " + "None (fully causal) and 'vision' (image-block bidirectional) " + "are implemented." + ) + self._use_bidirectional_attention = config.use_bidirectional_attention # Local (sliding window) config — full rotation, local rope_theta local_config = dataclasses.replace( @@ -1438,6 +1531,12 @@ def __init__(self, config: Gemma4Config): self._per_layer_dim = getattr(config, "hidden_size_per_layer_input", 0) self._hidden_size = config.hidden_size self._image_token_id: int = config.image_token_id or 0 + # The vision-block overlay keys on image_token_id. A 0/None id means the + # model has no image-placeholder token (e.g. the text-only backbone + # split), so no image tokens can appear — keep pure causal attention + # (GQA-eligible) and never build the overlay. This also avoids the + # footgun of a 0 fallback marking real token-0 positions as vision. + self._has_image_token: bool = bool(config.image_token_id) self._audio_token_id: int | None = ( config.audio.audio_token_id if config.audio is not None else None ) @@ -1513,6 +1612,7 @@ def forward( past_key_values: list | None = None, inputs_embeds: ir.Value | None = None, per_layer_inputs: ir.Value | None = None, + block_sequence_ids: ir.Value | None = None, ) -> tuple[ir.Value, list]: if inputs_embeds is not None: hidden_states = inputs_embeds @@ -1547,11 +1647,52 @@ def forward( caps = ep_capabilities() 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. + # + # ``block_sequence_ids`` [B, S] identifies contiguous image token + # blocks. It is derived from ``input_ids`` (image token spans only; + # audio keeps causal attention, matching HF). In the multimodal + # 3/4-model split the decoder receives ``input_ids`` alongside + # ``inputs_embeds`` and computes the overlay here, so it does not need a + # separate cross-model tensor (onnxruntime-genai forwards ``input_ids`` + # 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. + 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( + op, + input_ids, + image_token_id=self._image_token_id, + ) + use_block_overlay = bidirectional and block_sequence_ids is not None + use_gqa = ( attention_mask is not None and dtype in caps.gqa_dtypes 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 @@ -1615,12 +1756,14 @@ def forward( attention_mask=attention_mask, sliding_window=self.sliding_window, dtype=self._dtype, + block_sequence_ids=block_sequence_ids if use_block_overlay else None, ), "full_attention": create_attention_bias( op, 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 @@ -1666,6 +1809,7 @@ def forward( shared_kv_states=shared_kv_states, per_layer_input=per_layer_input, past_key_value=past_kv, + is_causal=attn_is_causal, ) # KV-shared layers borrow K,V from source layers — exclude from # present_key_values so the output has exactly num_kv_layers entries. @@ -1776,15 +1920,22 @@ def forward( position_ids: ir.Value, per_layer_inputs: ir.Value | None = None, past_key_values: list | None = None, + block_sequence_ids: ir.Value | None = None, + input_ids: ir.Value | None = None, ) -> tuple[ir.Value, list]: + # ``input_ids`` is forwarded alongside ``inputs_embeds`` so the text + # model can derive the bidirectional vision-block overlay internally + # (see Gemma4TextModel.forward). ``inputs_embeds`` still takes + # precedence for the actual token embeddings. hidden_states, present_key_values = self.model( op, - input_ids=None, + input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, per_layer_inputs=per_layer_inputs, + block_sequence_ids=block_sequence_ids, ) logits = self.lm_head(op, hidden_states) # Gemma4 applies final logit soft-capping: logit_cap * tanh(x / logit_cap) @@ -1982,7 +2133,17 @@ def forward( input_ids: ir.Value, image_features: ir.Value, audio_features: ir.Value | None = None, - ) -> ir.Value | tuple[ir.Value, ir.Value]: + ) -> dict[str, ir.Value]: + """Return a dict of named embedding outputs. + + Always contains ``inputs_embeds``. Contains ``per_layer_inputs`` when + ``hidden_size_per_layer_input > 0``. + + The vision-block bidirectional attention overlay is NOT emitted here: + the decoder derives it from ``input_ids`` directly (see + ``Gemma4TextModel.forward``), which avoids a cross-model tensor that + onnxruntime-genai cannot forward between sub-models. + """ # [B, S] → [B, S, hidden] hidden = self.embed_tokens(op, input_ids) @@ -2001,8 +2162,10 @@ def forward( op, hidden, input_ids, self.audio_token_id, audio_features ) + outputs: dict[str, ir.Value] = {"inputs_embeds": hidden} + if not self._per_layer_dim: - return hidden + return outputs # Compute per-layer input embeddings (moved from the decoder). # 1. Project hidden states → [B, S, L*D] and scale by hidden_size**-0.5 @@ -2046,8 +2209,9 @@ def forward( combined, op.Constant(value_ints=[0, 0, self._num_layers * self._per_layer_dim]), ) + outputs["per_layer_inputs"] = per_layer_inputs - return hidden, per_layer_inputs + return outputs def preprocess_weights( self, state_dict: dict[str, torch.Tensor] @@ -2134,6 +2298,256 @@ def forward( return self.projector(op, audio_features), downsampled_mask +# --------------------------------------------------------------------------- +# gemma4_unified (gemma-4-12B) encoder-free vision / audio embedders +# --------------------------------------------------------------------------- + + +class _F32Linear(Linear): + """Linear that computes its MatMul in float32 regardless of model dtype. + + Used by the gemma4_unified vision embedder's ``patch_dense`` projection + **only when the model dtype is float16**, whose output magnitude (~77000) + exceeds the float16 range (65504). The weights are stored in the model dtype; + activations and weights are upcast to float32 for the MatMul so the result + does not overflow to +inf. The output stays float32 (the following + ``_F32LayerNorm`` normalizes it back into a float16-safe range). bfloat16 and + float32 models have the range natively and use a plain :class:`Linear`. + """ + + def forward(self, op: OpBuilder, x: ir.Value) -> ir.Value: + w_t = op.Cast( + op.Transpose(self.weight, perm=[1, 0]), to=ir.DataType.FLOAT + ) # [in_features, out_features] + result = op.MatMul(op.Cast(x, to=ir.DataType.FLOAT), w_t) + if self.bias is not None: + result = op.Add(result, op.Cast(self.bias, to=ir.DataType.FLOAT)) + return result # float32 + + +class _F32LayerNorm(LayerNorm): + """LayerNorm that computes in float32 and returns a float32 output. + + Pairs with :class:`_F32Linear` in the gemma4_unified vision embedder, **only + for float16 models**, so the large (out-of-float16-range) ``patch_dense`` + output is normalized in float32 before being cast back to the model dtype. + bfloat16 and float32 models use a plain :class:`LayerNorm`. + """ + + def forward(self, op: OpBuilder, hidden_states: ir.Value) -> ir.Value: + return op.LayerNormalization( + op.Cast(hidden_states, to=ir.DataType.FLOAT), + op.Cast(self.weight, to=ir.DataType.FLOAT), + op.Cast(self.bias, to=ir.DataType.FLOAT), + epsilon=self.eps, + axis=-1, + ) # float32 + + +class _Gemma4UnifiedVisionEmbedderModel(nn.Module): + """Encoder-free vision embedder for ``gemma4_unified`` (gemma-4-12B). + + Unlike gemma4's SigLIP tower, the unified model has **no vision encoder**. + Raw merged pixel patches are projected directly into language-model space. + + Replicates HF ``Gemma4UnifiedVisionEmbedder``: + + patch_ln1 (LayerNorm, patch_dim) + → patch_dense (Linear patch_dim → mm_embed_dim) + → patch_ln2 (LayerNorm, mm_embed_dim) + → + factorized 2D positional embedding + → pos_norm (LayerNorm, mm_embed_dim) + → embedding_pre_projection_norm (scale-free RMSNorm) + → embedding_projection (Linear mm_embed_dim → text_hidden) + + ``patch_dim = (patch_size * pooling_kernel_size)^2 * 3`` (48*48*3 = 6912). + + Inputs: + - ``pixel_values [B, N, patch_dim]``: raw merged pixel patches. + - ``pixel_position_ids [B, N, 2]``: integer (x, y) patch coordinates; + ``(-1, -1)`` marks padding patch slots. + + Output: + - ``image_features [num_valid_patches, text_hidden_size]``: padding + patches (position == -1) are stripped so the output rows align 1:1 with + image placeholder tokens in the text sequence (matches HF, which selects + ``vision_outputs[~padding_mask]``). + """ + + def __init__(self, config: Gemma4Config): + super().__init__() + vc = config.vision # VisionConfig populated by the gemma4_unified hook + patch_size = (vc.patch_size if vc else None) or 16 + pooling = (vc.pooling_kernel_size if vc else None) or 3 + model_patch_size = patch_size * pooling + patch_dim = 3 * model_patch_size * model_patch_size + mm_embed_dim = (vc.hidden_size if vc else None) or config.hidden_size + posemb_size = (vc.position_embedding_size if vc else None) or 1120 + out_proj_dim = (vc.out_hidden_size if vc else None) or mm_embed_dim + eps = (vc.norm_eps if vc else None) or config.rms_norm_eps or 1e-6 + self._text_hidden_size = config.hidden_size + + self.patch_ln1 = LayerNorm(patch_dim, eps=eps) + # patch_dense produces activations whose magnitude (~77000, measured) is + # outside the float16 range (max 65504); HF runs this embedder in bfloat16 + # (max ~3.4e38). Only float16 actually overflows, so we upcast the dense + # projection + the following LayerNorm to float32 *only* when the model + # dtype is float16. bfloat16 and float32 have the range natively and keep + # their dtype (matching HF for bfloat16). See _F32Linear / _F32LayerNorm. + if config.dtype == ir.DataType.FLOAT16: + self.patch_dense = _F32Linear(patch_dim, mm_embed_dim, bias=True) + self.patch_ln2 = _F32LayerNorm(mm_embed_dim, eps=eps) + else: + self.patch_dense = Linear(patch_dim, mm_embed_dim, bias=True) + self.patch_ln2 = LayerNorm(mm_embed_dim, eps=eps) + # Factorized 2D positional embedding: HF stores a single + # [posemb_size, 2, mm_embed_dim] table looked up per axis. We split it + # into two [posemb_size, mm_embed_dim] tables (x and y) in + # preprocess_weights so each axis is a plain Gather. + self.pos_emb_x = Embedding(posemb_size, mm_embed_dim) + self.pos_emb_y = Embedding(posemb_size, mm_embed_dim) + self.pos_norm = LayerNorm(mm_embed_dim, eps=eps) + # Scale-free RMSNorm before the projection (HF + # embed_vision.multimodal_embedder.embedding_pre_projection_norm). + # The projection consumes the post-position-norm activations, whose + # last dim is mm_embed_dim (HF embedding_projection: mm_embed_dim → + # text_hidden). out_proj_dim is retained only as a sanity check. + assert out_proj_dim == mm_embed_dim, ( + "gemma4_unified vision projector expects output_proj_dims == " + f"mm_embed_dim, got {out_proj_dim} != {mm_embed_dim}" + ) + self.projector_norm = _Gemma4ScaleFreeRMSNorm(mm_embed_dim, eps=eps) + # HF embed_vision.multimodal_embedder.embedding_projection (no bias). + self.projector = Linear(mm_embed_dim, config.hidden_size, bias=False) + + def forward( + self, + op: OpBuilder, + pixel_values: ir.Value, + pixel_position_ids: ir.Value, + ) -> ir.Value: + # Patch embedding: LN → Dense → LN. [B, N, patch_dim] → [B, N, mm_embed_dim] + # For float16 models patch_dense + patch_ln2 run in float32 (see __init__ + # and _F32Linear / _F32LayerNorm): the dense projection produces activations + # whose magnitude exceeds the float16 range (measured absmax ~77000 > 65504; + # HF runs this embedder in bfloat16), so an f16 intermediate would overflow + # to +inf and patch_ln2 would emit NaN. patch_ln2 normalizes the result back + # into a float16-safe range. For bfloat16 / float32 the projection stays in + # the model dtype and CastLike below is a no-op. + h = self.patch_ln1(op, pixel_values) + h = self.patch_dense(op, h) # f16 model: f16 in → f32 out; else native dtype + h = self.patch_ln2(op, h) # f16 model: f32 in → f32 out (normalized) + h = op.CastLike(h, pixel_values) # back to the model dtype (no-op unless upcast) + + # Factorized positional embedding. Split (x, y) coords on the last axis. + x_ids = op.Gather(pixel_position_ids, op.Constant(value_int=0), axis=-1) # [B, N] + y_ids = op.Gather(pixel_position_ids, op.Constant(value_int=1), axis=-1) # [B, N] + # Padding patches carry -1: clamp to a valid index for the Gather, then + # zero their contribution via the validity mask. + zero = op.Constant(value_int=0) + clamped_x = op.Max(x_ids, zero) + clamped_y = op.Max(y_ids, zero) + neg_one = op.Constant(value_int=-1) + valid_x = op.Unsqueeze( + op.CastLike(op.Not(op.Equal(x_ids, neg_one)), h), [-1] + ) # [B, N, 1] + valid_y = op.Unsqueeze(op.CastLike(op.Not(op.Equal(y_ids, neg_one)), h), [-1]) + pos_x = op.Mul(self.pos_emb_x(op, clamped_x), valid_x) # [B, N, mm_embed_dim] + pos_y = op.Mul(self.pos_emb_y(op, clamped_y), valid_y) + h = op.Add(h, op.Add(pos_x, pos_y)) + h = self.pos_norm(op, h) + + # Scale-free RMSNorm → projection to text hidden size. + h = self.projector_norm(op, h) + h = self.projector(op, h) # [B, N, text_hidden] + + # Strip padding patches so output rows align 1:1 with placeholder tokens. + h_flat = op.Reshape(h, op.Constant(value_ints=[-1, self._text_hidden_size])) + keep = op.Reshape( + op.Not(op.Equal(x_ids, neg_one)), op.Constant(value_ints=[-1]) + ) # [B*N] BOOL + return _dtype_safe_compress(op, h_flat, keep, axis=0) # [num_valid, text_hidden] + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + renamed: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith("vision_embedder.pos_embedding"): + # HF [posemb_size, 2, mm_embed_dim] → two [posemb_size, mm_embed_dim] + renamed["pos_emb_x.weight"] = value[:, 0, :].contiguous() + renamed["pos_emb_y.weight"] = value[:, 1, :].contiguous() + elif key.startswith("vision_embedder."): + renamed[key[len("vision_embedder.") :]] = value + elif key.startswith("embed_vision.embedding_projection."): + renamed["projector." + key[len("embed_vision.embedding_projection.") :]] = ( + value + ) + # embed_vision.*.embedding_pre_projection_norm: scale-free, no weight. + return renamed + + +class _Gemma4UnifiedAudioEmbedderModel(nn.Module): + """Encoder-free audio embedder for ``gemma4_unified`` (gemma-4-12B). + + The unified model has **no Conformer audio tower**. Raw waveform-frame + features are projected directly into language-model space. + + Replicates HF ``Gemma4UnifiedMultimodalEmbedder`` (the ``embed_audio`` + branch): + + embedding_pre_projection_norm (scale-free RMSNorm, audio_embed_dim) + → embedding_projection (Linear audio_embed_dim → text_hidden) + + Inputs: + - ``input_features [B, T, audio_embed_dim]``: raw waveform-frame features. + - ``input_features_mask [B, T]``: BOOL mask, ``True`` for valid frames. + + Output: + - ``audio_features [num_valid_frames, text_hidden_size]``: padding frames + are stripped so output rows align 1:1 with audio placeholder tokens + (matches HF, which selects ``audio_features[audio_mask]``). + """ + + def __init__(self, config: Gemma4Config): + super().__init__() + ac = config.audio # Gemma4AudioConfig populated by the gemma4_unified hook + audio_embed_dim = (ac.hidden_size if ac else None) or 640 + # Prefer the audio config's own eps; fall back to the text decoder's. + eps = (ac.rms_norm_eps if ac else None) or config.rms_norm_eps or 1e-6 + self._text_hidden_size = config.hidden_size + # HF embed_audio.embedding_pre_projection_norm (scale-free RMSNorm). + self.projector_norm = _Gemma4ScaleFreeRMSNorm(audio_embed_dim, eps=eps) + # HF embed_audio.embedding_projection (no bias). + self.projector = Linear(audio_embed_dim, config.hidden_size, bias=False) + + def forward( + self, + op: OpBuilder, + input_features: ir.Value, + input_features_mask: ir.Value | None = None, + ) -> tuple[ir.Value, ir.Value | None]: + # [B, T, audio_embed_dim] → scale-free RMSNorm → [B, T, text_hidden] + h = self.projector_norm(op, input_features) + h = self.projector(op, h) + if input_features_mask is None: + return h, None + # Strip padding frames so output rows align 1:1 with placeholder tokens. + h_flat = op.Reshape(h, op.Constant(value_ints=[-1, self._text_hidden_size])) + keep = op.Reshape(input_features_mask, op.Constant(value_ints=[-1])) # [B*T] + return _dtype_safe_compress(op, h_flat, keep, axis=0), None # [num_valid, text_hidden] + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + renamed: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith("embed_audio.embedding_projection."): + renamed["projector." + key[len("embed_audio.embedding_projection.") :]] = value + # embed_audio.embedding_pre_projection_norm: scale-free, no weight. + return renamed + + # --------------------------------------------------------------------------- # Gemma4Model — unified vision-language (+ optional audio) model # --------------------------------------------------------------------------- @@ -2308,3 +2722,119 @@ def preprocess_weights( _remap_moe_expert_weights(renamed, self.config) return renamed + + +# --------------------------------------------------------------------------- +# Gemma4UnifiedModel — gemma-4-12B encoder-free multimodal model +# --------------------------------------------------------------------------- + + +class Gemma4UnifiedModel(nn.Module): + """Unified gemma-4-12B (``gemma4_unified``) multimodal model. + + Encoder-free counterpart to :class:`Gemma4Model`: it shares the gemma4 + text decoder and the multimodal-fusion embedding sub-model, but replaces + the SigLIP vision tower and Conformer audio tower with the lightweight + encoder-free embedders (:class:`_Gemma4UnifiedVisionEmbedderModel` and + :class:`_Gemma4UnifiedAudioEmbedderModel`). + + Builds a 3- or 4-model package (built by :class:`~mobius.tasks.Gemma4UnifiedTask`): + + Always produced: + - ``decoder``: gemma4 text decoder taking ``inputs_embeds`` (and + ``input_ids`` for the vision-block bidirectional mask, which it derives + internally; dual head_dim, k_eq_v) + - ``vision_encoder``: raw-patch vision embedder + - ``embedding``: scaled word embedding + multimodal feature fusion + + Added when ``config.audio is not None``: + - ``audio_encoder``: raw-frame audio embedder + + Registered as ``gemma4_unified``. + """ + + default_task: str = "gemma4-unified" + category: str = "Multimodal" + + def __init__(self, config: Gemma4Config): + super().__init__() + self.config = config + self.decoder = _Gemma4DecoderModel(config) + self.vision_encoder = _Gemma4UnifiedVisionEmbedderModel(config) + self.embedding = Gemma4EmbeddingModel(config) + self.audio_encoder: _Gemma4UnifiedAudioEmbedderModel | None = ( + _Gemma4UnifiedAudioEmbedderModel(config) if config.audio is not None else None + ) + + def forward(self, op: OpBuilder, **kwargs): + raise NotImplementedError( + "Gemma4UnifiedModel is a multi-model split; Gemma4UnifiedTask builds " + "each sub-module (decoder, vision_encoder, embedding, and optionally " + "audio_encoder) separately." + ) + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Rename HuggingFace ``gemma4_unified`` checkpoint keys to ONNX names. + + After stripping the leading ``model.`` prefix: + + - ``language_model.lm_head.*`` → ``decoder.lm_head.*`` + - ``language_model.*`` → ``decoder.model.*`` (token embedding is also + shared with ``embedding.embed_tokens.weight``) + - ``vision_embedder.*`` → ``vision_encoder.*`` (``pos_embedding`` is + split into ``pos_emb_x``/``pos_emb_y``) + - ``embed_vision.embedding_projection.*`` → ``vision_encoder.projector.*`` + - ``embed_audio.embedding_projection.*`` → ``audio_encoder.projector.*`` + - ``embed_{vision,audio}.*.embedding_pre_projection_norm.*`` → skip + (scale-free RMSNorm, no learnable weight) + """ + # Strip top-level "model." prefix used by HF multimodal checkpoints. + state_dict = { + (key[len("model.") :] if key.startswith("model.") else key): value + for key, value in state_dict.items() + } + + # Synthesize lm_head from embed_tokens when weights are tied. + if self.config.tie_word_embeddings: + embed_key = "language_model.embed_tokens.weight" + head_key = "language_model.lm_head.weight" + if head_key not in state_dict and embed_key in state_dict: + state_dict[head_key] = state_dict[embed_key] + + renamed: dict[str, torch.Tensor] = {} + for key, value in state_dict.items(): + if key.startswith("language_model."): + suffix = key[len("language_model.") :] + if suffix.startswith("lm_head"): + renamed["decoder." + suffix] = value + else: + renamed["decoder.model." + suffix] = value + if suffix == "embed_tokens.weight": + renamed["embedding.embed_tokens.weight"] = value + + elif key.startswith("vision_embedder.pos_embedding"): + # [posemb_size, 2, mm_embed_dim] → two [posemb_size, mm_embed_dim] + renamed["vision_encoder.pos_emb_x.weight"] = value[:, 0, :].contiguous() + renamed["vision_encoder.pos_emb_y.weight"] = value[:, 1, :].contiguous() + + elif key.startswith("vision_embedder."): + renamed["vision_encoder." + key[len("vision_embedder.") :]] = value + + elif key.startswith("embed_vision.embedding_projection."): + suffix = key[len("embed_vision.embedding_projection.") :] + renamed["vision_encoder.projector." + suffix] = value + + elif key.startswith("embed_audio.embedding_projection."): + suffix = key[len("embed_audio.embedding_projection.") :] + renamed["audio_encoder.projector." + suffix] = value + + elif key.startswith(("embed_vision.", "embed_audio.")): + # *.embedding_pre_projection_norm.*: scale-free RMSNorm, no weight. + pass + + else: + renamed[key] = value + + return renamed diff --git a/src/mobius/models/gemma4_test.py b/src/mobius/models/gemma4_test.py index cef7b3803..66e24e194 100644 --- a/src/mobius/models/gemma4_test.py +++ b/src/mobius/models/gemma4_test.py @@ -176,3 +176,295 @@ def test_vnorm_fp16_no_nan(self): np.ones_like(output, dtype=np.float32), atol=0.01, ) + + +class TestGemma4BlockSequenceIds: + """Vision-block bidirectional attention wiring (use_bidirectional_attention).""" + + def test_compute_block_sequence_ids_values(self): + """``_compute_block_sequence_ids`` matches HF get_block_sequence_ids_for_mask.""" + import numpy as np + + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _compute_block_sequence_ids + from mobius.tasks._base import _make_graph, _make_model + + graph, builder = _make_graph() + op = builder.op + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[1, "S"]) + out = _compute_block_sequence_ids(op, input_ids, image_token_id=255) + builder.add_output(out, "block_sequence_ids") + session = OnnxModelSession(_make_model(graph), device="cpu") + + # text img img text aud aud text img text + ids = np.array([[10, 255, 255, 11, 254, 254, 11, 255, 12]], dtype=np.int64) + result = session.run({"input_ids": ids})["block_sequence_ids"] + # Only image runs form blocks (groups 0, 1). Audio (254) and text -> -1, + # matching HF where audio is token-type 3 and excluded from is_vision. + expected = np.array([[-1, 0, 0, -1, -1, -1, -1, 1, -1]], dtype=np.int64) + np.testing.assert_array_equal(result, expected) + + def test_unsupported_bidirectional_mode_raises(self): + """``use_bidirectional_attention='all'`` is rejected, not silently causal.""" + import pytest + + from mobius.models.gemma4 import Gemma4TextModel + + config = _tiny_gemma4_config(use_bidirectional_attention="all") + with pytest.raises(NotImplementedError, match="use_bidirectional_attention"): + Gemma4TextModel(config) + + def test_audio_tokens_excluded_from_blocks(self): + """Audio placeholders never join a vision block (HF parity). + + HF derives ``is_vision`` from ``mm_token_type_ids`` as ``(==1)|(==2)`` + (image or video); audio is token-type ``3`` and is excluded, so an audio + run adjacent to an image run does NOT extend the block. + """ + import numpy as np + + from mobius._testing.ort_inference import OnnxModelSession + from mobius.models.gemma4 import _compute_block_sequence_ids + from mobius.tasks._base import _make_graph, _make_model + + graph, builder = _make_graph() + op = builder.op + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[1, "S"]) + out = _compute_block_sequence_ids(op, input_ids, image_token_id=255) + builder.add_output(out, "block_sequence_ids") + session = OnnxModelSession(_make_model(graph), device="cpu") + + ids = np.array([[10, 255, 255, 254, 254, 11]], dtype=np.int64) + result = session.run({"input_ids": ids})["block_sequence_ids"] + # Image run -> block 0; the adjacent audio run (254) stays -1 (causal). + expected = np.array([[-1, 0, 0, -1, -1, -1]], dtype=np.int64) + np.testing.assert_array_equal(result, expected) + + def test_package_wires_block_sequence_ids_end_to_end(self): + """Decoder takes input_ids and derives the vision-block overlay itself. + + With ``use_bidirectional_attention='vision'`` the embedding no longer + emits ``block_sequence_ids``; instead the decoder receives ``input_ids`` + (alongside ``inputs_embeds``) and computes the overlay internally. This + avoids a cross-model tensor that onnxruntime-genai cannot forward. + """ + from mobius.models.gemma4 import Gemma4Model + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_config(use_bidirectional_attention="vision", image_token_id=255) + pkg = Gemma4Task().build(Gemma4Model(config), config) + + emb_outputs = {o.name for o in pkg["embedding"].graph.outputs} + dec_inputs = {i.name for i in pkg["decoder"].graph.inputs} + assert "block_sequence_ids" not in emb_outputs + 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. + 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 + + def test_no_block_sequence_ids_when_causal(self): + """Without bidirectional attention, no input_ids/overlay is wired.""" + from mobius.models.gemma4 import Gemma4Model + from mobius.tasks._gemma4 import Gemma4Task + + config = _tiny_gemma4_config(use_bidirectional_attention=None) + pkg = Gemma4Task().build(Gemma4Model(config), config) + + emb_outputs = {o.name for o in pkg["embedding"].graph.outputs} + dec_inputs = {i.name for i in pkg["decoder"].graph.inputs} + assert "block_sequence_ids" not in emb_outputs + assert "block_sequence_ids" not in dec_inputs + assert "input_ids" not in dec_inputs + + +def _tiny_gemma4_unified_config(**overrides) -> Gemma4Config: + """Minimal gemma4_unified config (dense decoder + encoder-free embedders).""" + from mobius._configs import Gemma4AudioConfig, VisionConfig + + base = dict( + model_type="gemma4_unified", + enable_moe_block=False, + tie_word_embeddings=True, + use_bidirectional_attention="vision", + image_token_id=255, + vision=VisionConfig( + hidden_size=32, + position_embedding_size=1120, + patch_size=4, + pooling_kernel_size=3, + out_hidden_size=32, + norm_eps=1e-6, + ), + audio=Gemma4AudioConfig(hidden_size=16, output_proj_dims=16, audio_token_id=254), + ) + base.update(overrides) + return _tiny_gemma4_config(**base) + + +class TestGemma4UnifiedPreprocessWeights: + """Gemma4UnifiedModel.preprocess_weights — checkpoint name mapping.""" + + def test_full_checkpoint_rename(self): + from mobius.models.gemma4 import Gemma4UnifiedModel + + config = _tiny_gemma4_unified_config() + model = Gemma4UnifiedModel(config) + + # pos_embedding stored as [posemb, 2, mm_embed_dim]; x-axis = [:, 0, :], + # y-axis = [:, 1, :]. Use distinct constants to verify the split. + pos_embedding = torch.empty(1120, 2, 32) + pos_embedding[:, 0, :] = 1.0 + pos_embedding[:, 1, :] = 2.0 + + fake_sd = { + "model.language_model.embed_tokens.weight": torch.zeros(256, 64), + "model.language_model.layers.0.input_layernorm.weight": torch.zeros(64), + "model.vision_embedder.patch_ln1.weight": torch.zeros(432), + "model.vision_embedder.patch_dense.weight": torch.zeros(32, 432), + "model.vision_embedder.pos_embedding": pos_embedding, + "model.vision_embedder.pos_norm.weight": torch.zeros(32), + "model.embed_vision.embedding_projection.weight": torch.zeros(64, 32), + "model.embed_audio.embedding_projection.weight": torch.zeros(64, 16), + # Scale-free RMSNorms have no learnable weight in the checkpoint, but + # assert they are dropped even if a stray key appears. + "model.embed_vision.embedding_pre_projection_norm.weight": torch.zeros(32), + "model.embed_audio.embedding_pre_projection_norm.weight": torch.zeros(16), + } + result = model.preprocess_weights(fake_sd) + + # Text backbone → decoder.model.* and tied embedding/lm_head. + assert "decoder.model.embed_tokens.weight" in result + assert "embedding.embed_tokens.weight" in result + assert "decoder.lm_head.weight" in result # synthesized from tied embed + assert "decoder.model.layers.0.input_layernorm.weight" in result + + # Vision front-end → vision_encoder.*; pos_embedding split into x/y. + assert "vision_encoder.patch_ln1.weight" in result + assert "vision_encoder.patch_dense.weight" in result + assert "vision_encoder.pos_norm.weight" in result + assert result["vision_encoder.pos_emb_x.weight"].shape == (1120, 32) + assert result["vision_encoder.pos_emb_y.weight"].shape == (1120, 32) + assert torch.allclose(result["vision_encoder.pos_emb_x.weight"], torch.tensor(1.0)) + assert torch.allclose(result["vision_encoder.pos_emb_y.weight"], torch.tensor(2.0)) + + # Projections → vision_encoder.projector / audio_encoder.projector. + assert "vision_encoder.projector.weight" in result + assert "audio_encoder.projector.weight" in result + + # Scale-free pre-projection norms must be dropped (no graph initializer). + assert not any("embedding_pre_projection_norm" in k for k in result) + # The raw checkpoint module prefixes must not leak through. + assert not any(k.startswith("vision_embedder.") for k in result) + assert not any(k.startswith("embed_vision.") for k in result) + assert not any(k.startswith("language_model.") for k in result) + + def test_vision_embedder_preprocess_standalone(self): + from mobius.models.gemma4 import _Gemma4UnifiedVisionEmbedderModel + + config = _tiny_gemma4_unified_config() + embedder = _Gemma4UnifiedVisionEmbedderModel(config) + + pos_embedding = torch.empty(1120, 2, 32) + pos_embedding[:, 0, :] = 3.0 + pos_embedding[:, 1, :] = 4.0 + fake_sd = { + "vision_embedder.patch_ln1.weight": torch.zeros(432), + "vision_embedder.pos_embedding": pos_embedding, + "embed_vision.embedding_projection.weight": torch.zeros(64, 32), + "embed_vision.embedding_pre_projection_norm.weight": torch.zeros(32), + } + result = embedder.preprocess_weights(fake_sd) + + assert "patch_ln1.weight" in result + assert "projector.weight" in result + assert torch.allclose(result["pos_emb_x.weight"], torch.tensor(3.0)) + assert torch.allclose(result["pos_emb_y.weight"], torch.tensor(4.0)) + assert not any("embedding_pre_projection_norm" in k for k in result) + + def test_audio_embedder_preprocess_standalone(self): + from mobius.models.gemma4 import _Gemma4UnifiedAudioEmbedderModel + + config = _tiny_gemma4_unified_config() + embedder = _Gemma4UnifiedAudioEmbedderModel(config) + + fake_sd = { + "embed_audio.embedding_projection.weight": torch.zeros(64, 16), + "embed_audio.embedding_pre_projection_norm.weight": torch.zeros(16), + } + result = embedder.preprocess_weights(fake_sd) + + assert result["projector.weight"].shape == (64, 16) + assert not any("embedding_pre_projection_norm" in k for k in result) + + +class TestGemma4UnifiedConfigHooks: + """Config extraction hooks map unified vision/audio sub-configs.""" + + def test_vision_hook_maps_fields(self): + from types import SimpleNamespace + + from mobius._configs.per_model._gemma4_unified_vision import ( + _gemma4_unified_vision, + ) + + composite = SimpleNamespace( + model_type="gemma4_unified", + image_token_id=258880, + vision_config=SimpleNamespace( + mm_embed_dim=3840, + patch_size=16, + pooling_kernel_size=3, + mm_posemb_size=1120, + output_proj_dims=3840, + rms_norm_eps=1e-6, + ), + ) + fields: dict = {} + _gemma4_unified_vision(composite, None, "gemma4_unified", fields) + + assert fields["hidden_size"] == 3840 + assert fields["patch_size"] == 16 + assert fields["pooling_kernel_size"] == 3 + assert fields["position_embedding_size"] == 1120 + assert fields["out_hidden_size"] == 3840 + assert fields["image_token_id"] == 258880 + + def test_vision_hook_skips_unrelated_model(self): + from types import SimpleNamespace + + from mobius._configs.per_model._gemma4_unified_vision import ( + _gemma4_unified_vision, + ) + + composite = SimpleNamespace(model_type="qwen2_vl", vision_config=object()) + fields: dict = {} + result = _gemma4_unified_vision(composite, None, "qwen2_vl", fields) + assert result is None + assert fields == {} + + def test_audio_hook_maps_fields(self): + from types import SimpleNamespace + + from mobius._configs.per_model._gemma4_unified_audio import ( + _gemma4_unified_audio, + ) + + composite = SimpleNamespace( + model_type="gemma4_unified", + audio_token_id=258881, + audio_config=SimpleNamespace(audio_embed_dim=640), + ) + result = _gemma4_unified_audio(composite, None, "gemma4_unified", {}) + + assert result is not None + audio_cfg = result["audio"] + assert audio_cfg.hidden_size == 640 + assert audio_cfg.output_proj_dims == 640 + assert audio_cfg.audio_token_id == 258881 diff --git a/src/mobius/rewrite_rules/_group_query_attention.py b/src/mobius/rewrite_rules/_group_query_attention.py index d2f0a3435..df19cb06a 100644 --- a/src/mobius/rewrite_rules/_group_query_attention.py +++ b/src/mobius/rewrite_rules/_group_query_attention.py @@ -49,25 +49,6 @@ RewriteRuleSet, ) -# CUDA EP's GroupQueryAttention kernel historically enforced MAX_HEAD_SIZE = 256. -# Keep the rewrite-rule limit conservative until GQA fusion is gated on a -# runtime/EP capability check. This avoids emitting GroupQueryAttention nodes -# with head_dim=512 that can still fail on released ORT builds. -_MAX_GQA_HEAD_DIM = 256 - - -def _head_dim_exceeds_gqa_limit(past_key) -> int | None: - """Return the head_dim if it exceeds the GQA kernel limit, else None. - - ``past_key`` is expected to have shape ``(batch, kv_heads, seq, head_dim)``. - """ - if past_key is None or past_key.shape is None or len(past_key.shape) < 4: - return None - hd = past_key.shape[3] - if isinstance(hd, int) and hd > _MAX_GQA_HEAD_DIM: - return hd - return None - class RotaryAttentionToGQA(RewriteRuleClassBase): """Replace RotaryEmbedding + Attention with GroupQueryAttention. @@ -166,11 +147,6 @@ def check(self, context, attn_out, cos, sin, past_key, past_value, **_): if past_value.producer() is not None: return result.fail("past_value is not a graph input") - # Skip when head_dim exceeds CUDA GQA MAX_HEAD_SIZE (256). - hd = _head_dim_exceeds_gqa_limit(past_key) - if hd is not None: - return result.fail(f"head_dim={hd} exceeds GQA MAX_HEAD_SIZE={_MAX_GQA_HEAD_DIM}") - return result # ------------------------------------------------------------------ rewrite @@ -595,11 +571,6 @@ def check(self, context, attn_out, past_key, past_value, **_): if not any(gi.name == "attention_mask" for gi in graph.inputs): return result.fail("No attention_mask graph input — cannot build seqlens_k") - # Skip when head_dim exceeds CUDA GQA MAX_HEAD_SIZE (256). - hd = _head_dim_exceeds_gqa_limit(past_key) - if hd is not None: - return result.fail(f"head_dim={hd} exceeds GQA MAX_HEAD_SIZE={_MAX_GQA_HEAD_DIM}") - return result # ------------------------------------------------------------------ rewrite @@ -675,13 +646,23 @@ def group_query_attention_rules() -> RewriteRuleSet: position embeddings are not expressed as standard ``RotaryEmbedding`` ops (e.g. Qwen3.5 3D mRoPE via ``Where`` nodes). + GQA fusion is applied uniformly to every decoder ``Attention`` node + regardless of ``head_dim`` — there is no head-dim cap. Whether a given + runtime's GQA kernel supports a particular ``head_dim`` is a separate EP + concern, gated by :attr:`~mobius._execution_providers.EpCapabilities.gqa_dtypes`. + QKV packing is a separate optional pass; use :func:`pack_qkv_for_gqa_rules` for that. Returns: :class:`RewriteRuleSet` containing the GQA fusion rules. """ - return RewriteRuleSet([RotaryAttentionToGQA().rule(), AttentionToGQA().rule()]) + return RewriteRuleSet( + [ + RotaryAttentionToGQA.rule(), + AttentionToGQA.rule(), + ] + ) def pack_qkv_for_gqa_rules() -> RewriteRuleSet: diff --git a/src/mobius/rewrite_rules/_group_query_attention_test.py b/src/mobius/rewrite_rules/_group_query_attention_test.py index aa10d3dba..b67db66a8 100644 --- a/src/mobius/rewrite_rules/_group_query_attention_test.py +++ b/src/mobius/rewrite_rules/_group_query_attention_test.py @@ -130,6 +130,40 @@ def test_replaces_attention_with_gqa(self): assert counts_after.get("Attention", 0) == 0 assert counts_after["GroupQueryAttention"] == 28 + def test_fuses_large_head_dim(self): + """GQA fusion applies uniformly, including head_dim > 256. + + There is no head-dim cap: every decoder ``Attention`` node is fused to + ``GroupQueryAttention`` regardless of ``head_dim``. This covers + Gemma4-style global-attention layers (head_dim=512), which the CUDA GQA + kernel handles via its FP32-QK-accumulation unfused fallback. + """ + cfg = ArchitectureConfig( + hidden_size=1024, + intermediate_size=128, + num_attention_heads=2, + num_key_value_heads=1, + head_dim=512, + num_hidden_layers=2, + vocab_size=256, + max_position_embeddings=128, + hidden_act="silu", + rms_norm_eps=1e-6, + rope_type="default", + rope_theta=10000.0, + pad_token_id=0, + ) + model = registry.get("llama")(cfg) + pkg = build_from_module(model, cfg, execution_provider="default") + m = pkg["model"] + assert count_ops(m).get("Attention", 0) == 2 + + rewrite(m, pattern_rewrite_rules=group_query_attention_rules()) + + counts_after = count_ops(m) + assert counts_after.get("Attention", 0) == 0 + assert counts_after.get("GroupQueryAttention", 0) == 2 + def test_absorbs_rotary_embedding(self): """RotaryEmbedding ops are absorbed into GQA with do_rotary=1.""" pkg = build("Qwen/Qwen3-0.6B", load_weights=False) diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 0d2d9ec1c..52ca5b0da 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -31,6 +31,7 @@ "FeatureExtractionTask", "FunASRSpeechLanguageTask", "Gemma4Task", + "Gemma4UnifiedTask", "Gemma4TextCausalLMTask", "HybridCausalLMTask", "HybridQwenVLTask", @@ -83,6 +84,7 @@ from mobius.tasks._gemma4 import ( Gemma4Task, Gemma4TextCausalLMTask, + Gemma4UnifiedTask, ) from mobius.tasks._hunyuan_vl_mot import HunYuanVLMoTTask from mobius.tasks._image_classification import ImageClassificationTask @@ -134,6 +136,7 @@ "qwen3-vl-vision-language": Qwen3VLVisionLanguageTask, "gemma4": Gemma4Task, "gemma4-text-generation": Gemma4TextCausalLMTask, + "gemma4-unified": Gemma4UnifiedTask, "hunyuan-vl-mot": HunYuanVLMoTTask, "multimodal": MultiModalTask, "phi4mm-multimodal": Phi4MMMultiModalTask, diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py index 92754b604..965c19730 100644 --- a/src/mobius/tasks/_gemma4.py +++ b/src/mobius/tasks/_gemma4.py @@ -260,6 +260,21 @@ def _build_decoder( shape=[batch, seq_len, total_per_layer], ) + # Vision-block bidirectional attention: the decoder receives the raw + # ``input_ids`` (alongside ``inputs_embeds``) and derives the block + # overlay internally. This avoids a separate cross-model + # ``block_sequence_ids`` tensor, which onnxruntime-genai cannot forward + # between the embedding and decoder sub-models (it can forward + # ``input_ids``). Only models with + # ``use_bidirectional_attention == "vision"`` need it. + input_ids_val: ir.Value | None = None + if config.use_bidirectional_attention == "vision": + input_ids_val = builder.input( + "input_ids", + dtype=ir.DataType.INT64, + shape=[batch, seq_len], + ) + past_key_values = _make_gemma4_kv_cache_inputs(builder, config, batch, past_seq_len) logits, present_key_values = decoder( @@ -269,6 +284,7 @@ def _build_decoder( position_ids=position_ids, per_layer_inputs=per_layer_inputs_val, past_key_values=past_key_values, + input_ids=input_ids_val, ) builder.add_output(logits, "logits") @@ -418,13 +434,116 @@ def _build_embedding( audio_features=audio_features_val, ) - per_layer_dim = getattr(config, "hidden_size_per_layer_input", 0) - if per_layer_dim: - # Embedding model returns (inputs_embeds, per_layer_inputs) when - # per-layer input gating is enabled. - inputs_embeds, per_layer_inputs = result - builder.add_output(inputs_embeds, "inputs_embeds") - builder.add_output(per_layer_inputs, "per_layer_inputs") - else: - builder.add_output(result, "inputs_embeds") + # ``embedding`` returns a dict of named outputs: always + # ``inputs_embeds``; optionally ``per_layer_inputs`` (per-layer gating). + builder.add_output(result["inputs_embeds"], "inputs_embeds") + if "per_layer_inputs" in result: + builder.add_output(result["per_layer_inputs"], "per_layer_inputs") + return _make_model(graph) + + +class Gemma4UnifiedTask(Gemma4Task): + """Task for ``gemma4_unified`` (gemma-4-12B) encoder-free multimodal models. + + Reuses :class:`Gemma4Task`'s ``decoder`` and ``embedding`` builders (the + text decoder and multimodal fusion are identical to gemma4), but overrides + the vision and audio builders for the encoder-free embedders: + + - **vision** — raw merged pixel patches ``pixel_values [B, N, P^2*3]`` and + integer patch coordinates ``pixel_position_ids [B, N, 2]`` → + ``image_features [num_valid_patches, text_hidden]`` (padding patches, + with position ``-1``, are stripped inside the graph). + - **audio** — raw waveform-frame features ``input_features [B, T, D_a]`` + and a validity mask ``input_features_mask [B, T]`` → + ``audio_features [num_valid_frames, text_hidden]`` (padding frames are + stripped inside the graph; no separate mask output is needed). + """ + + def _build_vision( + self, + vision: nn.Module, + config: Gemma4Config, + ) -> ir.Model: + """Build the encoder-free vision embedder. + + Inputs: + - ``pixel_values [B, N, P^2*3]``: raw merged pixel patches, where + ``P = patch_size * pooling_kernel_size`` (48 for gemma-4-12B). + - ``pixel_position_ids [B, N, 2]``: integer (x, y) patch coordinates; + ``(-1, -1)`` marks padding slots. + + Output: + - ``image_features [num_valid_patches, text_hidden]``. + """ + batch = ir.SymbolicDim("batch") + num_patches = ir.SymbolicDim("num_patches") + vc = config.vision + patch_size = (vc.patch_size if vc else None) or 16 + pooling = (vc.pooling_kernel_size if vc else None) or 3 + model_patch_size = patch_size * pooling + pixel_dim = 3 * model_patch_size * model_patch_size + + graph, builder = _make_graph(name="vision_encoder") + op = builder.op + + pixel_values = builder.input( + "pixel_values", + dtype=config.dtype, + shape=[batch, num_patches, pixel_dim], + ) + pixel_position_ids = builder.input( + "pixel_position_ids", + dtype=ir.DataType.INT64, + shape=[batch, num_patches, 2], + ) + + image_features = vision( + op, + pixel_values=pixel_values, + pixel_position_ids=pixel_position_ids, + ) + builder.add_output(image_features, "image_features") + return _make_model(graph) + + def _build_audio( + self, + audio: nn.Module, + config: Gemma4Config, + ) -> ir.Model: + """Build the encoder-free audio embedder. + + Inputs: + - ``input_features [B, T, D_a]``: raw waveform-frame features + (``D_a = audio_embed_dim``, 640 for gemma-4-12B). + - ``input_features_mask [B, T]``: BOOL mask, ``True`` for valid frames. + + Output: + - ``audio_features [num_valid_frames, text_hidden]``. + """ + batch = ir.SymbolicDim("batch") + time = ir.SymbolicDim("time") + audio_embed_dim = ( + getattr(config.audio, "hidden_size", None) if config.audio else None + ) or 640 + + graph, builder = _make_graph(name="audio_encoder") + op = builder.op + + input_features = builder.input( + "input_features", + dtype=config.dtype, + shape=[batch, time, audio_embed_dim], + ) + input_features_mask = builder.input( + "input_features_mask", + dtype=ir.DataType.BOOL, + shape=[batch, time], + ) + + audio_features, _ = audio( + op, + input_features, + input_features_mask=input_features_mask, + ) + builder.add_output(audio_features, "audio_features") return _make_model(graph) diff --git a/testdata/cases/causal-lm/gemma-4-12b.yaml b/testdata/cases/causal-lm/gemma-4-12b.yaml new file mode 100644 index 000000000..0ba4f0a5d --- /dev/null +++ b/testdata/cases/causal-lm/gemma-4-12b.yaml @@ -0,0 +1,18 @@ +model_id: "google/gemma-4-12B" +model_type: "gemma4_unified_text" +revision: "main" +task_type: "text-generation" +dtype: "float32" + +inputs: + prompts: + - "Here is my poem:" + +level: "L4+L5" + +generation: + max_new_tokens: 20 + do_sample: false + +ci_skip_reason: "Model is 12B — too large for CI hardware. Runs locally (gated repo needs HF auth); golden files are generated and committed." +notes: "Gemma 4 12B unified (text-only path) via Gemma4CausalLMModel. Dual head_dim (separate attention/value head_dim), single shared KV head (k_eq_v), KV sharing, sliding+global attention, logit softcapping." diff --git a/testdata/cases/speech/gemma-4-12b-audio.yaml b/testdata/cases/speech/gemma-4-12b-audio.yaml new file mode 100644 index 000000000..258aabf27 --- /dev/null +++ b/testdata/cases/speech/gemma-4-12b-audio.yaml @@ -0,0 +1,21 @@ +model_id: "google/gemma-4-12B" +model_type: "gemma4_unified" +revision: "main" +task_type: "speech-language" +dtype: "float32" + +inputs: + prompts: + - "The audio says" + audio: + - "652-129742-0006.flac" + +level: "L4+L5" + +generation: + max_new_tokens: 30 + do_sample: false + eos_token_id: 1 + +ci_skip_reason: "Model is 12B — too large for CI hardware. Runs locally (gated repo needs HF auth); golden files are generated and committed." +notes: "Gemma 4 12B unified (encoder-free) audio input. Raw 640-dim waveform-frame embedder instead of a Conformer tower. 4-model split: decoder + vision_encoder + audio_encoder + embedding. Base checkpoint, so completion-style prompt + structural-token suppression (suppress_tokens 258882/258883)." diff --git a/testdata/cases/vision-language/gemma-4-12b.yaml b/testdata/cases/vision-language/gemma-4-12b.yaml new file mode 100644 index 000000000..edfa2bd88 --- /dev/null +++ b/testdata/cases/vision-language/gemma-4-12b.yaml @@ -0,0 +1,21 @@ +model_id: "google/gemma-4-12B" +model_type: "gemma4_unified" +revision: "main" +task_type: "image-text-to-text" +dtype: "float32" + +inputs: + prompts: + - "This image shows" + images: + - "pipeline-cat-chonk.jpeg" + +level: "L4+L5" + +generation: + max_new_tokens: 30 + do_sample: false + eos_token_id: 1 + +ci_skip_reason: "Model is 12B — too large for CI hardware. Runs locally (gated repo needs HF auth); golden files are generated and committed." +notes: "Gemma 4 12B unified (encoder-free) Image-Text-to-Text. Raw-patch vision embedder instead of a SigLIP tower. Split: decoder + vision_encoder + embedding (audio_encoder added when config.audio is set). Vision-block bidirectional attention in the decoder." diff --git a/testdata/golden/causal-lm/gemma-4-12b.json b/testdata/golden/causal-lm/gemma-4-12b.json new file mode 100644 index 000000000..70a072a5a --- /dev/null +++ b/testdata/golden/causal-lm/gemma-4-12b.json @@ -0,0 +1,42 @@ +{ + "top1_id": 108, + "top2_id": 107, + "top10_ids": [ + 108, + 107, + 109, + 236743, + 669, + 623, + 564, + 110, + 999, + 562 + ], + "top10_logits": [ + "0x1.8dfae80000000p+4", + "0x1.7303460000000p+4", + "0x1.633ea00000000p+4", + "0x1.5b4c2e0000000p+4", + "0x1.5149640000000p+4", + "0x1.50e1780000000p+4", + "0x1.4e8a920000000p+4", + "0x1.4c96080000000p+4", + "0x1.4902780000000p+4", + "0x1.43e3720000000p+4" + ], + "logits_summary": [ + "0x1.8dfae80000000p+4", + "-0x1.02eae60000000p+2", + "0x1.4483d256b5865p+2", + "0x1.aa558ca60eaaep+1" + ], + "input_ids": [ + 2, + 8291, + 563, + 1041, + 27355, + 236787 + ] +} diff --git a/testdata/golden/causal-lm/gemma-4-12b_generation.json b/testdata/golden/causal-lm/gemma-4-12b_generation.json new file mode 100644 index 000000000..df9a6252b --- /dev/null +++ b/testdata/golden/causal-lm/gemma-4-12b_generation.json @@ -0,0 +1,27 @@ +{ + "model_id": "google/gemma-4-12B", + "prompt": "Here is my poem:", + "generated_tokens": [ + 108, + 198, + 199, + 818, + 4109, + 208, + 207, + 108, + 199, + 818, + 1902, + 563, + 496, + 4148, + 1977, + 236764, + 208, + 108, + 199, + 4573 + ], + "generated_text": "\n\nThe World\n\nThe world is a beautiful place,\n\nBut" +} diff --git a/testdata/golden/speech/gemma-4-12b-audio.json b/testdata/golden/speech/gemma-4-12b-audio.json new file mode 100644 index 000000000..345d6ccc2 --- /dev/null +++ b/testdata/golden/speech/gemma-4-12b-audio.json @@ -0,0 +1,270 @@ +{ + "top1_id": 236764, + "top2_id": 600, + "top10_ids": [ + 236764, + 600, + 236743, + 506, + 625, + 531, + 611, + 106377, + 496, + 123051 + ], + "top10_logits": [ + "0x1.907d5e0000000p+4", + "0x1.8c4e8a0000000p+4", + "0x1.88c3200000000p+4", + "0x1.813b980000000p+4", + "0x1.805a560000000p+4", + "0x1.7df7de0000000p+4", + "0x1.7baa620000000p+4", + "0x1.771b680000000p+4", + "0x1.74eb920000000p+4", + "0x1.74654e0000000p+4" + ], + "logits_summary": [ + "0x1.907d5e0000000p+4", + "-0x1.10cde00000000p+3", + "0x1.4cddae0f7d301p+2", + "0x1.0e7ac0d3c9dcfp+2" + ], + "input_ids": [ + 2, + 256000, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258881, + 258883, + 818, + 9855, + 3189 + ] +} diff --git a/testdata/golden/speech/gemma-4-12b-audio_generation.json b/testdata/golden/speech/gemma-4-12b-audio_generation.json new file mode 100644 index 000000000..3a000cf56 --- /dev/null +++ b/testdata/golden/speech/gemma-4-12b-audio_generation.json @@ -0,0 +1,37 @@ +{ + "model_id": "google/gemma-4-12B", + "prompt": "The audio says", + "generated_tokens": [ + 236764, + 623, + 13751, + 7445, + 50890, + 106377, + 236764, + 2541, + 1131, + 15092, + 236764, + 8009, + 9551, + 236764, + 22300, + 236764, + 532, + 48399, + 531, + 3409, + 1781, + 236743, + 236743, + 2094, + 563, + 496, + 1237, + 8610, + 221509, + 699 + ], + "generated_text": ", \"Take cold boiled cauliflower, break into branches, adding salt, pepper, and vinegar to season.\" This is a public domain audiobook from" +} diff --git a/testdata/golden/vision-language/gemma-4-12b.json b/testdata/golden/vision-language/gemma-4-12b.json new file mode 100644 index 000000000..00f5d3648 --- /dev/null +++ b/testdata/golden/vision-language/gemma-4-12b.json @@ -0,0 +1,308 @@ +{ + "top1_id": 496, + "top2_id": 506, + "top10_ids": [ + 496, + 506, + 614, + 886, + 593, + 1041, + 1156, + 1217, + 1023, + 600 + ], + "top10_logits": [ + "0x1.9eb3ea0000000p+4", + "0x1.86a1a80000000p+4", + "0x1.7b8d680000000p+4", + "0x1.66dbf80000000p+4", + "0x1.55ba340000000p+4", + "0x1.51fa3c0000000p+4", + "0x1.5125500000000p+4", + "0x1.5011040000000p+4", + "0x1.4f05d00000000p+4", + "0x1.4d22900000000p+4" + ], + "logits_summary": [ + "0x1.9eb3ea0000000p+4", + "-0x1.f6c69c0000000p+2", + "0x1.07a0bbcac4a64p+2", + "0x1.0c1e1c730fc57p+2" + ], + "input_ids": [ + 2, + 255999, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258880, + 258882, + 2094, + 2471, + 3831 + ] +} diff --git a/testdata/golden/vision-language/gemma-4-12b_generation.json b/testdata/golden/vision-language/gemma-4-12b_generation.json new file mode 100644 index 000000000..f73dce970 --- /dev/null +++ b/testdata/golden/vision-language/gemma-4-12b_generation.json @@ -0,0 +1,15 @@ +{ + "model_id": "google/gemma-4-12B", + "prompt": "This image shows", + "generated_tokens": [ + 496, + 34455, + 9307, + 528, + 506, + 7613, + 236761, + 1 + ], + "generated_text": " a bobcat in the snow." +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 86c769014..6d7846686 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -336,6 +336,31 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: }, True, ), + ( + # gemma-4-12B text backbone (reuses Gemma4CausalLMModel): dual head_dim + # (local 16 / global 32), single global KV head with attention_k_eq_v, + # vision-block bidirectional attention. Built generically here; numeric + # parity is covered by the real-weight integration test (no HF model_type + # for this internal alias, so it is excluded from synthetic parity). + "gemma4_unified_text", + { + "_config_cls": Gemma4Config, + "hidden_act": "gelu_pytorch_tanh", + "attn_qk_norm": True, + "layer_types": ["sliding_attention", "full_attention"], + "sliding_window": 8, + "global_head_dim": 2 * TINY_HEAD_DIM, + "global_rope_theta": 1_000_000.0, + "global_partial_rotary_factor": 0.25, + "final_logit_softcapping": 30.0, + "hidden_size_per_layer_input": 0, + "num_global_key_value_heads": 1, + "attention_k_eq_v": True, + "use_bidirectional_attention": "vision", + "tie_word_embeddings": True, + }, + False, + ), ( "gemma3n_text", { @@ -2089,6 +2114,38 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: # --- Gemma4 Any-to-Any (4-model split: decoder + vision + speech + embedding) --- # Tested directly in test_gemma4_any_to_any_graph; omitted from parametrized suite # because it uses the unified "gemma4" registry key, same as the VL-only config above. + # --- Gemma4 unified gemma-4-12B (encoder-free 3-model split here; the 4-model + # audio path is exercised by test_gemma4_unified_multimodal_graph) --- + ( + "gemma4_unified", + { + "_config_cls": Gemma4Config, + "hidden_act": "gelu_pytorch_tanh", + "attn_qk_norm": True, + "layer_types": ["sliding_attention", "full_attention"], + "sliding_window": 8, + "global_head_dim": 2 * TINY_HEAD_DIM, + "global_rope_theta": 1_000_000.0, + "global_partial_rotary_factor": 0.25, + "final_logit_softcapping": 30.0, + "hidden_size_per_layer_input": 0, + "num_global_key_value_heads": 1, + "attention_k_eq_v": True, + "use_bidirectional_attention": "vision", + "image_token_id": 255999, + "tie_word_embeddings": True, + # Encoder-free vision embedder (raw patches → pooled image features). + "vision": VisionConfig( + hidden_size=48, + patch_size=4, + pooling_kernel_size=2, + position_embedding_size=64, + out_hidden_size=48, + norm_eps=1e-6, + ), + }, + False, + ), # --- Blip2 (ViT + Q-Former + LLM) --- ( "blip-2", diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index d6d2f9c88..f3e0d50b4 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -1249,6 +1249,57 @@ def test_gemma4_moe_graph(self): assert "input_ids" in input_names assert "logits" in output_names + def test_gemma4_unified_text_graph(self): + """Build the gemma4_unified (gemma-4-12B) text backbone via Gemma4CausalLMModel. + + ``gemma4_unified_text`` reuses Gemma4CausalLMModel. This exercises the + 12B-family text architecture: dual head_dim (local 16 / global 32), + ``attention_k_eq_v`` with a single global KV head, vision-block + bidirectional attention, and final-logit softcapping. + """ + from mobius._configs import Gemma4Config + + config = Gemma4Config( + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=256, + rms_norm_eps=1e-6, + hidden_act="gelu_pytorch_tanh", + attn_qk_norm=True, + layer_types=["sliding_attention", "full_attention"], + sliding_window=8, + global_head_dim=32, + global_rope_theta=1_000_000.0, + global_partial_rotary_factor=0.25, + final_logit_softcapping=30.0, + hidden_size_per_layer_input=0, + num_global_key_value_heads=1, + attention_k_eq_v=True, + use_bidirectional_attention="vision", + pad_token_id=0, + tie_word_embeddings=True, + ) + model_cls = registry.get("gemma4_unified_text") + module = model_cls(config) + task_name = _default_task_for_model("gemma4_unified_text") + task = get_task(task_name) + pkg = task.build(module, config) + + assert "model" in pkg + model = pkg["model"] + input_names = {i.name for i in model.graph.inputs} + output_names = {o.name for o in model.graph.outputs} + assert "input_ids" in input_names + assert "logits" in output_names + # Full-attention layer uses the single global KV head, so its cache + # entry has a different head_dim than the sliding layer. + assert "past_key_values.0.key" in input_names + assert "past_key_values.1.key" in input_names + def test_gemma4_any_to_any_graph(self): """Build Gemma4 Any-to-Any model (4-model split: decoder+vision+speech+embedding). @@ -1342,6 +1393,98 @@ def test_gemma4_any_to_any_graph(self): assert "present.1.key" not in decoder_output_names # shared layer excluded assert "present.1.value" not in decoder_output_names # shared layer excluded + def test_gemma4_unified_multimodal_graph(self): + """Build gemma4_unified (gemma-4-12B) encoder-free multimodal model. + + Produces a 4-model split (decoder + vision_encoder + audio_encoder + + embedding). The vision/audio encoders are encoder-free embedders + (no SigLIP/Conformer tower); the decoder uses vision-block + bidirectional attention, which it derives internally from + ``input_ids`` (the embedding model does *not* emit + ``block_sequence_ids``). + """ + from mobius._configs import Gemma4AudioConfig, Gemma4Config + + config = Gemma4Config( + num_hidden_layers=2, + hidden_size=64, + intermediate_size=128, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=256, + rms_norm_eps=1e-6, + hidden_act="gelu_pytorch_tanh", + attn_qk_norm=True, + layer_types=["sliding_attention", "full_attention"], + sliding_window=8, + global_head_dim=32, + global_rope_theta=1_000_000.0, + global_partial_rotary_factor=0.25, + final_logit_softcapping=30.0, + hidden_size_per_layer_input=0, + num_global_key_value_heads=1, + attention_k_eq_v=True, + use_bidirectional_attention="vision", + image_token_id=255999, + pad_token_id=0, + tie_word_embeddings=True, + vision=VisionConfig( + hidden_size=48, + patch_size=4, + pooling_kernel_size=2, + position_embedding_size=64, + out_hidden_size=48, + norm_eps=1e-6, + ), + audio=Gemma4AudioConfig( + hidden_size=40, + output_proj_dims=40, + audio_token_id=255998, + ), + ) + model_cls = registry.get("gemma4_unified") + module = model_cls(config) + task = get_task(_default_task_for_model("gemma4_unified")) + pkg = task.build(module, config) + + assert set(pkg.keys()) == { + "decoder", + "vision_encoder", + "audio_encoder", + "embedding", + }, f"gemma4_unified should produce 4 models, got: {set(pkg.keys())}" + + # Vision embedder: raw patches (no encoder layers) → image_features + vision = pkg["vision_encoder"] + v_inputs = {i.name for i in vision.graph.inputs} + assert v_inputs == {"pixel_values", "pixel_position_ids"} + assert "image_features" in {o.name for o in vision.graph.outputs} + + # Audio embedder: raw frames + mask → audio_features + audio = pkg["audio_encoder"] + a_inputs = {i.name for i in audio.graph.inputs} + assert a_inputs == {"input_features", "input_features_mask"} + assert "audio_features" in {o.name for o in audio.graph.outputs} + + # Embedding: fuses both modalities → inputs_embeds (no block_sequence_ids; + # the decoder derives the bidirectional overlay from input_ids itself) + embedding = pkg["embedding"] + e_inputs = {i.name for i in embedding.graph.inputs} + assert {"input_ids", "image_features", "audio_features"} <= e_inputs + e_outputs = {o.name for o in embedding.graph.outputs} + assert "inputs_embeds" in e_outputs + assert "block_sequence_ids" not in e_outputs + + # Decoder: consumes inputs_embeds + input_ids (for the vision-block + # bidirectional overlay, derived internally) + decoder = pkg["decoder"] + d_inputs = {i.name for i in decoder.graph.inputs} + assert "inputs_embeds" in d_inputs + assert "input_ids" in d_inputs + assert "block_sequence_ids" not in d_inputs + assert "logits" in {o.name for o in decoder.graph.outputs} + def test_gemma4_kv_shared_layer_tracing(self): """Verify all num_hidden_layers are traced and KV outputs = num_kv_layers. @@ -4351,6 +4494,8 @@ def test_jamba_preprocess_weights_moe_renames(self): "deepseek_vl_v2", "gemma3", "gemma4", + "gemma4_unified", + "gemma4_unified_text", "llava", "mllama", "phi4_multimodal", diff --git a/tests/e2e_golden_test.py b/tests/e2e_golden_test.py index 8006e04aa..a8748e9cc 100644 --- a/tests/e2e_golden_test.py +++ b/tests/e2e_golden_test.py @@ -68,6 +68,65 @@ def _get_test_device_kwargs() -> dict[str, str]: _IN_CI = os.environ.get("GITHUB_ACTIONS") == "true" +def _load_suppress_token_ids(model_id: str, trust_remote_code: bool = False) -> list[int]: + """Return ``generation_config.suppress_tokens`` for a model (empty if none). + + Mirrors HuggingFace ``generate()``: tokens in ``suppress_tokens`` are forced + to ``-inf`` at every decode step. Needed for base checkpoints (e.g. + ``google/gemma-4-12B``) whose generation_config suppresses the structural + ```` / ```` tokens — without it, greedy decode + degenerates (repeating those tokens) and diverges from the golden reference + produced by ``model.generate``. For models with no suppress_tokens this is a + no-op, so it is safe to apply unconditionally. + """ + import transformers + + try: + gen_config = transformers.GenerationConfig.from_pretrained( + model_id, trust_remote_code=trust_remote_code + ) + except Exception: + return [] + return [int(t) for t in (gen_config.suppress_tokens or [])] + + +def _suppress_logits(logits: np.ndarray, suppress_ids: list[int]) -> np.ndarray: + """Force ``suppress_ids`` columns of the last-position logits to ``-inf``.""" + if suppress_ids: + logits[..., suppress_ids] = -np.inf + return logits + + +def _build_mm_prompt( + processor: object, + base_prompt: str, + media_paths: list[str], + media_kind: str, +) -> str: + """Format a multimodal prompt, falling back when there is no chat template. + + Instruction-tuned processors expose a chat template that injects the right + media placeholder tokens. Base checkpoints (e.g. ``google/gemma-4-12B``) + ship none, so manually prepend one placeholder token per media item — the + processor then expands each into the correct number of soft tokens. + + ``media_kind`` is ``"image"`` or ``"audio"``. + """ + if getattr(processor, "chat_template", None): + content: list[dict[str, str]] = [] + for path in media_paths: + content.append({"type": media_kind, media_kind: str(_TESTDATA_DIR / path)}) + content.append({"type": "text", "text": base_prompt}) + messages = [{"role": "user", "content": content}] + return processor.apply_chat_template( # type: ignore[attr-defined] + messages, tokenize=False, add_generation_prompt=True + ) + placeholder = getattr(processor, f"{media_kind}_token", None) + if placeholder: + return placeholder * len(media_paths) + base_prompt + return base_prompt + + def _make_empty_kv_cache( session: OnnxModelSession, config: object, @@ -570,17 +629,8 @@ def _run_vision_language_prefill( ) image = Image.open(_TESTDATA_DIR / case.images[0]) - # Build chat template for models that need it - prompt_text = case.prompts[0] - if hasattr(processor, "apply_chat_template"): - content: list[dict[str, str]] = [] - for img_path in case.images: - content.append({"type": "image", "image": str(_TESTDATA_DIR / img_path)}) - content.append({"type": "text", "text": prompt_text}) - messages = [{"role": "user", "content": content}] - prompt_text = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) + # Build the prompt (chat template when available, else manual placeholder). + prompt_text = _build_mm_prompt(processor, case.prompts[0], case.images, "image") # Use PyTorch tensors then convert — some processors don't support np processed_pt = processor(text=prompt_text, images=[image], return_tensors="pt") @@ -764,16 +814,8 @@ def _run_vl_generation( ) image = Image.open(_TESTDATA_DIR / case.images[0]) - prompt_text = case.prompts[0] - if hasattr(processor, "apply_chat_template"): - content: list[dict[str, str]] = [] - for img_path in case.images: - content.append({"type": "image", "image": str(_TESTDATA_DIR / img_path)}) - content.append({"type": "text", "text": prompt_text}) - messages = [{"role": "user", "content": content}] - prompt_text = processor.apply_chat_template( - messages, tokenize=False, add_generation_prompt=True - ) + prompt_text = _build_mm_prompt(processor, case.prompts[0], case.images, "image") + suppress_ids = _load_suppress_token_ids(case.model_id, case.trust_remote_code) processed_pt = processor(text=prompt_text, images=[image], return_tensors="pt") processed: dict[str, np.ndarray] = { @@ -879,7 +921,7 @@ def _run_vl_generation( dec_feeds[name] = emb_out[name] prefill_out = dec_session.run(dec_feeds) - logits = prefill_out["logits"] + logits = _suppress_logits(prefill_out["logits"], suppress_ids) next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64) _update_vl_cache(past_cache, prefill_out, config) @@ -941,7 +983,7 @@ def _run_vl_generation( step_feeds[name] = step_emb_out[name] step_out = dec_session.run(step_feeds) - logits = step_out["logits"] + logits = _suppress_logits(step_out["logits"], suppress_ids) next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64) generated.append(next_token) _update_vl_cache(past_cache, step_out, config) @@ -1750,6 +1792,7 @@ def _run_speech_language_generation( ) # --- Step 1: audio encoder --- + suppress_ids = _load_suppress_token_ids(case.model_id, case.trust_remote_code) audio_session = OnnxModelSession(pkg["audio_encoder"], **device_kwargs) try: audio_feeds: dict[str, np.ndarray] = {} @@ -1861,7 +1904,7 @@ def _run_speech_language_generation( dec_feeds[name] = emb_out[name] prefill_out = dec_session.run(dec_feeds) - logits = prefill_out["logits"] + logits = _suppress_logits(prefill_out["logits"], suppress_ids) next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64) _update_vl_cache(past_cache, prefill_out, config) @@ -1910,7 +1953,7 @@ def _run_speech_language_generation( step_feeds[name] = step_emb_out[name] step_out = dec_session.run(step_feeds) - logits = step_out["logits"] + logits = _suppress_logits(step_out["logits"], suppress_ids) next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64) generated.append(next_token) _update_vl_cache(past_cache, step_out, config) diff --git a/tests/integration_test.py b/tests/integration_test.py index a14a4910a..d3f93710a 100644 --- a/tests/integration_test.py +++ b/tests/integration_test.py @@ -5344,3 +5344,404 @@ def test_gemma4_e2b_text_prefill_bf16(): ) assert_logits_close(onnx_logits, hf_logits, rtol=1e-2, atol=5e-3) + + +# --------------------------------------------------------------------------- +# Gemma 4 unified (gemma-4-12B) text backbone prefill parity +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +@pytest.mark.integration_slow +def test_gemma4_unified_12b_text_prefill(): + """gemma-4-12B (``gemma4_unified``) text backbone: ONNX logits match HF. + + Builds Gemma4CausalLMModel from the ``google/gemma-4-12B`` unified text + config and compares a single prefill forward pass against the HuggingFace + ``Gemma4UnifiedForConditionalGeneration`` text path. This exercises the + real 48-layer 12B architecture: dual head_dim (local 256 / global 512), + ``attention_k_eq_v`` with a single global KV head, dual RoPE, and + final-logit softcapping. + + Loads the 24 GB checkpoint in float32 (~48 GB host RAM). Runs on the + device from ``MOBIUS_TEST_DEVICE`` (set ``cuda`` for GPU). Tolerances + atol=2e-2 / rtol=5e-2 account for the deep (48-layer) network and CUDA + floating-point accumulation. + """ + import dataclasses + + import onnx_ir as ir + from transformers import AutoModelForImageTextToText + + from mobius import build_from_module + from mobius._configs import Gemma4Config + from mobius._weight_loading import apply_weights + from mobius.models.gemma4 import Gemma4CausalLMModel + + model_id = "google/gemma-4-12B" + + if not _model_accessible(model_id): + pytest.skip(f"{model_id} not accessible (requires HuggingFace authentication)") + + hf_config = transformers.AutoConfig.from_pretrained(model_id) + hf_full = AutoModelForImageTextToText.from_pretrained( + model_id, + dtype=torch.float32, + low_cpu_mem_usage=True, + ).eval() + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id) + + # Build Gemma4Config from the unified text sub-config (parent supplies + # boa/image/audio token ids and use_bidirectional_attention="vision"). + text_cfg = hf_config.text_config + gemma4_config = Gemma4Config.from_transformers(text_cfg, parent_config=hf_config) + gemma4_config = dataclasses.replace(gemma4_config, dtype=ir.DataType.FLOAT) + + onnx_module = Gemma4CausalLMModel(gemma4_config) + pkg = build_from_module(onnx_module, gemma4_config, task="gemma4-text-generation") + assert "model" in pkg + + preprocessed = onnx_module.preprocess_weights(dict(hf_full.state_dict())) + apply_weights(pkg["model"], preprocessed) + + prompt = "Hello, world!" + tokens = tokenizer(prompt, return_tensors="np") + input_ids = tokens["input_ids"].astype(np.int64) + attention_mask = tokens["attention_mask"].astype(np.int64) + seq_len = input_ids.shape[1] + position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :] + + with torch.no_grad(): + hf_out = hf_full( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + position_ids=torch.from_numpy(position_ids), + ) + hf_logits = hf_out.logits.detach().cpu().numpy() + + session = _make_session(pkg["model"]) + feeds = _make_gemma4_prefill_feeds(gemma4_config, input_ids, attention_mask, position_ids) + # gemma4_unified full-attention layers use a single global KV head; the + # generic feed helper assumes num_key_value_heads, so rebuild KV feeds with + # the per-layer-type KV head count. + lt = gemma4_config.layer_types + for i in range( + gemma4_config.num_hidden_layers - (gemma4_config.num_kv_shared_layers or 0) + ): + is_full = lt[i] == "full_attention" + hd = gemma4_config.global_head_dim if is_full else gemma4_config.head_dim + kvh = ( + gemma4_config.num_global_key_value_heads + if (is_full and gemma4_config.num_global_key_value_heads) + else gemma4_config.num_key_value_heads + ) + feeds[f"past_key_values.{i}.key"] = np.zeros((1, kvh, 0, hd), dtype=np.float32) + feeds[f"past_key_values.{i}.value"] = np.zeros((1, kvh, 0, hd), dtype=np.float32) + onnx_outputs = session.run(feeds) + session.close() + onnx_logits = onnx_outputs["logits"] + + max_diff = float(np.max(np.abs(onnx_logits - hf_logits))) + mean_diff = float(np.mean(np.abs(onnx_logits - hf_logits))) + print( + f"\nGemma4 unified 12B text prefill parity — " + f"max_abs_diff={max_diff:.6f}, mean_abs_diff={mean_diff:.6f}" + ) + assert not np.isnan(onnx_logits).any() + assert_logits_close(onnx_logits, hf_logits, rtol=5e-2, atol=2e-2) + + +@pytest.mark.integration +@pytest.mark.integration_slow +def test_gemma4_unified_12b_multimodal_prefill(): + """gemma-4-12B (``gemma4_unified``) full multimodal prefill parity vs HF. + + Exercises the complete encoder-free multimodal pipeline end to end and + compares against ``Gemma4UnifiedForConditionalGeneration``: + + 1. ``vision_encoder``: raw image patches → 3840-d image features + (compared against HF ``model.get_image_features`` ``pooler_output``). + 2. ``embedding``: scatters the image features into the scaled word + embeddings to produce ``inputs_embeds``. + 3. ``decoder``: 48-layer gemma4 text decoder with vision-block + bidirectional attention (``use_bidirectional_attention="vision"``). + It derives the vision-block ids internally from ``input_ids``. + + The bidirectional reference REQUIRES passing ``mm_token_type_ids`` to HF — + without it HF falls back to a purely causal mask. Loads the 24 GB + checkpoint in float32 (~48 GB host RAM). Runs on ``MOBIUS_TEST_DEVICE`` + (set ``cuda`` for GPU). + """ + import dataclasses + + import onnx_ir as ir + from PIL import Image + from transformers import AutoModelForImageTextToText, AutoProcessor + + from mobius._configs import Gemma4Config + from mobius._weight_loading import _download_weights, apply_weights + from mobius.models.gemma4 import Gemma4UnifiedModel + from mobius.tasks import TASK_REGISTRY + + model_id = "google/gemma-4-12B" + + if not _model_accessible(model_id): + pytest.skip(f"{model_id} not accessible (requires HuggingFace authentication)") + + hf_config = transformers.AutoConfig.from_pretrained(model_id) + hf_full = AutoModelForImageTextToText.from_pretrained( + model_id, + dtype=torch.float32, + low_cpu_mem_usage=True, + ).eval() + processor = AutoProcessor.from_pretrained(model_id) + + # Build a single-image multimodal input. The unified processor expands the + # `<|image|>` placeholder into the right number of soft tokens and emits + # pixel_values [B, N, P^2*3], image_position_ids [B, N, 2] and + # mm_token_type_ids [B, S] (1 = image span). + image = Image.new("RGB", (112, 112), (100, 150, 200)) + proc_inputs = processor( + text=[f"{processor.image_token} Describe the image."], + images=[image], + return_tensors="pt", + ) + input_ids = proc_inputs["input_ids"].numpy().astype(np.int64) + seq_len = input_ids.shape[1] + attention_mask = np.ones((1, seq_len), dtype=np.int64) + position_ids = np.arange(seq_len, dtype=np.int64)[np.newaxis, :] + pixel_values = proc_inputs["pixel_values"] + image_position_ids = proc_inputs["image_position_ids"] + mm_token_type_ids = proc_inputs["mm_token_type_ids"] + + # HF reference: full multimodal forward (mm_token_type_ids enables the + # vision-block bidirectional mask) and the isolated image features. + with torch.no_grad(): + hf_out = hf_full( + input_ids=torch.from_numpy(input_ids), + attention_mask=torch.from_numpy(attention_mask), + pixel_values=pixel_values, + image_position_ids=image_position_ids, + mm_token_type_ids=mm_token_type_ids, + ) + hf_image_features = ( + hf_full.model.get_image_features( + pixel_values, image_position_ids, return_dict=True + ) + .pooler_output.detach() + .cpu() + .numpy() + .astype(np.float32) + ) + hf_logits = hf_out.logits.detach().cpu().numpy() + + # Build the 4-model gemma4_unified package and load real weights. + gemma4_config = Gemma4Config.from_transformers( + hf_config.text_config, parent_config=hf_config + ) + gemma4_config = dataclasses.replace(gemma4_config, dtype=ir.DataType.FLOAT) + module = Gemma4UnifiedModel(gemma4_config) + pkg = TASK_REGISTRY["gemma4-unified"]().build(module, gemma4_config) + # Load the raw safetensors checkpoint (production weight path). The unified + # checkpoint stores `vision_embedder.*` / `embed_vision.embedding_projection.*` + # names, which differ from the runtime module names in hf_full.state_dict(); + # preprocess_weights maps the checkpoint names. + preprocessed = module.preprocess_weights(_download_weights(model_id)) + for name in ("decoder", "vision_encoder", "audio_encoder", "embedding"): + if name in pkg: + apply_weights(pkg[name], preprocessed) + + # Stage 1: vision embedder. + vision_session = _make_session(pkg["vision_encoder"]) + image_features = vision_session.run( + { + "pixel_values": pixel_values.numpy().astype(np.float32), + "pixel_position_ids": image_position_ids.numpy().astype(np.int64), + } + )["image_features"] + vision_session.close() + assert image_features.shape == hf_image_features.shape + vis_cos = float( + np.mean( + np.sum(image_features * hf_image_features, axis=1) + / ( + np.linalg.norm(image_features, axis=1) + * np.linalg.norm(hf_image_features, axis=1) + + 1e-9 + ) + ) + ) + print(f"\nGemma4 unified 12B vision cos_sim={vis_cos:.6f}") + assert vis_cos > 0.999 + + # Stage 2: embedding fusion → inputs_embeds. + embedding_session = _make_session(pkg["embedding"]) + emb_out = embedding_session.run( + { + "input_ids": input_ids, + "image_features": image_features, + "audio_features": np.zeros((0, gemma4_config.hidden_size), dtype=np.float32), + } + ) + embedding_session.close() + inputs_embeds = emb_out["inputs_embeds"] + + # Stage 3: decoder with vision-block bidirectional attention. The decoder + # derives ``block_sequence_ids`` internally from ``input_ids`` (forwarded + # alongside ``inputs_embeds``). + decoder_session = _make_session(pkg["decoder"]) + feeds = { + "inputs_embeds": inputs_embeds, + "attention_mask": attention_mask, + "position_ids": position_ids, + "input_ids": input_ids, + } + layer_types = gemma4_config.layer_types + for i in range(gemma4_config.num_hidden_layers): + is_full = layer_types[i] == "full_attention" + head_dim = gemma4_config.global_head_dim if is_full else gemma4_config.head_dim + kv_heads = ( + gemma4_config.num_global_key_value_heads + if (is_full and gemma4_config.num_global_key_value_heads) + else gemma4_config.num_key_value_heads + ) + feeds[f"past_key_values.{i}.key"] = np.zeros( + (1, kv_heads, 0, head_dim), dtype=np.float32 + ) + feeds[f"past_key_values.{i}.value"] = np.zeros( + (1, kv_heads, 0, head_dim), dtype=np.float32 + ) + onnx_logits = decoder_session.run(feeds)["logits"] + decoder_session.close() + + max_diff = float(np.max(np.abs(onnx_logits - hf_logits))) + last_cos = float( + np.dot(onnx_logits[0, -1], hf_logits[0, -1]) + / (np.linalg.norm(onnx_logits[0, -1]) * np.linalg.norm(hf_logits[0, -1]) + 1e-9) + ) + print( + f"Gemma4 unified 12B multimodal prefill parity — " + f"max_abs_diff={max_diff:.4f}, last_token_cos_sim={last_cos:.6f}" + ) + assert not np.isnan(onnx_logits).any() + assert last_cos > 0.999 + assert onnx_logits[0, -1].argmax() == hf_logits[0, -1].argmax() + + +# --------------------------------------------------------------------------- +# Gemma 4 bidirectional (vision-block) attention mask parity +# --------------------------------------------------------------------------- + + +@pytest.mark.integration +def test_gemma4_bidirectional_mask_parity(): + """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. + """ + import onnx_ir as ir + import torch + from transformers.masking_utils import ( + create_causal_mask, + create_sliding_window_causal_mask, + ) + + from mobius._testing import create_test_builder, create_test_input + from mobius._testing.ort_inference import OnnxModelSession + from mobius.components._common import create_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._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 + attention_mask = np.ones((1, seq_len), dtype=np.int64) + position_ids = np.arange(seq_len, dtype=np.int64)[None, :] + + # Mobius block_sequence_ids from input_ids. + graph, builder = _make_graph() + op = builder.op + iid = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[1, seq_len]) + bsid = _compute_block_sequence_ids(op, iid, image_token_id=image_token_id) + builder.add_output(bsid, "bsid") + 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), + ) + + def mobius_attended(sliding_window): + b, bop, g = create_test_builder() + 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, + sliding_window=sliding_window, + dtype=ir.DataType.FLOAT, + block_sequence_ids=bk, + ) + bias.name = "bias" + g.outputs.append(bias) + out = OnnxModelSession(ir.Model(g, ir_version=10), device="cpu").run( + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "block_sequence_ids": block_ids, + } + )["bias"] + return out[0, 0] > -1.0 # True where the position is attended + + 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( + text_cfg, + inputs_embeds, + torch.from_numpy(attention_mask), + None, + torch.from_numpy(position_ids), + block_sequence_ids=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) + ) diff --git a/tests/synthetic_parity_test.py b/tests/synthetic_parity_test.py index 7dc86c20c..b60b0c34a 100644 --- a/tests/synthetic_parity_test.py +++ b/tests/synthetic_parity_test.py @@ -234,6 +234,12 @@ "exaone", # real HF type is exaone4 "phi3small", # real HF type is phi3 "mistral3", # our implementation maps to mistral; real mistral3 is different + # gemma4_unified_text: mobius-internal alias for the gemma-4-12B text + # backbone (reuses Gemma4CausalLMModel). No matching HF model_type is + # registered with AutoModelForCausalLM, so a reference model cannot be + # constructed here. Text parity is covered by the real-weight + # integration test (test_gemma4_unified_12b_text_prefill). + "gemma4_unified_text", # falcon_h1: our ONNX uses FalconCausalLMModel (ALiBi attention), not the # real HF FalconH1 (Mamba2+SSM hybrid). Comparing against HF would be apples-to-oranges. "falcon_h1",