From 4c49ea005e00de703a8c3b14995f81fb94a582eb Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 24 Aug 2026 22:10:08 +0100 Subject: [PATCH 01/14] perf(sample): the token read spins ahead of the parking wait Signed-off-by: Feathbow (cherry picked from commit 5c8c6c14165d03c90526a9e8886c25f07ae1581d) Signed-off-by: Feathbow --- pegainfer-kernels/src/tensor.rs | 28 ++++++++++++++++++++++++++++ pegainfer-sample/src/lib.rs | 5 ++++- 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/pegainfer-kernels/src/tensor.rs b/pegainfer-kernels/src/tensor.rs index 74d39ddf6..bbecd6f31 100644 --- a/pegainfer-kernels/src/tensor.rs +++ b/pegainfer-kernels/src/tensor.rs @@ -79,6 +79,34 @@ pub fn active_cu_stream(ctx: &DeviceContext) -> CUstream { .unwrap_or_else(|| ctx.stream.cu_stream()) } +/// Bounded spin wait on the context stream. A blocking wait parks the caller +/// and pays the OS scheduler's wake-up — low milliseconds per call on a +/// loaded host — which a decode loop pays every step right before it reads +/// the sampled tokens. Polling `cuStreamQuery` keeps the wake on our side; +/// the cap falls back to the parking wait so a stuck stream still sleeps +/// instead of burning the core forever. +const STREAM_SPIN_WAIT_CAP: std::time::Duration = std::time::Duration::from_millis(5); + +pub fn stream_spin_wait(ctx: &DeviceContext) -> anyhow::Result<()> { + let stream = active_cu_stream(ctx); + let cap = std::time::Instant::now() + STREAM_SPIN_WAIT_CAP; + loop { + match unsafe { cudarc::driver::sys::cuStreamQuery(stream) } { + cudarc::driver::sys::CUresult::CUDA_SUCCESS => return Ok(()), + cudarc::driver::sys::CUresult::CUDA_ERROR_NOT_READY => { + if std::time::Instant::now() >= cap { + return ctx + .stream + .synchronize() + .map_err(|e| anyhow::anyhow!("stream sync after spin cap failed: {e}")); + } + std::hint::spin_loop(); + } + err => return Err(anyhow::anyhow!("cuStreamQuery failed: {err:?}")), + } + } +} + /// Marker trait for tensor metadata tags. pub trait NamedTag { const NAME: &'static str; diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs index 6adcdd391..52a5bb1a2 100644 --- a/pegainfer-sample/src/lib.rs +++ b/pegainfer-sample/src/lib.rs @@ -49,6 +49,7 @@ use pegainfer_kernels::ops::logprob_topk_batch_bf16_into; use pegainfer_kernels::tensor::DeviceContext; use pegainfer_kernels::tensor::HiddenStates; use pegainfer_kernels::tensor::has_stream_override; +use pegainfer_kernels::tensor::stream_spin_wait; /// Allocate-once device buffers for [`select_batch`], sized for `max_rows` × `vocab`. /// @@ -207,7 +208,9 @@ pub fn select_batch( .map_err(|e| anyhow!("select_batch D2H greedy tokens failed: {e}"))?; // Blocks on this copy's own event — which transitively covers the // argmax kernel queued before it on the same stream, so the wait is - // equivalent to the old full-stream sync for this path. + // equivalent to the old full-stream sync for this path. The spin + // ahead of it takes the scheduler wake-up out of the step loop. + stream_spin_wait(ctx)?; let out = scratch .argmax_host .as_slice() From 297c550b8117c35899e2507b64d2206d90874d4c Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 11:55:15 +0100 Subject: [PATCH 02/14] perf(gemma4): decode rounds pipeline the greedy sample turnaround Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 217 ++++++++++++++++++++++++++++++-- pegainfer-gemma4/src/serve.rs | 52 ++++++-- pegainfer-kernels/src/tensor.rs | 37 ++++++ pegainfer-sample/src/lib.rs | 105 ++++++++++++++++ 4 files changed, 390 insertions(+), 21 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index dfdb147f1..d644f0a21 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -697,6 +697,9 @@ struct Active { next: u32, emitted: usize, prompt_tokens: usize, + /// The row has finished while a speculative step over its slot is still + /// in flight and retires when that step drains. + stopping: bool, } impl Active { @@ -708,6 +711,52 @@ impl Active { ignore_eos: self.request.params.ignore_eos, } } + + fn settle_staged(&mut self, policy: &GenerationPolicy, token: u32) { + if self.stopping { + return; + } + if policy.stops(token, self.request.params.ignore_eos) { + let _ = self.request.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Stop, + prompt_tokens: self.prompt_tokens, + completion_tokens: self.emitted, + }); + self.stopping = true; + return; + } + self.emitted += 1; + if self + .request + .token_tx + .send(TokenEvent::Token { + id: token, + logprob: None, + }) + .is_err() + { + self.stopping = true; + return; + } + if self.emitted >= self.request.max_tokens { + let _ = self.request.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: self.prompt_tokens, + completion_tokens: self.emitted, + }); + self.stopping = true; + return; + } + self.next = token; + } +} + +const DECODE_PIPELINE_DEPTH: usize = 2; + +/// One staged decode step whose readback has not yet been collected. +struct PendingDecode { + rows: usize, + slot: usize, } enum Admitted { @@ -751,6 +800,8 @@ struct EngineState { /// mixed into the per-call seed; a request's own `params.seed` replays /// via (seed, step) regardless of it. sample_nonce: u64, + /// Present only while the active row order is frozen. + pipeline: Option, /// The overlap lane; `None` unless `PEGAINFER_ASYNC_PREFILL` opted in /// at startup. lane: Option, @@ -922,6 +973,7 @@ impl EngineState { suppress_ids, base_seed, sample_nonce: 0, + pipeline: None, lane, mix_chunk, max_context, @@ -935,6 +987,10 @@ impl EngineState { /// costs the streams a bounded number of prefills per token however /// deep the queue is. fn admit_from_queue(&mut self, pending: &mut VecDeque, active: &mut Vec) { + if pending.is_empty() { + return; + } + self.drain_pipeline(active); let mut attempts = 0; while attempts < self.slots && active.len() < self.slots { // With the lane busy, arrivals wait in `pending` while decode @@ -1302,6 +1358,13 @@ impl EngineState { /// release, capture into the prefix cache, and take the first-token /// flow the sync path uses. fn join_async_prefill(&mut self, active: &mut Vec) { + if self + .lane + .as_ref() + .is_some_and(|lane| lane.inflight.is_some()) + { + self.drain_pipeline(active); + } let Some(lane) = self.lane.as_mut() else { return; }; @@ -1675,6 +1738,101 @@ impl EngineState { } } + fn pipeline_eligible(&self, active: &[Active]) -> bool { + !active.is_empty() + && active.len() <= self.scratch.max_rows() + && active.iter().all(|entry| { + !entry.stopping + && entry.request.logprobs == 0 + && entry.request.max_tokens.saturating_sub(entry.emitted) + >= DECODE_PIPELINE_DEPTH + && pegainfer_sample::effectively_greedy( + &entry.request.params, + self.scratch.vocab(), + ) + }) + } + + /// Reserve the next token without retiring or reordering a row. + fn ready_rows_pinned(&self, active: &mut [Active]) -> bool { + active.iter_mut().all(|entry| { + !entry.request.token_tx.is_closed() + && admit_tokens( + &self.serve.local_pool, + &self.serve.global_pool, + &mut entry.kv, + 1, + ) + .is_ok() + }) + } + + /// Queue one decode and stage its greedy picks into the next embedding's + /// id buffer and one pinned readback slot. + fn launch_staged( + &mut self, + active: &mut [Active], + resident: bool, + slot: usize, + ) -> Result { + let rows = active.len(); + let tokens = (!resident).then(|| active.iter().map(|entry| entry.next).collect::>()); + { + let mut kvs: Vec<&mut GemmaKv> = active.iter_mut().map(|entry| &mut entry.kv).collect(); + if let Some(tokens) = tokens.as_deref() { + self.serve + .decode_batch_step(&self.ctx, &mut self.arena, &mut kvs, tokens)?; + } else { + self.serve + .decode_batch_step_resident(&self.ctx, &mut self.arena, &mut kvs)?; + } + } + let (logits, ids) = self.arena.logits_and_ids(); + ops::suppress_logits_bf16_in_place(&self.ctx, logits, &self.suppress_ids) + .context("suppression")?; + self.sample_nonce = self.sample_nonce.wrapping_add(1); + pegainfer_sample::greedy_stage_resident( + &self.ctx, + logits, + rows, + ids, + slot, + &mut self.scratch, + ) + .context("stage greedy picks")?; + Ok(rows) + } + + /// Deliver a staged step without changing the row order. Finished rows + /// stay pinned until the speculative successor drains. + fn collect_pending(&mut self, active: &mut [Active], pending: &PendingDecode) -> Result<()> { + anyhow::ensure!( + pending.rows == active.len(), + "pipeline collected {} rows against a batch of {}", + pending.rows, + active.len() + ); + let picked = pegainfer_sample::greedy_collect_resident( + pending.rows, + pending.slot, + &mut self.scratch, + )?; + for (entry, token) in active.iter_mut().zip(picked) { + entry.settle_staged(&self.policy, token); + } + Ok(()) + } + + fn drain_pipeline(&mut self, active: &mut Vec) { + let Some(pending) = self.pipeline.take() else { + return; + }; + if let Err(err) = self.collect_pending(active, &pending) { + return fail_active_batch(active, "pipelined decode drain", &err); + } + active.retain(|entry| !entry.stopping); + } + /// The mixed-admission tail of [`Self::admit_and_prefill`]: every /// gathered prompt and the live decode batch share one step, then one /// sampler call covers the newcomers' first tokens (logits rows `0..k`) @@ -1776,16 +1934,7 @@ impl EngineState { Admitted::Done } - /// One batched decode step: every active request advances a token, - /// sharing each layer's weight pass. A cancelled request and a request - /// the pools cannot grow for retire before the batch is built; a finished - /// one retires after its token lands. - fn decode_round(&mut self, active: &mut Vec) { - self.ready_decode_rows(active); - if active.is_empty() { - return; - } - + fn decode_round_collect(&mut self, active: &mut Vec) { let tokens: Vec = active.iter().map(|entry| entry.next).collect(); let logits = { let mut kvs: Vec<&mut GemmaKv> = active.iter_mut().map(|entry| &mut entry.kv).collect(); @@ -1815,6 +1964,53 @@ impl EngineState { Err(err) => fail_active_batch(active, "batched decode", &err), } } + + /// One batched decode step. Eligible greedy batches keep one speculative + /// successor in flight and drain before any row-order change. + fn decode_round(&mut self, active: &mut Vec) { + if let Some(pending) = self.pipeline.take() { + if self.pipeline_eligible(active) && self.ready_rows_pinned(active) { + let next_slot = (pending.slot + 1) % DECODE_PIPELINE_DEPTH; + match self.launch_staged(active, true, next_slot) { + Ok(rows) => { + if let Err(err) = self.collect_pending(active, &pending) { + fail_active_batch(active, "pipelined decode collect", &err); + return; + } + self.pipeline = Some(PendingDecode { + rows, + slot: next_slot, + }); + } + Err(err) => { + if let Err(collect_err) = self.collect_pending(active, &pending) { + log::error!("collect during launch failure failed: {collect_err:#}"); + } + fail_active_batch(active, "pipelined decode launch", &err); + } + } + return; + } + self.pipeline = Some(pending); + self.drain_pipeline(active); + if active.is_empty() { + return; + } + } + + self.ready_decode_rows(active); + if active.is_empty() { + return; + } + if self.pipeline_eligible(active) { + match self.launch_staged(active, false, 0) { + Ok(rows) => self.pipeline = Some(PendingDecode { rows, slot: 0 }), + Err(err) => fail_active_batch(active, "batched decode", &err), + } + return; + } + self.decode_round_collect(active); + } } /// Deliver one decode step's outcome to every active row and retire the @@ -1903,6 +2099,7 @@ fn settle_first_token( next, emitted: 1, prompt_tokens, + stopping: false, }) } diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 1cf70cad5..868c52da1 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -575,6 +575,11 @@ impl StepArena { ); self.tower.open(rows) } + + /// Logits and the id buffer consumed by the next decode embedding. + pub(crate) fn logits_and_ids(&mut self) -> (&mut HiddenStates, &mut CudaSlice) { + (&mut self.logits, &mut self.ids) + } } /// How many pseudo-requests the global decode read presents each request @@ -1811,6 +1816,27 @@ impl GemmaServe { arena: &'a mut StepArena, kvs: &mut [&mut GemmaKv], tokens: &[u32], + ) -> Result<&'a mut HiddenStates> { + self.decode_batch_step_inner(ctx, arena, kvs, Some(tokens)) + } + + /// Decode using token ids written into the arena by the prior staged + /// sampler. The caller keeps the row order and bucket unchanged. + pub(crate) fn decode_batch_step_resident<'a>( + &self, + ctx: &DeviceContext, + arena: &'a mut StepArena, + kvs: &mut [&mut GemmaKv], + ) -> Result<&'a mut HiddenStates> { + self.decode_batch_step_inner(ctx, arena, kvs, None) + } + + fn decode_batch_step_inner<'a>( + &self, + ctx: &DeviceContext, + arena: &'a mut StepArena, + kvs: &mut [&mut GemmaKv], + tokens: Option<&[u32]>, ) -> Result<&'a mut HiddenStates> { let batch = kvs.len(); let padded = self.prepare_decode_step(ctx, arena, kvs, tokens)?; @@ -2165,22 +2191,24 @@ impl GemmaServe { ctx: &DeviceContext, arena: &mut StepArena, kvs: &[&mut GemmaKv], - tokens: &[u32], + tokens: Option<&[u32]>, ) -> Result { let batch = kvs.len(); anyhow::ensure!(batch > 0, "a decode batch needs at least one request"); - anyhow::ensure!( - tokens.len() == batch, - "decode batch has {batch} requests but {} tokens", - tokens.len() - ); anyhow::ensure!( batch <= arena.max_rows, "decode batch of {batch} exceeds the arena's {} row ceiling", arena.max_rows ); self.check_stream(ctx)?; - validate_tokens(&self.weights, self.local_geom.hidden_size, tokens)?; + if let Some(tokens) = tokens { + anyhow::ensure!( + tokens.len() == batch, + "decode batch has {batch} requests but {} tokens", + tokens.len() + ); + validate_tokens(&self.weights, self.local_geom.hidden_size, tokens)?; + } // A step computes at its power-of-two bucket whether or not graphs // are on: padding is part of the numeric contract, which is what @@ -2189,9 +2217,11 @@ impl GemmaServe { let padded = batch.next_power_of_two().max(arena.min_bucket); arena.open(ctx, padded)?; self.plan_decode_batch(ctx, arena, kvs, padded)?; - arena.host.ids.extend_from_slice(tokens); - arena.host.ids.resize(padded, 0); - upload_prefix(ctx, &mut arena.ids, &arena.host.ids)?; + if let Some(tokens) = tokens { + arena.host.ids.extend_from_slice(tokens); + arena.host.ids.resize(padded, 0); + upload_prefix(ctx, &mut arena.ids, &arena.host.ids)?; + } Ok(padded) } @@ -2218,7 +2248,7 @@ impl GemmaServe { admit_tokens(&self.local_pool, &self.global_pool, &mut kv, 1)?; { let mut kvs: [&mut GemmaKv; 1] = [&mut kv]; - let padded = self.prepare_decode_step(ctx, arena, &kvs, &[0])?; + let padded = self.prepare_decode_step(ctx, arena, &kvs, Some(&[0]))?; let StepArena { tower, local_plan, diff --git a/pegainfer-kernels/src/tensor.rs b/pegainfer-kernels/src/tensor.rs index bbecd6f31..3288f756b 100644 --- a/pegainfer-kernels/src/tensor.rs +++ b/pegainfer-kernels/src/tensor.rs @@ -107,6 +107,43 @@ pub fn stream_spin_wait(ctx: &DeviceContext) -> anyhow::Result<()> { } } +/// Copy token ids from an i32 argmax buffer into a u32 embedding buffer. +pub fn memcpy_dtod_u32_from_i32( + ctx: &DeviceContext, + src: &CudaSlice, + dst: &mut CudaSlice, + count: usize, +) -> anyhow::Result<()> { + use cudarc::driver::DevicePtr; + use cudarc::driver::DevicePtrMut; + + anyhow::ensure!( + count <= src.len() && count <= dst.len(), + "dtod i32->u32 copy of {count} exceeds src {} or dst {}", + src.len(), + dst.len() + ); + if count == 0 { + return Ok(()); + } + let stream = active_cu_stream(ctx); + let (src_ptr, _src_guard) = src.device_ptr(&ctx.stream); + let (dst_ptr, _dst_guard) = dst.device_ptr_mut(&ctx.stream); + let result = unsafe { + cudarc::driver::sys::cuMemcpyDtoDAsync_v2( + dst_ptr, + src_ptr, + count * std::mem::size_of::(), + stream, + ) + }; + anyhow::ensure!( + result == cudarc::driver::sys::CUresult::CUDA_SUCCESS, + "cuMemcpyDtoDAsync i32->u32 failed: {result:?}" + ); + Ok(()) +} + /// Marker trait for tensor metadata tags. pub trait NamedTag { const NAME: &'static str; diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs index 52a5bb1a2..add3f78ef 100644 --- a/pegainfer-sample/src/lib.rs +++ b/pegainfer-sample/src/lib.rs @@ -51,6 +51,8 @@ use pegainfer_kernels::tensor::HiddenStates; use pegainfer_kernels::tensor::has_stream_override; use pegainfer_kernels::tensor::stream_spin_wait; +const STAGED_READBACK_SLOTS: usize = 2; + /// Allocate-once device buffers for [`select_batch`], sized for `max_rows` × `vocab`. /// /// Reused across decode steps — the decode path needs pointer-stable buffers, so @@ -73,6 +75,10 @@ pub struct SampleScratch { /// active stream (#704). Pinned keeps the D2H async; the reader blocks /// only on the copy's own event. argmax_host: PinnedHostSlice, + /// Alternate landing slot used while the prior readback is collected. + argmax_host_alt: PinnedHostSlice, + /// Identity row map for the all-greedy staged path. + identity_rows: CudaSlice, sampling: BatchSamplingScratch, /// Vocab width every buffer above was sized for; `select_batch` rejects a /// logits arena whose `hidden_dim` differs, since the sizes are baked in. @@ -110,6 +116,12 @@ impl SampleScratch { // but don't grow this buffer into anything read in a hot loop. argmax_host: unsafe { ctx.ctx.alloc_pinned::(max_rows) } .map_err(|e| anyhow!("SampleScratch pinned alloc failed: {e}"))?, + argmax_host_alt: unsafe { ctx.ctx.alloc_pinned::(max_rows) } + .map_err(|e| anyhow!("SampleScratch pinned alloc failed: {e}"))?, + identity_rows: ctx + .stream + .clone_htod(&(0..max_rows as i32).collect::>()) + .map_err(|e| anyhow!("SampleScratch identity upload failed: {e}"))?, sampling: BatchSamplingScratch::new(ctx, max_rows, vocab)?, vocab, max_rows, @@ -119,6 +131,10 @@ impl SampleScratch { pub fn max_rows(&self) -> usize { self.max_rows } + + pub fn vocab(&self) -> usize { + self.vocab + } } /// Pick the next token for every row of a logits arena. @@ -278,6 +294,95 @@ pub fn select_batch( Ok(tokens) } +fn validate_greedy_stage( + logits: &HiddenStates, + rows: usize, + slot: usize, + scratch: &SampleScratch, +) -> Result<()> { + ensure!(rows > 0, "greedy_stage_resident: empty batch"); + ensure!( + rows <= scratch.max_rows, + "greedy_stage_resident: {rows} rows exceeds scratch capacity {}", + scratch.max_rows + ); + ensure!( + logits.seq_len >= rows && logits.hidden_dim == scratch.vocab, + "greedy_stage_resident: logits shape {}x{} cannot serve {rows} rows x vocab {}", + logits.seq_len, + logits.hidden_dim, + scratch.vocab + ); + ensure!( + slot < STAGED_READBACK_SLOTS, + "greedy_stage_resident: slot {slot} out of range" + ) +} + +fn stage_greedy_readback( + ctx: &DeviceContext, + slot: usize, + scratch: &mut SampleScratch, +) -> Result<()> { + let host = if slot == 0 { + &mut scratch.argmax_host + } else { + &mut scratch.argmax_host_alt + }; + ctx.stream + .memcpy_dtoh(&scratch.argmax_out, host) + .map_err(|e| anyhow!("greedy_stage_resident D2H stage failed: {e}")) +} + +/// Argmax every row, leave the picks in the next embedding's id buffer, and +/// queue their pinned readback without waiting for it. +pub fn greedy_stage_resident( + ctx: &DeviceContext, + logits: &HiddenStates, + rows: usize, + ids_out: &mut CudaSlice, + slot: usize, + scratch: &mut SampleScratch, +) -> Result<()> { + validate_greedy_stage(logits, rows, slot, scratch)?; + argmax_batch_bf16_split_indexed_into( + ctx, + logits, + &scratch.identity_rows, + rows, + &mut scratch.argmax_partial_values, + &mut scratch.argmax_partial_indices, + &mut scratch.top1_values, + &mut scratch.argmax_out, + )?; + pegainfer_kernels::tensor::memcpy_dtod_u32_from_i32(ctx, &scratch.argmax_out, ids_out, rows)?; + stage_greedy_readback(ctx, slot, scratch) +} + +/// Collect a readback staged into one of the two pinned slots. +pub fn greedy_collect_resident( + rows: usize, + slot: usize, + scratch: &mut SampleScratch, +) -> Result> { + ensure!( + slot < STAGED_READBACK_SLOTS, + "greedy_collect_resident: slot {slot} out of range" + ); + ensure!( + rows <= scratch.max_rows, + "greedy_collect_resident: {rows} rows exceeds scratch capacity {}", + scratch.max_rows + ); + let landed = if slot == 0 { + scratch.argmax_host.as_slice() + } else { + scratch.argmax_host_alt.as_slice() + } + .map_err(|e| anyhow!("greedy_collect_resident D2H sync failed: {e}"))?; + Ok(landed[..rows].iter().map(|&token| token as u32).collect()) +} + /// SplitMix64 over (seed, step): a distinct, well-mixed philox seed per /// request step, deterministic across runs and batch layouts. Public for the /// models that drive their own greedy path (see the module docs) and must From a17fca9ff75200d4a96f31e5113fa1c0d92b126e Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 12:02:39 +0100 Subject: [PATCH 03/14] perf(gemma4): regular decode steps upload nothing Signed-off-by: Feathbow --- pegainfer-core/src/ops.rs | 1 + pegainfer-gemma4/src/engine.rs | 2 + pegainfer-gemma4/src/serve.rs | 92 ++++++++++++++++++++ pegainfer-kernels/csrc/shared/elementwise.cu | 31 +++++++ pegainfer-kernels/src/ffi/shared.rs | 10 +++ pegainfer-kernels/src/ops.rs | 1 + pegainfer-kernels/src/ops/elementwise.rs | 36 ++++++++ 7 files changed, 173 insertions(+) diff --git a/pegainfer-core/src/ops.rs b/pegainfer-core/src/ops.rs index 8f167d55e..a0a683663 100644 --- a/pegainfer-core/src/ops.rs +++ b/pegainfer-core/src/ops.rs @@ -24,6 +24,7 @@ pub use pegainfer_kernels::ops::SuppressIds; pub use pegainfer_kernels::ops::accumulate_bf16_token_scaled_to_f32_into; pub use pegainfer_kernels::ops::add_batch; pub use pegainfer_kernels::ops::add_batch_into; +pub use pegainfer_kernels::ops::advance_decode_metadata; pub use pegainfer_kernels::ops::argmax; pub use pegainfer_kernels::ops::argmax_batch_bf16_into; pub use pegainfer_kernels::ops::batch_prefill_paged_hd512_into; diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index d644f0a21..2dfec3f8f 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -990,6 +990,7 @@ impl EngineState { if pending.is_empty() { return; } + self.arena.invalidate_decode_fingerprint(); self.drain_pipeline(active); let mut attempts = 0; while attempts < self.slots && active.len() < self.slots { @@ -1363,6 +1364,7 @@ impl EngineState { .as_ref() .is_some_and(|lane| lane.inflight.is_some()) { + self.arena.invalidate_decode_fingerprint(); self.drain_pipeline(active); } let Some(lane) = self.lane.as_mut() else { diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 868c52da1..f1c5f8e46 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -483,6 +483,34 @@ fn hidden_pair(hidden: &mut [HiddenStates; 2], src: usize) -> (&HiddenStates, &m const GLOBAL_SPLIT_CHUNK_TOKENS: usize = 256; +struct SteadyDecode { + padded: usize, + rows: Vec, +} + +#[derive(Eq, PartialEq)] +struct SteadyRow { + kv_len: usize, + local_origin: usize, + local_pages: usize, + global_pages: usize, + global_chunks: usize, +} + +impl SteadyDecode { + fn advances_to(&self, next: &Self) -> bool { + self.padded == next.padded + && self.rows.len() == next.rows.len() + && self.rows.iter().zip(&next.rows).all(|(current, next)| { + next.kv_len == current.kv_len + 1 + && next.local_origin == current.local_origin + && next.local_pages == current.local_pages + && next.global_pages == current.global_pages + && next.global_chunks == current.global_chunks + }) + } +} + /// The global family's decode tables, uploaded per step: the per-request /// half feeds the prep, the factor-repeated half feeds the split-KV /// attention read over the pseudo-requests (see [`global_split_factor`]). @@ -529,6 +557,7 @@ pub(crate) struct StepArena { local_plan: PrefillPagedPlan, global_tables: GlobalTables, global_split: SplitKvState, + steady: Option, local_origins: CudaSlice, ids: CudaSlice, /// Mixed-step per-row prep metadata at step-stable pointers, covering @@ -580,6 +609,10 @@ impl StepArena { pub(crate) fn logits_and_ids(&mut self) -> (&mut HiddenStates, &mut CudaSlice) { (&mut self.logits, &mut self.ids) } + + pub(crate) fn invalidate_decode_fingerprint(&mut self) { + self.steady = None; + } } /// How many pseudo-requests the global decode read presents each request @@ -876,6 +909,7 @@ impl GemmaServe { .map_err(alloc("global split tmp_s"))?, cap: global_split_cap, }, + steady: None, local_origins: ctx.stream.alloc_zeros(max_rows).map_err(alloc("origins"))?, ids: ctx.stream.alloc_zeros(max_rows).map_err(alloc("ids"))?, mix_positions: ctx @@ -1648,6 +1682,38 @@ impl GemmaServe { Ok(()) } + fn decode_fingerprint(&self, kvs: &[&mut GemmaKv], padded: usize) -> Option { + let local_page = self.local_pool.layout().page_size; + let global_page = self.global_pool.layout().page_size; + let rows = kvs + .iter() + .map(|kv| { + let kv = &**kv; + let kv_len = kv.local.seq_len().checked_add(1)?; + let origin = kv.local.origin_pages(); + let resident_len = kv_len.checked_sub(origin.checked_mul(local_page)?)?; + if resident_len == 0 { + return None; + } + let local_pages = kv.local.held_pages(); + let global_pages = kv.global.held_pages(); + if local_pages != resident_len.div_ceil(local_page) + || global_pages != kv_len.div_ceil(global_page) + { + return None; + } + Some(SteadyRow { + kv_len, + local_origin: origin, + local_pages, + global_pages, + global_chunks: kv_len.div_ceil(GLOBAL_SPLIT_CHUNK_TOKENS), + }) + }) + .collect::>>()?; + Some(SteadyDecode { padded, rows }) + } + fn plan_decode_batch( &self, ctx: &DeviceContext, @@ -1656,6 +1722,20 @@ impl GemmaServe { padded: usize, ) -> Result<()> { let batch = kvs.len(); + let fresh = self.decode_fingerprint(kvs, padded); + let regular = arena + .steady + .as_ref() + .zip(fresh.as_ref()) + .is_some_and(|(current, next)| current.advances_to(next)); + if regular { + for kv in kvs { + self.check_step_bounds(kv, kv.local.seq_len() + 1)?; + } + arena.steady = fresh; + return Ok(()); + } + arena.steady = fresh; let StepArena { host, local_plan, @@ -1923,6 +2003,15 @@ impl GemmaServe { self.final_logit_softcapping, head_normed, logits, + )?; + ops::advance_decode_metadata( + ctx, + &global_tables.positions, + local_plan.last_page_len_d(), + &global_tables.pseudo_last, + local_plan.kv_chunk_size_d(), + rows, + self.global_split_factor, ) } @@ -1961,6 +2050,7 @@ impl GemmaServe { arena.max_rows ); self.check_stream(ctx)?; + arena.steady = None; for (_, prompt) in prefills.iter() { validate_tokens(&self.weights, self.local_geom.hidden_size, prompt)?; } @@ -2218,6 +2308,7 @@ impl GemmaServe { arena.open(ctx, padded)?; self.plan_decode_batch(ctx, arena, kvs, padded)?; if let Some(tokens) = tokens { + arena.host.ids.clear(); arena.host.ids.extend_from_slice(tokens); arena.host.ids.resize(padded, 0); upload_prefix(ctx, &mut arena.ids, &arena.host.ids)?; @@ -2245,6 +2336,7 @@ impl GemmaServe { let mut bucket = 1usize; while bucket <= arena.max_rows { arena.min_bucket = bucket; + arena.steady = None; admit_tokens(&self.local_pool, &self.global_pool, &mut kv, 1)?; { let mut kvs: [&mut GemmaKv; 1] = [&mut kv]; diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index 4b60e619f..97bc8fa83 100644 --- a/pegainfer-kernels/csrc/shared/elementwise.cu +++ b/pegainfer-kernels/csrc/shared/elementwise.cu @@ -518,6 +518,27 @@ __global__ void embedding_batched_vocab_shard_vec4_kernel( } } +// Advances the device tables a regular decode step will consume next. +static const int ADVANCE_DECODE_METADATA_BLOCK = 32; + +__global__ void advance_decode_metadata_kernel( + int *__restrict__ positions, + int *__restrict__ local_last, + int *__restrict__ pseudo_last, + int *__restrict__ kv_chunk, + int rows, + int factor) { + int row = blockIdx.x * blockDim.x + threadIdx.x; + if (row < rows) { + positions[row] += 1; + local_last[row] += 1; + kv_chunk[row] += 1; + for (int copy = 0; copy < factor; ++copy) { + pseudo_last[row * factor + copy] += 1; + } + } +} + extern "C" { CUresult add_cuda( @@ -904,4 +925,14 @@ CUresult embedding_batched_vocab_shard_cuda( return (CUresult)cudaGetLastError(); } +CUresult advance_decode_metadata_cuda( + int *positions, int *local_last, int *pseudo_last, int *kv_chunk, + int rows, int factor, cudaStream_t stream) { + int block = ADVANCE_DECODE_METADATA_BLOCK; + int grid = (rows + block - 1) / block; + advance_decode_metadata_kernel<<>>( + positions, local_last, pseudo_last, kv_chunk, rows, factor); + return (CUresult)cudaGetLastError(); +} + } // extern "C" diff --git a/pegainfer-kernels/src/ffi/shared.rs b/pegainfer-kernels/src/ffi/shared.rs index 651b7b377..b5757d639 100644 --- a/pegainfer-kernels/src/ffi/shared.rs +++ b/pegainfer-kernels/src/ffi/shared.rs @@ -82,6 +82,16 @@ unsafe extern "C" { stream: CUstream, ) -> CUresult; + pub fn advance_decode_metadata_cuda( + positions: *mut i32, + local_last: *mut i32, + pseudo_last: *mut i32, + kv_chunk: *mut i32, + rows: i32, + factor: i32, + stream: CUstream, + ) -> CUresult; + pub fn add_scaled_bf16_cuda( routed: *const Half, scale: f32, diff --git a/pegainfer-kernels/src/ops.rs b/pegainfer-kernels/src/ops.rs index f48b76f1c..085e9e53f 100644 --- a/pegainfer-kernels/src/ops.rs +++ b/pegainfer-kernels/src/ops.rs @@ -87,6 +87,7 @@ pub use elementwise::add_batch; pub use elementwise::add_batch_into; pub use elementwise::add_into; pub use elementwise::add_scaled_bf16_into; +pub use elementwise::advance_decode_metadata; pub use elementwise::bf16_bytes_to_f32_into; pub use elementwise::bf16_hidden_to_f32_into; pub use elementwise::copy_hidden_rows_into; diff --git a/pegainfer-kernels/src/ops/elementwise.rs b/pegainfer-kernels/src/ops/elementwise.rs index 8e9f38a1f..c13a8b951 100644 --- a/pegainfer-kernels/src/ops/elementwise.rs +++ b/pegainfer-kernels/src/ops/elementwise.rs @@ -49,6 +49,42 @@ pub fn add_batch_into( Ok(()) } +/// Advance the per-row decode tables written by a regular graph replay. +pub fn advance_decode_metadata( + ctx: &DeviceContext, + positions: &CudaSlice, + local_last: &CudaSlice, + pseudo_last: &CudaSlice, + kv_chunk: &CudaSlice, + rows: usize, + factor: usize, +) -> Result<()> { + anyhow::ensure!( + positions.len() >= rows + && local_last.len() >= rows + && kv_chunk.len() >= rows + && pseudo_last.len() >= rows * factor, + "advance_decode_metadata: {rows} rows x {factor} exceeds a table" + ); + let (positions_ptr, _positions_guard) = positions.device_ptr(&ctx.stream); + let (local_ptr, _local_guard) = local_last.device_ptr(&ctx.stream); + let (pseudo_ptr, _pseudo_guard) = pseudo_last.device_ptr(&ctx.stream); + let (chunk_ptr, _chunk_guard) = kv_chunk.device_ptr(&ctx.stream); + let result = unsafe { + ffi::advance_decode_metadata_cuda( + positions_ptr as *mut i32, + local_ptr as *mut i32, + pseudo_ptr as *mut i32, + chunk_ptr as *mut i32, + rows as i32, + factor as i32, + crate::tensor::active_cu_stream(ctx), + ) + }; + result.result()?; + Ok(()) +} + /// Element-wise add of `n` bf16 elements into a pre-allocated output /// (`out = a + b`). Slice-level twin of [`add_batch_into`] — same kernel — /// for callers whose buffers live in a persistent decode arena rather than From 7772c3804c4bd6bef12507cb65ed8dce11848113 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 12:05:18 +0100 Subject: [PATCH 04/14] perf(gemma4): the staged sampler chain rides two graph launches Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 60 ++++++++++++++++++++++++++-------- pegainfer-sample/src/lib.rs | 60 ++++++++++++++++++++-------------- 2 files changed, 82 insertions(+), 38 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 2dfec3f8f..3175ec52b 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -8,6 +8,7 @@ use std::path::Path; use anyhow::Context as AnyhowContext; use anyhow::Result; +use pegainfer_core::cuda_graph::CudaGraphState; use pegainfer_core::ops; use pegainfer_core::tensor::DeviceContext; use pegainfer_core::tensor::HiddenStates; @@ -802,6 +803,8 @@ struct EngineState { sample_nonce: u64, /// Present only while the active row order is frozen. pipeline: Option, + /// Captured suppression, argmax and id-copy chain per decode bucket. + sampler_graphs: Vec, /// The overlap lane; `None` unless `PEGAINFER_ASYNC_PREFILL` opted in /// at startup. lane: Option, @@ -956,10 +959,28 @@ impl EngineState { )) })?; let prefix_cache = cache_cap.map(|k| PrefixCache::new(k, sliding_window)); - let scratch = SampleScratch::new(&ctx, vocab, arena_rows)?; + let mut scratch = SampleScratch::new(&ctx, vocab, arena_rows)?; let mut arena = serve.alloc_step_arena(&ctx, arena_rows, graph_enabled)?; serve.precapture_decode_graphs(&ctx, &mut arena)?; let suppress_ids = ops::SuppressIds::upload(&ctx, &policy.suppress, vocab)?; + let mut sampler_graphs = Vec::new(); + if graph_enabled { + let (logits, ids) = arena.logits_and_ids(); + let mut bucket = 1usize; + while bucket <= arena_rows { + logits.seq_len = bucket; + ops::suppress_logits_bf16_in_place(&ctx, logits, &suppress_ids)?; + pegainfer_sample::greedy_argmax_ids(&ctx, logits, bucket, ids, &mut scratch)?; + let mut graph = CudaGraphState::new(); + graph.capture_only(&ctx, || { + ops::suppress_logits_bf16_in_place(&ctx, logits, &suppress_ids)?; + pegainfer_sample::greedy_argmax_ids(&ctx, logits, bucket, ids, &mut scratch) + })?; + sampler_graphs.push(graph); + bucket *= 2; + } + ctx.sync()?; + } let lane = lane_mode .map(|mode| AsyncPrefillLane::new(&ctx, mode)) .transpose()?; @@ -974,6 +995,7 @@ impl EngineState { base_seed, sample_nonce: 0, pipeline: None, + sampler_graphs, lane, mix_chunk, max_context, @@ -1789,19 +1811,29 @@ impl EngineState { .decode_batch_step_resident(&self.ctx, &mut self.arena, &mut kvs)?; } } - let (logits, ids) = self.arena.logits_and_ids(); - ops::suppress_logits_bf16_in_place(&self.ctx, logits, &self.suppress_ids) - .context("suppression")?; - self.sample_nonce = self.sample_nonce.wrapping_add(1); - pegainfer_sample::greedy_stage_resident( - &self.ctx, - logits, - rows, - ids, - slot, - &mut self.scratch, - ) - .context("stage greedy picks")?; + let graph_slot = rows.next_power_of_two().trailing_zeros() as usize; + if let Some(graph) = self.sampler_graphs.get_mut(graph_slot) { + graph + .launch_captured(&self.ctx) + .context("launch sampler graph")?; + self.sample_nonce = self.sample_nonce.wrapping_add(1); + pegainfer_sample::greedy_stage_readback(&self.ctx, slot, &mut self.scratch) + .context("stage greedy readback")?; + } else { + let (logits, ids) = self.arena.logits_and_ids(); + ops::suppress_logits_bf16_in_place(&self.ctx, logits, &self.suppress_ids) + .context("suppression")?; + self.sample_nonce = self.sample_nonce.wrapping_add(1); + pegainfer_sample::greedy_stage_resident( + &self.ctx, + logits, + rows, + ids, + slot, + &mut self.scratch, + ) + .context("stage greedy picks")?; + } Ok(rows) } diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs index add3f78ef..080450eba 100644 --- a/pegainfer-sample/src/lib.rs +++ b/pegainfer-sample/src/lib.rs @@ -294,36 +294,59 @@ pub fn select_batch( Ok(tokens) } -fn validate_greedy_stage( +fn validate_greedy_shape( logits: &HiddenStates, rows: usize, - slot: usize, scratch: &SampleScratch, ) -> Result<()> { - ensure!(rows > 0, "greedy_stage_resident: empty batch"); + ensure!(rows > 0, "greedy_argmax_ids: empty batch"); ensure!( rows <= scratch.max_rows, - "greedy_stage_resident: {rows} rows exceeds scratch capacity {}", + "greedy_argmax_ids: {rows} rows exceeds scratch capacity {}", scratch.max_rows ); ensure!( logits.seq_len >= rows && logits.hidden_dim == scratch.vocab, - "greedy_stage_resident: logits shape {}x{} cannot serve {rows} rows x vocab {}", + "greedy_argmax_ids: logits shape {}x{} cannot serve {rows} rows x vocab {}", logits.seq_len, logits.hidden_dim, scratch.vocab ); - ensure!( - slot < STAGED_READBACK_SLOTS, - "greedy_stage_resident: slot {slot} out of range" - ) + Ok(()) +} + +/// Capturable argmax and device copy into the next embedding's id buffer. +pub fn greedy_argmax_ids( + ctx: &DeviceContext, + logits: &HiddenStates, + rows: usize, + ids_out: &mut CudaSlice, + scratch: &mut SampleScratch, +) -> Result<()> { + validate_greedy_shape(logits, rows, scratch)?; + argmax_batch_bf16_split_indexed_into( + ctx, + logits, + &scratch.identity_rows, + rows, + &mut scratch.argmax_partial_values, + &mut scratch.argmax_partial_indices, + &mut scratch.top1_values, + &mut scratch.argmax_out, + )?; + pegainfer_kernels::tensor::memcpy_dtod_u32_from_i32(ctx, &scratch.argmax_out, ids_out, rows) } -fn stage_greedy_readback( +/// Queue the pinned readback outside the sampler graph. +pub fn greedy_stage_readback( ctx: &DeviceContext, slot: usize, scratch: &mut SampleScratch, ) -> Result<()> { + ensure!( + slot < STAGED_READBACK_SLOTS, + "greedy_stage_readback: slot {slot} out of range" + ); let host = if slot == 0 { &mut scratch.argmax_host } else { @@ -331,7 +354,7 @@ fn stage_greedy_readback( }; ctx.stream .memcpy_dtoh(&scratch.argmax_out, host) - .map_err(|e| anyhow!("greedy_stage_resident D2H stage failed: {e}")) + .map_err(|e| anyhow!("greedy_stage_readback D2H stage failed: {e}")) } /// Argmax every row, leave the picks in the next embedding's id buffer, and @@ -344,19 +367,8 @@ pub fn greedy_stage_resident( slot: usize, scratch: &mut SampleScratch, ) -> Result<()> { - validate_greedy_stage(logits, rows, slot, scratch)?; - argmax_batch_bf16_split_indexed_into( - ctx, - logits, - &scratch.identity_rows, - rows, - &mut scratch.argmax_partial_values, - &mut scratch.argmax_partial_indices, - &mut scratch.top1_values, - &mut scratch.argmax_out, - )?; - pegainfer_kernels::tensor::memcpy_dtod_u32_from_i32(ctx, &scratch.argmax_out, ids_out, rows)?; - stage_greedy_readback(ctx, slot, scratch) + greedy_argmax_ids(ctx, logits, rows, ids_out, scratch)?; + greedy_stage_readback(ctx, slot, scratch) } /// Collect a readback staged into one of the two pinned slots. From 1c2269e4b1b2eed0ee3b9d9288241470a41da103 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 17:34:45 +0100 Subject: [PATCH 05/14] fix(kernels): the metadata advance takes its tables mutably and checks them Signed-off-by: Feathbow --- pegainfer-core/src/ops/paged_plan.rs | 3 ++ pegainfer-gemma4/src/serve.rs | 13 ++++---- pegainfer-kernels/csrc/shared/elementwise.cu | 4 +++ pegainfer-kernels/src/ops/attention.rs | 3 ++ pegainfer-kernels/src/ops/elementwise.rs | 32 +++++++++++++------- 5 files changed, 38 insertions(+), 17 deletions(-) diff --git a/pegainfer-core/src/ops/paged_plan.rs b/pegainfer-core/src/ops/paged_plan.rs index 2a01206a1..6f2078229 100644 --- a/pegainfer-core/src/ops/paged_plan.rs +++ b/pegainfer-core/src/ops/paged_plan.rs @@ -223,6 +223,9 @@ impl PrefillPagedPlan { pub fn kv_chunk_size_d(&self) -> &CudaSlice { self.inner.kv_chunk_size_d() } + pub fn decode_metadata_d_mut(&mut self) -> (&mut CudaSlice, &mut CudaSlice) { + self.inner.decode_metadata_d_mut() + } pub fn total_num_rows_d(&self) -> &CudaSlice { self.inner.total_num_rows_d() } diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index f1c5f8e46..8a49895f8 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -1976,8 +1976,8 @@ impl GemmaServe { tower: &mut TowerScratch, ids: &CudaSlice, rows: usize, - local_plan: &PrefillPagedPlan, - global_tables: &GlobalTables, + local_plan: &mut PrefillPagedPlan, + global_tables: &mut GlobalTables, global_split: &mut SplitKvState, local_origins: &CudaSlice, head_normed: &mut HiddenStates, @@ -2004,12 +2004,13 @@ impl GemmaServe { head_normed, logits, )?; + let (local_last, kv_chunk) = local_plan.decode_metadata_d_mut(); ops::advance_decode_metadata( ctx, - &global_tables.positions, - local_plan.last_page_len_d(), - &global_tables.pseudo_last, - local_plan.kv_chunk_size_d(), + &mut global_tables.positions, + local_last, + &mut global_tables.pseudo_last, + kv_chunk, rows, self.global_split_factor, ) diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index 97bc8fa83..12c5c3679 100644 --- a/pegainfer-kernels/csrc/shared/elementwise.cu +++ b/pegainfer-kernels/csrc/shared/elementwise.cu @@ -928,6 +928,10 @@ CUresult embedding_batched_vocab_shard_cuda( CUresult advance_decode_metadata_cuda( int *positions, int *local_last, int *pseudo_last, int *kv_chunk, int rows, int factor, cudaStream_t stream) { + if (positions == nullptr || local_last == nullptr || pseudo_last == nullptr || + kv_chunk == nullptr || rows <= 0 || factor <= 0) { + return CUDA_ERROR_INVALID_VALUE; + } int block = ADVANCE_DECODE_METADATA_BLOCK; int grid = (rows + block - 1) / block; advance_decode_metadata_kernel<<>>( diff --git a/pegainfer-kernels/src/ops/attention.rs b/pegainfer-kernels/src/ops/attention.rs index fab0b6227..201a63ad1 100644 --- a/pegainfer-kernels/src/ops/attention.rs +++ b/pegainfer-kernels/src/ops/attention.rs @@ -110,6 +110,9 @@ impl PrefillPagedPlan { pub fn kv_chunk_size_d(&self) -> &CudaSlice { &self.kv_chunk_size_d } + pub fn decode_metadata_d_mut(&mut self) -> (&mut CudaSlice, &mut CudaSlice) { + (&mut self.last_page_len_d, &mut self.kv_chunk_size_d) + } pub fn total_num_rows_d(&self) -> &CudaSlice { &self.total_num_rows_d } diff --git a/pegainfer-kernels/src/ops/elementwise.rs b/pegainfer-kernels/src/ops/elementwise.rs index c13a8b951..db3353f64 100644 --- a/pegainfer-kernels/src/ops/elementwise.rs +++ b/pegainfer-kernels/src/ops/elementwise.rs @@ -52,32 +52,42 @@ pub fn add_batch_into( /// Advance the per-row decode tables written by a regular graph replay. pub fn advance_decode_metadata( ctx: &DeviceContext, - positions: &CudaSlice, - local_last: &CudaSlice, - pseudo_last: &CudaSlice, - kv_chunk: &CudaSlice, + positions: &mut CudaSlice, + local_last: &mut CudaSlice, + pseudo_last: &mut CudaSlice, + kv_chunk: &mut CudaSlice, rows: usize, factor: usize, ) -> Result<()> { + anyhow::ensure!(rows > 0, "advance_decode_metadata: rows must be positive"); + anyhow::ensure!( + factor > 0, + "advance_decode_metadata: factor must be positive" + ); + let pseudo_rows = rows + .checked_mul(factor) + .ok_or_else(|| anyhow!("advance_decode_metadata: rows x factor overflow"))?; anyhow::ensure!( positions.len() >= rows && local_last.len() >= rows && kv_chunk.len() >= rows - && pseudo_last.len() >= rows * factor, + && pseudo_last.len() >= pseudo_rows, "advance_decode_metadata: {rows} rows x {factor} exceeds a table" ); - let (positions_ptr, _positions_guard) = positions.device_ptr(&ctx.stream); - let (local_ptr, _local_guard) = local_last.device_ptr(&ctx.stream); - let (pseudo_ptr, _pseudo_guard) = pseudo_last.device_ptr(&ctx.stream); - let (chunk_ptr, _chunk_guard) = kv_chunk.device_ptr(&ctx.stream); + let rows = super::checked_i32(rows, "advance decode metadata rows")?; + let factor = super::checked_i32(factor, "advance decode metadata factor")?; + let (positions_ptr, _positions_guard) = positions.device_ptr_mut(&ctx.stream); + let (local_ptr, _local_guard) = local_last.device_ptr_mut(&ctx.stream); + let (pseudo_ptr, _pseudo_guard) = pseudo_last.device_ptr_mut(&ctx.stream); + let (chunk_ptr, _chunk_guard) = kv_chunk.device_ptr_mut(&ctx.stream); let result = unsafe { ffi::advance_decode_metadata_cuda( positions_ptr as *mut i32, local_ptr as *mut i32, pseudo_ptr as *mut i32, chunk_ptr as *mut i32, - rows as i32, - factor as i32, + rows, + factor, crate::tensor::active_cu_stream(ctx), ) }; From 49976d5a24e9f5f0944f2c85d706d7bbd2a68944 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 17:34:45 +0100 Subject: [PATCH 06/14] fix(kernels): the staged sampler path stays on the base stream Signed-off-by: Feathbow --- pegainfer-kernels/src/tensor.rs | 24 ++++++++++++------------ pegainfer-sample/src/lib.rs | 20 ++++++++++++++++---- 2 files changed, 28 insertions(+), 16 deletions(-) diff --git a/pegainfer-kernels/src/tensor.rs b/pegainfer-kernels/src/tensor.rs index 3288f756b..b2c423a6a 100644 --- a/pegainfer-kernels/src/tensor.rs +++ b/pegainfer-kernels/src/tensor.rs @@ -79,12 +79,7 @@ pub fn active_cu_stream(ctx: &DeviceContext) -> CUstream { .unwrap_or_else(|| ctx.stream.cu_stream()) } -/// Bounded spin wait on the context stream. A blocking wait parks the caller -/// and pays the OS scheduler's wake-up — low milliseconds per call on a -/// loaded host — which a decode loop pays every step right before it reads -/// the sampled tokens. Polling `cuStreamQuery` keeps the wake on our side; -/// the cap falls back to the parking wait so a stuck stream still sleeps -/// instead of burning the core forever. +/// Poll the active stream up to a fixed cap, then synchronize that stream. const STREAM_SPIN_WAIT_CAP: std::time::Duration = std::time::Duration::from_millis(5); pub fn stream_spin_wait(ctx: &DeviceContext) -> anyhow::Result<()> { @@ -95,10 +90,12 @@ pub fn stream_spin_wait(ctx: &DeviceContext) -> anyhow::Result<()> { cudarc::driver::sys::CUresult::CUDA_SUCCESS => return Ok(()), cudarc::driver::sys::CUresult::CUDA_ERROR_NOT_READY => { if std::time::Instant::now() >= cap { - return ctx - .stream - .synchronize() - .map_err(|e| anyhow::anyhow!("stream sync after spin cap failed: {e}")); + let result = unsafe { cudarc::driver::sys::cuStreamSynchronize(stream) }; + anyhow::ensure!( + result == cudarc::driver::sys::CUresult::CUDA_SUCCESS, + "stream sync after spin cap failed: {result:?}" + ); + return Ok(()); } std::hint::spin_loop(); } @@ -117,6 +114,10 @@ pub fn memcpy_dtod_u32_from_i32( use cudarc::driver::DevicePtr; use cudarc::driver::DevicePtrMut; + anyhow::ensure!( + !has_stream_override(), + "dtod i32->u32 copy runs on the base stream only" + ); anyhow::ensure!( count <= src.len() && count <= dst.len(), "dtod i32->u32 copy of {count} exceeds src {} or dst {}", @@ -126,7 +127,6 @@ pub fn memcpy_dtod_u32_from_i32( if count == 0 { return Ok(()); } - let stream = active_cu_stream(ctx); let (src_ptr, _src_guard) = src.device_ptr(&ctx.stream); let (dst_ptr, _dst_guard) = dst.device_ptr_mut(&ctx.stream); let result = unsafe { @@ -134,7 +134,7 @@ pub fn memcpy_dtod_u32_from_i32( dst_ptr, src_ptr, count * std::mem::size_of::(), - stream, + ctx.stream.cu_stream(), ) }; anyhow::ensure!( diff --git a/pegainfer-sample/src/lib.rs b/pegainfer-sample/src/lib.rs index 080450eba..071871082 100644 --- a/pegainfer-sample/src/lib.rs +++ b/pegainfer-sample/src/lib.rs @@ -315,7 +315,7 @@ fn validate_greedy_shape( Ok(()) } -/// Capturable argmax and device copy into the next embedding's id buffer. +/// Capturable base-stream-only argmax and device copy into the next embedding's id buffer. pub fn greedy_argmax_ids( ctx: &DeviceContext, logits: &HiddenStates, @@ -323,6 +323,10 @@ pub fn greedy_argmax_ids( ids_out: &mut CudaSlice, scratch: &mut SampleScratch, ) -> Result<()> { + ensure!( + !has_stream_override(), + "the staged sampler path runs on the base stream only" + ); validate_greedy_shape(logits, rows, scratch)?; argmax_batch_bf16_split_indexed_into( ctx, @@ -337,12 +341,16 @@ pub fn greedy_argmax_ids( pegainfer_kernels::tensor::memcpy_dtod_u32_from_i32(ctx, &scratch.argmax_out, ids_out, rows) } -/// Queue the pinned readback outside the sampler graph. +/// Queue the base-stream-only pinned readback outside the sampler graph. pub fn greedy_stage_readback( ctx: &DeviceContext, slot: usize, scratch: &mut SampleScratch, ) -> Result<()> { + ensure!( + !has_stream_override(), + "the staged sampler path runs on the base stream only" + ); ensure!( slot < STAGED_READBACK_SLOTS, "greedy_stage_readback: slot {slot} out of range" @@ -357,8 +365,8 @@ pub fn greedy_stage_readback( .map_err(|e| anyhow!("greedy_stage_readback D2H stage failed: {e}")) } -/// Argmax every row, leave the picks in the next embedding's id buffer, and -/// queue their pinned readback without waiting for it. +/// On the base stream, argmax every row, leave the picks in the next embedding's +/// id buffer, and queue their pinned readback without waiting for it. pub fn greedy_stage_resident( ctx: &DeviceContext, logits: &HiddenStates, @@ -367,6 +375,10 @@ pub fn greedy_stage_resident( slot: usize, scratch: &mut SampleScratch, ) -> Result<()> { + ensure!( + !has_stream_override(), + "the staged sampler path runs on the base stream only" + ); greedy_argmax_ids(ctx, logits, rows, ids_out, scratch)?; greedy_stage_readback(ctx, slot, scratch) } From 773846c7ef9843d82827f976ad7a16bef24774a7 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 17:34:45 +0100 Subject: [PATCH 07/14] fix(gemma4): the pipeline fences the device before it drops a failed batch Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 34 ++++++++++++++++++++++++++++------ pegainfer-gemma4/src/serve.rs | 10 +++++----- 2 files changed, 33 insertions(+), 11 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 3175ec52b..c17d8ac49 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -966,11 +966,13 @@ impl EngineState { let mut sampler_graphs = Vec::new(); if graph_enabled { let (logits, ids) = arena.logits_and_ids(); + // The warm pass lands lazy module loads outside capture. + logits.seq_len = arena_rows; + ops::suppress_logits_bf16_in_place(&ctx, logits, &suppress_ids)?; + pegainfer_sample::greedy_argmax_ids(&ctx, logits, arena_rows, ids, &mut scratch)?; let mut bucket = 1usize; while bucket <= arena_rows { logits.seq_len = bucket; - ops::suppress_logits_bf16_in_place(&ctx, logits, &suppress_ids)?; - pegainfer_sample::greedy_argmax_ids(&ctx, logits, bucket, ids, &mut scratch)?; let mut graph = CudaGraphState::new(); graph.capture_only(&ctx, || { ops::suppress_logits_bf16_in_place(&ctx, logits, &suppress_ids)?; @@ -1791,6 +1793,14 @@ impl EngineState { }) } + fn fence_or_abort(&self) { + let sync = unsafe { cudarc::driver::sys::cuStreamSynchronize(self.ctx.stream.cu_stream()) }; + if sync != cudarc::driver::sys::CUresult::CUDA_SUCCESS { + log::error!("FATAL: cuStreamSynchronize(decode) failed ({sync:?}); aborting"); + std::process::abort(); + } + } + /// Queue one decode and stage its greedy picks into the next embedding's /// id buffer and one pinned readback slot. fn launch_staged( @@ -1811,7 +1821,7 @@ impl EngineState { .decode_batch_step_resident(&self.ctx, &mut self.arena, &mut kvs)?; } } - let graph_slot = rows.next_power_of_two().trailing_zeros() as usize; + let graph_slot = crate::serve::decode_bucket_slot(rows); if let Some(graph) = self.sampler_graphs.get_mut(graph_slot) { graph .launch_captured(&self.ctx) @@ -1862,6 +1872,7 @@ impl EngineState { return; }; if let Err(err) = self.collect_pending(active, &pending) { + self.fence_or_abort(); return fail_active_batch(active, "pipelined decode drain", &err); } active.retain(|entry| !entry.stopping); @@ -1977,7 +1988,10 @@ impl EngineState { .decode_batch_step(&self.ctx, &mut self.arena, &mut kvs, &tokens) { Ok(logits) => logits, - Err(err) => return fail_active_batch(active, "batched decode", &err), + Err(err) => { + self.fence_or_abort(); + return fail_active_batch(active, "batched decode", &err); + } } }; let sampled = { @@ -1995,7 +2009,10 @@ impl EngineState { }; match sampled { Ok(mut sampled) => emit_decode_rows(active, &mut sampled, 0), - Err(err) => fail_active_batch(active, "batched decode", &err), + Err(err) => { + self.fence_or_abort(); + fail_active_batch(active, "batched decode", &err); + } } } @@ -2008,6 +2025,7 @@ impl EngineState { match self.launch_staged(active, true, next_slot) { Ok(rows) => { if let Err(err) = self.collect_pending(active, &pending) { + self.fence_or_abort(); fail_active_batch(active, "pipelined decode collect", &err); return; } @@ -2020,6 +2038,7 @@ impl EngineState { if let Err(collect_err) = self.collect_pending(active, &pending) { log::error!("collect during launch failure failed: {collect_err:#}"); } + self.fence_or_abort(); fail_active_batch(active, "pipelined decode launch", &err); } } @@ -2039,7 +2058,10 @@ impl EngineState { if self.pipeline_eligible(active) { match self.launch_staged(active, false, 0) { Ok(rows) => self.pipeline = Some(PendingDecode { rows, slot: 0 }), - Err(err) => fail_active_batch(active, "batched decode", &err), + Err(err) => { + self.fence_or_abort(); + fail_active_batch(active, "batched decode", &err); + } } return; } diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 8a49895f8..9b0192d89 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -468,8 +468,8 @@ impl TowerScratch { } } -fn bucket_slot(bucket: usize) -> usize { - bucket.trailing_zeros() as usize +pub(crate) fn decode_bucket_slot(rows: usize) -> usize { + rows.next_power_of_two().trailing_zeros() as usize } fn hidden_pair(hidden: &mut [HiddenStates; 2], src: usize) -> (&HiddenStates, &mut HiddenStates) { @@ -935,7 +935,7 @@ impl GemmaServe { mix_indptr, head_normed: HiddenStates::zeros(ctx, self.local_geom.hidden_size, max_rows)?, logits: HiddenStates::zeros(ctx, self.weights.embed_tokens.rows, max_rows)?, - graphs: (0..=bucket_slot(max_rows)) + graphs: (0..=decode_bucket_slot(max_rows)) .map(|_| CudaGraphState::new()) .collect(), graph_enabled, @@ -1934,7 +1934,7 @@ impl GemmaServe { .. } = arena; if *graph_enabled { - let graph = &mut graphs[bucket_slot(padded)]; + let graph = &mut graphs[decode_bucket_slot(padded)]; anyhow::ensure!( graph.is_captured(), "no captured graph for bucket {padded}; the pre-capture sweep must cover \ @@ -2366,7 +2366,7 @@ impl GemmaServe { head_normed, logits, )?; - graphs[bucket_slot(padded)].capture_only(ctx, || { + graphs[decode_bucket_slot(padded)].capture_only(ctx, || { self.decode_gpu_body( ctx, tower, From 872c5d70a405ed513d50d4ca90891643a94e0339 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 18:06:34 +0100 Subject: [PATCH 08/14] fix(kernels): the metadata advance bounds its index arithmetic to i32 Signed-off-by: Feathbow --- pegainfer-kernels/csrc/shared/elementwise.cu | 4 ++-- pegainfer-kernels/src/ops/elementwise.rs | 3 ++- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index 12c5c3679..cd8b6e58a 100644 --- a/pegainfer-kernels/csrc/shared/elementwise.cu +++ b/pegainfer-kernels/csrc/shared/elementwise.cu @@ -929,11 +929,11 @@ CUresult advance_decode_metadata_cuda( int *positions, int *local_last, int *pseudo_last, int *kv_chunk, int rows, int factor, cudaStream_t stream) { if (positions == nullptr || local_last == nullptr || pseudo_last == nullptr || - kv_chunk == nullptr || rows <= 0 || factor <= 0) { + kv_chunk == nullptr || rows <= 0 || factor <= 0 || rows > INT_MAX / factor) { return CUDA_ERROR_INVALID_VALUE; } int block = ADVANCE_DECODE_METADATA_BLOCK; - int grid = (rows + block - 1) / block; + int grid = 1 + (rows - 1) / block; advance_decode_metadata_kernel<<>>( positions, local_last, pseudo_last, kv_chunk, rows, factor); return (CUresult)cudaGetLastError(); diff --git a/pegainfer-kernels/src/ops/elementwise.rs b/pegainfer-kernels/src/ops/elementwise.rs index db3353f64..29fa7f56c 100644 --- a/pegainfer-kernels/src/ops/elementwise.rs +++ b/pegainfer-kernels/src/ops/elementwise.rs @@ -66,7 +66,8 @@ pub fn advance_decode_metadata( ); let pseudo_rows = rows .checked_mul(factor) - .ok_or_else(|| anyhow!("advance_decode_metadata: rows x factor overflow"))?; + .filter(|n| i32::try_from(*n).is_ok()) + .ok_or_else(|| anyhow!("advance_decode_metadata: rows x factor exceeds i32"))?; anyhow::ensure!( positions.len() >= rows && local_last.len() >= rows From 5b08b46e88fa987df26db79397ca2bac42a43ffa Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 18:34:36 +0100 Subject: [PATCH 09/14] fix(gemma4): the admission queue drains the pipeline only when it can admit Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index c17d8ac49..fbef2a16b 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -1014,9 +1014,8 @@ impl EngineState { if pending.is_empty() { return; } - self.arena.invalidate_decode_fingerprint(); - self.drain_pipeline(active); let mut attempts = 0; + let mut drained = false; while attempts < self.slots && active.len() < self.slots { // With the lane busy, arrivals wait in `pending` while decode // keeps stepping. @@ -1031,6 +1030,13 @@ impl EngineState { break; }; attempts += 1; + // Only an admission that can run changes the roster; a full + // batch with a waiting queue keeps its pipeline and fingerprint. + if !drained { + self.arena.invalidate_decode_fingerprint(); + self.drain_pipeline(active); + drained = true; + } let can_wait = !active.is_empty(); match self.admit_and_prefill(item, can_wait, active, pending, &mut attempts) { Admitted::Active(request) => active.push(*request), From ca46460e6cd92192bade6caedcbe3447475db787 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 18:34:36 +0100 Subject: [PATCH 10/14] docs(gemma4): the decode pipeline contract Signed-off-by: Feathbow --- docs/models/gemma4/serving.md | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 6c3570292..0052570c1 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -27,6 +27,16 @@ The register router accepts exactly 128 experts and from 1 through 32 picks. A n The checkpoint-backed `the_routed_block_matches_the_reference_formulas` gate owns the scratch-capacity, companion-route and coarse-block evidence, including a narrow block replay after a coarse block on one scratch. On shared rows, it proves that the 16-row and 64-row block pick the same experts with the same weight bits and produce the same gate, weighted-down and block bits. `router_topk_matches_the_exact_128_expert_contract` owns the register-router boundary and non-finite rows. The kernels-owned `kimi_marlin_align_boundary_matches_vllm_contract` oracle owns stable counts, offsets, padding and expert-local order on both sides of the alignment dispatch boundary. `scripts/gemma4_gates.sh` owns the Gemma crate's ignored gates and the kernels crate's Gemma router contract (the device-only test under the `gemma4` feature), holding both crates' ignored sets against its manifest; the Kimi alignment oracle needs the `kimi-k2` feature and an `sm_90` device and is run by hand. +## The decode pipeline + +Greedy decode rounds run a depth-two software pipeline on the base stream: a step's argmax writes its picks straight into the id buffer the next step's embedding reads, its readback lands in one of two pinned slots, and the emitted token stream lags compute by one step. The production invariant is that token ids, finish reasons and token counts are identical to collecting every step at the same batch composition, and that the kernels inside a step and their order do not change. + +A batch is eligible only when every row is greedy, scores no logprobs, sits at least two tokens from its length cap and is not already stopping. While a successor step is in flight the active row order is pinned: the successor's ids were written on the device for the current order, so no row may be retired, reordered or added until that step is collected. A stop found late marks the row stopping and it retires when the pipeline drains. The pipeline drains before anything that can change the roster — an admission that is actually about to run (a full batch with a waiting queue keeps its pipeline), a lane join, a cancelled row, or a batch that stopped being eligible — and the collect-every-step path takes over from there. + +A regular decode step, every row one token further with its page, chunk and split structure untouched, also skips its metadata rebuild and every upload: a kernel captured at the decode graph's tail advances the per-row tables in place, and a fingerprint of the previous step proves the advanced device state is what a full upload would write. Mixed steps, page turns, chunk boundaries, admissions and the precapture warm pass invalidate the fingerprint and rebuild as before. + +The staged sampler chain — suppression, argmax and the device copy of the picks — is captured per bucket at startup and launched as one graph beside the decode replay; the pinned readback stays outside so the collector keeps the copy's own event. The staged path runs on the base stream only and refuses a stream override. If a step fails after its device work was enqueued, the engine synchronizes the stream before it drops the batch and returns its pages, and aborts if the device cannot be synchronized, the policy the prefill lane already applies. + ## The two pools Gemma 4 runs two attention families with different KV shapes, so the budget is two budgets. With 16-token pages, `C = ceil(8192/16) = 512` context pages and `W = ceil(1024/16) + 1 = 65` window pages: From 286351f941e19eb91e512787c3ffebcf73a55ea9 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 18:46:28 +0100 Subject: [PATCH 11/14] test(gemma4): the full roster keeps its pipeline under a queue Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 69 ++++++++++++++++++++++++++++++++++ pegainfer-gemma4/src/serve.rs | 5 +++ scripts/gemma4_gates.sh | 1 + 3 files changed, 75 insertions(+) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index fbef2a16b..6b0ff4440 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -2757,6 +2757,75 @@ mod lane_tests { assert_eq!(drain(&mut rx_c, "queued third").tokens, 6); } + /// The production loop runs an intake pass before every decode round. + /// With both slots held and a request queued, that pass can admit + /// nothing, so it must leave the staged successor and the step + /// fingerprint alone; the third request joins only once a slot frees. + #[test] + #[ignore = "requires the pinned 12B checkpoint, a GPU, and --test-threads=1"] + fn the_full_roster_keeps_its_pipeline_under_a_queue() { + let dir = crate::testkit::model_path(); + let policy = super::generation_policy(&dir).expect("policy"); + let _env = scoped_engine_env(&[("PEGAINFER_DECODE_SLOTS", "2")]); + let mut state = + super::EngineState::load(&dir, 0, policy, 0x5EED, true).expect("engine state"); + let prompts = walk_prompts(); + let (req_a, mut rx_a) = walk_request(prompts[0].clone(), 24); + let (req_b, mut rx_b) = walk_request(prompts[1].clone(), 24); + let (req_c, mut rx_c) = walk_request(prompts[2].clone(), 6); + + let mut pending = std::collections::VecDeque::new(); + let mut active: Vec = Vec::new(); + pending.push_back((req_a, pegainfer_frontend::engine::KvPrefix::none())); + pending.push_back((req_b, pegainfer_frontend::engine::KvPrefix::none())); + state.admit_from_queue(&mut pending, &mut active); + assert_eq!(active.len(), 2, "both slots are held by live requests"); + pending.push_back((req_c, pegainfer_frontend::engine::KvPrefix::none())); + + state.decode_round(&mut active); + state.decode_round(&mut active); + assert!( + state.pipeline.is_some(), + "a greedy batch stages a successor" + ); + assert!( + state.arena.has_decode_fingerprint(), + "a regular step leaves its fingerprint" + ); + for round in 0..4 { + state.admit_from_queue(&mut pending, &mut active); + assert_eq!( + pending.len(), + 1, + "round {round}: full slots keep the third queued" + ); + assert!( + state.pipeline.is_some(), + "round {round}: an intake that admits nothing keeps the pipeline" + ); + assert!( + state.arena.has_decode_fingerprint(), + "round {round}: and the fingerprint" + ); + state.decode_round(&mut active); + } + + while active.len() == 2 { + state.admit_from_queue(&mut pending, &mut active); + state.decode_round(&mut active); + } + state.admit_from_queue(&mut pending, &mut active); + assert_eq!(active.len(), 2, "the freed slot admits the third request"); + assert!(pending.is_empty(), "the queue drained"); + while !active.is_empty() { + state.admit_from_queue(&mut pending, &mut active); + state.decode_round(&mut active); + } + assert_eq!(drain(&mut rx_a, "incumbent a").tokens, 24); + assert_eq!(drain(&mut rx_b, "incumbent b").tokens, 24); + assert_eq!(drain(&mut rx_c, "queued third").tokens, 6); + } + /// The chunked pool provisions one shared segment transient, so no /// walker may park pages ahead of its rounds: with the knob set before /// load — the reduced production pool, asserted against the provision diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 9b0192d89..8b1818c96 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -613,6 +613,11 @@ impl StepArena { pub(crate) fn invalidate_decode_fingerprint(&mut self) { self.steady = None; } + + #[cfg(test)] + pub(crate) fn has_decode_fingerprint(&self) -> bool { + self.steady.is_some() + } } /// How many pseudo-requests the global decode read presents each request diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 27aa46d21..61fb5262e 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -60,6 +60,7 @@ GATES_DENSE_AND_ROUTED=( GATES_SERVING_CONTRACT=( "gpu,ckpt engine::lane_tests::the_gathered_lifecycle_completes" "gpu,ckpt,prompts engine::lane_tests::the_raised_ceiling_and_slots_hold_at_the_roster_edge" + "gpu,ckpt,prompts engine::lane_tests::the_full_roster_keeps_its_pipeline_under_a_queue" "gpu,ckpt engine::lane_tests::the_raise_reaches_the_frontend" "ckpt engine::lane_tests::the_raise_refuses_without_its_prerequisites" ) From e5838c0ef3174653753cf3fcc9e0c56039f3b2de Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 19:11:34 +0100 Subject: [PATCH 12/14] test(gemma4): the roster gate retires one incumbent before the other Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 6b0ff4440..8a37ea986 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -2771,7 +2771,7 @@ mod lane_tests { super::EngineState::load(&dir, 0, policy, 0x5EED, true).expect("engine state"); let prompts = walk_prompts(); let (req_a, mut rx_a) = walk_request(prompts[0].clone(), 24); - let (req_b, mut rx_b) = walk_request(prompts[1].clone(), 24); + let (req_b, mut rx_b) = walk_request(prompts[1].clone(), 40); let (req_c, mut rx_c) = walk_request(prompts[2].clone(), 6); let mut pending = std::collections::VecDeque::new(); @@ -2822,7 +2822,7 @@ mod lane_tests { state.decode_round(&mut active); } assert_eq!(drain(&mut rx_a, "incumbent a").tokens, 24); - assert_eq!(drain(&mut rx_b, "incumbent b").tokens, 24); + assert_eq!(drain(&mut rx_b, "incumbent b").tokens, 40); assert_eq!(drain(&mut rx_c, "queued third").tokens, 6); } From 01744222c96eea670759b8117129185c253ff00e Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 19:11:36 +0100 Subject: [PATCH 13/14] fix(gemma4): the admission drains the pipeline only when it changes the roster Signed-off-by: Feathbow --- docs/models/gemma4/serving.md | 2 +- pegainfer-gemma4/src/engine.rs | 10 ++-------- 2 files changed, 3 insertions(+), 9 deletions(-) diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 0052570c1..4a75884f2 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -31,7 +31,7 @@ The checkpoint-backed `the_routed_block_matches_the_reference_formulas` gate own Greedy decode rounds run a depth-two software pipeline on the base stream: a step's argmax writes its picks straight into the id buffer the next step's embedding reads, its readback lands in one of two pinned slots, and the emitted token stream lags compute by one step. The production invariant is that token ids, finish reasons and token counts are identical to collecting every step at the same batch composition, and that the kernels inside a step and their order do not change. -A batch is eligible only when every row is greedy, scores no logprobs, sits at least two tokens from its length cap and is not already stopping. While a successor step is in flight the active row order is pinned: the successor's ids were written on the device for the current order, so no row may be retired, reordered or added until that step is collected. A stop found late marks the row stopping and it retires when the pipeline drains. The pipeline drains before anything that can change the roster — an admission that is actually about to run (a full batch with a waiting queue keeps its pipeline), a lane join, a cancelled row, or a batch that stopped being eligible — and the collect-every-step path takes over from there. +A batch is eligible only when every row is greedy, scores no logprobs, sits at least two tokens from its length cap and is not already stopping. While a successor step is in flight the active row order is pinned: the successor's ids were written on the device for the current order, so no row may be retired, reordered or added until that step is collected. A stop found late marks the row stopping and it retires when the pipeline drains. The pipeline drains before a synchronous or mixed admission that will use the step arena or change the roster, a lane join, a cancelled row, or a batch that stopped being eligible. An admission that turns out not viable — a closed or invalid request, a page-shortage requeue, or a failed Scheduled send — and an async lane launch keep it; the lane join drains. The collect-every-step path takes over after a drain. A regular decode step, every row one token further with its page, chunk and split structure untouched, also skips its metadata rebuild and every upload: a kernel captured at the decode graph's tail advances the per-row tables in place, and a fingerprint of the previous step proves the advanced device state is what a full upload would write. Mixed steps, page turns, chunk boundaries, admissions and the precapture warm pass invalidate the fingerprint and rebuild as before. diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 8a37ea986..17181ec39 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -1015,7 +1015,6 @@ impl EngineState { return; } let mut attempts = 0; - let mut drained = false; while attempts < self.slots && active.len() < self.slots { // With the lane busy, arrivals wait in `pending` while decode // keeps stepping. @@ -1030,13 +1029,6 @@ impl EngineState { break; }; attempts += 1; - // Only an admission that can run changes the roster; a full - // batch with a waiting queue keeps its pipeline and fingerprint. - if !drained { - self.arena.invalidate_decode_fingerprint(); - self.drain_pipeline(active); - drained = true; - } let can_wait = !active.is_empty(); match self.admit_and_prefill(item, can_wait, active, pending, &mut attempts) { Admitted::Active(request) => active.push(*request), @@ -1142,6 +1134,8 @@ impl EngineState { // weight scan — one step prefills every gathered newcomer and // advances every active row. if !active.is_empty() { + self.arena.invalidate_decode_fingerprint(); + self.drain_pipeline(active); self.ready_decode_rows(active); if !active.is_empty() { // Gather more admissible prompts into the same step. A From 7dbc296185f9be11c375db2c75cd2a68d46d680d Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 19:38:33 +0100 Subject: [PATCH 14/14] fix(gemma4): a solo admission drops the retired roster's fingerprint Signed-off-by: Feathbow --- docs/models/gemma4/serving.md | 2 +- pegainfer-gemma4/src/engine.rs | 85 +++++++++++++++++++++++++++++++++- scripts/gemma4_gates.sh | 1 + 3 files changed, 86 insertions(+), 2 deletions(-) diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 4a75884f2..e9f7ed397 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -31,7 +31,7 @@ The checkpoint-backed `the_routed_block_matches_the_reference_formulas` gate own Greedy decode rounds run a depth-two software pipeline on the base stream: a step's argmax writes its picks straight into the id buffer the next step's embedding reads, its readback lands in one of two pinned slots, and the emitted token stream lags compute by one step. The production invariant is that token ids, finish reasons and token counts are identical to collecting every step at the same batch composition, and that the kernels inside a step and their order do not change. -A batch is eligible only when every row is greedy, scores no logprobs, sits at least two tokens from its length cap and is not already stopping. While a successor step is in flight the active row order is pinned: the successor's ids were written on the device for the current order, so no row may be retired, reordered or added until that step is collected. A stop found late marks the row stopping and it retires when the pipeline drains. The pipeline drains before a synchronous or mixed admission that will use the step arena or change the roster, a lane join, a cancelled row, or a batch that stopped being eligible. An admission that turns out not viable — a closed or invalid request, a page-shortage requeue, or a failed Scheduled send — and an async lane launch keep it; the lane join drains. The collect-every-step path takes over after a drain. +A batch is eligible only when every row is greedy, scores no logprobs, sits at least two tokens from its length cap and is not already stopping. While a successor step is in flight the active row order is pinned: the successor's ids were written on the device for the current order, so no row may be retired, reordered or added until that step is collected. A stop found late marks the row stopping and it retires when the pipeline drains. The pipeline drains before a synchronous or mixed admission that will use the step arena or change the roster, a lane join, a cancelled row, or a batch that stopped being eligible. An admission that turns out not viable — a closed or invalid request, a page-shortage requeue, or a failed Scheduled send — and an async lane launch keep it; the lane join drains. A solo admission into an idle engine has nothing to drain but must drop the fingerprint the retired roster left, since it starts a new roster with new page identities. The collect-every-step path takes over after a drain. A regular decode step, every row one token further with its page, chunk and split structure untouched, also skips its metadata rebuild and every upload: a kernel captured at the decode graph's tail advances the per-row tables in place, and a fingerprint of the previous step proves the advanced device state is what a full upload would write. Mixed steps, page turns, chunk boundaries, admissions and the precapture warm pass invalidate the fingerprint and rebuild as before. diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 17181ec39..d4aabb0fa 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -1256,6 +1256,11 @@ impl EngineState { }); Admitted::Done }; + // A solo admission starts a new roster: the fingerprint the retired + // one left would otherwise pass a new request whose frontier and + // page structure happen to line up, and its first step would keep + // the old page tables. Nothing is in flight, so no drain is needed. + self.arena.invalidate_decode_fingerprint(); // Under the chunk knob a solo prompt walks its own segments too: // residency stays window plus segment whatever the prompt length. let stepped = if let Some(chunk) = self.mix_chunk { @@ -2340,14 +2345,19 @@ mod lane_tests { tokens: usize, cached: usize, finish: FinishReason, + ids: Vec, } fn drain(rx: &mut TokenStreamReceiver, name: &str) -> Drained { let mut tokens = 0; let mut cached = 0; + let mut ids = Vec::new(); loop { match rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Token { .. }) => tokens += 1, + Some(TokenEvent::Token { id, .. }) => { + tokens += 1; + ids.push(id); + } Some(TokenEvent::Scheduled { cached_tokens, .. }) => cached = cached_tokens, Some(TokenEvent::PromptTokens { .. } | TokenEvent::KvTransfer { .. }) => {} Some(TokenEvent::Finished { finish_reason, .. }) => { @@ -2355,6 +2365,7 @@ mod lane_tests { tokens, cached, finish: finish_reason, + ids, }; } Some(TokenEvent::Error { message, .. }) => panic!("{name}: error: {message}"), @@ -2820,6 +2831,78 @@ mod lane_tests { assert_eq!(drain(&mut rx_c, "queued third").tokens, 6); } + /// A roster that empties leaves its last fingerprint in the arena. The + /// solo admission that refills the idle engine must drop it: a new + /// request whose frontier sits one token past the retired one's, with + /// the same page structure, would otherwise take the regular-step skip + /// on its first decode and keep the retired roster's page tables. + #[test] + #[ignore = "requires the pinned 12B checkpoint, a GPU, and --test-threads=1"] + fn an_idle_refill_drops_the_retired_fingerprint() { + let dir = crate::testkit::model_path(); + let policy = super::generation_policy(&dir).expect("policy"); + let _env = scoped_engine_env(&[("PEGAINFER_DECODE_SLOTS", "2")]); + let prompts = walk_prompts(); + let first_len = 64usize; + let budget = 8usize; + let first: Vec = prompts[0].iter().cycle().copied().take(first_len).collect(); + // The retired request's last step ran at kv_len = first_len + budget - 1; + // this prompt's first decode step runs at exactly one more. + let second: Vec = prompts[1] + .iter() + .cycle() + .copied() + .take(first_len + budget - 1) + .collect(); + + let run_second_alone = |state: &mut super::EngineState| -> Drained { + let (req, mut rx) = walk_request(second.clone(), budget); + let mut pending = std::collections::VecDeque::new(); + let mut active: Vec = Vec::new(); + pending.push_back((req, pegainfer_frontend::engine::KvPrefix::none())); + state.admit_from_queue(&mut pending, &mut active); + assert_eq!(active.len(), 1, "the prompt is admitted"); + assert!( + !state.arena.has_decode_fingerprint(), + "a solo admission starts with no fingerprint" + ); + while !active.is_empty() { + state.admit_from_queue(&mut pending, &mut active); + state.decode_round(&mut active); + } + drain(&mut rx, "second") + }; + + let mut state = + super::EngineState::load(&dir, 0, policy, 0x5EED, true).expect("engine state"); + let (req_a, mut rx_a) = walk_request(first, budget); + let mut pending = std::collections::VecDeque::new(); + let mut active: Vec = Vec::new(); + pending.push_back((req_a, pegainfer_frontend::engine::KvPrefix::none())); + state.admit_from_queue(&mut pending, &mut active); + while !active.is_empty() { + state.admit_from_queue(&mut pending, &mut active); + state.decode_round(&mut active); + } + assert_eq!(drain(&mut rx_a, "first").tokens, budget); + assert!( + state.arena.has_decode_fingerprint(), + "the retired roster leaves its fingerprint behind" + ); + let refilled = run_second_alone(&mut state); + + drop(state); + let fresh_policy = super::generation_policy(&dir).expect("policy"); + let mut fresh = super::EngineState::load(&dir, 0, fresh_policy, 0x5EED, true) + .expect("fresh engine state"); + let alone = run_second_alone(&mut fresh); + assert_eq!(refilled.tokens, budget); + assert_eq!( + refilled.ids, alone.ids, + "the refill answers exactly as a fresh engine does" + ); + } + /// The chunked pool provisions one shared segment transient, so no /// walker may park pages ahead of its rounds: with the knob set before /// load — the reduced production pool, asserted against the provision diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 61fb5262e..268d216fb 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -61,6 +61,7 @@ GATES_SERVING_CONTRACT=( "gpu,ckpt engine::lane_tests::the_gathered_lifecycle_completes" "gpu,ckpt,prompts engine::lane_tests::the_raised_ceiling_and_slots_hold_at_the_roster_edge" "gpu,ckpt,prompts engine::lane_tests::the_full_roster_keeps_its_pipeline_under_a_queue" + "gpu,ckpt,prompts engine::lane_tests::an_idle_refill_drops_the_retired_fingerprint" "gpu,ckpt engine::lane_tests::the_raise_reaches_the_frontend" "ckpt engine::lane_tests::the_raise_refuses_without_its_prerequisites" )