perf(gemma4): optimize prefill and decode

.github/actions/build-hf-adapters/action.yaml

@@ -123,6 +123,18 @@ runs:
         head -40 "$HOME/.bashrc" || true
         echo "::endgroup::"
 
+        # Tests use `uv run --active --no-sync`. Install into that same
+        # environment, not a second .venv under the checkout. The image owns
+        # Torch's native build (in particular on ppc64le, without a CPU wheel).
+        # Resolve dependencies without an implicit sync, then leave Torch and
+        # unrelated image packages intact in the one explicit sync below.
+        : "${VIRTUAL_ENV:?The runner must activate its Python environment}"
+        export UV_PROJECT_ENVIRONMENT="$VIRTUAL_ENV"
+        CI_PYTHON="$VIRTUAL_ENV/bin/python"
+        test -x "$CI_PYTHON"
+        TORCH_FINGERPRINT='import hashlib, json, pathlib, torch; p = pathlib.Path(torch._C.__file__); print(json.dumps([torch.__version__, str(p.resolve()), hashlib.sha256(p.read_bytes()).hexdigest()]))'
+        TORCH_BEFORE=$(TORCH_DEVICE_BACKEND_AUTOLOAD=0 "$CI_PYTHON" -c "$TORCH_FINGERPRINT")
+
         # Build group flags from space-separated list
         GROUP_FLAGS=""
         for group in $UV_GROUPS; do
@@ -142,17 +154,37 @@ runs:
           git submodule update --init --recursive --force
           git log -1 --oneline
           cd "$CLONED_HF_ADAPTERS_DIR"
-          uv add torch-spyre /home/senuser/torch-spyre
-          uv sync --verbose --refresh $GROUP_FLAGS
+          uv add --no-sync torch-spyre /home/senuser/torch-spyre
         else
           cd "$CLONED_HF_ADAPTERS_DIR"
           # though we pin to main
           # frozen installation takes pinned commit when locked
           # we want to relax on torch-spyre to use floating version
           # and have other dependencies locked
-          uv add torch-spyre /home/senuser/torch-spyre
+          uv add --no-sync torch-spyre /home/senuser/torch-spyre
           uv lock --upgrade-package torch-spyre
-          uv sync --frozen --verbose --refresh $GROUP_FLAGS
         fi
+        uv sync --active --frozen --inexact --no-install-package torch --verbose --refresh $GROUP_FLAGS
+        TORCH_AFTER=$(TORCH_DEVICE_BACKEND_AUTOLOAD=0 "$CI_PYTHON" -c "$TORCH_FINGERPRINT")
+        if [[ "$TORCH_BEFORE" != "$TORCH_AFTER" ]]; then
+          echo "::error::Dependency installation changed the image's native Torch build"
+          exit 1
+        fi
+        uv pip check --python "$CI_PYTHON"
+        # Exercise exactly the interpreter-selection path used by Makefile.
+        # Fail during setup, rather than dispatching a suite with missing deps.
+        uv run --active --no-sync python - <<'PY'
+        import os
+        import pathlib
+        import sys
+        import torch
+        import transformers
+        import torch_spyre
+
+        assert pathlib.Path(sys.prefix).resolve() == pathlib.Path(os.environ["VIRTUAL_ENV"]).resolve()
+        if "oot" in os.environ["UV_GROUPS"].split():
+            import oot_framework
+        print(f"Test interpreter: {sys.executable}; torch: {torch.__version__}; transformers: {transformers.__version__}")
+        PY
         uv pip show torch
         uv pip freeze

.github/workflows/_test_matrix.yaml

@@ -1495,7 +1495,8 @@ jobs:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           cd "$CLONED_HF_ADAPTERS_DIR"
-          source .venv/bin/activate
+          # build-hf-adapters prepared the active image environment, not .venv.
+          : "${VIRTUAL_ENV:?The build action must prepare the active test environment}"
           make model-module-tests MODULE_CONFIG=${{ matrix.config }} PYTEST_ARGS='${{ inputs.extra_test_flags }}' \
             JUNIT_XML=1 RESULTS_DIR=.
           mv "model-module-${{ matrix.config }}.xml" "junit-model-module-${{ matrix.config }}.xml"
@@ -2434,7 +2435,8 @@ jobs:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
           cd "$CLONED_HF_ADAPTERS_DIR"
-          source .venv/bin/activate
+          # build-hf-adapters prepared the active image environment, not .venv.
+          : "${VIRTUAL_ENV:?The build action must prepare the active test environment}"
           make model-module-tests MODULE_CONFIG=${{ matrix.config }} PYTEST_ARGS='${{ inputs.extra_test_flags }}' \
             JUNIT_XML=1 RESULTS_DIR=.
           mv "model-module-${{ matrix.config }}.xml" "junit-model-module-${{ matrix.config }}.xml"

Makefile

@@ -153,7 +153,7 @@ model-module-tests: ## Run oot_framework module tests (suite key: model_module;
 	set -e; \
 	_run_test=$$(uv run --active --no-sync python3 -c \
 	  "import oot_framework, os; print(os.path.join(os.path.dirname(oot_framework.__file__), 'run_test.sh'))") || { \
-	  echo "ERROR: oot_framework is not installed in the active venv. Run 'uv sync --group oot' (see CLAUDE.md) and retry."; \
+	  echo "ERROR: oot_framework is not installed in the active venv. Install the oot group into $$VIRTUAL_ENV (see .github/actions/build-hf-adapters/action.yaml) and retry."; \
 	  exit 1; \
 	}; \
 	configs="$(MODULE_CONFIG)"; \

hf_adapters/hf_common.py

@@ -21,6 +21,7 @@
 compiled block functions.
 """
 
+import inspect
 import math
 import os
 import sys
@@ -1951,6 +1952,50 @@ def select_next_token(
     return (tokens, scores) if return_scores else tokens
 
 
+def _prefill_next_logits(logits, *, last_row_only=True):
+    """Copy the row generation consumes, without changing model forward.
+
+    On Spyre, transferring an offset view can convert its entire underlying
+    allocation. Materialize the selected row on device before the CPU copy.
+    Only movement changes; vocabulary cropping remains downstream. The optional
+    False override retains the full transfer for controlled comparisons.
+    """
+    if last_row_only:
+        return logits[:, -1:, :].clone().to("cpu")[:, 0, :]
+    return logits.to("cpu")[:, -1, :]
+
+
+def _generation_forward_options(run_forward_fn, last_hidden_row_only=None):
+    """Use the bounded head automatically when the driver supports it.
+
+    None selects automatically, False is a comparison opt-out, and an explicit
+    True still rejects an unsupported driver before generation touches caches.
+    """
+    if last_hidden_row_only is False:
+        return {}
+    if run_forward_fn is None:
+        if last_hidden_row_only is None:
+            return {}
+        raise ValueError("The forward driver must declare _last_hidden_row_only")
+    try:
+        parameter = inspect.signature(run_forward_fn).parameters.get(
+            "_last_hidden_row_only"
+        )
+    except (TypeError, ValueError):
+        # Some valid callable drivers expose no Python signature. Automatic
+        # selection must leave their existing calling convention unchanged.
+        parameter = None
+    if parameter is None or parameter.kind not in (
+        inspect.Parameter.POSITIONAL_OR_KEYWORD,
+        inspect.Parameter.KEYWORD_ONLY,
+    ):
+        # A **kwargs-only driver could silently ignore this request.
+        if last_hidden_row_only is None:
+            return {}
+        raise ValueError("The forward driver must declare _last_hidden_row_only")
+    return {"_last_hidden_row_only": True}
+
+
 def generate(
     run_forward_fn: Optional[Callable],
     model,
@@ -1967,6 +2012,8 @@ def generate(
     top_p=None,
     eos_token_id=_UNSET,
     timing=False,
+    _prefill_last_row_only=True,
+    _generation_last_hidden_row_only=None,
     prefill_fn: Optional[Callable] = None,
     decode_fn: Optional[Callable] = None,
     token_aligned_inputs: Optional[dict[str, tuple[torch.Tensor, Any]]] = None,
@@ -2026,6 +2073,12 @@ def generate(
             each prefill chunk. Falls back to the adapter's configured chunk
             size, or one-shot prefill when the adapter has no override.
     """
+    # Prefill branches are mutually exclusive: a custom callback replaces the
+    # text driver. Inspect the one that will actually receive the keyword.
+    prefill_driver = prefill_fn if prefill_fn is not None else run_forward_fn
+    forward_row_kwargs = _generation_forward_options(
+        prefill_driver, _generation_last_hidden_row_only
+    )
     overrides = {
         "max_new_tokens": max_new_tokens,
         "max_length": max_length,
@@ -2183,6 +2236,7 @@ def generate(
                     value_caches=prefill_value_caches,
                     cache_index=make_cache_index(0, padded_len, DEVICE),
                     **normalized_token_inputs,
+                    **forward_row_kwargs,
                 )
             else:
                 # Keep Lk fixed at the complete prefill extent while advancing
@@ -2208,9 +2262,12 @@ def generate(
                         cache_index=make_cache_index(
                             chunk_start, query_chunk_size, DEVICE
                         ),
+                        **forward_row_kwargs,
                     )
             # Only the last chunk's logits matter for next-token selection.
-            next_logits = logits.to("cpu")[:, -1, :]
+            next_logits = _prefill_next_logits(
+                logits, last_row_only=_prefill_last_row_only
+            )
             current_cache_len = padded_len
 
         else:

hf_adapters/hf_gemma4.py

@@ -756,8 +756,18 @@ def _run_forward(
     key_caches,
     value_caches,
     cache_index,
+    *,
+    _last_hidden_row_only=False,
 ):
-    """Gemma 4 causal-LM forward: backbone + LM head + logit softcap."""
+    """Gemma 4 forward; generation may request only the final head row.
+
+    The default returns every row. The opt-in leaves the backbone and all
+    cache updates intact, selecting a row before the existing vocabulary head
+    and softcap. Chunked generation requests one head row for each chunk,
+    while every chunk still runs its complete backbone and cache updates.
+    A differently shaped matmul may round differently; this is
+    not a promise of bitwise equality to the full-row device calculation.
+    """
     h = _run_backbone_forward(
         model,
         input_ids,
@@ -768,6 +778,8 @@ def _run_forward(
         cache_index,
     )
 
+    if _last_hidden_row_only:
+        h = h[:, -1:, :]
     logits = model.lm_head(h)
 
     cap = text_config(model.config).final_logit_softcapping