From 6472c73a3c619bf5c2e296a9f0381a3b49a18447 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 20:02:19 +0100 Subject: [PATCH 1/7] perf(gemma4): an admission coalesce door protects the decode tail Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 91 +++++++++++++++++++++++++++++++++- 1 file changed, 90 insertions(+), 1 deletion(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 96cfe478b..373000137 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -195,6 +195,36 @@ fn parse_kv_fp8(raw: Option<&str>) -> Result { } } +/// The admission coalesce door: `PEGAINFER_ADMIT_COALESCE_MS=N` holds +/// arrivals that would ride a live decode batch for up to `N` ms, so one +/// window's arrivals share a single mixed admission instead of invading +/// the batch once each — a live stream's tail gap prices the number of +/// admission events, not their size. An idle engine admits on sight; the +/// door only prices arrivals that have someone to disturb. Unset (or +/// `0`/`off`) admits on sight everywhere. +fn admit_coalesce_ms() -> Result> { + match std::env::var("PEGAINFER_ADMIT_COALESCE_MS") { + Ok(raw) => parse_admit_coalesce_ms(&raw), + Err(std::env::VarError::NotPresent) => Ok(None), + Err(std::env::VarError::NotUnicode(_)) => { + anyhow::bail!("PEGAINFER_ADMIT_COALESCE_MS is not valid UTF-8") + } + } +} + +fn parse_admit_coalesce_ms(raw: &str) -> Result> { + let value = raw.trim().to_ascii_lowercase(); + match value.as_str() { + "" | "0" | "off" => Ok(None), + other => match other.parse::() { + Ok(ms) if (1..=2000).contains(&ms) => Ok(Some(std::time::Duration::from_millis(ms))), + _ => anyhow::bail!( + "PEGAINFER_ADMIT_COALESCE_MS={raw:?} not recognized (off | N ms, 1 <= N <= 2000)" + ), + }, + } +} + pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result { let dir = model_path .to_str() @@ -235,6 +265,7 @@ pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result = VecDeque::new(); let mut active: Vec = Vec::new(); let mut disconnected = false; + let mut coalesce_since: Option = None; 'engine: loop { loop { match submit_rx.try_recv() { @@ -283,7 +314,38 @@ pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result= state.slots => + { + let cohort = MIX_MAX_PROMPTS + .min(state.slots.saturating_sub(active.len())) + .max(1); + let gathered = pending.len() >= cohort; + let since = *coalesce_since.get_or_insert_with(std::time::Instant::now); + let open = gathered || since.elapsed() >= window; + if open { + coalesce_since = None; + } + open + } + _ => { + coalesce_since = None; + true + } + }; + if door_open { + state.admit_from_queue(&mut pending, &mut active); + } if !active.is_empty() { state.decode_round(&mut active); } @@ -843,6 +905,9 @@ struct EngineState { /// The decode-slot count the pools are budgeted for; requests past it /// queue. slots: usize, + /// The admission coalesce window; `None` unless + /// `PEGAINFER_ADMIT_COALESCE_MS` opted in at startup. + admit_coalesce: Option, } impl EngineState { @@ -901,6 +966,7 @@ impl EngineState { let max_context = serving_context(config.max_position_embeddings)?; let lane_mode = async_prefill_mode()?; let mix_chunk = mix_chunk_tokens(max_context)?; + let admit_coalesce = admit_coalesce_ms()?; let slots = decode_slots()?; let local_kv_storage = kv_fp8_storage()?; if max_context > MAX_CONTEXT { @@ -1036,6 +1102,7 @@ impl EngineState { mix_chunk, max_context, slots, + admit_coalesce, }) } @@ -2275,6 +2342,28 @@ mod knob_tests { assert!(parse_kv_fp8(Some("global")).is_err()); } + #[test] + fn admit_coalesce_parses_or_refuses() { + for off in ["off", "0", ""] { + assert_eq!(parse_admit_coalesce_ms(off).unwrap(), None); + } + assert_eq!( + parse_admit_coalesce_ms("300").unwrap(), + Some(std::time::Duration::from_millis(300)) + ); + assert_eq!( + parse_admit_coalesce_ms("1").unwrap(), + Some(std::time::Duration::from_millis(1)) + ); + assert_eq!( + parse_admit_coalesce_ms("2000").unwrap(), + Some(std::time::Duration::from_millis(2000)) + ); + for bad in ["0x", "2001", "abc"] { + assert!(parse_admit_coalesce_ms(bad).is_err(), "{bad:?} must refuse"); + } + } + #[test] fn chunk_mode_parses_or_refuses() { for off in ["", "0", "off", " OFF "] { From 5f7432ee49b12206cbe6d217eb4334b9a6fd49ff Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:11:43 +0100 Subject: [PATCH 2/7] refactor(gemma4): the door knob reads the shared env plumbing admit_coalesce_ms hand-rolled std::env::var and spelled its variable three times, and SERVING_KNOBS omitted it, so the test-env guard neither cleared nor restored the door. The name becomes a const beside its siblings, the getter reads through read_env, and the guard covers all six knobs. Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 34 +++++++++++----------------------- 1 file changed, 11 insertions(+), 23 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 373000137..3340d4f31 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -46,6 +46,7 @@ const MIX_CHUNK_TOKENS_ENV: &str = "PEGAINFER_MIX_CHUNK_TOKENS"; const MAX_CONTEXT_ENV: &str = "PEGAINFER_MAX_CONTEXT"; const DECODE_SLOTS_ENV: &str = "PEGAINFER_DECODE_SLOTS"; const KV_FP8_ENV: &str = "PEGAINFER_KV_FP8"; +const ADMIT_COALESCE_ENV: &str = "PEGAINFER_ADMIT_COALESCE_MS"; const MIN_CONTEXT: usize = 1024; const MIN_CHUNK_TOKENS: usize = 64; const CEILING_DOMAIN: usize = i32::MAX as usize; @@ -194,22 +195,8 @@ fn parse_kv_fp8(raw: Option<&str>) -> Result { Some(value) => anyhow::bail!("PEGAINFER_KV_FP8 supports only \"local\", got {value:?}"), } } - -/// The admission coalesce door: `PEGAINFER_ADMIT_COALESCE_MS=N` holds -/// arrivals that would ride a live decode batch for up to `N` ms, so one -/// window's arrivals share a single mixed admission instead of invading -/// the batch once each — a live stream's tail gap prices the number of -/// admission events, not their size. An idle engine admits on sight; the -/// door only prices arrivals that have someone to disturb. Unset (or -/// `0`/`off`) admits on sight everywhere. fn admit_coalesce_ms() -> Result> { - match std::env::var("PEGAINFER_ADMIT_COALESCE_MS") { - Ok(raw) => parse_admit_coalesce_ms(&raw), - Err(std::env::VarError::NotPresent) => Ok(None), - Err(std::env::VarError::NotUnicode(_)) => { - anyhow::bail!("PEGAINFER_ADMIT_COALESCE_MS is not valid UTF-8") - } - } + read_env(ADMIT_COALESCE_ENV)?.map_or(Ok(None), |raw| parse_admit_coalesce_ms(&raw)) } fn parse_admit_coalesce_ms(raw: &str) -> Result> { @@ -219,7 +206,7 @@ fn parse_admit_coalesce_ms(raw: &str) -> Result> { other => match other.parse::() { Ok(ms) if (1..=2000).contains(&ms) => Ok(Some(std::time::Duration::from_millis(ms))), _ => anyhow::bail!( - "PEGAINFER_ADMIT_COALESCE_MS={raw:?} not recognized (off | N ms, 1 <= N <= 2000)" + "{ADMIT_COALESCE_ENV}={raw:?} not recognized (off | N ms, 1 <= N <= 2000)" ), }, } @@ -2573,13 +2560,14 @@ mod lane_tests { } } - const SERVING_KNOBS: [&str; 6] = [ - "PEGAINFER_ASYNC_PREFILL", - "PEGAINFER_PREFIX_CACHE", - "PEGAINFER_MIX_CHUNK_TOKENS", - "PEGAINFER_MAX_CONTEXT", - "PEGAINFER_DECODE_SLOTS", - "PEGAINFER_KV_FP8", + const SERVING_KNOBS: [&str; 7] = [ + super::ASYNC_PREFILL_ENV, + super::PREFIX_CACHE_ENV, + super::MIX_CHUNK_TOKENS_ENV, + super::MAX_CONTEXT_ENV, + super::DECODE_SLOTS_ENV, + super::ADMIT_COALESCE_ENV, + super::KV_FP8_ENV, ]; /// Clear every serving knob, set `overrides`, and hand back the guard From f44af04eff7eabb2462cb949cfa7c0d2e0e31876 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:11:43 +0100 Subject: [PATCH 3/7] fix(gemma4): the coalesce door and the async prefill lane refuse to combine The lane flies one prefill at a time and admission stops while it is busy, so a door in front of it can only delay serialized launches. The combination now refuses in EngineState::load before a device opens, and the early-refusal gate covers it. Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 3340d4f31..02a4f0634 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -956,6 +956,11 @@ impl EngineState { let admit_coalesce = admit_coalesce_ms()?; let slots = decode_slots()?; let local_kv_storage = kv_fp8_storage()?; + anyhow::ensure!( + admit_coalesce.is_none() || lane_mode.is_none(), + "{ADMIT_COALESCE_ENV} and {ASYNC_PREFILL_ENV} cannot combine: the lane flies one \ + prefill at a time, so the door could only delay it" + ); if max_context > MAX_CONTEXT { anyhow::ensure!( mix_chunk.is_some(), @@ -2788,9 +2793,8 @@ mod lane_tests { super::EngineState::load(&dir, 0, policy, 0x5EED, true).expect("engine state") } - /// A raise without the chunk knob, and a raise with the overlap lane, - /// both refuse before the multi-GiB load — the startup policy the - /// serving doc promises. + /// Invalid knob combinations refuse before the multi-GiB load — the + /// startup policy the serving doc promises. #[test] #[ignore = "requires the pinned 12B checkpoint and --test-threads=1"] fn the_raise_refuses_without_its_prerequisites() { @@ -2818,6 +2822,18 @@ mod lane_tests { format!("{err:#}").contains("unsupported over"), "unexpected refusal: {err:#}" ); + let err = load(&[ + (super::ADMIT_COALESCE_ENV, "300"), + (super::ASYNC_PREFILL_ENV, "green:35"), + ]) + .err() + .expect("the coalesce door and async lane must refuse"); + assert!( + format!("{err:#}").contains("the door could only delay it"), + "unexpected refusal: {err:#}" + ); + } + } /// The slots boundary, driven at the roster edge the engine loop owns From cf2901b7675f9e4d7714333f72538db68ed8f709 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:11:43 +0100 Subject: [PATCH 4/7] test(gemma4): the door decision is a pure function under its own gates The door block moves onto CoalesceDoor with an injected clock: the depth gate, the capacity-bounded cohort, the early release, the window timeout and the clears are unit-tested, and a checkpoint gate drives a live roster through the closed arm - two arrivals wait, the third releases the burst, one intake pass admits the cohort. The contract note moves onto the type and states what the door buys: a burst of admissions, not a shared weight scan. Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 215 ++++++++++++++++++++++++++++----- scripts/gemma4_gates.sh | 1 + 2 files changed, 185 insertions(+), 31 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 02a4f0634..8a5b0ec42 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -212,6 +212,46 @@ fn parse_admit_coalesce_ms(raw: &str) -> Result> { } } +/// Holds arrivals that would invade a live decode batch so one window's +/// arrivals land as a back-to-back burst of admissions: the stream's tail +/// gap prices the number of interruptions. One mixed step merges extra +/// prompts only with chunking or while its leader is under +/// `MIX_GATHER_ROWS`. The cohort bounds free-slot capacity, not a batch +/// across completions; idle engines admit on sight and shallow batches skip. +struct CoalesceDoor { + window: std::time::Duration, + since: Option, +} + +impl CoalesceDoor { + fn new(window: std::time::Duration) -> Self { + Self { + window, + since: None, + } + } + + fn opens( + &mut self, + pending: usize, + active: usize, + slots: usize, + now: std::time::Instant, + ) -> bool { + if pending == 0 || active == 0 || (active + pending) * 2 < slots { + self.since = None; + return true; + } + let cohort = MIX_MAX_PROMPTS.min(slots.saturating_sub(active)).max(1); + let since = *self.since.get_or_insert(now); + let open = pending >= cohort || now.duration_since(since) >= self.window; + if open { + self.since = None; + } + open + } +} + pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result { let dir = model_path .to_str() @@ -252,7 +292,7 @@ pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result = VecDeque::new(); let mut active: Vec = Vec::new(); let mut disconnected = false; - let mut coalesce_since: Option = None; + let mut coalesce_door = state.admit_coalesce.map(CoalesceDoor::new); 'engine: loop { loop { match submit_rx.try_recv() { @@ -301,35 +341,14 @@ pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result= state.slots => - { - let cohort = MIX_MAX_PROMPTS - .min(state.slots.saturating_sub(active.len())) - .max(1); - let gathered = pending.len() >= cohort; - let since = *coalesce_since.get_or_insert_with(std::time::Instant::now); - let open = gathered || since.elapsed() >= window; - if open { - coalesce_since = None; - } - open - } - _ => { - coalesce_since = None; - true - } - }; + let door_open = coalesce_door.as_mut().is_none_or(|door| { + door.opens( + pending.len(), + active.len(), + state.slots, + std::time::Instant::now(), + ) + }); if door_open { state.admit_from_queue(&mut pending, &mut active); } @@ -2356,6 +2375,81 @@ mod knob_tests { } } + fn test_door() -> CoalesceDoor { + CoalesceDoor::new(std::time::Duration::from_millis(100)) + } + + #[test] + fn coalesce_door_idle_opens_and_clears_the_timer() { + let now = std::time::Instant::now(); + let mut door = test_door(); + door.since = Some(now); + assert!(door.opens(1, 0, 8, now)); + assert_eq!(door.since, None); + } + + #[test] + fn coalesce_door_empty_queue_opens() { + let now = std::time::Instant::now(); + let mut door = test_door(); + assert!(door.opens(0, 4, 8, now)); + assert_eq!(door.since, None); + } + + #[test] + fn coalesce_door_shallow_batch_opens() { + let now = std::time::Instant::now(); + let mut door = test_door(); + assert!(door.opens(1, 1, 8, now)); + assert_eq!(door.since, None); + } + + #[test] + fn coalesce_door_deep_under_cohort_batch_closes_and_pins_the_timer() { + let now = std::time::Instant::now(); + let mut door = test_door(); + assert!(!door.opens(1, 3, 8, now)); + assert_eq!(door.since, Some(now)); + assert!(!door.opens(1, 3, 8, now + std::time::Duration::from_millis(1))); + assert_eq!(door.since, Some(now)); + } + + #[test] + fn coalesce_door_full_cohort_opens_and_clears() { + let now = std::time::Instant::now(); + let mut door = test_door(); + door.since = Some(now); + assert!(door.opens(4, 4, 8, now)); + assert_eq!(door.since, None); + } + + #[test] + fn coalesce_door_elapsed_window_opens() { + let now = std::time::Instant::now(); + let mut door = test_door(); + let window = door.window; + assert!(!door.opens(1, 3, 8, now)); + assert!(door.opens(1, 3, 8, now + window)); + assert_eq!(door.since, None); + } + + #[test] + fn coalesce_door_full_roster_opens_for_one_capacity_bounded_arrival() { + let now = std::time::Instant::now(); + let mut door = test_door(); + assert!(door.opens(1, 8, 8, now)); + assert_eq!(door.since, None); + } + + #[test] + fn coalesce_door_freed_slots_rederive_the_cohort() { + let now = std::time::Instant::now(); + let mut door = test_door(); + assert!(door.opens(2, 6, 8, now)); + assert!(!door.opens(2, 5, 8, now)); + assert_eq!(door.since, Some(now)); + } + #[test] fn chunk_mode_parses_or_refuses() { for off in ["", "0", "off", " OFF "] { @@ -2466,6 +2560,7 @@ mod lane_tests { struct Drained { tokens: usize, cached: usize, + scheduled: usize, finish: FinishReason, ids: Vec, } @@ -2473,6 +2568,7 @@ mod lane_tests { fn drain(rx: &mut TokenStreamReceiver, name: &str) -> Drained { let mut tokens = 0; let mut cached = 0; + let mut scheduled = 0; let mut ids = Vec::new(); loop { match rx.blocking_recv().map(|(_, event)| event) { @@ -2480,12 +2576,16 @@ mod lane_tests { tokens += 1; ids.push(id); } - Some(TokenEvent::Scheduled { cached_tokens, .. }) => cached = cached_tokens, + Some(TokenEvent::Scheduled { cached_tokens, .. }) => { + cached = cached_tokens; + scheduled += 1; + } Some(TokenEvent::PromptTokens { .. } | TokenEvent::KvTransfer { .. }) => {} Some(TokenEvent::Finished { finish_reason, .. }) => { return Drained { tokens, cached, + scheduled, finish: finish_reason, ids, }; @@ -2834,6 +2934,59 @@ mod lane_tests { ); } + /// The production intake decision is driven with an injected clock so + /// this gate observes scheduling events without sleeping: two arrivals + /// wait behind a live roster, the capacity-bounded cohort releases when + /// its third arrives, and one intake pass admits the whole burst. + #[test] + #[ignore = "requires the pinned 12B checkpoint, a GPU, and --test-threads=1"] + fn the_coalesce_door_releases_one_admission_burst() { + let dir = crate::testkit::model_path(); + let policy = super::generation_policy(&dir).expect("policy"); + let _env = scoped_engine_env(&[ + (super::ADMIT_COALESCE_ENV, "2000"), + (super::DECODE_SLOTS_ENV, "4"), + ]); + let mut state = + super::EngineState::load(&dir, 0, policy, 0x5EED, true).expect("engine state"); + let mut pending = std::collections::VecDeque::new(); + let mut active = Vec::new(); + let (incumbent, mut incumbent_rx) = walk_request(ids(40, 1), 16); + pending.push_back((incumbent, pegainfer_frontend::engine::KvPrefix::none())); + state.admit_from_queue(&mut pending, &mut active); + assert_eq!(active.len(), 1, "the live roster is pinned"); + + let (second, mut second_rx) = walk_request(ids(40, 2), 4); + let (third, mut third_rx) = walk_request(ids(40, 3), 4); + pending.push_back((second, pegainfer_frontend::engine::KvPrefix::none())); + pending.push_back((third, pegainfer_frontend::engine::KvPrefix::none())); + let now = std::time::Instant::now(); + let mut door = super::CoalesceDoor::new(state.admit_coalesce.expect("door enabled")); + assert!(!door.opens(pending.len(), active.len(), state.slots, now)); + assert!(second_rx.try_recv().is_err()); + assert!(third_rx.try_recv().is_err()); + + let (fourth, mut fourth_rx) = walk_request(ids(40, 4), 4); + pending.push_back((fourth, pegainfer_frontend::engine::KvPrefix::none())); + assert!(door.opens(pending.len(), active.len(), state.slots, now)); + state.admit_from_queue(&mut pending, &mut active); + assert!(pending.is_empty(), "the release drains the cohort"); + assert_eq!(active.len(), 4, "every request is admitted"); + while !active.is_empty() { + state.decode_round(&mut active); + } + + let incumbent = drain(&mut incumbent_rx, "incumbent"); + let second = drain(&mut second_rx, "second"); + let third = drain(&mut third_rx, "third"); + let fourth = drain(&mut fourth_rx, "fourth"); + assert_eq!(incumbent.tokens, 16); + assert_eq!((second.tokens, third.tokens, fourth.tokens), (4, 4, 4)); + assert_eq!( + second.scheduled + third.scheduled + fourth.scheduled, + 3, + "the released walk emits one admission event per request" + ); } /// The slots boundary, driven at the roster edge the engine loop owns diff --git a/scripts/gemma4_gates.sh b/scripts/gemma4_gates.sh index d54f4d1e6..51a5645c5 100755 --- a/scripts/gemma4_gates.sh +++ b/scripts/gemma4_gates.sh @@ -61,6 +61,7 @@ GATES_DENSE_AND_ROUTED=( # reads a weight, so that one needs the config and nothing else. GATES_SERVING_CONTRACT=( "gpu,ckpt engine::lane_tests::the_gathered_lifecycle_completes" + "gpu,ckpt engine::lane_tests::the_coalesce_door_releases_one_admission_burst" "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" From 675f876cb6a0909be548c35c050ad79afe492e32 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Sun, 30 Aug 2026 23:11:43 +0100 Subject: [PATCH 5/7] docs(gemma4): the coalesce door states what it actually buys The knob joins the Key env vars list and the serving doc gains its section: what the door prices, the gather-budget caveat, the lane refusal, and the measured trade-off that keeps it off by default. Signed-off-by: Feathbow --- CLAUDE.md | 1 + docs/models/gemma4/serving.md | 6 ++++++ 2 files changed, 7 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 8a6209655..24c6769ad 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -46,6 +46,7 @@ cargo run --release --features glm52 -- --model-path models/GLM5.2 - `PEGAINFER_NVCC_JOBS` — override parallel nvcc job count - `PEGAINFER_KV_FP8` — gemma4 opt-in fp8 KV: `local` stores the sliding family's K/V as e4m3 at scale 1.0 (lossy; halves the local pool; refuses an enabled prefix cache; unset = byte-identical serving) - `PEGAINFER_PREFIX_CACHE` — gemma4 opt-in conversation prefix cache: `K` entries of captured prompt state resume multi-turn prompts (pre-allocated page budget; unset = off, byte-identical serving) +- `PEGAINFER_ADMIT_COALESCE_MS` — gemma4 opt-in admission coalesce door: `N` ms in `1..=2000` (`off`/`0`/unset = admit on sight), holds arrivals that would invade a live decode batch so a window's arrivals land as one admission burst; refuses the async prefill lane; merging into one mixed step needs the chunked walk or a sub-budget prompt - `PEGAINFER_ASYNC_PREFILL` — gemma4 opt-in overlap lane: `green:NN` prefills live-batch admissions on an SM-capped stream to protect decode tails (`shared` for comparison; unset = off; bad values refuse to start) - `PEGAINFER_MIX_CHUNK_TOKENS` — gemma4 opt-in chunked walk: a mixed admission computes at most `N` prompt rows per step (`64 <= N <` the serving ceiling; unset = whole-prompt steps; bad values refuse to start) - `PEGAINFER_MAX_CONTEXT` — gemma4 serving ceiling raise (default 8192, up to the checkpoint's 262144; a raise past the default needs `PEGAINFER_MIX_CHUNK_TOKENS` and refuses the async lane) diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index f89f5b2a1..1f4f3b4e5 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -110,6 +110,12 @@ The cache brings its own page budget, added to the pool lines above at startup. At `PEGAINFER_PREFIX_CACHE=16` the idle footprint measured **39242 MiB** against the 33034 MiB baseline — the difference is the pre-allocated cache budget. +## The admission coalesce door (opt-in) + +`PEGAINFER_ADMIT_COALESCE_MS=N` (`1..=2000`; unset, `off` or `0` admits on sight) holds arrivals that would invade a live decode batch, then releases a window's arrivals as one back-to-back admission burst. It prices the number of admission interruptions, not their size: whole prompts beyond the 512-row gather budget still take separate weight scans unless `PEGAINFER_MIX_CHUNK_TOKENS` enables the chunked walk. An idle engine admits immediately, and a shallow roster skips the door when `(active + pending) * 2 < slots`. + +A deep roster releases when the window expires or the pending queue reaches `min(4, slots - active)`, with a floor of one. That cohort is a capacity bound over the currently free slots, not a cross-completion batch. The door refuses to combine with `PEGAINFER_ASYNC_PREFILL`, whose single in-flight prefill could only be delayed by it. Measured under sustained load, c16 median TPOT improves about 8.5% for about +288 ms median TTFT; c8 pays about 5.7% throughput and about +19 ms TTFT. P99 ITL is flat to slightly worse everywhere, so the door remains off by default. + ## The async prefill lane (opt-in) When `PEGAINFER_ASYNC_PREFILL` is unset, serving uses the normal mixed-step path; when set, a live-batch admission's prefill moves onto its own stream so decode steps keep replaying while the prompt computes. Dense and routed checkpoints share this path at the default context ceiling. `green:NN` pins the lane to roughly NN% of the SMs via a Green Context — the cap is the mechanism: a `shared` lane's full-width prefill grids starve decode steps, and is kept only for comparison. An unrecognized value or an unviable SM partition refuses to start rather than silently degrading. From e7c7eb98c5151016c2e6050fd64a55273566b8a3 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 00:35:00 +0100 Subject: [PATCH 6/7] refactor(gemma4): the intake turn is one function the loop and its gate share The dedicated gate drove a privately built door and called admission itself, so deleting the door conditional, never constructing the door, or admitting through a closed door would all have passed. The door construction and the decision-plus-conditional-admission move onto EngineState (coalesce_door, intake_turn); the loop calls them, and the gate drives the same functions with an injected clock through four scenarios: a closed wait that schedules nothing while decode advances the incumbent, a cohort release draining one burst with one Scheduled event per waiter, a timeout release, and a full roster admitting nothing until a slot frees. Signed-off-by: Feathbow --- pegainfer-gemma4/src/engine.rs | 221 +++++++++++++++++++++++++++------ 1 file changed, 186 insertions(+), 35 deletions(-) diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index 8a5b0ec42..ddfc6ac0a 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -292,7 +292,7 @@ pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result = VecDeque::new(); let mut active: Vec = Vec::new(); let mut disconnected = false; - let mut coalesce_door = state.admit_coalesce.map(CoalesceDoor::new); + let mut door = state.coalesce_door(); 'engine: loop { loop { match submit_rx.try_recv() { @@ -341,17 +341,12 @@ pub(crate) fn start(model_path: &Path, options: &EngineLoadOptions) -> Result Option { + self.admit_coalesce.map(CoalesceDoor::new) + } + + fn intake_turn( + &mut self, + door: &mut Option, + pending: &mut VecDeque, + active: &mut Vec, + now: std::time::Instant, + ) -> bool { + let open = door + .as_mut() + .is_none_or(|door| door.opens(pending.len(), active.len(), self.slots, now)); + if open { + self.admit_from_queue(pending, active); + } + open + } + fn reserve_with_eviction( &mut self, kv: &mut GemmaKv, @@ -2934,10 +2949,9 @@ mod lane_tests { ); } - /// The production intake decision is driven with an injected clock so - /// this gate observes scheduling events without sleeping: two arrivals - /// wait behind a live roster, the capacity-bounded cohort releases when - /// its third arrives, and one intake pass admits the whole burst. + /// Drive the production intake turn with an injected clock: closed turns + /// preserve the live stream, cohort and timeout releases each drain one + /// burst, and an open door still respects slot capacity. #[test] #[ignore = "requires the pinned 12B checkpoint, a GPU, and --test-threads=1"] fn the_coalesce_door_releases_one_admission_burst() { @@ -2951,41 +2965,178 @@ mod lane_tests { super::EngineState::load(&dir, 0, policy, 0x5EED, true).expect("engine state"); let mut pending = std::collections::VecDeque::new(); let mut active = Vec::new(); - let (incumbent, mut incumbent_rx) = walk_request(ids(40, 1), 16); + let mut door = state.coalesce_door(); + let now = std::time::Instant::now(); + let (incumbent, mut incumbent_rx) = walk_request(ids(40, 1), 64); pending.push_back((incumbent, pegainfer_frontend::engine::KvPrefix::none())); - state.admit_from_queue(&mut pending, &mut active); - assert_eq!(active.len(), 1, "the live roster is pinned"); + assert!( + state.intake_turn(&mut door, &mut pending, &mut active, now), + "closed-wait setup: an idle roster opens" + ); + assert_eq!(active.len(), 1, "closed-wait setup: the roster is live"); + let mut incumbent_tokens = 0; + while let Ok((_, event)) = incumbent_rx.try_recv() { + if matches!(event, TokenEvent::Token { .. }) { + incumbent_tokens += 1; + } + } let (second, mut second_rx) = walk_request(ids(40, 2), 4); let (third, mut third_rx) = walk_request(ids(40, 3), 4); pending.push_back((second, pegainfer_frontend::engine::KvPrefix::none())); pending.push_back((third, pegainfer_frontend::engine::KvPrefix::none())); - let now = std::time::Instant::now(); - let mut door = super::CoalesceDoor::new(state.admit_coalesce.expect("door enabled")); - assert!(!door.opens(pending.len(), active.len(), state.slots, now)); - assert!(second_rx.try_recv().is_err()); - assert!(third_rx.try_recv().is_err()); + assert!( + !state.intake_turn(&mut door, &mut pending, &mut active, now), + "closed wait: two arrivals stay behind the door" + ); + assert!( + second_rx.try_recv().is_err(), + "closed wait: the second arrival is not scheduled" + ); + assert!( + third_rx.try_recv().is_err(), + "closed wait: the third arrival is not scheduled" + ); + // The decode pipeline emits one step behind, so give the token a + // bounded number of rounds to surface. + let before_closed_decode = incumbent_tokens; + for _ in 0..3 { + state.decode_round(&mut active); + while let Ok((_, event)) = incumbent_rx.try_recv() { + if matches!(event, TokenEvent::Token { .. }) { + incumbent_tokens += 1; + } + } + if incumbent_tokens > before_closed_decode { + break; + } + } + assert!( + incumbent_tokens > before_closed_decode, + "closed wait: decode advances the incumbent between intake turns" + ); let (fourth, mut fourth_rx) = walk_request(ids(40, 4), 4); pending.push_back((fourth, pegainfer_frontend::engine::KvPrefix::none())); - assert!(door.opens(pending.len(), active.len(), state.slots, now)); - state.admit_from_queue(&mut pending, &mut active); - assert!(pending.is_empty(), "the release drains the cohort"); - assert_eq!(active.len(), 4, "every request is admitted"); - while !active.is_empty() { + assert!( + state.intake_turn(&mut door, &mut pending, &mut active, now), + "cohort release: the third arrival opens the door" + ); + assert!( + pending.is_empty(), + "cohort release: one turn drains the queue" + ); + assert_eq!(active.len(), 4, "cohort release: every waiter is admitted"); + while active.len() > 1 { state.decode_round(&mut active); } - - let incumbent = drain(&mut incumbent_rx, "incumbent"); let second = drain(&mut second_rx, "second"); let third = drain(&mut third_rx, "third"); let fourth = drain(&mut fourth_rx, "fourth"); - assert_eq!(incumbent.tokens, 16); - assert_eq!((second.tokens, third.tokens, fourth.tokens), (4, 4, 4)); assert_eq!( - second.scheduled + third.scheduled + fourth.scheduled, + (second.scheduled, third.scheduled, fourth.scheduled), + (1, 1, 1), + "cohort release: every waiter carries exactly one Scheduled event" + ); + assert_eq!( + active.len(), + 1, + "cohort release: the incumbent keeps its stream" + ); + + let timeout_start = now + std::time::Duration::from_secs(3); + let (timeout_a, mut timeout_a_rx) = walk_request(ids(40, 5), 4); + let (timeout_b, mut timeout_b_rx) = walk_request(ids(40, 6), 4); + pending.push_back((timeout_a, pegainfer_frontend::engine::KvPrefix::none())); + pending.push_back((timeout_b, pegainfer_frontend::engine::KvPrefix::none())); + assert!( + !state.intake_turn(&mut door, &mut pending, &mut active, timeout_start), + "timeout release: a fresh sub-cohort starts closed" + ); + let timeout_release = timeout_start + state.admit_coalesce.expect("door enabled"); + assert!( + state.intake_turn(&mut door, &mut pending, &mut active, timeout_release), + "timeout release: the window opens the next intake turn" + ); + assert!( + pending.is_empty(), + "timeout release: one turn drains the waiters" + ); + while active.len() > 1 { + state.decode_round(&mut active); + } + assert_eq!( + ( + drain(&mut timeout_a_rx, "timeout a").scheduled, + drain(&mut timeout_b_rx, "timeout b").scheduled, + ), + (1, 1), + "timeout release: every waiter carries one Scheduled event" + ); + + // The injected clock stays monotonic past the timeout release. + let free_slot_now = timeout_release + std::time::Duration::from_secs(1); + let (short, mut short_rx) = walk_request(ids(40, 7), 2); + let (long_a, mut long_a_rx) = walk_request(ids(40, 8), 8); + let (long_b, mut long_b_rx) = walk_request(ids(40, 9), 8); + pending.push_back((short, pegainfer_frontend::engine::KvPrefix::none())); + pending.push_back((long_a, pegainfer_frontend::engine::KvPrefix::none())); + pending.push_back((long_b, pegainfer_frontend::engine::KvPrefix::none())); + assert!( + state.intake_turn(&mut door, &mut pending, &mut active, free_slot_now), + "one free slot setup: a full cohort opens" + ); + assert_eq!(active.len(), 4, "one free slot setup: the roster is full"); + let (replacement, mut replacement_rx) = walk_request(ids(40, 10), 4); + pending.push_back((replacement, pegainfer_frontend::engine::KvPrefix::none())); + assert!( + state.intake_turn(&mut door, &mut pending, &mut active, free_slot_now), + "one free slot: the capacity-bounded floor opens the turn" + ); + assert_eq!( + pending.len(), + 1, + "one free slot: a full roster admits nothing" + ); + assert!( + replacement_rx.try_recv().is_err(), + "one free slot: the replacement is not scheduled while full" + ); + state.decode_round(&mut active); + assert_eq!( + active.len(), 3, - "the released walk emits one admission event per request" + "one free slot: exactly one incumbent finishes" + ); + assert!( + state.intake_turn(&mut door, &mut pending, &mut active, free_slot_now), + "one free slot: the next turn opens for the replacement" + ); + assert!( + pending.is_empty(), + "one free slot: the replacement is admitted" + ); + while !active.is_empty() { + state.decode_round(&mut active); + } + assert_eq!( + drain(&mut replacement_rx, "replacement").scheduled, + 1, + "one free slot: the replacement carries one Scheduled event" + ); + let short = drain(&mut short_rx, "short incumbent"); + let long_a = drain(&mut long_a_rx, "long incumbent a"); + let long_b = drain(&mut long_b_rx, "long incumbent b"); + let incumbent = drain(&mut incumbent_rx, "incumbent"); + assert_eq!( + incumbent_tokens + incumbent.tokens, + 64, + "one free slot: the original incumbent keeps its whole stream" + ); + assert_eq!( + (short.tokens, long_a.tokens, long_b.tokens), + (2, 8, 8), + "one free slot: every incumbent keeps its stream" ); } From 7c96bdbc23f37078f5c4b2238fde678b1bffd4e3 Mon Sep 17 00:00:00 2001 From: Feathbow Date: Mon, 31 Aug 2026 00:35:00 +0100 Subject: [PATCH 7/7] docs(gemma4): the door's release granularity and the index row The serving doc states the release lands at the first intake after the window elapses (checked once per engine iteration, so the wait can exceed N by up to one decode round), and the index row for the serving doc gains the knob. Signed-off-by: Feathbow --- docs/index.md | 2 +- docs/models/gemma4/serving.md | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/index.md b/docs/index.md index 63a5950b1..021ca864a 100644 --- a/docs/index.md +++ b/docs/index.md @@ -64,7 +64,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | | `models/gemma4/tokenizer.md` | Gemma 4 tokenizer/chat-template contracts, gated against a Hugging Face reference (token ids plus all five chat renders, content flattened to strings): BOS comes only from the standalone `chat_template.jinja` (which opens a thought channel and accepts a native system role), EOS is declared three times with three values, the published defaults are sampled rather than greedy, image/audio tokens encode straight from user text so text-only serving must reject them at admission, and one divergence stays open — the server's default content format adds a trailing space to system turns. | -| `models/gemma4/serving.md` | What the engine promises under load: iteration level scheduling with a ceiling on prompt plus output (8192 by default), prompts prefilled whole at a step boundary by default, requests beyond the configured decode slots (16 by default) queued rather than refused, and the two KV families budgeted separately (7.27 GiB sliding + 2.00 GiB global at 12B with the chunk knob off; the sliding budget shrinks to window plus segment under it). The default configuration needs a 48 GiB card: 32.2 GiB resident before the first request. Clients must send `` themselves or the model degenerates. A row is bit-identical when its companions' content and lengths change under a fixed batch width trajectory, and moves when arrivals or retirements change that trajectory, so greedy output is reproducible for a workload rather than across workloads. An opt-in conversation prefix cache (`PEGAINFER_PREFIX_CACHE=K`) resumes multi-turn prompts from captured prompt state at a pre-allocated page cost. An opt-in chunked walk (`PEGAINFER_MIX_CHUNK_TOKENS=N`) walks admissions through shared segment steps with round-by-round page reservation, and `PEGAINFER_MAX_CONTEXT` raises the ceiling to the checkpoint's 262144 while `PEGAINFER_DECODE_SLOTS` trades concurrency for the memory that buys. Open: no cross-request prefix sharing, single GPU. | +| `models/gemma4/serving.md` | What the engine promises under load: iteration level scheduling with a ceiling on prompt plus output (8192 by default), prompts prefilled whole at a step boundary by default, requests beyond the configured decode slots (16 by default) queued rather than refused, and the two KV families budgeted separately (7.27 GiB sliding + 2.00 GiB global at 12B with the chunk knob off; the sliding budget shrinks to window plus segment under it). The default configuration needs a 48 GiB card: 32.2 GiB resident before the first request. Clients must send `` themselves or the model degenerates. A row is bit-identical when its companions' content and lengths change under a fixed batch width trajectory, and moves when arrivals or retirements change that trajectory, so greedy output is reproducible for a workload rather than across workloads. An opt-in conversation prefix cache (`PEGAINFER_PREFIX_CACHE=K`) resumes multi-turn prompts from captured prompt state at a pre-allocated page cost. An opt-in chunked walk (`PEGAINFER_MIX_CHUNK_TOKENS=N`) walks admissions through shared segment steps with round-by-round page reservation, and `PEGAINFER_MAX_CONTEXT` raises the ceiling to the checkpoint's 262144 while `PEGAINFER_DECODE_SLOTS` trades concurrency for the memory that buys. `PEGAINFER_ADMIT_COALESCE_MS=N` batches arrivals that would interrupt a live decode batch into one admission burst, defaults off, and refuses the async lane. Open: no cross-request prefix sharing, single GPU. | | `models/gemma4/hf-golden.md` | Three Hugging Face references for 12B. The base fixture carries layer-boundary activations at both ends of both layer types plus top-64 logprobs, over a single-token, a nine-token and a 1024-token (exactly the sliding window) case, and pins three facts the forward path has to match — the embedding scale is bf16 62.0 rather than `sqrt(3840)`, text attention is causal, and `layer_scalar` applies to the layer output after both residual adds. The window fixture goes past the window (1023/1024/1025/4096, teacher-forced) under both sdpa and eager. The long-context fixture takes the same comparison to 16384/32768 for the raised ceiling, sdpa-only, with the window fixture's dual-backend floor on loan. Regeneration is byte-identical and checked with sha256. | ## models / glm52 diff --git a/docs/models/gemma4/serving.md b/docs/models/gemma4/serving.md index 1f4f3b4e5..7cfe99951 100644 --- a/docs/models/gemma4/serving.md +++ b/docs/models/gemma4/serving.md @@ -114,7 +114,7 @@ At `PEGAINFER_PREFIX_CACHE=16` the idle footprint measured **39242 MiB** against `PEGAINFER_ADMIT_COALESCE_MS=N` (`1..=2000`; unset, `off` or `0` admits on sight) holds arrivals that would invade a live decode batch, then releases a window's arrivals as one back-to-back admission burst. It prices the number of admission interruptions, not their size: whole prompts beyond the 512-row gather budget still take separate weight scans unless `PEGAINFER_MIX_CHUNK_TOKENS` enables the chunked walk. An idle engine admits immediately, and a shallow roster skips the door when `(active + pending) * 2 < slots`. -A deep roster releases when the window expires or the pending queue reaches `min(4, slots - active)`, with a floor of one. That cohort is a capacity bound over the currently free slots, not a cross-completion batch. The door refuses to combine with `PEGAINFER_ASYNC_PREFILL`, whose single in-flight prefill could only be delayed by it. Measured under sustained load, c16 median TPOT improves about 8.5% for about +288 ms median TTFT; c8 pays about 5.7% throughput and about +19 ms TTFT. P99 ITL is flat to slightly worse everywhere, so the door remains off by default. +A deep roster releases when the window expires or the pending queue reaches `min(4, slots - active)`, with a floor of one. A full cohort releases before `N`; the timeout release lands no earlier than `N`, at the first intake turn after the window elapses — the engine drains its submission channel before each intake, so there is no hard bound on how much later. That cohort is a capacity bound over the currently free slots, not a cross-completion batch. The door refuses to combine with `PEGAINFER_ASYNC_PREFILL`, whose single in-flight prefill could only be delayed by it. Measured under sustained load, c16 median TPOT improves about 8.5% for about +288 ms median TTFT; c8 pays about 5.7% throughput and about +19 ms TTFT. P99 ITL is flat to slightly worse everywhere, so the door remains off by default. ## The async prefill lane (opt-in)