diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index ddfc6ac0..f579ff6f 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -631,6 +631,22 @@ enum ReservationDecision { Refused(String), } +type Newcomer = (GenerateRequest, GemmaKv, Option); + +#[derive(Clone, Copy)] +struct NewcomerOptions { + reserve_whole: bool, + evict_cache: bool, + can_wait: bool, + max_new_tokens: Option, +} + +enum PreparedNewcomer { + Ready(Newcomer, usize), + Done, + Requeue(Submitted), +} + /// One row of a step's sampler call. A mid-walk segment's row is sampled and /// discarded, so it carries `ignore_eos` whatever its request asked and a /// `logprobs` of 0: it never stops and is never scored. @@ -806,38 +822,15 @@ impl Active { 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 { + let stop = policy.stops(token, self.request.params.ignore_eos); + self.stopping = deliver_decode_row( + self, + DecodeToken { 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; + stop, + }, + ); } } @@ -869,6 +862,17 @@ fn send_scheduled(request: &GenerateRequest, prompt_tokens: usize, cached_tokens .is_ok() } +fn reject_newcomer(request: &GenerateRequest, message: String) { + let prompt_tokens = request.prompt_tokens.len(); + if send_scheduled(request, prompt_tokens, 0) { + let _ = request.token_tx.send(TokenEvent::Rejected { + message, + prompt_tokens, + completion_tokens: 0, + }); + } +} + /// Everything the engine thread owns for the life of the process. CUDA state /// is thread-affine, so it is built here rather than handed in: a context or /// cuBLAS handle minted on the caller thread fails with invalid-handle on the @@ -936,6 +940,7 @@ impl EngineState { &mut self, kv: &mut GemmaKv, need: AdmissionNeed, + evict_cache: bool, can_wait: bool, ) -> ReservationDecision { loop { @@ -958,10 +963,11 @@ impl EngineState { let Some(message) = refusal else { return ReservationDecision::Ready; }; - if self - .prefix_cache - .as_mut() - .is_some_and(PrefixCache::evict_lru) + if evict_cache + && self + .prefix_cache + .as_mut() + .is_some_and(PrefixCache::evict_lru) { continue; } @@ -973,6 +979,65 @@ impl EngineState { } } + fn resolve_newcomer_kv(&mut self, request: &GenerateRequest) -> (GemmaKv, Option) { + match self + .prefix_cache + .as_mut() + .and_then(|cache| cache.resolve(&request.prompt_tokens)) + { + Some((entry, t)) => match self.serve.restore_from_checkpoint(&self.ctx, entry, t) { + Ok(kv) => (kv, Some(entry.id)), + Err(err) => { + log::warn!("gemma4 prefix-cache restore failed (falling back): {err:#}"); + (self.serve.alloc_kv(), None) + } + }, + None => (self.serve.alloc_kv(), None), + } + } + + fn prepare_newcomer(&mut self, item: Submitted, options: NewcomerOptions) -> PreparedNewcomer { + let (request, prefix) = item; + if request.token_tx.is_closed() { + return PreparedNewcomer::Done; + } + let context_len = match validate_request(&request, prefix.hit_tokens(), self.max_context) { + Ok(len) => len, + Err(message) => { + reject_newcomer(&request, message); + return PreparedNewcomer::Done; + } + }; + let (mut kv, resumed) = self.resolve_newcomer_kv(&request); + let new_tokens = request.prompt_tokens.len() - kv.local.seq_len(); + if options + .max_new_tokens + .is_some_and(|limit| new_tokens > limit) + { + return PreparedNewcomer::Requeue((request, prefix)); + } + let need = if options.reserve_whole { + AdmissionNeed::Tokens(new_tokens) + } else { + AdmissionNeed::GlobalPages(global_account_pages(context_len)) + }; + match self.reserve_with_eviction(&mut kv, need, options.evict_cache, options.can_wait) { + ReservationDecision::Ready => {} + ReservationDecision::Requeue => { + return PreparedNewcomer::Requeue((request, prefix)); + } + ReservationDecision::Refused(message) => { + reject_newcomer(&request, message); + return PreparedNewcomer::Done; + } + } + let prompt_tokens = request.prompt_tokens.len(); + if !send_scheduled(&request, prompt_tokens, kv.local.seq_len()) { + return PreparedNewcomer::Done; + } + PreparedNewcomer::Ready((request, kv, resumed), new_tokens) + } + fn load( dir: &str, device: usize, @@ -1180,47 +1245,6 @@ impl EngineState { pending: &mut VecDeque, attempts: &mut usize, ) -> Admitted { - let (request, prefix) = item; - let sink = request.token_tx.clone(); - if sink.is_closed() { - return Admitted::Done; - } - let prompt_tokens = request.prompt_tokens.len(); - // Scheduled is paired with whatever ends the request, so a refusal - // emits it first rather than leaving the client with no lifecycle. - let reject = |message: String| { - if send_scheduled(&request, prompt_tokens, 0) { - let _ = sink.send(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens: 0, - }); - } - Admitted::Done - }; - let context_len = match validate_request(&request, prefix.hit_tokens(), self.max_context) { - Ok(len) => len, - Err(message) => return reject(message), - }; - - let mut resumed = None; - let mut kv = match self - .prefix_cache - .as_mut() - .and_then(|cache| cache.resolve(&request.prompt_tokens)) - { - Some((entry, t)) => match self.serve.restore_from_checkpoint(&self.ctx, entry, t) { - Ok(kv) => { - resumed = Some(entry.id); - kv - } - Err(err) => { - log::warn!("gemma4 prefix-cache restore failed (falling back): {err:#}"); - self.serve.alloc_kv() - } - }, - None => self.serve.alloc_kv(), - }; // The lane prefills whole on its own stream, so a lane-bound // admission still reserves everything up front. A chunked // admission reserves nothing here: every segment admits its own @@ -1228,25 +1252,19 @@ impl EngineState { // — parked first segments across several walkers would exhaust // the one shared segment transient the pool provisions. let lane_takes = self.lane.is_some() && !active.is_empty(); - let need = if self.mix_chunk.is_none() || lane_takes { - AdmissionNeed::Tokens(prompt_tokens - kv.local.seq_len()) - } else { - // A chunked admission reserves per segment; the whole-account - // door (see `global_account_pages`) still answers up front. - AdmissionNeed::GlobalPages(global_account_pages(context_len)) + let options = NewcomerOptions { + reserve_whole: self.mix_chunk.is_none() || lane_takes, + evict_cache: true, + can_wait, + max_new_tokens: None, }; - match self.reserve_with_eviction(&mut kv, need, can_wait) { - ReservationDecision::Ready => {} - ReservationDecision::Requeue => { - return Admitted::Requeue(Box::new((request, prefix))); - } - ReservationDecision::Refused(message) => return reject(message), - } - // A restored prefix is what the bridge reports as cached: the - // resumed KV's frontier is exactly the token count served from it. - if !send_scheduled(&request, prompt_tokens, kv.local.seq_len()) { - return Admitted::Done; - } + let (request, mut kv, resumed) = match self.prepare_newcomer(item, options) { + PreparedNewcomer::Ready(newcomer, _) => newcomer, + PreparedNewcomer::Done => return Admitted::Done, + PreparedNewcomer::Requeue(item) => return Admitted::Requeue(Box::new(item)), + }; + let prompt_tokens = request.prompt_tokens.len(); + let sink = request.token_tx.clone(); // Overlapped admission: the prefill launches onto the lane stream // and this call returns immediately — decode steps continue while it @@ -1274,8 +1292,7 @@ impl EngineState { // queue of dead submissions cannot stall the decode round. // The row pricing runs after the prefix-cache resolve: a // warm candidate costs the step only its unseen suffix. - let mut newcomers: Vec<(GenerateRequest, GemmaKv, Option)> = - vec![(request, kv, resumed)]; + let mut newcomers: Vec = vec![(request, kv, resumed)]; let mut rows_budget = { let (_, kv, _) = &newcomers[0]; prompt_tokens - kv.local.seq_len() @@ -1285,91 +1302,32 @@ impl EngineState { && newcomers.len() + active.len() < self.slots && *attempts < self.slots { - let Some((cand, cand_prefix)) = pending.pop_front() else { + let Some(candidate) = pending.pop_front() else { break; }; *attempts += 1; - let cand_sink = cand.token_tx.clone(); - if cand_sink.is_closed() { - continue; - } - let cand_context_len = - match validate_request(&cand, cand_prefix.hit_tokens(), self.max_context) { - Ok(len) => len, - Err(message) => { - let n = cand.prompt_tokens.len(); - if send_scheduled(&cand, n, 0) { - let _ = cand_sink.send(TokenEvent::Rejected { - message, - prompt_tokens: n, - completion_tokens: 0, - }); - } - continue; - } - }; - let cand_len = cand.prompt_tokens.len(); - let mut cand_resumed = None; - let mut cand_kv = match self - .prefix_cache - .as_mut() - .and_then(|cache| cache.resolve(&cand.prompt_tokens)) - { - Some((entry, t)) => { - match self.serve.restore_from_checkpoint(&self.ctx, entry, t) { - Ok(kv) => { - cand_resumed = Some(entry.id); - kv - } - Err(err) => { - log::warn!( - "gemma4 prefix-cache restore failed (falling back): {err:#}" - ); - self.serve.alloc_kv() - } - } - } - None => self.serve.alloc_kv(), + let options = NewcomerOptions { + reserve_whole: self.mix_chunk.is_none(), + evict_cache: false, + can_wait: true, + // Lazy: a chunked gather's budget can exceed the + // gather rows, and the bound is unused there. + max_new_tokens: self + .mix_chunk + .is_none() + .then(|| MIX_GATHER_ROWS - rows_budget), }; - let new_tokens = cand_len - cand_kv.local.seq_len(); - if self.mix_chunk.is_some() { - // A chunked candidate reserves nothing locally: its - // segments admit their own pages inside the walk; - // the whole-account door still answers here. - let cand_global_want = global_account_pages(cand_context_len); - if cand_global_want - > cand_kv.global.held_pages() + self.serve.global_pool.available_pages() - { - pending.push_front((cand, cand_prefix)); - break; + match self.prepare_newcomer(candidate, options) { + PreparedNewcomer::Ready(newcomer, new_tokens) => { + rows_budget += new_tokens; + newcomers.push(newcomer); } - if !send_scheduled(&cand, cand_len, cand_kv.local.seq_len()) { - continue; + PreparedNewcomer::Done => {} + PreparedNewcomer::Requeue(candidate) => { + pending.push_front(candidate); + break; } - rows_budget += new_tokens; - newcomers.push((cand, cand_kv, cand_resumed)); - continue; } - if rows_budget + new_tokens > MIX_GATHER_ROWS { - pending.push_front((cand, cand_prefix)); - break; - } - if admit_tokens( - &self.serve.local_pool, - &self.serve.global_pool, - &mut cand_kv, - new_tokens, - ) - .is_err() - { - pending.push_front((cand, cand_prefix)); - break; - } - if !send_scheduled(&cand, cand_len, cand_kv.local.seq_len()) { - continue; - } - rows_budget += new_tokens; - newcomers.push((cand, cand_kv, cand_resumed)); } return self.mixed_admission(newcomers, active); } @@ -1654,7 +1612,7 @@ impl EngineState { fn mixed_walk( &mut self, chunk: usize, - newcomers: Vec<(GenerateRequest, GemmaKv, Option)>, + newcomers: Vec, active: &mut Vec, ) -> Admitted { let mut walkers: Vec = newcomers @@ -2018,14 +1976,13 @@ impl EngineState { /// `Done`. fn mixed_admission( &mut self, - mut newcomers: Vec<(GenerateRequest, GemmaKv, Option)>, + mut newcomers: Vec, active: &mut Vec, ) -> Admitted { if let Some(chunk) = self.mix_chunk { return self.mixed_walk(chunk, newcomers, active); } - let fail_newcomers = |newcomers: &mut Vec<(GenerateRequest, GemmaKv, Option)>, - message: &str| { + let fail_newcomers = |newcomers: &mut Vec, message: &str| { for (request, _, _) in newcomers.drain(..) { let _ = request.token_tx.send(TokenEvent::Error { message: message.to_string(), @@ -2201,6 +2158,46 @@ impl EngineState { } } +/// One decode row's sampled outcome. +struct DecodeToken { + id: u32, + logprob: Option, + stop: bool, +} + +fn deliver_decode_row(entry: &mut Active, token: DecodeToken) -> bool { + if token.stop { + let _ = entry.request.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Stop, + prompt_tokens: entry.prompt_tokens, + completion_tokens: entry.emitted, + }); + return true; + } + entry.emitted += 1; + if entry + .request + .token_tx + .send(TokenEvent::Token { + id: token.id, + logprob: token.logprob, + }) + .is_err() + { + return true; + } + if entry.emitted >= entry.request.max_tokens { + let _ = entry.request.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + prompt_tokens: entry.prompt_tokens, + completion_tokens: entry.emitted, + }); + return true; + } + entry.next = token.id; + false +} + /// Deliver one decode step's outcome to every active row and retire the /// finished ones — the event flow both the pure decode round and the mixed /// admission share; `row_base` is the row's offset into the step's logits @@ -2209,39 +2206,17 @@ impl EngineState { fn emit_decode_rows(active: &mut Vec, sampled: &mut SampledRows, row_base: usize) { let mut retire: Vec = Vec::new(); for (row, entry) in active.iter_mut().enumerate() { - if sampled.stops[row + row_base] { - let _ = entry.request.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: entry.prompt_tokens, - completion_tokens: entry.emitted, - }); - retire.push(row); - continue; - } - let token = sampled.picked[row + row_base]; - entry.emitted += 1; - if entry - .request - .token_tx - .send(TokenEvent::Token { - id: token, - logprob: sampled.logprobs[row + row_base].take(), - }) - .is_err() - { - retire.push(row); - continue; - } - if entry.emitted >= entry.request.max_tokens { - let _ = entry.request.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: entry.prompt_tokens, - completion_tokens: entry.emitted, - }); + let index = row + row_base; + if deliver_decode_row( + entry, + DecodeToken { + id: sampled.picked[index], + logprob: sampled.logprobs[index].take(), + stop: sampled.stops[index], + }, + ) { retire.push(row); - continue; } - entry.next = token; } for row in retire.into_iter().rev() { active.swap_remove(row); diff --git a/pegainfer-gemma4/src/serve.rs b/pegainfer-gemma4/src/serve.rs index 9deeed64..05f471be 100644 --- a/pegainfer-gemma4/src/serve.rs +++ b/pegainfer-gemma4/src/serve.rs @@ -548,6 +548,92 @@ struct SplitKvState { cap: usize, } +struct SplitKvSpec { + label: &'static str, + slots: usize, + rows: usize, + heads: usize, + head_dim: usize, + cap: usize, + chunk_size_d: CudaSlice, +} + +struct SplitKvLaunch<'a> { + metadata: ops::Hd512DecodeMetadata<'a>, + o_indptr_d: &'a CudaSlice, + valid_mask_d: &'a CudaSlice, + tmp_v: &'a mut CudaSlice, + tmp_s: &'a mut CudaSlice, + cap: usize, +} + +impl SplitKvState { + fn new(ctx: &DeviceContext, spec: SplitKvSpec) -> Result { + let label = spec.label; + let alloc = |what: &'static str| { + move |err| anyhow::anyhow!("{label} split {what} alloc failed: {err}") + }; + Ok(Self { + request_indices_d: ctx + .stream + .alloc_zeros(spec.slots) + .map_err(alloc("request indices"))?, + kv_tile_indices_d: ctx + .stream + .alloc_zeros(spec.slots) + .map_err(alloc("tile indices"))?, + chunk_size_d: spec.chunk_size_d, + o_indptr_d: ctx + .stream + .alloc_zeros(spec.rows + 1) + .map_err(alloc("o_indptr"))?, + valid_mask_d: ctx + .stream + .alloc_zeros(spec.slots) + .map_err(alloc("valid mask"))?, + tmp_v: ctx + .stream + .alloc_zeros(spec.slots * spec.heads * spec.head_dim) + .map_err(alloc("tmp_v"))?, + tmp_s: ctx + .stream + .alloc_zeros(spec.slots * spec.heads) + .map_err(alloc("tmp_s"))?, + cap: spec.cap, + }) + } + + fn upload_csr(&mut self, ctx: &DeviceContext, csr: &ops::SplitKvCsr) -> Result<()> { + upload_prefix(ctx, &mut self.request_indices_d, &csr.request_indices)?; + upload_prefix(ctx, &mut self.kv_tile_indices_d, &csr.kv_tile_indices)?; + upload_prefix(ctx, &mut self.o_indptr_d, &csr.o_indptr)?; + upload_prefix(ctx, &mut self.valid_mask_d, &csr.block_valid_mask) + } + + fn metadata<'a>( + &'a mut self, + page_indices: &'a CudaSlice, + page_indptr: &'a CudaSlice, + last_page_len: &'a CudaSlice, + ) -> SplitKvLaunch<'a> { + SplitKvLaunch { + metadata: ops::Hd512DecodeMetadata::new( + page_indices, + page_indptr, + last_page_len, + &self.request_indices_d, + &self.kv_tile_indices_d, + &self.chunk_size_d, + ), + o_indptr_d: &self.o_indptr_d, + valid_mask_d: &self.valid_mask_d, + tmp_v: &mut self.tmp_v, + tmp_s: &mut self.tmp_s, + cap: self.cap, + } + } +} + /// The host buffers a split-KV pseudo expansion fills, borrowed apart so one /// builder serves every step shape. struct PseudoTables<'a> { @@ -888,36 +974,18 @@ impl GemmaServe { .alloc_zeros(factor * max_rows) .map_err(alloc("global pseudo last-page lens"))?, }, - global_split: SplitKvState { - request_indices_d: ctx - .stream - .alloc_zeros(global_split_slots) - .map_err(alloc("global split request indices"))?, - kv_tile_indices_d: ctx - .stream - .alloc_zeros(global_split_slots) - .map_err(alloc("global split tile indices"))?, - chunk_size_d: global_chunk, - o_indptr_d: ctx - .stream - .alloc_zeros(factor * max_rows + 1) - .map_err(alloc("global split o_indptr"))?, - valid_mask_d: ctx - .stream - .alloc_zeros(global_split_slots) - .map_err(alloc("global split valid mask"))?, - tmp_v: ctx - .stream - .alloc_zeros( - global_split_slots * global_split_heads * self.global_geom.head_dim, - ) - .map_err(alloc("global split tmp_v"))?, - tmp_s: ctx - .stream - .alloc_zeros(global_split_slots * global_split_heads) - .map_err(alloc("global split tmp_s"))?, - cap: global_split_cap, - }, + global_split: SplitKvState::new( + ctx, + SplitKvSpec { + label: "global", + slots: global_split_slots, + rows: factor * max_rows, + heads: global_split_heads, + head_dim: self.global_geom.head_dim, + cap: global_split_cap, + chunk_size_d: global_chunk, + }, + )?, 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"))?, @@ -1501,13 +1569,10 @@ impl GemmaServe { scratch.q_prep.seq_len = factor * seq_len; scratch.attn.hidden_dim = q_dim / factor; scratch.attn.seq_len = factor * seq_len; - let meta = ops::Hd512DecodeMetadata::new( + let launch = split.metadata( &global_tables.pseudo_pages, &global_tables.pseudo_indptr, &global_tables.pseudo_last, - &split.request_indices_d, - &split.kv_tile_indices_d, - &split.chunk_size_d, ); ops::paged_attention_batch_decode_split_kv_hd512_into( ctx, @@ -1516,12 +1581,12 @@ impl GemmaServe { self.global_pool.buffer(), &self.global_pool.layout().kernel_layout(), family_layer, - &meta, - &split.o_indptr_d, - &split.valid_mask_d, - &mut split.tmp_v, - &mut split.tmp_s, - factor * seq_len * split.cap, + &launch.metadata, + launch.o_indptr_d, + launch.valid_mask_d, + launch.tmp_v, + launch.tmp_s, + factor * seq_len * launch.cap, &mut scratch.attn, geom.num_q_heads / factor, 1.0, @@ -1590,13 +1655,10 @@ impl GemmaServe { scratch.q_prep.seq_len = factor * seq_len; scratch.attn.hidden_dim = q_dim / factor; scratch.attn.seq_len = factor * seq_len; - let meta = ops::Hd512DecodeMetadata::new( + let launch = split.metadata( &global_tables.pseudo_pages, &global_tables.pseudo_indptr, &global_tables.pseudo_last, - &split.request_indices_d, - &split.kv_tile_indices_d, - &split.chunk_size_d, ); ops::paged_attention_batch_decode_split_kv_hd512_into( ctx, @@ -1605,12 +1667,12 @@ impl GemmaServe { self.global_pool.buffer(), &self.global_pool.layout().kernel_layout(), family_layer, - &meta, - &split.o_indptr_d, - &split.valid_mask_d, - &mut split.tmp_v, - &mut split.tmp_s, - factor * batch * split.cap, + &launch.metadata, + launch.o_indptr_d, + launch.valid_mask_d, + launch.tmp_v, + launch.tmp_s, + factor * batch * launch.cap, &mut scratch.attn, geom.num_q_heads / factor, 1.0, @@ -1664,19 +1726,7 @@ impl GemmaServe { upload_prefix(ctx, &mut global_tables.pseudo_pages, pages)?; upload_prefix(ctx, &mut global_tables.pseudo_indptr, indptr)?; upload_prefix(ctx, &mut global_tables.pseudo_last, last_lens)?; - upload_prefix( - ctx, - &mut global_split.request_indices_d, - &csr.request_indices, - )?; - upload_prefix( - ctx, - &mut global_split.kv_tile_indices_d, - &csr.kv_tile_indices, - )?; - upload_prefix(ctx, &mut global_split.o_indptr_d, &csr.o_indptr)?; - upload_prefix(ctx, &mut global_split.valid_mask_d, &csr.block_valid_mask)?; - Ok(()) + global_split.upload_csr(ctx, &csr) } fn decode_fingerprint(&self, kvs: &[&mut GemmaKv], padded: usize) -> Option {