Description
The 26B's sliding-window family carries almost all of a decode step's KV bytes (the global family reads a fraction of a gigabyte per step against several for the window), and it stores them as bf16. The reference engine serves this checkpoint with an e4m3 KV cache at scale 1.0 by default, which is where its decode-step advantage on this box came from once the kernel accounts were corrected: the same window read at half the bytes. Our pool has no element width but bf16, so the sliding family cannot take the same trade.
Proposed Plan
KvLayout/KvPool carry elem_bytes (2 or 1); strides stay in elements, the backing length converts by bytes, and the geometry gate accounts for the width.- The shared launch header from slice 7 takes the KV element type as a template parameter; a feature-gated
csrc/gemma4/local_kv_fp8.cu instantiates the two hd256 windowed reads (batch prefill, split-KV decode) over e4m3, and the hd256 prep kernel takes a store type so it writes e4m3 at scale 1.0. The bf16 instantiations and bindings do not move.
PEGAINFER_KV_FP8=local opts the local pool in; unset keeps every byte as before. Wrappers without an fp8 twin refuse a one-byte layout at the operator boundary; the prefix-cache combination is refused at startup because its page copies index bf16 elements.
Acceptance Criteria
- With the variable unset, the pinned-composition prompts answer with the same bytes as
main.
- Under fp8, the 12B window-crossing gate and the greedy HF generation gate pass; the long-context waypoint gate stays within its logprob tolerance; and a distribution gate over a fixed prompt set shows the fp8 pool's argmax agreement with the bf16 pool at or above the bf16 pool's own run-to-run baseline.
- Serving A/B, fp8 pool against bf16 pool, shows the decode-step gain at c1/c4/c16 with first-token time not worse; unset A/B against
main is flat.
- fmt, all-target Clippy at deny-warnings, lib tests and the kernel GPU suites stay green; the
qwen3 default build (no gemma4 feature) still builds.
Description
The 26B's sliding-window family carries almost all of a decode step's KV bytes (the global family reads a fraction of a gigabyte per step against several for the window), and it stores them as bf16. The reference engine serves this checkpoint with an e4m3 KV cache at scale 1.0 by default, which is where its decode-step advantage on this box came from once the kernel accounts were corrected: the same window read at half the bytes. Our pool has no element width but bf16, so the sliding family cannot take the same trade.
Proposed Plan
KvLayout/KvPoolcarryelem_bytes(2 or 1); strides stay in elements, the backing length converts by bytes, and the geometry gate accounts for the width.csrc/gemma4/local_kv_fp8.cuinstantiates the two hd256 windowed reads (batch prefill, split-KV decode) over e4m3, and the hd256 prep kernel takes a store type so it writes e4m3 at scale 1.0. The bf16 instantiations and bindings do not move.PEGAINFER_KV_FP8=localopts the local pool in; unset keeps every byte as before. Wrappers without an fp8 twin refuse a one-byte layout at the operator boundary; the prefix-cache combination is refused at startup because its page copies index bf16 elements.Acceptance Criteria
main.mainis flat.qwen3default build (nogemma4feature) still builds.