From e198a68ff4189717950a397c9110c49f42c79514 Mon Sep 17 00:00:00 2001 From: rui-ren Date: Wed, 6 May 2026 22:53:23 +0000 Subject: [PATCH 1/4] update-genai integration --- src/mobius/_model_package.py | 7 + src/mobius/integrations/ort_genai/__init__.py | 6 +- .../integrations/ort_genai/auto_export.py | 172 ++++++++++++++---- .../ort_genai/auto_export_test.py | 112 ++++++++++++ 4 files changed, 263 insertions(+), 34 deletions(-) diff --git a/src/mobius/_model_package.py b/src/mobius/_model_package.py index 7c099296..d79f1703 100644 --- a/src/mobius/_model_package.py +++ b/src/mobius/_model_package.py @@ -73,6 +73,13 @@ def save( ``model.onnx`` in *directory*. When multiple models are present, each is saved in its own subfolder as ``{name}/model.onnx``. + .. note:: + This method writes ONNX files only. If you need a directory that + ``onnxruntime-genai`` can load (i.e. with ``genai_config.json`` and + tokenizer files), use + :func:`mobius.integrations.ort_genai.export_package` instead — it + wraps :meth:`save` with the ORT-GenAI config-generation step. + Args: directory: Path to the output directory (created if needed). external_data: External data format. ``"onnx"`` (default) saves diff --git a/src/mobius/integrations/ort_genai/__init__.py b/src/mobius/integrations/ort_genai/__init__.py index ba240796..97454462 100644 --- a/src/mobius/integrations/ort_genai/__init__.py +++ b/src/mobius/integrations/ort_genai/__init__.py @@ -7,7 +7,10 @@ layers remain runtime-agnostic. """ -from mobius.integrations.ort_genai.auto_export import write_ort_genai_config +from mobius.integrations.ort_genai.auto_export import ( + export_package, + write_ort_genai_config, +) from mobius.integrations.ort_genai.ep_config import ( make_genai_decoder_config, make_kv_cache_dim_name, @@ -20,6 +23,7 @@ __all__ = [ "GenaiConfigGenerator", + "export_package", "write_ort_genai_config", "make_genai_decoder_config", "make_kv_cache_dim_name", diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 3dd6ac51..42dbc3f9 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -3,34 +3,42 @@ """Auto-export pipeline for onnxruntime-genai. -Two entry points: +Three entry points, in order of increasing convenience: -- :func:`write_ort_genai_config` — programmatic API. Takes an already-built - :class:`~mobius._model_package.ModelPackage` (with weights) and writes the - ORT-GenAI config artifacts (``genai_config.json``, tokenizer files, - ``processor_config.json`` / ``image_processor.json``) alongside the ONNX models. +- :func:`write_ort_genai_config` — config-only API. Takes an already-built + :class:`~mobius._model_package.ModelPackage` (with weights already saved + separately) and writes the ORT-GenAI config artifacts + (``genai_config.json``, tokenizer files, ``processor_config.json`` / + ``image_processor.json``) into a directory. -- :func:`auto_export` — end-to-end convenience function. Builds the model - from a HuggingFace ID, saves the ONNX files, then calls - :func:`write_ort_genai_config` to write the config artifacts. +- :func:`export_package` — save+config API. Takes an already-built + ``ModelPackage`` and writes both the ONNX models AND the ORT-GenAI config + artifacts in one call. Use this when you built the package manually + (e.g. with custom dtype / quantization). -Both functions produce a directory that ``onnxruntime-genai`` can load -directly. +- :func:`auto_export` — end-to-end API. Builds the model from a HuggingFace + ID and calls :func:`export_package`. Use this for the common + HF-model-id → ORT-GenAI-directory case. + +All three produce a directory that ``onnxruntime-genai`` can load directly. Example:: - # Programmatic API — build first, then export configs - from mobius import build + # Config-only — when ONNX is already on disk from mobius.integrations.ort_genai import write_ort_genai_config + write_ort_genai_config(pkg, "/output/qwen3", hf_model_id="Qwen/Qwen3-0.6B") + + # Save + config — when you have a built package in memory + from mobius import build + from mobius.integrations.ort_genai import export_package pkg = build("Qwen/Qwen3-0.6B", load_weights=True) - pkg.save("/output/qwen3") - write_ort_genai_config(pkg, "/output/qwen3", hf_model_id="Qwen/Qwen3-0.6B") + export_package(pkg, "/output/qwen3", hf_model_id="Qwen/Qwen3-0.6B", ep="cuda") - # End-to-end convenience + # End-to-end — when you only have an HF model id from mobius.integrations.ort_genai.auto_export import auto_export - auto_export("Qwen/Qwen3-0.6B", "/output/qwen3") + auto_export("Qwen/Qwen3-0.6B", "/output/qwen3", ep="cuda") """ from __future__ import annotations @@ -852,6 +860,116 @@ def write_ort_genai_config( return result +def export_package( + pkg: ModelPackage, + output_dir: str, + *, + hf_model_id: str | None = None, + ep: str = "cpu", + context_length: int = 4096, + local_config_dir: str | None = None, + external_data: str = "onnx", + progress_bar: bool = True, +) -> dict[str, str]: + """Save an already-built ModelPackage as a complete ORT-GenAI directory. + + This is the convenience function for users who built a ``ModelPackage`` + themselves (e.g. with custom dtype / quantization / weight overrides) and + want a single call that produces an ``onnxruntime-genai``-loadable + directory. It calls :meth:`ModelPackage.save` followed by + :func:`write_ort_genai_config`. + + For the end-to-end case where you start from a HuggingFace model id, use + :func:`auto_export` instead — it builds the package for you. + + Args: + pkg: Already-built :class:`~mobius._model_package.ModelPackage` with + weights applied and ``config`` set. Must contain all components + you want exported; partial exports are not supported because the + generated ``genai_config.json`` would reference components that + do not exist on disk. Build a separate filtered package if you + need a subset. + output_dir: Output directory (created if needed). + hf_model_id: HuggingFace model ID for tokenizer download / token-id + resolution. When ``None``, token IDs are read from ``pkg.config`` + and tokenizer files are not copied (unless ``local_config_dir`` + is provided). + ep: Execution provider written to ``session_options`` in + ``genai_config.json`` (e.g. ``"cpu"``, ``"cuda"``, ``"dml"``, + ``"webgpu"``, ``"trt-rtx"``). + context_length: Minimum context length written to ``genai_config.json``. + Overridden upward by ``pkg.config.max_position_embeddings`` when + larger. + local_config_dir: Local model directory to copy tokenizer files from + when ``hf_model_id`` is ``None``. + external_data: External-data format passed to :meth:`ModelPackage.save` + (``"onnx"`` or ``"safetensors"``). + progress_bar: Whether to show the save progress bar. + + Returns: + Manifest dict mapping artifact names to paths:: + + { + "model": "/output/model.onnx", # or per-component paths + "genai_config": "/output/genai_config.json", + "tokenizer.json": "/output/tokenizer.json", + ... + } + + Raises: + ValueError: If ``pkg.config`` is ``None`` (required for genai_config + generation; e.g. diffusion models have no config and are not + supported). + + Example:: + + from mobius import build + from mobius.integrations.ort_genai import export_package + + pkg = build("Qwen/Qwen3-0.6B", load_weights=True) + export_package(pkg, "/output/qwen3", hf_model_id="Qwen/Qwen3-0.6B", ep="cuda") + """ + # Preflight: fail fast before writing ONNX so the user doesn't end up + # with a half-exported directory containing only the model file. + if getattr(pkg, "config", None) is None: + raise ValueError( + "export_package requires ModelPackage.config to be set. " + "This is set automatically when building with mobius.build(). " + "Diffusion models (which have no config) are not supported — " + "use ModelPackage.save() directly for those." + ) + + os.makedirs(output_dir, exist_ok=True) + + # 1. Save ONNX models + weights + logger.info("Saving ONNX models to %s", output_dir) + pkg.save( + output_dir, + external_data=external_data, + progress_bar=progress_bar, + ) + + # 2. Write ORT-GenAI config artifacts + result = write_ort_genai_config( + pkg, + output_dir, + hf_model_id=hf_model_id, + ep=ep, + context_length=context_length, + local_config_dir=local_config_dir, + ) + + # 3. Add ONNX paths to the manifest + if len(pkg) == 1: + result["model"] = os.path.join(output_dir, "model.onnx") + else: + for name in pkg: + result[name] = os.path.join(output_dir, name, "model.onnx") + + logger.info("Export complete: %d artifacts", len(result)) + return result + + def auto_export( model_id: str, output_dir: str, @@ -918,28 +1036,16 @@ def auto_export( "Diffusion models are not yet supported." ) - # Save ONNX models - logger.info("Saving ONNX models to %s", output_dir) - pkg.save( - output_dir, - external_data=external_data, - progress_bar=progress_bar, - ) - - # Write ORT-GenAI config artifacts (genai_config.json, tokenizer, processor) - result = write_ort_genai_config( + # Delegate save + config generation to the integration helper + result = export_package( pkg, output_dir, hf_model_id=model_id, ep=ep, context_length=context_length, + external_data=external_data, + progress_bar=progress_bar, ) - # Add ONNX model paths to manifest - if len(pkg) == 1: - result["model"] = os.path.join(output_dir, "model.onnx") - else: - for name in pkg: - result[name] = os.path.join(output_dir, name, "model.onnx") - logger.info("Export complete: %d artifacts", len(result)) + return result diff --git a/src/mobius/integrations/ort_genai/auto_export_test.py b/src/mobius/integrations/ort_genai/auto_export_test.py index 0339386e..558fedcc 100644 --- a/src/mobius/integrations/ort_genai/auto_export_test.py +++ b/src/mobius/integrations/ort_genai/auto_export_test.py @@ -910,6 +910,118 @@ class FakeConfig: assert data["model"]["eos_token_id"] == [1, 106] +class TestExportPackage: + """Tests for export_package() — the save+config integration helper.""" + + @staticmethod + def _make_pkg(): + return _make_fake_llm_pkg("qwen2") + + def test_writes_both_onnx_and_genai_config(self, tmp_path, monkeypatch): + """export_package calls pkg.save AND writes genai_config.json.""" + from mobius.integrations.ort_genai.auto_export import export_package + + pkg = self._make_pkg() + save_calls = [] + + def fake_save(self, directory, **kwargs): + save_calls.append((directory, kwargs)) + + monkeypatch.setattr(pkg.__class__, "save", fake_save) + + result = export_package(pkg, str(tmp_path)) + + # pkg.save called exactly once with the output dir + assert len(save_calls) == 1 + assert save_calls[0][0] == str(tmp_path) + # genai_config artifact is in the manifest + assert "genai_config" in result + assert os.path.isfile(result["genai_config"]) + # ONNX path is in the manifest (single-component package) + assert result["model"] == os.path.join(str(tmp_path), "model.onnx") + + def test_propagates_save_kwargs(self, tmp_path, monkeypatch): + """external_data and progress_bar are forwarded to pkg.save.""" + from mobius.integrations.ort_genai.auto_export import export_package + + pkg = self._make_pkg() + save_calls = [] + + def fake_save(self, directory, **kwargs): + save_calls.append(kwargs) + + monkeypatch.setattr(pkg.__class__, "save", fake_save) + + export_package( + pkg, + str(tmp_path), + external_data="safetensors", + progress_bar=False, + ) + + assert save_calls[0]["external_data"] == "safetensors" + assert save_calls[0]["progress_bar"] is False + + def test_propagates_genai_config_kwargs(self, tmp_path, monkeypatch): + """ep, context_length, and hf_model_id reach genai_config.json.""" + from mobius.integrations.ort_genai.auto_export import export_package + + pkg = self._make_pkg() + monkeypatch.setattr(pkg.__class__, "save", lambda self, d, **kw: None) + + result = export_package( + pkg, + str(tmp_path), + ep="cuda", + context_length=8192, + ) + + with open(result["genai_config"]) as f: + data = json.load(f) + # ep="cuda" should produce CUDA provider_options + provider_opts = data["model"]["decoder"]["session_options"]["provider_options"] + assert any("cuda" in po for po in provider_opts) + # context_length should bump max_length + assert data["search"]["max_length"] == 8192 + + def test_preflights_missing_config(self, tmp_path, monkeypatch): + """Raises ValueError BEFORE writing any ONNX when pkg.config is None. + + This avoids leaving a half-exported directory with model.onnx but + no genai_config.json. + """ + from mobius._model_package import ModelPackage + from mobius.integrations.ort_genai.auto_export import export_package + + pkg = ModelPackage({"model": mock.MagicMock()}, config=None) + save_called = [] + + def fake_save(self, *a, **kw): + save_called.append(True) + + monkeypatch.setattr(pkg.__class__, "save", fake_save) + + with pytest.raises(ValueError, match="config"): + export_package(pkg, str(tmp_path)) + + # save was NOT called — preflight failed before any I/O + assert save_called == [] + + def test_returns_manifest_with_all_artifacts(self, tmp_path, monkeypatch): + """Returned manifest contains ONNX paths AND config artifacts.""" + from mobius.integrations.ort_genai.auto_export import export_package + + pkg = self._make_pkg() + monkeypatch.setattr(pkg.__class__, "save", lambda self, d, **kw: None) + + result = export_package(pkg, str(tmp_path)) + + # Manifest is non-empty and includes both kinds of artifacts + assert isinstance(result, dict) + assert "model" in result + assert "genai_config" in result + + class TestGemma4GenaiConfig: """Tests for Gemma4-specific genai_config generation via graph introspection.""" From b65192dddaae9385618a2345695cfdc525fdf993 Mon Sep 17 00:00:00 2001 From: rui-ren Date: Fri, 8 May 2026 00:56:43 +0000 Subject: [PATCH 2/4] LM-head-prune --- src/mobius/_flags.py | 39 ++++++++++++++++++++ src/mobius/models/_models_test.py | 59 ++++++++++++++++++++++++++++++- src/mobius/models/base.py | 13 +++++++ 3 files changed, 110 insertions(+), 1 deletion(-) diff --git a/src/mobius/_flags.py b/src/mobius/_flags.py index 3e61872c..832f9dcb 100644 --- a/src/mobius/_flags.py +++ b/src/mobius/_flags.py @@ -89,6 +89,13 @@ class _Flags: - ``False`` - Use GQA for KV-shared layers on CUDA EP. Requires ORT GQA new_kv_length=0 support. + * - ``prune_lm_head`` + - ``MOBIUS_PRUNE_LM_HEAD`` + - ``False`` + - Insert ``Gather(axis=1, indices=[-1])`` before the LM head MatMul + so logits are only computed for the last token. Speeds up prefill + for chat-only workloads but breaks logprob scoring and + speculative decoding. """ suppress_dedup_warning: bool = dataclasses.field( @@ -137,6 +144,38 @@ class _Flags: opset 24 kernel support. """ + prune_lm_head: bool = dataclasses.field( + default_factory=lambda: _env_bool("MOBIUS_PRUNE_LM_HEAD", False) + ) + """Insert ``Gather(axis=1, indices=[-1])`` before the LM head MatMul to + select only the last token's hidden state. + + When ``False`` (default), the LM head computes logits for every token + in the sequence (output shape ``[B, S, vocab]``). This preserves all + use cases including logprob scoring, perplexity evaluation, + speculative-decoding verification, and multi-token-at-a-time generation. + + When ``True``, only the last token's logits are computed (output shape + ``[B, 1, vocab]``), avoiding the full ``[B, S, vocab]`` MatMul during + prefill. This is a meaningful prefill speedup for chat / single-token + generation workloads on models with large vocabularies (e.g. Qwen at + 151936 tokens), but **breaks** any workflow that needs per-token logits. + + Mirrors the ``prune_lm_head`` extra option in onnxruntime-genai's + Model Builder. Set ``MOBIUS_PRUNE_LM_HEAD=1`` to enable for chat-only + deployments. + + .. note:: + **Compatibility:** This flag only takes effect for models that use + the base :class:`~mobius.models.base.CausalLMModel.forward` + implementation — Llama, Mistral, Qwen2 / 2.5 / 3, and other + standard decoder-only architectures. Models that override + ``forward()`` (GPT-J, CodeGen, Phi, NanoChat, DOGE, Gemma3n, + Apertus, GPT-2 family, Mamba, InternLM2, Qwen3.5, etc.) **silently + ignore this flag** because they assemble logits via their own code + paths. Coverage will be extended in follow-up work as needed. + """ + # Global singleton — import and use this directly. flags = _Flags() diff --git a/src/mobius/models/_models_test.py b/src/mobius/models/_models_test.py index 17e2ab8e..f3d40013 100644 --- a/src/mobius/models/_models_test.py +++ b/src/mobius/models/_models_test.py @@ -138,7 +138,64 @@ def test_build_with_task_string(self): assert model.graph.num_nodes() > 0 -class TestModelRegistry: +class TestPruneLmHead: + """Tests for the ``prune_lm_head`` flag in CausalLMModel. + + When ``True``, a ``Gather + Unsqueeze`` is inserted before the LM head + MatMul so logits are produced for only the last sequence position. + Mirrors onnxruntime-genai Model Builder's ``prune_lm_head`` opt-in. + """ + + def _build(self) -> ir.Model: + config = make_config() + module = CausalLMModel(config) + return build_from_module(module, config)["model"] + + def test_default_emits_no_lm_head_pruning(self): + """Default flag: graph contains no Gather feeding the lm_head MatMul.""" + from mobius._flags import override_flags + + with override_flags(prune_lm_head=False): + model = self._build() + + # The logits output should track full sequence_length in dim 1 + logits = next(v for v in model.graph.outputs if v.name == "logits") + # Full path: shape[1] is a symbolic dim ("sequence_length"), not 1 + seq_dim = logits.shape[1] + # symbolic dims either expose an integer value of None or a symbolic + # name; the default (un-pruned) path must NOT be the literal int 1 + assert seq_dim != 1, ( + f"Expected dynamic sequence_length in logits dim 1, got {seq_dim!r}" + ) + + def test_prune_emits_gather_before_lm_head(self): + """With prune_lm_head=True, logits dim 1 collapses to 1.""" + from mobius._flags import override_flags + + with override_flags(prune_lm_head=True): + model = self._build() + + logits = next(v for v in model.graph.outputs if v.name == "logits") + seq_dim = logits.shape[1] + # Pruned path: shape[1] is the literal integer 1 + assert seq_dim == 1, ( + f"Expected logits dim 1 to be 1 after pruning, got {seq_dim!r}" + ) + + def test_prune_does_not_change_input_shapes(self): + """Pruning only affects output; input_ids / attention_mask still + have dynamic sequence_length so the model accepts arbitrary prompts.""" + from mobius._flags import override_flags + + with override_flags(prune_lm_head=True): + model = self._build() + + input_ids = next(v for v in model.graph.inputs if v.name == "input_ids") + # input dim 1 (sequence_length) should still be dynamic, not 1 + assert input_ids.shape[1] != 1 + + + def test_registry_not_empty(self): assert len(registry) > 0 diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index c44b7e59..ecbccbfd 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -20,6 +20,7 @@ from mobius._build_context import ep_capabilities, get_build_dtype from mobius._configs import ArchitectureConfig, CausalLMConfig +from mobius._flags import flags from mobius._weight_utils import ( preprocess_awq_weights, preprocess_gptq_weights, @@ -223,6 +224,18 @@ def forward( position_ids=position_ids, past_key_values=past_key_values, ) + if flags.prune_lm_head: + # Select only the last token's hidden state before the LM head. + # Hidden shape goes [B, S, H] → [B, H] → [B, 1, H], turning the + # downstream MatMul output from [B, S, vocab] into [B, 1, vocab]. + # Saves the [B, (S-1), vocab] portion of work during prefill. + # + # Mirrors onnxruntime-genai Model Builder's `prune_lm_head` flag + # (see make_lm_head() in builders/base.py). Breaks workflows that + # need per-token logits — see flag docstring. + indices = op.Constant(value_ints=[-1]) + last_hidden = op.Gather(hidden_states, indices, axis=1) + hidden_states = op.Unsqueeze(last_hidden, op.Constant(value_ints=[1])) logits = self.lm_head(op, hidden_states) return logits, present_key_values From 573edb8f059205066eaf44eb8f8e3988faf2c68c Mon Sep 17 00:00:00 2001 From: rui-ren Date: Fri, 8 May 2026 02:23:50 +0000 Subject: [PATCH 3/4] conflict --- .../integrations/ort_genai/auto_export.py | 1 - src/mobius/models/_models_test.py | 29 ++++++++++++++----- src/mobius/models/base.py | 16 ++++++---- 3 files changed, 31 insertions(+), 15 deletions(-) diff --git a/src/mobius/integrations/ort_genai/auto_export.py b/src/mobius/integrations/ort_genai/auto_export.py index 336a7753..869cc9a8 100644 --- a/src/mobius/integrations/ort_genai/auto_export.py +++ b/src/mobius/integrations/ort_genai/auto_export.py @@ -1048,5 +1048,4 @@ def auto_export( progress_bar=progress_bar, ) - return result diff --git a/src/mobius/models/_models_test.py b/src/mobius/models/_models_test.py index f3d40013..35284a14 100644 --- a/src/mobius/models/_models_test.py +++ b/src/mobius/models/_models_test.py @@ -152,35 +152,48 @@ def _build(self) -> ir.Model: return build_from_module(module, config)["model"] def test_default_emits_no_lm_head_pruning(self): - """Default flag: graph contains no Gather feeding the lm_head MatMul.""" + """Default flag: graph emits full [B, S, vocab] logits.""" from mobius._flags import override_flags with override_flags(prune_lm_head=False): model = self._build() - # The logits output should track full sequence_length in dim 1 logits = next(v for v in model.graph.outputs if v.name == "logits") - # Full path: shape[1] is a symbolic dim ("sequence_length"), not 1 + # Logits must remain rank-3 [B, S, vocab] + assert len(logits.shape) == 3, ( + f"Expected rank-3 logits [B, S, V], got rank {len(logits.shape)}: " + f"shape={list(logits.shape)!r}" + ) + # The full path: shape[1] is a symbolic dim ("sequence_length"), not 1 seq_dim = logits.shape[1] - # symbolic dims either expose an integer value of None or a symbolic - # name; the default (un-pruned) path must NOT be the literal int 1 assert seq_dim != 1, ( f"Expected dynamic sequence_length in logits dim 1, got {seq_dim!r}" ) + # Last dim is the vocabulary size + config = make_config() + assert logits.shape[2] == config.vocab_size def test_prune_emits_gather_before_lm_head(self): - """With prune_lm_head=True, logits dim 1 collapses to 1.""" + """With prune_lm_head=True, logits dim 1 collapses to 1 (still rank-3).""" from mobius._flags import override_flags with override_flags(prune_lm_head=True): model = self._build() logits = next(v for v in model.graph.outputs if v.name == "logits") + # Logits must still be rank-3 [B, 1, vocab] (NOT rank-4 [B, 1, 1, V]) + assert len(logits.shape) == 3, ( + f"Expected rank-3 logits [B, 1, V] after pruning, got rank " + f"{len(logits.shape)}: shape={list(logits.shape)!r}" + ) + # Pruned: dim 1 must be the literal integer 1 seq_dim = logits.shape[1] - # Pruned path: shape[1] is the literal integer 1 assert seq_dim == 1, ( f"Expected logits dim 1 to be 1 after pruning, got {seq_dim!r}" ) + # Last dim is still the vocabulary size + config = make_config() + assert logits.shape[2] == config.vocab_size def test_prune_does_not_change_input_shapes(self): """Pruning only affects output; input_ids / attention_mask still @@ -195,7 +208,7 @@ def test_prune_does_not_change_input_shapes(self): assert input_ids.shape[1] != 1 - +class TestModelRegistry: def test_registry_not_empty(self): assert len(registry) > 0 diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index ecbccbfd..6dcde73d 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -226,16 +226,20 @@ def forward( ) if flags.prune_lm_head: # Select only the last token's hidden state before the LM head. - # Hidden shape goes [B, S, H] → [B, H] → [B, 1, H], turning the - # downstream MatMul output from [B, S, vocab] into [B, 1, vocab]. - # Saves the [B, (S-1), vocab] portion of work during prefill. + # Use a scalar (rank-0) index so Gather collapses dim 1: + # [B, S, H] --Gather(axis=1, idx=-1, scalar)--> [B, H] + # --Unsqueeze(axis=1)--> [B, 1, H] + # The downstream MatMul then produces [B, 1, vocab] instead of + # [B, S, vocab], saving the [B, S-1, vocab] of work during prefill. # # Mirrors onnxruntime-genai Model Builder's `prune_lm_head` flag # (see make_lm_head() in builders/base.py). Breaks workflows that # need per-token logits — see flag docstring. - indices = op.Constant(value_ints=[-1]) - last_hidden = op.Gather(hidden_states, indices, axis=1) - hidden_states = op.Unsqueeze(last_hidden, op.Constant(value_ints=[1])) + last_idx = op.Constant(value_int=-1) # scalar (rank-0) INT64 + last_hidden = op.Gather(hidden_states, last_idx, axis=1) # [B, H] + hidden_states = op.Unsqueeze( + last_hidden, op.Constant(value_ints=[1]) + ) # [B, 1, H] logits = self.lm_head(op, hidden_states) return logits, present_key_values From bc1f6e8c1cbdf8e5c418dd5289bf699c6b7c4492 Mon Sep 17 00:00:00 2001 From: rui-ren Date: Fri, 8 May 2026 02:58:02 +0000 Subject: [PATCH 4/4] source --- src/mobius/models/_models_test.py | 7 +++---- src/mobius/models/base.py | 4 +--- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/src/mobius/models/_models_test.py b/src/mobius/models/_models_test.py index 35284a14..90f0482d 100644 --- a/src/mobius/models/_models_test.py +++ b/src/mobius/models/_models_test.py @@ -188,16 +188,15 @@ def test_prune_emits_gather_before_lm_head(self): ) # Pruned: dim 1 must be the literal integer 1 seq_dim = logits.shape[1] - assert seq_dim == 1, ( - f"Expected logits dim 1 to be 1 after pruning, got {seq_dim!r}" - ) + assert seq_dim == 1, f"Expected logits dim 1 to be 1 after pruning, got {seq_dim!r}" # Last dim is still the vocabulary size config = make_config() assert logits.shape[2] == config.vocab_size def test_prune_does_not_change_input_shapes(self): """Pruning only affects output; input_ids / attention_mask still - have dynamic sequence_length so the model accepts arbitrary prompts.""" + have dynamic sequence_length so the model accepts arbitrary prompts. + """ from mobius._flags import override_flags with override_flags(prune_lm_head=True): diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py index 6dcde73d..80c8c87b 100644 --- a/src/mobius/models/base.py +++ b/src/mobius/models/base.py @@ -237,9 +237,7 @@ def forward( # need per-token logits — see flag docstring. last_idx = op.Constant(value_int=-1) # scalar (rank-0) INT64 last_hidden = op.Gather(hidden_states, last_idx, axis=1) # [B, H] - hidden_states = op.Unsqueeze( - last_hidden, op.Constant(value_ints=[1]) - ) # [B, 1, H] + hidden_states = op.Unsqueeze(last_hidden, op.Constant(value_ints=[1])) # [B, 1, H] logits = self.lm_head(op, hidden_states) return logits, present_key_values