diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 6c3570292..e9f7ed397 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 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. + +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: 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-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/engine.rs b/pegainfer-gemma4/src/engine.rs index dfdb147f1..d4aabb0fa 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; @@ -697,6 +698,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 +712,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 +801,10 @@ 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, + /// 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, @@ -905,10 +959,30 @@ 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(); + // 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; + 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()?; @@ -922,6 +996,8 @@ impl EngineState { suppress_ids, base_seed, sample_nonce: 0, + pipeline: None, + sampler_graphs, lane, mix_chunk, max_context, @@ -935,6 +1011,9 @@ 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; + } let mut attempts = 0; while attempts < self.slots && active.len() < self.slots { // With the lane busy, arrivals wait in `pending` while decode @@ -1055,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 @@ -1175,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 { @@ -1302,6 +1388,14 @@ 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.arena.invalidate_decode_fingerprint(); + self.drain_pipeline(active); + } let Some(lane) = self.lane.as_mut() else { return; }; @@ -1675,6 +1769,120 @@ 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() + }) + } + + 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( + &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 graph_slot = crate::serve::decode_bucket_slot(rows); + 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) + } + + /// 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) { + self.fence_or_abort(); + 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 +1984,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(); @@ -1794,7 +1993,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 = { @@ -1812,8 +2014,63 @@ 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); + } + } + } + + /// 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) { + self.fence_or_abort(); + 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:#}"); + } + self.fence_or_abort(); + 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) => { + self.fence_or_abort(); + fail_active_batch(active, "batched decode", &err); + } + } + return; + } + self.decode_round_collect(active); } } @@ -1903,6 +2160,7 @@ fn settle_first_token( next, emitted: 1, prompt_tokens, + stopping: false, }) } @@ -2087,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, .. }) => { @@ -2102,6 +2365,7 @@ mod lane_tests { tokens, cached, finish: finish_reason, + ids, }; } Some(TokenEvent::Error { message, .. }) => panic!("{name}: error: {message}"), @@ -2498,6 +2762,147 @@ 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(), 40); + 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, 40); + 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/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 1cf70cad5..8b1818c96 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) { @@ -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 @@ -575,6 +604,20 @@ 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) + } + + 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 @@ -871,6 +914,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 @@ -896,7 +940,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, @@ -1643,6 +1687,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, @@ -1651,6 +1727,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, @@ -1811,6 +1901,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)?; @@ -1828,7 +1939,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 \ @@ -1870,8 +1981,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, @@ -1897,6 +2008,16 @@ impl GemmaServe { self.final_logit_softcapping, head_normed, logits, + )?; + let (local_last, kv_chunk) = local_plan.decode_metadata_d_mut(); + ops::advance_decode_metadata( + ctx, + &mut global_tables.positions, + local_last, + &mut global_tables.pseudo_last, + kv_chunk, + rows, + self.global_split_factor, ) } @@ -1935,6 +2056,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)?; } @@ -2165,22 +2287,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 +2313,12 @@ 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.clear(); + arena.host.ids.extend_from_slice(tokens); + arena.host.ids.resize(padded, 0); + upload_prefix(ctx, &mut arena.ids, &arena.host.ids)?; + } Ok(padded) } @@ -2215,10 +2342,11 @@ 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]; - 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, @@ -2243,7 +2371,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, diff --git a/pegainfer-kernels/csrc/shared/elementwise.cu b/pegainfer-kernels/csrc/shared/elementwise.cu index 4b60e619f..cd8b6e58a 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,18 @@ 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) { + if (positions == nullptr || local_last == nullptr || pseudo_last == nullptr || + kv_chunk == nullptr || rows <= 0 || factor <= 0 || rows > INT_MAX / factor) { + return CUDA_ERROR_INVALID_VALUE; + } + int block = ADVANCE_DECODE_METADATA_BLOCK; + int grid = 1 + (rows - 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/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 8e9f38a1f..29fa7f56c 100644 --- a/pegainfer-kernels/src/ops/elementwise.rs +++ b/pegainfer-kernels/src/ops/elementwise.rs @@ -49,6 +49,53 @@ 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: &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) + .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 + && kv_chunk.len() >= rows + && pseudo_last.len() >= pseudo_rows, + "advance_decode_metadata: {rows} rows x {factor} exceeds a table" + ); + 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, + factor, + 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 diff --git a/pegainfer-kernels/src/tensor.rs b/pegainfer-kernels/src/tensor.rs index 74d39ddf6..b2c423a6a 100644 --- a/pegainfer-kernels/src/tensor.rs +++ b/pegainfer-kernels/src/tensor.rs @@ -79,6 +79,71 @@ pub fn active_cu_stream(ctx: &DeviceContext) -> CUstream { .unwrap_or_else(|| ctx.stream.cu_stream()) } +/// 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<()> { + 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 { + 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(); + } + err => return Err(anyhow::anyhow!("cuStreamQuery failed: {err:?}")), + } + } +} + +/// 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!( + !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 {}", + src.len(), + dst.len() + ); + if count == 0 { + return Ok(()); + } + 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::(), + ctx.stream.cu_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 6adcdd391..071871082 100644 --- a/pegainfer-sample/src/lib.rs +++ b/pegainfer-sample/src/lib.rs @@ -49,6 +49,9 @@ 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; + +const STAGED_READBACK_SLOTS: usize = 2; /// Allocate-once device buffers for [`select_batch`], sized for `max_rows` × `vocab`. /// @@ -72,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. @@ -109,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, @@ -118,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. @@ -207,7 +224,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() @@ -275,6 +294,119 @@ pub fn select_batch( Ok(tokens) } +fn validate_greedy_shape( + logits: &HiddenStates, + rows: usize, + scratch: &SampleScratch, +) -> Result<()> { + ensure!(rows > 0, "greedy_argmax_ids: empty batch"); + ensure!( + rows <= scratch.max_rows, + "greedy_argmax_ids: {rows} rows exceeds scratch capacity {}", + scratch.max_rows + ); + ensure!( + logits.seq_len >= rows && logits.hidden_dim == scratch.vocab, + "greedy_argmax_ids: logits shape {}x{} cannot serve {rows} rows x vocab {}", + logits.seq_len, + logits.hidden_dim, + scratch.vocab + ); + Ok(()) +} + +/// Capturable base-stream-only 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<()> { + 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, + 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) +} + +/// 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" + ); + 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_readback D2H stage failed: {e}")) +} + +/// 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, + rows: usize, + ids_out: &mut CudaSlice, + 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) +} + +/// 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 diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index 27aa46d21..268d216fb 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -60,6 +60,8 @@ 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,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" )