perf(gemma4): optimize prefill and decode

[AdnanHoque]

Summary

Consolidates Gemma performance work from #485, #486, #487, #488, #489 and #490 into one review and merge unit, following Ariel's request. #485 is the surviving PR; the other five are superseded. The compiler changes remain separate.

Refreshed head: 47ec92a3b46be3a0d9ac40e5c0990a7b2c1f0c5e, incorporating HF main 1b3c1a6d163977940779564cfb1a53f5aef9014c (September 10). Signed merge ab550856 preserves the published feature history; e484f49b clarifies the read-copy capability wording. Separate commits 1c0f425 and 47ec92a fix the CI environment mismatch at installation and both module-test callers. Nine files, +719/−38: three model/generation files (+365/−30), three existing test files (+313/−1), and three CI/setup files (+41/−7). No new test file or public dependency change.

All supported optimizations are automatic by default. Unsupported shapes, dtypes or compiler capabilities retain the existing path. Private controls remain only for comparisons and explicit support checks; no new public model API, generation loop, weight bank, cache representation or backend operation is introduced.

What changes

Prefill output work

Generation materializes only the final logit row before copying it to the CPU. Materialization matters: copying an offset view alone can convert the entire underlying allocation.

For Gemma's forward driver, generation also selects the last hidden-state row before the vocabulary projection and existing softcap. The backbone still processes all tokens and updates all caches. Ordinary forward still returns all rows, and decode is unchanged. Selection follows the actual prefill callback, including custom callbacks and chunked text prefill; callables without an inspectable signature retain their existing calling convention.

These two changes overlap: once the head computes one row, there is no full prefill-logit allocation for the copy optimization to remove.

In chunked text prefill the head processes one row per chunk, not just on the final chunk. Only the final chunk supplies generation's first token. The recorded one-shot timings do not measure this chunked case.

Decode weight delivery

The one-token/eight-route case pairs eight route owners with the compiler's consumer-compatible indexed-selection layout. The layout and assignment must be enabled together.

• Gate/up: slice each unchanged expert bank into four 704-term reduction blocks before selecting experts. Compute gate then up for each block and add partial products in the existing experimental order. This changes rounding versus one full 2816-term dot product.
• Down: slice output columns before expert selection: widths 1024, 1024, 768 cover the original 2816 columns exactly. Each result still uses the complete 704-term reduction; concatenate output columns, not selected weights.
• Both paths support matching FP16/BF16 H2816/F704 inputs, one token/eight routes and the supported compiler. Other inputs retain ordinary decode. Their full selected-weight gathers are independently guarded out when blocking is active.
• No output/reduction core split is forced for these indexed loads. The compiler retains ownership, layout and capacity decisions.

Prefill expert assignments

When the required compiler composition is available for matching FP16/BF16 E128/T512/H2816/F704, use gate/up token-by-hidden 8×4, down token-by-output 16×2, and retain 32 token groups around the persistent expert loop. Routing, GELU, scaling, original banks and expert accumulation are unchanged.

Configuration is scoped to the existing prefill/decode compile calls. Explicit unsupported prefill requests fail before attention/cache writes; named dimensions are cleaned up on errors.

Dependencies

Included feature	Required compiler work	Missing support
Final-row transfer and vocabulary projection	No new compiler PR	Unsupported forward drivers keep full-row behavior
Decode route assignment and both weight-blocking changes	torch-spyre#4349, indexed-selection layout choice	Keep ordinary decode; passing fallback tests does not validate the optimized path
Prefill expert assignments	LX ownership/transfers, completed-sum support and reader-compatible input-stage selection in torch-spyre#4347; torch-spyre#4350 additionally enables the recorded copy-removal performance	Missing schedule capabilities retain ordinary persistent prefill; a missing direct-read proof retains the copy

Dependency check, September 10: #4283/#4284 are merged, but compiler main 447d19e1 still lacks the required decode and input-stage capabilities. Refreshed #4349 at 3bf56334 supplies the decode capability. Refreshed #4347 at a9a4fd15 supplies consumer_compatible_input_staging through the existing chooser and separates generated-copy hint applicability from source history; it does not rewrite assignments after selection. Its source chain remains #3440 → #4152 → #4153 → #4347, with sibling #3955's completed-sum integration. Refreshed #4350 at bede205d is independently main-based. The existing read_copy_elision option proves neither that #4350 is present nor that any copy was removed: absent preservation, the existing final proof retains the copy safely. No extra capability flag or LX/core-assignment source is added to the adapter. Composed expert-function device qualification now demonstrates P8 engagement and LX activation staging; this is not full-model acceptance.

Decode's historical placement results included both blocking changes. Source independence of the two controls is not a promise of independent on-chip residency or additive gains.

The fresh P8-only measurement in the composed expert function is 29.726537 → 28.997294 ms wall time, 2.4831% median paired saving across five improving pairs, with exact OFF/ON outputs and final activation LX placement. This is an incremental compiler-candidate result on the complete expert function, not new whole-model or decode evidence for this HF PR; details and controls are recorded in #4347.

Recorded performance — historical, not measurements of this consolidated head

Batch 1; original source/runtime pins and raw receipts remain in the preparation package. Each row has its own immediate control.

Change	Boundary and recorded comparison	Meaning / limitation
Final-logit-row transfer (CP11)	First-token window 2376.765 → 1769.939 ms, 25.5% lower	Five warmed serial runs/arm; numerical disposition remains open
Paired route assignment + indexed-selection layout	Decode 288.236 → 241.821 ms, 16.1% lower	Pair only; predates #4349's corrected cost chooser and requires remeasurement
Gate/up blocking (CP16)	Decode 222.823 → 187.674 ms, 15.8% lower; layer 4.043 → 2.956 ms	Control already included the route/layout pair and a different down assignment; changed arithmetic
Down blocking (CP18)	Decode 187.397 → 163.515 ms, 12.7% lower; expert layer 2.948 → 2.501 ms	Control already included gate/up blocking; not an independent percentage
Final-hidden-row head (CP21)	1685.716 ms prefill ON versus prior CP19 1741.601 ms	Apparent 3.2% is not a paired gain: fresh OFF control faulted
Prefill down 16×2 assignment	Expert function 31.364 → 28.912 ms, 7.8% lower	Five interleaved pairs; control already had gate/up 8×4, relayout, copy removal and private stage selection; not whole-model prefill

Do not add these percentages. The preceding compiler copy-removal result (41.521 → 31.406 ms, 24.4% lower expert-function time) belongs to #4350, not to this HF PR. Historical accepted 4× full-model prefill (6.4 → 1.6 s) and 8.8× expert-FFN results belong to the earlier merged campaign, not this diff.

Smaller, integrated test coverage

Removed seven new standalone test/helper files (1,504 lines), including source-parsing/AST scaffolding. Original commits and superseded branches retain them for recovery. No existing model test is removed or weakened.

Replacement coverage lives in existing suites:

• tests/cpu/test_generate_cpu.py: actual generation and imported Gemma forward; one-shot/chunked/custom callback paths; unchanged cache-write positions and decode rows; vocabulary cropping, returned logits/scores, bounded head and materialized row storage; unsupported/uninspectable drivers.
• tests/cpu/test_adapter_cpu_accuracy.py: imported MoE implementation; both block controls together and separately for FP16/BF16; actual dispatch recording of every indexed-selection source and BMM shape; repeated indices/tails; capability/dtype/shape fallback and explicit prefill rejection before cache writes.
• tests/spyre/test_e2e_token_compare_spyre.py: reuse the already-loaded Gemma model and CPU reference to also check the real generation path, with fresh request caches. The old forward-only loop bypassed both generation optimizations. Existing model smoke, forward and token checks remain.

Refreshed local result: 61 selected checks pass (18 focused Gemma/generation cases plus 43 existing input-normalization/cache cases). Ruff 0.15.6, Black 26.3.1, whitespace and mypy over all 43 adapter source files pass. These are host checks using Python 3.12, PyTorch 2.14.0 and Transformers 5.15.0, not a match to the earlier device-CI runtime. CPU/meta checks prove dispatch and algebra, not generated device placement, reduction numerics or current-head performance. The CPU test matrix is currently disabled upstream; these local checks are not presented as completed device CI.

Review and remaining acceptance

Two independent source-review passes found no remaining blocker after clarifying that the existing read-copy control is not a probe for #4350. The review covered the callable-driver fix, all four decode block combinations and pre-cache prefill validation. This is AI-assisted source review, not human maintainer approval or device acceptance. Optimization-specific validation can assert the existing eligibility predicate and inspect the emitted program; no additional production enablement flag is added.

Features remain traceable in the original signed source commits, including the rounding-changing gate/up block operation. This refresh preserves that history and adds upstream main without replaying or duplicating the optimization changes. The six-file feature diff and separate three-file CI repair remain independently reviewable.

CI environment repair

The e484f49b run failed before model execution across the sampled suite families: dependency installation created a checkout .venv, while uv run --active --no-sync ran tests from the image's /home/senuser/.venv. Missing transformers/oot_framework were the result. Repeating a fresh pod did not change that contract mismatch.

Commit 1c0f425 sets the project environment to the already-active environment, makes both uv add calls resolve without implicit installation, and performs one explicit active sync. It preserves image Torch and unrelated image-only packages, compares Torch's version/path/native-binary hash before and after, checks dependency compatibility, and imports the required packages through the exact test invocation. The Makefile's missing-oot guidance points to that setup. No shared installation is changed by the local validation, and no test or tolerance is disabled.

The real uv command passed an isolated wheel-fixture check: image-owned package unchanged, unrelated image package retained, new dependency importable from the test command, and no second environment created. Fresh runner CI then passed this installation/probe, printing the intended image interpreter, Torch 2.13.0+cpu and Transformers 5.15.0. It exposed two module-test callers that still activated the absent checkout .venv; 47ec92a removes those environment switches and requires the prepared active environment instead. The extended fixture audits every build-action caller and executes both module caller prefixes successfully without a checkout .venv. Unrelated workflows that deliberately create their own environment are untouched. The latest full CI run remains required.

Fresh combined device CI and capability-present model acceptance are required. A green run without the required compiler capabilities validates fallbacks, not blocked decode/prefill performance.

The previous CI run 34242808842 is not a result for this refresh: Gemma token comparison, including actual generation, passed, but its smoke test failed first in attention with a bus-fence error and on retry in the 64-token prefill FFN with a device compute error. Ministral and RoBERTa passed their retries. No treatment-specific source defect was identified, but the Gemma fault remains unattributed, not proven to be infrastructure. The new head needs fresh CI; no device failure or numerical threshold is suppressed by this refresh.

The historical gate/up comparison recorded maximum absolute difference 0.09375 and relative L2 about 0.00288048; that is a measurement, not an accepted error limit. CP11 also has an unresolved seeded all-NaN receipt. Neither consolidation nor passing host checks closes the correctness owner's disposition. No numerical limit is widened here.

Before performance acceptance, verify exact compiler/adapter revisions, real checkpoint logits/tokens, intended on-chip blocks/original-bank addresses and repeated whole-model timing. No new device speedup is claimed in this consolidation.

[AdnanHoque]

Implemented change: consolidate the Gemma performance work into one HF PR

Agreed concern: Six interdependent HF PRs and seven standalone test/helper files made review and merging unnecessarily fragmented.

Implemented approach:

1. This PR now contains all six adapter optimizations: final-row logit copying, final-row vocabulary projection, paired decode route scheduling, gate/up reduction blocks, down output blocks, and capability-gated prefill assignments. perf(gemma4): pair decode route scheduling with indexed-selection layouts #486–490 are closed as superseded, not abandoned; their branches and feature commits are preserved.
2. Remove the seven standalone test/helper files (1,504 lines). Retain targeted coverage in the existing CPU generation/adapter suites, and extend the existing Gemma device token-comparison test to exercise real generation as well as ordinary forward. No new test file.
3. Keep all eligible optimizations automatic. Preserve fallback/safety checks, scoped compiler settings, ordinary forward and cache semantics. Fix automatic selection for callable drivers without an inspectable Python signature.
4. Keep compiler responsibilities separate: decode requires perf(moe): offer consumer-compatible layouts for indexed selections torch-spyre#4349; enhanced prefill needs the LX/completed-sum composition, perf(moe): preserve direct weight reads through compiler rewrites torch-spyre#4350, and the still-unpublished input-stage contribution assigned to LX planning: choose work divisions using placement cost torch-spyre#4347. Current fallback runs do not validate those optimized paths.

Validation and review: 48 selected host checks pass, plus Ruff/Black and mypy over 41 source files. Claude's independent combined-source review found no blocking code defect; its nonblocking points were resolved or documented. Fresh combined device CI and capability-present numerical/placement acceptance remain separate. The description preserves each historical speedup with its actual immediate control; percentages are not additive, and none is claimed as a new measurement of this head.

Scope: Six files, +676/−31; three production files and three existing test files.

Addressed in 6b01e51.

[AdnanHoque]

Author update — refreshed source and corrected CI environment

Implemented change

Head 47ec92a retains the main refresh and the consolidated Gemma optimizations. Separate CI commits fix a demonstrated setup mismatch: dependencies were installed into a checkout .venv, while the test command ran from the active image environment. That caused missing transformers/oot_framework before model execution across the sampled suite families.

Implemented approach

1. Resolve without implicit installation, then install the requested groups once into the same active environment the tests use.
2. Preserve image Torch and unrelated image-only packages. Check Torch's version, location and native-binary hash before/after installation; fail if they change.
3. Check dependency compatibility and import required packages through the exact test command before running suites. Correct the Makefile's missing-dependency guidance.
4. Keep both module-test callers in that environment; neither reactivates a checkout .venv. Unrelated workflows that create their own environment are unchanged.

This fixes environment identity at setup; it does not skip tests, weaken tolerances or modify shared development installations. Broader platform-specific image ownership is not invented by this change.

Scope and dependencies

The PR now has nine files, +719/−38: the unchanged six-file optimization/test diff plus three CI/setup files. No new test file or public model/backend interface.

• Final-row transfer and vocabulary-head work need no new compiler PR.
• Decode route assignment and selected-weight blocking require torch-spyre#4349 at 3bf56334.
• Prefill assignments require #4347 at a9a4fd15, preserving #3440 → #4152 → #4153 → #4347 and sibling #3955's completed-sum support. P8 uses the existing chooser; it does not override a committed assignment.
• Independently main-based #4350 at bede205d enables the recorded copy-removal performance. Without a proven direct read, the copy remains safely.

Validation and performance

The earlier refreshed model source passed 61 host checks and its pinned lint/type checks; no model source changed in this CI repair. The new setup passed shell syntax and a real-uv isolated fixture proving one environment, preservation of image packages and importability of a newly installed dependency. Fresh runner CI then passed the actual install/import probe, before exposing the two stale module callers. The extended fixture audits all build-action callers and successfully executes both module caller prefixes without a checkout .venv. Independent source review now approves both the installation and caller fixes. Current 47ec92a CI has progressed into the actual model/component/module tests; complete CI is not yet counted as green.

The composed expert-function device check confirms P8 engages: the activation stage changes from HBM to LX with matching 8×4 reader ownership, and ON/OFF outputs match exactly within the fixed CPU-reference screen. Five interleaved pairs now measure 29.726537 → 28.997294 ms expert-function wall time; median paired saving 2.4831%, with all five pairs improving beyond the control noise gate. This belongs to #4347's P8 increment, not full-model or decode acceptance of this HF PR. Fresh CI and capability-present model qualification remain open, so this stays draft. Historical gains and their distinct controls remain in the description; no percentages are added together or relabeled as fresh measurements.

[AdnanHoque]

/spyre-test

[spyre-ci]

❌ spyre-test: failure

Triggered by: /spyre-test comment: #485 (comment)

Plan (build waves + dependencies, per arch)

amd64

flowchart LR
  subgraph Lamd64_0["amd64 L0 · 1 parallel"]
    n_amd64_torch_spyre_torch_spyre_dev["torch-spyre/torch-spyre-dev 🔴<br/>image · fc1752087216"]
  end
  subgraph Lamd64_1["amd64 L1 · 1 parallel"]
    n_amd64_hf_adapters_hf_adapters_dev["hf-adapters/hf-adapters-dev 🟢<br/>image · c923b5d43e4c"]
  end
  subgraph Lamd64_2["amd64 L2 · 1 parallel"]
    n_amd64_spyre_inference_spyre_inference_dev["spyre-inference/spyre-inference-dev 🔴<br/>image · 83521da27449"]
  end
  n_amd64_torch_spyre_torch_spyre_dev --> n_amd64_hf_adapters_hf_adapters_dev
  n_amd64_hf_adapters_hf_adapters_dev --> n_amd64_spyre_inference_spyre_inference_dev
  classDef sPending fill:#eceff1,stroke:#90a4ae,color:#37474f
  classDef sBuilding fill:#fff8e1,stroke:#f9a825,color:#5d4037,stroke-width:2px
  classDef sOk fill:#e8f5e9,stroke:#43a047,color:#1b5e20
  classDef sReused fill:#e3f2fd,stroke:#1e88e5,color:#0d47a1
  classDef sFailed fill:#ffebee,stroke:#e53935,color:#b71c1c,stroke-width:2px
  classDef sDropped fill:#f5f5f5,stroke:#bdbdbd,color:#9e9e9e
  class n_amd64_torch_spyre_torch_spyre_dev sOk;
  class n_amd64_hf_adapters_hf_adapters_dev sOk;
  class n_amd64_spyre_inference_spyre_inference_dev sOk;

Loading

✅ orch trigger-pr-validation — green · arches amd64 · fp amd64=caad41d3

level	component	arch	build	smoke	unit	integration	trunk	regression	perf
L0	torch-spyre/torch-spyre-dev	amd64	✅ ok	⏳	·	🟡 gha	·	·	·
L1	hf-adapters/hf-adapters-dev	amd64	✅ ok	🟢	·	🟢 gha	·	·	·
L2	spyre-inference/spyre-inference-dev	amd64	✅ ok	🟢	·	🟡 gha	·	·	·

GHA test runs:

• [hf-adapters/amd64] integration: SUCCESS https://github.com/torch-spyre/hf-adapters/actions/runs/34545430444
• [torch-spyre/amd64] integration: FAILURE https://github.com/torch-spyre/torch-spyre/actions/runs/34545316703
• [spyre-inference/amd64] integration: FAILURE https://github.com/torch-spyre/spyre-inference/actions/runs/34545796909

---

⚠️ advisory failures only — mergeable, see below

Build: built 3

Tests: passed 3 · blocking 0 · advisory 2 · infra/inconclusive 0 · no signal 0

Failures by kind

• ⚠️ advisory (does not block) · torch-spyre/amd64 integration: UNSTABLE
• ⚠️ advisory (does not block) · torch-spyre/amd64 gha:integration: FAILURE
• ℹ️ informational · spyre-inference/amd64 gha:integration: FAILURE

Before merging, consider:

• torch-spyre/amd64 gha:integration, torch-spyre/amd64 integration failed with gating: "unstable" — advisory. It set the build UNSTABLE but does not block the merge. Worth a look, not a stop.
• spyre-inference/amd64 gha:integration failed with gating: false — informational only, no merge impact.

[spyre-ci]

GHA test runs:

• [torch-spyre/x86_64] integration: DISPATCHED https://github.com/torch-spyre/torch-spyre/actions/runs/34545316703

[spyre-ci]

GHA test runs:

• [hf-adapters/x86_64] integration: DISPATCHED https://github.com/torch-spyre/hf-adapters/actions/runs/34545430444

[spyre-ci]

GHA test runs:

• [spyre-inference/x86_64] integration: DISPATCHED https://github.com/torch-spyre/spyre-inference/actions/runs/34545796909