update lm head last token pruning

CHANGELOG.md

@@ -15,15 +15,26 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   single Rust/cargo-style option. Accepts a comma-separated list and may be
   repeated (`--features fp8-kv-cache,static-cache` or `--features fp8-kv-cache
   --features static-cache`). Available features: `static-cache`, `fp8-kv-cache`,
-  `text-only`. Unknown feature names are rejected with an error listing the
-  valid set.
+  `prune-lm-head`, `text-only`. Unknown feature names are rejected with an error
+  listing the valid set.
 
 #### Changed
 
 - The boolean flags `--static-cache`, `--fp8-kv-cache`, and `--text-only` have
   been **removed** in favor of the equivalent `--features` value. Companion
   value args (`--max-seq-len`, `--kv-cache-scale-file`) are unchanged.
 
+### Final-token LM-head pruning (`--features prune-lm-head`)
+
+#### Added
+
+- `build(prune_lm_head=True)` and `mobius build --features prune-lm-head`
+  select the final hidden-state position before the LM-head projection, reducing
+  prefill logits from `[B, S, vocab]` to `[B, 1, vocab]`. This avoids computing
+  unused per-token logits for single-token autoregressive generation. Models
+  with custom forward paths that do not support pre-projection pruning fail
+  explicitly instead of silently producing an unoptimized graph.
+
 ### FP8 (E4M3) KV-cache export (`--features fp8-kv-cache`)
 
 #### Added

docs/cli_reference.md

@@ -175,6 +175,7 @@ option. Pass a comma-separated list (and/or repeat the flag):
 
 ```
 --features fp8-kv-cache,static-cache
+--features prune-lm-head
 --features text-only
 ```
 
@@ -184,6 +185,7 @@ Available features:
 |---------|--------|
 | `static-cache` | Pre-allocate fixed-size KV cache buffers using `TensorScatter` (pair with `--max-seq-len N`). Requires `DecoderLayer` / `MoEDecoderLayer` models. Cannot combine with `--task`. |
 | `fp8-kv-cache` | Store the `GroupQueryAttention` KV cache as `FLOAT8E4M3FN` (per-tensor E4M3), halving KV-cache memory. Requires a GQA build (e.g. `--ep cuda --dtype f16`) and an ORT runtime with the FP8 KV-cache kernel (SM89+). Pair with `--kv-cache-scale-file` for calibrated scales. |
+| `prune-lm-head` | Select the final hidden-state position before the LM-head projection and emit logits shaped `[B, 1, vocab]`. Supported by models using the base `CausalLMModel.forward()` path; unsupported custom forwards fail explicitly. Use only when the downstream workflow does not need per-token logits. |
 | `text-only` | Export the text backbone of a multimodal checkpoint as a standalone decoder-only LLM (see below). |
 
 The legacy boolean flags `--static-cache`, `--fp8-kv-cache`, and
@@ -195,6 +197,9 @@ mobius build --model meta-llama/Llama-3.2-1B output/ \
 
 mobius build --model Qwen/Qwen2.5-0.5B output/ \
     --ep cuda --dtype f16 --features fp8-kv-cache
+
+mobius build --model meta-llama/Llama-3.2-1B output/ \
+    --features prune-lm-head
 ```
 
 ### Static Cache (`--features static-cache`)

src/mobius/__main__.py

@@ -39,6 +39,7 @@
 _BUILD_FEATURES: dict[str, str] = {
     "static-cache": "static_cache",
     "fp8-kv-cache": "fp8_kv_cache",
+    "prune-lm-head": "prune_lm_head",
     "text-only": "text_only",
 }
 
@@ -222,6 +223,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
     # FP8 KV cache: resolve the optional per-layer scale file up front so both
     # the --config and --model build paths can pass the same scales.
     fp8_kv_cache = getattr(args, "fp8_kv_cache", False)
+    prune_lm_head = getattr(args, "prune_lm_head", False)
     kv_cache_scales: dict[int, tuple[float, float]] | None = None
     scale_file = getattr(args, "kv_cache_scale_file", None)
     if scale_file is not None and not fp8_kv_cache:
@@ -311,6 +313,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             execution_provider=execution_provider,
             fp8_kv_cache=fp8_kv_cache,
             kv_cache_scales=kv_cache_scales,
+            prune_lm_head=prune_lm_head,
         )
         for name, model in pkg.items():
             model.graph.name = f"{config_path}/{name}"
@@ -340,6 +343,7 @@ def _resolve_static_cache_task(model_type: str) -> ModelTask:
             text_only=args.text_only,
             fp8_kv_cache=fp8_kv_cache,
             kv_cache_scales=kv_cache_scales,
+            prune_lm_head=prune_lm_head,
         )
 
     _save_package(pkg, output_dir, args, optimize, component_filter)

src/mobius/_build_context.py

@@ -35,6 +35,8 @@
     "build_context",
     "ep_capabilities",
     "get_build_dtype",
+    "is_lm_head_pruning_enabled",
+    "lm_head_pruning",
 ]
 
 _DEFAULT_CAPABILITIES = EpCapabilities(name="default")
@@ -45,6 +47,9 @@
 _current_dtype: contextvars.ContextVar[ir.DataType] = contextvars.ContextVar(
     "mobius_build_dtype", default=ir.DataType.FLOAT
 )
+_prune_lm_head: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "mobius_prune_lm_head", default=False
+)
 
 
 @contextmanager
@@ -113,3 +118,18 @@ def get_build_dtype() -> ir.DataType:
             ...
     """
     return _current_dtype.get()
+
+
+@contextmanager
+def lm_head_pruning(enabled: bool) -> Iterator[None]:
+    """Enable or disable pre-projection LM-head pruning during graph construction."""
+    token = _prune_lm_head.set(enabled)
+    try:
+        yield
+    finally:
+        _prune_lm_head.reset(token)
+
+
+def is_lm_head_pruning_enabled() -> bool:
+    """Return whether the active causal-LM task requests final-token logits only."""
+    return _prune_lm_head.get()

src/mobius/_build_context_test.py

@@ -8,7 +8,13 @@
 import onnx_ir as ir
 import pytest
 
-from mobius._build_context import build_context, ep_capabilities, get_build_dtype
+from mobius._build_context import (
+    build_context,
+    ep_capabilities,
+    get_build_dtype,
+    is_lm_head_pruning_enabled,
+    lm_head_pruning,
+)
 from mobius._execution_providers import EpCapabilities, ep_registry
 
 _CUDA_CAPABILITIES = EpCapabilities(
@@ -31,6 +37,9 @@ def test_default_dtype_is_float32(self):
         """No context active → returns FLOAT."""
         assert get_build_dtype() == ir.DataType.FLOAT
 
+    def test_lm_head_pruning_is_disabled(self):
+        assert not is_lm_head_pruning_enabled()
+
     def test_default_capabilities_has_no_fusions(self):
         """Default EP has no GQA dtypes (portable ONNX)."""
         capabilities = ep_capabilities()
@@ -65,6 +74,11 @@ def test_capabilities_restored_on_exception(self):
         assert ep_capabilities().name == "default"
         assert get_build_dtype() == ir.DataType.FLOAT
 
+    def test_lm_head_pruning_restored_after_context(self):
+        with lm_head_pruning(True):
+            assert is_lm_head_pruning_enabled()
+        assert not is_lm_head_pruning_enabled()
+
 
 class TestBuildContextNesting:
     def test_inner_context_shadows_outer(self):

src/mobius/_builder.py

@@ -106,6 +106,28 @@ def _cast_module_dtype(module: nn.Module, dtype: ir.DataType) -> None:
             param.const_value = tensor_adapters.TorchTensor(cast_tensor)
 
 
+def _enable_pruned_lm_head_task(task: str | ModelTask) -> str | ModelTask:
+    """Return a task equivalent to *task* with LM-head pruning enabled."""
+    from mobius.tasks import CausalLMTask, HybridCausalLMTask
+
+    if task == "text-generation":
+        return CausalLMTask(prune_lm_head=True)
+    if task == "hybrid-text-generation":
+        return HybridCausalLMTask(prune_lm_head=True)
+    if isinstance(task, CausalLMTask):
+        return CausalLMTask(
+            static_cache=getattr(task, "_static_cache", False),
+            max_seq_len=getattr(task, "_max_seq_len", None),
+            prune_lm_head=True,
+        )
+    if isinstance(task, HybridCausalLMTask):
+        return HybridCausalLMTask(prune_lm_head=True)
+    raise ValueError(
+        "prune_lm_head=True is only supported for text-generation and "
+        "hybrid-text-generation tasks."
+    )
+
+
 # Map ModelPackage entry names to semantic model roles.
 # GQA fusion is only applied to "decoder" role models.
 _MODEL_ROLE_MAP: dict[str, str] = {
@@ -137,6 +159,7 @@ def build_from_module(
     trace_optimization: bool = False,
     fp8_kv_cache: bool = False,
     kv_cache_scales: dict[int, tuple[float, float]] | None = None,
+    prune_lm_head: bool = False,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a module instance and config.
 
@@ -177,6 +200,10 @@ def build_from_module(
             per-tensor FP8 scales (from offline calibration), used only when
             ``fp8_kv_cache`` is ``True``. Layers absent from the map use a unit
             scale of ``1.0``.
+        prune_lm_head: When ``True``, reduce decoder logits to the final token
+            via the causal-LM task so runtimes can avoid full prefill LM-head
+            projection. Only supported by ``text-generation`` and
+            ``hybrid-text-generation`` tasks.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -210,6 +237,8 @@ def forward(self, op, input_ids, attention_mask,
     # are included — their graph inputs are kept at f32 (matching GenAI's
     # image processor output) with a Cast at the graph entry.
     _cast_module_dtype(module, dtype)
+    if prune_lm_head:
+        task = _enable_pruned_lm_head_task(task)
     resolved_task = get_task(task)
     capabilities = ep_registry.require(execution_provider)
     with build_context(capabilities, dtype):
@@ -371,6 +400,7 @@ def build(
     text_only: bool = False,
     fp8_kv_cache: bool = False,
     kv_cache_scales: dict[int, tuple[float, float]] | None = None,
+    prune_lm_head: bool = False,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID.
 
@@ -446,6 +476,11 @@ def build(
             per-tensor FP8 scales (from offline calibration), used only when
             ``fp8_kv_cache`` is ``True``. Layers absent from the map use a unit
             scale of ``1.0``.
+        prune_lm_head: When ``True``, build supported causal-LM tasks so the
+            exported ``logits`` output contains only the final token
+            (``[B, 1, vocab]``). This is intended for single-token
+            autoregressive generation and is incompatible with workflows that
+            need per-token logits.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -624,6 +659,7 @@ def build(
         trace_optimization=trace_optimization,
         fp8_kv_cache=fp8_kv_cache,
         kv_cache_scales=kv_cache_scales,
+        prune_lm_head=prune_lm_head,
     )
 
     for name, model in pkg.items():

src/mobius/models/_models_test.py

@@ -137,6 +137,16 @@ def test_build_with_task_string(self):
         assert isinstance(model, ir.Model)
         assert model.graph.num_nodes() > 0
 
+    def test_build_with_prune_lm_head_feature(self):
+        config = make_config()
+        module = CausalLMModel(config)
+        model = build_from_module(module, config, prune_lm_head=True)["model"]
+
+        logits = next(v for v in model.graph.outputs if v.name == "logits")
+        assert len(logits.shape) == 3
+        assert logits.shape[1] == 1
+        assert logits.shape[2] == config.vocab_size
+
     def test_build_with_output_layer_indices(self):
         config = make_config(num_hidden_layers=4, output_layer_indices=[1, 2])
         module = CausalLMModel(config)
@@ -181,6 +191,98 @@ def test_textmodel_output_layer_indices_set(self):
         assert model.output_layer_indices == [0, 3]
 
 
+class TestPruneLmHead:
+    """Tests for the ``prune_lm_head`` option in :class:`CausalLMTask`.
+
+    When ``True``, a ``Gather + Unsqueeze`` is inserted after the LM head
+    so logits are produced for only the last sequence position.
+    Mirrors onnxruntime-genai Model Builder's ``prune_lm_head`` opt-in.
+    """
+
+    def _build(self, prune_lm_head: bool = False) -> ir.Model:
+        config = make_config()
+        module = CausalLMModel(config)
+        task = CausalLMTask(prune_lm_head=prune_lm_head)
+        return build_from_module(module, config, task=task)["model"]
+
+    def test_default_emits_no_lm_head_pruning(self):
+        """Default (prune_lm_head=False): graph emits full [B, S, vocab] logits."""
+        model = self._build(prune_lm_head=False)
+
+        logits = next(v for v in model.graph.outputs if v.name == "logits")
+        # 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]
+        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_on_logits(self):
+        """Pruning selects the final hidden state before the LM-head MatMul."""
+        model = self._build(prune_lm_head=True)
+
+        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]
+        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
+        lm_head = logits.producer()
+        assert lm_head is not None and lm_head.op_type == "MatMul"
+        unsqueeze = lm_head.inputs[0].producer()
+        assert unsqueeze is not None and unsqueeze.op_type == "Unsqueeze"
+        gather = unsqueeze.inputs[0].producer()
+        assert gather is not None and gather.op_type == "Gather"
+
+    def test_prune_does_not_change_input_shapes(self):
+        """Pruning only affects output.
+
+        input_ids still has dynamic sequence_length
+        so the model accepts arbitrary prompts.
+        """
+        model = self._build(prune_lm_head=True)
+
+        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_custom_forward_that_ignores_pruning_fails(self):
+        class UnsupportedCausalLM(CausalLMModel):
+            def forward(
+                self,
+                op,
+                input_ids,
+                attention_mask,
+                position_ids,
+                past_key_values=None,
+            ):
+                hidden_states, present = self.model(
+                    op,
+                    input_ids=input_ids,
+                    attention_mask=attention_mask,
+                    position_ids=position_ids,
+                    past_key_values=past_key_values,
+                )
+                return self.lm_head(op, hidden_states), present
+
+        config = make_config()
+        with pytest.raises(ValueError, match="does not support prune_lm_head"):
+            build_from_module(UnsupportedCausalLM(config), config, prune_lm_head=True)
+
+
 class TestDeepStackCaptureOrdering:
     """``output_layer_indices`` must capture the post-DeepStack-injection state.