Fix prefill pruning for Gemma 4 models

CHANGELOG.md

@@ -45,7 +45,7 @@ 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`,
-  `prune-lm-head`, `text-only`. Unknown feature names are rejected with an error
+  `prune-prefill-prefix`, `text-only`. Unknown feature names are rejected with an error
   listing the valid set.
 
 #### Changed
@@ -54,16 +54,15 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   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`)
+### Prefill token-prefix pruning (`--features prune-prefill-prefix`)
 
 #### 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.
+- `build(prune_prefill_prefix=True)` and
+  `mobius build --features prune-prefill-prefix` discard prefill token positions
+  before the final token after required KV states have been produced. Generic
+  causal models prune immediately before the LM head; Gemma 4 also prunes its
+  KV-sharing layer suffix and per-layer inputs.
 
 ### FP8 (E4M3) KV-cache export (`--features fp8-kv-cache`)
 

README.md

@@ -113,12 +113,12 @@ mobius build --model openai/whisper-tiny output_dir/
 ```
 
 Build-mode toggles use the cargo-style `--features` option. Available features
-are `static-cache`, `fp8-kv-cache`, `prune-lm-head`, and `text-only`. Pass them
+are `static-cache`, `fp8-kv-cache`, `prune-prefill-prefix`, and `text-only`. Pass them
 as a comma-separated list or repeat the option:
 
 ```sh
 mobius build --model meta-llama/Llama-3.2-1B output_dir/ \
-    --features static-cache,prune-lm-head --max-seq-len 2048
+      --features static-cache,prune-prefill-prefix --max-seq-len 2048
 ```
 
 See the [CLI Reference](https://onnxruntime.github.io/mobius/cli_reference.html) for all subcommands and flags.

docs/cli_reference.md

@@ -175,7 +175,7 @@ option. Pass a comma-separated list (and/or repeat the flag):
 
 ```
 --features fp8-kv-cache,static-cache
---features prune-lm-head
+--features prune-prefill-prefix
 --features text-only
 ```
 
@@ -185,7 +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. |
+| `prune-prefill-prefix` | Emit logits shaped `[B, 1, vocab]` by selecting the final token before the LM head. Gemma 4 also prunes its KV-sharing layer suffix and per-layer inputs to reduce prefill compute. |
 | `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
@@ -199,7 +199,7 @@ 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
+    --features prune-prefill-prefix
 ```
 
 ### Static Cache (`--features static-cache`)

src/mobius/__main__.py

@@ -39,7 +39,7 @@
 _BUILD_FEATURES: dict[str, str] = {
     "static-cache": "static_cache",
     "fp8-kv-cache": "fp8_kv_cache",
-    "prune-lm-head": "prune_lm_head",
+    "prune-prefill-prefix": "prune_prefill_prefix",
     "text-only": "text_only",
 }
 
@@ -223,7 +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)
+    prune_prefill_prefix = getattr(args, "prune_prefill_prefix", 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:
@@ -313,7 +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,
+            prune_prefill_prefix=prune_prefill_prefix,
         )
         for name, model in pkg.items():
             model.graph.name = f"{config_path}/{name}"
@@ -343,7 +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,
+            prune_prefill_prefix=prune_prefill_prefix,
         )
 
     _save_package(pkg, output_dir, args, optimize, component_filter)

src/mobius/_build_context.py

@@ -35,8 +35,8 @@
     "build_context",
     "ep_capabilities",
     "get_build_dtype",
-    "is_lm_head_pruning_enabled",
-    "lm_head_pruning",
+    "is_prefill_prefix_pruning_enabled",
+    "prefill_prefix_pruning",
 ]
 
 _DEFAULT_CAPABILITIES = EpCapabilities(name="default")
@@ -47,8 +47,8 @@
 _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
+_prune_prefill_prefix: contextvars.ContextVar[bool] = contextvars.ContextVar(
+    "mobius_prune_prefill_prefix", default=False
 )
 
 
@@ -121,15 +121,15 @@ def get_build_dtype() -> ir.DataType:
 
 
 @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)
+def prefill_prefix_pruning(enabled: bool) -> Iterator[None]:
+    """Enable or disable prefill token-prefix pruning during graph construction."""
+    token = _prune_prefill_prefix.set(enabled)
     try:
         yield
     finally:
-        _prune_lm_head.reset(token)
+        _prune_prefill_prefix.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()
+def is_prefill_prefix_pruning_enabled() -> bool:
+    """Return whether the active task discards prefill tokens before the final token."""
+    return _prune_prefill_prefix.get()

src/mobius/_build_context_test.py

@@ -12,8 +12,8 @@
     build_context,
     ep_capabilities,
     get_build_dtype,
-    is_lm_head_pruning_enabled,
-    lm_head_pruning,
+    is_prefill_prefix_pruning_enabled,
+    prefill_prefix_pruning,
 )
 from mobius._execution_providers import EpCapabilities, ep_registry
 
@@ -37,8 +37,8 @@ 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_prefill_prefix_pruning_is_disabled(self):
+        assert not is_prefill_prefix_pruning_enabled()
 
     def test_default_capabilities_has_no_fusions(self):
         """Default EP has no GQA dtypes (portable ONNX)."""
@@ -74,10 +74,10 @@ 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()
+    def test_prefill_prefix_pruning_restored_after_context(self):
+        with prefill_prefix_pruning(True):
+            assert is_prefill_prefix_pruning_enabled()
+        assert not is_prefill_prefix_pruning_enabled()
 
 
 class TestBuildContextNesting:

src/mobius/_builder.py

@@ -106,25 +106,46 @@ 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
+def _enable_prefill_prefix_pruning_task(task: str | ModelTask) -> str | ModelTask:
+    """Return a task equivalent to *task* with prefill-prefix pruning enabled."""
+    from mobius.tasks import (
+        CausalLMTask,
+        Gemma4Task,
+        Gemma4TextCausalLMTask,
+        HybridCausalLMTask,
+    )
 
     if task == "text-generation":
-        return CausalLMTask(prune_lm_head=True)
+        return CausalLMTask(prune_prefill_prefix=True)
     if task == "hybrid-text-generation":
-        return HybridCausalLMTask(prune_lm_head=True)
+        return HybridCausalLMTask(prune_prefill_prefix=True)
+    if task == "gemma4-text-generation":
+        return Gemma4TextCausalLMTask(prune_prefill_prefix=True)
+    if task == "gemma4":
+        return Gemma4Task(prune_prefill_prefix=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,
+            prune_prefill_prefix=True,
         )
     if isinstance(task, HybridCausalLMTask):
-        return HybridCausalLMTask(prune_lm_head=True)
+        return HybridCausalLMTask(prune_prefill_prefix=True)
+    if isinstance(task, Gemma4TextCausalLMTask):
+        return Gemma4TextCausalLMTask(
+            static_cache=getattr(task, "_static_cache", False),
+            max_seq_len=getattr(task, "_max_seq_len", None),
+            prune_prefill_prefix=True,
+        )
+    if isinstance(task, Gemma4Task):
+        return Gemma4Task(
+            static_cache=getattr(task, "_static_cache", False),
+            max_seq_len=getattr(task, "_max_seq_len", None),
+            prune_prefill_prefix=True,
+        )
     raise ValueError(
-        "prune_lm_head=True is only supported for text-generation and "
-        "hybrid-text-generation tasks."
+        "prune_prefill_prefix=True is only supported for text-generation, "
+        "hybrid-text-generation, gemma4-text-generation, and gemma4 tasks."
     )
 
 
@@ -159,7 +180,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,
+    prune_prefill_prefix: bool = False,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a module instance and config.
 
@@ -200,10 +221,9 @@ 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.
+        prune_prefill_prefix: When ``True``, discard prefill token positions
+            before the final token from the remaining decoder computation and
+            logits. Only supported by causal generation tasks.
 
     Returns:
         A :class:`ModelPackage` containing the built model(s).
@@ -237,8 +257,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)
+    if prune_prefill_prefix:
+        task = _enable_prefill_prefix_pruning_task(task)
     resolved_task = get_task(task)
     capabilities = ep_registry.require(execution_provider)
     with build_context(capabilities, dtype):
@@ -400,7 +420,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,
+    prune_prefill_prefix: bool = False,
 ) -> ModelPackage:
     """Build an ONNX :class:`ModelPackage` from a HuggingFace model ID.
 
@@ -476,7 +496,7 @@ 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
+        prune_prefill_prefix: 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
@@ -659,7 +679,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,
+        prune_prefill_prefix=prune_prefill_prefix,
     )
 
     for name, model in pkg.items():

src/mobius/_builder_test.py

@@ -13,7 +13,12 @@
 import onnx_ir as ir
 import pytest
 
-from mobius._builder import _graph_requires_opset24, _maybe_apply_opset_lowering, flags
+from mobius._builder import (
+    _enable_prefill_prefix_pruning_task,
+    _graph_requires_opset24,
+    _maybe_apply_opset_lowering,
+    flags,
+)
 from mobius._model_package import ModelPackage
 
 
@@ -57,6 +62,16 @@ def _standard_nodes() -> list[ir.Node]:
     return [ir.Node("", "Reshape", inputs=[_make_value("x"), _make_value("shape")])]
 
 
+def test_prefill_prefix_pruning_error_lists_supported_tasks() -> None:
+    with pytest.raises(
+        ValueError,
+        match=(
+            "text-generation, hybrid-text-generation, gemma4-text-generation, and gemma4 tasks"
+        ),
+    ):
+        _enable_prefill_prefix_pruning_task("feature-extraction")
+
+
 def test_graph_requires_opset24_tensor_scatter() -> None:
     # A TensorScatter node (opset-24-only) must force opset 24 retention.
     node = ir.Node(

src/mobius/integrations/ort_genai/ep_config.py

@@ -33,6 +33,8 @@
 
 def make_provider_options(
     ep: str,
+    *,
+    graph_capture: bool | None = None,
 ) -> list[dict[str, dict[str, str]]]:
     """Build the ``provider_options`` list for genai_config.json.
 
@@ -57,7 +59,8 @@ def make_provider_options(
 
     # Graph capture comes from the EP's registered capability flag (the registry
     # is the single source of truth). Translate it into the EP-specific option.
-    graph_capture = bool(caps and caps.enable_graph_capture)
+    if graph_capture is None:
+        graph_capture = bool(caps and caps.enable_graph_capture)
 
     if ep == "webgpu":
         options["enableGraphCapture"] = "1" if graph_capture else "0"