feat(gemma4): the sliding family serves an opt in fp8 KV pool - #1012
Merged
Conversation
Signed-off-by: Feathbow <feathbow@gmail.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: fe393f7564
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| Err(std::env::VarError::NotPresent) => Ok(2), | ||
| Ok(value) if value == "local" => { | ||
| anyhow::ensure!( | ||
| std::env::var_os("PEGAINFER_PREFIX_CACHE").is_none(), |
There was a problem hiding this comment.
When FP8 is enabled in an environment that explicitly disables caching with PEGAINFER_PREFIX_CACHE=0, off, or an empty string, this raw presence check rejects startup even though parse_prefix_cache_cap in engine.rs normalizes all three spellings to None. This breaks an otherwise supported deployment configuration; base the incompatibility check on the parsed cache capacity rather than whether the variable exists.
Useful? React with 👍 / 👎.
Comment on lines
+658
to
+659
| fn local_kv_elem_bytes() -> Result<usize> { | ||
| match std::env::var("PEGAINFER_KV_FP8") { |
There was a problem hiding this comment.
This introduces a production environment variable, a prefix-cache incompatibility, changed pool memory consumption, and approximate-output behavior without updating docs/models/gemma4/serving.md or the repository's key-env-var list. Operators therefore cannot discover or safely configure the feature from the repository documentation; document the knob, supported value, incompatibilities, memory envelope, and accuracy tradeoff alongside the existing Gemma serving controls.
AGENTS.md reference: AGENTS.md:L176-L178
Useful? React with 👍 / 👎.
…riable Signed-off-by: Feathbow <feathbow@gmail.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
PagedKvLayout and KvLayout carry KvStorage::{Bf16, E4m3} instead of a
public byte count, so an illegal width is unrepresentable and the
geometry gate dispatches on the enum rather than defaulting everything
that is not exactly one byte to bf16. The three fp8 dispatch sites name
the entry point that actually ran when they fail. PEGAINFER_KV_FP8
parses once in the engine beside the other knobs, before the checkpoint
loads, and GemmaServe receives the storage as a parameter; the second
raw-env interpretation in serve.rs is gone, and the enabled-prefix-cache
refusal keeps its message.
Signed-off-by: Feathbow <feathbow@gmail.com>
…inned Signed-off-by: Feathbow <feathbow@gmail.com>
…gine's own floor Signed-off-by: Feathbow <feathbow@gmail.com>
FeathBow
force-pushed
the
feat/gemma4-fp8-kv
branch
from
August 30, 2026 20:49
a40bef6 to
552de78
Compare
Since the knob parse moved into the engine, stack_with pinned Bf16, so PEGAINFER_KV_FP8=local stopped reaching the broad checkpoint gates: the runner fp8 arm silently measured bf16. The stack now asks the engine parsed knob, prints the storage it built, and the runner banner echoes the variable when set. Signed-off-by: Feathbow <feathbow@gmail.com>
…rossing case The floor gate ran one prompt with no degeneracy guard. The body moves into agreement_case; the gate runs the window-edge prompt at stride 2 and the 4096 prompt truncated to 2048 positions at stride 8, teacher forcing well past the window, and every case demands a non-degenerate floor before comparing fp8 against it. Signed-off-by: Feathbow <feathbow@gmail.com>
The read instruments used an all-zero Q over per-page-identical values, so equal scores hid wrong K reads and wrong page picks, and only the unbounded window ran. The pools now carry distinct e4m3-representable values under a non-zero Q (identical post-load floats keep the bf16 comparison bitwise), a finite-window arm proves the mask live, the decode prep store lands its exact bytes, a varied-value geometry probe joins the constant one, and the core pool constructor byte capacity is pinned. Signed-off-by: Feathbow <feathbow@gmail.com>
The top-1 bar asserted inside gate_waypoint, so the first failing case hid every case behind it - under the fp8 pool the w1023 failure silenced the other seven. The shortfall now joins the collected failures beside the tolerance overage, both batches collect before the one verdict, and the per-case line still prints, so a failing run shows the whole table. The structural asserts (fixture positions, shifted coverage, page accounting) stay in place. Signed-off-by: Feathbow <feathbow@gmail.com>
…table The varied pool byte pattern had period 120 against a page step of 1024, so pages aliased every 15 pages and the 151-page probe target was byte-identical to page 0 - a wrong page-list base passed anyway. Each page now draws a distinct e4m3 code offset (aliasing would need 254 pages), the packed prefix table is reversed so a logical-index confusion changes what is read, and the probe asserts the target page differs from every other page before comparing. Signed-off-by: Feathbow <feathbow@gmail.com>
…suite joins the runner SERVING_KNOBS omitted PEGAINFER_KV_FP8, so an operator shell fp8 value could steer lifecycle tests into the fp8 refusals. The runner now owns the kernels integration binary the way it owns everything else: a manifest of its seven tests, a membership check against the binary own listing, and an execution arm that demands the device (PEGAINFER_REQUIRE_GPU=1) so a missing GPU fails instead of skipping. Signed-off-by: Feathbow <feathbow@gmail.com>
mixed_step_matches_serial demands bit-exact equality between admission schedules and prefix_restore_matches_cold_path exercises page copies that are bf16-only, but both built their stacks through the env-honoring constructor, so an fp8 arm turned contract statements into permanent failures. Both now pin bf16 and name where the fp8 pool is judged instead; the same-pool invariance gates stay storage-honoring. Signed-off-by: Feathbow <feathbow@gmail.com>
Signed-off-by: Feathbow <feathbow@gmail.com>
A set-but-empty PEGAINFER_KV_FP8 slipped past the ambient refusal while the production parser refuses it, and a set-but-empty PEGAINFER_GATE_STORAGE silently became bf16 although only unset, bf16 and fp8 are accepted. Presence now checks +x and the default applies only when the variable is truly unset, so an empty value dies. Signed-off-by: Feathbow <feathbow@gmail.com>
…eterminism The determinism argument behind the cross-shape floor rested on a serving-probe observation. The gate now replays the bf16 incremental arm on its own prompt and schedule and asserts the replay agrees at every position, so the run-to-run baseline is measured where the floor is used: deterministic, hence degenerate for judging a lossy storage. Signed-off-by: Feathbow <feathbow@gmail.com>
… invariants Pinning the bit-exact gates to bf16 left mixed_prefill_decode_step unexercised under e4m3. A new gate builds the e4m3 stack explicitly and drives the same mixed machinery through the plain and the window-crossing walks, asserting structure instead of bit equality: finite mixed-step logits, every lane reaching exactly its budget, and the window-family page accounting holding right after the mixed step that releases the front. Registered in the admission set and the fp8 profile; no tolerance constants. Signed-off-by: Feathbow <feathbow@gmail.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Closes #1010
PEGAINFER_KV_FP8=localstores the sliding family's K/V as e4m3 at scale 1.0, halving the decode step's dominant KV read; unset serves byte-identically to before. The reference engine defaults to exactly this scheme for the checkpoint, so the two engines finally read the same bytes per step.The pool layout carries an element width, the geometry gate converts the backing length by it, and the shared launch header takes the KV element type as a template parameter: a gemma4 translation unit instantiates the two hd256 windowed reads over e4m3 while the bf16 instantiations and bindings stay where they are, and the hd256 prep kernel takes a store type the same way. Wrappers without an fp8 twin refuse a one-byte layout at the operator boundary, and the prefix-cache combination is refused at startup because its page copies index bf16 elements.
Test Env
Single GPU (
sm_89, 48 GiB, x86_64).Verification
Results from the box run on the carve; the accuracy legs decide whether this opens.
cargo fmt --all -- --checkandgit diff --checkagainstmainpass; the Gemma 4 release server builds; theqwen3default build (nogemma4feature) builds; all-target kernels/Gemma/core Clippy passes with-D warnings; release lib tests pass; the kernel GPU suites pass 21/21./v1/completionsrequests (31-, 35-, 33- and 1421-token prompts, 64 forced tokens) with bytes identical tomain's binary, md5219de1e3c09c.PEGAINFER_KV_FP8=localwithPEGAINFER_PREFIX_CACHEset exits with "the prefix cache copies pool pages in bf16 element units";PEGAINFER_KV_FP8=globalexits with "supports only "local"".PEGAINFER_KV_FP8=local: the greedy HF generation gate passes token for token; the ragged-batch row-order gate and the eviction gate pass; the 16K waypoint gate fails at the window-edge waypoint (w1023: top-1 5/9 below backend bar 6/9; every other waypoint within tolerance); the mixed-step-against-serial gate, bit-exact under bf16, diverges under fp8 on the window-crossing partner; the shared-lane lifecycle gate configures the prefix cache and is refused at startup by design.vllm bench serverandom 1024 in / 128 out at 8 / 24 / 64 prompts, two cooled alternating rounds, output tok/s as mean [range]: c1 106.3 [106.0–106.5] → 105.1 [105.0–105.2]; c4 279.8 [279.1–280.4] → 285.0 [284.2–285.7]; c16 595.0 [591.4–598.5] → 629.1 [627.5–630.6]; median TPOT at c16 23.92 → 22.10 ms.