Fix prefill pruning for Gemma 4 models

[sushraja-msft]

Gemma 4 models have shared KV cache

• For the E2B model the KV cache of the first 15 layers are shared with the remaining 20 layers.
• This means during prefill there is no need to compute beyond 15 layers, only decode needs to run the full graph. Effectively 2B model is a 0.6B model for prefill purposes.
• Without this optimization prefill perf will suffer on CPU/GPU.
  This change extends the prune_lm_head concept, now instead of pruning at the lm_head level - gemma models prune at a much higher layer, at level 15 onwards.

The brings the prefill TTFT from 845ms on a Nvidia 4070 to 183ms.

Mobius is also fixed to generate graph capture compatible models for Gemma 4.

[CLAassistant]

All committers have signed the CLA.

[Copilot]

Pull request overview

This PR introduces a more general “prefill pruning” build feature and applies it to Gemma 4 models that have KV-cache sharing across a layer suffix, enabling prefill to stop carrying full-sequence activations beyond the first KV-sharing layer while keeping decode behavior intact.

Changes:

• Replace/rename the previous LM-head-only pruning option with prune_prefill_prefix, wiring it through tasks, build APIs, CLI --features, and docs.
• Implement Gemma 4–specific pruning for KV-sharing layer suffix + per-layer inputs, plus weight splitting for the per-layer projection to support the new structure.
• Refine ORT-GenAI config generation so WebGPU graph capture is enabled for decoder only in multimodal packages.

Reviewed changes

Copilot reviewed 17 out of 17 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/gemma4_prefill_prefix_test.py	Adds coverage validating Gemma 4 shared-layer prefix pruning behavior and weight splitting.
tests/cli_test.py	Updates CLI --features plumbing tests to use prune-prefill-prefix.
src/mobius/tasks/_gemma4.py	Adds prune_prefill_prefix support and scopes pruning via build context.
src/mobius/tasks/_causal_lm.py	Renames task option and switches to the new pruning context manager.
src/mobius/models/gemma4.py	Implements Gemma 4 KV-sharing suffix pruning, typed scalar constants, and per-layer projection weight split.
src/mobius/models/base.py	Renames pruning helper and ties pruning behavior to the new build context flag.
src/mobius/models/_models_test.py	Updates pruning feature tests to prune_prefill_prefix.
src/mobius/integrations/ort_genai/genai_config.py	Allows session options to override graph capture per submodel.
src/mobius/integrations/ort_genai/genai_config_test.py	Adds regression test ensuring graph capture is decoder-only in multimodal exports.
src/mobius/integrations/ort_genai/ep_config.py	Adds graph_capture override to provider options generation.
src/mobius/_builder.py	Renames the public build flag and extends task enablement for Gemma 4 tasks.
src/mobius/_build_context.py	Renames build context flag and context manager to prefill-prefix pruning.
src/mobius/_build_context_test.py	Updates context tests for the renamed pruning flag/context manager.
src/mobius/main.py	Renames build feature key to prune-prefill-prefix and plumbs through to build().
README.md	Updates listed CLI features and example to prune-prefill-prefix.
docs/cli_reference.md	Updates feature documentation for prune-prefill-prefix.
CHANGELOG.md	Documents the new prune-prefill-prefix feature and its Gemma 4 implications.

Suppressed comments (3)

src/mobius/_builder.py:226

• The build_from_module() docstring overstates the behavior of prune_prefill_prefix for the generic CausalLMTask: it only selects the final token before the LM head (logits shape [B, 1, vocab]). Discarding earlier tokens from remaining decoder computation only happens in certain model/task implementations (e.g. Gemma 4 KV-sharing suffix).

        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.

src/mobius/models/gemma4.py:2134

• total_seq_len is documented (and elsewhere implemented) as attention_mask.shape[1], but this code now derives it from ReduceSum(attention_mask)[0]. That changes semantics when attention_mask contains padding zeros and hard-codes batch element 0. Consider restoring the Shape()-based scalar for normal builds and using the ReduceSum/Gather workaround only for EPs that require graph-capture rewrites.

            position_embeddings_dict: dict = {
                "sliding_attention": None,
                "full_attention": None,
            }
            # In hybrid mode, static-cache full-attention layers need RoPE

src/mobius/integrations/ort_genai/ep_config.py:43

• The docstring says graph capture is driven entirely by the EP registry flag, but this function now accepts a graph_capture override (used to disable capture for non-decoder submodels). The docstring should reflect the actual behavior.

    """Build the ``provider_options`` list for genai_config.json.

    Graph capture is driven entirely by the EP's registered
    ``EpCapabilities.enable_graph_capture`` flag (the single source of truth).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.