feat(gemma4): the sliding family serves an opt in fp8 KV pool

CLAUDE.md

@@ -44,6 +44,7 @@ cargo run --release --features glm52 -- --model-path models/GLM5.2
 - `PEGAINFER_TEST_MODEL_PATH` — override test model path (default: `models/Qwen3-4B`)
 - `PEGAINFER_BUILD_TIMING=1` — print per-phase build timings (nvcc, Triton AOT, etc.)
 - `PEGAINFER_NVCC_JOBS` — override parallel nvcc job count
+- `PEGAINFER_KV_FP8` — gemma4 opt-in fp8 KV: `local` stores the sliding family's K/V as e4m3 at scale 1.0 (lossy; halves the local pool; refuses an enabled prefix cache; unset = byte-identical serving)
 - `PEGAINFER_PREFIX_CACHE` — gemma4 opt-in conversation prefix cache: `K` entries of captured prompt state resume multi-turn prompts (pre-allocated page budget; unset = off, byte-identical serving)
 - `PEGAINFER_ASYNC_PREFILL` — gemma4 opt-in overlap lane: `green:NN` prefills live-batch admissions on an SM-capped stream to protect decode tails (`shared` for comparison; unset = off; bad values refuse to start)
 - `PEGAINFER_MIX_CHUNK_TOKENS` — gemma4 opt-in chunked walk: a mixed admission computes at most `N` prompt rows per step (`64 <= N <` the serving ceiling; unset = whole-prompt steps; bad values refuse to start)

docs/models/gemma4/serving.md

@@ -118,6 +118,10 @@ One prefill is in flight at most; further arrivals wait while decode keeps stepp
 
 Measured (a streaming request, then sixteen ~1900-token prompts admitted at once; two runs per arm): the stream's worst inter-token gap under the flood drops from 387-452 ms — one mixed step at that prompt length — to 75-76 ms with `green:35`, p99 385-432 → 39-40 ms, while the flood's own TTFT p50 grows 3.3-3.7 → 9.8-10.3 s and its wall about 2.4×. The quiet stream and idle footprint are unchanged, so an idle lane costs nothing. That trade is the positioning: a high-concurrency, decode-tail-sensitive profile, not a default — at light load the capped lane only costs TTFT.
 
+## The fp8 KV pool (opt-in)
+
+`PEGAINFER_KV_FP8=local` stores the sliding family's K/V as e4m3 at scale 1.0 — the scheme the reference engine defaults to for this checkpoint — halving the local pool's bytes (the global family stays bf16) and the decode step's dominant KV read; at c16 that is worth several percent of throughput, at c1 nothing. Unset serves byte-identically; `local` is the only accepted value. The output is approximate by construction: greedy generation still matches HF token for token on the fixture prompts, but the window-edge waypoint sits below the dual-backend top-1 bar and serving is no longer bit-equal across batch compositions where bf16 was. The prefix cache cannot be combined (its page copies index bf16 elements): an enabled `PEGAINFER_PREFIX_CACHE` refuses before the checkpoint loads, while a disabled one (`unset`, `0`, `off`) is fine. Operators wanting bit-exact serving leave it unset.
+
 ## The chunked walk (opt-in)
 
 `PEGAINFER_MIX_CHUNK_TOKENS=N` (64 <= N, below the serving ceiling; unset, `off` or `0` keeps whole-prompt steps; anything else refuses startup) bounds how many prompt rows a mixed admission computes per step. The effective step rounds down to whole 128-row tiles — GEMM and attention consume full tiles, so an unaligned width pays the whole tile on every full segment — which keeps "at most N rows" true while a width under one tile stays exact. Gathered prompts walk shared segment steps: each round fills one N-row budget across the walkers in admission order, every active stream advances one token per round, and a mid-walk segment's sampled row is discarded — no token, no logprob, no stop — until the prompt's final segment produces its first token, emitted at that round's boundary as the walker joins the decode batch. A walker whose client disconnects mid-walk is dropped between rounds. The knob owns every scan: a drained roster's tails and a prompt arriving with nothing active walk their own segments too, paying one ~27 ms step floor per segment where a whole scan paid one — the price of holding window plus segment instead of the full prompt. The exception is the async prefill lane: a live-batch admission goes to the lane and prefills whole. With the knob set, the gather's 512-row ceiling no longer applies: the per-round budget bounds each step instead.

pegainfer-core/src/kv_pool.rs

@@ -4,6 +4,7 @@ use anyhow::Result;
 use anyhow::bail;
 use cudarc::driver::CudaSlice;
 use half::bf16;
+pub use pegainfer_kernels::paged_kv::KvStorage;
 
 use crate::page_pool::OwnedPagePermit;
 use crate::page_pool::PageId;
@@ -25,6 +26,7 @@ pub struct KvLayout {
     pub layer_stride: usize,
     /// Elements per page (all layers): num_layers × layer_stride.
     pub page_stride: usize,
+    pub storage: KvStorage,
 }
 
 impl KvLayout {
@@ -33,6 +35,22 @@ impl KvLayout {
         num_kv_heads: usize,
         head_dim: usize,
         page_size: usize,
+    ) -> anyhow::Result<Self> {
+        Self::with_storage(
+            num_layers,
+            num_kv_heads,
+            head_dim,
+            page_size,
+            KvStorage::Bf16,
+        )
+    }
+
+    pub fn with_storage(
+        num_layers: usize,
+        num_kv_heads: usize,
+        head_dim: usize,
+        page_size: usize,
+        storage: KvStorage,
     ) -> anyhow::Result<Self> {
         let strides = || -> Option<(usize, usize, usize)> {
             let kv_block_len = page_size.checked_mul(num_kv_heads)?.checked_mul(head_dim)?;
@@ -54,6 +72,7 @@ impl KvLayout {
             kv_block_len,
             layer_stride,
             page_stride,
+            storage,
         })
     }
 
@@ -66,6 +85,7 @@ impl KvLayout {
             kv_block_len: self.kv_block_len,
             layer_stride: self.layer_stride,
             page_stride: self.page_stride,
+            storage: self.storage,
         }
     }
 }
@@ -105,7 +125,29 @@ impl KvPool {
         page_size: usize,
         num_pages: usize,
     ) -> Result<Self> {
-        let layout = KvLayout::new(num_layers, num_kv_heads, head_dim, page_size)?;
+        Self::with_storage(
+            ctx,
+            num_layers,
+            num_kv_heads,
+            head_dim,
+            page_size,
+            num_pages,
+            KvStorage::Bf16,
+        )
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn with_storage(
+        ctx: &DeviceContext,
+        num_layers: usize,
+        num_kv_heads: usize,
+        head_dim: usize,
+        page_size: usize,
+        num_pages: usize,
+        storage: KvStorage,
+    ) -> Result<Self> {
+        let layout =
+            KvLayout::with_storage(num_layers, num_kv_heads, head_dim, page_size, storage)?;
         let total_elements = num_pages.checked_mul(layout.page_stride).ok_or_else(|| {
             anyhow::anyhow!(
                 "KvPool geometry overflows: {num_pages} pages x {} elements per page",
@@ -114,17 +156,18 @@ impl KvPool {
         })?;
         // The allocator multiplies by the element size unchecked; answer
         // for the byte domain here, before it does.
-        total_elements
-            .checked_mul(std::mem::size_of::<bf16>())
+        let total_bytes = total_elements
+            .checked_mul(layout.storage.elem_bytes())
             .ok_or_else(|| {
                 anyhow::anyhow!(
-                    "KvPool geometry overflows the byte domain: {total_elements} bf16 elements"
+                    "KvPool geometry overflows the byte domain: {total_elements} elements"
                 )
             })?;
+        let backing_slots = total_bytes.div_ceil(std::mem::size_of::<bf16>());
 
         let buffer: CudaSlice<bf16> = ctx
             .stream
-            .alloc_zeros(total_elements)
+            .alloc_zeros(backing_slots)
             .map_err(|e| anyhow::anyhow!("KvPool alloc failed: {e}"))?;
 
         let pool = PagePool::new(num_pages);
@@ -359,6 +402,21 @@ mod tests {
         KvPool::new(&ctx, 1, 1, 1, page_size, num_pages).expect("KvPool::new failed")
     }
 
+    #[test]
+    fn storage_width_sets_backing_capacity() {
+        let ctx = DeviceContext::new().expect("GPU required for kv_pool tests");
+        let pages = 3;
+        let fp8 =
+            KvPool::with_storage(&ctx, 1, 1, 1, 16, pages, KvStorage::E4m3).expect("fp8 KvPool");
+        let bf16 =
+            KvPool::with_storage(&ctx, 1, 1, 1, 16, pages, KvStorage::Bf16).expect("bf16 KvPool");
+        assert_eq!(
+            fp8.buffer().len(),
+            (pages * fp8.layout().page_stride).div_ceil(2)
+        );
+        assert_eq!(bf16.buffer().len(), pages * bf16.layout().page_stride);
+    }
+
     #[test]
     fn stride_geometry_qwen35() {
         // Qwen3.5-4B: 8 full attn layers, 4 KV heads, head_dim=256, page_size=16

pegainfer-gemma4/src/engine.rs

@@ -9,6 +9,7 @@ use std::path::Path;
 use anyhow::Context as AnyhowContext;
 use anyhow::Result;
 use pegainfer_core::cuda_graph::CudaGraphState;
+use pegainfer_core::kv_pool::KvStorage;
 use pegainfer_core::ops;
 use pegainfer_core::tensor::DeviceContext;
 use pegainfer_core::tensor::HiddenStates;
@@ -44,6 +45,7 @@ const PREFIX_CACHE_ENV: &str = "PEGAINFER_PREFIX_CACHE";
 const MIX_CHUNK_TOKENS_ENV: &str = "PEGAINFER_MIX_CHUNK_TOKENS";
 const MAX_CONTEXT_ENV: &str = "PEGAINFER_MAX_CONTEXT";
 const DECODE_SLOTS_ENV: &str = "PEGAINFER_DECODE_SLOTS";
+const KV_FP8_ENV: &str = "PEGAINFER_KV_FP8";
 const MIN_CONTEXT: usize = 1024;
 const MIN_CHUNK_TOKENS: usize = 64;
 const CEILING_DOMAIN: usize = i32::MAX as usize;
@@ -154,7 +156,7 @@ fn parse_mix_chunk_tokens(raw: &str, max_context: usize) -> Result<Option<usize>
     }
 }
 
-fn prefix_cache_cap() -> Result<Option<usize>> {
+pub(crate) fn prefix_cache_cap() -> Result<Option<usize>> {
     read_env(PREFIX_CACHE_ENV)?.map_or(Ok(None), |raw| parse_prefix_cache_cap(&raw))
 }
 
@@ -169,6 +171,30 @@ fn parse_prefix_cache_cap(raw: &str) -> Result<Option<usize>> {
     }
 }
 
+pub(crate) fn kv_fp8_storage() -> Result<KvStorage> {
+    let storage = match std::env::var(KV_FP8_ENV) {
+        Err(std::env::VarError::NotPresent) => parse_kv_fp8(None),
+        Ok(raw) => parse_kv_fp8(Some(&raw)),
+        Err(err) => anyhow::bail!("PEGAINFER_KV_FP8 is not unicode: {err}"),
+    }?;
+    if storage == KvStorage::E4m3 {
+        anyhow::ensure!(
+            prefix_cache_cap()?.is_none(),
+            "PEGAINFER_KV_FP8 and PEGAINFER_PREFIX_CACHE cannot combine: the prefix cache \
+             copies pool pages in bf16 element units"
+        );
+    }
+    Ok(storage)
+}
+
+fn parse_kv_fp8(raw: Option<&str>) -> Result<KvStorage> {
+    match raw {
+        None => Ok(KvStorage::Bf16),
+        Some("local") => Ok(KvStorage::E4m3),
+        Some(value) => anyhow::bail!("PEGAINFER_KV_FP8 supports only \"local\", got {value:?}"),
+    }
+}
+
 pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result<EngineHandle> {
     let dir = model_path
         .to_str()
@@ -876,6 +902,7 @@ impl EngineState {
         let lane_mode = async_prefill_mode()?;
         let mix_chunk = mix_chunk_tokens(max_context)?;
         let slots = decode_slots()?;
+        let local_kv_storage = kv_fp8_storage()?;
         if max_context > MAX_CONTEXT {
             anyhow::ensure!(
                 mix_chunk.is_some(),
@@ -950,14 +977,21 @@ impl EngineState {
              derives page or row counts past the i32 metadata domain (the global family's pseudo \
              tables carry {global_split} copies of every page)"
         );
-        let serve = GemmaServe::new(&ctx, weights, max_context, local_pages, global_pages)
-            .map_err(|err| {
-                err.context(format!(
-                    "a {max_context} token ceiling, {slots} decode slots and {cache_entries} \
+        let serve = GemmaServe::new(
+            &ctx,
+            weights,
+            max_context,
+            local_kv_storage,
+            local_pages,
+            global_pages,
+        )
+        .map_err(|err| {
+            err.context(format!(
+                "a {max_context} token ceiling, {slots} decode slots and {cache_entries} \
                      cache entries sized the pools to {local_pages} local / {global_pages} \
                      global pages"
-                ))
-            })?;
+            ))
+        })?;
         let prefix_cache = cache_cap.map(|k| PrefixCache::new(k, sliding_window));
         let mut scratch = SampleScratch::new(&ctx, vocab, arena_rows)?;
         let mut arena = serve.alloc_step_arena(&ctx, arena_rows, graph_enabled)?;
@@ -2234,6 +2268,13 @@ mod knob_tests {
         }
     }
 
+    #[test]
+    fn fp8_knob_parses_or_refuses() {
+        assert_eq!(parse_kv_fp8(None).unwrap(), KvStorage::Bf16);
+        assert_eq!(parse_kv_fp8(Some("local")).unwrap(), KvStorage::E4m3);
+        assert!(parse_kv_fp8(Some("global")).is_err());
+    }
+
     #[test]
     fn chunk_mode_parses_or_refuses() {
         for off in ["", "0", "off", " OFF "] {
@@ -2443,12 +2484,13 @@ mod lane_tests {
         }
     }
 
-    const SERVING_KNOBS: [&str; 5] = [
+    const SERVING_KNOBS: [&str; 6] = [
         "PEGAINFER_ASYNC_PREFILL",
         "PEGAINFER_PREFIX_CACHE",
         "PEGAINFER_MIX_CHUNK_TOKENS",
         "PEGAINFER_MAX_CONTEXT",
         "PEGAINFER_DECODE_SLOTS",
+        "PEGAINFER_KV_FP8",
     ];
 
     /// Clear every serving knob, set `overrides`, and hand back the guard

pegainfer-gemma4/src/serve.rs

@@ -19,6 +19,7 @@ use cudarc::driver::CudaSlice;
 use half::bf16;
 use pegainfer_core::cuda_graph::CudaGraphState;
 use pegainfer_core::kv_pool::KvPool;
+use pegainfer_core::kv_pool::KvStorage;
 use pegainfer_core::ops;
 use pegainfer_core::ops::PrefillPagedPlan;
 use pegainfer_core::rope::RopeTableSpec;
@@ -313,6 +314,11 @@ fn copy_pool_pages(
     dst: &[i32],
 ) -> Result<()> {
     use cudarc::driver::DevicePtr;
+    // Page copies index bf16 elements, so only a bf16 pool may reach them.
+    anyhow::ensure!(
+        layout.storage == KvStorage::Bf16,
+        "pool page copies index bf16 elements; the fp8 pool must not reach them"
+    );
     anyhow::ensure!(
         src.len() == dst.len(),
         "page copy list mismatch: {} src vs {} dst",
@@ -714,6 +720,7 @@ impl GemmaServe {
         ctx: &DeviceContext,
         weights: Gemma4Weights,
         max_context: usize,
+        local_kv_storage: KvStorage,
         local_pages: usize,
         global_pages: usize,
     ) -> Result<Self> {
@@ -744,13 +751,14 @@ impl GemmaServe {
                 }
             })
             .collect();
-        let local_pool = KvPool::new(
+        let local_pool = KvPool::with_storage(
             ctx,
             locals,
             config.num_key_value_heads,
             config.head_dim,
             PAGE_SIZE,
             local_pages,
+            local_kv_storage,
         )?;
         let global_pool = KvPool::new(
             ctx,