feat(gemma4): the sliding family serves an opt in fp8 KV pool - #1012
feat(gemma4): the sliding family serves an opt in fp8 KV pool#1012FeathBow wants to merge 6 commits into
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 👍 / 👎.
| 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>
a40bef6 to
552de78
Compare
1 participant
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.