update lm head last token pruning

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()

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

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",

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