perf(gemma4): select decode down weights in output blocks

[AdnanHoque]

Consolidated into HF #485 — perf(gemma4): optimize prefill and decode. This PR is superseded following the author's request for one HF review/merge unit. Its optimization is retained in the combined source; current dependencies, smaller integrated tests and validation status are documented there. The original branch and historical body below are preserved for recovery and evidence. Old CI and timings do not validate the new combined head.

---

Default-on/automatic selection update requested by the author. Safety checks remain. Review readiness is based on CI for the current head, not the earlier opt-out head; numerical and performance acceptance remain separate.

CI review — 2026-09-08

Current head d7de86d7: run 34208548617 is complete with 132 successful and 14 skipped checks, no failures or pending checks. The PR is mergeable and was marked ready for review after that verification. Source was unchanged in this sweep. This is review readiness, not a new speedup claim: P3 is the code parent, P2 is the compiler capability, and the recorded timing/placement control also included P5. Those dependency and numerical boundaries below remain unchanged.

Automatic behavior and current validation

Current defaults-update head: d7de86d75166ca1437d0eb3808a7001ff3a3c647. This replaces the initial opt-out/default-off head below.

The default down-output block width is 1024. It activates only with the eligible paired P3 schedule and matching FP16 or BF16 H2816/F704 inputs/weights (BF16 is the checkpoint's actual dtype). Other shapes and dtypes retain the ordinary full-reduction BMM. None remains the control; unimplemented block-width overrides are rejected. The branch includes the updated P3 parent and preserves the sibling relationship with P5.

9 host tests pass including inherited P3 coverage. Default dispatch executes the five expected BMM shapes on BF16 meta tensors; CPU repeated-index and tail coverage remains. These checks do not prove device LX placement.

No new device timings were collected for this defaults change. Historical measurements below retain their original boundaries and must not be attributed to this new head without a matched rerun.

Claude's independent review caught an FP16-only eligibility condition that would have excluded the actual BF16 checkpoint. It is corrected: both host 16-bit formats use Spyre's SEN169_FP16 device format; inputs must have matching dtypes, and FP32/mixed inputs decline. Default-path host tests now cover BF16. This is a dispatch correction, not new numerical or performance acceptance.

A separate local composition check passed all four feature combinations (neither, down only, gate/up only, both) plus both full-gather guard checks: 6 host checks. Both enabled emits eight 704-wide gate/up BMMs followed by down widths 1024/1024/768. All full down selections remain inside down_panel is None; full gate/up selections remain in the no-gate/up-block branch. This proves source dispatch/shape composition, not device residency or numerical acceptance. Matching FP16/BF16 pairs are allowed; mixed 16-bit formats are explicitly tested to decline.

Current GitHub diff against main: 4 files, +493/−12, including the documented prerequisites. Earlier sizes and feature comparisons below describe the original publication, not this updated head.

Original publication boundary

The initial publication used the prepared source and commit recorded below; the automatic/default-on revision is recorded separately above.

Feature parent: 6cf881d301ab590f037fbdbec2b42fb5d3dcfafc. Feature-only diff: 2 files; production +42/−2; total +170/−2 including tests/docs/CI.

GitHub targets main, so its diff also contains the route-schedule parent (P3 / #486). This is a sibling of the other weight-blocking change, not a flattened independent implementation.

Change and reason

Instead of selecting each expert's entire down matrix first, slice columns from the original weight bank and then select experts. For [128,704,2816], the blocks are columns 0–1023, 1024–2047 and 2048–2815: widths 1024, 1024 and 768.

Each block computes a complete 704-term dot product. Concatenate columns in their original order. Unlike gate/up reduction blocking in P5 / #488, this does not introduce partial-sum additions.

Default-on _DECODE_DOWN_OUTPUT_PANEL=1024 requires P3 / #486, one token/eight routes and the supported hidden/expert widths 2816/704. The down assignment is eight route owners with no output or reduction split (R:8,H:1). The full down selection is not constructed on this path. Unsupported input shapes/dtypes fall back; unimplemented width overrides raise.

Dependencies: code versus placement evidence

Code depends on adapter P3 / #486 and compiler P2 / #4349. This extracted branch deliberately does not require P5 / #488's arithmetic change. However, the historical winning experiment DID include P5 / #488, and smaller gate/up buffers can affect whether down blocks fit on chip.

Therefore sibling placement is proposed, not device-proven. Before merge or performance acceptance, validate this extraction with ordinary gate/up and with P5 / #488. If on-chip placement requires P5 / #488, make that an explicit placement dependency and stack it; do not claim independence from a CPU geometry test.

No pending LX PR, alternate bank, four-core reduction split, new backend operation or capacity-check bypass.

Current compiler dependency: P2 / #4349's corrected cost-selection source, with explicit legal-target conversion costing and ordinary-layout tie preference. The adapter now selects its supported optimization automatically; its historical placement and gain must still be re-established with that corrected chooser. A configuration name alone does not prove the same selected layout. See the correction receipt (retained in the local preparation package).

Recorded performance

CP18's whole-token immediate control already included P5 / #488:

Boundary	Immediate control	Down blocks	Change
Decode token	187.397 ms	163.515 ms	23.882 ms; 12.7% lower
Complete expert layer	2.948 ms	2.501 ms	five paired improvements

The whole-token result used five serial warmed runs per arm and matching tokens; it is not a freshly interleaved acceptance run on this extracted sibling. The change also moved the prior down assignment to eight routes/no output split. A separate fixed-H:1 layer control measured about 3.038 → 2.489 ms, supporting that blocking itself matters.

Frozen emitted-program evidence showed the three intended LX panels, per-owner bytes 1,441,792 / 1,441,792 / 1,081,344, with original-bank column offsets 0/1024/2048. Modeled removed temporary writes and rereads total 63,438,848 bytes per layer; this is modeled traffic, not a hardware counter.

Validation and limits

Four added CPU tests pass (nine including inherited P3 tests): exact column/tail coverage, repeated and boundary experts, slicing before selection, full reduction width, requested assignment, no unconditional full gather, and rejection without P3 / #486.

Recorded layer captures were exactly equal for the down-block comparisons. The extracted sibling still needs native compile, original-bank address/LX placement inspection and existing HF token/logit tests. CPU tests do not prove storage reuse, capacity, emitted ownership or full-model numerical acceptance.

Scope

One adapter production file and one incremental test file. P3 / #486's inherited helper/tests are not counted again in this PR's size.

Relationship to the added prefill drafts

P7 / #4350 (compiler weight-copy proofs), P8 (pending in #4347) (reader-compatible input staging), and P9 / #490 (prefill matmul divisions) form a separate relayout-dependent track. They are not new prerequisites of this draft. Their 24.4% and 7.8% increments are expert-layer measurements on a different frozen composition; do not add them to this draft's historical full-model result. See the shared inventory (retained in the local preparation package) for code versus measured-performance dependencies and the posted relayout PRs. Existing code/acceptance boundaries above are unchanged.

Companion PRs and dependencies

Feature	PR	Dependency / review boundary
P1	torch-spyre/hf-adapters#485	P6 overlaps its avoided output-copy work; no compiler prerequisite.
P2	torch-spyre/torch-spyre#4349	Compiler capability used by P3/P4/P5; old paired gain must be remeasured after the chooser correction.
P3	torch-spyre/hf-adapters#486	Requires P2; parent of P4 and P5.
P4	torch-spyre/hf-adapters#487	Code child of P3; P5 was also present in the recorded placement/timing control.
P5	torch-spyre/hf-adapters#488	Code sibling of P4 on P3; changed arithmetic remains unaccepted.
P6	torch-spyre/hf-adapters#489	Independent output-head change; saved transfer work overlaps P1.
P7	torch-spyre/torch-spyre#4350	Frozen relayout composition; P8 is a measured placement prerequisite on a different parent.
P8	#4347 — existing owner, P8 pending	Included in automatic-selection owner #4347; contribution still local, not present in that published head.
P9	torch-spyre/hf-adapters#490	Requires relayout capabilities plus P7 and the P8 contribution; automatic only with the required compiler composition; P8 still pending.

P4/P5 are code siblings, not cumulative speedups. P7/P8/P9 require a fresh exact-head compiler/adapter composition. The historical 4× full-model prefill and 8.8× expert-FFN gains are separate already-landed work.

GitHub diff and merge status at publication

At initial publication GitHub reported 4 files, +398/−12 against main. This includes the unmerged prerequisite content described above; use the feature-only comparison for this optimization's incremental scope. Draft/ready status is tracked on each current PR; old CI does not apply to a changed head. CI, final-base reconciliation and device acceptance are separate merge gates; source publication does not make them green.

[AdnanHoque]

Consolidated into #485, now published at 0cf04e7 with all six HF optimizations. Following the request for one review unit, this PR is superseded rather than dropped. The combined description records compiler dependencies, historical per-feature results, reduced tests and remaining numerical/device acceptance. Original feature commits and this branch are preserved; no branch deletion.