diff --git a/CLAUDE.md b/CLAUDE.md index c6494c566..8a6209655 100644 --- a/CLAUDE.md +++ b/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) diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index e9f7ed397..f89f5b2a1 100644 --- a/docs/models/gemma4/serving.md +++ b/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. diff --git a/pegainfer-core/src/kv_pool.rs b/pegainfer-core/src/kv_pool.rs index 543ced8c7..6ba1fac84 100644 --- a/pegainfer-core/src/kv_pool.rs +++ b/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::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 { 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 { - 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 { + 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::()) + 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::()); let buffer: CudaSlice = 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 diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index d4aabb0fa..96cfe478b 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/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 } } -fn prefix_cache_cap() -> Result> { +pub(crate) fn prefix_cache_cap() -> Result> { 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> { } } +pub(crate) fn kv_fp8_storage() -> Result { + 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 { + 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 { 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 diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 8b1818c96..d6c2e3d69 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/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 { @@ -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, diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index fdb616e22..ac6eea6ce 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -3,6 +3,7 @@ //! different admission shape. use anyhow::Result; +use pegainfer_core::kv_pool::KvStorage; use super::*; use crate::kv::admit_tokens; @@ -15,12 +16,25 @@ use crate::testkit::model_path; use crate::testkit::u32_tensor; fn stack_with(max_context: usize, pages: usize) -> (DeviceContext, GemmaServe, String) { + stack_with_storage( + max_context, + pages, + crate::engine::kv_fp8_storage().expect("PEGAINFER_KV_FP8"), + ) +} + +fn stack_with_storage( + max_context: usize, + pages: usize, + storage: KvStorage, +) -> (DeviceContext, GemmaServe, String) { let dir = model_path(); let config = Gemma4Config::from_file(&dir).expect("config"); let (weights, _) = Gemma4Weights::from_safetensors(&dir, 0, config).expect("load checkpoint weights"); let ctx = DeviceContext::new_with_device(0).expect("device context"); - let serve = GemmaServe::new(&ctx, weights, max_context, pages, pages).expect("serve"); + let serve = GemmaServe::new(&ctx, weights, max_context, storage, pages, pages).expect("serve"); + eprintln!("oracle stack storage: {storage:?}"); (ctx, serve, dir) } @@ -223,7 +237,7 @@ fn gate_waypoint( serve: &GemmaServe, fixture: &safetensors::SafeTensors<'_>, point: Waypoint<'_>, -) -> Option { +) -> Vec { let label = match point.chunk { 0 => point.case.to_string(), _ => format!("{}-chunked", point.case), @@ -237,10 +251,6 @@ fn gate_waypoint( "{label}: shifted multi-token coverage" ); let (max_abs, top1) = score_rows(&run.rows, &ids, &lps, top_k, &label); - assert!( - top1 >= backend_top1, - "{label}: top-1 {top1}/{positions} below backend bar {backend_top1}/{positions}" - ); let page = serve.local_pool.layout().page_size; let released = run.kv_len.saturating_sub(serve.sliding_window) / page; assert_eq!(run.local_pages, run.kv_len.div_ceil(page) - released); @@ -250,7 +260,16 @@ fn gate_waypoint( {top1}/{positions}, local pages {}, global {}", run.local_pages, run.global_pages ); - (max_abs > tolerance).then(|| format!("{label} ({max_abs} > {tolerance})")) + let mut failures = Vec::new(); + if top1 < backend_top1 { + failures.push(format!( + "{label}: top-1 {top1}/{positions} below backend bar {backend_top1}/{positions}" + )); + } + if max_abs > tolerance { + failures.push(format!("{label} ({max_abs} > {tolerance})")); + } + failures } fn gate_waypoints( @@ -258,15 +277,11 @@ fn gate_waypoints( serve: &GemmaServe, fixture: &safetensors::SafeTensors<'_>, points: &[Waypoint<'_>], -) { - let over: Vec = points +) -> Vec { + points .iter() - .filter_map(|&point| gate_waypoint(ctx, serve, fixture, point)) - .collect(); - assert!( - over.is_empty(), - "cases over their calibrated floor: {over:?}" - ); + .flat_map(|&point| gate_waypoint(ctx, serve, fixture, point)) + .collect() } fn validate_waypoint_provenance(dir: &str, window_bytes: &[u8], long_bytes: &[u8]) { @@ -350,8 +365,128 @@ fn context_waypoints_match_hf() { floor: Some(floor), }, ]; - gate_waypoints(&ctx, &serve, &window, &window_points); - gate_waypoints(&ctx, &serve, &long, &long_points); + let mut failures = gate_waypoints(&ctx, &serve, &window, &window_points); + failures.extend(gate_waypoints(&ctx, &serve, &long, &long_points)); + assert!( + failures.is_empty(), + "cases over their calibrated floor: {failures:?}" + ); +} + +fn incremental_argmaxes(ctx: &DeviceContext, serve: &GemmaServe, prompt: &[u32]) -> Vec { + let mut kv = serve.alloc_kv(); + let mut arena = serve + .alloc_step_arena(ctx, 1, false) + .expect("oracle step arena"); + admit_tokens(&serve.local_pool, &serve.global_pool, &mut kv, 1).expect("admit first token"); + let first = serve.step(ctx, &mut kv, &prompt[..1]).expect("first step"); + let mut choices = vec![argmax(&first.to_host(ctx).expect("first logits D2H"))]; + for &token in &prompt[1..] { + let row = decode_serving(serve, ctx, &mut arena, &mut kv, token).expect("decode"); + choices.push(argmax(&row)); + } + choices +} + +fn recomputed_argmaxes( + ctx: &DeviceContext, + serve: &GemmaServe, + prompt: &[u32], + positions: &[usize], +) -> Vec { + positions + .iter() + .map(|&position| { + let mut kv = serve.alloc_kv(); + let prefix = &prompt[..=position]; + admit_tokens(&serve.local_pool, &serve.global_pool, &mut kv, prefix.len()) + .expect("admit recompute prefix"); + let logits = serve.step(ctx, &mut kv, prefix).expect("recompute step"); + let host = logits.to_host(ctx).expect("recompute logits D2H"); + let vocab = logits.hidden_dim; + argmax(&host[(logits.seq_len - 1) * vocab..]) + }) + .collect() +} + +fn agreement(left: &[usize], right: &[usize]) -> usize { + left.iter().zip(right).filter(|(a, b)| a == b).count() +} + +fn agreement_case( + fixture: &safetensors::SafeTensors<'_>, + tensor_name: &str, + truncate_to: usize, + stride: usize, +) -> (usize, usize, usize) { + let (_, mut prompt) = u32_tensor(fixture, tensor_name); + prompt.truncate(truncate_to); + let sampled: Vec = (0..prompt.len()).step_by(stride).collect(); + let max_context = prompt.len().div_ceil(crate::kv::PAGE_SIZE) * crate::kv::PAGE_SIZE; + // One request at the case depth, plus each pool's padding page. + let pages = max_context.div_ceil(crate::kv::PAGE_SIZE) + 2; + + let (ctx, bf16, _) = stack_with_storage(max_context, pages, KvStorage::Bf16); + let bf16_incremental = incremental_argmaxes(&ctx, &bf16, &prompt); + // The same-schedule run-to-run baseline, measured on this exact prompt + // and schedule: a replay must agree everywhere, which is why a lossy + // storage is judged against the cross-shape floor below instead. + let bf16_replay = incremental_argmaxes(&ctx, &bf16, &prompt); + let replay_matches = agreement(&bf16_incremental, &bf16_replay); + eprintln!( + "{tensor_name}: same-schedule bf16 run-to-run agreement {replay_matches}/{}", + bf16_incremental.len() + ); + assert_eq!( + replay_matches, + bf16_incremental.len(), + "{tensor_name}: the same-schedule bf16 replay must be deterministic" + ); + let bf16_recomputed = recomputed_argmaxes(&ctx, &bf16, &prompt, &sampled); + let bf16_sampled: Vec = sampled.iter().map(|&pos| bf16_incremental[pos]).collect(); + let floor_matches = agreement(&bf16_sampled, &bf16_recomputed); + drop(bf16); + drop(ctx); + + let (ctx, fp8, _) = stack_with_storage(max_context, pages, KvStorage::E4m3); + let fp8_incremental = incremental_argmaxes(&ctx, &fp8, &prompt); + let fp8_sampled: Vec = sampled.iter().map(|&pos| fp8_incremental[pos]).collect(); + let fp8_matches = agreement(&fp8_sampled, &bf16_sampled); + let samples = sampled.len(); + let samples_f64 = f64::from(u32::try_from(samples).expect("sample count fits u32")); + let floor_rate = + f64::from(u32::try_from(floor_matches).expect("match count fits u32")) / samples_f64; + let fp8_rate = + f64::from(u32::try_from(fp8_matches).expect("match count fits u32")) / samples_f64; + eprintln!( + "{tensor_name}: argmax agreement: bf16 incremental/recompute {floor_rate:.6} \ + ({floor_matches}/{samples}), fp8/bf16 incremental {fp8_rate:.6} \ + ({fp8_matches}/{samples})" + ); + (floor_matches, fp8_matches, samples) +} + +#[test] +#[ignore = "requires the pinned 12B checkpoint, fixtures, and a GPU"] +fn fp8_argmax_agreement_meets_the_bf16_floor() { + let bytes = std::fs::read(WINDOW_FIXTURE).expect("read window fixture"); + let fixture = safetensors::SafeTensors::deserialize(&bytes).expect("window fixture"); + // Recompute is quadratic; these strides keep both window depths sampled. + // Every case measures before any verdict, so a failing first case cannot + // hide the second case's numbers. + let results = [("w1023_prompt", usize::MAX, 2), ("w4096_prompt", 2048, 8)] + .map(|(name, cut, stride)| (name, agreement_case(&fixture, name, cut, stride))); + for (name, (floor, fp8, samples)) in results { + assert!( + floor * 2 > samples, + "{name}: degenerate bf16 incremental/recompute floor {floor}/{samples}" + ); + assert!( + fp8 >= floor, + "{name}: fp8/bf16 argmax agreement {fp8}/{samples} is below the bf16 \ + incremental/recompute floor {floor}/{samples}" + ); + } } /// `window_left` masks out-of-window keys whether or not their pages are @@ -452,8 +587,10 @@ fn greedy_matches_hf_generate() { ); } -fn mixed_gate_argmax(host: &[f32], row: usize, vocab: usize) -> u32 { - u32::try_from(argmax(&host[row * vocab..(row + 1) * vocab])).expect("token id") +/// The production sampler draws from the true vocabulary; a padded lm_head +/// column is never a candidate, so neither is it here. +fn mixed_gate_argmax(host: &[f32], row: usize, stride: usize, bound: usize) -> u32 { + u32::try_from(argmax(&host[row * stride..row * stride + bound])).expect("token id") } /// Reserve a lane's whole prompt. These gates drive the serving primitives @@ -494,7 +631,8 @@ fn gate_host_logits(ctx: &DeviceContext, logits: &HiddenStates) -> (usize, Vec, produced: &mut [Vec], @@ -502,7 +640,7 @@ fn settle_gate_lanes( ) { let mut retire: Vec = Vec::new(); for (row, (req, _, next)) in lanes.iter_mut().enumerate() { - let token = mixed_gate_argmax(host, row + row_base, vocab); + let token = mixed_gate_argmax(host, row + row_base, stride, bound); produced[*req].push(token); if produced[*req].len() >= budgets[*req] { retire.push(row); @@ -536,7 +674,15 @@ fn mixed_gate_decode_rounds( .expect("batched decode"); gate_host_logits(ctx, logits) }; - settle_gate_lanes(&host, vocab, 0, lanes, produced, budgets); + settle_gate_lanes( + &host, + vocab, + serve.weights.config.vocab_size, + 0, + lanes, + produced, + budgets, + ); } } @@ -596,11 +742,19 @@ fn assert_mixed_admissions_match_serial(ctx: &DeviceContext, serve: &GemmaServe) .expect("k=2 mixed step"); gate_host_logits(ctx, logits) }; - settle_gate_lanes(&host, vocab, 2, &mut lanes, &mut produced, &budgets); - let first_b = mixed_gate_argmax(&host, 0, vocab); + settle_gate_lanes( + &host, + vocab, + serve.weights.config.vocab_size, + 2, + &mut lanes, + &mut produced, + &budgets, + ); + let first_b = mixed_gate_argmax(&host, 0, vocab, serve.weights.config.vocab_size); produced[1].push(first_b); lanes.push((1, kv_b, first_b)); - let first_c = mixed_gate_argmax(&host, 1, vocab); + let first_c = mixed_gate_argmax(&host, 1, vocab, serve.weights.config.vocab_size); produced[2].push(first_c); lanes.push((2, kv_c, first_c)); mixed_gate_decode_rounds( @@ -681,8 +835,16 @@ fn assert_mixed_window_crossing_matches_serial(ctx: &DeviceContext, serve: &Gemm "the mixed prefill must have released its window front (origin {})", kv.local.origin_pages() ); - settle_gate_lanes(&host, vocab, 1, &mut lanes, &mut produced, &budgets); - mixed_gate_argmax(&host, 0, vocab) + settle_gate_lanes( + &host, + vocab, + serve.weights.config.vocab_size, + 1, + &mut lanes, + &mut produced, + &budgets, + ); + mixed_gate_argmax(&host, 0, vocab, serve.weights.config.vocab_size) } else { let logits = serve .step(ctx, &mut kv, &long_prompt) @@ -731,11 +893,174 @@ fn assert_mixed_window_crossing_matches_serial(ctx: &DeviceContext, serve: &Gemm #[test] #[ignore = "requires the pinned 12B checkpoint and a GPU"] fn mixed_step_matches_serial() { - let (ctx, serve, _dir) = stack_with(2048, 512); + // This is a bf16 bit-exactness contract; distribution and waypoint gates judge fp8. + let (ctx, serve, _dir) = stack_with_storage(2048, 512, KvStorage::Bf16); assert_mixed_admissions_match_serial(&ctx, &serve); assert_mixed_window_crossing_matches_serial(&ctx, &serve); } +fn assert_finite_gate_logits(host: &[f32], what: &str) { + assert!( + host.iter().all(|value| value.is_finite()), + "{what}: mixed step produced non-finite logits" + ); +} + +fn assert_gate_page_accounting(serve: &GemmaServe, kv: &GemmaKv, what: &str) { + let page = serve.local_pool.layout().page_size; + let kv_len = kv.local.seq_len(); + assert_eq!(kv.global.seq_len(), kv_len, "{what}: KV lengths"); + let released = kv_len.saturating_sub(serve.sliding_window) / page; + assert_eq!( + kv.local.held_pages(), + kv_len.div_ceil(page) - released, + "{what}: local pages" + ); + assert_eq!( + kv.global.held_pages(), + kv_len.div_ceil(page), + "{what}: global pages" + ); +} + +fn fp8_plain_mixed_walk(ctx: &DeviceContext, serve: &GemmaServe) { + let prompts = crate::testkit::generate_fixture_prompts(); + let budgets = [50usize, 37, 44]; + let mut arena = serve.alloc_step_arena(ctx, 4, false).expect("step arena"); + let mut lanes = Vec::new(); + let mut produced = vec![Vec::new(); prompts.len()]; + + let (kv_a, first_a) = gate_open_lane(ctx, serve, &prompts[0], "prompt a"); + produced[0].push(first_a); + lanes.push((0, kv_a, first_a)); + mixed_gate_decode_rounds( + ctx, + serve, + &mut arena, + &mut lanes, + &mut produced, + &budgets, + 3, + ); + + let mut kv_b = gate_admit_kv(serve, &prompts[1], "prompt b"); + let mut kv_c = gate_admit_kv(serve, &prompts[2], "prompt c"); + let tokens = gate_step_tokens(serve, &mut lanes); + let (vocab, host) = { + let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); + let mut prefills = [ + (&mut kv_b, prompts[1].as_slice()), + (&mut kv_c, prompts[2].as_slice()), + ]; + let logits = serve + .mixed_prefill_decode_step(ctx, &mut arena, &mut prefills, &mut kvs, &tokens) + .expect("k=2 mixed step"); + gate_host_logits(ctx, logits) + }; + assert_finite_gate_logits(&host, "plain fp8 walk"); + settle_gate_lanes( + &host, + vocab, + serve.weights.config.vocab_size, + 2, + &mut lanes, + &mut produced, + &budgets, + ); + for (req, kv, row) in [(1, kv_b, 0), (2, kv_c, 1)] { + let first = mixed_gate_argmax(&host, row, vocab, serve.weights.config.vocab_size); + produced[req].push(first); + lanes.push((req, kv, first)); + } + mixed_gate_decode_rounds( + ctx, + serve, + &mut arena, + &mut lanes, + &mut produced, + &budgets, + 3, + ); + mixed_gate_decode_rounds( + ctx, + serve, + &mut arena, + &mut lanes, + &mut produced, + &budgets, + usize::MAX, + ); + for (tokens, budget) in produced.iter().zip(budgets) { + assert_eq!(tokens.len(), budget, "fp8 mixed lane token budget"); + } +} + +fn fp8_window_mixed_walk(ctx: &DeviceContext, serve: &GemmaServe) { + let partner: Vec = (0..40u32).map(|i| 1000 + i * 31).collect(); + let long_prompt: Vec = (0..1500u32).map(|i| 1000 + (i * 37) % 50000).collect(); + let budgets = [24usize, 20]; + let mut arena = serve.alloc_step_arena(ctx, 2, false).expect("step arena"); + let mut produced = vec![Vec::new(); 2]; + let (kv_partner, first_partner) = gate_open_lane(ctx, serve, &partner, "partner"); + let mut lanes = vec![(0, kv_partner, first_partner)]; + produced[0].push(first_partner); + mixed_gate_decode_rounds( + ctx, + serve, + &mut arena, + &mut lanes, + &mut produced, + &budgets, + 3, + ); + + let mut kv_long = gate_admit_kv(serve, &long_prompt, "long prompt"); + let tokens = gate_step_tokens(serve, &mut lanes); + let (vocab, host) = { + let mut kvs: Vec<&mut GemmaKv> = lanes.iter_mut().map(|(_, kv, _)| kv).collect(); + let mut prefills = [(&mut kv_long, long_prompt.as_slice())]; + let logits = serve + .mixed_prefill_decode_step(ctx, &mut arena, &mut prefills, &mut kvs, &tokens) + .expect("window-crossing mixed step"); + gate_host_logits(ctx, logits) + }; + assert_finite_gate_logits(&host, "window-crossing fp8 walk"); + assert_gate_page_accounting(serve, &kv_long, "long prompt mixed step"); + settle_gate_lanes( + &host, + vocab, + serve.weights.config.vocab_size, + 1, + &mut lanes, + &mut produced, + &budgets, + ); + let first_long = mixed_gate_argmax(&host, 0, vocab, serve.weights.config.vocab_size); + produced[1].push(first_long); + lanes.push((1, kv_long, first_long)); + + for (req, mut kv, mut next) in lanes { + while produced[req].len() < budgets[req] { + let host = decode_serving(serve, ctx, &mut arena, &mut kv, next).expect("decode"); + let bound = serve.weights.config.vocab_size.min(host.len()); + next = u32::try_from(argmax(&host[..bound])).expect("token id"); + produced[req].push(next); + } + assert_gate_page_accounting(serve, &kv, ["partner", "long prompt"][req]); + } + for (tokens, budget) in produced.iter().zip(budgets) { + assert_eq!(tokens.len(), budget, "fp8 window lane token budget"); + } +} + +#[test] +#[ignore = "requires the pinned 12B checkpoint, generate prompts, and a GPU"] +fn fp8_mixed_walk_holds_its_structure() { + let (ctx, serve, _dir) = stack_with_storage(2048, 512, KvStorage::E4m3); + fp8_plain_mixed_walk(&ctx, &serve); + fp8_window_mixed_walk(&ctx, &serve); +} + /// The overlap-safe prefill under a lane-stream override must be bit-equal /// to the sync step: identical row-0 logits, the same released-window /// shape after the deferred release, and greedy decode over the @@ -850,7 +1175,8 @@ fn overlapped_prefill_matches_the_sync_step() { #[ignore = "requires the pinned 12B checkpoint and a GPU"] fn prefix_restore_matches_cold_path() { use crate::prefix_cache::PrefixCache; - let (ctx, serve, _dir) = stack_with(4096, 512); + // The cache's page copies are bf16-only; distribution and waypoint gates judge fp8. + let (ctx, serve, _dir) = stack_with_storage(4096, 512, KvStorage::Bf16); let window = serve.weights.config.sliding_window; let mut arena = serve.alloc_step_arena(&ctx, 1, false).expect("step arena"); let budget = 16usize; diff --git a/pegainfer-kernels/build.rs b/pegainfer-kernels/build.rs index 949be5c06..78ebc5420 100644 --- a/pegainfer-kernels/build.rs +++ b/pegainfer-kernels/build.rs @@ -2171,6 +2171,7 @@ fn main() { || stem == "flashinfer_sampling" || stem == "flashinfer_top1" || stem == "glm52_topk" + || stem == "local_kv_fp8" { for dir in &flashinfer.cccl { nvcc_args.extend(["-I".to_string(), dir.to_string_lossy().to_string()]); diff --git a/pegainfer-kernels/csrc/gemma4/local_kv_fp8.cu b/pegainfer-kernels/csrc/gemma4/local_kv_fp8.cu new file mode 100644 index 000000000..703de76bb --- /dev/null +++ b/pegainfer-kernels/csrc/gemma4/local_kv_fp8.cu @@ -0,0 +1,27 @@ +#include "../shared/paged_launch.cuh" + +#include + +extern "C" { + +int gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda( + void* q, void* output, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + int32_t* page_indices, int32_t* page_indptr, int32_t* last_page_len_d, + int32_t* q_indptr, int32_t* request_indices, int32_t* qo_tile_indices, + int32_t* kv_tile_indices, int32_t* kv_chunk_size_ptr, uint32_t* total_num_rows, + int32_t num_qo_heads, int32_t num_kv_heads, int32_t head_dim, + int32_t page_size, int32_t seq_len, int32_t batch_size, + int32_t padded_batch_size, int64_t stride_page, float sm_scale, + int32_t cta_tile_q_override, int32_t window_left, void* stream) +{ + return prefill_paged_launch<256, WindowVariant, __nv_fp8_e4m3>( + q, output, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indptr, last_page_len_d, q_indptr, + request_indices, qo_tile_indices, kv_tile_indices, kv_chunk_size_ptr, + total_num_rows, num_qo_heads, num_kv_heads, head_dim, page_size, + seq_len, batch_size, padded_batch_size, stride_page, sm_scale, + cta_tile_q_override, window_left, stream); +} + +} // extern "C" diff --git a/pegainfer-kernels/csrc/shared/paged_launch.cuh b/pegainfer-kernels/csrc/shared/paged_launch.cuh index 7371d95af..8701e6fec 100644 --- a/pegainfer-kernels/csrc/shared/paged_launch.cuh +++ b/pegainfer-kernels/csrc/shared/paged_launch.cuh @@ -17,8 +17,6 @@ using namespace flashinfer; using DType = __nv_bfloat16; using IdType = int32_t; -using ParamsT = BatchDecodeParams; -using BatchPrefillParamsT = BatchPrefillPagedParams; using PrefillParamsT = SinglePrefillParams; using Variant = DefaultAttention; -static paged_kv_t make_paged_kv( +template +static paged_kv_t make_paged_kv( void* kv_data, int64_t k_offset_elems, int64_t v_offset_elems, @@ -45,8 +44,8 @@ static paged_kv_t make_paged_kv( int32_t batch_size, int64_t stride_page) { - DType* k_data = reinterpret_cast(kv_data) + k_offset_elems; - DType* v_data = reinterpret_cast(kv_data) + v_offset_elems; + KvT* k_data = reinterpret_cast(kv_data) + k_offset_elems; + KvT* v_data = reinterpret_cast(kv_data) + v_offset_elems; // kv_strides[0] = stride_page, [1] = stride for NHD-n, [2] = stride for NHD-h int64_t kv_strides[3] = { @@ -55,7 +54,7 @@ static paged_kv_t make_paged_kv( static_cast(head_dim), }; - return paged_kv_t( + return paged_kv_t( num_kv_heads, page_size, head_dim, batch_size, QKVLayout::kNHD, k_data, v_data, kv_strides, @@ -63,7 +62,7 @@ static paged_kv_t make_paged_kv( /*rope_pos_offset=*/nullptr); } -template +template static int decode_launch( // Q and output void* q, // [num_qo_heads * head_dim] bf16, device @@ -93,7 +92,8 @@ static int decode_launch( void* stream) { PEGAINFER_FFI_GUARD_BEGIN - auto paged_kv = make_paged_kv( + using ParamsT = BatchDecodeParams; + auto paged_kv = make_paged_kv( kv_data, k_offset_elems, v_offset_elems, page_indices, page_indptr, last_page_len_d, num_kv_heads, head_dim, page_size, batch_size, stride_page); @@ -148,7 +148,7 @@ static int decode_launch( // // tmp_v/tmp_s hold per-chunk partial states and are merged by FlashInfer. // --------------------------------------------------------------------------- -template +template static int decode_split_kv_launch( void* q, // [batch_size * num_qo_heads * head_dim] bf16, device void* output, // [batch_size * num_qo_heads * head_dim] bf16, device @@ -177,7 +177,8 @@ static int decode_split_kv_launch( void* stream) { PEGAINFER_FFI_GUARD_BEGIN - auto paged_kv = make_paged_kv( + using ParamsT = BatchDecodeParams; + auto paged_kv = make_paged_kv( kv_data, k_offset_elems, v_offset_elems, page_indices, page_indptr, last_page_len_d, num_kv_heads, head_dim, page_size, batch_size, stride_page); @@ -235,7 +236,7 @@ static uint32_t resolve_prefill_cta_tile_q( return 0; } -template +template static int prefill_paged_launch( // Q and output (HiddenStates col-major: [q_dim, total_seq_len]) void* q, @@ -271,7 +272,8 @@ static int prefill_paged_launch( void* stream) { PEGAINFER_FFI_GUARD_BEGIN - auto paged_kv = make_paged_kv( + using BatchPrefillParamsT = BatchPrefillPagedParams; + auto paged_kv = make_paged_kv( kv_data, k_offset_elems, v_offset_elems, page_indices, page_indptr, last_page_len_d, num_kv_heads, head_dim, page_size, batch_size, stride_page); diff --git a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu b/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu index b9965829b..70b1caa76 100644 --- a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu +++ b/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu @@ -20,6 +20,7 @@ #include "common.cuh" #include "ffi_guard.cuh" #include "qk_prep.cuh" +#include #define HD256_PLAIN 256 #define THREADS_HD256_PLAIN 256 @@ -129,11 +130,22 @@ __global__ void qk_norm_rope_prefill_hd256_plain_kernel( // and rotated; V is weightless-normed over its own head vector (v_proj // output — a separate reduction, unlike the hd512 K=V fork) and never // rotated. K and V write straight into the pool's per-layer K/V blocks. +template +__device__ __forceinline__ KvT kv_store_cast(__nv_bfloat16 x); +template <> +__device__ __forceinline__ __nv_bfloat16 kv_store_cast(__nv_bfloat16 x) { + return x; +} +template <> +__device__ __forceinline__ __nv_fp8_e4m3 kv_store_cast(__nv_bfloat16 x) { + return __nv_fp8_e4m3(__bfloat162float(x)); +} + // PER_TOKEN_META = true is the batched-decode form: token t is its own // request, so its absolute position, its page-table window (page_indices + // page_indptr[t]) and its released-front origin ride per-token arrays and // the scalar start_pos/page_origin are ignored. -template +template __global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( const __nv_bfloat16* __restrict__ q_batch, // [q_dim, seq_len] const __nv_bfloat16* __restrict__ k_batch, // [kv_dim, seq_len] @@ -143,7 +155,7 @@ __global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( const __nv_bfloat16* __restrict__ cos_cache, // [max_seq * rotary_dim] const __nv_bfloat16* __restrict__ sin_cache, __nv_bfloat16* __restrict__ q_batch_out, // [q_dim, seq_len] - __nv_bfloat16* __restrict__ kv_data, // paged KV pool + KvT* __restrict__ kv_data, // paged KV pool int64_t k_offset_elems, int64_t v_offset_elems, const int* __restrict__ page_indices, // resident page row(s) @@ -228,7 +240,8 @@ __global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( int64_t dst = paged_kv_offset( page_id, v_offset_elems, stride_page, page_size, num_kv_heads, pos, head_local, d); - kv_data[dst] = __float2bfloat16(__bfloat162float(x) * inv_rms); + kv_data[dst] = kv_store_cast( + __float2bfloat16(__bfloat162float(x) * inv_rms)); return; } @@ -256,8 +269,8 @@ __global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( int64_t dst = paged_kv_offset( page_id, k_offset_elems, stride_page, page_size, num_kv_heads, pos, head_local, d); - kv_data[dst] = lo; - kv_data[dst + half_rotary] = hi; + kv_data[dst] = kv_store_cast(lo); + kv_data[dst + half_rotary] = kv_store_cast(hi); } } @@ -269,7 +282,7 @@ __global__ void qkv_norm_rope_paged_prefill_hd256_plain_kernel( int64_t dst = paged_kv_offset( page_id, k_offset_elems, stride_page, page_size, num_kv_heads, pos, head_local, d); - kv_data[dst] = smem[d]; + kv_data[dst] = kv_store_cast(smem[d]); } } } @@ -346,7 +359,10 @@ int qk_norm_rope_prefill_hd256_plain_cuda( PEGAINFER_FFI_GUARD_END(-1) } -int qkv_norm_rope_paged_prefill_hd256_plain_cuda( +} // extern "C" + +template +static int qkv_prep_paged_prefill_launch( const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, const __nv_bfloat16* v_batch, @@ -355,7 +371,7 @@ int qkv_norm_rope_paged_prefill_hd256_plain_cuda( const __nv_bfloat16* cos_cache, const __nv_bfloat16* sin_cache, __nv_bfloat16* q_batch_out, - __nv_bfloat16* kv_data, + void* kv_data, int64_t k_offset_elems, int64_t v_offset_elems, const int* page_indices, @@ -409,7 +425,7 @@ int qkv_norm_rope_paged_prefill_hd256_plain_cuda( return -1; } dim3 prep_grid(seq_len, num_q_heads + 2 * num_kv_heads); - qkv_norm_rope_paged_prefill_hd256_plain_kernel + qkv_norm_rope_paged_prefill_hd256_plain_kernel <<>>( q_batch, k_batch, @@ -419,7 +435,7 @@ int qkv_norm_rope_paged_prefill_hd256_plain_cuda( cos_cache, sin_cache, q_batch_out, - kv_data, + reinterpret_cast(kv_data), k_offset_elems, v_offset_elems, page_indices, @@ -447,7 +463,8 @@ int qkv_norm_rope_paged_prefill_hd256_plain_cuda( PEGAINFER_FFI_GUARD_END(-1) } -int qkv_norm_rope_paged_decode_hd256_plain_cuda( +template +static int qkv_prep_paged_decode_launch( const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, const __nv_bfloat16* v_batch, @@ -456,7 +473,7 @@ int qkv_norm_rope_paged_decode_hd256_plain_cuda( const __nv_bfloat16* cos_cache, const __nv_bfloat16* sin_cache, __nv_bfloat16* q_batch_out, - __nv_bfloat16* kv_data, + void* kv_data, int64_t k_offset_elems, int64_t v_offset_elems, const int* page_indices, @@ -501,7 +518,7 @@ int qkv_norm_rope_paged_decode_hd256_plain_cuda( return -1; } dim3 prep_grid(batch, num_q_heads + 2 * num_kv_heads); - qkv_norm_rope_paged_prefill_hd256_plain_kernel + qkv_norm_rope_paged_prefill_hd256_plain_kernel <<>>( q_batch, k_batch, @@ -511,7 +528,7 @@ int qkv_norm_rope_paged_decode_hd256_plain_cuda( cos_cache, sin_cache, q_batch_out, - kv_data, + reinterpret_cast(kv_data), k_offset_elems, v_offset_elems, page_indices, @@ -539,4 +556,98 @@ int qkv_norm_rope_paged_decode_hd256_plain_cuda( PEGAINFER_FFI_GUARD_END(-1) } +extern "C" { + +int qkv_norm_rope_paged_prefill_hd256_plain_cuda( + const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, + const __nv_bfloat16* v_batch, + const __nv_bfloat16* q_norm_weight, const __nv_bfloat16* k_norm_weight, + const __nv_bfloat16* cos_cache, const __nv_bfloat16* sin_cache, + __nv_bfloat16* q_batch_out, __nv_bfloat16* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + const int* page_indices, int page_indices_len, int page_origin, + int num_q_heads, int num_kv_heads, int seq_len, int start_pos, + int cos_max_pos, int rotary_dim, float rms_eps, + int page_size, int num_pages, int64_t stride_page, cudaStream_t stream) +{ + return qkv_prep_paged_prefill_launch<__nv_bfloat16>( + q_batch, k_batch, v_batch, + q_norm_weight, k_norm_weight, cos_cache, sin_cache, + q_batch_out, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indices_len, page_origin, + num_q_heads, num_kv_heads, seq_len, start_pos, + cos_max_pos, rotary_dim, rms_eps, page_size, num_pages, + stride_page, stream); +} + +// E4m3 KV twin. +int qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda( + const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, + const __nv_bfloat16* v_batch, + const __nv_bfloat16* q_norm_weight, const __nv_bfloat16* k_norm_weight, + const __nv_bfloat16* cos_cache, const __nv_bfloat16* sin_cache, + __nv_bfloat16* q_batch_out, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + const int* page_indices, int page_indices_len, int page_origin, + int num_q_heads, int num_kv_heads, int seq_len, int start_pos, + int cos_max_pos, int rotary_dim, float rms_eps, + int page_size, int num_pages, int64_t stride_page, cudaStream_t stream) +{ + return qkv_prep_paged_prefill_launch<__nv_fp8_e4m3>( + q_batch, k_batch, v_batch, + q_norm_weight, k_norm_weight, cos_cache, sin_cache, + q_batch_out, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indices_len, page_origin, + num_q_heads, num_kv_heads, seq_len, start_pos, + cos_max_pos, rotary_dim, rms_eps, page_size, num_pages, + stride_page, stream); +} + +int qkv_norm_rope_paged_decode_hd256_plain_cuda( + const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, + const __nv_bfloat16* v_batch, + const __nv_bfloat16* q_norm_weight, const __nv_bfloat16* k_norm_weight, + const __nv_bfloat16* cos_cache, const __nv_bfloat16* sin_cache, + __nv_bfloat16* q_batch_out, __nv_bfloat16* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + const int* page_indices, int page_indices_len, + const int* page_indptr, const int* page_origins, const int* positions, + int num_q_heads, int num_kv_heads, int batch, + int cos_max_pos, int rotary_dim, float rms_eps, + int page_size, int num_pages, int64_t stride_page, cudaStream_t stream) +{ + return qkv_prep_paged_decode_launch<__nv_bfloat16>( + q_batch, k_batch, v_batch, + q_norm_weight, k_norm_weight, cos_cache, sin_cache, + q_batch_out, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indices_len, page_indptr, page_origins, + positions, num_q_heads, num_kv_heads, batch, + cos_max_pos, rotary_dim, rms_eps, page_size, num_pages, + stride_page, stream); +} + +// E4m3 KV twin. +int qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda( + const __nv_bfloat16* q_batch, const __nv_bfloat16* k_batch, + const __nv_bfloat16* v_batch, + const __nv_bfloat16* q_norm_weight, const __nv_bfloat16* k_norm_weight, + const __nv_bfloat16* cos_cache, const __nv_bfloat16* sin_cache, + __nv_bfloat16* q_batch_out, void* kv_data, + int64_t k_offset_elems, int64_t v_offset_elems, + const int* page_indices, int page_indices_len, + const int* page_indptr, const int* page_origins, const int* positions, + int num_q_heads, int num_kv_heads, int batch, + int cos_max_pos, int rotary_dim, float rms_eps, + int page_size, int num_pages, int64_t stride_page, cudaStream_t stream) +{ + return qkv_prep_paged_decode_launch<__nv_fp8_e4m3>( + q_batch, k_batch, v_batch, + q_norm_weight, k_norm_weight, cos_cache, sin_cache, + q_batch_out, kv_data, k_offset_elems, v_offset_elems, + page_indices, page_indices_len, page_indptr, page_origins, + positions, num_q_heads, num_kv_heads, batch, + cos_max_pos, rotary_dim, rms_eps, page_size, num_pages, + stride_page, stream); +} + } // extern "C" diff --git a/pegainfer-kernels/src/ffi/gemma4.rs b/pegainfer-kernels/src/ffi/gemma4.rs index af4421b04..ed5e964aa 100644 --- a/pegainfer-kernels/src/ffi/gemma4.rs +++ b/pegainfer-kernels/src/ffi/gemma4.rs @@ -97,4 +97,35 @@ unsafe extern "C" { max_m_blocks: i32, stream: CUstream, ) -> CUresult; + + /// fp8-KV twin of `batch_prefill_paged_window_hd256_cuda`: same launch + /// body over an e4m3 pool (scale 1.0), Q and output bf16. + pub fn gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda( + q: *const Half, + output: *mut Half, + kv_data: *const core::ffi::c_void, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indptr: *const i32, + last_page_len_d: *const i32, + q_indptr: *const i32, + request_indices: *const i32, + qo_tile_indices: *const i32, + kv_tile_indices: *const i32, + kv_chunk_size_ptr: *const i32, + total_num_rows: *const u32, + num_qo_heads: i32, + num_kv_heads: i32, + head_dim: i32, + page_size: i32, + seq_len: i32, + batch_size: i32, + padded_batch_size: i32, + stride_page: i64, + sm_scale: f32, + cta_tile_q_override: i32, + window_left: i32, + stream: CUstream, + ) -> i32; } diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index b5757d639..08f159707 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/pegainfer-kernels/src/ffi/shared.rs @@ -1216,6 +1216,65 @@ unsafe extern "C" { stride_page: i64, stream: CUstream, ) -> i32; + + /// E4m3 KV twin. + pub fn qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_origin: i32, + num_q_heads: i32, + num_kv_heads: i32, + seq_len: i32, + start_pos: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; + + /// E4m3 KV twin. + pub fn qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda( + q_batch: *const Half, + k_batch: *const Half, + v_batch: *const Half, + q_norm_weight: *const Half, + k_norm_weight: *const Half, + cos_cache: *const Half, + sin_cache: *const Half, + q_batch_out: *mut Half, + kv_data: *mut Half, + k_offset_elems: i64, + v_offset_elems: i64, + page_indices: *const i32, + page_indices_len: i32, + page_indptr: *const i32, + page_origins: *const i32, + positions: *const i32, + num_q_heads: i32, + num_kv_heads: i32, + batch: i32, + cos_max_pos: i32, + rotary_dim: i32, + rms_eps: f32, + page_size: i32, + num_pages: i32, + stride_page: i64, + stream: CUstream, + ) -> i32; } // hd512 QK-norm + partial RoPE prep (Gemma 4 global layers): diff --git a/pegainfer-kernels/src/ops/attention.rs b/pegainfer-kernels/src/ops/attention.rs index 201a63ad1..87d6b5ce3 100644 --- a/pegainfer-kernels/src/ops/attention.rs +++ b/pegainfer-kernels/src/ops/attention.rs @@ -5,6 +5,7 @@ use cudarc::driver::DevicePtrMut; use half::bf16; use crate::ffi; +use crate::paged_kv::KvStorage; use crate::paged_kv::PagedKvLayout; use crate::tensor::DeviceContext; use crate::tensor::DeviceVec; @@ -617,6 +618,7 @@ pub fn prefill_attention_paged_into( layer, head_dim, num_kv_heads, + false, )?; let (q_ptr, _gq) = q_batch.data.device_ptr_mut(&ctx.stream); @@ -1262,6 +1264,7 @@ pub fn paged_attention_batch_decode_into( layer, head_dim, num_kv_heads, + false, )?; let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); @@ -1422,6 +1425,7 @@ pub fn paged_attention_batch_decode_split_kv_into( layer, head_dim, num_kv_heads, + false, )?; let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); @@ -1550,6 +1554,7 @@ fn scatter_decode_kv_into_paged( layer, head_dim, num_kv_heads, + false, )?; let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); @@ -1632,6 +1637,7 @@ pub fn paged_attention_batch_decode_hd256_into( layer, head_dim, num_kv_heads, + false, )?; let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); @@ -1749,6 +1755,7 @@ pub fn paged_attention_batch_decode_via_prefill_hd256_into( layer, head_dim, num_kv_heads, + false, )?; let sm_scale = 1.0f32 / (head_dim as f32).sqrt(); @@ -1910,6 +1917,7 @@ pub fn paged_attention_batch_decode_split_kv_hd512_into( layer, 512, num_kv_heads, + false, )?; anyhow::ensure!( row_offset < q.seq_len, @@ -2199,6 +2207,7 @@ pub fn paged_attention_batch_decode_via_prefill_hd512_into( layer, head_dim, num_kv_heads, + false, )?; let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); @@ -2361,6 +2370,7 @@ pub fn batch_prefill_paged_hd512_into( layer, head_dim, num_kv_heads, + false, )?; let (buf_ptr, _gbuf) = kv_buffer.device_ptr(&ctx.stream); @@ -2451,6 +2461,7 @@ pub fn batch_prefill_paged_window_hd256_into( layer, 256, num_kv_heads, + true, )?; anyhow::ensure!( plan.max_page_index < geometry.num_pages, @@ -2551,38 +2562,80 @@ pub fn batch_prefill_paged_window_hd256_into( let (kcs_ptr, _gkcs) = plan.kv_chunk_size_d.device_ptr(&ctx.stream); let (tnr_ptr, _gtnr) = plan.total_num_rows_d.device_ptr(&ctx.stream); - let result = unsafe { - ffi::batch_prefill_paged_window_cuda_hd256( - q_ptr as *const ffi::Half, - out_ptr as *mut ffi::Half, - buf_ptr as *const ffi::Half, - geometry.k_offset_elems, - geometry.v_offset_elems, - pi_ptr as *const i32, - pip_ptr as *const i32, - lpl_ptr as *const i32, - qi_ptr as *const i32, - ri_ptr as *const i32, - qti_ptr as *const i32, - kti_ptr as *const i32, - kcs_ptr as *const i32, - tnr_ptr as *const u32, - num_qo_heads_i32, - num_kv_heads_i32, - 256, - geometry.page_size, - total_tokens_i32, - plan.batch_size(), - plan.num_tiles, - geometry.stride_page, - sm_scale, - window_left, - crate::tensor::active_cu_stream(ctx), - ) + let (result, entry_point) = match layout.storage { + KvStorage::E4m3 => { + #[cfg(feature = "gemma4")] + let result = unsafe { + ffi::gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const core::ffi::c_void, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + qi_ptr as *const i32, + ri_ptr as *const i32, + qti_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + tnr_ptr as *const u32, + num_qo_heads_i32, + num_kv_heads_i32, + 256, + geometry.page_size, + total_tokens_i32, + plan.batch_size(), + plan.num_tiles, + geometry.stride_page, + sm_scale, + 0, + window_left, + crate::tensor::active_cu_stream(ctx), + ) + }; + #[cfg(not(feature = "gemma4"))] + anyhow::bail!("hd256 windowed batch prefill: fp8 KV needs the gemma4 feature"); + #[cfg(feature = "gemma4")] + (result, "gemma4_batch_prefill_paged_window_hd256_fp8kv_cuda") + } + KvStorage::Bf16 => ( + unsafe { + ffi::batch_prefill_paged_window_cuda_hd256( + q_ptr as *const ffi::Half, + out_ptr as *mut ffi::Half, + buf_ptr as *const ffi::Half, + geometry.k_offset_elems, + geometry.v_offset_elems, + pi_ptr as *const i32, + pip_ptr as *const i32, + lpl_ptr as *const i32, + qi_ptr as *const i32, + ri_ptr as *const i32, + qti_ptr as *const i32, + kti_ptr as *const i32, + kcs_ptr as *const i32, + tnr_ptr as *const u32, + num_qo_heads_i32, + num_kv_heads_i32, + 256, + geometry.page_size, + total_tokens_i32, + plan.batch_size(), + plan.num_tiles, + geometry.stride_page, + sm_scale, + window_left, + crate::tensor::active_cu_stream(ctx), + ) + }, + "batch_prefill_paged_window_cuda_hd256", + ), }; if result != 0 { anyhow::bail!( - "batch_prefill_paged_window_cuda_hd256 failed for layer {layer}, \ + "{entry_point} failed for layer {layer}, \ bs={}, tiles={}, qo_heads={num_qo_heads}, kv_heads={num_kv_heads}, \ window_left={window_left}: {result}{}", plan.batch_size(), @@ -2847,7 +2900,21 @@ fn checked_paged_geometry( layer: usize, head_dim: usize, num_kv_heads: usize, + fp8_capable: bool, ) -> Result { + // `pool_len` counts the pool's bf16 backing slots; an e4m3 pool packs two + // elements per slot. Only wrappers with an fp8 kernel twin may see one. + anyhow::ensure!( + fp8_capable || layout.storage == KvStorage::Bf16, + "{what} has no fp8 KV path; the layout carries {}-byte elements", + layout.storage.elem_bytes() + ); + let pool_len = match layout.storage { + KvStorage::E4m3 => pool_len + .checked_mul(2) + .ok_or_else(|| anyhow::anyhow!("{what} fp8 pool element count overflows"))?, + KvStorage::Bf16 => pool_len, + }; anyhow::ensure!( layout.head_dim == head_dim, "{what} layout.head_dim {} != {head_dim}", @@ -3047,6 +3114,7 @@ pub fn qkv_norm_rope_paged_prefill_hd256_plain_into( layer, 256, num_kv_heads, + true, )?; ensure_vec_backed(q_norm_weight, "hd256 paged prep q_norm_weight")?; ensure_vec_backed(k_norm_weight, "hd256 paged prep k_norm_weight")?; @@ -3152,8 +3220,21 @@ pub fn qkv_norm_rope_paged_prefill_hd256_plain_into( let (pi_ptr, _gpi) = page_indices.device_ptr(&ctx.stream); let pi_ptr = pi_ptr + (pages_offset * std::mem::size_of::()) as u64; + // Two `if`s, not a tuple-building match: a fn item only coerces to a + // fn pointer across plain `if` arms. + let fp8 = layout.storage == KvStorage::E4m3; + let launch = if fp8 { + ffi::qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda + } else { + ffi::qkv_norm_rope_paged_prefill_hd256_plain_cuda + }; + let entry_point = if fp8 { + "qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda" + } else { + "qkv_norm_rope_paged_prefill_hd256_plain_cuda" + }; let result = unsafe { - ffi::qkv_norm_rope_paged_prefill_hd256_plain_cuda( + launch( q_ptr as *const ffi::Half, k_ptr as *const ffi::Half, v_ptr as *const ffi::Half, @@ -3183,8 +3264,7 @@ pub fn qkv_norm_rope_paged_prefill_hd256_plain_into( }; if result != 0 { anyhow::bail!( - "qkv_norm_rope_paged_prefill_hd256_plain_cuda failed with error \ - {result}{}", + "{entry_point} failed with error {result}{}", crate::ops::ffi_exception_message(result) ); } @@ -3263,6 +3343,7 @@ pub fn qkv_norm_rope_paged_decode_hd256_plain_into( layer, 256, num_kv_heads, + true, )?; ensure_vec_backed(q_norm_weight, "hd256 paged decode prep q_norm_weight")?; ensure_vec_backed(k_norm_weight, "hd256 paged decode prep k_norm_weight")?; @@ -3337,8 +3418,21 @@ pub fn qkv_norm_rope_paged_decode_hd256_plain_into( let (og_ptr, _gog) = page_origins.device_ptr(&ctx.stream); let (ps_ptr, _gps) = positions.device_ptr(&ctx.stream); + // Two `if`s, not a tuple-building match: a fn item only coerces to a + // fn pointer across plain `if` arms. + let fp8 = layout.storage == KvStorage::E4m3; + let launch = if fp8 { + ffi::qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda + } else { + ffi::qkv_norm_rope_paged_decode_hd256_plain_cuda + }; + let entry_point = if fp8 { + "qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda" + } else { + "qkv_norm_rope_paged_decode_hd256_plain_cuda" + }; let result = unsafe { - ffi::qkv_norm_rope_paged_decode_hd256_plain_cuda( + launch( q_ptr as *const ffi::Half, k_ptr as *const ffi::Half, v_ptr as *const ffi::Half, @@ -3369,8 +3463,7 @@ pub fn qkv_norm_rope_paged_decode_hd256_plain_into( }; if result != 0 { anyhow::bail!( - "qkv_norm_rope_paged_decode_hd256_plain_cuda failed with error \ - {result}{}", + "{entry_point} failed with error {result}{}", crate::ops::ffi_exception_message(result) ); } @@ -3593,6 +3686,7 @@ pub fn qk_norm_partial_rope_paged_prefill_hd512_into( layer, 512, num_kv_heads, + false, )?; anyhow::ensure!( q_norm_weight.len == 512, @@ -3778,6 +3872,7 @@ pub fn qk_norm_partial_rope_paged_decode_hd512_into( layer, 512, num_kv_heads, + false, )?; ensure_vec_backed(q_norm_weight, "hd512 paged decode prep q_norm_weight")?; ensure_vec_backed(k_norm_weight, "hd512 paged decode prep k_norm_weight")?; diff --git a/pegainfer-kernels/src/paged_kv.rs b/pegainfer-kernels/src/paged_kv.rs index f99d6bd48..b2f4ea457 100644 --- a/pegainfer-kernels/src/paged_kv.rs +++ b/pegainfer-kernels/src/paged_kv.rs @@ -1,3 +1,18 @@ +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum KvStorage { + Bf16, + E4m3, +} + +impl KvStorage { + pub const fn elem_bytes(self) -> usize { + match self { + Self::Bf16 => 2, + Self::E4m3 => 1, + } + } +} + /// Page-first geometry used by paged-KV kernels. /// /// This is kernel-facing shape metadata only. Pool allocation, page ownership, @@ -14,10 +29,27 @@ pub struct PagedKvLayout { pub layer_stride: usize, /// Elements per page (all layers): num_layers x layer_stride. pub page_stride: usize, + pub storage: KvStorage, } impl PagedKvLayout { pub fn new(num_layers: usize, num_kv_heads: usize, head_dim: usize, page_size: usize) -> 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, + ) -> Self { let kv_block_len = page_size * num_kv_heads * head_dim; let layer_stride = 2 * kv_block_len; let page_stride = num_layers * layer_stride; @@ -29,6 +61,7 @@ impl PagedKvLayout { kv_block_len, layer_stride, page_stride, + storage, } } } diff --git a/pegainfer-kernels/tests/hd256_fp8_pool.rs b/pegainfer-kernels/tests/hd256_fp8_pool.rs new file mode 100644 index 000000000..61b82a2b9 --- /dev/null +++ b/pegainfer-kernels/tests/hd256_fp8_pool.rs @@ -0,0 +1,470 @@ +//! Device gates for Gemma 4's e4m3 paged-KV storage contract. + +#![cfg(feature = "gemma4")] + +mod common; + +use cudarc::driver::CudaSlice; +use half::bf16; +use pegainfer_kernels::ops::PrefillPagedPlan; +use pegainfer_kernels::ops::batch_prefill_paged_window_hd256_into; +use pegainfer_kernels::ops::paged_attention_batch_decode_hd256_into; +use pegainfer_kernels::ops::qkv_norm_rope_paged_decode_hd256_plain_into; +use pegainfer_kernels::ops::qkv_norm_rope_paged_prefill_hd256_plain_into; +use pegainfer_kernels::paged_kv::KvStorage; +use pegainfer_kernels::paged_kv::PagedKvLayout; +use pegainfer_kernels::tensor::DeviceContext; +use pegainfer_kernels::tensor::DeviceVec; +use pegainfer_kernels::tensor::HiddenStates; + +const HD: usize = 256; +const PAGE_SIZE: usize = 2; +const NUM_LAYERS: usize = 3; + +fn packed_fp8(bytes: &[u8]) -> Vec { + (0..bytes.len() / 2) + .map(|slot| bf16::from_bits(u16::from_le_bytes([bytes[2 * slot], bytes[2 * slot + 1]]))) + .collect() +} + +fn raw_bytes(ctx: &DeviceContext, pool: &CudaSlice) -> Vec { + ctx.stream + .clone_dtoh(pool) + .expect("pool D2H") + .into_iter() + .flat_map(|value| value.to_bits().to_le_bytes()) + .collect() +} + +fn constant_states(ctx: &DeviceContext, value: f32, rows: usize) -> HiddenStates { + HiddenStates::from_host(ctx, &vec![bf16::from_f32(value); HD * rows], HD, rows) + .expect("states H2D") +} + +fn identity_rope(ctx: &DeviceContext, rows: usize) -> (DeviceVec, DeviceVec) { + let cos = vec![bf16::ONE; HD * rows]; + let sin = vec![bf16::ZERO; HD * rows]; + ( + DeviceVec::from_host(ctx, &cos).expect("cos H2D"), + DeviceVec::from_host(ctx, &sin).expect("sin H2D"), + ) +} + +#[test] +fn fp8_prep_stores_exact_bytes_at_layout_offsets() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let layout = PagedKvLayout::with_storage(NUM_LAYERS, 1, HD, PAGE_SIZE, KvStorage::E4m3); + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride * 3 / 2) + .expect("pool alloc"); + let q = constant_states(&ctx, 1.0, 3); + let k = constant_states(&ctx, 1.0, 3); + let v = constant_states(&ctx, 0.5, 3); + let mut q_out = HiddenStates::zeros(&ctx, HD, 3).expect("q_out alloc"); + let weights = DeviceVec::from_host(&ctx, &vec![bf16::from_f32(2.0); HD]).expect("weights H2D"); + let (cos, sin) = identity_rope(&ctx, 3); + let pages = ctx.stream.clone_htod(&[2i32, 0]).expect("pages H2D"); + qkv_norm_rope_paged_prefill_hd256_plain_into( + &ctx, &q, &k, &v, &mut q_out, 0, &pool, &layout, &weights, &weights, &cos, &sin, 1, &pages, + 0, 0, 0, 3, 1, 1, HD, 0.0, + ) + .expect("fp8 prep"); + let got = raw_bytes(&ctx, &pool); + let mut expected = vec![0u8; layout.page_stride * 3]; + for (token, page) in [2usize, 2, 0].into_iter().enumerate() { + let slot = token % PAGE_SIZE; + let layer = page * layout.page_stride + layout.layer_stride; + let k = layer + slot * HD; + let v = layer + layout.kv_block_len + slot * HD; + expected[k..k + HD].fill(0x40); + expected[v..v + HD].fill(0x38); + } + assert_eq!(got, expected); +} + +#[test] +fn fp8_decode_prep_stores_exact_bytes_at_layout_offsets() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let layout = PagedKvLayout::with_storage(NUM_LAYERS, 1, HD, PAGE_SIZE, KvStorage::E4m3); + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride * 3 / 2) + .expect("pool alloc"); + let q = constant_states(&ctx, 1.0, 3); + let k = constant_states(&ctx, 1.0, 3); + let v = constant_states(&ctx, 0.5, 3); + let mut q_out = HiddenStates::zeros(&ctx, HD, 3).expect("q_out alloc"); + let weights = DeviceVec::from_host(&ctx, &vec![bf16::from_f32(2.0); HD]).expect("weights H2D"); + let (cos, sin) = identity_rope(&ctx, 4); + let pages = ctx.stream.clone_htod(&[2i32]).expect("pages H2D"); + let indptr = ctx.stream.clone_htod(&[0i32, 1]).expect("indptr H2D"); + // Released-front page of the row window: resident row = pos / page_size - origin. + let origins = ctx.stream.clone_htod(&[1i32]).expect("origins H2D"); + let positions = ctx.stream.clone_htod(&[3i32]).expect("positions H2D"); + qkv_norm_rope_paged_decode_hd256_plain_into( + &ctx, &q, &k, &v, &mut q_out, 2, &pool, &layout, &weights, &weights, &cos, &sin, 1, &pages, + &indptr, &origins, &positions, 4, 1, 1, HD, 0.0, + ) + .expect("fp8 decode prep"); + let got = raw_bytes(&ctx, &pool); + let mut expected = vec![0u8; layout.page_stride * 3]; + let layer = 2 * layout.page_stride + layout.layer_stride; + let k = layer + HD; + let v = layer + layout.kv_block_len + HD; + expected[k..k + HD].fill(0x40); + expected[v..v + HD].fill(0x38); + assert_eq!(got, expected); +} + +fn constant_pool(ctx: &DeviceContext, storage: KvStorage, pages: usize) -> CudaSlice { + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, storage); + let values = [1.0f32, 2.0, 0.5, -1.0]; + match storage { + KvStorage::Bf16 => { + let mut host = vec![bf16::ZERO; layout.page_stride * pages]; + for page in 0..pages { + for slot in 0..PAGE_SIZE { + let base = page * layout.page_stride + slot * HD; + host[base..base + HD].fill(bf16::from_f32(values[slot])); + let v = base + layout.kv_block_len; + host[v..v + HD].fill(bf16::from_f32(values[slot + 2])); + } + } + ctx.stream.clone_htod(&host).expect("bf16 pool H2D") + } + KvStorage::E4m3 => { + let mut bytes = vec![0u8; layout.page_stride * pages]; + for page in 0..pages { + for slot in 0..PAGE_SIZE { + let base = page * layout.page_stride + slot * HD; + bytes[base..base + HD].fill([0x38, 0x40][slot]); + let v = base + layout.kv_block_len; + bytes[v..v + HD].fill([0x30, 0xb8][slot]); + } + } + ctx.stream + .clone_htod(&packed_fp8(&bytes)) + .expect("fp8 pool H2D") + } + } +} + +fn e4m3_to_f32(byte: u8) -> f32 { + let sign = if byte & 0x80 == 0 { 1.0 } else { -1.0 }; + let exponent = i32::from((byte >> 3) & 0x0f); + let mantissa = f32::from(byte & 0x07) / 8.0; + if exponent == 0 { + sign * mantissa * 2.0f32.powi(-6) + } else { + sign * (1.0 + mantissa) * 2.0f32.powi(exponent - 7) + } +} + +fn valid_e4m3_code(ordinal: usize) -> u8 { + let code = ordinal % 254; + u8::try_from(code + usize::from(code >= 0x7f)).expect("e4m3 code") +} + +fn varied_values(pages: usize) -> Vec { + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, KvStorage::Bf16); + let mut values = vec![0.0; layout.page_stride * pages]; + for page in 0..pages { + for slot in 0..PAGE_SIZE { + for element in 0..HD { + for block in 0..2 { + let index = page * layout.page_stride + + block * layout.kv_block_len + + slot * HD + + element; + let seed = slot * HD * 2 + element * 2 + block; + let variation = ((5 + seed % 5) << 3) + | (seed % 8) + | if seed.is_multiple_of(3) { 0x80 } else { 0 }; + let byte = valid_e4m3_code(page + variation); + values[index] = e4m3_to_f32(byte); + } + } + } + } + values +} + +fn exact_e4m3(value: f32) -> u8 { + (0u8..=254) + .find(|&byte| byte != 0x7f && e4m3_to_f32(byte).to_bits() == value.to_bits()) + .expect("value must be exactly e4m3-representable") +} + +fn semantic_pool(ctx: &DeviceContext, storage: KvStorage, pages: usize) -> CudaSlice { + let values = varied_values(pages); + match storage { + KvStorage::Bf16 => ctx + .stream + .clone_htod( + &values + .iter() + .map(|&value| bf16::from_f32(value)) + .collect::>(), + ) + .expect("bf16 pool H2D"), + KvStorage::E4m3 => ctx + .stream + .clone_htod(&packed_fp8( + &values + .iter() + .map(|&value| exact_e4m3(value)) + .collect::>(), + )) + .expect("fp8 pool H2D"), + } +} + +fn varied_q(ctx: &DeviceContext, rows: usize) -> HiddenStates { + let host: Vec = (0..HD * rows) + .map(|index| bf16::from_f32(0.25 + (index % HD) as f32 / 512.0)) + .collect(); + HiddenStates::from_host(ctx, &host, HD, rows).expect("q H2D") +} + +fn attend(ctx: &DeviceContext, storage: KvStorage, kv_len: usize, window_left: i32) -> Vec { + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, storage); + let pages = kv_len.div_ceil(PAGE_SIZE); + let pool = semantic_pool(ctx, storage, pages); + let page_indices: Vec = (0..pages as i32).collect(); + let plan = PrefillPagedPlan::new_with_cta_tile_q( + ctx, + &page_indices, + (kv_len - 1) % PAGE_SIZE + 1, + kv_len - 1, + 1, + 1, + 1, + HD, + 0, + ) + .expect("plan"); + let q = varied_q(ctx, 1); + let mut output = HiddenStates::zeros(ctx, HD, 1).expect("output alloc"); + batch_prefill_paged_window_hd256_into( + ctx, + &q, + &pool, + &layout, + 0, + &plan, + &mut output, + 1, + 1.0, + window_left, + ) + .expect("attention"); + output + .to_host(ctx) + .expect("output D2H") + .into_iter() + .map(|value| bf16::from_f32(value).to_bits()) + .collect() +} + +#[test] +fn fp8_window_read_matches_bf16_for_exact_values() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + assert_eq!( + attend(&ctx, KvStorage::E4m3, 2, -1), + attend(&ctx, KvStorage::Bf16, 2, -1) + ); +} + +#[test] +fn fp8_finite_window_read_matches_bf16_and_changes_the_result() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let full = attend(&ctx, KvStorage::E4m3, 6, -1); + let windowed = attend(&ctx, KvStorage::E4m3, 6, 2); + assert_eq!(windowed, attend(&ctx, KvStorage::Bf16, 6, 2)); + assert_ne!( + windowed, full, + "finite window must change the varied-pool output" + ); +} + +fn geometry_probe(ctx: &DeviceContext, prefix_rows: usize) -> Vec { + let pages = prefix_rows.div_ceil(PAGE_SIZE) + 1; + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, KvStorage::E4m3); + let pool = constant_pool(ctx, KvStorage::E4m3, pages); + let (page_lists, starts, lengths, lasts) = if prefix_rows == 0 { + (vec![vec![0]], vec![1], vec![1], vec![2]) + } else { + let prefix_pages: Vec = (0..pages as i32 - 1).collect(); + ( + vec![prefix_pages, vec![pages as i32 - 1]], + vec![0, 1], + vec![prefix_rows, 1], + vec![(prefix_rows - 1) % PAGE_SIZE + 1, 2], + ) + }; + let plan = PrefillPagedPlan::new_batch_with_cta_tile_q( + ctx, + &page_lists, + &lasts, + &starts, + &lengths, + 1, + 1, + HD, + 0, + ) + .expect("batch plan"); + let q = constant_states(ctx, 0.0, prefix_rows + 1); + let mut output = HiddenStates::zeros(ctx, HD, prefix_rows + 1).expect("output alloc"); + batch_prefill_paged_window_hd256_into( + ctx, + &q, + &pool, + &layout, + 0, + &plan, + &mut output, + 1, + 1.0, + -1, + ) + .expect("attention"); + let host = output.to_host(ctx).expect("output D2H"); + host[prefix_rows * HD..] + .iter() + .map(|&value| bf16::from_f32(value).to_bits()) + .collect() +} + +#[test] +fn fp8_window_read_is_geometry_invariant_for_the_probed_row() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let lone = geometry_probe(&ctx, 0); + let packed = geometry_probe(&ctx, 300); + if let Some(index) = lone.iter().zip(&packed).position(|(a, b)| a != b) { + panic!( + "geometry dependence at output[{index}]: lone={:#06x}, packed={:#06x}", + lone[index], packed[index] + ); + } +} + +fn varied_geometry_probe(ctx: &DeviceContext, prefix_rows: usize, pages: usize) -> Vec { + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, KvStorage::E4m3); + let pool = semantic_pool(ctx, KvStorage::E4m3, pages); + let (page_lists, starts, lengths, lasts) = if prefix_rows == 0 { + (vec![vec![pages as i32 - 1]], vec![1], vec![1], vec![2]) + } else { + let prefix_pages: Vec = (0..pages as i32 - 1).rev().collect(); + ( + vec![prefix_pages, vec![pages as i32 - 1]], + vec![0, 1], + vec![prefix_rows, 1], + vec![(prefix_rows - 1) % PAGE_SIZE + 1, 2], + ) + }; + let plan = PrefillPagedPlan::new_batch_with_cta_tile_q( + ctx, + &page_lists, + &lasts, + &starts, + &lengths, + 1, + 1, + HD, + 0, + ) + .expect("batch plan"); + let q = varied_q(ctx, prefix_rows + 1); + let mut output = HiddenStates::zeros(ctx, HD, prefix_rows + 1).expect("output alloc"); + batch_prefill_paged_window_hd256_into( + ctx, + &q, + &pool, + &layout, + 0, + &plan, + &mut output, + 1, + 1.0, + -1, + ) + .expect("attention"); + let host = output.to_host(ctx).expect("output D2H"); + host[prefix_rows * HD..] + .iter() + .map(|&value| bf16::from_f32(value).to_bits()) + .collect() +} + +#[test] +fn varied_fp8_window_read_is_geometry_invariant_for_the_probed_row() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let packed_rows = 300; + let pages = (packed_rows + 1usize).div_ceil(PAGE_SIZE); + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, KvStorage::E4m3); + let bytes: Vec = varied_values(pages).into_iter().map(exact_e4m3).collect(); + let probed = &bytes[(pages - 1) * layout.page_stride..pages * layout.page_stride]; + for page in 0..pages - 1 { + let other = &bytes[page * layout.page_stride..(page + 1) * layout.page_stride]; + assert_ne!(probed, other, "probed page aliases physical page {page}"); + } + let lone = varied_geometry_probe(&ctx, 0, pages); + let packed = varied_geometry_probe(&ctx, packed_rows, pages); + if let Some(index) = lone.iter().zip(&packed).position(|(a, b)| a != b) { + panic!( + "varied-pool geometry dependence at output[{index}]: lone={:#06x}, packed={:#06x}", + lone[index], packed[index] + ); + } +} + +#[test] +fn decode_wrapper_without_fp8_twin_refuses_e4m3() { + let Some(ctx) = common::device_or_skip() else { + return; + }; + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, KvStorage::E4m3); + let pool: CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride / 2) + .expect("pool"); + let state = constant_states(&ctx, 0.0, 1); + let mut output = HiddenStates::zeros(&ctx, HD, 1).expect("output"); + let meta = ctx.stream.clone_htod(&[0i32]).expect("metadata"); + let indptr = ctx.stream.clone_htod(&[0i32, 1]).expect("indptr"); + let err = paged_attention_batch_decode_hd256_into( + &ctx, + &state, + &state, + &state, + &pool, + &layout, + 0, + &meta, + &indptr, + &meta, + &meta, + &meta, + &meta, + &meta, + &mut output, + 1, + 1, + ) + .expect_err("unsupported fp8 wrapper must reject"); + assert!(err.to_string().contains("has no fp8 KV path"), "{err}"); +} diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 268d216fb..d54f4d1e6 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -39,10 +39,12 @@ GPU_LOCK_ROOT=/tmp # run a gate without producing the whole suite's prerequisites. GATES_NUMERIC_PARITY=( "gpu,ckpt,fixtures serve::oracle::context_waypoints_match_hf" + "gpu,ckpt,fixtures serve::oracle::fp8_argmax_agreement_meets_the_bf16_floor" "gpu,ckpt,fixtures serve::oracle::greedy_matches_hf_generate" ) GATES_ADMISSION=( "gpu,ckpt,prompts serve::oracle::mixed_step_matches_serial" + "gpu,ckpt,prompts serve::oracle::fp8_mixed_walk_holds_its_structure" "gpu,ckpt,prompts engine::lane_tests::the_gathered_walk_matches_the_serial_path" "gpu,ckpt,prompts engine::lane_tests::the_gathered_transient_leaves_headroom" ) @@ -89,6 +91,15 @@ GATES_KERNELS=( "gpu ops::norm::parity::the_epilogue_norm_pair_matches_its_parts" "gpu ops::norm::parity::the_moe_combine_tail_matches_its_parts" ) +GATES_KERNELS_HD256_FP8_POOL=( + "gpu fp8_prep_stores_exact_bytes_at_layout_offsets" + "gpu fp8_decode_prep_stores_exact_bytes_at_layout_offsets" + "gpu fp8_window_read_matches_bf16_for_exact_values" + "gpu fp8_finite_window_read_matches_bf16_and_changes_the_result" + "gpu fp8_window_read_is_geometry_invariant_for_the_probed_row" + "gpu varied_fp8_window_read_is_geometry_invariant_for_the_probed_row" + "gpu decode_wrapper_without_fp8_twin_refuses_e4m3" +) MANIFEST_LIB=( "${GATES_NUMERIC_PARITY[@]}" "${GATES_ADMISSION[@]}" @@ -99,6 +110,15 @@ MANIFEST_LIB=( "${GATES_DEVICE[@]}" "${GATES_ROUTED[@]}" ) +GATES_FP8_PROFILE=( + "serve::oracle::context_waypoints_match_hf" + "serve::oracle::greedy_matches_hf_generate" + "serve::oracle::fp8_argmax_agreement_meets_the_bf16_floor" + "serve::oracle::fp8_mixed_walk_holds_its_structure" + "serve::oracle::a_ragged_batch_does_not_depend_on_row_order" + "serve::oracle::eviction_is_footprint_only" + "serve::oracle::overlapped_prefill_matches_the_sync_step" +) # Integration gates live in their own binaries, which `--lib` cannot see. One # array per binary, named GATES_; the target list itself is held @@ -122,9 +142,27 @@ CHAT_GOLDEN=test_data/gemma4-tokenizer-golden.json die() { echo "gemma4 gates: $*" >&2; exit 1; } +gate_is_in() { + local wanted=$1 candidate + shift + for candidate in "$@"; do + [ "$candidate" = "$wanted" ] && return 0 + done + return 1 +} + root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) cd "$root" || die "cannot enter the repository root" +[ -z "${PEGAINFER_KV_FP8+x}" ] || die \ + "PEGAINFER_KV_FP8 is ambient; PEGAINFER_GATE_STORAGE is the only storage switch" +gate_storage=${PEGAINFER_GATE_STORAGE-bf16} +case "$gate_storage" in + bf16) ;; + fp8) export PEGAINFER_KV_FP8=local ;; + *) die "PEGAINFER_GATE_STORAGE must be unset, bf16, or fp8" ;; +esac + # --- prerequisites, one refusal per tier ---------------------------------- # Each is demanded only when a selected gate declares it, so a focused run # carries the cost of what it runs: the device-only gates need no checkpoint, @@ -176,6 +214,10 @@ require_gpu() { exec {gpu_lock_fd}<"$lock_path" || die "cannot open device lock $lock_path" flock -n "$gpu_lock_fd" || die "GPU $gpu_uuid is already owned by another Gemma 4 gate runner" echo "gemma4 gates: claimed GPU $gpu_uuid (selector $selector)" + echo "gemma4 gates: storage profile $gate_storage" + if [ -n "${PEGAINFER_KV_FP8:-}" ]; then + echo "gemma4 gates: PEGAINFER_KV_FP8=$PEGAINFER_KV_FP8" + fi } require_ckpt() { @@ -263,6 +305,13 @@ ignored_in() { --ignored --list 2>/dev/null | sed -n 's/^\(.*\): test$/\1/p' | sort } +listed_in() { + local crate=$1 + shift + cargo test --release -p "$crate" --features "$FEATURE" "$@" -- \ + --list 2>/dev/null | sed -n 's/^\(.*\): test$/\1/p' | sort +} + check_membership() { local what=$1 listing=$2 expected=$3 missing extra missing=$(comm -13 <(printf '%s\n' "$listing") <(printf '%s\n' "$expected")) @@ -276,6 +325,10 @@ lib_listing=$(ignored_in "$CRATE" --lib) lib_names=() for entry in "${MANIFEST_LIB[@]}"; do lib_names+=("${entry##* }"); done check_membership "library" "$lib_listing" "$(printf '%s\n' "${lib_names[@]}" | sort)" +for gate in "${GATES_FP8_PROFILE[@]}"; do + gate_is_in "$gate" "${lib_names[@]}" || die \ + "the fp8 storage profile names a gate outside the library manifest: $gate" +done kernels_listing=$(ignored_in "$KERNELS_CRATE" --lib) [ -n "$kernels_listing" ] || die "could not list the kernels library's ignored gates" @@ -284,6 +337,13 @@ for entry in "${GATES_KERNELS[@]}"; do kernels_names+=("${entry##* }"); done check_membership "kernels library" "$kernels_listing" \ "$(printf '%s\n' "${kernels_names[@]}" | sort)" +kernels_pool_listing=$(listed_in "$KERNELS_CRATE" --test hd256_fp8_pool) +[ -n "$kernels_pool_listing" ] || die "could not list the kernels hd256_fp8_pool integration gates" +kernels_pool_names=() +for entry in "${GATES_KERNELS_HD256_FP8_POOL[@]}"; do kernels_pool_names+=("${entry##* }"); done +check_membership "kernels integration binary hd256_fp8_pool" "$kernels_pool_listing" \ + "$(printf '%s\n' "${kernels_pool_names[@]}" | sort)" + # The integration binaries the crate actually has, so adding one without a # manifest entry fails here instead of leaving its gates unowned. discovered=$(find "$CRATE/tests" -maxdepth 1 -name '*.rs' -exec basename {} .rs \; 2>/dev/null | sort) @@ -320,6 +380,9 @@ done for entry in "${GATES_KERNELS[@]}"; do append_gate "${entry%% *}" kernels "${entry##* }" done +for entry in "${GATES_KERNELS_HD256_FP8_POOL[@]}"; do + append_gate "${entry%% *}" kernels:hd256_fp8_pool "${entry##* }" +done manifest_gate_count=${#all_gates[@]} for entry in "${GATES_DENSE_AND_ROUTED[@]}"; do routed_needs=${entry%% *} @@ -330,6 +393,10 @@ done filter=${1:-} selected=() for entry in "${all_gates[@]}"; do + if [ "$gate_storage" = fp8 ]; then + gate=${entry##*|} + gate_is_in "$gate" "${GATES_FP8_PROFILE[@]}" || continue + fi [ -z "$filter" ] || [[ ${entry##*|} == *"$filter"* ]] || continue selected+=("$entry") done @@ -360,9 +427,16 @@ failed=() for entry in "${selected[@]}"; do IFS='|' read -r _needs target profile gate <<<"$entry" test_crate=$CRATE + require_gpu_env=0 + ignored_args=(--ignored) if [ "$target" = kernels ]; then test_crate=$KERNELS_CRATE target_args=(--lib) + elif [[ $target == kernels:* ]]; then + test_crate=$KERNELS_CRATE + target_args=(--test "${target#kernels:}") + require_gpu_env=1 + ignored_args=() elif [ "$target" = lib ]; then target_args=(--lib) else @@ -375,10 +449,13 @@ for entry in "${selected[@]}"; do device) ;; *) die "unknown execution profile $profile" ;; esac + if [ "$require_gpu_env" -eq 1 ]; then + model_env=(env PEGAINFER_REQUIRE_GPU=1) + fi echo "--- [$profile] $gate" if "${model_env[@]}" cargo test --release -p "$test_crate" --features "$FEATURE" \ "${target_args[@]}" -- \ - --ignored --exact "$gate" --test-threads=1 --nocapture 2>&1 | tail -20; then + "${ignored_args[@]}" --exact "$gate" --test-threads=1 --nocapture 2>&1 | tail -20; then completed=$((completed + 1)) else failed+=("[$profile] $gate")