From fe393f7564590e7cd90303c920ecf2a2cb8707ab Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 18:18:10 +0100 Subject: [PATCH 01/17] feat(gemma4): the sliding family serves an opt in fp8 KV pool Signed-off-by: Feathbow --- pegainfer-core/src/kv_pool.rs | 51 ++++++- pegainfer-gemma4/src/serve.rs | 25 +++- pegainfer-kernels/build.rs | 1 + pegainfer-kernels/csrc/gemma4/local_kv_fp8.cu | 27 ++++ .../csrc/shared/paged_launch.cuh | 26 ++-- .../shared/prefill_attention_hd256_plain.cu | 139 ++++++++++++++++-- pegainfer-kernels/src/ffi/gemma4.rs | 31 ++++ pegainfer-kernels/src/ffi/shared.rs | 59 ++++++++ pegainfer-kernels/src/ops/attention.rs | 134 +++++++++++++---- pegainfer-kernels/src/paged_kv.rs | 17 +++ 10 files changed, 448 insertions(+), 62 deletions(-) create mode 100644 pegainfer-kernels/csrc/gemma4/local_kv_fp8.cu diff --git a/pegainfer-core/src/kv_pool.rs b/pegainfer-core/src/kv_pool.rs index 543ced8c7..966b99228 100644 --- a/pegainfer-core/src/kv_pool.rs +++ b/pegainfer-core/src/kv_pool.rs @@ -25,6 +25,8 @@ pub struct KvLayout { pub layer_stride: usize, /// Elements per page (all layers): num_layers × layer_stride. pub page_stride: usize, + /// Bytes per stored element. Strides remain in elements. + pub elem_bytes: usize, } impl KvLayout { @@ -34,6 +36,20 @@ impl KvLayout { head_dim: usize, page_size: usize, ) -> anyhow::Result { + Self::with_elem_bytes(num_layers, num_kv_heads, head_dim, page_size, 2) + } + + pub fn with_elem_bytes( + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + page_size: usize, + elem_bytes: usize, + ) -> anyhow::Result { + anyhow::ensure!( + elem_bytes == 1 || elem_bytes == 2, + "paged KV elements are bf16 (2 bytes) or e4m3 (1 byte), not {elem_bytes} bytes" + ); let strides = || -> Option<(usize, usize, usize)> { let kv_block_len = page_size.checked_mul(num_kv_heads)?.checked_mul(head_dim)?; let layer_stride = kv_block_len.checked_mul(2)?; @@ -54,6 +70,7 @@ impl KvLayout { kv_block_len, layer_stride, page_stride, + elem_bytes, }) } @@ -66,6 +83,7 @@ impl KvLayout { kv_block_len: self.kv_block_len, layer_stride: self.layer_stride, page_stride: self.page_stride, + elem_bytes: self.elem_bytes, } } } @@ -105,7 +123,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_elem_bytes( + ctx, + num_layers, + num_kv_heads, + head_dim, + page_size, + num_pages, + 2, + ) + } + + #[allow(clippy::too_many_arguments)] + pub fn with_elem_bytes( + ctx: &DeviceContext, + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + page_size: usize, + num_pages: usize, + elem_bytes: usize, + ) -> Result { + let layout = + KvLayout::with_elem_bytes(num_layers, num_kv_heads, head_dim, page_size, elem_bytes)?; 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 +154,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.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); diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 8b1818c96..e1fc62295 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -313,6 +313,12 @@ fn copy_pool_pages( dst: &[i32], ) -> Result<()> { use cudarc::driver::DevicePtr; + // Startup refuses the prefix cache alongside the fp8 pool, so a one-byte + // layout here means a byte-domain bug, not a configuration. + anyhow::ensure!( + layout.elem_bytes == 2, + "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", @@ -649,6 +655,22 @@ pub(crate) fn global_split_factor(config: &Gemma4Config) -> Result { } /// Everything a serving step needs that outlives requests. +fn local_kv_elem_bytes() -> Result { + match std::env::var("PEGAINFER_KV_FP8") { + Err(std::env::VarError::NotPresent) => Ok(2), + Ok(value) if value == "local" => { + anyhow::ensure!( + std::env::var_os("PEGAINFER_PREFIX_CACHE").is_none(), + "PEGAINFER_KV_FP8 and PEGAINFER_PREFIX_CACHE cannot combine: the prefix cache \ + copies pool pages in bf16 element units" + ); + Ok(1) + } + Ok(value) => anyhow::bail!("PEGAINFER_KV_FP8 supports only \"local\", got {value:?}"), + Err(err) => anyhow::bail!("PEGAINFER_KV_FP8 is not unicode: {err}"), + } +} + pub(crate) struct GemmaServe { /// The weights these pools, rope tables and layer numbering were built /// for. Holding them is what makes a step's model identity structural: @@ -744,13 +766,14 @@ impl GemmaServe { } }) .collect(); - let local_pool = KvPool::new( + let local_pool = KvPool::with_elem_bytes( ctx, locals, config.num_key_value_heads, config.head_dim, PAGE_SIZE, local_pages, + local_kv_elem_bytes()?, )?; let global_pool = KvPool::new( ctx, 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..95430f8c0 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); +} + +// fp8 KV pool twin: same math, e4m3 stores at scale 1.0. +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); +} + +// fp8 KV pool twin: same math, e4m3 stores at scale 1.0. +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..1a96d3f70 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; + + /// fp8-KV twin: e4m3 stores at scale 1.0 into a one-byte-element pool. + 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; + + /// fp8-KV twin: e4m3 stores at scale 1.0 into a one-byte-element pool. + 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..f7f433901 100644 --- a/pegainfer-kernels/src/ops/attention.rs +++ b/pegainfer-kernels/src/ops/attention.rs @@ -617,6 +617,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 +1263,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 +1424,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 +1553,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 +1636,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 +1754,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 +1916,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 +2206,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 +2369,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 +2460,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,34 +2561,70 @@ 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 = if layout.elem_bytes == 1 { + #[cfg(feature = "gemma4")] + 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") + } else { + 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), + ) + } }; if result != 0 { anyhow::bail!( @@ -2847,7 +2893,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.elem_bytes == 2, + "{what} has no fp8 KV path; the layout carries {}-byte elements", + layout.elem_bytes + ); + let pool_len = match layout.elem_bytes { + 1 => pool_len + .checked_mul(2) + .ok_or_else(|| anyhow::anyhow!("{what} fp8 pool element count overflows"))?, + _ => pool_len, + }; anyhow::ensure!( layout.head_dim == head_dim, "{what} layout.head_dim {} != {head_dim}", @@ -3047,6 +3107,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 +3213,13 @@ 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; + let launch = if layout.elem_bytes == 1 { + ffi::qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda + } else { + ffi::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, @@ -3263,6 +3329,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 +3404,13 @@ 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); + let launch = if layout.elem_bytes == 1 { + ffi::qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda + } else { + ffi::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, @@ -3593,6 +3665,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 +3851,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..7b9f7fb03 100644 --- a/pegainfer-kernels/src/paged_kv.rs +++ b/pegainfer-kernels/src/paged_kv.rs @@ -14,10 +14,26 @@ pub struct PagedKvLayout { pub layer_stride: usize, /// Elements per page (all layers): num_layers x layer_stride. pub page_stride: usize, + /// Bytes per stored KV element. Strides remain in elements. + pub elem_bytes: usize, } impl PagedKvLayout { pub fn new(num_layers: usize, num_kv_heads: usize, head_dim: usize, page_size: usize) -> Self { + Self::with_elem_bytes(num_layers, num_kv_heads, head_dim, page_size, 2) + } + + pub fn with_elem_bytes( + num_layers: usize, + num_kv_heads: usize, + head_dim: usize, + page_size: usize, + elem_bytes: usize, + ) -> Self { + assert!( + elem_bytes == 1 || elem_bytes == 2, + "paged KV elements are bf16 (2 bytes) or e4m3 (1 byte), not {elem_bytes} bytes" + ); 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 +45,7 @@ impl PagedKvLayout { kv_block_len, layer_stride, page_stride, + elem_bytes, } } } From 82c42796c5a6e82e022e853e99ac163fe7114561 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 21:11:22 +0100 Subject: [PATCH 02/17] fix(gemma4): the fp8 pool reads the prefix cache capacity, not the variable Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 2 +- pegainfer-gemma4/src/serve.rs | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index d4aabb0fa..5ff3814cc 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -154,7 +154,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)) } diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index e1fc62295..57a66bc34 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -659,8 +659,10 @@ fn local_kv_elem_bytes() -> Result { match std::env::var("PEGAINFER_KV_FP8") { Err(std::env::VarError::NotPresent) => Ok(2), Ok(value) if value == "local" => { + // The parsed capacity, not the variable's presence: "0", "off" + // and an empty value all mean the cache is disabled. anyhow::ensure!( - std::env::var_os("PEGAINFER_PREFIX_CACHE").is_none(), + crate::engine::prefix_cache_cap()?.is_none(), "PEGAINFER_KV_FP8 and PEGAINFER_PREFIX_CACHE cannot combine: the prefix cache \ copies pool pages in bf16 element units" ); From 192f9a2c0ac2834bfa6bd0b6b52b2d2c52f09bdb Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 21:11:22 +0100 Subject: [PATCH 03/17] docs(gemma4): the fp8 KV pool serving contract Signed-off-by: Feathbow --- CLAUDE.md | 1 + docs/models/gemma4/serving.md | 4 ++++ 2 files changed, 5 insertions(+) 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..88df17896 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 at startup, 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. From 24a353f1c47137b61375a1d3400bb6e12a090ea5 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 21:26:59 +0100 Subject: [PATCH 04/17] fix(gemma4): the KV storage width is a type, parsed once before the load PagedKvLayout and KvLayout carry KvStorage::{Bf16, E4m3} instead of a public byte count, so an illegal width is unrepresentable and the geometry gate dispatches on the enum rather than defaulting everything that is not exactly one byte to bf16. The three fp8 dispatch sites name the entry point that actually ran when they fail. PEGAINFER_KV_FP8 parses once in the engine beside the other knobs, before the checkpoint loads, and GemmaServe receives the storage as a parameter; the second raw-env interpretation in serve.rs is gone, and the enabled-prefix-cache refusal keeps its message. Signed-off-by: Feathbow --- docs/models/gemma4/serving.md | 2 +- pegainfer-core/src/kv_pool.rs | 36 ++-- pegainfer-gemma4/src/engine.rs | 53 +++++- pegainfer-gemma4/src/serve.rs | 29 +-- pegainfer-gemma4/src/serve_oracle.rs | 10 +- .../shared/prefill_attention_hd256_plain.cu | 4 +- pegainfer-kernels/src/ffi/shared.rs | 4 +- pegainfer-kernels/src/ops/attention.rs | 171 ++++++++++-------- pegainfer-kernels/src/paged_kv.rs | 36 +++- 9 files changed, 208 insertions(+), 137 deletions(-) diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 88df17896..f89f5b2a1 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -120,7 +120,7 @@ Measured (a streaming request, then sixteen ~1900-token prompts admitted at once ## 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 at startup, while a disabled one (`unset`, `0`, `off`) is fine. Operators wanting bit-exact serving leave it unset. +`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) diff --git a/pegainfer-core/src/kv_pool.rs b/pegainfer-core/src/kv_pool.rs index 966b99228..edce1602b 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,8 +26,7 @@ pub struct KvLayout { pub layer_stride: usize, /// Elements per page (all layers): num_layers × layer_stride. pub page_stride: usize, - /// Bytes per stored element. Strides remain in elements. - pub elem_bytes: usize, + pub storage: KvStorage, } impl KvLayout { @@ -36,20 +36,22 @@ impl KvLayout { head_dim: usize, page_size: usize, ) -> anyhow::Result { - Self::with_elem_bytes(num_layers, num_kv_heads, head_dim, page_size, 2) + Self::with_storage( + num_layers, + num_kv_heads, + head_dim, + page_size, + KvStorage::Bf16, + ) } - pub fn with_elem_bytes( + pub fn with_storage( num_layers: usize, num_kv_heads: usize, head_dim: usize, page_size: usize, - elem_bytes: usize, + storage: KvStorage, ) -> anyhow::Result { - anyhow::ensure!( - elem_bytes == 1 || elem_bytes == 2, - "paged KV elements are bf16 (2 bytes) or e4m3 (1 byte), not {elem_bytes} bytes" - ); let strides = || -> Option<(usize, usize, usize)> { let kv_block_len = page_size.checked_mul(num_kv_heads)?.checked_mul(head_dim)?; let layer_stride = kv_block_len.checked_mul(2)?; @@ -70,7 +72,7 @@ impl KvLayout { kv_block_len, layer_stride, page_stride, - elem_bytes, + storage, }) } @@ -83,7 +85,7 @@ impl KvLayout { kv_block_len: self.kv_block_len, layer_stride: self.layer_stride, page_stride: self.page_stride, - elem_bytes: self.elem_bytes, + storage: self.storage, } } } @@ -123,29 +125,29 @@ impl KvPool { page_size: usize, num_pages: usize, ) -> Result { - Self::with_elem_bytes( + Self::with_storage( ctx, num_layers, num_kv_heads, head_dim, page_size, num_pages, - 2, + KvStorage::Bf16, ) } #[allow(clippy::too_many_arguments)] - pub fn with_elem_bytes( + pub fn with_storage( ctx: &DeviceContext, num_layers: usize, num_kv_heads: usize, head_dim: usize, page_size: usize, num_pages: usize, - elem_bytes: usize, + storage: KvStorage, ) -> Result { let layout = - KvLayout::with_elem_bytes(num_layers, num_kv_heads, head_dim, page_size, elem_bytes)?; + 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", @@ -155,7 +157,7 @@ impl KvPool { // The allocator multiplies by the element size unchecked; answer // for the byte domain here, before it does. let total_bytes = total_elements - .checked_mul(layout.elem_bytes) + .checked_mul(layout.storage.elem_bytes()) .ok_or_else(|| { anyhow::anyhow!( "KvPool geometry overflows the byte domain: {total_elements} elements" diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 5ff3814cc..582e9031f 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; @@ -169,6 +171,30 @@ fn parse_prefix_cache_cap(raw: &str) -> Result> { } } +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 "] { diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 57a66bc34..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,10 +314,9 @@ fn copy_pool_pages( dst: &[i32], ) -> Result<()> { use cudarc::driver::DevicePtr; - // Startup refuses the prefix cache alongside the fp8 pool, so a one-byte - // layout here means a byte-domain bug, not a configuration. + // Page copies index bf16 elements, so only a bf16 pool may reach them. anyhow::ensure!( - layout.elem_bytes == 2, + layout.storage == KvStorage::Bf16, "pool page copies index bf16 elements; the fp8 pool must not reach them" ); anyhow::ensure!( @@ -655,24 +655,6 @@ pub(crate) fn global_split_factor(config: &Gemma4Config) -> Result { } /// Everything a serving step needs that outlives requests. -fn local_kv_elem_bytes() -> Result { - match std::env::var("PEGAINFER_KV_FP8") { - Err(std::env::VarError::NotPresent) => Ok(2), - Ok(value) if value == "local" => { - // The parsed capacity, not the variable's presence: "0", "off" - // and an empty value all mean the cache is disabled. - anyhow::ensure!( - crate::engine::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(1) - } - Ok(value) => anyhow::bail!("PEGAINFER_KV_FP8 supports only \"local\", got {value:?}"), - Err(err) => anyhow::bail!("PEGAINFER_KV_FP8 is not unicode: {err}"), - } -} - pub(crate) struct GemmaServe { /// The weights these pools, rope tables and layer numbering were built /// for. Holding them is what makes a step's model identity structural: @@ -738,6 +720,7 @@ impl GemmaServe { ctx: &DeviceContext, weights: Gemma4Weights, max_context: usize, + local_kv_storage: KvStorage, local_pages: usize, global_pages: usize, ) -> Result { @@ -768,14 +751,14 @@ impl GemmaServe { } }) .collect(); - let local_pool = KvPool::with_elem_bytes( + let local_pool = KvPool::with_storage( ctx, locals, config.num_key_value_heads, config.head_dim, PAGE_SIZE, local_pages, - local_kv_elem_bytes()?, + 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..9138ef863 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -20,7 +20,15 @@ fn stack_with(max_context: usize, pages: usize) -> (DeviceContext, GemmaServe, S 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, + pegainfer_core::kv_pool::KvStorage::Bf16, + pages, + pages, + ) + .expect("serve"); (ctx, serve, dir) } diff --git a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu b/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu index 95430f8c0..70b1caa76 100644 --- a/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu +++ b/pegainfer-kernels/csrc/shared/prefill_attention_hd256_plain.cu @@ -580,7 +580,7 @@ int qkv_norm_rope_paged_prefill_hd256_plain_cuda( stride_page, stream); } -// fp8 KV pool twin: same math, e4m3 stores at scale 1.0. +// 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, @@ -626,7 +626,7 @@ int qkv_norm_rope_paged_decode_hd256_plain_cuda( stride_page, stream); } -// fp8 KV pool twin: same math, e4m3 stores at scale 1.0. +// 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, diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index 1a96d3f70..08f159707 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/pegainfer-kernels/src/ffi/shared.rs @@ -1217,7 +1217,7 @@ unsafe extern "C" { stream: CUstream, ) -> i32; - /// fp8-KV twin: e4m3 stores at scale 1.0 into a one-byte-element pool. + /// E4m3 KV twin. pub fn qkv_norm_rope_paged_prefill_hd256_plain_fp8kv_cuda( q_batch: *const Half, k_batch: *const Half, @@ -1246,7 +1246,7 @@ unsafe extern "C" { stream: CUstream, ) -> i32; - /// fp8-KV twin: e4m3 stores at scale 1.0 into a one-byte-element pool. + /// E4m3 KV twin. pub fn qkv_norm_rope_paged_decode_hd256_plain_fp8kv_cuda( q_batch: *const Half, k_batch: *const Half, diff --git a/pegainfer-kernels/src/ops/attention.rs b/pegainfer-kernels/src/ops/attention.rs index f7f433901..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; @@ -2561,74 +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 = if layout.elem_bytes == 1 { - #[cfg(feature = "gemma4")] - 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") - } else { - 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(), @@ -2898,15 +2905,15 @@ fn checked_paged_geometry( // `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.elem_bytes == 2, + fp8_capable || layout.storage == KvStorage::Bf16, "{what} has no fp8 KV path; the layout carries {}-byte elements", - layout.elem_bytes + layout.storage.elem_bytes() ); - let pool_len = match layout.elem_bytes { - 1 => pool_len + let pool_len = match layout.storage { + KvStorage::E4m3 => pool_len .checked_mul(2) .ok_or_else(|| anyhow::anyhow!("{what} fp8 pool element count overflows"))?, - _ => pool_len, + KvStorage::Bf16 => pool_len, }; anyhow::ensure!( layout.head_dim == head_dim, @@ -3213,11 +3220,19 @@ 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; - let launch = if layout.elem_bytes == 1 { + // 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 { launch( q_ptr as *const ffi::Half, @@ -3249,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) ); } @@ -3404,11 +3418,19 @@ 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); - let launch = if layout.elem_bytes == 1 { + // 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 { launch( q_ptr as *const ffi::Half, @@ -3441,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) ); } diff --git a/pegainfer-kernels/src/paged_kv.rs b/pegainfer-kernels/src/paged_kv.rs index 7b9f7fb03..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,26 +29,27 @@ pub struct PagedKvLayout { pub layer_stride: usize, /// Elements per page (all layers): num_layers x layer_stride. pub page_stride: usize, - /// Bytes per stored KV element. Strides remain in elements. - pub elem_bytes: 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_elem_bytes(num_layers, num_kv_heads, head_dim, page_size, 2) + Self::with_storage( + num_layers, + num_kv_heads, + head_dim, + page_size, + KvStorage::Bf16, + ) } - pub fn with_elem_bytes( + pub fn with_storage( num_layers: usize, num_kv_heads: usize, head_dim: usize, page_size: usize, - elem_bytes: usize, + storage: KvStorage, ) -> Self { - assert!( - elem_bytes == 1 || elem_bytes == 2, - "paged KV elements are bf16 (2 bytes) or e4m3 (1 byte), not {elem_bytes} bytes" - ); 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; @@ -45,7 +61,7 @@ impl PagedKvLayout { kv_block_len, layer_stride, page_stride, - elem_bytes, + storage, } } } From f7a21a1b29595b7c62ee136539fce10f776e9a07 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 21:36:29 +0100 Subject: [PATCH 05/17] test(kernels): the e4m3 pool's bytes, offsets and window wiring are pinned Signed-off-by: Feathbow --- pegainfer-kernels/tests/hd256_fp8_pool.rs | 256 ++++++++++++++++++++++ 1 file changed, 256 insertions(+) create mode 100644 pegainfer-kernels/tests/hd256_fp8_pool.rs diff --git a/pegainfer-kernels/tests/hd256_fp8_pool.rs b/pegainfer-kernels/tests/hd256_fp8_pool.rs new file mode 100644 index 000000000..1160fe9b4 --- /dev/null +++ b/pegainfer-kernels/tests/hd256_fp8_pool.rs @@ -0,0 +1,256 @@ +//! 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_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 { + bytes + .chunks_exact(2) + .map(|pair| bf16::from_bits(u16::from_le_bytes([pair[0], pair[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); +} + +fn semantic_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 attend(ctx: &DeviceContext, storage: KvStorage) -> Vec { + let layout = PagedKvLayout::with_storage(1, 1, HD, PAGE_SIZE, storage); + let pool = semantic_pool(ctx, storage, 1); + let plan = + PrefillPagedPlan::new_with_cta_tile_q(ctx, &[0], 2, 1, 1, 1, 1, HD, 0).expect("plan"); + let q = constant_states(ctx, 0.0, 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, + -1, + ) + .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), attend(&ctx, KvStorage::Bf16)); +} + +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 = semantic_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] + ); + } +} + +#[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}"); +} From 552de7889de56955bcef24fd679f5ef012900378 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 21:36:29 +0100 Subject: [PATCH 06/17] test(gemma4): the fp8 pool's argmax agreement is gated against the engine's own floor Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 99 +++++++++++++++++++++++++--- scripts/gemma4_gates.sh | 1 + 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index 9138ef863..a5016302b 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,20 +16,20 @@ 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, KvStorage::Bf16) +} + +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, - pegainfer_core::kv_pool::KvStorage::Bf16, - pages, - pages, - ) - .expect("serve"); + let serve = GemmaServe::new(&ctx, weights, max_context, storage, pages, pages).expect("serve"); (ctx, serve, dir) } @@ -362,6 +363,86 @@ fn context_waypoints_match_hf() { gate_waypoints(&ctx, &serve, &long, &long_points); } +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() +} + +#[test] +#[ignore = "requires the pinned 12B checkpoint, fixtures, and a GPU"] +fn fp8_argmax_agreement_meets_the_bf16_floor() { + const SAMPLE_STRIDE: usize = 8; + let bytes = std::fs::read(WINDOW_FIXTURE).expect("read window fixture"); + let fixture = safetensors::SafeTensors::deserialize(&bytes).expect("window fixture"); + let (_, prompt) = u32_tensor(&fixture, "w1023_prompt"); + let sampled: Vec = (0..prompt.len()).step_by(SAMPLE_STRIDE).collect(); + + let (ctx, bf16, _) = stack_with_storage(1024, 66, KvStorage::Bf16); + let bf16_incremental = incremental_argmaxes(&ctx, &bf16, &prompt); + // Recompute is quadratic, so every eighth position defines the sampled floor. + 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(1024, 66, 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!( + "argmax agreement: bf16 incremental/recompute {floor_rate:.6} \ + ({floor_matches}/{samples}), fp8/bf16 incremental {fp8_rate:.6} \ + ({fp8_matches}/{samples})" + ); + assert!( + fp8_matches >= floor_matches, + "fp8/bf16 argmax agreement {fp8_matches}/{samples} is below the bf16 \ + incremental/recompute floor {floor_matches}/{samples}" + ); +} + /// `window_left` masks out-of-window keys whether or not their pages are /// still resident, so releasing them need not change a single generated /// token. diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 268d216fb..f43385e62 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -39,6 +39,7 @@ 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=( From 7e23164633950ff6d1274bf2804293d3f5c43032 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 22:59:14 +0100 Subject: [PATCH 07/17] fix(gemma4): the oracle stack honors the serving fp8 knob Since the knob parse moved into the engine, stack_with pinned Bf16, so PEGAINFER_KV_FP8=local stopped reaching the broad checkpoint gates: the runner fp8 arm silently measured bf16. The stack now asks the engine parsed knob, prints the storage it built, and the runner banner echoes the variable when set. Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 2 +- pegainfer-gemma4/src/serve_oracle.rs | 7 ++++++- scripts/gemma4_gates.sh | 3 +++ 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 582e9031f..a18a386dd 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -171,7 +171,7 @@ fn parse_prefix_cache_cap(raw: &str) -> Result> { } } -fn kv_fp8_storage() -> 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)), diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index a5016302b..2939ca31e 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -16,7 +16,11 @@ 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, KvStorage::Bf16) + stack_with_storage( + max_context, + pages, + crate::engine::kv_fp8_storage().expect("PEGAINFER_KV_FP8"), + ) } fn stack_with_storage( @@ -30,6 +34,7 @@ fn stack_with_storage( 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, storage, pages, pages).expect("serve"); + eprintln!("oracle stack storage: {storage:?}"); (ctx, serve, dir) } diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index f43385e62..cf21a4986 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -177,6 +177,9 @@ 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)" + if [ -n "${PEGAINFER_KV_FP8:-}" ]; then + echo "gemma4 gates: PEGAINFER_KV_FP8=$PEGAINFER_KV_FP8" + fi } require_ckpt() { From 8712399d2a32c50af12459ec51db8ac1a99ce85b Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 22:59:14 +0100 Subject: [PATCH 08/17] test(gemma4): the agreement floor gains a guard and a second window-crossing case The floor gate ran one prompt with no degeneracy guard. The body moves into agreement_case; the gate runs the window-edge prompt at stride 2 and the 4096 prompt truncated to 2048 positions at stride 8, teacher forcing well past the window, and every case demands a non-degenerate floor before comparing fp8 against it. Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 58 +++++++++++++++++++--------- 1 file changed, 40 insertions(+), 18 deletions(-) diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index 2939ca31e..37f2282ec 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -408,25 +408,28 @@ fn agreement(left: &[usize], right: &[usize]) -> usize { left.iter().zip(right).filter(|(a, b)| a == b).count() } -#[test] -#[ignore = "requires the pinned 12B checkpoint, fixtures, and a GPU"] -fn fp8_argmax_agreement_meets_the_bf16_floor() { - const SAMPLE_STRIDE: usize = 8; - let bytes = std::fs::read(WINDOW_FIXTURE).expect("read window fixture"); - let fixture = safetensors::SafeTensors::deserialize(&bytes).expect("window fixture"); - let (_, prompt) = u32_tensor(&fixture, "w1023_prompt"); - let sampled: Vec = (0..prompt.len()).step_by(SAMPLE_STRIDE).collect(); - - let (ctx, bf16, _) = stack_with_storage(1024, 66, KvStorage::Bf16); +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); - // Recompute is quadratic, so every eighth position defines the sampled floor. 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(1024, 66, KvStorage::E4m3); + 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); @@ -437,15 +440,34 @@ fn fp8_argmax_agreement_meets_the_bf16_floor() { let fp8_rate = f64::from(u32::try_from(fp8_matches).expect("match count fits u32")) / samples_f64; eprintln!( - "argmax agreement: bf16 incremental/recompute {floor_rate:.6} \ + "{tensor_name}: argmax agreement: bf16 incremental/recompute {floor_rate:.6} \ ({floor_matches}/{samples}), fp8/bf16 incremental {fp8_rate:.6} \ ({fp8_matches}/{samples})" ); - assert!( - fp8_matches >= floor_matches, - "fp8/bf16 argmax agreement {fp8_matches}/{samples} is below the bf16 \ - incremental/recompute floor {floor_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 From 8222fecd824c64931d74df7bd7d3397037975655 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 22:59:14 +0100 Subject: [PATCH 09/17] test(kernels): varied values give the e4m3 instruments discrimination The read instruments used an all-zero Q over per-page-identical values, so equal scores hid wrong K reads and wrong page picks, and only the unbounded window ran. The pools now carry distinct e4m3-representable values under a non-zero Q (identical post-load floats keep the bf16 comparison bitwise), a finite-window arm proves the mask live, the decode prep store lands its exact bytes, a varied-value geometry probe joins the constant one, and the core pool constructor byte capacity is pinned. Signed-off-by: Feathbow --- pegainfer-core/src/kv_pool.rs | 15 ++ pegainfer-kernels/tests/hd256_fp8_pool.rs | 225 ++++++++++++++++++++-- 2 files changed, 228 insertions(+), 12 deletions(-) diff --git a/pegainfer-core/src/kv_pool.rs b/pegainfer-core/src/kv_pool.rs index edce1602b..6ba1fac84 100644 --- a/pegainfer-core/src/kv_pool.rs +++ b/pegainfer-core/src/kv_pool.rs @@ -402,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-kernels/tests/hd256_fp8_pool.rs b/pegainfer-kernels/tests/hd256_fp8_pool.rs index 1160fe9b4..3c0c1dcb3 100644 --- a/pegainfer-kernels/tests/hd256_fp8_pool.rs +++ b/pegainfer-kernels/tests/hd256_fp8_pool.rs @@ -9,6 +9,7 @@ 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; @@ -21,9 +22,8 @@ const PAGE_SIZE: usize = 2; const NUM_LAYERS: usize = 3; fn packed_fp8(bytes: &[u8]) -> Vec { - bytes - .chunks_exact(2) - .map(|pair| bf16::from_bits(u16::from_le_bytes([pair[0], pair[1]]))) + (0..bytes.len() / 2) + .map(|slot| bf16::from_bits(u16::from_le_bytes([bytes[2 * slot], bytes[2 * slot + 1]]))) .collect() } @@ -85,7 +85,43 @@ fn fp8_prep_stores_exact_bytes_at_layout_offsets() { assert_eq!(got, expected); } -fn semantic_pool(ctx: &DeviceContext, storage: KvStorage, pages: usize) -> CudaSlice { +#[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 { @@ -118,12 +154,95 @@ fn semantic_pool(ctx: &DeviceContext, storage: KvStorage, pages: usize) -> CudaS } } -fn attend(ctx: &DeviceContext, storage: KvStorage) -> Vec { +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 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 = page * PAGE_SIZE * HD * 2 + slot * HD * 2 + element * 2 + block; + let byte = ((5 + seed % 5) << 3) as u8 + | (seed % 8) as u8 + | if seed.is_multiple_of(3) { 0x80 } else { 0 }; + 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 pool = semantic_pool(ctx, storage, 1); - let plan = - PrefillPagedPlan::new_with_cta_tile_q(ctx, &[0], 2, 1, 1, 1, 1, HD, 0).expect("plan"); - let q = constant_states(ctx, 0.0, 1); + 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, @@ -135,7 +254,7 @@ fn attend(ctx: &DeviceContext, storage: KvStorage) -> Vec { &mut output, 1, 1.0, - -1, + window_left, ) .expect("attention"); output @@ -151,13 +270,30 @@ fn fp8_window_read_matches_bf16_for_exact_values() { let Some(ctx) = common::device_or_skip() else { return; }; - assert_eq!(attend(&ctx, KvStorage::E4m3), attend(&ctx, KvStorage::Bf16)); + 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 = semantic_pool(ctx, KvStorage::E4m3, pages); + 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 { @@ -218,6 +354,71 @@ fn fp8_window_read_is_geometry_invariant_for_the_probed_row() { } } +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).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 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 { From 4eab9d900aca065ebe8be435fd6c7fbfefb4c314 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:30:40 +0100 Subject: [PATCH 10/17] test(gemma4): the waypoint gate measures every case before its verdict The top-1 bar asserted inside gate_waypoint, so the first failing case hid every case behind it - under the fp8 pool the w1023 failure silenced the other seven. The shortfall now joins the collected failures beside the tolerance overage, both batches collect before the one verdict, and the per-case line still prints, so a failing run shows the whole table. The structural asserts (fixture positions, shifted coverage, page accounting) stay in place. Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 37 ++++++++++++++++------------ 1 file changed, 21 insertions(+), 16 deletions(-) diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index 37f2282ec..af7004f0e 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -237,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), @@ -251,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); @@ -264,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( @@ -272,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]) { @@ -364,8 +365,12 @@ 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 { From d5b41897a8cdb9c033549302971f37bb7513831f Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:30:40 +0100 Subject: [PATCH 11/17] test(kernels): every page carries a unique sentinel under a permuted table The varied pool byte pattern had period 120 against a page step of 1024, so pages aliased every 15 pages and the 151-page probe target was byte-identical to page 0 - a wrong page-list base passed anyway. Each page now draws a distinct e4m3 code offset (aliasing would need 254 pages), the packed prefix table is reversed so a logical-index confusion changes what is read, and the probe asserts the target page differs from every other page before comparing. Signed-off-by: Feathbow --- pegainfer-kernels/tests/hd256_fp8_pool.rs | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/pegainfer-kernels/tests/hd256_fp8_pool.rs b/pegainfer-kernels/tests/hd256_fp8_pool.rs index 3c0c1dcb3..61b82a2b9 100644 --- a/pegainfer-kernels/tests/hd256_fp8_pool.rs +++ b/pegainfer-kernels/tests/hd256_fp8_pool.rs @@ -165,6 +165,11 @@ fn e4m3_to_f32(byte: u8) -> f32 { } } +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]; @@ -176,10 +181,11 @@ fn varied_values(pages: usize) -> Vec { + block * layout.kv_block_len + slot * HD + element; - let seed = page * PAGE_SIZE * HD * 2 + slot * HD * 2 + element * 2 + block; - let byte = ((5 + seed % 5) << 3) as u8 - | (seed % 8) as u8 + 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); } } @@ -360,7 +366,7 @@ fn varied_geometry_probe(ctx: &DeviceContext, prefix_rows: usize, pages: usize) 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).collect(); + let prefix_pages: Vec = (0..pages as i32 - 1).rev().collect(); ( vec![prefix_pages, vec![pages as i32 - 1]], vec![0, 1], @@ -409,6 +415,13 @@ fn varied_fp8_window_read_is_geometry_invariant_for_the_probed_row() { }; 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) { From a0af53c04e3314c6434fbe61330c8e890b447694 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:30:40 +0100 Subject: [PATCH 12/17] test(gemma4): the fp8 knob joins the serving-knob guard and the pool suite joins the runner SERVING_KNOBS omitted PEGAINFER_KV_FP8, so an operator shell fp8 value could steer lifecycle tests into the fp8 refusals. The runner now owns the kernels integration binary the way it owns everything else: a manifest of its seven tests, a membership check against the binary own listing, and an execution arm that demands the device (PEGAINFER_REQUIRE_GPU=1) so a missing GPU fails instead of skipping. Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 3 ++- scripts/gemma4_gates.sh | 38 +++++++++++++++++++++++++++++++++- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index a18a386dd..96cfe478b 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -2484,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/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index cf21a4986..9ee3a608c 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -90,6 +90,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[@]}" @@ -267,6 +276,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")) @@ -288,6 +304,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) @@ -324,6 +347,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%% *} @@ -364,9 +390,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 @@ -379,10 +412,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") From 9d7655fb91359dd75ca22f98477492267c8973eb Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 00:33:16 +0100 Subject: [PATCH 13/17] test(gemma4): the bit-exact gates pin their bf16 storage contract mixed_step_matches_serial demands bit-exact equality between admission schedules and prefix_restore_matches_cold_path exercises page copies that are bf16-only, but both built their stacks through the env-honoring constructor, so an fp8 arm turned contract statements into permanent failures. Both now pin bf16 and name where the fp8 pool is judged instead; the same-pool invariance gates stay storage-honoring. Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index af7004f0e..b801bec51 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -852,7 +852,8 @@ 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); } @@ -971,7 +972,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; From 233305b65c39024f7640cfd7f1f249462c32ab20 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 00:33:17 +0100 Subject: [PATCH 14/17] test(gemma4): the runner gains an explicit fp8 storage arm Signed-off-by: Feathbow --- scripts/gemma4_gates.sh | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 9ee3a608c..d2b020357 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -109,6 +109,14 @@ 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::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 @@ -132,9 +140,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:-}" ] || 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, @@ -186,6 +212,7 @@ 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 @@ -296,6 +323,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" @@ -360,6 +391,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 From db1711dd035fb85caa3c73505ce1818ccf663a38 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 00:57:37 +0100 Subject: [PATCH 15/17] test(gemma4): the runner treats a present-but-empty knob as present A set-but-empty PEGAINFER_KV_FP8 slipped past the ambient refusal while the production parser refuses it, and a set-but-empty PEGAINFER_GATE_STORAGE silently became bf16 although only unset, bf16 and fp8 are accepted. Presence now checks +x and the default applies only when the variable is truly unset, so an empty value dies. Signed-off-by: Feathbow --- scripts/gemma4_gates.sh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index d2b020357..bb5ade988 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -152,9 +152,9 @@ gate_is_in() { root=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd) cd "$root" || die "cannot enter the repository root" -[ -z "${PEGAINFER_KV_FP8:-}" ] || die \ +[ -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} +gate_storage=${PEGAINFER_GATE_STORAGE-bf16} case "$gate_storage" in bf16) ;; fp8) export PEGAINFER_KV_FP8=local ;; From b3e79479d81883eee4e632f75874b7ec340e36e0 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 00:57:43 +0100 Subject: [PATCH 16/17] test(gemma4): the floor gate measures its own schedule's run-to-run determinism The determinism argument behind the cross-shape floor rested on a serving-probe observation. The gate now replays the bf16 incremental arm on its own prompt and schedule and asserts the replay agrees at every position, so the run-to-run baseline is measured where the floor is used: deterministic, hence degenerate for judging a lossy storage. Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index b801bec51..f6dc9452e 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -428,6 +428,20 @@ fn agreement_case( 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); From bb0c07d6de8ecd32303eed057c87072d535a5885 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 01:02:04 +0100 Subject: [PATCH 17/17] test(gemma4): the fp8 pool walks the live mixed path under structural invariants Pinning the bit-exact gates to bf16 left mixed_prefill_decode_step unexercised under e4m3. A new gate builds the e4m3 stack explicitly and drives the same mixed machinery through the plain and the window-crossing walks, asserting structure instead of bit equality: finite mixed-step logits, every lane reaching exactly its budget, and the window-family page accounting holding right after the mixed step that releases the front. Registered in the admission set and the fp8 profile; no tolerance constants. Signed-off-by: Feathbow --- pegainfer-gemma4/src/serve_oracle.rs | 209 +++++++++++++++++++++++++-- scripts/gemma4_gates.sh | 2 + 2 files changed, 201 insertions(+), 10 deletions(-) diff --git a/pegainfer-gemma4/src/serve_oracle.rs b/pegainfer-gemma4/src/serve_oracle.rs index f6dc9452e..ac6eea6ce 100644 --- a/pegainfer-gemma4/src/serve_oracle.rs +++ b/pegainfer-gemma4/src/serve_oracle.rs @@ -587,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 @@ -629,7 +631,8 @@ fn gate_host_logits(ctx: &DeviceContext, logits: &HiddenStates) -> (usize, Vec, produced: &mut [Vec], @@ -637,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); @@ -671,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, + ); } } @@ -731,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( @@ -816,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) @@ -872,6 +899,168 @@ fn mixed_step_matches_serial() { 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 diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index bb5ade988..d54f4d1e6 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -44,6 +44,7 @@ GATES_NUMERIC_PARITY=( ) 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" ) @@ -113,6 +114,7 @@ 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"