diff --git a/CHANGELOG.md b/CHANGELOG.md
index 29b8d268d..d31205dd2 100644
--- a/CHANGELOG.md
+++ b/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`)
 
diff --git a/README.md b/README.md
index f9f2320b6..c2c90ed01 100644
--- a/README.md
+++ b/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.
diff --git a/docs/cli_reference.md b/docs/cli_reference.md
index 989608091..efb2b4137 100644
--- a/docs/cli_reference.md
+++ b/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`)
diff --git a/src/mobius/__main__.py b/src/mobius/__main__.py
index 821490e8f..d943c9167 100644
--- a/src/mobius/__main__.py
+++ b/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)
diff --git a/src/mobius/_build_context.py b/src/mobius/_build_context.py
index 8288dae7d..03a0b4317 100644
--- a/src/mobius/_build_context.py
+++ b/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()
diff --git a/src/mobius/_build_context_test.py b/src/mobius/_build_context_test.py
index d86efeed0..8d835088c 100644
--- a/src/mobius/_build_context_test.py
+++ b/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:
diff --git a/src/mobius/_builder.py b/src/mobius/_builder.py
index a02f6884d..0ace94769 100644
--- a/src/mobius/_builder.py
+++ b/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():
diff --git a/src/mobius/_builder_test.py b/src/mobius/_builder_test.py
index 7f262efa0..354308ed1 100644
--- a/src/mobius/_builder_test.py
+++ b/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(
diff --git a/src/mobius/integrations/ort_genai/ep_config.py b/src/mobius/integrations/ort_genai/ep_config.py
index 3f57d7086..f7ca9b6b9 100644
--- a/src/mobius/integrations/ort_genai/ep_config.py
+++ b/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"
diff --git a/src/mobius/integrations/ort_genai/genai_config.py b/src/mobius/integrations/ort_genai/genai_config.py
index 36c6761ee..dd2a9a7f2 100644
--- a/src/mobius/integrations/ort_genai/genai_config.py
+++ b/src/mobius/integrations/ort_genai/genai_config.py
@@ -104,7 +104,7 @@ def _default_search_params(
     }
 
 
-def _make_session_options(ep: str) -> dict[str, Any]:
+def _make_session_options(ep: str, *, graph_capture: bool | None = None) -> dict[str, Any]:
     """Return session options with EP-specific provider_options.
 
     Args:
@@ -115,7 +115,7 @@ def _make_session_options(ep: str) -> dict[str, Any]:
 
     return {
         "log_id": "onnxruntime-genai",
-        "provider_options": make_provider_options(ep),
+        "provider_options": make_provider_options(ep, graph_capture=graph_capture),
     }
 
 
@@ -322,7 +322,7 @@ def with_vision(
             "config_filename": config_filename,
             "inputs": input_names,
             "outputs": output_names,
-            "session_options": _make_session_options(self.ep),
+            "session_options": _make_session_options(self.ep, graph_capture=False),
         }
         if spatial_merge_size is not None:
             self._vision["spatial_merge_size"] = spatial_merge_size
@@ -341,7 +341,7 @@ def with_vision(
             else {
                 "inputs_embeds": "inputs_embeds",
             },
-            "session_options": _make_session_options(self.ep),
+            "session_options": _make_session_options(self.ep, graph_capture=False),
         }
         self._vlm_token_ids["image_token_id"] = image_token_id
         if vision_start_token_id is not None:
@@ -391,7 +391,7 @@ def with_audio(
             "config_filename": config_filename,
             "inputs": input_names,
             "outputs": output_names,
-            "session_options": _make_session_options(self.ep),
+            "session_options": _make_session_options(self.ep, graph_capture=False),
         }
 
         if audio_token_id is not None:
diff --git a/src/mobius/integrations/ort_genai/genai_config_test.py b/src/mobius/integrations/ort_genai/genai_config_test.py
index 8ee45d18f..6ae6f49ef 100644
--- a/src/mobius/integrations/ort_genai/genai_config_test.py
+++ b/src/mobius/integrations/ort_genai/genai_config_test.py
@@ -232,6 +232,27 @@ def test_webgpu_graph_capture_propagates_to_session_options(self):
         assert webgpu["enableGraphCapture"] == "1"
         assert webgpu["validationMode"] == "disabled"
 
+    def test_webgpu_graph_capture_is_decoder_only_for_multimodal(self):
+        gen = GenaiConfigGenerator(
+            "gemma4",
+            vocab_size=256,
+            hidden_size=64,
+            num_hidden_layers=2,
+            num_attention_heads=4,
+            num_key_value_heads=1,
+            head_dim=16,
+            ep="webgpu",
+        )
+        config = gen.with_vision(image_token_id=255999).with_audio().generate()
+
+        model = config["model"]
+        decoder_webgpu = model["decoder"]["session_options"]["provider_options"][0]["webgpu"]
+        assert decoder_webgpu["enableGraphCapture"] == "1"
+        for component in ("vision", "embedding", "speech"):
+            webgpu = model[component]["session_options"]["provider_options"][0]["webgpu"]
+            assert webgpu["enableGraphCapture"] == "0"
+            assert webgpu["validationMode"] == "basic"
+
     def test_search_params_custom_ep_with_share_buffer(self):
         """A custom EP registered with supports_past_present_share_buffer=True gets the flag set.
 
diff --git a/src/mobius/models/_models_test.py b/src/mobius/models/_models_test.py
index 61352b9f7..36d064af2 100644
--- a/src/mobius/models/_models_test.py
+++ b/src/mobius/models/_models_test.py
@@ -137,10 +137,10 @@ 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):
+    def test_build_with_prune_prefill_prefix_feature(self):
         config = make_config()
         module = CausalLMModel(config)
-        model = build_from_module(module, config, prune_lm_head=True)["model"]
+        model = build_from_module(module, config, prune_prefill_prefix=True)["model"]
 
         logits = next(v for v in model.graph.outputs if v.name == "logits")
         assert len(logits.shape) == 3
@@ -191,23 +191,23 @@ 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`.
+class TestPrunePrefillPrefix:
+    """Tests for the ``prune_prefill_prefix`` 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.
+    The generic causal task applies the optimization immediately before the LM head.
     """
 
-    def _build(self, prune_lm_head: bool = False) -> ir.Model:
+    def _build(self, prune_prefill_prefix: bool = False) -> ir.Model:
         config = make_config()
         module = CausalLMModel(config)
-        task = CausalLMTask(prune_lm_head=prune_lm_head)
+        task = CausalLMTask(prune_prefill_prefix=prune_prefill_prefix)
         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)
+    def test_default_emits_no_prefill_prefix_pruning(self):
+        """Default graph emits full [B, S, vocab] logits."""
+        model = self._build(prune_prefill_prefix=False)
 
         logits = next(v for v in model.graph.outputs if v.name == "logits")
         # Logits must remain rank-3 [B, S, vocab]
@@ -226,7 +226,7 @@ def test_default_emits_no_lm_head_pruning(self):
 
     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)
+        model = self._build(prune_prefill_prefix=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])
@@ -253,7 +253,7 @@ def test_prune_does_not_change_input_shapes(self):
         input_ids still has dynamic sequence_length
         so the model accepts arbitrary prompts.
         """
-        model = self._build(prune_lm_head=True)
+        model = self._build(prune_prefill_prefix=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
@@ -279,8 +279,8 @@ def forward(
                 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)
+        with pytest.raises(ValueError, match="does not support prune_prefill_prefix"):
+            build_from_module(UnsupportedCausalLM(config), config, prune_prefill_prefix=True)
 
 
 class TestDeepStackCaptureOrdering:
diff --git a/src/mobius/models/base.py b/src/mobius/models/base.py
index 8590d41cc..2e40ebf17 100644
--- a/src/mobius/models/base.py
+++ b/src/mobius/models/base.py
@@ -22,7 +22,7 @@
 from mobius._build_context import (
     ep_capabilities,
     get_build_dtype,
-    is_lm_head_pruning_enabled,
+    is_prefill_prefix_pruning_enabled,
 )
 from mobius._configs import ArchitectureConfig, CausalLMConfig
 from mobius._flags import flags
@@ -448,11 +448,11 @@ def forward(
         )
         if len(result) == 3:
             hidden_states, present_key_values, intermediate_hidden_states = result
-            hidden_states = _prune_lm_head_hidden_states(op, hidden_states)
+            hidden_states = _retain_last_sequence_token(op, hidden_states)
             logits = self.lm_head(op, hidden_states)
             return logits, present_key_values, intermediate_hidden_states
         hidden_states, present_key_values = result
-        hidden_states = _prune_lm_head_hidden_states(op, hidden_states)
+        hidden_states = _retain_last_sequence_token(op, hidden_states)
         logits = self.lm_head(op, hidden_states)
         return logits, present_key_values
 
@@ -499,9 +499,9 @@ def preprocess_weights(
         return state_dict
 
 
-def _prune_lm_head_hidden_states(op: OpBuilder, hidden_states: ir.Value) -> ir.Value:
-    """Select the final sequence position before the LM-head projection."""
-    if not is_lm_head_pruning_enabled():
+def _retain_last_sequence_token(op: OpBuilder, hidden_states: ir.Value) -> ir.Value:
+    """Retain only the final sequence position when prefill-prefix pruning is active."""
+    if not is_prefill_prefix_pruning_enabled():
         return hidden_states
     last_hidden = op.Gather(hidden_states, op.Constant(value_int=-1), axis=1)
     return op.Unsqueeze(last_hidden, op.Constant(value_ints=[1]))
diff --git a/src/mobius/models/gemma4.py b/src/mobius/models/gemma4.py
index 265f8b0f2..1bc0922c4 100644
--- a/src/mobius/models/gemma4.py
+++ b/src/mobius/models/gemma4.py
@@ -50,7 +50,7 @@
 from mobius.components._activations import get_activation
 from mobius.components._gemma4_audio import Gemma4AudioEncoder
 from mobius.components._mlp import GatedMLP
-from mobius.models.base import CausalLMModel
+from mobius.models.base import CausalLMModel, _retain_last_sequence_token
 from mobius.models.gemma3_text import Gemma3TextScaledWordEmbedding
 
 if TYPE_CHECKING:
@@ -70,6 +70,36 @@
 # ---------------------------------------------------------------------------
 
 
+def _split_per_layer_projection_weight(
+    state_dict: dict[str, torch.Tensor],
+    prefix: str,
+    config: Gemma4Config,
+) -> None:
+    """Split the packed PLE projection at the first KV-sharing layer."""
+    shared_layers = config.num_kv_shared_layers
+    per_layer_dim = config.hidden_size_per_layer_input
+    if not shared_layers or not per_layer_dim:
+        return
+
+    weight_key = f"{prefix}per_layer_model_projection.weight"
+    weight = state_dict.get(weight_key)
+    if weight is None:
+        return
+
+    expected_rows = config.num_hidden_layers * per_layer_dim
+    if weight.shape[0] != expected_rows:
+        raise ValueError(f"{weight_key} dim 0 expected {expected_rows}, got {weight.shape[0]}")
+    producer_rows = (config.num_hidden_layers - shared_layers) * per_layer_dim
+    producer, consumer = weight.split([producer_rows, expected_rows - producer_rows], dim=0)
+    state_dict[weight_key] = producer.contiguous()
+    state_dict[f"{prefix}per_layer_model_projection_consumer.weight"] = consumer.contiguous()
+
+
+def _typed_scalar_constant(op: OpBuilder, value: float, dtype: ir.DataType) -> ir.Value:
+    """Create a scalar constant directly in the model compute dtype."""
+    return op.Constant(value=ir.tensor(np.asarray(value, dtype=dtype.numpy())))
+
+
 def _text_quantization_config(config: Gemma4Config):
     """Return the active weight-quantization config, or ``None`` when off."""
     quantization_config = getattr(config, "quantization", None)
@@ -143,7 +173,7 @@ def _make_scaled_word_embedding(
             block_size=quantization_config.group_size,
             has_zero_point=not quantization_config.sym,
         )
-    return Gemma3TextScaledWordEmbedding(
+    return Gemma4ScaledWordEmbedding(
         num_embeddings,
         embedding_dim,
         config.pad_token_id,
@@ -166,6 +196,17 @@ def _make_lm_head(config: Gemma4Config) -> nn.Module:
     return Linear(config.hidden_size, config.vocab_size, bias=False)
 
 
+class Gemma4ScaledWordEmbedding(Gemma3TextScaledWordEmbedding):
+    """Gemma embedding with a typed scale constant suitable for graph capture."""
+
+    def forward(self, op: OpBuilder, input_ids: ir.Value) -> ir.Value:
+        embeddings = Embedding.forward(self, op, input_ids)
+        scale = op.Constant(
+            value=ir.tensor(np.asarray(self.embed_scale, dtype=self.weight.dtype.numpy()))
+        )
+        return op.Mul(embeddings, scale)
+
+
 class Gemma4ScaledQuantizedWordEmbedding(QuantizedEmbedding):
     """GatherBlockQuantized token embedding scaled by ``embed_scale``.
 
@@ -197,7 +238,10 @@ def __init__(
 
     def forward(self, op: OpBuilder, input_ids: ir.Value) -> ir.Value:
         embeddings = super().forward(op, input_ids)
-        return op.Mul(embeddings, self.embed_scale)
+        scale = op.Constant(
+            value=ir.tensor(np.asarray(self.embed_scale, dtype=self.scales.dtype.numpy()))
+        )
+        return op.Mul(embeddings, scale)
 
 
 def _dtype_safe_compress(
@@ -1773,6 +1817,8 @@ def __init__(self, config: Gemma4Config):
         # embedding model and passed as per_layer_inputs. In single-model
         # (text-only) mode, they are computed here from input_ids.
         self._per_layer_dim = getattr(config, "hidden_size_per_layer_input", 0)
+        self._num_layers = config.num_hidden_layers
+        self._first_kv_shared_layer = self._num_layers - config.num_kv_shared_layers
         self._hidden_size = config.hidden_size
         self._image_token_id: int = config.image_token_id or 0
         # The vision-block overlay keys on image_token_id. A 0/None id means the
@@ -1785,7 +1831,6 @@ def __init__(self, config: Gemma4Config):
             config.audio.audio_token_id if config.audio is not None else None
         )
         if self._per_layer_dim:
-            self._num_layers = config.num_hidden_layers
             vocab_per_layer = getattr(config, "vocab_size_per_layer_input", 0)
             # Fused [V, L*D] table — used when split_per_layer_embedding is False.
             # Requires ORT >= 1.27 for CUDA Gather int64 index support (onnxruntime#28107).
@@ -1802,7 +1847,7 @@ def __init__(self, config: Gemma4Config):
             # ONNX initializer, so the unused one adds no graph weight.
             self.embed_tokens_per_layer_split = nn.ModuleList(
                 [
-                    Gemma3TextScaledWordEmbedding(
+                    Gemma4ScaledWordEmbedding(
                         vocab_per_layer,
                         self._per_layer_dim,
                         config.pad_token_id,
@@ -1813,9 +1858,15 @@ def __init__(self, config: Gemma4Config):
             )
             self.per_layer_model_projection = Linear(
                 config.hidden_size,
-                config.num_hidden_layers * self._per_layer_dim,
+                self._first_kv_shared_layer * self._per_layer_dim,
                 bias=False,
             )
+            if self._first_kv_shared_layer < self._num_layers:
+                self.per_layer_model_projection_consumer = Linear(
+                    config.hidden_size,
+                    (self._num_layers - self._first_kv_shared_layer) * self._per_layer_dim,
+                    bias=False,
+                )
             self.per_layer_projection_norm = RMSNorm(
                 self._per_layer_dim, eps=config.rms_norm_eps
             )
@@ -1859,12 +1910,26 @@ def _compute_per_layer_inputs(
         inputs_embeds: ir.Value,
     ) -> list[ir.Value]:
         """Compute per-layer input embeddings for single-model (text-only) mode."""
-        proj = self.per_layer_model_projection(op, inputs_embeds)
-        proj = op.Mul(proj, float(self._hidden_size**-0.5))
-        proj = op.Reshape(
-            proj, op.Constant(value_ints=[0, 0, self._num_layers, self._per_layer_dim])
+        producer_count = self._first_kv_shared_layer
+        producer_proj = self.per_layer_model_projection(op, inputs_embeds)
+        producer_proj = op.Mul(producer_proj, float(self._hidden_size**-0.5))
+        producer_proj = op.Reshape(
+            producer_proj,
+            op.Constant(value_ints=[0, 0, producer_count, self._per_layer_dim]),
         )
-        proj = self.per_layer_projection_norm(op, proj)
+        producer_proj = self.per_layer_projection_norm(op, producer_proj)
+
+        consumer_count = self._num_layers - producer_count
+        consumer_proj: ir.Value | None = None
+        if consumer_count:
+            consumer_inputs = _retain_last_sequence_token(op, inputs_embeds)
+            consumer_proj = self.per_layer_model_projection_consumer(op, consumer_inputs)
+            consumer_proj = op.Mul(consumer_proj, float(self._hidden_size**-0.5))
+            consumer_proj = op.Reshape(
+                consumer_proj,
+                op.Constant(value_ints=[0, 0, consumer_count, self._per_layer_dim]),
+            )
+            consumer_proj = self.per_layer_projection_norm(op, consumer_proj)
 
         pad = op.Constant(value_int=0)
         masked_ids = input_ids
@@ -1884,25 +1949,46 @@ def _compute_per_layer_inputs(
         if self.config.split_per_layer_embedding:
             # L separate Gathers on [V, D] tables — each fits within the EP's
             # max_buffer_size (e.g. WebGPU's 256 MiB limit).
-            per_layer_embs = [
-                op.Unsqueeze(self.embed_tokens_per_layer_split[i](op, masked_ids), [2])
-                for i in range(self._num_layers)
-            ]
-            fused_emb = op.Concat(*per_layer_embs, axis=2)  # [B, S, L, D]
+            per_layer_embs = []
+            for layer_idx in range(self._num_layers):
+                embedding = self.embed_tokens_per_layer_split[layer_idx](op, masked_ids)
+                if layer_idx >= producer_count:
+                    embedding = _retain_last_sequence_token(op, embedding)
+                per_layer_embs.append(op.Unsqueeze(embedding, [2]))
+            producer_emb = op.Concat(*per_layer_embs[:producer_count], axis=2)
+            consumer_emb = (
+                op.Concat(*per_layer_embs[producer_count:], axis=2) if consumer_count else None
+            )
         else:
             fused_emb = self.embed_tokens_per_layer(op, masked_ids)
             fused_emb = op.Reshape(
                 fused_emb,
                 op.Constant(value_ints=[0, 0, self._num_layers, self._per_layer_dim]),
             )
+            producer_emb = op.Slice(fused_emb, starts=[0], ends=[producer_count], axes=[2])
+            consumer_emb = None
+            if consumer_count:
+                consumer_emb = op.Slice(
+                    fused_emb,
+                    starts=[producer_count],
+                    ends=[self._num_layers],
+                    axes=[2],
+                )
+                consumer_emb = _retain_last_sequence_token(op, consumer_emb)
 
-        combined = op.Add(proj, fused_emb)
-        combined = op.Mul(combined, float(0.5**0.5))
-
-        return [
-            op.Gather(combined, op.Constant(value_int=i), axis=2)
-            for i in range(self._num_layers)
+        producer_combined = op.Mul(op.Add(producer_proj, producer_emb), float(0.5**0.5))
+        per_layer_inputs = [
+            op.Gather(producer_combined, op.Constant(value_int=i), axis=2)
+            for i in range(producer_count)
         ]
+        if consumer_count:
+            assert consumer_proj is not None and consumer_emb is not None
+            consumer_combined = op.Mul(op.Add(consumer_proj, consumer_emb), float(0.5**0.5))
+            per_layer_inputs.extend(
+                op.Gather(consumer_combined, op.Constant(value_int=i), axis=2)
+                for i in range(consumer_count)
+            )
+        return per_layer_inputs
 
     def forward(
         self,
@@ -1934,6 +2020,10 @@ def forward(
                 op.Squeeze(op.Slice(per_layer_4d, starts=[i], ends=[i + 1], axes=[2]), [2])
                 for i in range(num_layers)
             ]
+            for layer_idx in range(self._first_kv_shared_layer, num_layers):
+                per_layer_list[layer_idx] = _retain_last_sequence_token(
+                    op, per_layer_list[layer_idx]
+                )
         elif self._per_layer_dim and input_ids is not None:
             # Text-only: compute per-layer inputs from input_ids
             per_layer_list = self._compute_per_layer_inputs(op, input_ids, hidden_states)
@@ -2013,20 +2103,12 @@ def forward(
                 op.Sub(reduce_sum, one_i32),
                 to=ir.DataType.INT32,
             )
-            if caps.requires_graph_capture_rewrite:
-                # Support graph capture for shared-KV layer models on WebGPU EP.
-                # Derive total_seq_len from reduce_sum (already computed) as a
-                # scalar INT32 via Gather index 0 (valid because graph capture
-                # requires batch=1).
-                total_seq_len = op.Gather(
-                    op.Cast(reduce_sum, to=ir.DataType.INT32),
-                    op.Constant(value_int=0),
-                )
-            else:
-                total_seq_len = op.Cast(
-                    op.Gather(op.Shape(attention_mask), 1),
-                    to=ir.DataType.INT32,
-                )
+            # Graph capture requires batch=1, so the first reduced sequence
+            # length is also the scalar total sequence length expected by GQA.
+            total_seq_len = op.Gather(
+                op.Cast(reduce_sum, to=ir.DataType.INT32),
+                op.Constant(value_int=0),
+            )
 
             # Per-layer-type GQA contexts with appropriate cos/sin caches
             # and local_window_size for sliding layers.
@@ -2157,6 +2239,8 @@ def forward(
         for i, (layer, layer_type, past_kv) in enumerate(
             zip(self.layers, self.layer_types, past_kvs)
         ):
+            if i == self._first_kv_shared_layer and i < len(self.layers):
+                hidden_states = _retain_last_sequence_token(op, hidden_states)
             per_layer_input = per_layer_list[i] if per_layer_list is not None else None
 
             # Per-layer cache/attention dispatch:
@@ -2243,12 +2327,12 @@ def forward(
             past_key_values=past_key_values,
             inputs_embeds=inputs_embeds,
         )
+        hidden_states = _retain_last_sequence_token(op, hidden_states)
         logits = self.lm_head(op, hidden_states)
         # Optional final logit soft-capping (tanh scaled): logit_cap * tanh(x / logit_cap)
         if self.config.final_logit_softcapping:
-            cap = op.CastLike(
-                self.config.final_logit_softcapping,
-                logits,
+            cap = _typed_scalar_constant(
+                op, self.config.final_logit_softcapping, self.config.dtype
             )
             logits = op.Mul(op.Tanh(op.Div(logits, cap)), cap)
         return logits, present_key_values
@@ -2268,6 +2352,7 @@ def preprocess_weights(
         # (For WebGPU, splitting is handled by _Gemma4DecoderModel.preprocess_weights.)
         # Map HF expert weight names and fold router scale
         _remap_moe_expert_weights(state_dict, self.config)
+        _split_per_layer_projection_weight(state_dict, "model.", self.config)
         return super().preprocess_weights(state_dict)
 
     def static_kv_cache_specs(self) -> list[tuple[int, int]]:
@@ -2321,12 +2406,12 @@ def forward(
             per_layer_inputs=per_layer_inputs,
             block_sequence_ids=block_sequence_ids,
         )
+        hidden_states = _retain_last_sequence_token(op, hidden_states)
         logits = self.lm_head(op, hidden_states)
         # Gemma4 applies final logit soft-capping: logit_cap * tanh(x / logit_cap)
         if self.config.final_logit_softcapping:
-            cap = op.CastLike(
-                self.config.final_logit_softcapping,
-                logits,
+            cap = _typed_scalar_constant(
+                op, self.config.final_logit_softcapping, self.config.dtype
             )
             logits = op.Mul(op.Tanh(op.Div(logits, cap)), cap)
         return logits, present_key_values
@@ -2335,6 +2420,7 @@ def preprocess_weights(
         self, state_dict: dict[str, torch.Tensor]
     ) -> dict[str, torch.Tensor]:
         state_dict = vlm_decoder_weights(state_dict, tie=self.config.tie_word_embeddings)
+        _split_per_layer_projection_weight(state_dict, "model.", self.config)
         # For WebGPU: split the fused [V, L*D] per-layer embedding into L separate [V, D] tables.
         per_layer_dim = self.config.hidden_size_per_layer_input
         if per_layer_dim and self.config.split_per_layer_embedding:
@@ -3160,6 +3246,7 @@ def preprocess_weights(
                     )
                 ):
                     renamed[k.replace("embedding.", "decoder.model.", 1)] = renamed.pop(k)
+            _split_per_layer_projection_weight(renamed, "decoder.model.", self.config)
 
         return renamed
 
diff --git a/src/mobius/tasks/_causal_lm.py b/src/mobius/tasks/_causal_lm.py
index cc665c9fb..954f3975c 100644
--- a/src/mobius/tasks/_causal_lm.py
+++ b/src/mobius/tasks/_causal_lm.py
@@ -8,7 +8,7 @@
 import onnx_ir as ir
 from onnxscript import GraphBuilder, nn
 
-from mobius._build_context import lm_head_pruning
+from mobius._build_context import prefill_prefix_pruning
 from mobius._configs import ArchitectureConfig
 from mobius._model_package import ModelPackage
 from mobius.components._attention import StaticCacheState
@@ -76,7 +76,7 @@ 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)``
+        prune_prefill_prefix: 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
@@ -84,8 +84,7 @@ class CausalLMTask(ModelTask):
             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.
+            multi-token generation).
     """
 
     def __init__(
@@ -93,11 +92,11 @@ def __init__(
         *,
         static_cache: bool = False,
         max_seq_len: int | None = None,
-        prune_lm_head: bool = False,
+        prune_prefill_prefix: bool = False,
     ):
         self._static_cache = static_cache
         self._max_seq_len = max_seq_len
-        self._prune_lm_head = prune_lm_head
+        self._prune_prefill_prefix = prune_prefill_prefix
 
     def build(
         self,
@@ -189,7 +188,7 @@ def build(
                 value_head_dim=kv_value_head_dim,
             )
 
-        with lm_head_pruning(self._prune_lm_head):
+        with prefill_prefix_pruning(self._prune_prefill_prefix):
             result = module(
                 op,
                 input_ids=input_ids,
@@ -203,7 +202,7 @@ def build(
         else:
             logits, present_key_values = result
 
-        _validate_pruned_logits(logits, self._prune_lm_head, module)
+        _validate_pruned_logits(logits, self._prune_prefill_prefix, module)
 
         builder.add_output(logits, "logits")
 
@@ -258,13 +257,13 @@ class HybridCausalLMTask(ModelTask):
         - present.{i}.{key|value|conv_state|recurrent_state}: FLOAT
 
     Args:
-        prune_lm_head: If ``True``, insert ``Gather(axis=1, index=-1)``
+        prune_prefill_prefix: If ``True``, insert ``Gather(axis=1, index=-1)``
             before 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 __init__(self, *, prune_prefill_prefix: bool = False):
+        self._prune_prefill_prefix = prune_prefill_prefix
 
     def build(
         self,
@@ -296,7 +295,7 @@ def build(
             past_seq_len,
         )
 
-        with lm_head_pruning(self._prune_lm_head):
+        with prefill_prefix_pruning(self._prune_prefill_prefix):
             result = module(
                 op,
                 input_ids=input_ids,
@@ -310,7 +309,7 @@ def build(
         else:
             logits, present_key_values = result
 
-        _validate_pruned_logits(logits, self._prune_lm_head, module)
+        _validate_pruned_logits(logits, self._prune_prefill_prefix, module)
 
         builder.add_output(logits, "logits")
         _register_hybrid_cache_outputs(
@@ -331,16 +330,16 @@ def build(
 
 def _validate_pruned_logits(
     logits: ir.Value,
-    prune_lm_head: bool,
+    prune_prefill_prefix: bool,
     module: nn.Module,
 ) -> None:
     """Fail when a custom model forward ignores the task's pruning request."""
-    if not prune_lm_head:
+    if not prune_prefill_prefix:
         return
     shape = logits.shape
     if shape is None or len(shape) != 3 or shape[1] != 1:
         raise ValueError(
-            f"{type(module).__name__} does not support prune_lm_head. "
+            f"{type(module).__name__} does not support prune_prefill_prefix. "
             "The model must select the final hidden state before its LM-head projection."
         )
 
diff --git a/src/mobius/tasks/_gemma4.py b/src/mobius/tasks/_gemma4.py
index a8233a790..7e1d4fc55 100644
--- a/src/mobius/tasks/_gemma4.py
+++ b/src/mobius/tasks/_gemma4.py
@@ -24,7 +24,7 @@
 import onnx_ir as ir
 from onnxscript import GraphBuilder, nn
 
-from mobius._build_context import ep_capabilities
+from mobius._build_context import ep_capabilities, prefill_prefix_pruning
 from mobius._configs import Gemma4Config
 from mobius._model_package import ModelPackage
 from mobius._pipeline_contract import (
@@ -276,9 +276,11 @@ def __init__(
         *,
         static_cache: bool = False,
         max_seq_len: int | None = None,
+        prune_prefill_prefix: bool = False,
     ):
         self._static_cache = static_cache
         self._max_seq_len = max_seq_len
+        self._prune_prefill_prefix = prune_prefill_prefix
 
     def build(
         self,
@@ -351,13 +353,14 @@ def build(
                 past_seq_len,
             )
 
-        logits, present_key_values = module(
-            op,
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            past_key_values=past_key_values,
-        )
+        with prefill_prefix_pruning(self._prune_prefill_prefix):
+            logits, present_key_values = module(
+                op,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                past_key_values=past_key_values,
+            )
         builder.add_output(logits, "logits")
 
         if static:
@@ -416,9 +419,11 @@ def __init__(
         *,
         static_cache: bool = False,
         max_seq_len: int | None = None,
+        prune_prefill_prefix: bool = False,
     ):
         self._static_cache = static_cache
         self._max_seq_len = max_seq_len
+        self._prune_prefill_prefix = prune_prefill_prefix
 
     def build(
         self,
@@ -551,15 +556,16 @@ def _build_decoder(
                 past_seq_len,
             )
 
-        logits, present_key_values = decoder(
-            op,
-            inputs_embeds=inputs_embeds,
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            per_layer_inputs=per_layer_inputs_val,
-            past_key_values=past_key_values,
-            input_ids=input_ids_val,
-        )
+        with prefill_prefix_pruning(self._prune_prefill_prefix):
+            logits, present_key_values = decoder(
+                op,
+                inputs_embeds=inputs_embeds,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                per_layer_inputs=per_layer_inputs_val,
+                past_key_values=past_key_values,
+                input_ids=input_ids_val,
+            )
 
         builder.add_output(logits, "logits")
         if static:
diff --git a/tests/cli_test.py b/tests/cli_test.py
index 7a3e7cf6e..1b932d996 100644
--- a/tests/cli_test.py
+++ b/tests/cli_test.py
@@ -244,8 +244,8 @@ def test_features_fp8_kv_cache_passed_through(self):
             )
         assert mock_build.call_args.kwargs.get("fp8_kv_cache") is True
 
-    def test_features_prune_lm_head_passed_through(self):
-        """--features prune-lm-head sets prune_lm_head on the build() call."""
+    def test_features_prune_prefill_prefix_passed_through(self):
+        """--features prune-prefill-prefix sets the build option."""
         with (
             tempfile.TemporaryDirectory() as tmpdir,
             mock.patch(
@@ -263,10 +263,10 @@ def test_features_prune_lm_head_passed_through(self):
                     tmpdir,
                     "--no-weights",
                     "--features",
-                    "prune-lm-head",
+                    "prune-prefill-prefix",
                 ]
             )
-        assert mock_build.call_args.kwargs.get("prune_lm_head") is True
+        assert mock_build.call_args.kwargs.get("prune_prefill_prefix") is True
 
     def test_features_comma_separated_multiple(self):
         """A single --features accepts a comma-separated list."""
@@ -287,13 +287,13 @@ def test_features_comma_separated_multiple(self):
                     tmpdir,
                     "--no-weights",
                     "--features",
-                    "text-only,fp8-kv-cache,prune-lm-head",
+                    "text-only,fp8-kv-cache,prune-prefill-prefix",
                 ]
             )
         kwargs = mock_build.call_args.kwargs
         assert kwargs.get("text_only") is True
         assert kwargs.get("fp8_kv_cache") is True
-        assert kwargs.get("prune_lm_head") is True
+        assert kwargs.get("prune_prefill_prefix") is True
 
     def test_features_unknown_errors(self):
         """An unrecognised feature name is rejected with a clear error."""
diff --git a/tests/gemma4_prefill_prefix_test.py b/tests/gemma4_prefill_prefix_test.py
new file mode 100644
index 000000000..dbccb285d
--- /dev/null
+++ b/tests/gemma4_prefill_prefix_test.py
@@ -0,0 +1,135 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from __future__ import annotations
+
+import torch
+
+from mobius import build_from_module
+from mobius._configs import Gemma4Config, VisionConfig
+from mobius._registry import registry
+from mobius.models.gemma4 import _split_per_layer_projection_weight
+
+
+def _make_config(*, with_vision: bool = False) -> Gemma4Config:
+    return Gemma4Config(
+        num_hidden_layers=4,
+        hidden_size=64,
+        intermediate_size=128,
+        num_attention_heads=4,
+        num_key_value_heads=1,
+        head_dim=16,
+        vocab_size=256,
+        rms_norm_eps=1e-6,
+        hidden_act="silu",
+        attn_qk_norm=True,
+        layer_types=[
+            "sliding_attention",
+            "full_attention",
+            "sliding_attention",
+            "full_attention",
+        ],
+        sliding_window=8,
+        global_head_dim=16,
+        global_rope_theta=10_000.0,
+        global_partial_rotary_factor=0.25,
+        final_logit_softcapping=30.0,
+        hidden_size_per_layer_input=8,
+        vocab_size_per_layer_input=64,
+        split_per_layer_embedding=True,
+        image_token_id=255999 if with_vision else None,
+        pad_token_id=0,
+        tie_word_embeddings=False,
+        num_kv_shared_layers=2,
+        vision=(
+            VisionConfig(
+                hidden_size=32,
+                intermediate_size=64,
+                num_hidden_layers=1,
+                num_attention_heads=2,
+                patch_size=16,
+                norm_eps=1e-6,
+            )
+            if with_vision
+            else None
+        ),
+    )
+
+
+def test_prunes_gemma4_shared_layer_prefix() -> None:
+    config = _make_config()
+    module = registry.get("gemma4_text")(config)
+    model = build_from_module(
+        module,
+        config,
+        task="gemma4-text-generation",
+        execution_provider="webgpu",
+        prune_prefill_prefix=True,
+    )["model"]
+
+    producer = next(
+        node
+        for node in model.graph
+        if node.op_type == "MatMul" and "/per_layer_model_projection/" in node.name
+    )
+    consumer = next(
+        node
+        for node in model.graph
+        if node.op_type == "MatMul" and "/per_layer_model_projection_consumer/" in node.name
+    )
+    assert producer.outputs[0].shape[1] != 1
+    assert producer.outputs[0].shape[2] == 16
+    assert consumer.inputs[0].shape[1:] == (1, 64)
+    assert consumer.outputs[0].shape[1:] == (1, 16)
+
+    first_shared_norm = next(
+        node for node in model.graph if "layers.2/input_layernorm" in node.name
+    )
+    assert first_shared_norm.inputs[0].shape[1:] == (1, 64)
+    logits = next(value for value in model.graph.outputs if value.name == "logits")
+    assert logits.shape[1:] == (1, config.vocab_size)
+
+    consumer_embedding_scale = next(
+        node
+        for node in model.graph
+        if node.op_type == "Mul" and "embed_tokens_per_layer_split.2" in node.name
+    )
+    assert any(
+        node.op_type == "Gather"
+        and any(
+            input_value is consumer_embedding_scale.outputs[0] for input_value in node.inputs
+        )
+        for node in model.graph
+    )
+    assert not any(node.op_type == "CastLike" for node in model.graph)
+
+
+def test_multimodal_task_prunes_decoder_prefix() -> None:
+    config = _make_config(with_vision=True)
+    module = registry.get("gemma4")(config)
+    package = build_from_module(
+        module,
+        config,
+        task="gemma4",
+        execution_provider="webgpu",
+        prune_prefill_prefix=True,
+    )
+
+    logits = next(
+        value for value in package["decoder"].graph.outputs if value.name == "logits"
+    )
+    assert logits.shape[1:] == (1, config.vocab_size)
+
+
+def test_splits_per_layer_projection_weight() -> None:
+    config = _make_config()
+    original = torch.arange(32 * 64, dtype=torch.float32).reshape(32, 64)
+    state_dict = {"model.per_layer_model_projection.weight": original.clone()}
+
+    _split_per_layer_projection_weight(state_dict, "model.", config)
+
+    assert torch.equal(state_dict["model.per_layer_model_projection.weight"], original[:16])
+    assert torch.equal(
+        state_dict["model.per_layer_model_projection_consumer.weight"],
+        original[16:],
+    )