update lm head last token pruning

src/mobius/models/_models_test.py

@@ -181,6 +181,67 @@
         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):
+        """With prune_lm_head=True, logits dim 1 collapses to 1 (still rank-3)."""
+        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
+
+    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
+
+
 class TestModelRegistry:
     def test_registry_not_empty(self):
         assert len(registry) > 0

src/mobius/tasks/_causal_lm.py

@@ -75,16 +75,28 @@ class CausalLMTask(ModelTask):
         max_seq_len: Maximum sequence length for static cache buffers.
             Only used when ``static_cache=True``.  Defaults to
             ``config.max_position_embeddings``.
+        prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)``
+            before the LM head so only the last token's hidden state is
+            projected to logits.  Output logits shape becomes ``[B, 1,
+            vocab]`` instead of ``[B, S, vocab]``, reducing prefill cost
+            for large-vocabulary models.  Set this when the downstream
+            runtime only needs the final token's logits (single-token
+            autoregressive generation).  Breaks workflows that require
+            per-token logits (logprob scoring, speculative decoding,
+            multi-token generation).  Mirrors the ``prune_lm_head`` extra
+            option in onnxruntime-genai's Model Builder.
     """
 
     def __init__(
         self,
         *,
         static_cache: bool = False,
         max_seq_len: int | None = None,
+        prune_lm_head: bool = False,
     ):
         self._static_cache = static_cache
         self._max_seq_len = max_seq_len
+        self._prune_lm_head = prune_lm_head
 
     def build(
         self,
@@ -189,6 +201,19 @@ def build(
         else:
             logits, present_key_values = result
 
+        if self._prune_lm_head:
+            # Select only the last token's logits: [B, S, vocab] -> [B, 1, vocab].
+            # Use a scalar (rank-0) index so Gather collapses dim 1:
+            #   [B, S, vocab] --Gather(axis=1, idx=-1)--> [B, vocab]
+            #   --Unsqueeze(axis=1)--> [B, 1, vocab]
+            # ONNX Runtime's graph optimizer can push this Gather backward
+            # through the LM head MatMul, avoiding the full [B, S, vocab]
+            # computation during prefill.  Mirrors onnxruntime-genai Model
+            # Builder's prune_lm_head option.
+            last_idx = op.Constant(value_int=-1)  # scalar (rank-0) INT64
+            last_token_logits = op.Gather(logits, last_idx, axis=1)  # [B, vocab]
+            logits = op.Unsqueeze(last_token_logits, op.Constant(value_ints=[1]))  # [B, 1, vocab]
+
         builder.add_output(logits, "logits")
 
         # --- Output registration (static vs dynamic) ---
@@ -240,8 +265,16 @@ class HybridCausalLMTask(ModelTask):
     Outputs:
         - logits: FLOAT
         - present.{i}.{key|value|conv_state|recurrent_state}: FLOAT
+
+    Args:
+        prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)``
+            after the LM head so only the last token's logits are emitted.
+            See :class:`CausalLMTask` for full documentation.
     """
 
+    def __init__(self, *, prune_lm_head: bool = False):
+        self._prune_lm_head = prune_lm_head
+
     def build(
         self,
         module: nn.Module,
@@ -285,6 +318,19 @@ def build(
         else:
             logits, present_key_values = result
 
+        if self._prune_lm_head:
+            # Select only the last token's logits: [B, S, vocab] -> [B, 1, vocab].
+            # Use a scalar (rank-0) index so Gather collapses dim 1:
+            #   [B, S, vocab] --Gather(axis=1, idx=-1)--> [B, vocab]
+            #   --Unsqueeze(axis=1)--> [B, 1, vocab]
+            # ONNX Runtime's graph optimizer can push this Gather backward
+            # through the LM head MatMul, avoiding the full [B, S, vocab]
+            # computation during prefill.  Mirrors onnxruntime-genai Model
+            # Builder's prune_lm_head option.
+            last_idx = op.Constant(value_int=-1)  # scalar (rank-0) INT64
+            last_token_logits = op.Gather(logits, last_idx, axis=1)  # [B, vocab]
+            logits = op.Unsqueeze(last_token_logits, op.Constant(value_ints=[1]))  # [B, 1, vocab]
+
         builder.add_output(logits, "logits")
         _register_hybrid_cache_outputs(
             builder,