From 20a0beff1dd166a6fa32119ac20b653edac9c63c Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Wed, 26 Aug 2026 19:19:09 +0800 Subject: [PATCH 01/15] fix(qwen3): preserve explicit stop-token causes Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-frontend/src/engine/driver.rs | 3 + pegainfer-frontend/src/engine/ledger.rs | 35 ++++- pegainfer-frontend/src/engine/mod.rs | 2 + pegainfer-frontend/src/engine/step.rs | 7 + pegainfer-frontend/src/engine/stop.rs | 104 +++++++++++++++ pegainfer-frontend/src/vllm/bridge/stepped.rs | 122 +++++++++++++++++- pegainfer-frontend/src/vllm/wire.rs | 38 +++++- pegainfer-k3/src/scheduler/tests.rs | 3 + pegainfer-qwen3/src/executor.rs | 24 +++- pegainfer-qwen3/src/executor/spec.rs | 78 ++++++++++- pegainfer-qwen3/src/frontend_adapter.rs | 62 +++------ pegainfer-qwen3/src/frontend_adapter/tests.rs | 121 +++++++++++++++++ pegainfer-qwen3/src/scheduler.rs | 4 + pegainfer-qwen3/src/scheduler/effects.rs | 18 +-- pegainfer-qwen3/src/scheduler/plan.rs | 60 +++++---- pegainfer-qwen3/src/scheduler/resolve.rs | 57 ++++++-- pegainfer-qwen3/src/scheduler/test_support.rs | 111 +++++++++++++++- pegainfer-qwen3/src/scheduler/tests.rs | 79 +++++++++--- pegainfer-qwen3/src/speculative.rs | 9 +- pegainfer-qwen3/tests/common/harness.rs | 2 + pegainfer-sim/src/lib.rs | 3 + 21 files changed, 810 insertions(+), 132 deletions(-) create mode 100644 pegainfer-frontend/src/engine/stop.rs diff --git a/pegainfer-frontend/src/engine/driver.rs b/pegainfer-frontend/src/engine/driver.rs index b69619b69..aed54654b 100644 --- a/pegainfer-frontend/src/engine/driver.rs +++ b/pegainfer-frontend/src/engine/driver.rs @@ -108,6 +108,7 @@ mod tests { use super::super::step::Request; use super::super::step::RequestId; use super::super::step::Terminal; + use super::super::stop::StopPolicy; use super::*; use crate::engine::FinishReason; @@ -163,6 +164,7 @@ mod tests { Request { prompt_tokens: vec![1, 2], params: crate::sampler::SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -197,6 +199,7 @@ mod tests { terminal, Some(Terminal::Finished { reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 2, completion_tokens: 3, }) diff --git a/pegainfer-frontend/src/engine/ledger.rs b/pegainfer-frontend/src/engine/ledger.rs index 7926addd5..d3bed79e3 100644 --- a/pegainfer-frontend/src/engine/ledger.rs +++ b/pegainfer-frontend/src/engine/ledger.rs @@ -38,6 +38,7 @@ use super::step::RequestUpdate; use super::step::ScheduledInfo; use super::step::StepOutputs; use super::step::Terminal; +use super::stop::StopCause; /// One open account: the request's admission facts and running tally. The /// payload is not here — it went to the scheduler at `submit`; the account is @@ -239,12 +240,23 @@ impl RequestLedger { /// Finish the request. Token counts come from the ledger's tally. pub fn finish(&mut self, id: RequestId, reason: FinishReason) { + self.finish_with_cause(id, reason, None); + } + + /// Finish a request while preserving a typed token-level stop cause. + pub fn finish_with_cause( + &mut self, + id: RequestId, + reason: FinishReason, + stop_cause: Option, + ) { let account = self.close(id); let AccountState::Active { completion_tokens } = account.state else { panic!("finish on {id} before admission"); }; self.statement.entry(id).terminal = Some(Terminal::Finished { reason, + stop_cause, prompt_tokens: account.prompt_len, completion_tokens, }); @@ -279,6 +291,16 @@ impl RequestLedger { /// this step — tokens included — folds into the returned message, so late /// delivery cannot reorder against the step stream. pub fn defer_finish(&mut self, id: RequestId, reason: FinishReason) -> DeferredFinish { + self.defer_finish_with_cause(id, reason, None) + } + + /// Defer a finish while preserving a typed token-level stop cause. + pub fn defer_finish_with_cause( + &mut self, + id: RequestId, + reason: FinishReason, + stop_cause: Option, + ) -> DeferredFinish { let account = self.close(id); let AccountState::Active { completion_tokens } = account.state else { panic!("defer_finish on {id} before admission"); @@ -289,6 +311,7 @@ impl RequestLedger { .unwrap_or_else(|| RequestUpdate::empty(id)); update.terminal = Some(Terminal::Finished { reason, + stop_cause, prompt_tokens: account.prompt_len, completion_tokens, }); @@ -374,6 +397,7 @@ mod tests { use super::super::request_lifecycle::StepReceiver; use super::super::step::Request; use super::super::step::Terminal; + use super::super::stop::StopPolicy; use super::super::wiring::SchedulerHandle; use super::super::wiring::scheduler_pair; use super::*; @@ -382,6 +406,7 @@ mod tests { Request { prompt_tokens: prompt, params: crate::sampler::SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 8, lora_adapter: None, kv_transfer_params: None, @@ -402,7 +427,9 @@ mod tests { backend.ledger.admit(id); backend.ledger.push_tokens(id, &[10, 11], &[]); backend.ledger.set_cached_tokens(id, 2); - backend.ledger.finish(id, FinishReason::Stop); + backend + .ledger + .finish_with_cause(id, FinishReason::Stop, Some(StopCause::Token(11))); backend.ledger.commit_step(); let mut steps = handle_steps(handle); @@ -419,6 +446,7 @@ mod tests { update.terminal, Some(Terminal::Finished { reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(11)), prompt_tokens: 3, completion_tokens: 2, }) @@ -479,7 +507,9 @@ mod tests { let id = backend.ledger.register(envelope).id; backend.ledger.admit(id); backend.ledger.push_tokens(id, &[7], &[]); - let deferred = backend.ledger.defer_finish(id, FinishReason::Length); + let deferred = backend + .ledger + .defer_finish_with_cause(id, FinishReason::Length, None); backend.ledger.commit_step(); let mut steps = handle_steps(handle); @@ -498,6 +528,7 @@ mod tests { update.terminal, Some(Terminal::Finished { reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 2, completion_tokens: 1, }) diff --git a/pegainfer-frontend/src/engine/mod.rs b/pegainfer-frontend/src/engine/mod.rs index 2c69a3d63..a2a8382fd 100644 --- a/pegainfer-frontend/src/engine/mod.rs +++ b/pegainfer-frontend/src/engine/mod.rs @@ -39,6 +39,7 @@ mod request; mod request_lifecycle; mod sink; mod step; +mod stop; mod wiring; pub use control::*; @@ -52,4 +53,5 @@ pub use request::*; pub use request_lifecycle::*; pub use sink::*; pub use step::*; +pub use stop::*; pub use wiring::*; diff --git a/pegainfer-frontend/src/engine/step.rs b/pegainfer-frontend/src/engine/step.rs index d36d80e46..923bda288 100644 --- a/pegainfer-frontend/src/engine/step.rs +++ b/pegainfer-frontend/src/engine/step.rs @@ -14,6 +14,8 @@ use std::time::Instant; use super::event::FinishReason; use super::event::TokenLogprob; +use super::stop::StopCause; +use super::stop::StopPolicy; /// In-process routing id for one generate request, minted by /// [`super::SchedulerHandle::submit`] from a per-scheduler counter. `Copy` and @@ -49,6 +51,7 @@ impl std::fmt::Display for RequestId { pub struct Request { pub prompt_tokens: Vec, pub params: crate::sampler::SamplingParams, + pub stop_policy: StopPolicy, pub max_tokens: usize, pub lora_adapter: Option, /// Opaque router/P-D metadata from the request's @@ -233,6 +236,10 @@ impl fmt::Display for RejectReason { pub enum Terminal { Finished { reason: FinishReason, + /// Present for token-driven stop finishes. The triggering token remains + /// in `RequestUpdate.tokens`, with its real logprob in the matching + /// `RequestUpdate.logprobs` entry. + stop_cause: Option, prompt_tokens: usize, completion_tokens: usize, }, diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs new file mode 100644 index 000000000..064088ba6 --- /dev/null +++ b/pegainfer-frontend/src/engine/stop.rs @@ -0,0 +1,104 @@ +/// How a request treats end-of-sequence tokens. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum EosPolicy { + /// Do not stop on model EOS tokens. + Ignore, + /// Use the model executor's configured EOS set. + #[default] + ModelDefault, + /// Stop only on this protocol-provided primary EOS token. + Token(u32), +} + +/// Request-scoped token stopping policy. +/// +/// EOS is kept separate from caller stop tokens because the vLLM protocol +/// reports them differently: EOS has no 'stop_reason', while a request stop +/// reports the actual matching token ID. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct StopPolicy { + pub eos: EosPolicy, + pub token_ids: Vec, +} + +impl StopPolicy { + /// Classify a token using vLLM's priority: EOS first, then the request's + /// explicit stop-token set. + #[must_use] + pub fn classify( + &self, + token_id: u32, + is_model_eos: impl FnOnce(u32) -> bool, + ) -> Option { + let is_eos = match self.eos { + EosPolicy::Ignore => false, + EosPolicy::ModelDefault => is_model_eos(token_id), + EosPolicy::Token(eos_token_id) => token_id == eos_token_id, + }; + + if is_eos { + Some(StopCause::Eos(token_id)) + } else if self.token_ids.contains(&token_id) { + Some(StopCause::Token(token_id)) + } else { + None + } + } +} + +/// The token-level cause of a normal stop finish. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StopCause { + /// A primary or model-default EOS token. + Eos(u32), + /// A token from the request's explicit stop-token set. + Token(u32), +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_default_classifies_model_eos() { + let policy = StopPolicy::default(); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Eos(99)) + ); + } + + #[test] + fn ignored_eos_does_not_disable_an_explicit_stop() { + let policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![99], + }; + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Token(99)) + ); + } + + #[test] + fn eos_wins_when_the_same_id_is_also_an_explicit_stop() { + let policy = StopPolicy { + eos: EosPolicy::Token(99), + token_ids: vec![99], + }; + + assert_eq!(policy.classify(99, |_| false), Some(StopCause::Eos(99))); + } + + #[test] + fn unmatched_token_does_not_stop() { + let policy = StopPolicy { + eos: EosPolicy::Token(99), + token_ids: vec![42], + }; + + assert_eq!(policy.classify(7, |_| false), None); + } +} diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 48179e9cc..fee78413e 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -55,9 +55,11 @@ use crate::engine::RequestId; use crate::engine::RequestUpdate; use crate::engine::SchedulerHandle; use crate::engine::StepOutputs; +use crate::engine::StopCause; use crate::engine::Terminal; use crate::vllm::wire::convert_finish_reason; use crate::vllm::wire::convert_sampling; +use crate::vllm::wire::convert_stop_policy; use crate::vllm::wire::lora_adapter_from_sampling_params; use crate::vllm::wire::requested_logprobs; use crate::vllm::wire::requested_prompt_logprobs; @@ -364,6 +366,7 @@ impl SteppedEngineBridge { None, ); } + let lora_adapter = match lora_adapter_from_sampling_params(&sampling_params) { Ok(adapter) => adapter, Err(error) => { @@ -385,6 +388,10 @@ impl SteppedEngineBridge { .as_ref() .and_then(|args| args.get("kv_transfer_params")) .cloned(); + // Older stepped model producers still suppress their terminal token + // and report only `FinishReason::Stop`. Keep the legacy sentinel for + // that producer shape; typed stop causes carry the real token and do + // not need a synthetic suffix. let stop_sentinel_id = stop_sentinel_id( sampling_params.eos_token_id, &sampling_params.stop_token_ids, @@ -400,9 +407,14 @@ impl SteppedEngineBridge { Span::noop() }; let trace_parent = SpanContext::from_span(&trace_root); + let control = self.scheduler.submit(Request { prompt_tokens, + // Keep the legacy SamplingParams lowering unchanged for stepped + // producers that have not migrated to StopPolicy. Qwen3 uses the + // independent policy below for stop classification. params: convert_sampling(&sampling_params), + stop_policy: convert_stop_policy(&sampling_params), max_tokens: sampling_params.max_tokens as usize, lora_adapter, kv_transfer_params, @@ -437,8 +449,9 @@ struct SteppedStream { /// P/D handoff metadata can arrive in an update with no token or /// terminal, so retain it until the next output carries it to the router. kv_transfer_params: Option, - /// The vLLM text decoder removes the final token from a stop-finished - /// output. Keep an EOS or explicit stop token as that removable sentinel. + /// Compatibility sentinel for stepped producers that predate typed + /// [`StopCause`]. New producers must include their triggering token in the + /// update and therefore bypass this fallback. stop_sentinel_id: Option, /// Request-lifetime root span; held only for its `Drop`, which closes the /// trace when the stream state is removed. @@ -556,11 +569,11 @@ fn reduce_update( let mut terminated = false; match update.terminal { None => {} - Some(Terminal::Finished { reason, .. }) => { - // PegaInfer suppresses EOS before emitting tokens, while vLLM's - // text decoder expects the terminal Stop output to contain EOS - // and unconditionally removes its final token. + Some(Terminal::Finished { + reason, stop_cause, .. + }) => { if reason == FinishReason::Stop + && stop_cause.is_none() && let Some(stop_sentinel_id) = state.stop_sentinel_id { token_ids.push(stop_sentinel_id); @@ -568,6 +581,10 @@ fn reduce_update( entries: Vec::new(), }); } + if let Some(StopCause::Token(token_id)) = stop_cause { + stop_reason = Some(StopReason::TokenId(token_id)); + } + finish_reason = Some(convert_finish_reason(reason)); terminated = true; } @@ -634,6 +651,8 @@ impl UnixAnchor { #[cfg(test)] mod tests { + use std::sync::atomic::AtomicBool; + use vllm_engine_core_client::protocol::output::EngineCoreOutputs; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; @@ -647,6 +666,7 @@ mod tests { Request { prompt_tokens: vec![1, 2], params: crate::sampler::SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 1, lora_adapter: None, kv_transfer_params: None, @@ -790,4 +810,94 @@ mod tests { ); assert!(backend.ledger.is_aborted(queued.id)); } + + #[test] + fn request_stop_maps_the_actual_token_and_preserves_its_logprob() { + let id = RequestId::new(7); + let control = RequestControl::new(id, Arc::new(AtomicBool::new(false))); + let mut state = SteppedStream::new("request-7".to_string(), control, Span::noop(), None); + + let mut update = RequestUpdate::empty(id); + update.tokens = vec![11, 43]; + update.logprobs = vec![ + None, + Some(TokenLogprob { + rank: 1, + logprob: -0.25, + top_logprobs: vec![(43, -0.25), (44, -1.0)], + }), + ]; + update.terminal = Some(Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(43)), + prompt_tokens: 16, + completion_tokens: 2, + }); + + let (output, terminated) = reduce_update(&mut state, update, &UnixAnchor::now()); + let output = output.expect("terminal output"); + + assert!(terminated); + assert_eq!(output.new_token_ids, vec![11, 43]); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); + assert_eq!(output.stop_reason, Some(StopReason::TokenId(43))); + + let direct = match output.new_logprobs.expect("stop-token logprob") { + MaybeWireLogprobs::Direct(direct) => direct, + MaybeWireLogprobs::Wire(_) => panic!("expected direct logprobs"), + }; + + assert_eq!(direct.positions.len(), 2); + assert_eq!(direct.positions[1].entries[0].token_id, 43); + assert!((direct.positions[1].entries[0].logprob + 0.25).abs() < f32::EPSILON); + } + + #[test] + fn model_eos_has_no_wire_stop_reason() { + let id = RequestId::new(8); + let control = RequestControl::new(id, Arc::new(AtomicBool::new(false))); + let mut state = SteppedStream::new("request-8".to_string(), control, Span::noop(), None); + + let mut update = RequestUpdate::empty(id); + update.tokens = vec![2]; + update.logprobs = vec![None]; + update.terminal = Some(Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(2)), + prompt_tokens: 16, + completion_tokens: 1, + }); + + let (output, terminated) = reduce_update(&mut state, update, &UnixAnchor::now()); + let output = output.expect("terminal output"); + + assert!(terminated); + assert_eq!(output.new_token_ids, vec![2]); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); + assert_eq!(output.stop_reason, None); + } + + #[test] + fn legacy_stop_without_typed_cause_keeps_the_wire_sentinel() { + let id = RequestId::new(9); + let control = RequestControl::new(id, Arc::new(AtomicBool::new(false))); + let mut state = + SteppedStream::new("request-9".to_string(), control, Span::noop(), Some(99)); + + let mut update = RequestUpdate::empty(id); + update.tokens = vec![11]; + update.terminal = Some(Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: None, + prompt_tokens: 16, + completion_tokens: 1, + }); + + let (output, terminated) = reduce_update(&mut state, update, &UnixAnchor::now()); + let output = output.expect("terminal output"); + + assert!(terminated); + assert_eq!(output.new_token_ids, vec![11, 99]); + assert_eq!(output.stop_reason, None); + } } diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index 2a8b2105a..7d1585f1c 100644 --- a/pegainfer-frontend/src/vllm/wire.rs +++ b/pegainfer-frontend/src/vllm/wire.rs @@ -9,8 +9,10 @@ use vllm_engine_core_client::protocol::logprobs::TokenLogprob as WireTokenLogpro use vllm_engine_core_client::protocol::output::EngineCoreFinishReason; use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use crate::engine::EosPolicy; use crate::engine::FinishReason; use crate::engine::PromptEcho; +use crate::engine::StopPolicy; use crate::engine::TokenLogprob; use crate::sampler::SamplingParams; @@ -79,8 +81,8 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar // None`, but `_all_stop_token_ids` always carries the model EOS set (it // exists for min_tokens masking, not stop detection). Deriving ignore_eos // from all_stop_token_ids would therefore void every ignore_eos request on - // models with a real EOS. Only `_eos_token_id` and the client's explicit - // `stop_token_ids` express a stop intent. + // models with a real EOS. Only _eos_token_id and the client's explicit + // stop_token_ids express the legacy scheduler's stop intent. let ignore_eos = params.eos_token_id.is_none() && params.stop_token_ids.is_empty(); if params.temperature <= 0.0 { return SamplingParams { @@ -111,7 +113,23 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar } } -/// Reject request parameters the engine would otherwise silently ignore. +pub(crate) fn convert_stop_policy(params: &EngineCoreSamplingParams) -> StopPolicy { + StopPolicy { + // Qwen3 owns the complete model EOS set in generation_config. The + // protocol's optional primary ID only tells us whether EOS is active; + // using it as a singleton would miss secondary model EOS IDs. + eos: params + .eos_token_id + .map_or(EosPolicy::Ignore, |_| EosPolicy::ModelDefault), + token_ids: params.stop_token_ids.clone(), + } +} + +/// Reject request parameters the frontend cannot represent faithfully. +/// +/// The stepped contract carries explicit request stop IDs independently in +/// [`StopPolicy`]; this helper only validates unrelated sampling/transfer +/// fields that would otherwise be silently ignored. /// Returns the offending description; `None` means the request is servable. /// /// The float comparisons are exact on purpose: they detect "the client sent @@ -233,13 +251,23 @@ mod tests { params.eos_token_id = Some(163_586); assert!(!convert_sampling(¶ms).ignore_eos); - // Explicit client stop tokens keep EOS detection on even when the - // frontend dropped _eos_token_id. + // The legacy scheduler keeps EOS active when an explicit stop token is + // present; the stepped bridge carries explicit stops in StopPolicy. params.eos_token_id = None; params.stop_token_ids = vec![42]; assert!(!convert_sampling(¶ms).ignore_eos); } + #[test] + fn convert_stop_policy_keeps_eos_and_explicit_stops_independent() { + let mut params = EngineCoreSamplingParams::for_test(); + params.eos_token_id = Some(99); + params.stop_token_ids = vec![11]; + + assert_eq!(convert_stop_policy(¶ms).eos, EosPolicy::ModelDefault); + assert_eq!(convert_stop_policy(¶ms).token_ids, vec![11]); + } + #[test] fn convert_sampling_passes_min_p_and_never_seed() { let mut params = EngineCoreSamplingParams::for_test(); diff --git a/pegainfer-k3/src/scheduler/tests.rs b/pegainfer-k3/src/scheduler/tests.rs index 6e779a206..052f7bb9b 100644 --- a/pegainfer-k3/src/scheduler/tests.rs +++ b/pegainfer-k3/src/scheduler/tests.rs @@ -23,6 +23,7 @@ use pegainfer_frontend::engine::Request; use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestUpdate; use pegainfer_frontend::engine::StepReceiver; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; @@ -175,6 +176,7 @@ fn request(prompt_len: usize, max_tokens: usize) -> Request { Request { prompt_tokens: vec![7; prompt_len], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -313,6 +315,7 @@ fn admitted_request_streams_its_tokens_and_finishes_at_max_tokens() { reason: FinishReason::Length, prompt_tokens: 4, completion_tokens: 3, + .. } ), "{terminal:?}" diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index be7744da1..fee03f956 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -19,6 +19,7 @@ use pegainfer_core::weight_loader::load_shard_info; use pegainfer_frontend::engine::DeferredFinish; use pegainfer_frontend::engine::LoadLoraAdapterRequest; use pegainfer_frontend::engine::SpecDecodeCounters; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::engine::panic_message; @@ -494,6 +495,7 @@ fn execute_step_on_lane( StepCommand::SpeculativeVerify { requests, kv_views, + stop_policies, sample_seed, } => { // One target forward over each request's K+1 draft span with a @@ -502,7 +504,8 @@ fn execute_step_on_lane( // token at each span position) and captures the target hidden states // (at the DFlash layers) to seed the next draft — all into reused, // pointer-stable scratch (`VerifyGraphBuffers`). - let result = lane.execute_dflash_verify(requests, kv_views, *sample_seed)?; + let result = + lane.execute_dflash_verify(requests, kv_views, stop_policies, *sample_seed)?; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -3605,8 +3608,15 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], + stop_policies: &[StopPolicy], sample_seed: u64, ) -> Result { + anyhow::ensure!( + stop_policies.len() == requests.len(), + "DFlash verify received {} stop policies for {} requests", + stop_policies.len(), + requests.len() + ); let capture_layer_ids = self.dflash_capture_layer_ids().ok_or_else(|| { anyhow::anyhow!("DFlash verify requested but no draft model is loaded") })?; @@ -3691,8 +3701,13 @@ impl LocalQwen3Lane { .flat_map(|req| std::iter::repeat_n(&req.params, req.as_slice().len())) .collect(); let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, sample_seed)?; - let request_results = build_verify_results(requests, &target_tokens)?; - + let mut request_results = build_verify_results(requests, &target_tokens)?; + // Apply the request policy before recording target hidden states; + // otherwise a suffix discarded by terminal handling would leak + // into the next DFlash draft context and acceptance counters. + for (policy, result) in stop_policies.iter().zip(&mut request_results) { + spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + } self.record_verify_dflash_context( requests, &request_results, @@ -3804,6 +3819,9 @@ enum StepCommand { SpeculativeVerify { requests: Vec, kv_views: Vec, + /// Request-local stop policies used before DFlash context recording; + /// these are host metadata and never enter the GPU batch. + stop_policies: Vec, sample_seed: u64, }, /// Speculative draft: roll the DFlash draft model forward one block per diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 488f261f0..a55ef96cb 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -6,6 +6,7 @@ //! module owns only the KV bookkeeping the executor thread is responsible for. use anyhow::Result; +use pegainfer_frontend::engine::StopPolicy; use super::Qwen3Executor; use super::RequestId; @@ -14,8 +15,31 @@ use super::WorkerStepOutcome; use crate::speculative::DraftPlan; use crate::speculative::DraftResult; use crate::speculative::VerifyPlan; +use crate::speculative::VerifyRequestResult; use crate::speculative::VerifyResult; +/// Remove accepted tokens after the first request-terminal token before the +/// speculative KV transaction commits. The scheduler classifies the same +/// retained trigger later to produce the typed protocol stop cause. +pub(super) fn truncate_after_terminal( + result: &mut VerifyRequestResult, + policy: &StopPolicy, + model_eos: &[u32], +) { + // Keep this helper idempotent: the worker trims before DFlash context is + // recorded, and the executor repeats the invariant before KV commit. + let Some(keep) = result.accepted_tokens.iter().position(|&token| { + policy + .classify(token, |id| model_eos.contains(&id)) + .is_some() + }) else { + return; + }; + let keep = keep + 1; + result.accepted_tokens.truncate(keep); + result.matched_draft_tokens = result.matched_draft_tokens.min(keep); +} + impl Qwen3Executor { pub(super) fn execute_speculative_verify_impl( &mut self, @@ -25,6 +49,12 @@ impl Qwen3Executor { self.speculative.is_some(), "speculative verification requested but no draft model is loaded" ); + anyhow::ensure!( + plan.stop_policies.len() == plan.requests.len(), + "speculative verify received {} stop policies for {} requests", + plan.stop_policies.len(), + plan.requests.len() + ); for req in plan.requests { anyhow::ensure!( !req.as_slice().is_empty(), @@ -71,6 +101,7 @@ impl Qwen3Executor { let step = StepCommand::SpeculativeVerify { requests: plan.requests.to_vec(), kv_views, + stop_policies: plan.stop_policies.to_vec(), sample_seed: plan.sample_seed, }; let outcome = match self.run_step(&step) { @@ -80,7 +111,7 @@ impl Qwen3Executor { return Err(e); } }; - let result = match outcome { + let mut result = match outcome { WorkerStepOutcome::SpeculativeVerify(result) => result, other => { self.revert_speculative_schedules(&scheduled); @@ -108,6 +139,12 @@ impl Qwen3Executor { )); } } + // The worker returns the mathematically accepted span. Apply the + // request contract before touching RequestKv so a terminal token's + // speculative suffix is rolled back with the unused reservation. + for (policy, req_result) in plan.stop_policies.iter().zip(&mut result.requests) { + truncate_after_terminal(req_result, policy, &self.metadata.stop_token_ids); + } // Commit the accepted prefix of each request's KV and free the rest. // On a mid-loop failure, only the not-yet-applied requests roll back @@ -191,3 +228,42 @@ impl Qwen3Executor { } } } + +#[cfg(test)] +mod tests { + use pegainfer_frontend::engine::EosPolicy; + + use super::*; + + #[test] + fn explicit_stop_truncates_the_kv_commit_after_the_trigger() { + let policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![7], + }; + let mut result = VerifyRequestResult { + request_id: RequestId::new(1), + matched_draft_tokens: 3, + accepted_tokens: vec![5, 7, 8, 9], + }; + + truncate_after_terminal(&mut result, &policy, &[99]); + + assert_eq!(result.accepted_tokens, vec![5, 7]); + assert_eq!(result.matched_draft_tokens, 2); + } + + #[test] + fn model_eos_truncates_but_keeps_the_trigger() { + let mut result = VerifyRequestResult { + request_id: RequestId::new(2), + matched_draft_tokens: 3, + accepted_tokens: vec![5, 99, 8, 9], + }; + + truncate_after_terminal(&mut result, &StopPolicy::default(), &[99]); + + assert_eq!(result.accepted_tokens, vec![5, 99]); + assert_eq!(result.matched_draft_tokens, 2); + } +} diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index e02125ea5..c89d8aec4 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -42,6 +42,7 @@ use pegainfer_frontend::engine::RejectReason; use pegainfer_frontend::engine::RequestLedger; use pegainfer_frontend::engine::Scheduler; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::engine::spawn_scheduler; use pegainfer_kernels::ops::NumericPolicy; use pegainfer_kernels::ops::numeric_policy; @@ -343,7 +344,7 @@ impl Qwen3Scheduler { // terminal rides the committed step, which the driver ships after // publishing metrics — the finishing batch's send-time stats then // read the drained occupancy instead of racing the publish. - let mut finishes: Vec<(RequestId, FinishReason)> = Vec::new(); + let mut finishes: Vec<(RequestId, FinishReason, Option)> = Vec::new(); for cached in effects.cached { if ledger.is_active(cached.request_id) { @@ -367,32 +368,11 @@ impl Qwen3Scheduler { for effect in effects.decode { match effect { DecodeEffect::Finish { - request_id, - finish_reason, - } => { - let Some(index) = self - .active - .iter() - .position(|req| req.request_id == request_id) - else { - continue; - }; - if ledger.is_active(request_id) { - if ledger.is_aborted(request_id) { - ledger.retire(request_id); - } else { - finishes.push((request_id, finish_reason)); - } - } - self.tracker.finish(request_id); - let _ = self.executor.drop_request(request_id); - to_retire.push(index); - } - DecodeEffect::EmitAndFinish { request_id, token, logprob, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -406,14 +386,14 @@ impl Qwen3Scheduler { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &[token], &[logprob]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); let _ = self.executor.drop_request(request_id); to_retire.push(index); } - DecodeEffect::EmitAndContinue { + DecodeEffect::Continue { request_id, token, logprob, @@ -436,7 +416,7 @@ impl Qwen3Scheduler { req.generated_count = completion_tokens; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id, tokens, completion_tokens, @@ -460,10 +440,11 @@ impl Qwen3Scheduler { } } } - DecodeEffect::EmitManyAndFinish { + DecodeEffect::FinishMany { request_id, tokens, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -477,7 +458,7 @@ impl Qwen3Scheduler { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &tokens, &[]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); @@ -507,32 +488,19 @@ impl Qwen3Scheduler { continued.push(req); } } - PendingEffect::Finish { - request_id, - finish_reason, - } => { - if ledger.is_active(request_id) { - if ledger.is_aborted(request_id) { - ledger.retire(request_id); - } else { - finishes.push((request_id, finish_reason)); - } - } - self.tracker.finish(request_id); - let _ = self.executor.drop_request(request_id); - } PendingEffect::EmitAndFinish { request_id, token, logprob, finish_reason, + stop_cause, } => { if ledger.is_active(request_id) { if ledger.is_aborted(request_id) { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &[token], &[logprob]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); @@ -560,12 +528,14 @@ impl Qwen3Scheduler { if self.executor.withholds_finishes() { let withheld: Vec = finishes .into_iter() - .map(|(request_id, reason)| ledger.defer_finish(request_id, reason)) + .map(|(request_id, reason, stop_cause)| { + ledger.defer_finish_with_cause(request_id, reason, stop_cause) + }) .collect(); self.executor.release_finished_events(withheld); } else { - for (request_id, reason) in finishes { - ledger.finish(request_id, reason); + for (request_id, reason, stop_cause) in finishes { + ledger.finish_with_cause(request_id, reason, stop_cause); } } } diff --git a/pegainfer-qwen3/src/frontend_adapter/tests.rs b/pegainfer-qwen3/src/frontend_adapter/tests.rs index 35441bc7d..0a89aac2b 100644 --- a/pegainfer-qwen3/src/frontend_adapter/tests.rs +++ b/pegainfer-qwen3/src/frontend_adapter/tests.rs @@ -89,6 +89,28 @@ impl StepCollector { } } + fn collect_terminal_with_logprobs( + &mut self, + id: RequestId, + ) -> ( + Vec, + Vec>, + Terminal, + ) { + let mut tokens = Vec::new(); + let mut logprobs = Vec::new(); + + loop { + let update = self.next_for(id); + tokens.extend_from_slice(&update.tokens); + logprobs.extend(update.logprobs); + + if let Some(terminal) = update.terminal { + return (tokens, logprobs, terminal); + } + } + } + /// Drain the remaining stream (until the scheduler is gone) and return /// every terminal seen for `id`. For asserting silence after an abort. fn drain_terminals_for(&mut self, id: RequestId) -> Vec { @@ -123,6 +145,105 @@ fn wait_until(timeout: Duration, mut predicate: impl FnMut() -> bool) -> bool { false } +#[test] +fn request_stop_token_beats_length_during_prefill_when_eos_is_ignored() { + let executor = FakeExecutor::new(4, Arc::new(Mutex::new(Vec::new()))).with_logprobs(); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 1); + req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; + req.stop_policy.token_ids = vec![100]; + + let control = partition.handle.submit(req); + let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); + + assert_eq!(tokens, vec![100]); + assert_eq!(logprobs.len(), 1); + assert!((logprobs[0].as_ref().expect("stop-token logprob").logprob + 0.1).abs() < f32::EPSILON); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(100)), + completion_tokens: 1, + .. + } + )); +} + +#[test] +fn request_stop_token_beats_length_during_decode_when_eos_is_ignored() { + let executor = FakeExecutor::new(4, Arc::new(Mutex::new(Vec::new()))).with_logprobs(); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 2); + req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; + req.stop_policy.token_ids = vec![200]; + + let control = partition.handle.submit(req); + let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); + + assert_eq!(tokens, vec![100, 200]); + assert_eq!(logprobs.len(), 2); + assert!((logprobs[1].as_ref().expect("stop-token logprob").logprob + 0.2).abs() < f32::EPSILON); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(200)), + completion_tokens: 2, + .. + } + )); +} + +#[test] +fn speculative_request_stop_beats_length_midspan_and_cleans_up() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(4, Arc::clone(&dropped)) + .with_stop_token(12) + .with_speculative_accepted_tokens(&[10, 11, 12, 13]); + + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 4); + req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; + req.stop_policy.token_ids = vec![12]; + + let control = partition.handle.submit(req); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!( + tokens, + vec![100, 10, 11, 12], + "the trigger is retained and the accepted suffix is discarded" + ); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(12)), + completion_tokens: 4, + .. + } + )); + + assert!( + wait_until(Duration::from_secs(1), || { + dropped.lock().unwrap().contains(&0) + }), + "stopped speculative request state should be dropped" + ); + + assert!( + wait_until(Duration::from_secs(1), || { + let metrics = partition.handle.metrics(); + metrics.num_running_reqs == 0 && metrics.kv_used_blocks == 0 + }), + "stopped speculative request should release scheduler state and KV blocks" + ); +} + #[test] fn unknown_lora_request_is_rejected_without_blocking_base_request() { let dropped = Arc::new(Mutex::new(Vec::new())); diff --git a/pegainfer-qwen3/src/scheduler.rs b/pegainfer-qwen3/src/scheduler.rs index efb212b18..5e915e706 100644 --- a/pegainfer-qwen3/src/scheduler.rs +++ b/pegainfer-qwen3/src/scheduler.rs @@ -18,6 +18,7 @@ use std::collections::HashSet; use log::debug; use log::warn; use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use crate::executor::ModelExecutor; @@ -34,6 +35,7 @@ pub(crate) struct ActiveRequestState { pub(crate) max_tokens: usize, pub(crate) prompt_len: usize, pub(crate) params: SamplingParams, + pub(crate) stop_policy: StopPolicy, /// Completion top-k count; None disables scoring, Some(0) scores only the chosen token. pub(crate) logprobs: Option, } @@ -47,6 +49,7 @@ pub(crate) struct PendingRequest { pub(crate) lora_adapter: Option, pub(crate) prompt_tokens: Vec, pub(crate) params: SamplingParams, + pub(crate) stop_policy: StopPolicy, pub(crate) max_tokens: usize, pub(crate) logprobs: Option, pub(crate) prompt_logprobs: Option, @@ -73,6 +76,7 @@ impl PendingRequest { lora_adapter: req.lora_adapter, prompt_tokens: req.prompt_tokens, params: req.params, + stop_policy: req.stop_policy, max_tokens: req.max_tokens, logprobs: req.logprobs, prompt_logprobs: req.prompt_logprobs, diff --git a/pegainfer-qwen3/src/scheduler/effects.rs b/pegainfer-qwen3/src/scheduler/effects.rs index 8e37a70e3..69d78411c 100644 --- a/pegainfer-qwen3/src/scheduler/effects.rs +++ b/pegainfer-qwen3/src/scheduler/effects.rs @@ -7,6 +7,7 @@ //! resolve logic stay a pure function of executor results. use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::engine::TokenLogprob; use super::ActiveRequestState; @@ -27,15 +28,12 @@ pub(crate) struct PromptEchoEffect { } pub(crate) enum PendingEffect { - Finish { - request_id: RequestId, - finish_reason: FinishReason, - }, EmitAndFinish { request_id: RequestId, token: u32, logprob: Option, finish_reason: FinishReason, + stop_cause: Option, }, Promote { state: ActiveRequestState, @@ -49,16 +47,13 @@ pub(crate) enum PendingEffect { pub(crate) enum DecodeEffect { Finish { - request_id: RequestId, - finish_reason: FinishReason, - }, - EmitAndFinish { request_id: RequestId, token: u32, logprob: Option, finish_reason: FinishReason, + stop_cause: Option, }, - EmitAndContinue { + Continue { request_id: RequestId, token: u32, logprob: Option, @@ -68,17 +63,18 @@ pub(crate) enum DecodeEffect { completion_tokens: usize, }, /// Commit several accepted speculative tokens and keep the request running. - EmitManyAndContinue { + ContinueMany { request_id: RequestId, tokens: Vec, completion_tokens: usize, }, /// Commit several accepted speculative tokens, then finish — a stop token or /// the max-output budget was hit partway through the accepted span. - EmitManyAndFinish { + FinishMany { request_id: RequestId, tokens: Vec, finish_reason: FinishReason, + stop_cause: Option, }, } diff --git a/pegainfer-qwen3/src/scheduler/plan.rs b/pegainfer-qwen3/src/scheduler/plan.rs index 89d82ccf3..cf472d4d9 100644 --- a/pegainfer-qwen3/src/scheduler/plan.rs +++ b/pegainfer-qwen3/src/scheduler/plan.rs @@ -118,9 +118,11 @@ pub(crate) fn execute_plan( requests: &draft_requests, })?; draft.requests.sort_by_key(|result| result.request_id); - let verify_requests = build_speculative_verify_items(active, &draft.requests); + let (verify_requests, stop_policies) = + build_speculative_verify_items(active, &draft.requests); let mut verify = executor.execute_speculative_verify(VerifyPlan { requests: &verify_requests, + stop_policies: &stop_policies, sample_seed: rand::RngExt::random(rng), })?; verify.requests.sort_by_key(|result| result.request_id); @@ -182,27 +184,35 @@ fn build_speculative_draft_items(active: &[ActiveRequestState]) -> Vec Vec { - draft_results - .iter() - .map(|draft| { - let active = active - .iter() - .find(|req| req.request_id == draft.request_id) - .expect("draft request_id must exist in active set"); - // Clamp the verify span to the request's remaining output budget so - // a long accepted run can't overshoot max_tokens. - let remaining = active.max_tokens.saturating_sub(active.generated_count); - // A continuing active request always has budget left (resolve emits - // EmitManyAndFinish the moment generated_count hits max_tokens), so - // this is a true invariant, not a runtime condition — don't crash the - // scheduler thread in release on a state we've proven unreachable. - debug_assert!(remaining > 0, "active request must have output budget"); - let mut token_ids = draft.token_ids.clone(); - token_ids.truncate(remaining); - VerifyStepItem::new(draft.request_id, token_ids, active.params) - }) - .collect() +) -> ( + Vec, + Vec, +) { + let mut requests = Vec::with_capacity(draft_results.len()); + let mut stop_policies = Vec::with_capacity(draft_results.len()); + for draft in draft_results { + let active = active + .iter() + .find(|req| req.request_id == draft.request_id) + .expect("draft request_id must exist in active set"); + // Clamp the verify span to the request's remaining output budget so + // a long accepted run can't overshoot max_tokens. + let remaining = active.max_tokens.saturating_sub(active.generated_count); + // A continuing active request always has budget left (resolve emits + // FinishMany the moment generated_count hits max_tokens), so + // this is a true invariant, not a runtime condition — don't crash the + // scheduler thread in release on a state we've proven unreachable. + debug_assert!(remaining > 0, "active request must have output budget"); + let mut token_ids = draft.token_ids.clone(); + token_ids.truncate(remaining); + requests.push(VerifyStepItem::new( + draft.request_id, + token_ids, + active.params, + )); + stop_policies.push(active.stop_policy.clone()); + } + (requests, stop_policies) } fn build_prefill_items(pending: &[PendingRequest], indices: &[usize]) -> Vec { @@ -253,6 +263,7 @@ fn sort_decode_results(results: &mut [crate::executor::DecodeRequestResult]) { #[cfg(test)] mod tests { + use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use super::*; @@ -264,6 +275,7 @@ mod tests { lora_adapter: None, prompt_tokens: vec![1, 2, 3], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 8, logprobs: None, prompt_logprobs: None, @@ -283,6 +295,7 @@ mod tests { max_tokens, prompt_len: 10, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: None, } } @@ -307,12 +320,13 @@ mod tests { token_ids: (0..16).collect(), }; - let verify = build_speculative_verify_items(&active, &[draft]); + let (verify, stop_policies) = build_speculative_verify_items(&active, &[draft]); assert_eq!(verify.len(), 1); // 32 - 24 = 8 remaining → the 16-token span truncates to 8. assert_eq!(verify[0].as_slice().len(), 8); assert_eq!(verify[0].as_slice(), (0..8).collect::>()); + assert_eq!(stop_policies, vec![StopPolicy::default()]); } // The plan selector is the whole batch-formation policy: what the scheduler diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 4cba623bb..75da390ae 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -1,4 +1,5 @@ use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; use super::ActiveRequestState; use super::PendingRequest; @@ -13,6 +14,24 @@ use crate::executor::ModelExecutor; use crate::executor::PrefillRequestResult; use crate::speculative::VerifyRequestResult; +fn stop_cause( + executor: &impl ModelExecutor, + req: &ActiveRequestState, + token: u32, +) -> Option { + req.stop_policy + .classify(token, |token_id| executor.is_stop_token(token_id)) +} + +fn pending_stop_cause( + executor: &impl ModelExecutor, + req: &PendingRequest, + token: u32, +) -> Option { + req.stop_policy + .classify(token, |token_id| executor.is_stop_token(token_id)) +} + pub(crate) fn resolve_step( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -60,26 +79,30 @@ pub(crate) fn resolve_speculative_outputs( .expect("speculative request_id must exist in active set"); let mut emitted = Vec::new(); let mut completion_tokens = req.generated_count; + for &token in &result.accepted_tokens { completion_tokens += 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(token); - if is_eos { - return DecodeEffect::EmitManyAndFinish { + emitted.push(token); + + if let Some(stop_cause) = stop_cause(executor, req, token) { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }; } - emitted.push(token); + if completion_tokens >= req.max_tokens { - return DecodeEffect::EmitManyAndFinish { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Length, + stop_cause: None, }; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id: result.request_id, tokens: emitted, completion_tokens, @@ -124,10 +147,13 @@ fn resolve_prefill_outputs( }); } - if !req.params.ignore_eos && executor.is_stop_token(result.first_token) { - effects.pending.push(PendingEffect::Finish { + if let Some(stop_cause) = pending_stop_cause(executor, &req, result.first_token) { + effects.pending.push(PendingEffect::EmitAndFinish { request_id: req.request_id, + token: result.first_token, + logprob: result.first_token_logprob, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }); continue; } @@ -138,6 +164,7 @@ fn resolve_prefill_outputs( token: result.first_token, logprob: result.first_token_logprob, finish_reason: FinishReason::Length, + stop_cause: None, }); continue; } @@ -152,6 +179,7 @@ fn resolve_prefill_outputs( max_tokens: req.max_tokens, prompt_len, params: req.params, + stop_policy: req.stop_policy, logprobs: req.logprobs, }, first_token: result.first_token, @@ -175,22 +203,27 @@ fn resolve_decode_outputs( .find(|req| req.request_id == result.request_id) .expect("decode request_id must exist in active set"); let completion_tokens = req.generated_count + 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(result.token); + let stop_cause = stop_cause(executor, req, result.token); let at_limit = completion_tokens >= req.max_tokens; - if is_eos { + + if let Some(stop_cause) = stop_cause { DecodeEffect::Finish { request_id: result.request_id, + token: result.token, + logprob: result.logprob.clone(), finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), } } else if at_limit { - DecodeEffect::EmitAndFinish { + DecodeEffect::Finish { request_id: result.request_id, token: result.token, logprob: result.logprob.clone(), finish_reason: FinishReason::Length, + stop_cause: None, } } else { - DecodeEffect::EmitAndContinue { + DecodeEffect::Continue { request_id: result.request_id, token: result.token, logprob: result.logprob.clone(), diff --git a/pegainfer-qwen3/src/scheduler/test_support.rs b/pegainfer-qwen3/src/scheduler/test_support.rs index e91ad8230..10dffd88b 100644 --- a/pegainfer-qwen3/src/scheduler/test_support.rs +++ b/pegainfer-qwen3/src/scheduler/test_support.rs @@ -9,6 +9,7 @@ use std::time::Duration; use anyhow::Result; use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::sampler::SamplingParams; @@ -23,6 +24,12 @@ use crate::executor::PrefillStepItem; use crate::executor::RequestId; use crate::executor::UnifiedPlan; use crate::executor::UnifiedResult; +use crate::speculative::DraftPlan; +use crate::speculative::DraftRequestResult; +use crate::speculative::DraftResult; +use crate::speculative::VerifyPlan; +use crate::speculative::VerifyRequestResult; +use crate::speculative::VerifyResult; pub(crate) struct FakeExecutor { pub(crate) block_size: usize, @@ -39,6 +46,8 @@ pub(crate) struct FakeExecutor { pub(crate) dropped: Arc>>, pub(crate) prefetch_offers: Arc>>, stop_token: Option, + emit_logprobs: bool, + speculative_accepted_tokens: Option>, } impl FakeExecutor { @@ -56,6 +65,8 @@ impl FakeExecutor { dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, + emit_logprobs: false, + speculative_accepted_tokens: None, } } @@ -64,6 +75,20 @@ impl FakeExecutor { self } + pub(crate) fn with_logprobs(mut self) -> Self { + self.emit_logprobs = true; + self + } + + pub(crate) fn with_speculative_accepted_tokens(mut self, tokens: &[u32]) -> Self { + assert!( + !tokens.is_empty(), + "fake speculative span must make progress" + ); + self.speculative_accepted_tokens = Some(tokens.to_vec()); + self + } + pub(crate) fn with_decode_failure(mut self) -> Self { self.fail_decode_once = true; self @@ -99,7 +124,13 @@ impl FakeExecutor { PrefillRequestResult { request_id: req.request_id, first_token: 100 + req.request_id.raw() as u32, - first_token_logprob: None, + first_token_logprob: self.emit_logprobs.then(|| { + let token = 100 + req.request_id.raw() as u32; + pegainfer_frontend::engine::TokenLogprob { + logprob: -0.1, + top_logprobs: vec![(token, -0.1)], + } + }), prompt_logprobs: None, cached_tokens: 0, completed, @@ -224,7 +255,13 @@ impl ModelExecutor for FakeExecutor { .map(|req| DecodeRequestResult { request_id: req.request_id, token: 200 + req.request_id.raw() as u32, - logprob: None, + logprob: self.emit_logprobs.then(|| { + let token = 200 + req.request_id.raw() as u32; + pegainfer_frontend::engine::TokenLogprob { + logprob: -0.2, + top_logprobs: vec![(token, -0.2)], + } + }), }) .collect(), }) @@ -260,6 +297,75 @@ impl ModelExecutor for FakeExecutor { .collect(), }) } + + fn execute_speculative_draft(&mut self, plan: DraftPlan<'_>) -> Result { + let accepted_tokens = self + .speculative_accepted_tokens + .as_ref() + .ok_or_else(|| anyhow::anyhow!("fake speculative decoding is disabled"))?; + + Ok(DraftResult { + requests: plan + .requests + .iter() + .map(|req| { + let mut token_ids = Vec::with_capacity(accepted_tokens.len()); + token_ids.push(req.current_token); + token_ids.extend( + accepted_tokens + .iter() + .copied() + .take(accepted_tokens.len().saturating_sub(1)), + ); + + DraftRequestResult { + request_id: req.request_id, + token_ids, + } + }) + .collect(), + }) + } + + fn execute_speculative_verify(&mut self, plan: VerifyPlan<'_>) -> Result { + let configured = self + .speculative_accepted_tokens + .clone() + .ok_or_else(|| anyhow::anyhow!("fake speculative decoding is disabled"))?; + + let mut requests = Vec::with_capacity(plan.requests.len()); + + for req in plan.requests { + let span_len = req.as_slice().len(); + anyhow::ensure!(span_len > 0, "fake speculative verify span is empty"); + + let accepted_tokens = configured[..configured.len().min(span_len)].to_vec(); + + let current_tokens = self + .held_tokens + .get(&req.request_id) + .copied() + .ok_or_else(|| anyhow::anyhow!("missing fake request state"))?; + + self.ensure_request_tokens(req.request_id, current_tokens + accepted_tokens.len())?; + + requests.push(VerifyRequestResult { + request_id: req.request_id, + matched_draft_tokens: accepted_tokens.len().saturating_sub(1), + accepted_tokens, + }); + } + + Ok(VerifyResult { requests }) + } + + fn speculative_enabled(&self) -> bool { + self.speculative_accepted_tokens.is_some() + } + + fn speculative_request_ready(&self, request_id: RequestId) -> bool { + self.speculative_accepted_tokens.is_some() && self.held_tokens.contains_key(&request_id) + } } /// A minimal contract request: `prompt_len` filler tokens, default sampling. @@ -267,6 +373,7 @@ pub(crate) fn request(prompt_len: usize, max_tokens: usize) -> Request { Request { prompt_tokens: vec![1; prompt_len], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen3/src/scheduler/tests.rs b/pegainfer-qwen3/src/scheduler/tests.rs index 774744714..6df3914c1 100644 --- a/pegainfer-qwen3/src/scheduler/tests.rs +++ b/pegainfer-qwen3/src/scheduler/tests.rs @@ -6,7 +6,10 @@ use std::sync::Arc; use std::sync::Mutex; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_kv_cache::BlockPool; use super::test_support::FakeExecutor; @@ -23,6 +26,7 @@ fn active_state(request_id: u64, generated_count: usize, max_tokens: usize) -> A max_tokens, prompt_len: 16, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: None, } } @@ -518,6 +522,14 @@ fn spec_active( ignore_eos, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: if ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + token_ids: Vec::new(), + }, ..active_state(id, generated_count, max_tokens) } } @@ -540,7 +552,7 @@ fn speculative_full_span_accept_continues() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndContinue { + effects::DecodeEffect::ContinueMany { request_id, tokens, completion_tokens, @@ -554,7 +566,7 @@ fn speculative_full_span_accept_continues() { "completion = prior generated + span len" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), } } @@ -567,41 +579,41 @@ fn speculative_stop_token_midspan_finishes_and_suppresses_eos() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, + stop_cause, .. }, ] => { - assert_eq!( - tokens, - &vec![10, 11], - "EOS itself is suppressed from emission" - ); + assert_eq!(tokens, &vec![10, 11, SPEC_EOS],); assert!(matches!(finish_reason, FinishReason::Stop)); + assert_eq!(*stop_cause, Some(StopCause::Eos(SPEC_EOS))); } - _ => panic!("expected EmitManyAndFinish(Stop)"), + _ => panic!("expected FinishMany(Stop)"), } } #[test] -fn speculative_stop_token_at_span_start_emits_nothing() { +fn speculative_stop_token_at_span_start_retains_the_token() { let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); let active = [spec_active(1, 5, 100, false)]; let results = [spec_result(1, vec![SPEC_EOS, 11, 12])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, + stop_cause, .. }, ] => { - assert!(tokens.is_empty(), "stop at position 0 emits no tokens"); + assert_eq!(tokens, &vec![SPEC_EOS]); assert!(matches!(finish_reason, FinishReason::Stop)); + assert_eq!(*stop_cause, Some(StopCause::Eos(SPEC_EOS))); } - _ => panic!("expected EmitManyAndFinish(Stop)"), + _ => panic!("expected FinishMany(Stop)"), } } @@ -614,7 +626,7 @@ fn speculative_max_tokens_truncates_midspan() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, .. @@ -627,7 +639,7 @@ fn speculative_max_tokens_truncates_midspan() { ); assert!(matches!(finish_reason, FinishReason::Length)); } - _ => panic!("expected EmitManyAndFinish(Length)"), + _ => panic!("expected FinishMany(Length)"), } } @@ -638,14 +650,43 @@ fn speculative_ignore_eos_does_not_stop() { let results = [spec_result(1, vec![SPEC_EOS, SPEC_EOS])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { - [effects::DecodeEffect::EmitManyAndContinue { tokens, .. }] => { + [effects::DecodeEffect::ContinueMany { tokens, .. }] => { assert_eq!( tokens, &vec![SPEC_EOS, SPEC_EOS], "ignore_eos passes stop tokens through" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), + } +} + +#[test] +fn speculative_request_stop_truncates_the_span_when_eos_is_ignored() { + let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); + + let mut request = spec_active(1, 0, 100, true); + request.stop_policy.token_ids = vec![12]; + + let active = [request]; + let results = [spec_result(1, vec![10, 11, 12, 13])]; + + let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); + + match &effects[..] { + [ + effects::DecodeEffect::FinishMany { + tokens, + finish_reason, + stop_cause, + .. + }, + ] => { + assert_eq!(tokens, &vec![10, 11, 12]); + assert_eq!(*finish_reason, FinishReason::Stop); + assert_eq!(*stop_cause, Some(StopCause::Token(12))); + } + _ => panic!("expected request stop to finish the speculative span"), } } @@ -660,11 +701,11 @@ fn speculative_resolves_each_request_independently() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); assert!(matches!( &effects[0], - effects::DecodeEffect::EmitManyAndContinue { request_id, .. } if *request_id == RequestId::new(1) + effects::DecodeEffect::ContinueMany { request_id, .. } if *request_id == RequestId::new(1) )); assert!(matches!( &effects[1], - effects::DecodeEffect::EmitManyAndFinish { request_id, finish_reason: FinishReason::Stop, .. } + effects::DecodeEffect::FinishMany { request_id, finish_reason: FinishReason::Stop, .. } if *request_id == RequestId::new(2) )); } diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index bea499611..a8050b42a 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -24,6 +24,7 @@ //! target distribution; acceptance only decides how many ride one step. use anyhow::Result; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use crate::executor::RequestId; @@ -56,6 +57,9 @@ impl VerifyStepItem { #[derive(Clone, Copy)] pub(crate) struct VerifyPlan<'a> { pub requests: &'a [VerifyStepItem], + /// Request-local stop policies in the same order as `requests`. They stay + /// executor-side and are not copied into the worker command or GPU batch. + pub stop_policies: &'a [StopPolicy], /// Engine step seed for the verify rows' sampler pass (same contract as /// decode: fresh per step; seeded rows re-mix their own request seed). pub sample_seed: u64, @@ -69,8 +73,9 @@ pub(crate) struct VerifyRequestResult { /// Tokens to commit: the accepted draft prefix followed by the target's /// posterior token at the first mismatch (or the block-end continuation /// when every draft is accepted). Always `1..=K + 1` tokens, so a verify - /// step always makes at least one token of progress. The scheduler still - /// owns stop-token suppression before client emission. + /// step always makes at least one token of progress. Before KV commit the + /// executor truncates this span after the first request-terminal token; + /// the scheduler retains ownership of typed stop-cause emission. pub accepted_tokens: Vec, } diff --git a/pegainfer-qwen3/tests/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index ea27a5abe..7baf4b126 100644 --- a/pegainfer-qwen3/tests/common/harness.rs +++ b/pegainfer-qwen3/tests/common/harness.rs @@ -25,6 +25,7 @@ use pegainfer_frontend::engine::RequestControl; use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestUpdate; use pegainfer_frontend::engine::SchedulerHandle; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::sampler::SamplingParams; @@ -39,6 +40,7 @@ pub(crate) fn request( Request { prompt_tokens, params, + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-sim/src/lib.rs b/pegainfer-sim/src/lib.rs index d7dae39cf..73ce591fb 100644 --- a/pegainfer-sim/src/lib.rs +++ b/pegainfer-sim/src/lib.rs @@ -304,6 +304,7 @@ fn duration_from_ms(ms: f64) -> Duration { #[cfg(test)] mod tests { use pegainfer_frontend::engine::Request; + use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; @@ -313,6 +314,7 @@ mod tests { Request { prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -418,6 +420,7 @@ mod tests { reason: FinishReason::Length, prompt_tokens: 2, completion_tokens: 3, + .. } )); } From 35aec766b322ab784d03fd655a57c0f6a60bebcc Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Wed, 26 Aug 2026 20:04:08 +0800 Subject: [PATCH 02/15] fix(qwen3): honor ignore-eos in test harness Signed-off-by: RicardoMin <17879681016@163.com> --- .../frontend/frontend-architecture.md | 33 +++++++++++-------- pegainfer-frontend/src/engine/stop.rs | 10 ------ pegainfer-qwen3/src/scheduler/resolve.rs | 5 +-- pegainfer-qwen3/src/scheduler/tests.rs | 4 +-- pegainfer-qwen3/src/speculative.rs | 4 +-- pegainfer-qwen3/tests/common/harness.rs | 12 ++++++- 6 files changed, 38 insertions(+), 30 deletions(-) diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index 239e2a61b..e314e004e 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -12,14 +12,19 @@ An engine is a set of schedulers, each a `Scheduler` implementation driven by th ``` pegainfer-frontend/src/engine/ -├── step.rs # the wire: RequestId, Request, QueuedRequest, -│ # StepOutputs { Vec }, -│ # RequestUpdate { scheduled, tokens, logprobs, cached_tokens, -│ # prompt_echo, kv_transfer, terminal }, Terminal -├── request_lifecycle.rs # submission envelope, abort control and step sender plumbing; -│ # DeferredFinish remains available for P/D handoff -├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, -│ # prompt/completion tallies, one merged update per touched id +├── step.rs # the wire: RequestId, Request, StepOutputs { Vec }, +│ # Request { ..., stop_policy }, RequestUpdate { scheduled, tokens, +│ # logprobs, cached_tokens, prompt_echo, kv_transfer, terminal }, +│ # Terminal { ..., stop_cause } +├── request_lifecycle.rs # typestate handles: QueuedRequest ─admit→ +│ # ActiveRequest ─finish/fail/defer→ consumed; every +│ # transition is by-move, a dropped handle emits Failed +│ # (drop bomb); DeferredFinish; RequestControl +├── emitter.rs # StepEmitter: the single writer of the per-step buffer; stamps +│ # timestamps, tallies prompt/completion counts, folds each +│ # request's step into one RequestUpdate; commit_step sends once +├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, +│ # prompt/completion tallies, one merged update per touched id ├── wiring.rs # scheduler_pair, SchedulerHandle (submit/take_steps/load), │ # Engine { schedulers, info, lora }, LiveScheduler, │ # EngineInfo, LaunchedEngine { Handle | Stepped } @@ -32,10 +37,12 @@ pegainfer-frontend/src/engine/ Design decisions worth knowing before touching it: - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. -- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. -- **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. -- **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. -- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. +- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. This is what makes `defer_finish` safe: a P/D prefill executor can withhold a request's `Finished` until its KV saves are peer-visible and send it later from any thread — the deferred message carries the request's entire buffered update, so late delivery cannot reorder. +- **Independent stop policy and cause.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in `RequestUpdate` and carries `Terminal::Finished.stop_cause`; decode paths include its real logprob when available, so the stepped bridge can report the actual explicit stop ID without reconstructing it. `ignore_eos` affects only model EOS. Legacy producers may leave the cause empty while they are migrated individually. +- **Typestate lifecycle.** `QueuedRequest` (queued) and `ActiveRequest` (streaming) are owned tokens; admit/reject/retire/finish/fail consume them, so "terminal exactly once, nothing after it" cannot be miscoded — it does not compile. A handle dropped without a transition emits `Failed` from its `Drop`, which is also how a crashed scheduler answers every in-flight request: the driver drops the scheduler, the handles fall, the terminals ship. +- **Emitter as single writer.** Schedulers never touch the channel; they call `StepEmitter` methods against their handles. The emitter stamps `ScheduledInfo` at admission, tallies token counts (terminal counts derive from the tally, never from model-side arithmetic), and `commit_step` publishes the whole step in one send. +- **Pure polling driver.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish load, commit. No idle/park distinction — the scheduler owns the GPU and spinning on it costs nothing anyone else could use; async KV I/O (prefetch, decode-overlap prefill) is naturally absorbed by polling. An idle iteration ends in a `spin_loop` hint (relaxes the core's issue slots, no latency cost — busy iterations never pause). The loop exits when the frontend drops the handle and the queue drains. +- **Gemma 4 async prefill exception.** While asynchronous prefill is the only remaining work, Gemma 4 drains and joins that lane rather than hot-polling its completion; decode or queued work keeps the normal polling path. - **Abort is a flag, not channel teardown.** `SchedulerHandle::submit` returns a `RequestControl`; the frontend flips its boolean abort flag and the scheduler retires the request silently on its next touch (no terminal — the frontend already dropped its state for that id). - **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::load()` — pull-only by design, "notify me on load change" is deliberately unrepresentable (the driver busy-polls, so a subscription edge would fire per spin). All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit. - **Control plane lives outside the contract.** `Scheduler` has no control method and the contract carries no control channel. A capability like LoRA is a private channel the model crate mints *before* `spawn_scheduler` — the scheduler closes over the receiver, the `LoraClient` sender surfaces as `Engine.lora: Option`, and the `Option` *is* the capability (no `bool` flag, no registry until a second capability exists). The vocabulary (`LoraControl`, `LoraClient`) is still defined in the frontend crate because the frontend must speak it without holding model structs; only the wiring is the model's business. @@ -91,7 +98,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Protocol stacks -**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` appends the stop sentinel token, which is how usage keeps counting the suppressed EOS). +**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a typed `StopCause` reports the actual request stop token, while the synthetic sentinel remains only as a compatibility fallback for producers that provide no cause). **`dynamo` (planned second stack).** dynamo's `lib/llm` in-process path removes the wire protocol entirely (`EngineConfig::InProcessTokens` + `run_input`). The step contract was shaped so this stack can consume `StepOutputs` directly without impersonation overhead. Decision gate: prototype, A/B against the vllm stack, let TTFT/step-overhead numbers pick the default. diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs index 064088ba6..099b1372b 100644 --- a/pegainfer-frontend/src/engine/stop.rs +++ b/pegainfer-frontend/src/engine/stop.rs @@ -91,14 +91,4 @@ mod tests { assert_eq!(policy.classify(99, |_| false), Some(StopCause::Eos(99))); } - - #[test] - fn unmatched_token_does_not_stop() { - let policy = StopPolicy { - eos: EosPolicy::Token(99), - token_ids: vec![42], - }; - - assert_eq!(policy.classify(7, |_| false), None); - } } diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 75da390ae..14993861d 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -63,8 +63,9 @@ pub(crate) fn resolve_step( /// Turn each request's accepted speculative span into a decode effect. A span /// commits 1..=K+1 tokens at once; we walk it in order so a stop token or the -/// max-output budget truncates exactly where it lands (the executor already -/// suppressed nothing — stop handling lives here, mirroring single-token decode). +/// max-output budget lands exactly where expected. The executor has already +/// truncated any suffix after a request-terminal token to keep speculative +/// state consistent; the resolver classifies its typed cause here. pub(crate) fn resolve_speculative_outputs( executor: &impl ModelExecutor, active: &[ActiveRequestState], diff --git a/pegainfer-qwen3/src/scheduler/tests.rs b/pegainfer-qwen3/src/scheduler/tests.rs index 6df3914c1..7d931d1c3 100644 --- a/pegainfer-qwen3/src/scheduler/tests.rs +++ b/pegainfer-qwen3/src/scheduler/tests.rs @@ -571,10 +571,10 @@ fn speculative_full_span_accept_continues() { } #[test] -fn speculative_stop_token_midspan_finishes_and_suppresses_eos() { +fn speculative_stop_token_midspan_finishes_and_retains_the_trigger() { let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); let active = [spec_active(1, 5, 100, false)]; - // EOS lands at span position 2; tokens before it are emitted, EOS is not. + // EOS lands at span position 2; the trigger is retained and the suffix is not. let results = [spec_result(1, vec![10, 11, SPEC_EOS, 13])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index a8050b42a..0614433b4 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -57,8 +57,8 @@ impl VerifyStepItem { #[derive(Clone, Copy)] pub(crate) struct VerifyPlan<'a> { pub requests: &'a [VerifyStepItem], - /// Request-local stop policies in the same order as `requests`. They stay - /// executor-side and are not copied into the worker command or GPU batch. + /// Request-local stop policies in the same order as `requests`. They remain + /// host-side and are not copied into GPU buffers. pub stop_policies: &'a [StopPolicy], /// Engine step seed for the verify rows' sampler pass (same contract as /// decode: fresh per step; seeded rows re-mix their own request seed). diff --git a/pegainfer-qwen3/tests/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index 7baf4b126..45d48e057 100644 --- a/pegainfer-qwen3/tests/common/harness.rs +++ b/pegainfer-qwen3/tests/common/harness.rs @@ -18,6 +18,7 @@ use std::sync::Mutex; use pegainfer_frontend::engine::Engine; use pegainfer_frontend::engine::EngineInfo; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::LoraClient; use pegainfer_frontend::engine::PromptEcho; use pegainfer_frontend::engine::Request; @@ -37,10 +38,19 @@ pub(crate) fn request( params: SamplingParams, max_tokens: usize, ) -> Request { + let stop_policy = StopPolicy { + eos: if params.ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + ..StopPolicy::default() + }; + Request { prompt_tokens, params, - stop_policy: StopPolicy::default(), + stop_policy, max_tokens, lora_adapter: None, kv_transfer_params: None, From a3fac5209ac57a98acde67b7201476f83e6b2340 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Fri, 4 Sep 2026 13:12:56 +0800 Subject: [PATCH 03/15] test(frontend): import StopPolicy in stepped bridge tests Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-frontend/src/vllm/bridge/stepped.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index fee78413e..9c4927ea5 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -659,6 +659,7 @@ mod tests { use super::*; use crate::engine::PromptEcho; use crate::engine::RejectReason; + use crate::engine::StopPolicy; use crate::engine::TokenLogprob; use crate::engine::scheduler_pair; From d3cf62d3f2967d3376c6bfbb34d06897a83d8cfb Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Fri, 4 Sep 2026 13:36:50 +0800 Subject: [PATCH 04/15] fix(gemma4): initialize stop policy in test harness Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-gemma4/src/engine/lane_tests.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/pegainfer-gemma4/src/engine/lane_tests.rs b/pegainfer-gemma4/src/engine/lane_tests.rs index f0b46250e..34bcd030e 100644 --- a/pegainfer-gemma4/src/engine/lane_tests.rs +++ b/pegainfer-gemma4/src/engine/lane_tests.rs @@ -55,6 +55,7 @@ impl Harness { ignore_eos: true, ..pegainfer_frontend::sampler::SamplingParams::default() }, + stop_policy: pegainfer_frontend::engine::StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, From 1ac42df872517d6b237ca9ba72e2ac4a003279e5 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sun, 6 Sep 2026 22:46:56 +0800 Subject: [PATCH 05/15] fix(qwen3): finalize stop handling before speculative commit Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-frontend/src/engine/stop.rs | 58 +++++++++---- pegainfer-frontend/src/vllm/wire.rs | 36 +++++--- pegainfer-qwen3/src/executor.rs | 23 ++++- pegainfer-qwen3/src/executor/spec.rs | 22 +---- pegainfer-qwen3/src/frontend_adapter/tests.rs | 59 ++----------- pegainfer-qwen3/src/scheduler/resolve.rs | 23 ++--- pegainfer-qwen3/src/scheduler/test_support.rs | 86 ------------------ pegainfer-qwen3/src/scheduler/tests.rs | 10 +-- .../tests/dflash_speculative_gate.rs | 87 +++++++++++++++++++ 9 files changed, 198 insertions(+), 206 deletions(-) diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs index 099b1372b..6aff2acf2 100644 --- a/pegainfer-frontend/src/engine/stop.rs +++ b/pegainfer-frontend/src/engine/stop.rs @@ -1,3 +1,5 @@ +use std::sync::Arc; + /// How a request treats end-of-sequence tokens. #[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] pub enum EosPolicy { @@ -6,8 +8,6 @@ pub enum EosPolicy { /// Use the model executor's configured EOS set. #[default] ModelDefault, - /// Stop only on this protocol-provided primary EOS token. - Token(u32), } /// Request-scoped token stopping policy. @@ -18,10 +18,31 @@ pub enum EosPolicy { #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct StopPolicy { pub eos: EosPolicy, - pub token_ids: Vec, + /// Sorted and deduplicated explicit stop IDs. + /// + /// Requests clone this policy while building a plan and sending it to + /// worker ranks. Keeping the normalized set behind an `Arc` makes those + /// clones cheap and lets classification use binary search for large stop + /// sets without regressing the common one-ID case. + pub token_ids: Arc<[u32]>, } impl StopPolicy { + /// Build a policy from wire-provided stop IDs. + /// + /// Normalization happens once at the request boundary. Internal copies can + /// then share the immutable slice instead of repeatedly sorting, deduping, + /// or cloning the caller's vector. + #[must_use] + pub fn new(eos: EosPolicy, mut token_ids: Vec) -> Self { + token_ids.sort_unstable(); + token_ids.dedup(); + Self { + eos, + token_ids: token_ids.into(), + } + } + /// Classify a token using vLLM's priority: EOS first, then the request's /// explicit stop-token set. #[must_use] @@ -33,12 +54,11 @@ impl StopPolicy { let is_eos = match self.eos { EosPolicy::Ignore => false, EosPolicy::ModelDefault => is_model_eos(token_id), - EosPolicy::Token(eos_token_id) => token_id == eos_token_id, }; if is_eos { Some(StopCause::Eos(token_id)) - } else if self.token_ids.contains(&token_id) { + } else if self.token_ids.binary_search(&token_id).is_ok() { Some(StopCause::Token(token_id)) } else { None @@ -71,10 +91,7 @@ mod tests { #[test] fn ignored_eos_does_not_disable_an_explicit_stop() { - let policy = StopPolicy { - eos: EosPolicy::Ignore, - token_ids: vec![99], - }; + let policy = StopPolicy::new(EosPolicy::Ignore, vec![99]); assert_eq!( policy.classify(99, |token_id| token_id == 99), @@ -83,12 +100,23 @@ mod tests { } #[test] - fn eos_wins_when_the_same_id_is_also_an_explicit_stop() { - let policy = StopPolicy { - eos: EosPolicy::Token(99), - token_ids: vec![99], - }; + fn normalizes_stop_sets_across_common_sizes() { + let empty = StopPolicy::default(); + assert!(empty.classify(7, |_| false).is_none()); + + let single = StopPolicy::new(EosPolicy::Ignore, vec![42]); + assert_eq!(single.classify(42, |_| false), Some(StopCause::Token(42))); + assert!(single.classify(43, |_| false).is_none()); + + let policy = StopPolicy::new(EosPolicy::Ignore, vec![7, 3, 7, 1]); + assert_eq!(policy.token_ids.as_ref(), [1, 3, 7]); + assert!(policy.classify(7, |_| false).is_some()); + assert!(policy.classify(8, |_| false).is_none()); - assert_eq!(policy.classify(99, |_| false), Some(StopCause::Eos(99))); + let full = StopPolicy::new(EosPolicy::Ignore, (0..151_936).collect()); + assert_eq!(full.token_ids.len(), 151_936); + assert!(full.classify(0, |_| false).is_some()); + assert!(full.classify(151_935, |_| false).is_some()); + assert!(full.classify(151_936, |_| false).is_none()); } } diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index 7d1585f1c..22706015c 100644 --- a/pegainfer-frontend/src/vllm/wire.rs +++ b/pegainfer-frontend/src/vllm/wire.rs @@ -79,10 +79,10 @@ pub(crate) fn to_wire_prompt_logprobs(prompt: PromptEcho) -> Result SamplingParams { // The vLLM frontend lowers a client `ignore_eos=true` to `_eos_token_id: // None`, but `_all_stop_token_ids` always carries the model EOS set (it - // exists for min_tokens masking, not stop detection). Deriving ignore_eos - // from all_stop_token_ids would therefore void every ignore_eos request on - // models with a real EOS. Only _eos_token_id and the client's explicit - // stop_token_ids express the legacy scheduler's stop intent. + // exists for min_tokens masking, not stop detection). This conversion feeds + // the legacy SamplingParams contract, which has no field for explicit + // request stop IDs; keep its historical lowering until each producer is + // migrated to StopPolicy. Qwen3 receives the independent policy below. let ignore_eos = params.eos_token_id.is_none() && params.stop_token_ids.is_empty(); if params.temperature <= 0.0 { return SamplingParams { @@ -114,15 +114,15 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar } pub(crate) fn convert_stop_policy(params: &EngineCoreSamplingParams) -> StopPolicy { - StopPolicy { + StopPolicy::new( // Qwen3 owns the complete model EOS set in generation_config. The // protocol's optional primary ID only tells us whether EOS is active; // using it as a singleton would miss secondary model EOS IDs. - eos: params + params .eos_token_id .map_or(EosPolicy::Ignore, |_| EosPolicy::ModelDefault), - token_ids: params.stop_token_ids.clone(), - } + params.stop_token_ids.clone(), + ) } /// Reject request parameters the frontend cannot represent faithfully. @@ -143,6 +143,12 @@ pub(crate) fn unsupported_request_params(params: &EngineCoreSamplingParams) -> O { return Some("negative logprob counts are not supported".into()); } + if params.min_tokens != 0 { + return Some(format!( + "min_tokens={} is not supported by current engine contracts", + params.min_tokens + )); + } if !(0.0..1.0).contains(¶ms.min_p) || !params.min_p.is_finite() { return Some(format!("min_p {} outside [0, 1)", params.min_p)); } @@ -251,8 +257,9 @@ mod tests { params.eos_token_id = Some(163_586); assert!(!convert_sampling(¶ms).ignore_eos); - // The legacy scheduler keeps EOS active when an explicit stop token is - // present; the stepped bridge carries explicit stops in StopPolicy. + // The legacy SamplingParams contract keeps EOS active when explicit + // stop IDs are present; the stepped Qwen3 path carries those IDs in + // StopPolicy instead. params.eos_token_id = None; params.stop_token_ids = vec![42]; assert!(!convert_sampling(¶ms).ignore_eos); @@ -265,7 +272,7 @@ mod tests { params.stop_token_ids = vec![11]; assert_eq!(convert_stop_policy(¶ms).eos, EosPolicy::ModelDefault); - assert_eq!(convert_stop_policy(¶ms).token_ids, vec![11]); + assert_eq!(convert_stop_policy(¶ms).token_ids.as_ref(), [11]); } #[test] @@ -292,6 +299,13 @@ mod tests { params.repetition_penalty = 1.0; assert_eq!(unsupported_request_params(¶ms), None); + params.min_tokens = 1; + assert_eq!( + unsupported_request_params(¶ms).as_deref(), + Some("min_tokens=1 is not supported by current engine contracts") + ); + params.min_tokens = 0; + params.min_p = 0.2; assert_eq!(unsupported_request_params(¶ms), None); params.min_p = 1.5; diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index fee03f956..7c3d563a9 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -3428,6 +3428,7 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], + stop_policies: &[StopPolicy], capture_layer_ids: &[usize], sample_seed: u64, bufs: &mut VerifyGraphBuffers, @@ -3535,9 +3536,22 @@ impl LocalQwen3Lane { .flat_map(|req| std::iter::repeat_n(&req.params, req.as_slice().len())) .collect(); let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, sample_seed)?; - let all_results = build_verify_results(&expanded, &target_tokens)?; - let (results_a, results_b) = all_results.split_at(requests.len()); - + let mut all_results = build_verify_results(&expanded, &target_tokens)?; + // A terminal token ends the request even when it appears in the middle + // of a speculative span. Normalize every candidate before selecting a + // winner; otherwise a discarded suffix can win the hedge and advance + // the wrong KV/hidden state and acceptance statistics. + let (results_a, results_b) = all_results.split_at_mut(requests.len()); + for (result, policy) in results_a.iter_mut().zip(stop_policies) { + spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + } + for (slot, (idx, _)) in hedge_spans.iter().enumerate() { + spec::truncate_after_terminal( + &mut results_b[slot], + &stop_policies[*idx], + &self.model.config().stop_token_ids, + ); + } // Per request keep the best-accepting chain; ties keep chain A (no // copies). A later chain of the same request only replaces the // running winner when strictly better, so the final page/hidden @@ -3566,7 +3580,7 @@ impl LocalQwen3Lane { cudarc::driver::result::memcpy_dtod_async( dst, src, - span_len * hidden_dim * elem, + res_b.accepted_tokens.len() * hidden_dim * elem, ctx.stream.cu_stream(), ) } @@ -3666,6 +3680,7 @@ impl LocalQwen3Lane { if let Some(result) = self.try_execute_hedged_verify( requests, kv_views, + stop_policies, &capture_layer_ids, sample_seed, &mut bufs, diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index a55ef96cb..27b411ec8 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -37,6 +37,9 @@ pub(super) fn truncate_after_terminal( }; let keep = keep + 1; result.accepted_tokens.truncate(keep); + // A trigger in the accepted draft prefix is itself a matched draft. If the + // trigger is the posterior token, the original count is already `keep - 1`; + // clamping to the retained prefix handles both cases. result.matched_draft_tokens = result.matched_draft_tokens.min(keep); } @@ -237,10 +240,7 @@ mod tests { #[test] fn explicit_stop_truncates_the_kv_commit_after_the_trigger() { - let policy = StopPolicy { - eos: EosPolicy::Ignore, - token_ids: vec![7], - }; + let policy = StopPolicy::new(EosPolicy::Ignore, vec![7]); let mut result = VerifyRequestResult { request_id: RequestId::new(1), matched_draft_tokens: 3, @@ -252,18 +252,4 @@ mod tests { assert_eq!(result.accepted_tokens, vec![5, 7]); assert_eq!(result.matched_draft_tokens, 2); } - - #[test] - fn model_eos_truncates_but_keeps_the_trigger() { - let mut result = VerifyRequestResult { - request_id: RequestId::new(2), - matched_draft_tokens: 3, - accepted_tokens: vec![5, 99, 8, 9], - }; - - truncate_after_terminal(&mut result, &StopPolicy::default(), &[99]); - - assert_eq!(result.accepted_tokens, vec![5, 99]); - assert_eq!(result.matched_draft_tokens, 2); - } } diff --git a/pegainfer-qwen3/src/frontend_adapter/tests.rs b/pegainfer-qwen3/src/frontend_adapter/tests.rs index 0a89aac2b..1d1b2244d 100644 --- a/pegainfer-qwen3/src/frontend_adapter/tests.rs +++ b/pegainfer-qwen3/src/frontend_adapter/tests.rs @@ -151,8 +151,10 @@ fn request_stop_token_beats_length_during_prefill_when_eos_is_ignored() { let (partition, _lora, mut steps) = launch(executor, false); let mut req = request(16, 1); - req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; - req.stop_policy.token_ids = vec![100]; + req.stop_policy = pegainfer_frontend::engine::StopPolicy::new( + pegainfer_frontend::engine::EosPolicy::Ignore, + vec![100], + ); let control = partition.handle.submit(req); let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); @@ -177,8 +179,10 @@ fn request_stop_token_beats_length_during_decode_when_eos_is_ignored() { let (partition, _lora, mut steps) = launch(executor, false); let mut req = request(16, 2); - req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; - req.stop_policy.token_ids = vec![200]; + req.stop_policy = pegainfer_frontend::engine::StopPolicy::new( + pegainfer_frontend::engine::EosPolicy::Ignore, + vec![200], + ); let control = partition.handle.submit(req); let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); @@ -197,53 +201,6 @@ fn request_stop_token_beats_length_during_decode_when_eos_is_ignored() { )); } -#[test] -fn speculative_request_stop_beats_length_midspan_and_cleans_up() { - let dropped = Arc::new(Mutex::new(Vec::new())); - let executor = FakeExecutor::new(4, Arc::clone(&dropped)) - .with_stop_token(12) - .with_speculative_accepted_tokens(&[10, 11, 12, 13]); - - let (partition, _lora, mut steps) = launch(executor, false); - - let mut req = request(16, 4); - req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; - req.stop_policy.token_ids = vec![12]; - - let control = partition.handle.submit(req); - let (tokens, terminal) = steps.collect_terminal(control.id()); - - assert_eq!( - tokens, - vec![100, 10, 11, 12], - "the trigger is retained and the accepted suffix is discarded" - ); - assert!(matches!( - terminal, - Terminal::Finished { - reason: FinishReason::Stop, - stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(12)), - completion_tokens: 4, - .. - } - )); - - assert!( - wait_until(Duration::from_secs(1), || { - dropped.lock().unwrap().contains(&0) - }), - "stopped speculative request state should be dropped" - ); - - assert!( - wait_until(Duration::from_secs(1), || { - let metrics = partition.handle.metrics(); - metrics.num_running_reqs == 0 && metrics.kv_used_blocks == 0 - }), - "stopped speculative request should release scheduler state and KV blocks" - ); -} - #[test] fn unknown_lora_request_is_rejected_without_blocking_base_request() { let dropped = Arc::new(Mutex::new(Vec::new())); diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 14993861d..e9ff31b41 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -1,5 +1,6 @@ use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use super::ActiveRequestState; use super::PendingRequest; @@ -14,22 +15,12 @@ use crate::executor::ModelExecutor; use crate::executor::PrefillRequestResult; use crate::speculative::VerifyRequestResult; -fn stop_cause( +fn classify_stop( executor: &impl ModelExecutor, - req: &ActiveRequestState, + policy: &StopPolicy, token: u32, ) -> Option { - req.stop_policy - .classify(token, |token_id| executor.is_stop_token(token_id)) -} - -fn pending_stop_cause( - executor: &impl ModelExecutor, - req: &PendingRequest, - token: u32, -) -> Option { - req.stop_policy - .classify(token, |token_id| executor.is_stop_token(token_id)) + policy.classify(token, |token_id| executor.is_stop_token(token_id)) } pub(crate) fn resolve_step( @@ -85,7 +76,7 @@ pub(crate) fn resolve_speculative_outputs( completion_tokens += 1; emitted.push(token); - if let Some(stop_cause) = stop_cause(executor, req, token) { + if let Some(stop_cause) = classify_stop(executor, &req.stop_policy, token) { return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, @@ -148,7 +139,7 @@ fn resolve_prefill_outputs( }); } - if let Some(stop_cause) = pending_stop_cause(executor, &req, result.first_token) { + if let Some(stop_cause) = classify_stop(executor, &req.stop_policy, result.first_token) { effects.pending.push(PendingEffect::EmitAndFinish { request_id: req.request_id, token: result.first_token, @@ -204,7 +195,7 @@ fn resolve_decode_outputs( .find(|req| req.request_id == result.request_id) .expect("decode request_id must exist in active set"); let completion_tokens = req.generated_count + 1; - let stop_cause = stop_cause(executor, req, result.token); + let stop_cause = classify_stop(executor, &req.stop_policy, result.token); let at_limit = completion_tokens >= req.max_tokens; if let Some(stop_cause) = stop_cause { diff --git a/pegainfer-qwen3/src/scheduler/test_support.rs b/pegainfer-qwen3/src/scheduler/test_support.rs index 10dffd88b..cfd639df7 100644 --- a/pegainfer-qwen3/src/scheduler/test_support.rs +++ b/pegainfer-qwen3/src/scheduler/test_support.rs @@ -24,12 +24,6 @@ use crate::executor::PrefillStepItem; use crate::executor::RequestId; use crate::executor::UnifiedPlan; use crate::executor::UnifiedResult; -use crate::speculative::DraftPlan; -use crate::speculative::DraftRequestResult; -use crate::speculative::DraftResult; -use crate::speculative::VerifyPlan; -use crate::speculative::VerifyRequestResult; -use crate::speculative::VerifyResult; pub(crate) struct FakeExecutor { pub(crate) block_size: usize, @@ -47,7 +41,6 @@ pub(crate) struct FakeExecutor { pub(crate) prefetch_offers: Arc>>, stop_token: Option, emit_logprobs: bool, - speculative_accepted_tokens: Option>, } impl FakeExecutor { @@ -66,7 +59,6 @@ impl FakeExecutor { prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, emit_logprobs: false, - speculative_accepted_tokens: None, } } @@ -80,15 +72,6 @@ impl FakeExecutor { self } - pub(crate) fn with_speculative_accepted_tokens(mut self, tokens: &[u32]) -> Self { - assert!( - !tokens.is_empty(), - "fake speculative span must make progress" - ); - self.speculative_accepted_tokens = Some(tokens.to_vec()); - self - } - pub(crate) fn with_decode_failure(mut self) -> Self { self.fail_decode_once = true; self @@ -297,75 +280,6 @@ impl ModelExecutor for FakeExecutor { .collect(), }) } - - fn execute_speculative_draft(&mut self, plan: DraftPlan<'_>) -> Result { - let accepted_tokens = self - .speculative_accepted_tokens - .as_ref() - .ok_or_else(|| anyhow::anyhow!("fake speculative decoding is disabled"))?; - - Ok(DraftResult { - requests: plan - .requests - .iter() - .map(|req| { - let mut token_ids = Vec::with_capacity(accepted_tokens.len()); - token_ids.push(req.current_token); - token_ids.extend( - accepted_tokens - .iter() - .copied() - .take(accepted_tokens.len().saturating_sub(1)), - ); - - DraftRequestResult { - request_id: req.request_id, - token_ids, - } - }) - .collect(), - }) - } - - fn execute_speculative_verify(&mut self, plan: VerifyPlan<'_>) -> Result { - let configured = self - .speculative_accepted_tokens - .clone() - .ok_or_else(|| anyhow::anyhow!("fake speculative decoding is disabled"))?; - - let mut requests = Vec::with_capacity(plan.requests.len()); - - for req in plan.requests { - let span_len = req.as_slice().len(); - anyhow::ensure!(span_len > 0, "fake speculative verify span is empty"); - - let accepted_tokens = configured[..configured.len().min(span_len)].to_vec(); - - let current_tokens = self - .held_tokens - .get(&req.request_id) - .copied() - .ok_or_else(|| anyhow::anyhow!("missing fake request state"))?; - - self.ensure_request_tokens(req.request_id, current_tokens + accepted_tokens.len())?; - - requests.push(VerifyRequestResult { - request_id: req.request_id, - matched_draft_tokens: accepted_tokens.len().saturating_sub(1), - accepted_tokens, - }); - } - - Ok(VerifyResult { requests }) - } - - fn speculative_enabled(&self) -> bool { - self.speculative_accepted_tokens.is_some() - } - - fn speculative_request_ready(&self, request_id: RequestId) -> bool { - self.speculative_accepted_tokens.is_some() && self.held_tokens.contains_key(&request_id) - } } /// A minimal contract request: `prompt_len` filler tokens, default sampling. diff --git a/pegainfer-qwen3/src/scheduler/tests.rs b/pegainfer-qwen3/src/scheduler/tests.rs index 7d931d1c3..4ac09de6a 100644 --- a/pegainfer-qwen3/src/scheduler/tests.rs +++ b/pegainfer-qwen3/src/scheduler/tests.rs @@ -522,14 +522,14 @@ fn spec_active( ignore_eos, ..SamplingParams::default() }, - stop_policy: StopPolicy { - eos: if ignore_eos { + stop_policy: StopPolicy::new( + if ignore_eos { EosPolicy::Ignore } else { EosPolicy::ModelDefault }, - token_ids: Vec::new(), - }, + Vec::new(), + ), ..active_state(id, generated_count, max_tokens) } } @@ -666,7 +666,7 @@ fn speculative_request_stop_truncates_the_span_when_eos_is_ignored() { let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); let mut request = spec_active(1, 0, 100, true); - request.stop_policy.token_ids = vec![12]; + request.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![12]); let active = [request]; let results = [spec_result(1, vec![10, 11, 12, 13])]; diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 16e07781b..a1d21a184 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -44,6 +44,10 @@ use std::path::PathBuf; use std::process::Command; use std::time::Duration; +use pegainfer_frontend::engine::EosPolicy; +use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; use pegainfer_qwen3::DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES; @@ -620,6 +624,88 @@ fn dflash_concurrent_heterogeneous_is_lossless() { ); } +/// Production hedge regression: an explicit stop in the middle of a verify +/// span must be applied before the hedge winner is selected and committed. +/// The trigger is retained, while every accepted suffix token is discarded. +#[test] +fn dflash_hedged_midspan_stop_retains_trigger() { + common::harness::init_capture_logging(); + let (Some(model_path), Some(draft_path)) = (target_path_or_skip(), draft_path_or_skip()) else { + return; + }; + if std::env::var_os("PEGAINFER_SPEC_HEDGE").is_none() { + eprintln!( + "skipping hedged mid-span stop gate: run it through hedged_ladder_passes_the_lossless_gates" + ); + return; + } + let _gpu = GPU + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + let prompt = "Write a short paragraph about a blue bicycle:"; + let tokenizer = common::load_tokenizer(&model_path); + let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); + let draft_config: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(Path::new(&draft_path).join("config.json")) + .expect("read draft config"), + ) + .expect("parse draft config"); + let block_size = draft_config["block_size"] + .as_u64() + .expect("draft block_size") as usize; + let min_index = (block_size / 2).max(1); + + let engine = EngineHarness::new( + pegainfer_qwen3::launch( + Path::new(&model_path), + launch_options(Some(PathBuf::from(&draft_path))), + ) + .expect("failed to start speculative engine"), + ); + let mut baseline_params = SamplingParams::default(); + baseline_params.ignore_eos = true; + let baseline = engine + .submit(request( + prompt_tokens.clone(), + baseline_params, + GENERATED_TOKENS, + )) + .expect_finished() + .tokens; + let Some((stop_index, &stop_id)) = baseline + .iter() + .enumerate() + .skip(min_index) + .find(|(index, token)| !baseline[..*index].contains(token)) + else { + panic!("baseline did not produce a unique token after the middle of the verify span"); + }; + assert!( + stop_index < block_size, + "selected stop token at position {stop_index}, expected inside the first {block_size}-token verify span" + ); + + let mut stopped_params = SamplingParams::default(); + stopped_params.ignore_eos = true; + let mut stopped = request(prompt_tokens, stopped_params, GENERATED_TOKENS); + stopped.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![stop_id]); + let outcome = engine.submit(stopped).expect_finished(); + + assert_eq!(outcome.tokens.len(), stop_index + 1); + assert_eq!(outcome.tokens.last(), Some(&stop_id)); + assert!(!outcome.tokens[..stop_index].contains(&stop_id)); + assert!(matches!( + outcome.terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(id)), + completion_tokens, + .. + } if id == stop_id && completion_tokens == stop_index + 1 + )); +} + /// P2 regression: a request that fits the target context window but lands in the /// draft's `block_size` in-fill headroom (`max_pos - block_size < prompt + /// max_tokens <= max_pos`) must be rejected cleanly at admission. Before the @@ -760,6 +846,7 @@ fn hedged_ladder_passes_the_lossless_gates() { for child_test in [ "dflash_speculative_greedy_matches_plain_greedy", "dflash_concurrent_heterogeneous_is_lossless", + "dflash_hedged_midspan_stop_retains_trigger", ] { let output = Command::new(&exe) .args(["--exact", child_test, "--test-threads=1", "--nocapture"]) From 26848a3fc5d0c737d255b3293bba8cea444358cb Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sun, 6 Sep 2026 23:51:52 +0800 Subject: [PATCH 06/15] fix(qwen3): satisfy CUDA Clippy gate Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-qwen3/tests/dflash_speculative_gate.rs | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index a1d21a184..05d81d66a 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -663,8 +663,10 @@ fn dflash_hedged_midspan_stop_retains_trigger() { ) .expect("failed to start speculative engine"), ); - let mut baseline_params = SamplingParams::default(); - baseline_params.ignore_eos = true; + let baseline_params = SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }; let baseline = engine .submit(request( prompt_tokens.clone(), @@ -686,8 +688,10 @@ fn dflash_hedged_midspan_stop_retains_trigger() { "selected stop token at position {stop_index}, expected inside the first {block_size}-token verify span" ); - let mut stopped_params = SamplingParams::default(); - stopped_params.ignore_eos = true; + let stopped_params = SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }; let mut stopped = request(prompt_tokens, stopped_params, GENERATED_TOKENS); stopped.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![stop_id]); let outcome = engine.submit(stopped).expect_finished(); From 77f8a4b3097c46c356fa97462e1c9d53efbf04d1 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Mon, 7 Sep 2026 21:01:14 +0800 Subject: [PATCH 07/15] fix(qwen3): enforce stop ordering in hedged verify Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-frontend/src/engine/stop.rs | 33 ++- pegainfer-frontend/src/vllm/wire.rs | 11 +- pegainfer-qwen3/src/executor.rs | 125 ++++++++++++ pegainfer-qwen3/src/executor/dflash_lane.rs | 9 + pegainfer-qwen3/src/executor/spec.rs | 14 +- pegainfer-qwen3/src/speculative.rs | 4 + pegainfer-qwen3/tests/common/harness.rs | 8 +- .../tests/dflash_speculative_gate.rs | 191 +++++++++++++----- 8 files changed, 320 insertions(+), 75 deletions(-) diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs index 6aff2acf2..83c62b3fc 100644 --- a/pegainfer-frontend/src/engine/stop.rs +++ b/pegainfer-frontend/src/engine/stop.rs @@ -17,14 +17,14 @@ pub enum EosPolicy { /// reports the actual matching token ID. #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct StopPolicy { - pub eos: EosPolicy, + eos: EosPolicy, /// Sorted and deduplicated explicit stop IDs. /// /// Requests clone this policy while building a plan and sending it to /// worker ranks. Keeping the normalized set behind an `Arc` makes those /// clones cheap and lets classification use binary search for large stop /// sets without regressing the common one-ID case. - pub token_ids: Arc<[u32]>, + token_ids: Arc<[u32]>, } impl StopPolicy { @@ -100,23 +100,22 @@ mod tests { } #[test] - fn normalizes_stop_sets_across_common_sizes() { - let empty = StopPolicy::default(); - assert!(empty.classify(7, |_| false).is_none()); - - let single = StopPolicy::new(EosPolicy::Ignore, vec![42]); - assert_eq!(single.classify(42, |_| false), Some(StopCause::Token(42))); - assert!(single.classify(43, |_| false).is_none()); - + fn normalizes_unsorted_duplicate_stop_ids() { let policy = StopPolicy::new(EosPolicy::Ignore, vec![7, 3, 7, 1]); - assert_eq!(policy.token_ids.as_ref(), [1, 3, 7]); - assert!(policy.classify(7, |_| false).is_some()); + + assert_eq!(policy.classify(1, |_| false), Some(StopCause::Token(1))); + assert_eq!(policy.classify(3, |_| false), Some(StopCause::Token(3))); + assert_eq!(policy.classify(7, |_| false), Some(StopCause::Token(7))); assert!(policy.classify(8, |_| false).is_none()); + } + + #[test] + fn model_eos_has_priority_over_explicit_stop() { + let policy = StopPolicy::new(EosPolicy::ModelDefault, vec![99]); - let full = StopPolicy::new(EosPolicy::Ignore, (0..151_936).collect()); - assert_eq!(full.token_ids.len(), 151_936); - assert!(full.classify(0, |_| false).is_some()); - assert!(full.classify(151_935, |_| false).is_some()); - assert!(full.classify(151_936, |_| false).is_none()); + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Eos(99)) + ); } } diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index 22706015c..299a16ab9 100644 --- a/pegainfer-frontend/src/vllm/wire.rs +++ b/pegainfer-frontend/src/vllm/wire.rs @@ -271,8 +271,15 @@ mod tests { params.eos_token_id = Some(99); params.stop_token_ids = vec![11]; - assert_eq!(convert_stop_policy(¶ms).eos, EosPolicy::ModelDefault); - assert_eq!(convert_stop_policy(¶ms).token_ids.as_ref(), [11]); + let policy = convert_stop_policy(¶ms); + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(crate::engine::StopCause::Eos(99)) + ); + assert_eq!( + policy.classify(11, |_| false), + Some(crate::engine::StopCause::Token(11)) + ); } #[test] diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 7c3d563a9..7036e7a5b 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -3542,6 +3542,29 @@ impl LocalQwen3Lane { // winner; otherwise a discarded suffix can win the hedge and advance // the wrong KV/hidden state and acceptance statistics. let (results_a, results_b) = all_results.split_at_mut(requests.len()); + anyhow::ensure!( + results_b.len() == hedge_spans.len(), + "hedge returned {} B results for {} hedge spans", + results_b.len(), + hedge_spans.len() + ); + let trace = std::env::var_os("PEGAINFER_TEST_LOG").is_some(); + let raw_a_lengths: Vec = trace + .then(|| { + results_a + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + }) + .unwrap_or_default(); + let raw_b_lengths: Vec = trace + .then(|| { + results_b + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + }) + .unwrap_or_default(); for (result, policy) in results_a.iter_mut().zip(stop_policies) { spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); } @@ -3552,6 +3575,54 @@ impl LocalQwen3Lane { &self.model.config().stop_token_ids, ); } + let ( + retained_a_lengths, + retained_b_lengths, + raw_best_lengths, + raw_best_is_b, + retained_best_lengths, + retained_best_is_b, + ) = if trace { + let retained_a_lengths: Vec = results_a + .iter() + .map(|result| result.accepted_tokens.len()) + .collect(); + let retained_b_lengths: Vec = results_b + .iter() + .map(|result| result.accepted_tokens.len()) + .collect(); + let mut raw_best_lengths = raw_a_lengths.clone(); + let mut raw_best_is_b = vec![false; requests.len()]; + let mut retained_best_lengths = retained_a_lengths.clone(); + let mut retained_best_is_b = vec![false; requests.len()]; + for (slot, (idx, _)) in hedge_spans.iter().enumerate() { + if raw_b_lengths[slot] > raw_best_lengths[*idx] { + raw_best_lengths[*idx] = raw_b_lengths[slot]; + raw_best_is_b[*idx] = true; + } + if retained_b_lengths[slot] > retained_best_lengths[*idx] { + retained_best_lengths[*idx] = retained_b_lengths[slot]; + retained_best_is_b[*idx] = true; + } + } + ( + retained_a_lengths, + retained_b_lengths, + raw_best_lengths, + raw_best_is_b, + retained_best_lengths, + retained_best_is_b, + ) + } else { + ( + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + Vec::new(), + ) + }; // Per request keep the best-accepting chain; ties keep chain A (no // copies). A later chain of the same request only replaces the // running winner when strictly better, so the final page/hidden @@ -3561,6 +3632,7 @@ impl LocalQwen3Lane { let elem = std::mem::size_of::(); let mut final_requests: Vec = requests.to_vec(); let mut final_results: Vec = results_a.to_vec(); + let mut selected_is_b = trace.then(|| vec![false; requests.len()]); let mut b_wins = 0usize; let mut b_row_offset = a_total_rows; for (slot, (idx, replaced)) in hedge_spans.iter().enumerate() { @@ -3586,6 +3658,9 @@ impl LocalQwen3Lane { } .map_err(|e| anyhow::anyhow!("hedge hidden compaction failed: {e}"))?; b_wins += 1; + if let Some(selected) = selected_is_b.as_mut() { + selected[*idx] = true; + } final_requests[*idx] = expanded[requests.len() + slot].clone(); final_results[*idx] = res_b.clone(); } @@ -3608,9 +3683,57 @@ impl LocalQwen3Lane { &final_requests, &final_results, Some(bufs.captured_hidden()), + true, )?; + if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + for idx in 0..requests.len() { + let has_hedge = hedge_spans + .iter() + .any(|(request_idx, _)| *request_idx == idx); + if !has_hedge { + continue; + } + let raw_winner = if raw_best_is_b[idx] { 'B' } else { 'A' }; + let retained_winner = if retained_best_is_b[idx] { 'B' } else { 'A' }; + let selected = if selected_is_b.as_ref().is_some_and(|selected| selected[idx]) { + 'B' + } else { + 'A' + }; + let raw_b_lens = hedge_spans + .iter() + .enumerate() + .filter(|(_, (request_idx, _))| *request_idx == idx) + .map(|(slot, _)| raw_b_lengths[slot].to_string()) + .collect::>() + .join(","); + let retained_b_lens = hedge_spans + .iter() + .enumerate() + .filter(|(_, (request_idx, _))| *request_idx == idx) + .map(|(slot, _)| retained_b_lengths[slot].to_string()) + .collect::>() + .join(","); + log::debug!( + "Qwen3 DFlash hedge detail request={} raw_a={} raw_b_lens={} raw_best_len={} kept_a={} kept_b_lens={} retained_best_len={} raw_winner={} retained_winner={} selected={} selected_len={} matched_draft={}", + requests[idx].request_id, + raw_a_lengths[idx], + raw_b_lens, + raw_best_lengths[idx], + retained_a_lengths[idx], + retained_b_lens, + retained_best_lengths[idx], + raw_winner, + retained_winner, + selected, + final_results[idx].accepted_tokens.len(), + final_results[idx].matched_draft_tokens, + ); + } + } Ok(Some(VerifyResult { requests: final_results, + hedged: true, })) } @@ -3727,9 +3850,11 @@ impl LocalQwen3Lane { requests, &request_results, Some(bufs.captured_hidden()), + false, )?; Ok(VerifyResult { requests: request_results, + hedged: false, }) })(); self.verify_bufs = Some(bufs); diff --git a/pegainfer-qwen3/src/executor/dflash_lane.rs b/pegainfer-qwen3/src/executor/dflash_lane.rs index 4165643d3..6a47c4d47 100644 --- a/pegainfer-qwen3/src/executor/dflash_lane.rs +++ b/pegainfer-qwen3/src/executor/dflash_lane.rs @@ -201,6 +201,7 @@ impl LocalQwen3Lane { requests: &[VerifyStepItem], results: &[VerifyRequestResult], captured_hidden: Option<&HiddenStates>, + hedged: bool, ) -> Result<()> { let Some(captured_hidden) = captured_hidden else { anyhow::bail!("DFlash verify context capture requested but no hidden states returned"); @@ -242,6 +243,14 @@ impl LocalQwen3Lane { token_offset, result.accepted_tokens.len(), )?; + if hedged && std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + log::debug!( + "Qwen3 DFlash hedge context request={} appended={} matched_draft={}", + req.request_id, + result.accepted_tokens.len(), + result.matched_draft_tokens, + ); + } dflash.requests.insert(req.request_id, state); dflash.verified_draft_tokens += req.token_ids.len().saturating_sub(1); dflash.accepted_draft_tokens += result.matched_draft_tokens; diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 27b411ec8..0a0f0b36f 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -142,9 +142,10 @@ impl Qwen3Executor { )); } } - // The worker returns the mathematically accepted span. Apply the - // request contract before touching RequestKv so a terminal token's - // speculative suffix is rolled back with the unused reservation. + // The worker normally applies the request contract before copying + // worker-side state. Recheck the same invariant here before touching + // RequestKv so a terminal suffix is rolled back with its reservation, + // including legacy workers that return an untrimmed span. for (policy, req_result) in plan.stop_policies.iter().zip(&mut result.requests) { truncate_after_terminal(req_result, policy, &self.metadata.stop_token_ids); } @@ -170,6 +171,13 @@ impl Qwen3Executor { req_result.request_id )); } + if result.hedged && std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + log::debug!( + "Qwen3 DFlash hedge commit request={} accepted_len={}", + req_result.request_id, + req_result.accepted_tokens.len(), + ); + } applied.push(req_result.request_id); } for req_result in &result.requests { diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index 0614433b4..67a5ed4f4 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -81,6 +81,10 @@ pub(crate) struct VerifyRequestResult { pub(crate) struct VerifyResult { pub requests: Vec, + /// Whether this result came from the expanded hedge path. Test-only + /// tracing uses this marker to pair worker context and KV commit records + /// without confusing them with a plain verify round. + pub hedged: bool, } /// One request's draft request: the proposer continues from `current_token`. diff --git a/pegainfer-qwen3/tests/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index 45d48e057..f66adbd5e 100644 --- a/pegainfer-qwen3/tests/common/harness.rs +++ b/pegainfer-qwen3/tests/common/harness.rs @@ -38,14 +38,14 @@ pub(crate) fn request( params: SamplingParams, max_tokens: usize, ) -> Request { - let stop_policy = StopPolicy { - eos: if params.ignore_eos { + let stop_policy = StopPolicy::new( + if params.ignore_eos { EosPolicy::Ignore } else { EosPolicy::ModelDefault }, - ..StopPolicy::default() - }; + Vec::new(), + ); Request { prompt_tokens, diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 05d81d66a..3b580db6f 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -39,6 +39,9 @@ //! `PEGAINFER_TEST_MODEL_PATH` (target) and `PEGAINFER_DFLASH_TEST_MODEL_PATH` //! (drafter); skips cleanly when either is absent. +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; use std::path::Path; use std::path::PathBuf; use std::process::Command; @@ -626,7 +629,8 @@ fn dflash_concurrent_heterogeneous_is_lossless() { /// Production hedge regression: an explicit stop in the middle of a verify /// span must be applied before the hedge winner is selected and committed. -/// The trigger is retained, while every accepted suffix token is discarded. +/// The parent gate also checks the request-local worker trace, because the +/// final stream alone is protected by the executor's legacy safety truncation. #[test] fn dflash_hedged_midspan_stop_retains_trigger() { common::harness::init_capture_logging(); @@ -643,7 +647,7 @@ fn dflash_hedged_midspan_stop_retains_trigger() { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); - let prompt = "Write a short paragraph about a blue bicycle:"; + let prompt = "Write a short paragraph about a blue"; let tokenizer = common::load_tokenizer(&model_path); let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); let draft_config: serde_json::Value = serde_json::from_str( @@ -675,39 +679,55 @@ fn dflash_hedged_midspan_stop_retains_trigger() { )) .expect_finished() .tokens; - let Some((stop_index, &stop_id)) = baseline + let candidate_stops: Vec<(usize, u32)> = baseline .iter() .enumerate() .skip(min_index) - .find(|(index, token)| !baseline[..*index].contains(token)) - else { - panic!("baseline did not produce a unique token after the middle of the verify span"); - }; + .take(block_size.saturating_sub(min_index)) + .filter(|(index, token)| !baseline[..*index].contains(token)) + .map(|(index, token)| (index, *token)) + .collect(); assert!( - stop_index < block_size, - "selected stop token at position {stop_index}, expected inside the first {block_size}-token verify span" + !candidate_stops.is_empty(), + "baseline did not produce a unique token inside the first verify span" ); - let stopped_params = SamplingParams { - ignore_eos: true, - ..SamplingParams::default() - }; - let mut stopped = request(prompt_tokens, stopped_params, GENERATED_TOKENS); - stopped.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![stop_id]); - let outcome = engine.submit(stopped).expect_finished(); - - assert_eq!(outcome.tokens.len(), stop_index + 1); - assert_eq!(outcome.tokens.last(), Some(&stop_id)); - assert!(!outcome.tokens[..stop_index].contains(&stop_id)); - assert!(matches!( - outcome.terminal, - Terminal::Finished { - reason: FinishReason::Stop, - stop_cause: Some(StopCause::Token(id)), - completion_tokens, - .. - } if id == stop_id && completion_tokens == stop_index + 1 - )); + let mut stopped_cases = 0usize; + for (baseline_index, stop_id) in candidate_stops { + let stopped_params = SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }; + let mut stopped = request(prompt_tokens.clone(), stopped_params, GENERATED_TOKENS); + stopped.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![stop_id]); + let stream = engine.submit(stopped); + let request_id = stream.id(); + eprintln!("hedge stop request={request_id}"); + let outcome = stream.expect_finished(); + + if outcome.tokens.last() != Some(&stop_id) { + continue; + } + assert!(!outcome.tokens[..outcome.tokens.len() - 1].contains(&stop_id)); + assert!(matches!( + outcome.terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(id)), + completion_tokens, + .. + } if id == stop_id && completion_tokens == outcome.tokens.len() + )); + eprintln!( + "hedge stop candidate baseline_index={baseline_index} token={stop_id} retained_len={}", + outcome.tokens.len() + ); + stopped_cases += 1; + } + assert!( + stopped_cases > 0, + "none of the candidate tokens produced a mid-span explicit stop" + ); } /// P2 regression: a request that fits the target context window but lands in the @@ -790,15 +810,14 @@ fn dflash_request_in_draft_headroom_is_rejected_not_panicked() { } } -/// Execution gate for the hedge ladder: re-runs the two losslessness tests +/// Execution gate for the hedge ladder: re-runs the hedge child suites /// above in child processes, since the hedge config is read once per process /// and cannot be toggled in-process. /// -/// Scope: this proves the copy-back and discard branches were ENTERED, not -/// that a later round consumes the winner's state correctly — `check_lossless` -/// returns at the first benign tie flip, so the suffix past it is never -/// compared. A tie is also a non-win (the win test is strictly-greater), so -/// the counters cannot separate a tie from a shorter chain. +/// The stop child additionally checks the request-local raw/retained winner, +/// selected winner, context append, and KV commit lengths. The losslessness +/// children retain their numerical tie tolerance and only require that the +/// configured hedge path actually ran. /// /// Strict token equality against an unhedged run is NOT a valid contract: /// hedged rounds change the verify batch shape, which legally flips bf16 ties, @@ -840,12 +859,11 @@ fn hedged_ladder_passes_the_lossless_gates() { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let exe = std::env::current_exe().expect("test binary path"); - // One child per lossless suite, each with a single exact filter, so the - // invocation shape is beyond dispute on any libtest version. + // One child per suite, each with a single exact filter, so the invocation + // shape is beyond dispute on any libtest version. // A hedge-free child would make its lossless pass vacuous, so each child // must show expanded spans in the executor's per-round trace. let mut total_spans = 0usize; - let mut total_wins = 0usize; let mut total_rounds = 0usize; for child_test in [ "dflash_speculative_greedy_matches_plain_greedy", @@ -868,7 +886,6 @@ fn hedged_ladder_passes_the_lossless_gates() { ); let mut rounds = 0usize; let mut spans = 0usize; - let mut wins = 0usize; for line in child_stderr.lines() { let Some(rest) = line.split("DFlash hedge: ").nth(1) else { continue; @@ -878,23 +895,99 @@ fn hedged_ladder_passes_the_lossless_gates() { .filter(|tok| !tok.is_empty()) .map(|tok| tok.parse::().expect("hedge trace number")); spans += nums.next().expect("span count"); - wins += nums.next().expect("win count"); + let _ = nums.next().expect("win count"); rounds += 1; } assert!( rounds > 0 && spans > 0, "child '{child_test}' executed no hedged verify round:\n{child_stderr}" ); + if child_test == "dflash_hedged_midspan_stop_retains_trigger" { + let stop_requests: HashSet = child_stderr + .lines() + .filter_map(|line| line.strip_prefix("hedge stop request=")) + .map(str::to_owned) + .collect(); + assert!( + !stop_requests.is_empty(), + "stop child emitted no request marker" + ); + let mut context_by_request: HashMap> = HashMap::new(); + let mut commit_by_request: HashMap> = HashMap::new(); + for line in child_stderr.lines() { + let fields: Vec<_> = line.split_whitespace().collect(); + let request = fields + .iter() + .find_map(|field| field.strip_prefix("request=")); + if let Some(request) = request { + if let Some(appended) = fields + .iter() + .find_map(|field| field.strip_prefix("appended=")) + .and_then(|value| value.parse::().ok()) + { + context_by_request + .entry(request.to_string()) + .or_default() + .push_back(appended); + } + if let Some(accepted_len) = fields + .iter() + .find_map(|field| field.strip_prefix("accepted_len=")) + .and_then(|value| value.parse::().ok()) + { + commit_by_request + .entry(request.to_string()) + .or_default() + .push_back(accepted_len); + } + } + } + let mut worker_side_truncation = false; + for line in child_stderr.lines() { + if !line.contains("Qwen3 DFlash hedge detail ") { + continue; + } + let value = |name: &str| { + line.split_whitespace() + .find_map(|field| field.strip_prefix(name)) + }; + let retained_winner = value("retained_winner="); + let selected = value("selected="); + let raw_winner = value("raw_winner="); + let request = value("request="); + let raw_best_len = value("raw_best_len=").and_then(|v| v.parse().ok()); + let retained_best_len = value("retained_best_len=").and_then(|v| v.parse().ok()); + let selected_len = value("selected_len=").and_then(|v| v.parse().ok()); + if !request.is_some_and(|id| stop_requests.contains(id)) { + continue; + } + let context_len = request + .and_then(|id| context_by_request.get_mut(id)) + .and_then(VecDeque::pop_front); + let commit_len = request + .and_then(|id| commit_by_request.get_mut(id)) + .and_then(VecDeque::pop_front); + if raw_winner == Some("B") + && retained_winner == Some("A") + && selected == Some("A") + && retained_best_len + .is_some_and(|retained| raw_best_len.is_some_and(|raw| retained < raw)) + && selected_len.is_some() + && selected_len == retained_best_len + && selected_len == context_len + && selected_len == commit_len + { + worker_side_truncation = true; + break; + } + } + assert!( + worker_side_truncation, + "stop hedge never truncated a candidate before winner/context commit:\n{child_stderr}" + ); + } total_rounds += rounds; total_spans += spans; - total_wins += wins; } - assert!( - total_wins > 0, - "no hedge chain ever won across {total_rounds} hedged rounds" - ); - assert!( - total_spans > total_wins, - "every hedge span won ({total_wins}/{total_spans}) — the discard path never executed" - ); + assert!(total_rounds > 0 && total_spans > 0); } From 8e6d585f309dfafff5e95bd2f9a477c9b4a79004 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Mon, 7 Sep 2026 21:36:36 +0800 Subject: [PATCH 08/15] fix(qwen3): satisfy clippy readability gate Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-qwen3/src/executor.rs | 32 ++++++++++++++++---------------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 7036e7a5b..7d10b0903 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -3549,22 +3549,22 @@ impl LocalQwen3Lane { hedge_spans.len() ); let trace = std::env::var_os("PEGAINFER_TEST_LOG").is_some(); - let raw_a_lengths: Vec = trace - .then(|| { - results_a - .iter() - .map(|result| result.accepted_tokens.len()) - .collect() - }) - .unwrap_or_default(); - let raw_b_lengths: Vec = trace - .then(|| { - results_b - .iter() - .map(|result| result.accepted_tokens.len()) - .collect() - }) - .unwrap_or_default(); + let raw_a_lengths: Vec = if trace { + results_a + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + } else { + Vec::new() + }; + let raw_b_lengths: Vec = if trace { + results_b + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + } else { + Vec::new() + }; for (result, policy) in results_a.iter_mut().zip(stop_policies) { spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); } From b748fccd349ed8e4f67e4d4c0bdcfba922af3e71 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Mon, 7 Sep 2026 23:46:30 +0800 Subject: [PATCH 09/15] test(qwen3): tighten hedge coverage and simplify tracing Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-qwen3/src/executor.rs | 75 ++----------------- pegainfer-qwen3/src/executor/dflash_lane.rs | 9 --- pegainfer-qwen3/src/executor/spec.rs | 4 +- pegainfer-qwen3/src/speculative.rs | 4 - .../tests/dflash_speculative_gate.rs | 45 +++++++---- 5 files changed, 37 insertions(+), 100 deletions(-) diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 7d10b0903..9a5d45463 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -3575,54 +3575,6 @@ impl LocalQwen3Lane { &self.model.config().stop_token_ids, ); } - let ( - retained_a_lengths, - retained_b_lengths, - raw_best_lengths, - raw_best_is_b, - retained_best_lengths, - retained_best_is_b, - ) = if trace { - let retained_a_lengths: Vec = results_a - .iter() - .map(|result| result.accepted_tokens.len()) - .collect(); - let retained_b_lengths: Vec = results_b - .iter() - .map(|result| result.accepted_tokens.len()) - .collect(); - let mut raw_best_lengths = raw_a_lengths.clone(); - let mut raw_best_is_b = vec![false; requests.len()]; - let mut retained_best_lengths = retained_a_lengths.clone(); - let mut retained_best_is_b = vec![false; requests.len()]; - for (slot, (idx, _)) in hedge_spans.iter().enumerate() { - if raw_b_lengths[slot] > raw_best_lengths[*idx] { - raw_best_lengths[*idx] = raw_b_lengths[slot]; - raw_best_is_b[*idx] = true; - } - if retained_b_lengths[slot] > retained_best_lengths[*idx] { - retained_best_lengths[*idx] = retained_b_lengths[slot]; - retained_best_is_b[*idx] = true; - } - } - ( - retained_a_lengths, - retained_b_lengths, - raw_best_lengths, - raw_best_is_b, - retained_best_lengths, - retained_best_is_b, - ) - } else { - ( - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - ) - }; // Per request keep the best-accepting chain; ties keep chain A (no // copies). A later chain of the same request only replaces the // running winner when strictly better, so the final page/hidden @@ -3632,7 +3584,11 @@ impl LocalQwen3Lane { let elem = std::mem::size_of::(); let mut final_requests: Vec = requests.to_vec(); let mut final_results: Vec = results_a.to_vec(); - let mut selected_is_b = trace.then(|| vec![false; requests.len()]); + let mut selected_is_b = if trace { + Some(vec![false; requests.len()]) + } else { + None + }; let mut b_wins = 0usize; let mut b_row_offset = a_total_rows; for (slot, (idx, replaced)) in hedge_spans.iter().enumerate() { @@ -3683,7 +3639,6 @@ impl LocalQwen3Lane { &final_requests, &final_results, Some(bufs.captured_hidden()), - true, )?; if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { for idx in 0..requests.len() { @@ -3693,8 +3648,6 @@ impl LocalQwen3Lane { if !has_hedge { continue; } - let raw_winner = if raw_best_is_b[idx] { 'B' } else { 'A' }; - let retained_winner = if retained_best_is_b[idx] { 'B' } else { 'A' }; let selected = if selected_is_b.as_ref().is_some_and(|selected| selected[idx]) { 'B' } else { @@ -3707,24 +3660,11 @@ impl LocalQwen3Lane { .map(|(slot, _)| raw_b_lengths[slot].to_string()) .collect::>() .join(","); - let retained_b_lens = hedge_spans - .iter() - .enumerate() - .filter(|(_, (request_idx, _))| *request_idx == idx) - .map(|(slot, _)| retained_b_lengths[slot].to_string()) - .collect::>() - .join(","); log::debug!( - "Qwen3 DFlash hedge detail request={} raw_a={} raw_b_lens={} raw_best_len={} kept_a={} kept_b_lens={} retained_best_len={} raw_winner={} retained_winner={} selected={} selected_len={} matched_draft={}", + "Qwen3 DFlash hedge detail request={} raw_a={} raw_b_lens={} selected={} selected_len={} matched_draft={}", requests[idx].request_id, raw_a_lengths[idx], raw_b_lens, - raw_best_lengths[idx], - retained_a_lengths[idx], - retained_b_lens, - retained_best_lengths[idx], - raw_winner, - retained_winner, selected, final_results[idx].accepted_tokens.len(), final_results[idx].matched_draft_tokens, @@ -3733,7 +3673,6 @@ impl LocalQwen3Lane { } Ok(Some(VerifyResult { requests: final_results, - hedged: true, })) } @@ -3850,11 +3789,9 @@ impl LocalQwen3Lane { requests, &request_results, Some(bufs.captured_hidden()), - false, )?; Ok(VerifyResult { requests: request_results, - hedged: false, }) })(); self.verify_bufs = Some(bufs); diff --git a/pegainfer-qwen3/src/executor/dflash_lane.rs b/pegainfer-qwen3/src/executor/dflash_lane.rs index 6a47c4d47..4165643d3 100644 --- a/pegainfer-qwen3/src/executor/dflash_lane.rs +++ b/pegainfer-qwen3/src/executor/dflash_lane.rs @@ -201,7 +201,6 @@ impl LocalQwen3Lane { requests: &[VerifyStepItem], results: &[VerifyRequestResult], captured_hidden: Option<&HiddenStates>, - hedged: bool, ) -> Result<()> { let Some(captured_hidden) = captured_hidden else { anyhow::bail!("DFlash verify context capture requested but no hidden states returned"); @@ -243,14 +242,6 @@ impl LocalQwen3Lane { token_offset, result.accepted_tokens.len(), )?; - if hedged && std::env::var_os("PEGAINFER_TEST_LOG").is_some() { - log::debug!( - "Qwen3 DFlash hedge context request={} appended={} matched_draft={}", - req.request_id, - result.accepted_tokens.len(), - result.matched_draft_tokens, - ); - } dflash.requests.insert(req.request_id, state); dflash.verified_draft_tokens += req.token_ids.len().saturating_sub(1); dflash.accepted_draft_tokens += result.matched_draft_tokens; diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 0a0f0b36f..1fffd0d35 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -171,9 +171,9 @@ impl Qwen3Executor { req_result.request_id )); } - if result.hedged && std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { log::debug!( - "Qwen3 DFlash hedge commit request={} accepted_len={}", + "Qwen3 DFlash commit request={} accepted_len={}", req_result.request_id, req_result.accepted_tokens.len(), ); diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index 67a5ed4f4..0614433b4 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -81,10 +81,6 @@ pub(crate) struct VerifyRequestResult { pub(crate) struct VerifyResult { pub requests: Vec, - /// Whether this result came from the expanded hedge path. Test-only - /// tracing uses this marker to pair worker context and KV commit records - /// without confusing them with a plain verify round. - pub hedged: bool, } /// One request's draft request: the proposer continues from `current_token`. diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 3b580db6f..42acb6567 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -814,10 +814,11 @@ fn dflash_request_in_draft_headroom_is_rejected_not_panicked() { /// above in child processes, since the hedge config is read once per process /// and cannot be toggled in-process. /// -/// The stop child additionally checks the request-local raw/retained winner, -/// selected winner, context append, and KV commit lengths. The losslessness -/// children retain their numerical tie tolerance and only require that the -/// configured hedge path actually ran. +/// The stop child additionally checks that an untrimmed B candidate is not +/// selected over the truncated A candidate, and that context append and KV +/// commit use the same retained length. The losslessness children retain their +/// numerical tie tolerance and only require that the configured hedge path +/// actually ran. /// /// Strict token equality against an unhedged run is NOT a valid contract: /// hedged rounds change the verify batch shape, which legally flips bf16 ties, @@ -865,6 +866,7 @@ fn hedged_ladder_passes_the_lossless_gates() { // must show expanded spans in the executor's per-round trace. let mut total_spans = 0usize; let mut total_rounds = 0usize; + let mut total_wins = 0usize; for child_test in [ "dflash_speculative_greedy_matches_plain_greedy", "dflash_concurrent_heterogeneous_is_lossless", @@ -895,7 +897,10 @@ fn hedged_ladder_passes_the_lossless_gates() { .filter(|tok| !tok.is_empty()) .map(|tok| tok.parse::().expect("hedge trace number")); spans += nums.next().expect("span count"); - let _ = nums.next().expect("win count"); + let wins = nums.next().expect("win count"); + if child_test != "dflash_hedged_midspan_stop_retains_trigger" { + total_wins += wins; + } rounds += 1; } assert!( @@ -951,12 +956,15 @@ fn hedged_ladder_passes_the_lossless_gates() { line.split_whitespace() .find_map(|field| field.strip_prefix(name)) }; - let retained_winner = value("retained_winner="); let selected = value("selected="); - let raw_winner = value("raw_winner="); let request = value("request="); - let raw_best_len = value("raw_best_len=").and_then(|v| v.parse().ok()); - let retained_best_len = value("retained_best_len=").and_then(|v| v.parse().ok()); + let raw_a = value("raw_a=").and_then(|v| v.parse().ok()); + let raw_b_lens = value("raw_b_lens=").map(|v| { + v.split(',') + .filter(|value| !value.is_empty()) + .map(|value| value.parse::().expect("raw B length")) + .collect::>() + }); let selected_len = value("selected_len=").and_then(|v| v.parse().ok()); if !request.is_some_and(|id| stop_requests.contains(id)) { continue; @@ -967,13 +975,14 @@ fn hedged_ladder_passes_the_lossless_gates() { let commit_len = request .and_then(|id| commit_by_request.get_mut(id)) .and_then(VecDeque::pop_front); - if raw_winner == Some("B") - && retained_winner == Some("A") + let raw_b_max = raw_b_lens + .as_ref() + .and_then(|lengths| lengths.iter().copied().max()); + if raw_b_max.is_some_and(|raw_b| raw_a.is_some_and(|raw_a| raw_b > raw_a)) && selected == Some("A") - && retained_best_len - .is_some_and(|retained| raw_best_len.is_some_and(|raw| retained < raw)) + && selected_len + .is_some_and(|selected| raw_b_max.is_some_and(|raw| selected < raw)) && selected_len.is_some() - && selected_len == retained_best_len && selected_len == context_len && selected_len == commit_len { @@ -986,8 +995,12 @@ fn hedged_ladder_passes_the_lossless_gates() { "stop hedge never truncated a candidate before winner/context commit:\n{child_stderr}" ); } - total_rounds += rounds; - total_spans += spans; + if child_test != "dflash_hedged_midspan_stop_retains_trigger" { + total_rounds += rounds; + total_spans += spans; + } } assert!(total_rounds > 0 && total_spans > 0); + assert!(total_wins > 0); + assert!(total_spans > total_wins); } From 397bf0822f1428e9f58c656436e4812c17e51ca5 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sat, 12 Sep 2026 13:33:42 +0800 Subject: [PATCH 10/15] fix(qwen3): finalize stop policy before hedged commit Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-frontend/src/engine/stop.rs | 17 +++- pegainfer-frontend/src/vllm/wire.rs | 25 ++++- pegainfer-qwen3/src/executor.rs | 23 ++++- pegainfer-qwen3/src/executor/dflash_lane.rs | 8 +- pegainfer-qwen3/src/executor/spec.rs | 6 +- pegainfer-qwen3/src/scheduler/test_support.rs | 2 + .../tests/dflash_speculative_gate.rs | 98 ++++++++++--------- 7 files changed, 122 insertions(+), 57 deletions(-) diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs index 83c62b3fc..a80f23408 100644 --- a/pegainfer-frontend/src/engine/stop.rs +++ b/pegainfer-frontend/src/engine/stop.rs @@ -8,6 +8,9 @@ pub enum EosPolicy { /// Use the model executor's configured EOS set. #[default] ModelDefault, + /// Use this protocol-provided primary EOS ID; secondary EOS IDs arrive + /// in the request's explicit stop-token set. + Token(u32), } /// Request-scoped token stopping policy. @@ -51,12 +54,13 @@ impl StopPolicy { token_id: u32, is_model_eos: impl FnOnce(u32) -> bool, ) -> Option { - let is_eos = match self.eos { + let is_primary_eos = match self.eos { EosPolicy::Ignore => false, EosPolicy::ModelDefault => is_model_eos(token_id), + EosPolicy::Token(eos_token_id) => token_id == eos_token_id, }; - if is_eos { + if is_primary_eos { Some(StopCause::Eos(token_id)) } else if self.token_ids.binary_search(&token_id).is_ok() { Some(StopCause::Token(token_id)) @@ -118,4 +122,13 @@ mod tests { Some(StopCause::Eos(99)) ); } + + #[test] + fn primary_eos_keeps_secondary_model_eos_as_token_stop() { + let policy = StopPolicy::new(EosPolicy::Token(1), vec![1, 2]); + + assert_eq!(policy.classify(1, |_| true), Some(StopCause::Eos(1))); + assert_eq!(policy.classify(2, |_| true), Some(StopCause::Token(2))); + assert_eq!(policy.classify(3, |_| true), None); + } } diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index 299a16ab9..98db32a0a 100644 --- a/pegainfer-frontend/src/vllm/wire.rs +++ b/pegainfer-frontend/src/vllm/wire.rs @@ -115,12 +115,12 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar pub(crate) fn convert_stop_policy(params: &EngineCoreSamplingParams) -> StopPolicy { StopPolicy::new( - // Qwen3 owns the complete model EOS set in generation_config. The - // protocol's optional primary ID only tells us whether EOS is active; - // using it as a singleton would miss secondary model EOS IDs. + // vLLM lowers secondary model EOS IDs into stop_token_ids. Keep its + // primary EOS separate so secondary stops retain their wire reason; + // all_stop_token_ids is only for min_tokens masking. params .eos_token_id - .map_or(EosPolicy::Ignore, |_| EosPolicy::ModelDefault), + .map_or(EosPolicy::Ignore, EosPolicy::Token), params.stop_token_ids.clone(), ) } @@ -269,7 +269,7 @@ mod tests { fn convert_stop_policy_keeps_eos_and_explicit_stops_independent() { let mut params = EngineCoreSamplingParams::for_test(); params.eos_token_id = Some(99); - params.stop_token_ids = vec![11]; + params.stop_token_ids = vec![11, 100]; let policy = convert_stop_policy(¶ms); assert_eq!( @@ -280,6 +280,21 @@ mod tests { policy.classify(11, |_| false), Some(crate::engine::StopCause::Token(11)) ); + assert_eq!( + policy.classify(100, |_| true), + Some(crate::engine::StopCause::Token(100)), + "secondary EOS IDs must retain a token stop reason" + ); + + params.eos_token_id = None; + params.stop_token_ids = vec![99]; + params.all_stop_token_ids = BTreeSet::from([99, 100]); + let policy = convert_stop_policy(¶ms); + assert_eq!( + policy.classify(99, |_| true), + Some(crate::engine::StopCause::Token(99)) + ); + assert_eq!(policy.classify(100, |_| true), None); } #[test] diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 9a5d45463..5bdff65b5 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -497,6 +497,7 @@ fn execute_step_on_lane( kv_views, stop_policies, sample_seed, + verify_round, } => { // One target forward over each request's K+1 draft span with a // speculative KV view. The fixed-buffer verify path computes all- @@ -504,8 +505,13 @@ fn execute_step_on_lane( // token at each span position) and captures the target hidden states // (at the DFlash layers) to seed the next draft — all into reused, // pointer-stable scratch (`VerifyGraphBuffers`). - let result = - lane.execute_dflash_verify(requests, kv_views, stop_policies, *sample_seed)?; + let result = lane.execute_dflash_verify( + requests, + kv_views, + stop_policies, + *sample_seed, + *verify_round, + )?; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -945,6 +951,8 @@ pub struct Qwen3Executor { /// [`enable_decode_overlap`] to create overlap streams on the correct /// device (the model, KV cache, and compute stream all live here). device_ordinal: usize, + /// Monotonic ID for correlating per-round speculative verify diagnostics. + verify_round: u64, } /// One request's in-flight CPU-tier KV prefetch. @@ -1147,6 +1155,7 @@ impl Qwen3Executor { spec_decode_counters: None, dflash_ready_requests: HashSet::new(), device_ordinal, + verify_round: 0, }) } @@ -1539,6 +1548,7 @@ impl Qwen3Executor { spec_decode_counters: None, dflash_ready_requests: HashSet::new(), device_ordinal: device_ordinals[0], + verify_round: 0, }) } @@ -3431,6 +3441,7 @@ impl LocalQwen3Lane { stop_policies: &[StopPolicy], capture_layer_ids: &[usize], sample_seed: u64, + verify_round: u64, bufs: &mut VerifyGraphBuffers, ) -> Result> { let page_size = self.layout.page_size; @@ -3639,6 +3650,7 @@ impl LocalQwen3Lane { &final_requests, &final_results, Some(bufs.captured_hidden()), + verify_round, )?; if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { for idx in 0..requests.len() { @@ -3661,7 +3673,8 @@ impl LocalQwen3Lane { .collect::>() .join(","); log::debug!( - "Qwen3 DFlash hedge detail request={} raw_a={} raw_b_lens={} selected={} selected_len={} matched_draft={}", + "Qwen3 DFlash hedge detail round={} request={} raw_a={} raw_b_lens={} selected={} selected_len={} matched_draft={}", + verify_round, requests[idx].request_id, raw_a_lengths[idx], raw_b_lens, @@ -3686,6 +3699,7 @@ impl LocalQwen3Lane { kv_views: &[KvView], stop_policies: &[StopPolicy], sample_seed: u64, + verify_round: u64, ) -> Result { anyhow::ensure!( stop_policies.len() == requests.len(), @@ -3745,6 +3759,7 @@ impl LocalQwen3Lane { stop_policies, &capture_layer_ids, sample_seed, + verify_round, &mut bufs, )? { return Ok(result); @@ -3789,6 +3804,7 @@ impl LocalQwen3Lane { requests, &request_results, Some(bufs.captured_hidden()), + verify_round, )?; Ok(VerifyResult { requests: request_results, @@ -3900,6 +3916,7 @@ enum StepCommand { /// these are host metadata and never enter the GPU batch. stop_policies: Vec, sample_seed: u64, + verify_round: u64, }, /// Speculative draft: roll the DFlash draft model forward one block per /// request. Uses the draft's own KV — no target KV views. diff --git a/pegainfer-qwen3/src/executor/dflash_lane.rs b/pegainfer-qwen3/src/executor/dflash_lane.rs index 4165643d3..4b17c45a8 100644 --- a/pegainfer-qwen3/src/executor/dflash_lane.rs +++ b/pegainfer-qwen3/src/executor/dflash_lane.rs @@ -201,6 +201,7 @@ impl LocalQwen3Lane { requests: &[VerifyStepItem], results: &[VerifyRequestResult], captured_hidden: Option<&HiddenStates>, + verify_round: u64, ) -> Result<()> { let Some(captured_hidden) = captured_hidden else { anyhow::bail!("DFlash verify context capture requested but no hidden states returned"); @@ -242,6 +243,7 @@ impl LocalQwen3Lane { token_offset, result.accepted_tokens.len(), )?; + let context_len = state.pending_context_len().unwrap_or(0); dflash.requests.insert(req.request_id, state); dflash.verified_draft_tokens += req.token_ids.len().saturating_sub(1); dflash.accepted_draft_tokens += result.matched_draft_tokens; @@ -251,10 +253,12 @@ impl LocalQwen3Lane { dflash.accepted_draft_tokens as f64 / dflash.verified_draft_tokens as f64 }; log::debug!( - "Qwen3 DFlash request={} accepted_draft={} committed_tokens={} cumulative_accept_rate={:.3}", - req.request_id.raw(), + "Qwen3 DFlash context round={} request={} accepted_draft={} committed_tokens={} context_len={} cumulative_accept_rate={:.3}", + verify_round, + req.request_id, result.matched_draft_tokens, result.accepted_tokens.len(), + context_len, rate, ); token_offset += req.token_ids.len(); diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 1fffd0d35..3ecb1e7f2 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -48,6 +48,8 @@ impl Qwen3Executor { &mut self, plan: VerifyPlan<'_>, ) -> Result { + let verify_round = self.verify_round; + self.verify_round = self.verify_round.wrapping_add(1); anyhow::ensure!( self.speculative.is_some(), "speculative verification requested but no draft model is loaded" @@ -106,6 +108,7 @@ impl Qwen3Executor { kv_views, stop_policies: plan.stop_policies.to_vec(), sample_seed: plan.sample_seed, + verify_round, }; let outcome = match self.run_step(&step) { Ok(outcome) => outcome, @@ -173,7 +176,8 @@ impl Qwen3Executor { } if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { log::debug!( - "Qwen3 DFlash commit request={} accepted_len={}", + "Qwen3 DFlash commit round={} request={} accepted_len={}", + verify_round, req_result.request_id, req_result.accepted_tokens.len(), ); diff --git a/pegainfer-qwen3/src/scheduler/test_support.rs b/pegainfer-qwen3/src/scheduler/test_support.rs index cfd639df7..76a597f6a 100644 --- a/pegainfer-qwen3/src/scheduler/test_support.rs +++ b/pegainfer-qwen3/src/scheduler/test_support.rs @@ -111,6 +111,7 @@ impl FakeExecutor { let token = 100 + req.request_id.raw() as u32; pegainfer_frontend::engine::TokenLogprob { logprob: -0.1, + rank: 0, top_logprobs: vec![(token, -0.1)], } }), @@ -242,6 +243,7 @@ impl ModelExecutor for FakeExecutor { let token = 200 + req.request_id.raw() as u32; pegainfer_frontend::engine::TokenLogprob { logprob: -0.2, + rank: 0, top_logprobs: vec![(token, -0.2)], } }), diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 42acb6567..73c22f724 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -41,7 +41,6 @@ use std::collections::HashMap; use std::collections::HashSet; -use std::collections::VecDeque; use std::path::Path; use std::path::PathBuf; use std::process::Command; @@ -917,35 +916,35 @@ fn hedged_ladder_passes_the_lossless_gates() { !stop_requests.is_empty(), "stop child emitted no request marker" ); - let mut context_by_request: HashMap> = HashMap::new(); - let mut commit_by_request: HashMap> = HashMap::new(); + let mut context_by_round_request: HashMap<(u64, String), usize> = HashMap::new(); + let mut commit_by_round_request: HashMap<(u64, String), usize> = HashMap::new(); for line in child_stderr.lines() { - let fields: Vec<_> = line.split_whitespace().collect(); - let request = fields - .iter() - .find_map(|field| field.strip_prefix("request=")); - if let Some(request) = request { - if let Some(appended) = fields - .iter() - .find_map(|field| field.strip_prefix("appended=")) - .and_then(|value| value.parse::().ok()) - { - context_by_request - .entry(request.to_string()) - .or_default() - .push_back(appended); - } - if let Some(accepted_len) = fields - .iter() - .find_map(|field| field.strip_prefix("accepted_len=")) - .and_then(|value| value.parse::().ok()) - { - commit_by_request - .entry(request.to_string()) - .or_default() - .push_back(accepted_len); - } - } + let (length_field, lengths) = if line.contains("Qwen3 DFlash context ") { + ("context_len=", &mut context_by_round_request) + } else if line.contains("Qwen3 DFlash commit ") { + ("accepted_len=", &mut commit_by_round_request) + } else { + continue; + }; + let value = |name: &str| { + line.split_whitespace() + .find_map(|field| field.strip_prefix(name)) + }; + let round = value("round=") + .expect("verify round") + .parse::() + .expect("numeric verify round"); + let request = value("request=").expect("verify request"); + let length = value(length_field) + .expect("verify length") + .parse::() + .expect("numeric verify length"); + assert!( + lengths + .insert((round, request.to_string()), length) + .is_none(), + "duplicate {length_field} for round={round} request={request}" + ); } let mut worker_side_truncation = false; for line in child_stderr.lines() { @@ -957,7 +956,11 @@ fn hedged_ladder_passes_the_lossless_gates() { .find_map(|field| field.strip_prefix(name)) }; let selected = value("selected="); - let request = value("request="); + let round = value("round=") + .expect("hedge round") + .parse::() + .expect("numeric hedge round"); + let request = value("request=").expect("hedge request"); let raw_a = value("raw_a=").and_then(|v| v.parse().ok()); let raw_b_lens = value("raw_b_lens=").map(|v| { v.split(',') @@ -965,29 +968,36 @@ fn hedged_ladder_passes_the_lossless_gates() { .map(|value| value.parse::().expect("raw B length")) .collect::>() }); - let selected_len = value("selected_len=").and_then(|v| v.parse().ok()); - if !request.is_some_and(|id| stop_requests.contains(id)) { + let selected_len = value("selected_len=") + .expect("selected length") + .parse::() + .expect("numeric selected length"); + if !stop_requests.contains(request) { continue; } - let context_len = request - .and_then(|id| context_by_request.get_mut(id)) - .and_then(VecDeque::pop_front); - let commit_len = request - .and_then(|id| commit_by_request.get_mut(id)) - .and_then(VecDeque::pop_front); + let key = (round, request.to_string()); + let context_len = context_by_round_request + .get(&key) + .expect("same-round DFlash context record"); + let commit_len = commit_by_round_request + .get(&key) + .expect("same-round KV commit record"); + assert_eq!( + selected_len, *context_len, + "hedge/context mismatch for round={round} request={request}" + ); + assert_eq!( + selected_len, *commit_len, + "hedge/commit mismatch for round={round} request={request}" + ); let raw_b_max = raw_b_lens .as_ref() .and_then(|lengths| lengths.iter().copied().max()); if raw_b_max.is_some_and(|raw_b| raw_a.is_some_and(|raw_a| raw_b > raw_a)) && selected == Some("A") - && selected_len - .is_some_and(|selected| raw_b_max.is_some_and(|raw| selected < raw)) - && selected_len.is_some() - && selected_len == context_len - && selected_len == commit_len + && raw_b_max.is_some_and(|raw| selected_len < raw) { worker_side_truncation = true; - break; } } assert!( From 8def99771ba31d91610218cb1dda78580dcf588a Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sun, 13 Sep 2026 14:04:17 +0800 Subject: [PATCH 11/15] docs: drop unrelated frontend architecture changes Signed-off-by: RicardoMin <17879681016@163.com> --- .../frontend/frontend-architecture.md | 33 ++++++++----------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index e314e004e..239e2a61b 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -12,19 +12,14 @@ An engine is a set of schedulers, each a `Scheduler` implementation driven by th ``` pegainfer-frontend/src/engine/ -├── step.rs # the wire: RequestId, Request, StepOutputs { Vec }, -│ # Request { ..., stop_policy }, RequestUpdate { scheduled, tokens, -│ # logprobs, cached_tokens, prompt_echo, kv_transfer, terminal }, -│ # Terminal { ..., stop_cause } -├── request_lifecycle.rs # typestate handles: QueuedRequest ─admit→ -│ # ActiveRequest ─finish/fail/defer→ consumed; every -│ # transition is by-move, a dropped handle emits Failed -│ # (drop bomb); DeferredFinish; RequestControl -├── emitter.rs # StepEmitter: the single writer of the per-step buffer; stamps -│ # timestamps, tallies prompt/completion counts, folds each -│ # request's step into one RequestUpdate; commit_step sends once -├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, -│ # prompt/completion tallies, one merged update per touched id +├── step.rs # the wire: RequestId, Request, QueuedRequest, +│ # StepOutputs { Vec }, +│ # RequestUpdate { scheduled, tokens, logprobs, cached_tokens, +│ # prompt_echo, kv_transfer, terminal }, Terminal +├── request_lifecycle.rs # submission envelope, abort control and step sender plumbing; +│ # DeferredFinish remains available for P/D handoff +├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, +│ # prompt/completion tallies, one merged update per touched id ├── wiring.rs # scheduler_pair, SchedulerHandle (submit/take_steps/load), │ # Engine { schedulers, info, lora }, LiveScheduler, │ # EngineInfo, LaunchedEngine { Handle | Stepped } @@ -37,12 +32,10 @@ pegainfer-frontend/src/engine/ Design decisions worth knowing before touching it: - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. -- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. This is what makes `defer_finish` safe: a P/D prefill executor can withhold a request's `Finished` until its KV saves are peer-visible and send it later from any thread — the deferred message carries the request's entire buffered update, so late delivery cannot reorder. -- **Independent stop policy and cause.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in `RequestUpdate` and carries `Terminal::Finished.stop_cause`; decode paths include its real logprob when available, so the stepped bridge can report the actual explicit stop ID without reconstructing it. `ignore_eos` affects only model EOS. Legacy producers may leave the cause empty while they are migrated individually. -- **Typestate lifecycle.** `QueuedRequest` (queued) and `ActiveRequest` (streaming) are owned tokens; admit/reject/retire/finish/fail consume them, so "terminal exactly once, nothing after it" cannot be miscoded — it does not compile. A handle dropped without a transition emits `Failed` from its `Drop`, which is also how a crashed scheduler answers every in-flight request: the driver drops the scheduler, the handles fall, the terminals ship. -- **Emitter as single writer.** Schedulers never touch the channel; they call `StepEmitter` methods against their handles. The emitter stamps `ScheduledInfo` at admission, tallies token counts (terminal counts derive from the tally, never from model-side arithmetic), and `commit_step` publishes the whole step in one send. -- **Pure polling driver.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish load, commit. No idle/park distinction — the scheduler owns the GPU and spinning on it costs nothing anyone else could use; async KV I/O (prefetch, decode-overlap prefill) is naturally absorbed by polling. An idle iteration ends in a `spin_loop` hint (relaxes the core's issue slots, no latency cost — busy iterations never pause). The loop exits when the frontend drops the handle and the queue drains. -- **Gemma 4 async prefill exception.** While asynchronous prefill is the only remaining work, Gemma 4 drains and joins that lane rather than hot-polling its completion; decode or queued work keeps the normal polling path. +- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. +- **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. +- **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. +- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. - **Abort is a flag, not channel teardown.** `SchedulerHandle::submit` returns a `RequestControl`; the frontend flips its boolean abort flag and the scheduler retires the request silently on its next touch (no terminal — the frontend already dropped its state for that id). - **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::load()` — pull-only by design, "notify me on load change" is deliberately unrepresentable (the driver busy-polls, so a subscription edge would fire per spin). All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit. - **Control plane lives outside the contract.** `Scheduler` has no control method and the contract carries no control channel. A capability like LoRA is a private channel the model crate mints *before* `spawn_scheduler` — the scheduler closes over the receiver, the `LoraClient` sender surfaces as `Engine.lora: Option`, and the `Option` *is* the capability (no `bool` flag, no registry until a second capability exists). The vocabulary (`LoraControl`, `LoraClient`) is still defined in the frontend crate because the frontend must speak it without holding model structs; only the wiring is the model's business. @@ -98,7 +91,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Protocol stacks -**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a typed `StopCause` reports the actual request stop token, while the synthetic sentinel remains only as a compatibility fallback for producers that provide no cause). +**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` appends the stop sentinel token, which is how usage keeps counting the suppressed EOS). **`dynamo` (planned second stack).** dynamo's `lib/llm` in-process path removes the wire protocol entirely (`EngineConfig::InProcessTokens` + `run_input`). The step contract was shaped so this stack can consume `StepOutputs` directly without impersonation overhead. Decision gate: prototype, A/B against the vllm stack, let TTFT/step-overhead numbers pick the default. From 22e28aeefec2be62b556b818abbee1ddef6f3ea7 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sun, 13 Sep 2026 21:28:09 +0800 Subject: [PATCH 12/15] fix(qwen3): validate hedged stop candidates and harden stop-contract probe Signed-off-by: RicardoMin <17879681016@163.com> --- docs/models/qwen3/model-crate.md | 2 +- .../frontend/frontend-architecture.md | 1 + pegainfer-frontend/src/vllm/wire.rs | 17 +- pegainfer-qwen3/src/executor.rs | 48 +- pegainfer-qwen3/src/executor/spec.rs | 11 +- pegainfer-qwen3/src/frontend_adapter/tests.rs | 10 +- pegainfer-qwen3/src/scheduler/plan.rs | 18 +- pegainfer-qwen3/src/speculative.rs | 28 +- .../tests/dflash_speculative_gate.rs | 55 +- scripts/qwen3_stop_contract_probe.py | 853 ++++++++++++++++++ 10 files changed, 954 insertions(+), 89 deletions(-) create mode 100644 scripts/qwen3_stop_contract_probe.py diff --git a/docs/models/qwen3/model-crate.md b/docs/models/qwen3/model-crate.md index 2f3468ffa..840b07f08 100644 --- a/docs/models/qwen3/model-crate.md +++ b/docs/models/qwen3/model-crate.md @@ -157,7 +157,7 @@ pub fn start_engine( ### Step 7: Retire ModelForward and Fix Length Limit - Deleted `pegainfer_core::model::{ModelForward, GenerationState}` and removed the root `src/model.rs` re-export. - Deleted the Qwen3 `forward.rs` compatibility path. Qwen3 tests that used it now build their baselines from `batch_prefill(bs=1)` plus `batch_decode(bs=1)`, so they exercise the same phase APIs as production. -- Fixed Qwen3 decode length-limit handling by adding `DecodeEffect::EmitAndFinish`. EOS behavior is unchanged: EOS finishes without emitting the stop token. Length limit now emits the sampled final token, then sends `Finished { finish_reason: Length }`. +- Fixed Qwen3 decode length-limit handling by adding `DecodeEffect::EmitAndFinish`. Length limit now emits the sampled final token, then sends `Finished { finish_reason: Length }`. Under the typed stop contract the triggering token — model EOS or an explicit stop ID — is retained and emitted as the final token; model EOS carries no wire `stop_reason`, while an explicit stop reports its ID. - Regenerated `test_data/Qwen3-4B.json` because every length-limited golden output now includes the final requested token. - Re-ran `bench_serving snapshot` on the CUDA validation host and pulled back `bench_snapshots/rtx-5090/qwen3-4b.json`; `decode_heavy (1024,256)` now records `generated_tokens min=max=avg=256`. - Performance stayed within noise on RTX 5090: diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index 239e2a61b..90fc57b92 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -33,6 +33,7 @@ Design decisions worth knowing before touching it: - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. - **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. +- **Typed stop contract.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in the emitted output and carries `Terminal::Finished.stop_cause`: `Eos(id)` has no wire `stop_reason`, while `Token(id)` reports the actual matching ID. `ignore_eos` disables only model EOS. Legacy producers may leave the cause empty while they are migrated individually. - **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. - **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. - **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index 98db32a0a..a1a75e1ed 100644 --- a/pegainfer-frontend/src/vllm/wire.rs +++ b/pegainfer-frontend/src/vllm/wire.rs @@ -266,30 +266,29 @@ mod tests { } #[test] - fn convert_stop_policy_keeps_eos_and_explicit_stops_independent() { + fn convert_stop_policy_keeps_primary_and_secondary_eos_distinct() { let mut params = EngineCoreSamplingParams::for_test(); params.eos_token_id = Some(99); - params.stop_token_ids = vec![11, 100]; + params.stop_token_ids = vec![100]; let policy = convert_stop_policy(¶ms); + // The primary protocol EOS is carried as `EosPolicy::Token`, so it is + // still an EOS stop even when the model predicate is unavailable. assert_eq!( - policy.classify(99, |token_id| token_id == 99), + policy.classify(99, |_| false), Some(crate::engine::StopCause::Eos(99)) ); - assert_eq!( - policy.classify(11, |_| false), - Some(crate::engine::StopCause::Token(11)) - ); + // Secondary model EOS IDs must keep their token stop reason. assert_eq!( policy.classify(100, |_| true), - Some(crate::engine::StopCause::Token(100)), - "secondary EOS IDs must retain a token stop reason" + Some(crate::engine::StopCause::Token(100)) ); params.eos_token_id = None; params.stop_token_ids = vec![99]; params.all_stop_token_ids = BTreeSet::from([99, 100]); let policy = convert_stop_policy(¶ms); + // EOS disabled: the primary ID stops only when explicitly requested. assert_eq!( policy.classify(99, |_| true), Some(crate::engine::StopCause::Token(99)) diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 5bdff65b5..290380276 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -19,7 +19,6 @@ use pegainfer_core::weight_loader::load_shard_info; use pegainfer_frontend::engine::DeferredFinish; use pegainfer_frontend::engine::LoadLoraAdapterRequest; use pegainfer_frontend::engine::SpecDecodeCounters; -use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::engine::panic_message; @@ -495,7 +494,6 @@ fn execute_step_on_lane( StepCommand::SpeculativeVerify { requests, kv_views, - stop_policies, sample_seed, verify_round, } => { @@ -505,13 +503,8 @@ fn execute_step_on_lane( // token at each span position) and captures the target hidden states // (at the DFlash layers) to seed the next draft — all into reused, // pointer-stable scratch (`VerifyGraphBuffers`). - let result = lane.execute_dflash_verify( - requests, - kv_views, - stop_policies, - *sample_seed, - *verify_round, - )?; + let result = + lane.execute_dflash_verify(requests, kv_views, *sample_seed, *verify_round)?; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -3438,7 +3431,6 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], - stop_policies: &[StopPolicy], capture_layer_ids: &[usize], sample_seed: u64, verify_round: u64, @@ -3519,7 +3511,12 @@ impl LocalQwen3Lane { replaced.push((orig, scratch_page)); } next_scratch += span_pages; - expanded.push(VerifyStepItem::new(req.request_id, ids.clone(), req.params)); + expanded.push(VerifyStepItem::new( + req.request_id, + ids.clone(), + req.params, + req.stop_policy.clone(), + )); views.push(KvView::new(pages, v.seq_len(), page_size)); hedge_spans.push((idx, replaced)); added.push(ids); @@ -3576,13 +3573,17 @@ impl LocalQwen3Lane { } else { Vec::new() }; - for (result, policy) in results_a.iter_mut().zip(stop_policies) { - spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + for (result, req) in results_a.iter_mut().zip(requests) { + spec::truncate_after_terminal( + result, + &req.stop_policy, + &self.model.config().stop_token_ids, + ); } for (slot, (idx, _)) in hedge_spans.iter().enumerate() { spec::truncate_after_terminal( &mut results_b[slot], - &stop_policies[*idx], + &requests[*idx].stop_policy, &self.model.config().stop_token_ids, ); } @@ -3697,16 +3698,9 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], - stop_policies: &[StopPolicy], sample_seed: u64, verify_round: u64, ) -> Result { - anyhow::ensure!( - stop_policies.len() == requests.len(), - "DFlash verify received {} stop policies for {} requests", - stop_policies.len(), - requests.len() - ); let capture_layer_ids = self.dflash_capture_layer_ids().ok_or_else(|| { anyhow::anyhow!("DFlash verify requested but no draft model is loaded") })?; @@ -3756,7 +3750,6 @@ impl LocalQwen3Lane { if let Some(result) = self.try_execute_hedged_verify( requests, kv_views, - stop_policies, &capture_layer_ids, sample_seed, verify_round, @@ -3797,8 +3790,12 @@ impl LocalQwen3Lane { // Apply the request policy before recording target hidden states; // otherwise a suffix discarded by terminal handling would leak // into the next DFlash draft context and acceptance counters. - for (policy, result) in stop_policies.iter().zip(&mut request_results) { - spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + for (req, result) in requests.iter().zip(&mut request_results) { + spec::truncate_after_terminal( + result, + &req.stop_policy, + &self.model.config().stop_token_ids, + ); } self.record_verify_dflash_context( requests, @@ -3912,9 +3909,6 @@ enum StepCommand { SpeculativeVerify { requests: Vec, kv_views: Vec, - /// Request-local stop policies used before DFlash context recording; - /// these are host metadata and never enter the GPU batch. - stop_policies: Vec, sample_seed: u64, verify_round: u64, }, diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 3ecb1e7f2..967dc6772 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -54,12 +54,6 @@ impl Qwen3Executor { self.speculative.is_some(), "speculative verification requested but no draft model is loaded" ); - anyhow::ensure!( - plan.stop_policies.len() == plan.requests.len(), - "speculative verify received {} stop policies for {} requests", - plan.stop_policies.len(), - plan.requests.len() - ); for req in plan.requests { anyhow::ensure!( !req.as_slice().is_empty(), @@ -106,7 +100,6 @@ impl Qwen3Executor { let step = StepCommand::SpeculativeVerify { requests: plan.requests.to_vec(), kv_views, - stop_policies: plan.stop_policies.to_vec(), sample_seed: plan.sample_seed, verify_round, }; @@ -149,8 +142,8 @@ impl Qwen3Executor { // worker-side state. Recheck the same invariant here before touching // RequestKv so a terminal suffix is rolled back with its reservation, // including legacy workers that return an untrimmed span. - for (policy, req_result) in plan.stop_policies.iter().zip(&mut result.requests) { - truncate_after_terminal(req_result, policy, &self.metadata.stop_token_ids); + for (req, req_result) in plan.requests.iter().zip(&mut result.requests) { + truncate_after_terminal(req_result, &req.stop_policy, &self.metadata.stop_token_ids); } // Commit the accepted prefix of each request's KV and free the rest. diff --git a/pegainfer-qwen3/src/frontend_adapter/tests.rs b/pegainfer-qwen3/src/frontend_adapter/tests.rs index 1d1b2244d..b4a5be33c 100644 --- a/pegainfer-qwen3/src/frontend_adapter/tests.rs +++ b/pegainfer-qwen3/src/frontend_adapter/tests.rs @@ -79,14 +79,8 @@ impl StepCollector { /// Fold this request's stream to its end: all tokens in order plus the /// terminal. Panics if the stream closes without a terminal. fn collect_terminal(&mut self, id: RequestId) -> (Vec, Terminal) { - let mut tokens = Vec::new(); - loop { - let update = self.next_for(id); - tokens.extend_from_slice(&update.tokens); - if let Some(terminal) = update.terminal { - return (tokens, terminal); - } - } + let (tokens, _, terminal) = self.collect_terminal_with_logprobs(id); + (tokens, terminal) } fn collect_terminal_with_logprobs( diff --git a/pegainfer-qwen3/src/scheduler/plan.rs b/pegainfer-qwen3/src/scheduler/plan.rs index cf472d4d9..c9c2d2d4c 100644 --- a/pegainfer-qwen3/src/scheduler/plan.rs +++ b/pegainfer-qwen3/src/scheduler/plan.rs @@ -118,11 +118,9 @@ pub(crate) fn execute_plan( requests: &draft_requests, })?; draft.requests.sort_by_key(|result| result.request_id); - let (verify_requests, stop_policies) = - build_speculative_verify_items(active, &draft.requests); + let verify_requests = build_speculative_verify_items(active, &draft.requests); let mut verify = executor.execute_speculative_verify(VerifyPlan { requests: &verify_requests, - stop_policies: &stop_policies, sample_seed: rand::RngExt::random(rng), })?; verify.requests.sort_by_key(|result| result.request_id); @@ -184,12 +182,8 @@ fn build_speculative_draft_items(active: &[ActiveRequestState]) -> Vec ( - Vec, - Vec, -) { +) -> Vec { let mut requests = Vec::with_capacity(draft_results.len()); - let mut stop_policies = Vec::with_capacity(draft_results.len()); for draft in draft_results { let active = active .iter() @@ -209,10 +203,10 @@ fn build_speculative_verify_items( draft.request_id, token_ids, active.params, + active.stop_policy.clone(), )); - stop_policies.push(active.stop_policy.clone()); } - (requests, stop_policies) + requests } fn build_prefill_items(pending: &[PendingRequest], indices: &[usize]) -> Vec { @@ -320,13 +314,13 @@ mod tests { token_ids: (0..16).collect(), }; - let (verify, stop_policies) = build_speculative_verify_items(&active, &[draft]); + let verify = build_speculative_verify_items(&active, &[draft]); assert_eq!(verify.len(), 1); // 32 - 24 = 8 remaining → the 16-token span truncates to 8. assert_eq!(verify[0].as_slice().len(), 8); assert_eq!(verify[0].as_slice(), (0..8).collect::>()); - assert_eq!(stop_policies, vec![StopPolicy::default()]); + assert_eq!(verify[0].stop_policy, StopPolicy::default()); } // The plan selector is the whole batch-formation policy: what the scheduler diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index 0614433b4..1606b0ea4 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -38,14 +38,21 @@ pub(crate) struct VerifyStepItem { pub(crate) request_id: RequestId, pub(crate) token_ids: Vec, pub(crate) params: SamplingParams, + pub(crate) stop_policy: StopPolicy, } impl VerifyStepItem { - pub(crate) fn new(request_id: RequestId, token_ids: Vec, params: SamplingParams) -> Self { + pub(crate) fn new( + request_id: RequestId, + token_ids: Vec, + params: SamplingParams, + stop_policy: StopPolicy, + ) -> Self { Self { request_id, token_ids, params, + stop_policy, } } @@ -57,9 +64,6 @@ impl VerifyStepItem { #[derive(Clone, Copy)] pub(crate) struct VerifyPlan<'a> { pub requests: &'a [VerifyStepItem], - /// Request-local stop policies in the same order as `requests`. They remain - /// host-side and are not copied into GPU buffers. - pub stop_policies: &'a [StopPolicy], /// Engine step seed for the verify rows' sampler pass (same contract as /// decode: fresh per step; seeded rows re-mix their own request seed). pub sample_seed: u64, @@ -262,6 +266,7 @@ mod tests { RequestId::new(7), vec![10, 11, 12, 13], SamplingParams::default(), + StopPolicy::default(), ); let results = build_verify_results(&[req], &[11, 12, 99, 100]).expect("verify results"); assert_eq!(results.len(), 1); @@ -276,6 +281,7 @@ mod tests { RequestId::new(8), vec![20, 21, 22], SamplingParams::default(), + StopPolicy::default(), ); let results = build_verify_results(&[req], &[21, 22, 23]).expect("verify results"); assert_eq!(results[0].matched_draft_tokens, 2); @@ -284,8 +290,18 @@ mod tests { #[test] fn batched_multi_request_splits_columns_by_span() { - let a = VerifyStepItem::new(RequestId::new(1), vec![5, 6], SamplingParams::default()); - let b = VerifyStepItem::new(RequestId::new(2), vec![7, 8, 9], SamplingParams::default()); + let a = VerifyStepItem::new( + RequestId::new(1), + vec![5, 6], + SamplingParams::default(), + StopPolicy::default(), + ); + let b = VerifyStepItem::new( + RequestId::new(2), + vec![7, 8, 9], + SamplingParams::default(), + StopPolicy::default(), + ); // a: posterior [6, 100] -> accept draft 6, bonus 100. b: posterior [8, 77, 0] // -> accept draft 8, correction 77. let results = build_verify_results(&[a, b], &[6, 100, 8, 77, 0]).expect("verify results"); diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 73c22f724..2be6b4fd0 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -692,6 +692,7 @@ fn dflash_hedged_midspan_stop_retains_trigger() { ); let mut stopped_cases = 0usize; + let mut no_stop_cases = 0usize; for (baseline_index, stop_id) in candidate_stops { let stopped_params = SamplingParams { ignore_eos: true, @@ -704,29 +705,49 @@ fn dflash_hedged_midspan_stop_retains_trigger() { eprintln!("hedge stop request={request_id}"); let outcome = stream.expect_finished(); - if outcome.tokens.last() != Some(&stop_id) { - continue; + // Do not `continue` past a candidate: a mid-span stop must still be + // the final token, and a run that never sampled the stop must terminate + // at the length limit. Filtering either shape would hide a suffix leak. + if let Some(position) = outcome.tokens.iter().position(|&token| token == stop_id) { + assert_eq!( + position, + outcome.tokens.len() - 1, + "stop {stop_id} was followed by {} more token(s) in candidate baseline_index={baseline_index}: {:?}", + outcome.tokens.len() - 1 - position, + outcome.tokens + ); + assert!(matches!( + outcome.terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(id)), + completion_tokens, + .. + } if id == stop_id && completion_tokens == outcome.tokens.len() + )); + eprintln!( + "hedge stop candidate baseline_index={baseline_index} token={stop_id} retained_len={}", + outcome.tokens.len() + ); + stopped_cases += 1; + } else { + assert!(matches!( + outcome.terminal, + Terminal::Finished { + reason: FinishReason::Length, + stop_cause: None, + completion_tokens, + .. + } if completion_tokens == GENERATED_TOKENS && completion_tokens == outcome.tokens.len() + )); + no_stop_cases += 1; } - assert!(!outcome.tokens[..outcome.tokens.len() - 1].contains(&stop_id)); - assert!(matches!( - outcome.terminal, - Terminal::Finished { - reason: FinishReason::Stop, - stop_cause: Some(StopCause::Token(id)), - completion_tokens, - .. - } if id == stop_id && completion_tokens == outcome.tokens.len() - )); - eprintln!( - "hedge stop candidate baseline_index={baseline_index} token={stop_id} retained_len={}", - outcome.tokens.len() - ); - stopped_cases += 1; } assert!( stopped_cases > 0, "none of the candidate tokens produced a mid-span explicit stop" ); + eprintln!("hedge stop child: {stopped_cases} stopped case(s), {no_stop_cases} no-stop case(s)"); } /// P2 regression: a request that fits the target context window but lands in the diff --git a/scripts/qwen3_stop_contract_probe.py b/scripts/qwen3_stop_contract_probe.py new file mode 100644 index 000000000..841081c18 --- /dev/null +++ b/scripts/qwen3_stop_contract_probe.py @@ -0,0 +1,853 @@ +#!/usr/bin/env python3 +"""Compare the Qwen3 typed stop contract with an un-migrated Qwen3.5 path. + +The script talks to already-running OpenAI-compatible PegaInfer servers. It +does not start a server and deliberately contains no model-specific launch +flags. By default each target receives a full-vocabulary explicit stop set; +this guarantees that the first sampled token exercises the explicit-stop path. +Use --stop-token-id when a smaller, known stop set is preferred. + +Every check validates the wire shape, not merely the presence of a field: + +- `stop_reason` must be a numeric token ID from the requested stop set (or + absent for a model-EOS finish); +- under a full-vocabulary stop set the first token always matches, so the + explicit-stop runs must report exactly one completion token; +- the triggering token must carry a non-null numeric logprob; +- the streaming case must terminate with `[DONE]` and must not emit content + after the finish event; +- `/v1/models` must actually serve the requested model name. + +Run `--self-check` to exercise the checks against a local mock service that +serves the malformed-response shapes these checks exist to reject; no GPU or +server binary is required. Both targets are reported side by side, and +`--require-legacy-gap` turns the expected adapted-passes/legacy-fails outcome +into an exit-code assertion. + +Example: + python3 scripts/qwen3_stop_contract_probe.py \ + --qwen3-url http://127.0.0.1:18081 --qwen3-model qwen3-adapted \ + --qwen35-url http://127.0.0.1:18082 --qwen35-model qwen35-legacy \ + --out stop-contract-ab.json +""" + +from __future__ import annotations + +import argparse +import json +import math +import sys +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + +DEFAULT_PROMPT = "The capital of France is" +DEFAULT_QWEN3_URL = "http://127.0.0.1:18081" +DEFAULT_QWEN35_URL = "http://127.0.0.1:18082" +DEFAULT_QWEN3_MODEL = "qwen3-adapted" +DEFAULT_QWEN35_MODEL = "qwen35-legacy" +DEFAULT_QWEN3_VOCAB = 151_936 +DEFAULT_QWEN35_VOCAB = 248_320 + + +def endpoint(base_url: str, suffix: str) -> str: + return base_url.rstrip("/") + suffix + + +def http_json(url: str, payload: dict[str, Any] | None, timeout: float) -> dict[str, Any]: + data = None + headers = {"Accept": "application/json"} + method = "GET" + if payload is not None: + data = json.dumps(payload, separators=(",", ":")).encode("utf-8") + headers["Content-Type"] = "application/json" + method = "POST" + request = Request(url, data=data, headers=headers, method=method) + started = time.perf_counter() + try: + with urlopen(request, timeout=timeout) as response: + raw = response.read() + status = response.status + except HTTPError as error: + raw = error.read() + try: + body: Any = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError): + body = raw.decode("utf-8", errors="replace") + return { + "http_status": error.code, + "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), + "error": body, + } + except (TimeoutError, URLError, OSError) as error: + return { + "http_status": None, + "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), + "error": str(error), + } + try: + body = json.loads(raw.decode("utf-8")) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + return { + "http_status": status, + "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), + "error": f"invalid JSON response: {error}", + } + return { + "http_status": status, + "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), + "body": body, + } + + +def is_finite_number(value: Any) -> bool: + if isinstance(value, bool) or not isinstance(value, (int, float)): + return False + return math.isfinite(value) + + +def last_token_logprob(choice: dict[str, Any]) -> float | None: + """Logprob of the last token in one choice, or None when absent. + + Accepts both the legacy ``{"tokens": [...], "token_logprobs": [...]}`` + shape and the OpenAI content shape ``{"content": [{"logprob": ...}]}``. A + null entry is not a logprob and must fail the trigger check. + """ + logprobs = choice.get("logprobs") + if not isinstance(logprobs, dict): + return None + token_logprobs = logprobs.get("token_logprobs") + if isinstance(token_logprobs, list) and token_logprobs: + value = token_logprobs[-1] + return value if is_finite_number(value) else None + content = logprobs.get("content") + if isinstance(content, list) and content: + last = content[-1] + if isinstance(last, dict): + value = last.get("logprob") + return value if is_finite_number(value) else None + return None + + +def stream_token_count(choice: dict[str, Any]) -> int: + """Number of tokens a streaming choice emits (via logprobs or bare text).""" + logprobs = choice.get("logprobs") + if isinstance(logprobs, dict): + tokens = logprobs.get("tokens") + if isinstance(tokens, list): + return len(tokens) + content = logprobs.get("content") + if isinstance(content, list): + return len(content) + return 0 + + +def first_choice(body: dict[str, Any]) -> dict[str, Any]: + choices = body.get("choices") + if not isinstance(choices, list) or not choices: + return {} + choice = choices[0] + return choice if isinstance(choice, dict) else {} + + +def completion_summary(response: dict[str, Any]) -> dict[str, Any]: + body = response.get("body") + if not isinstance(body, dict): + return { + **response, + "finish_reason": None, + "stop_reason": None, + "completion_tokens": None, + "trigger_logprob": None, + } + choice = first_choice(body) + usage = body.get("usage") if isinstance(body.get("usage"), dict) else {} + stop_reason = choice.get("stop_reason", body.get("stop_reason")) + return { + "http_status": response.get("http_status"), + "elapsed_ms": response.get("elapsed_ms"), + "finish_reason": choice.get("finish_reason"), + "stop_reason": stop_reason, + "completion_tokens": usage.get("completion_tokens"), + "trigger_logprob": last_token_logprob(choice), + } + + +def completion_payload( + model: str, + prompt: str, + max_tokens: int, + ignore_eos: bool, + stop_token_ids: list[int] | None, + logprobs: int, + stream: bool, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "model": model, + "prompt": prompt, + "temperature": 0.0, + "max_tokens": max_tokens, + "ignore_eos": ignore_eos, + "stream": stream, + } + if stop_token_ids is not None: + payload["stop_token_ids"] = stop_token_ids + if logprobs > 0: + payload["logprobs"] = logprobs + return payload + + +def request_completion( + base_url: str, + model: str, + prompt: str, + max_tokens: int, + timeout: float, + *, + ignore_eos: bool, + stop_token_ids: list[int] | None = None, + logprobs: int = 0, +) -> dict[str, Any]: + payload = completion_payload(model, prompt, max_tokens, ignore_eos, stop_token_ids, logprobs, stream=False) + return completion_summary(http_json(endpoint(base_url, "/v1/completions"), payload, timeout)) + + +def request_stream( + base_url: str, + model: str, + prompt: str, + max_tokens: int, + timeout: float, + *, + ignore_eos: bool, + stop_token_ids: list[int] | None = None, + logprobs: int = 0, +) -> dict[str, Any]: + payload = completion_payload(model, prompt, max_tokens, ignore_eos, stop_token_ids, logprobs, stream=True) + data = json.dumps(payload, separators=(",", ":")).encode("utf-8") + request = Request( + endpoint(base_url, "/v1/completions"), + data=data, + headers={"Accept": "text/event-stream", "Content-Type": "application/json"}, + method="POST", + ) + started = time.perf_counter() + events = 0 + emitted_tokens = 0 + last_stream_logprob = None + finish_reason: Any = None + stop_reason: Any = None + done_seen = False + content_after_finish = False + malformed_chunks = 0 + finished = False + status = None + error = None + try: + with urlopen(request, timeout=timeout) as response: + status = response.status + for raw_line in response: + line = raw_line.decode("utf-8", errors="replace").strip() + if not line.startswith("data:"): + continue + data_line = line[5:].strip() + if data_line == "[DONE]": + done_seen = True + break + try: + chunk = json.loads(data_line) + except json.JSONDecodeError: + malformed_chunks += 1 + continue + if not isinstance(chunk, dict): + malformed_chunks += 1 + continue + choices = chunk.get("choices") + if not isinstance(choices, list) or not choices: + # A usage-only frame carries no content and is legal. + if isinstance(chunk.get("usage"), dict): + continue + malformed_chunks += 1 + continue + choice = choices[0] if isinstance(choices[0], dict) else {} + text = choice.get("text") + token_count = stream_token_count(choice) + has_content = bool(text) or token_count > 0 + this_finish = choice.get("finish_reason") + this_stop = chunk.get("stop_reason", choice.get("stop_reason")) + terminal_event = this_finish is not None or this_stop is not None + if finished: + if has_content or terminal_event: + content_after_finish = True + continue + if terminal_event: + if this_finish is not None: + finish_reason = this_finish + if stop_reason is None and this_stop is not None: + stop_reason = this_stop + finished = True + continue + if has_content: + events += 1 + emitted_tokens += token_count if token_count else 1 + value = last_token_logprob(choice) + if value is not None: + last_stream_logprob = value + except HTTPError as exc: + status = exc.code + error = exc.read().decode("utf-8", errors="replace") + except (TimeoutError, URLError, OSError) as exc: + error = str(exc) + return { + "http_status": status, + "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), + "stream_events": events, + "emitted_tokens": emitted_tokens, + "trigger_logprob": last_stream_logprob, + "finish_reason": finish_reason, + "stop_reason": stop_reason, + "completion_tokens": None, + "done_seen": done_seen, + "content_after_finish": content_after_finish, + "malformed_chunks": malformed_chunks, + "error": error, + } + + +def probe_models(base_url: str, timeout: float, expected_model: str) -> dict[str, Any]: + result = http_json(endpoint(base_url, "/v1/models"), None, timeout) + body = result.get("body") + ids: list[str] = [] + if isinstance(body, dict) and isinstance(body.get("data"), list): + ids = [ + item.get("id") + for item in body["data"] + if isinstance(item, dict) and item.get("id") + ] + result["model_ids"] = ids + result["expected_model_present"] = expected_model in ids + return result + + +def numeric_stop_reason(result: dict[str, Any]) -> bool: + return isinstance(result.get("stop_reason"), int) + + +def is_stop_finish(result: dict[str, Any]) -> bool: + return result.get("http_status") == 200 and result.get("finish_reason") == "stop" + + +def typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: bool) -> bool: + """A typed request stop: integer stop_reason from the requested set.""" + if not is_stop_finish(result): + return False + if not numeric_stop_reason(result) or result["stop_reason"] not in stop_ids: + return False + tokens = result.get("completion_tokens") + if not isinstance(tokens, int) or tokens < 1: + return False + if full_vocab: + # The first sampled token is in the full-vocabulary stop set, so the + # trigger must be the only completion token. + return tokens == 1 + return True + + +def eos_or_typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: bool) -> bool: + """EOS-enabled explicit-stop case: a model EOS may win over the stop set.""" + if not is_stop_finish(result): + return False + stop_reason = result.get("stop_reason") + if stop_reason is not None and ( + not isinstance(stop_reason, int) or stop_reason not in stop_ids + ): + return False + tokens = result.get("completion_tokens") + if not isinstance(tokens, int) or tokens < 1: + return False + if full_vocab: + return tokens == 1 + return True + + +def length_control(result: dict[str, Any], max_tokens: int) -> bool: + return ( + result.get("http_status") == 200 + and result.get("finish_reason") == "length" + and result.get("stop_reason") is None + and result.get("completion_tokens") == max_tokens + ) + + +def build_stop_ids(vocab_size: int, selected_id: int | None) -> list[int]: + if selected_id is not None: + if selected_id < 0 or selected_id >= vocab_size: + raise ValueError(f"--stop-token-id must be in [0, {vocab_size}), got {selected_id}") + return [selected_id] + return list(range(vocab_size)) + + +def run_target( + name: str, + url: str, + model: str, + vocab_size: int, + args: argparse.Namespace, +) -> dict[str, Any]: + stop_ids = build_stop_ids(vocab_size, args.stop_token_id) + full_vocab = args.stop_token_id is None + model_probe = probe_models(url, args.timeout, model) + + def call(ignore_eos: bool, ids: list[int] | None, lp: int = 0) -> dict[str, Any]: + return request_completion( + url, + model, + args.prompt, + args.max_tokens, + args.timeout, + ignore_eos=ignore_eos, + stop_token_ids=ids, + logprobs=lp, + ) + + cases: dict[str, Any] = {} + cases["control"] = call(True, None) + cases["explicit_stop_ignore_eos"] = call(True, stop_ids) + cases["explicit_stop_eos_enabled"] = call(False, stop_ids) + cases["stop_ascending"] = call(True, stop_ids) + cases["stop_descending"] = call(True, list(reversed(stop_ids))) + cases["trigger_logprob"] = call(True, stop_ids, lp=1) + cases["streaming"] = request_stream( + url, + model, + args.prompt, + args.max_tokens, + args.timeout, + ignore_eos=True, + stop_token_ids=stop_ids, + logprobs=1, + ) + + mixed_jobs: list[tuple[str, bool]] = [("control", False)] * 3 + [("explicit", True)] * 3 + + def mixed_call(job: tuple[str, bool]) -> dict[str, Any]: + _, explicit = job + return request_completion( + url, + model, + args.prompt, + args.max_tokens, + args.timeout, + ignore_eos=True, + stop_token_ids=stop_ids if explicit else None, + ) + + with ThreadPoolExecutor(max_workers=len(mixed_jobs)) as pool: + mixed_results = list(pool.map(mixed_call, mixed_jobs)) + cases["mixed"] = [ + {"kind": kind, "result": result} + for (kind, _), result in zip(mixed_jobs, mixed_results) + ] + + mixed_controls = [item["result"] for item in cases["mixed"] if item["kind"] == "control"] + mixed_explicit = [item["result"] for item in cases["mixed"] if item["kind"] == "explicit"] + streaming = cases["streaming"] + checks = { + "model_present": bool(model_probe.get("expected_model_present")), + "baseline_control": length_control(cases["control"], args.max_tokens), + "explicit_stop_ignore_eos": typed_stop(cases["explicit_stop_ignore_eos"], stop_ids, full_vocab), + "explicit_stop_eos_enabled": eos_or_typed_stop(cases["explicit_stop_eos_enabled"], stop_ids, full_vocab), + "stop_set_order_invariant": ( + typed_stop(cases["stop_ascending"], stop_ids, full_vocab) + and typed_stop(cases["stop_descending"], stop_ids, full_vocab) + and cases["stop_ascending"].get("stop_reason") + == cases["stop_descending"].get("stop_reason") + and cases["stop_ascending"].get("completion_tokens") + == cases["stop_descending"].get("completion_tokens") + ), + "trigger_logprob_preserved": ( + typed_stop(cases["trigger_logprob"], stop_ids, full_vocab) + and is_finite_number(cases["trigger_logprob"].get("trigger_logprob")) + ), + "stream_reports_typed_stop": ( + streaming.get("http_status") == 200 + and streaming.get("finish_reason") == "stop" + and numeric_stop_reason(streaming) + and streaming.get("stop_reason") in stop_ids + and streaming.get("stream_events", 0) > 0 + and streaming.get("done_seen") is True + and not streaming.get("content_after_finish") + and streaming.get("malformed_chunks", 0) == 0 + and is_finite_number(streaming.get("trigger_logprob")) + and ( + streaming.get("emitted_tokens") == 1 + if full_vocab + else ( + isinstance(streaming.get("emitted_tokens"), int) + and streaming.get("emitted_tokens") >= 1 + ) + ) + ), + "mixed_controls_pass_3_of_3": sum(length_control(item, args.max_tokens) for item in mixed_controls) == 3, + "mixed_explicit_stops_pass_3_of_3": sum(typed_stop(item, stop_ids, full_vocab) for item in mixed_explicit) == 3, + } + return { + "name": name, + "url": url, + "model": model, + "vocab_size": vocab_size, + "stop_set": "single" if args.stop_token_id is not None else "full-vocabulary", + "stop_set_size": len(stop_ids), + "model_probe": model_probe, + "cases": cases, + "checks": checks, + "new_contract_passed": all(checks.values()), + } + + +def print_target(target: dict[str, Any]) -> None: + print(f"\n{target['name']} ({target['model']})") + print("case finish stop_reason tokens http") + for name, result in target["cases"].items(): + if name == "mixed": + continue + if name == "streaming": + print( + f"{name:28} {str(result.get('finish_reason')):7} " + f"{str(result.get('stop_reason')):12} {'n/a':>7} " + f"{str(result.get('http_status'))} events={result.get('stream_events')} " + f"done={result.get('done_seen')} tail={result.get('content_after_finish')}" + ) + continue + print( + f"{name:28} {str(result.get('finish_reason')):7} " + f"{str(result.get('stop_reason')):12} {str(result.get('completion_tokens')):>7} " + f"{str(result.get('http_status'))}" + ) + print("checks:") + for name, passed in target["checks"].items(): + print(f" [{'PASS' if passed else 'FAIL'}] {name}") + print(f"overall new contract: {'PASS' if target['new_contract_passed'] else 'FAIL'}") + + +def print_comparison(qwen3: dict[str, Any], qwen35: dict[str, Any]) -> None: + print("\ncontract checks (adapted vs legacy):") + print(f"{'check':28} {'qwen3':7} {'qwen35':7}") + for name in qwen3["checks"]: + left = "PASS" if qwen3["checks"][name] else "FAIL" + right = "PASS" if qwen35["checks"][name] else "FAIL" + print(f"{name:28} {left:7} {right:7}") + print( + f"{'overall new contract':28} " + f"{'PASS' if qwen3['new_contract_passed'] else 'FAIL':7} " + f"{'PASS' if qwen35['new_contract_passed'] else 'FAIL':7}" + ) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--qwen3-url", default=DEFAULT_QWEN3_URL) + parser.add_argument("--qwen3-model", default=DEFAULT_QWEN3_MODEL) + parser.add_argument("--qwen3-vocab-size", type=int, default=DEFAULT_QWEN3_VOCAB) + parser.add_argument("--qwen35-url", default=DEFAULT_QWEN35_URL) + parser.add_argument("--qwen35-model", default=DEFAULT_QWEN35_MODEL) + parser.add_argument("--qwen35-vocab-size", type=int, default=DEFAULT_QWEN35_VOCAB) + parser.add_argument("--prompt", default=DEFAULT_PROMPT) + parser.add_argument("--max-tokens", type=int, default=8) + parser.add_argument("--timeout", type=float, default=300.0) + parser.add_argument( + "--stop-token-id", + type=int, + help="Use one explicit stop ID instead of the default full-vocabulary set", + ) + parser.add_argument("--out", type=Path, help="Write the complete result as JSON") + parser.add_argument( + "--strict-both", + action="store_true", + help="Return failure unless both targets satisfy every new-contract check", + ) + parser.add_argument( + "--self-check", + action="store_true", + help="Run the checks against a local mock service that serves malformed " + "responses; verify each is rejected. Requires no server or GPU.", + ) + parser.add_argument( + "--require-legacy-gap", + action="store_true", + help="Return failure unless the adapted target passes every check and the " + "legacy target fails at least one (the expected A/B outcome).", + ) + return parser.parse_args() + + +SELF_CHECK_MODEL = "qwen3-adapted" +SELF_CHECK_STOP_ID = 12095 + +_MODE = "valid" +_MODE_LOCK = threading.Lock() + + +def set_mock_mode(mode: str) -> None: + global _MODE + with _MODE_LOCK: + _MODE = mode + + +def mock_mode() -> str: + with _MODE_LOCK: + return _MODE + + +class MockHandler(BaseHTTPRequestHandler): + def log_message(self, *args: Any) -> None: + pass + + def _json(self, payload: dict[str, Any]) -> None: + raw = json.dumps(payload).encode("utf-8") + self.send_response(200) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(raw))) + self.end_headers() + self.wfile.write(raw) + + def do_GET(self) -> None: + if self.path.rstrip("/").endswith("/v1/models"): + if mock_mode() == "wrong_model": + self._json({"data": [{"id": "different-model"}]}) + else: + self._json({"data": [{"id": SELF_CHECK_MODEL}]}) + return + self.send_error(404) + + def do_POST(self) -> None: + length = int(self.headers.get("Content-Length", "0")) + try: + body = json.loads(self.rfile.read(length) or b"{}") + except json.JSONDecodeError: + self.send_error(400) + return + if not self.path.rstrip("/").endswith("/v1/completions"): + self.send_error(404) + return + mode = mock_mode() + if body.get("stream"): + self._stream(mode) + else: + self._completion(body, mode) + + def _completion(self, body: dict[str, Any], mode: str) -> None: + explicit = body.get("stop_token_ids") is not None + max_tokens = body.get("max_tokens", 8) + if not explicit: + self._json( + { + "choices": [ + {"text": " word", "index": 0, "finish_reason": "length", "logprobs": None} + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": max_tokens, + "total_tokens": 5 + max_tokens, + }, + } + ) + return + stop_reason: Any = SELF_CHECK_STOP_ID if mode != "string_stop_reason" else "oops" + tokens = max_tokens if mode == "extra_tokens" else 1 + choice: dict[str, Any] = { + "text": " stop", + "index": 0, + "finish_reason": "stop", + "logprobs": None, + } + if (body.get("logprobs") or 0) > 0: + if mode == "null_logprobs": + choice["logprobs"] = {"content": [{"token": "", "logprob": None}]} + else: + choice["logprobs"] = {"content": [{"token": "", "logprob": -0.53125}]} + choice["stop_reason"] = stop_reason + self._json( + { + "choices": [choice], + "usage": { + "prompt_tokens": 5, + "completion_tokens": tokens, + "total_tokens": 5 + tokens, + }, + } + ) + + def _stream(self, mode: str) -> None: + stop_reason: Any = SELF_CHECK_STOP_ID if mode != "string_stop_reason" else "oops" + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.send_header("Cache-Control", "no-cache") + self.end_headers() + + def emit(chunk: dict[str, Any]) -> None: + self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode()) + + emit( + { + "id": "cmpl-1", + "choices": [ + { + "text": "", + "index": 0, + "finish_reason": None, + "logprobs": { + "tokens": [" stop"], + "token_logprobs": [-0.53125], + "top_logprobs": [], + }, + } + ], + } + ) + emit( + { + "id": "cmpl-1", + "choices": [{"text": "", "index": 0, "finish_reason": "stop", "logprobs": None}], + "stop_reason": stop_reason, + } + ) + if mode == "stream_tail": + emit( + { + "id": "cmpl-1", + "choices": [ + { + "text": "", + "index": 0, + "finish_reason": None, + "logprobs": { + "tokens": [" tail"], + "token_logprobs": [-1.0], + "top_logprobs": [], + }, + } + ], + } + ) + self.wfile.write(b"data: [DONE]\n\n") + + +def run_self_check() -> int: + server = ThreadingHTTPServer(("127.0.0.1", 0), MockHandler) + port = server.server_address[1] + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + base_url = f"http://127.0.0.1:{port}" + + class Args: + prompt = DEFAULT_PROMPT + max_tokens = 8 + timeout = 10.0 + stop_token_id = None + + cases = [ + ("valid", None, True, "valid service passes every check"), + ("string_stop_reason", SELF_CHECK_STOP_ID, False, "string stop_reason must be rejected"), + ("null_logprobs", SELF_CHECK_STOP_ID, False, "null trigger logprob must be rejected"), + ("stream_tail", SELF_CHECK_STOP_ID, False, "content after the finish event must be rejected"), + ("wrong_model", SELF_CHECK_STOP_ID, False, "model-name mismatch must be rejected"), + ( + "extra_tokens", + None, + False, + "extra completion tokens under a full-vocabulary stop set must be rejected", + ), + ] + failures = 0 + for mode, stop_id, expect_pass, label in cases: + set_mock_mode(mode) + args = Args() + args.stop_token_id = stop_id + target = run_target("qwen3_adapted", base_url, SELF_CHECK_MODEL, DEFAULT_QWEN3_VOCAB, args) + passed = target["new_contract_passed"] + ok = passed == expect_pass + if not ok: + failures += 1 + for check_name, check_value in target["checks"].items(): + if not check_value: + print(f" unexpected check state: {check_name}=FAIL") + print( + f"[{'PASS' if ok else 'FAIL'}] self-check {mode}: " + f"overall={'PASS' if passed else 'FAIL'}, expected={'PASS' if expect_pass else 'FAIL'} ({label})" + ) + server.shutdown() + thread.join(timeout=5.0) + return 1 if failures else 0 + + +def main() -> int: + args = parse_args() + if args.self_check: + return run_self_check() + if args.max_tokens <= 0: + print("--max-tokens must be positive", file=sys.stderr) + return 2 + try: + qwen3 = run_target( + "qwen3_adapted", + args.qwen3_url, + args.qwen3_model, + args.qwen3_vocab_size, + args, + ) + qwen35 = run_target( + "qwen35_legacy", + args.qwen35_url, + args.qwen35_model, + args.qwen35_vocab_size, + args, + ) + except ValueError as error: + print(str(error), file=sys.stderr) + return 2 + + comparison = { + "qwen3_new_contract_passed": qwen3["new_contract_passed"], + "qwen35_new_contract_passed": qwen35["new_contract_passed"], + "legacy_gap_observed": ( + qwen3["new_contract_passed"] and not qwen35["new_contract_passed"] + ), + } + report = { + "schema_version": 2, + "config": { + "prompt": args.prompt, + "max_tokens": args.max_tokens, + "stop_mode": "single" if args.stop_token_id is not None else "full-vocabulary", + "stop_token_id": args.stop_token_id, + }, + "targets": {"qwen3_adapted": qwen3, "qwen35_legacy": qwen35}, + "comparison": comparison, + } + print_target(qwen3) + print_target(qwen35) + print_comparison(qwen3, qwen35) + print("\ncomparison:") + print(json.dumps(comparison, indent=2)) + if args.out: + args.out.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") + print(f"wrote {args.out}") + + if not qwen3["new_contract_passed"]: + return 1 + if args.strict_both and not qwen35["new_contract_passed"]: + return 1 + if args.require_legacy_gap and not comparison["legacy_gap_observed"]: + print("expected adapted-passes/legacy-fails gap was not observed", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 9b91ad318fed2bc01c7245d157e92168c8e8239e Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Sun, 13 Sep 2026 21:32:59 +0800 Subject: [PATCH 13/15] fix(probe): reject boolean wire values in numeric stop checks Signed-off-by: RicardoMin <17879681016@163.com> --- scripts/qwen3_stop_contract_probe.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/scripts/qwen3_stop_contract_probe.py b/scripts/qwen3_stop_contract_probe.py index 841081c18..a56301272 100644 --- a/scripts/qwen3_stop_contract_probe.py +++ b/scripts/qwen3_stop_contract_probe.py @@ -334,8 +334,14 @@ def probe_models(base_url: str, timeout: float, expected_model: str) -> dict[str return result +def is_int(value: Any) -> bool: + # JSON `true` parses to Python True, which is also an `int` subclass; + # exclude it so a boolean wire value can never satisfy a numeric check. + return isinstance(value, int) and not isinstance(value, bool) + + def numeric_stop_reason(result: dict[str, Any]) -> bool: - return isinstance(result.get("stop_reason"), int) + return is_int(result.get("stop_reason")) def is_stop_finish(result: dict[str, Any]) -> bool: @@ -349,7 +355,7 @@ def typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: bool) -> if not numeric_stop_reason(result) or result["stop_reason"] not in stop_ids: return False tokens = result.get("completion_tokens") - if not isinstance(tokens, int) or tokens < 1: + if not is_int(tokens) or tokens < 1: return False if full_vocab: # The first sampled token is in the full-vocabulary stop set, so the @@ -363,12 +369,10 @@ def eos_or_typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: b if not is_stop_finish(result): return False stop_reason = result.get("stop_reason") - if stop_reason is not None and ( - not isinstance(stop_reason, int) or stop_reason not in stop_ids - ): + if stop_reason is not None and (not is_int(stop_reason) or stop_reason not in stop_ids): return False tokens = result.get("completion_tokens") - if not isinstance(tokens, int) or tokens < 1: + if not is_int(tokens) or tokens < 1: return False if full_vocab: return tokens == 1 @@ -380,6 +384,7 @@ def length_control(result: dict[str, Any], max_tokens: int) -> bool: result.get("http_status") == 200 and result.get("finish_reason") == "length" and result.get("stop_reason") is None + and is_int(result.get("completion_tokens")) and result.get("completion_tokens") == max_tokens ) From d57ca1135f815a5f350935f7eed6c4e284feb42d Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Mon, 14 Sep 2026 01:51:01 +0800 Subject: [PATCH 14/15] fix(qwen3): verify exact stop triggers and assert worker normalization Signed-off-by: RicardoMin <17879681016@163.com> --- pegainfer-qwen3/src/executor/spec.rs | 39 ++- pegainfer-qwen3/src/speculative.rs | 8 +- .../tests/dflash_speculative_gate.rs | 3 +- scripts/qwen3_stop_contract_probe.py | 325 ++++++++++++++---- 4 files changed, 294 insertions(+), 81 deletions(-) diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 967dc6772..5d309f280 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -27,7 +27,7 @@ pub(super) fn truncate_after_terminal( model_eos: &[u32], ) { // Keep this helper idempotent: the worker trims before DFlash context is - // recorded, and the executor repeats the invariant before KV commit. + // recorded, and the executor later asserts the same invariant at commit. let Some(keep) = result.accepted_tokens.iter().position(|&token| { policy .classify(token, |id| model_eos.contains(&id)) @@ -110,7 +110,7 @@ impl Qwen3Executor { return Err(e); } }; - let mut result = match outcome { + let result = match outcome { WorkerStepOutcome::SpeculativeVerify(result) => result, other => { self.revert_speculative_schedules(&scheduled); @@ -138,12 +138,35 @@ impl Qwen3Executor { )); } } - // The worker normally applies the request contract before copying - // worker-side state. Recheck the same invariant here before touching - // RequestKv so a terminal suffix is rolled back with its reservation, - // including legacy workers that return an untrimmed span. - for (req, req_result) in plan.requests.iter().zip(&mut result.requests) { - truncate_after_terminal(req_result, &req.stop_policy, &self.metadata.stop_token_ids); + // The worker must have normalized each span before recording DFlash + // context (both the hedged and plain verify paths truncate at the first + // terminal token). Re-assert that invariant at commit: a broken worker + // fails loudly here instead of being silently truncated a second time. + for (req, req_result) in plan.requests.iter().zip(&result.requests) { + let terminal = |token: u32| { + req.stop_policy + .classify(token, |id| self.metadata.stop_token_ids.contains(&id)) + .is_some() + }; + let terminal_position = req_result + .accepted_tokens + .iter() + .position(|&token| terminal(token)); + if let Some(position) = terminal_position + && position + 1 != req_result.accepted_tokens.len() + { + // Nothing has committed yet: roll back every reservation before + // surfacing the worker defect, mirroring the request-mismatch + // failures above. + self.revert_speculative_schedules(&scheduled); + return Err(anyhow::anyhow!( + "speculative worker returned an untruncated span for {:?}: \ + terminal token at position {} of {}", + req_result.request_id, + position, + req_result.accepted_tokens.len() + )); + } } // Commit the accepted prefix of each request's KV and free the rest. diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index 1606b0ea4..dc1215223 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -77,9 +77,11 @@ pub(crate) struct VerifyRequestResult { /// Tokens to commit: the accepted draft prefix followed by the target's /// posterior token at the first mismatch (or the block-end continuation /// when every draft is accepted). Always `1..=K + 1` tokens, so a verify - /// step always makes at least one token of progress. Before KV commit the - /// executor truncates this span after the first request-terminal token; - /// the scheduler retains ownership of typed stop-cause emission. + /// step always makes at least one token of progress. The worker normalizes + /// this span after the first request-terminal token before DFlash context is + /// recorded; the executor asserts that invariant at commit rather than + /// re-truncating. The scheduler retains ownership of typed stop-cause + /// emission. pub accepted_tokens: Vec, } diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 2be6b4fd0..a117d2b2d 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -629,7 +629,8 @@ fn dflash_concurrent_heterogeneous_is_lossless() { /// Production hedge regression: an explicit stop in the middle of a verify /// span must be applied before the hedge winner is selected and committed. /// The parent gate also checks the request-local worker trace, because the -/// final stream alone is protected by the executor's legacy safety truncation. +/// final stream alone cannot prove the winner was selected from a truncated +/// span (a broken worker is only rejected at commit, after the fact). #[test] fn dflash_hedged_midspan_stop_retains_trigger() { common::harness::init_capture_logging(); diff --git a/scripts/qwen3_stop_contract_probe.py b/scripts/qwen3_stop_contract_probe.py index a56301272..e2bc0edb2 100644 --- a/scripts/qwen3_stop_contract_probe.py +++ b/scripts/qwen3_stop_contract_probe.py @@ -14,6 +14,10 @@ - under a full-vocabulary stop set the first token always matches, so the explicit-stop runs must report exactly one completion token; - the triggering token must carry a non-null numeric logprob; +- the exact trigger is verified by requesting `return_token_ids` and comparing + the final returned token ID against `stop_reason`; +- when `--qwen3-eos-token-id` / `--qwen35-eos-token-id` is given, an EOS + finish must end on exactly that trigger ID; - the streaming case must terminate with `[DONE]` and must not emit content after the finish event; - `/v1/models` must actually serve the requested model name. @@ -111,22 +115,33 @@ def is_finite_number(value: Any) -> bool: return math.isfinite(value) -def last_token_logprob(choice: dict[str, Any]) -> float | None: - """Logprob of the last token in one choice, or None when absent. +def token_ids_from_choice(choice: dict[str, Any]) -> list[int] | None: + ids = choice.get("token_ids") + if isinstance(ids, list) and all(is_int(item) for item in ids): + return ids + return None + - Accepts both the legacy ``{"tokens": [...], "token_logprobs": [...]}`` - shape and the OpenAI content shape ``{"content": [{"logprob": ...}]}``. A - null entry is not a logprob and must fail the trigger check. +def logprob_at_last(choice: dict[str, Any], token_count: int) -> float | None: + """Finite logprob of the final emitted token, or None when absent. + + The logprob table must cover exactly the emitted tokens: a table shorter + than the token sequence (a missing trigger entry) or a null final entry is + a failure rather than an inherited value from an earlier token. """ logprobs = choice.get("logprobs") - if not isinstance(logprobs, dict): + if not isinstance(logprobs, dict) or token_count == 0: return None token_logprobs = logprobs.get("token_logprobs") - if isinstance(token_logprobs, list) and token_logprobs: + if isinstance(token_logprobs, list): + if len(token_logprobs) != token_count: + return None value = token_logprobs[-1] return value if is_finite_number(value) else None content = logprobs.get("content") - if isinstance(content, list) and content: + if isinstance(content, list): + if len(content) != token_count: + return None last = content[-1] if isinstance(last, dict): value = last.get("logprob") @@ -135,7 +150,10 @@ def last_token_logprob(choice: dict[str, Any]) -> float | None: def stream_token_count(choice: dict[str, Any]) -> int: - """Number of tokens a streaming choice emits (via logprobs or bare text).""" + """Number of tokens a streaming choice emits (token_ids, logprobs, or text).""" + ids = token_ids_from_choice(choice) + if ids is not None: + return len(ids) logprobs = choice.get("logprobs") if isinstance(logprobs, dict): tokens = logprobs.get("tokens") @@ -163,18 +181,21 @@ def completion_summary(response: dict[str, Any]) -> dict[str, Any]: "finish_reason": None, "stop_reason": None, "completion_tokens": None, + "token_ids": None, "trigger_logprob": None, } choice = first_choice(body) usage = body.get("usage") if isinstance(body.get("usage"), dict) else {} stop_reason = choice.get("stop_reason", body.get("stop_reason")) + token_ids = token_ids_from_choice(choice) return { "http_status": response.get("http_status"), "elapsed_ms": response.get("elapsed_ms"), "finish_reason": choice.get("finish_reason"), "stop_reason": stop_reason, "completion_tokens": usage.get("completion_tokens"), - "trigger_logprob": last_token_logprob(choice), + "token_ids": token_ids, + "trigger_logprob": logprob_at_last(choice, len(token_ids) if token_ids is not None else 0), } @@ -194,6 +215,7 @@ def completion_payload( "max_tokens": max_tokens, "ignore_eos": ignore_eos, "stream": stream, + "return_token_ids": True, } if stop_token_ids is not None: payload["stop_token_ids"] = stop_token_ids @@ -239,7 +261,9 @@ def request_stream( started = time.perf_counter() events = 0 emitted_tokens = 0 - last_stream_logprob = None + stream_token_ids: list[int] = [] + last_content_logprob: float | None = None + trigger_logprob_valid = False finish_reason: Any = None stop_reason: Any = None done_seen = False @@ -276,7 +300,8 @@ def request_stream( continue choice = choices[0] if isinstance(choices[0], dict) else {} text = choice.get("text") - token_count = stream_token_count(choice) + frame_ids = token_ids_from_choice(choice) + token_count = len(frame_ids) if frame_ids is not None else stream_token_count(choice) has_content = bool(text) or token_count > 0 this_finish = choice.get("finish_reason") this_stop = chunk.get("stop_reason", choice.get("stop_reason")) @@ -285,19 +310,23 @@ def request_stream( if has_content or terminal_event: content_after_finish = True continue + if has_content: + # Count content even when it shares a frame with the finish + # metadata; a terminal frame must not hide emitted tokens. + events += 1 + emitted_tokens += token_count if token_count else 1 + if frame_ids is not None: + stream_token_ids.extend(frame_ids) + value = logprob_at_last(choice, token_count) + trigger_logprob_valid = value is not None + if value is not None: + last_content_logprob = value if terminal_event: if this_finish is not None: finish_reason = this_finish if stop_reason is None and this_stop is not None: stop_reason = this_stop finished = True - continue - if has_content: - events += 1 - emitted_tokens += token_count if token_count else 1 - value = last_token_logprob(choice) - if value is not None: - last_stream_logprob = value except HTTPError as exc: status = exc.code error = exc.read().decode("utf-8", errors="replace") @@ -308,7 +337,8 @@ def request_stream( "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), "stream_events": events, "emitted_tokens": emitted_tokens, - "trigger_logprob": last_stream_logprob, + "token_ids": stream_token_ids, + "trigger_logprob": last_content_logprob if trigger_logprob_valid else None, "finish_reason": finish_reason, "stop_reason": stop_reason, "completion_tokens": None, @@ -349,41 +379,64 @@ def is_stop_finish(result: dict[str, Any]) -> bool: def typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: bool) -> bool: - """A typed request stop: integer stop_reason from the requested set.""" + """A typed request stop: the reported stop_reason is the emitted trigger.""" if not is_stop_finish(result): return False if not numeric_stop_reason(result) or result["stop_reason"] not in stop_ids: return False + ids = result.get("token_ids") + if not isinstance(ids, list) or not ids: + return False tokens = result.get("completion_tokens") - if not is_int(tokens) or tokens < 1: + if not is_int(tokens) or tokens != len(ids) or tokens < 1: + return False + if ids[-1] != result["stop_reason"]: return False if full_vocab: # The first sampled token is in the full-vocabulary stop set, so the # trigger must be the only completion token. - return tokens == 1 + return len(ids) == 1 return True -def eos_or_typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: bool) -> bool: +def eos_or_typed_stop( + result: dict[str, Any], + stop_ids: list[int], + full_vocab: bool, + eos_id: int | None, +) -> bool: """EOS-enabled explicit-stop case: a model EOS may win over the stop set.""" if not is_stop_finish(result): return False - stop_reason = result.get("stop_reason") - if stop_reason is not None and (not is_int(stop_reason) or stop_reason not in stop_ids): + ids = result.get("token_ids") + if not isinstance(ids, list) or not ids: return False tokens = result.get("completion_tokens") - if not is_int(tokens) or tokens < 1: + if not is_int(tokens) or tokens != len(ids) or tokens < 1: return False - if full_vocab: - return tokens == 1 - return True + stop_reason = result.get("stop_reason") + if stop_reason is None: + # A model EOS won. With a configured EOS ID the final token must be + # exactly that trigger; otherwise only the count and terminal shape + # are verifiable. + if eos_id is not None and ids[-1] != eos_id: + return False + return not full_vocab or len(ids) == 1 + if not is_int(stop_reason) or stop_reason not in stop_ids: + return False + if ids[-1] != stop_reason: + return False + return not full_vocab or len(ids) == 1 def length_control(result: dict[str, Any], max_tokens: int) -> bool: + ids = result.get("token_ids") return ( result.get("http_status") == 200 and result.get("finish_reason") == "length" and result.get("stop_reason") is None + and isinstance(ids, list) + and len(ids) == max_tokens and is_int(result.get("completion_tokens")) and result.get("completion_tokens") == max_tokens ) @@ -403,6 +456,7 @@ def run_target( model: str, vocab_size: int, args: argparse.Namespace, + eos_token_id: int | None = None, ) -> dict[str, Any]: stop_ids = build_stop_ids(vocab_size, args.stop_token_id) full_vocab = args.stop_token_id is None @@ -466,7 +520,9 @@ def mixed_call(job: tuple[str, bool]) -> dict[str, Any]: "model_present": bool(model_probe.get("expected_model_present")), "baseline_control": length_control(cases["control"], args.max_tokens), "explicit_stop_ignore_eos": typed_stop(cases["explicit_stop_ignore_eos"], stop_ids, full_vocab), - "explicit_stop_eos_enabled": eos_or_typed_stop(cases["explicit_stop_eos_enabled"], stop_ids, full_vocab), + "explicit_stop_eos_enabled": eos_or_typed_stop( + cases["explicit_stop_eos_enabled"], stop_ids, full_vocab, eos_token_id + ), "stop_set_order_invariant": ( typed_stop(cases["stop_ascending"], stop_ids, full_vocab) and typed_stop(cases["stop_descending"], stop_ids, full_vocab) @@ -489,13 +545,14 @@ def mixed_call(job: tuple[str, bool]) -> dict[str, Any]: and not streaming.get("content_after_finish") and streaming.get("malformed_chunks", 0) == 0 and is_finite_number(streaming.get("trigger_logprob")) + and isinstance(streaming.get("token_ids"), list) + and len(streaming["token_ids"]) > 0 + and streaming["token_ids"][-1] == streaming.get("stop_reason") + and streaming.get("emitted_tokens") == len(streaming["token_ids"]) and ( streaming.get("emitted_tokens") == 1 if full_vocab - else ( - isinstance(streaming.get("emitted_tokens"), int) - and streaming.get("emitted_tokens") >= 1 - ) + else (is_int(streaming.get("emitted_tokens")) and streaming["emitted_tokens"] >= 1) ) ), "mixed_controls_pass_3_of_3": sum(length_control(item, args.max_tokens) for item in mixed_controls) == 3, @@ -562,6 +619,18 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--qwen35-url", default=DEFAULT_QWEN35_URL) parser.add_argument("--qwen35-model", default=DEFAULT_QWEN35_MODEL) parser.add_argument("--qwen35-vocab-size", type=int, default=DEFAULT_QWEN35_VOCAB) + parser.add_argument( + "--qwen3-eos-token-id", + type=int, + default=None, + help="Verify an EOS finish ends on exactly this token ID (optional)", + ) + parser.add_argument( + "--qwen35-eos-token-id", + type=int, + default=None, + help="Verify an EOS finish ends on exactly this token ID (optional)", + ) parser.add_argument("--prompt", default=DEFAULT_PROMPT) parser.add_argument("--max-tokens", type=int, default=8) parser.add_argument("--timeout", type=float, default=300.0) @@ -593,6 +662,7 @@ def parse_args() -> argparse.Namespace: SELF_CHECK_MODEL = "qwen3-adapted" SELF_CHECK_STOP_ID = 12095 +SELF_CHECK_EOS_ID = 151645 _MODE = "valid" _MODE_LOCK = threading.Lock() @@ -650,11 +720,17 @@ def _completion(self, body: dict[str, Any], mode: str) -> None: explicit = body.get("stop_token_ids") is not None max_tokens = body.get("max_tokens", 8) if not explicit: + control_choice: dict[str, Any] = { + "text": " word", + "index": 0, + "finish_reason": "length", + "logprobs": None, + } + if mode != "missing_token_ids": + control_choice["token_ids"] = list(range(max_tokens)) self._json( { - "choices": [ - {"text": " word", "index": 0, "finish_reason": "length", "logprobs": None} - ], + "choices": [control_choice], "usage": { "prompt_tokens": 5, "completion_tokens": max_tokens, @@ -663,17 +739,50 @@ def _completion(self, body: dict[str, Any], mode: str) -> None: } ) return + if mode == "eos_win" and not body.get("ignore_eos"): + self._json( + { + "choices": [ + { + "text": "", + "index": 0, + "finish_reason": "stop", + "logprobs": None, + "stop_reason": None, + "token_ids": [SELF_CHECK_EOS_ID], + } + ], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + }, + } + ) + return stop_reason: Any = SELF_CHECK_STOP_ID if mode != "string_stop_reason" else "oops" tokens = max_tokens if mode == "extra_tokens" else 1 + token_ids = list(range(tokens)) if mode == "extra_tokens" else [SELF_CHECK_STOP_ID] + if mode == "missing_trigger_logprob": + tokens = 2 + token_ids = [999, SELF_CHECK_STOP_ID] + if mode == "wrong_trigger_id": + stop_reason = 17 + token_ids = [SELF_CHECK_STOP_ID] choice: dict[str, Any] = { "text": " stop", "index": 0, "finish_reason": "stop", "logprobs": None, } + if mode != "missing_token_ids": + choice["token_ids"] = token_ids if (body.get("logprobs") or 0) > 0: if mode == "null_logprobs": choice["logprobs"] = {"content": [{"token": "", "logprob": None}]} + elif mode == "missing_trigger_logprob": + # One logprob entry short of the two emitted tokens. + choice["logprobs"] = {"tokens": [""], "token_logprobs": [-0.5]} else: choice["logprobs"] = {"content": [{"token": "", "logprob": -0.53125}]} choice["stop_reason"] = stop_reason @@ -698,30 +807,45 @@ def _stream(self, mode: str) -> None: def emit(chunk: dict[str, Any]) -> None: self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode()) - emit( - { - "id": "cmpl-1", - "choices": [ - { - "text": "", - "index": 0, - "finish_reason": None, - "logprobs": { - "tokens": [" stop"], - "token_logprobs": [-0.53125], - "top_logprobs": [], - }, - } - ], - } - ) - emit( - { - "id": "cmpl-1", - "choices": [{"text": "", "index": 0, "finish_reason": "stop", "logprobs": None}], - "stop_reason": stop_reason, + first_choice: dict[str, Any] = { + "text": "", + "index": 0, + "finish_reason": None, + "token_ids": [SELF_CHECK_STOP_ID], + } + if mode != "stream_missing_trigger_logprob": + first_choice["logprobs"] = { + "tokens": [" stop"], + "token_logprobs": [-0.53125], + "top_logprobs": [], } - ) + emit({"id": "cmpl-1", "choices": [first_choice]}) + if mode in ("terminal_frame_extra_token", "terminal_frame_extra_token_full_vocab"): + # The finish frame smuggles a second token next to the metadata. + smuggled = [17] if mode == "terminal_frame_extra_token" else [SELF_CHECK_STOP_ID] + emit( + { + "id": "cmpl-1", + "choices": [ + { + "text": "", + "index": 0, + "finish_reason": "stop", + "logprobs": None, + "token_ids": smuggled, + } + ], + "stop_reason": stop_reason, + } + ) + else: + emit( + { + "id": "cmpl-1", + "choices": [{"text": "", "index": 0, "finish_reason": "stop", "logprobs": None}], + "stop_reason": stop_reason, + } + ) if mode == "stream_tail": emit( { @@ -736,11 +860,13 @@ def emit(chunk: dict[str, Any]) -> None: "token_logprobs": [-1.0], "top_logprobs": [], }, + "token_ids": [17], } ], } ) - self.wfile.write(b"data: [DONE]\n\n") + if mode != "stream_missing_done": + self.wfile.write(b"data: [DONE]\n\n") def run_self_check() -> int: @@ -757,24 +883,81 @@ class Args: stop_token_id = None cases = [ - ("valid", None, True, "valid service passes every check"), - ("string_stop_reason", SELF_CHECK_STOP_ID, False, "string stop_reason must be rejected"), - ("null_logprobs", SELF_CHECK_STOP_ID, False, "null trigger logprob must be rejected"), - ("stream_tail", SELF_CHECK_STOP_ID, False, "content after the finish event must be rejected"), - ("wrong_model", SELF_CHECK_STOP_ID, False, "model-name mismatch must be rejected"), + ("valid", None, None, True, "valid service passes every check"), + ( + "eos_win", + SELF_CHECK_STOP_ID, + SELF_CHECK_EOS_ID, + True, + "an EOS finish must end on the configured EOS trigger", + ), + ("string_stop_reason", SELF_CHECK_STOP_ID, None, False, "string stop_reason must be rejected"), + ("null_logprobs", SELF_CHECK_STOP_ID, None, False, "null trigger logprob must be rejected"), + ("stream_tail", SELF_CHECK_STOP_ID, None, False, "content after the finish event must be rejected"), + ("wrong_model", SELF_CHECK_STOP_ID, None, False, "model-name mismatch must be rejected"), ( "extra_tokens", None, + None, False, "extra completion tokens under a full-vocabulary stop set must be rejected", ), + ( + "missing_trigger_logprob", + SELF_CHECK_STOP_ID, + None, + False, + "a missing final logprob must not inherit the previous token's value", + ), + ( + "wrong_trigger_id", + None, + None, + False, + "stop_reason must match the actual final token ID", + ), + ( + "terminal_frame_extra_token", + SELF_CHECK_STOP_ID, + None, + False, + "tokens sharing a frame with finish metadata must be counted", + ), + ( + "terminal_frame_extra_token_full_vocab", + None, + None, + False, + "a duplicated trigger in the finish frame must still be counted", + ), + ( + "missing_token_ids", + None, + None, + False, + "responses without token IDs cannot satisfy the trigger checks", + ), + ( + "stream_missing_done", + SELF_CHECK_STOP_ID, + None, + False, + "a stream that never sends [DONE] must be rejected", + ), + ( + "stream_missing_trigger_logprob", + SELF_CHECK_STOP_ID, + None, + False, + "the trigger frame must carry its own finite logprob", + ), ] failures = 0 - for mode, stop_id, expect_pass, label in cases: + for mode, stop_id, eos_id, expect_pass, label in cases: set_mock_mode(mode) args = Args() args.stop_token_id = stop_id - target = run_target("qwen3_adapted", base_url, SELF_CHECK_MODEL, DEFAULT_QWEN3_VOCAB, args) + target = run_target("qwen3_adapted", base_url, SELF_CHECK_MODEL, DEFAULT_QWEN3_VOCAB, args, eos_id) passed = target["new_contract_passed"] ok = passed == expect_pass if not ok: @@ -805,6 +988,7 @@ def main() -> int: args.qwen3_model, args.qwen3_vocab_size, args, + args.qwen3_eos_token_id, ) qwen35 = run_target( "qwen35_legacy", @@ -812,6 +996,7 @@ def main() -> int: args.qwen35_model, args.qwen35_vocab_size, args, + args.qwen35_eos_token_id, ) except ValueError as error: print(str(error), file=sys.stderr) @@ -831,6 +1016,8 @@ def main() -> int: "max_tokens": args.max_tokens, "stop_mode": "single" if args.stop_token_id is not None else "full-vocabulary", "stop_token_id": args.stop_token_id, + "qwen3_eos_token_id": args.qwen3_eos_token_id, + "qwen35_eos_token_id": args.qwen35_eos_token_id, }, "targets": {"qwen3_adapted": qwen3, "qwen35_legacy": qwen35}, "comparison": comparison, From 3e1b392ba8d17868411c293f3ec0832034e3d3f5 Mon Sep 17 00:00:00 2001 From: RicardoMin <17879681016@163.com> Date: Mon, 14 Sep 2026 15:28:37 +0800 Subject: [PATCH 15/15] fix(qwen3): validate complete stop sequences and simplify verification Signed-off-by: RicardoMin <17879681016@163.com> --- docs/models/qwen3/model-crate.md | 2 +- .../frontend/frontend-architecture.md | 4 +- pegainfer-frontend/src/engine/step.rs | 3 +- pegainfer-frontend/src/engine/stop.rs | 5 +- pegainfer-qwen3/src/executor.rs | 47 +- pegainfer-qwen3/src/executor/spec.rs | 21 +- pegainfer-qwen3/src/scheduler/resolve.rs | 6 +- pegainfer-qwen3/src/speculative.rs | 15 +- .../tests/dflash_speculative_gate.rs | 13 +- scripts/qwen3_stop_contract_probe.py | 1740 ++++++++--------- 10 files changed, 883 insertions(+), 973 deletions(-) diff --git a/docs/models/qwen3/model-crate.md b/docs/models/qwen3/model-crate.md index 840b07f08..cd98d4f30 100644 --- a/docs/models/qwen3/model-crate.md +++ b/docs/models/qwen3/model-crate.md @@ -157,7 +157,7 @@ pub fn start_engine( ### Step 7: Retire ModelForward and Fix Length Limit - Deleted `pegainfer_core::model::{ModelForward, GenerationState}` and removed the root `src/model.rs` re-export. - Deleted the Qwen3 `forward.rs` compatibility path. Qwen3 tests that used it now build their baselines from `batch_prefill(bs=1)` plus `batch_decode(bs=1)`, so they exercise the same phase APIs as production. -- Fixed Qwen3 decode length-limit handling by adding `DecodeEffect::EmitAndFinish`. Length limit now emits the sampled final token, then sends `Finished { finish_reason: Length }`. Under the typed stop contract the triggering token — model EOS or an explicit stop ID — is retained and emitted as the final token; model EOS carries no wire `stop_reason`, while an explicit stop reports its ID. +- Fixed Qwen3 decode length-limit handling by adding `DecodeEffect::EmitAndFinish`. Length limit now emits the sampled final token, then sends `Finished { finish_reason: Length }`. Under the typed stop contract the triggering token — model EOS or an explicit stop ID — is retained and emitted as the final token; the vLLM primary EOS carries no wire `stop_reason`, while explicit stops and secondary EOS IDs report the matching ID. - Regenerated `test_data/Qwen3-4B.json` because every length-limited golden output now includes the final requested token. - Re-ran `bench_serving snapshot` on the CUDA validation host and pulled back `bench_snapshots/rtx-5090/qwen3-4b.json`; `decode_heavy (1024,256)` now records `generated_tokens min=max=avg=256`. - Performance stayed within noise on RTX 5090: diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index 90fc57b92..acef4e8b8 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -1,6 +1,6 @@ # Frontend architecture: pegainfer-frontend and the engine boundary -**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). The vLLM Rust dependencies are pinned to `295ac4e5`, including the DeepSeek V4/V4.1 tool-argument encoding fix. **Next step: migrate glm52, then delete the legacy contract.** +**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). The vLLM Rust dependency revision is recorded in `Cargo.toml` and `Cargo.lock`. **Next step: migrate glm52, then delete the legacy contract.** Last touched: 2026-09 @@ -33,7 +33,7 @@ Design decisions worth knowing before touching it: - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. - **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. -- **Typed stop contract.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in the emitted output and carries `Terminal::Finished.stop_cause`: `Eos(id)` has no wire `stop_reason`, while `Token(id)` reports the actual matching ID. `ignore_eos` disables only model EOS. Legacy producers may leave the cause empty while they are migrated individually. +- **Typed stop contract.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in the emitted output and carries `Terminal::Finished.stop_cause`: The vLLM primary EOS maps to `Eos(id)` with no wire `stop_reason`; secondary EOS IDs join the explicit stop set and map to `Token(id)`, reporting the actual matching ID. `ignore_eos` disables only model EOS. Legacy producers may leave the cause empty while they are migrated individually. - **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. - **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. - **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. diff --git a/pegainfer-frontend/src/engine/step.rs b/pegainfer-frontend/src/engine/step.rs index 923bda288..e6c68ab53 100644 --- a/pegainfer-frontend/src/engine/step.rs +++ b/pegainfer-frontend/src/engine/step.rs @@ -237,8 +237,7 @@ pub enum Terminal { Finished { reason: FinishReason, /// Present for token-driven stop finishes. The triggering token remains - /// in `RequestUpdate.tokens`, with its real logprob in the matching - /// `RequestUpdate.logprobs` entry. + /// in `RequestUpdate.tokens`, with its matching logprob when requested. stop_cause: Option, prompt_tokens: usize, completion_tokens: usize, diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs index a80f23408..f7e34afca 100644 --- a/pegainfer-frontend/src/engine/stop.rs +++ b/pegainfer-frontend/src/engine/stop.rs @@ -15,9 +15,8 @@ pub enum EosPolicy { /// Request-scoped token stopping policy. /// -/// EOS is kept separate from caller stop tokens because the vLLM protocol -/// reports them differently: EOS has no 'stop_reason', while a request stop -/// reports the actual matching token ID. +/// The vLLM primary EOS has no `stop_reason`; other matches, including +/// secondary EOS IDs lowered into the explicit stop set, report the token ID. #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct StopPolicy { eos: EosPolicy, diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 290380276..195a39686 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -3545,48 +3545,31 @@ impl LocalQwen3Lane { .collect(); let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, sample_seed)?; let mut all_results = build_verify_results(&expanded, &target_tokens)?; - // A terminal token ends the request even when it appears in the middle - // of a speculative span. Normalize every candidate before selecting a - // winner; otherwise a discarded suffix can win the hedge and advance - // the wrong KV/hidden state and acceptance statistics. - let (results_a, results_b) = all_results.split_at_mut(requests.len()); - anyhow::ensure!( - results_b.len() == hedge_spans.len(), - "hedge returned {} B results for {} hedge spans", - results_b.len(), - hedge_spans.len() - ); let trace = std::env::var_os("PEGAINFER_TEST_LOG").is_some(); - let raw_a_lengths: Vec = if trace { - results_a - .iter() - .map(|result| result.accepted_tokens.len()) - .collect() - } else { - Vec::new() - }; - let raw_b_lengths: Vec = if trace { - results_b + let raw_lengths: Vec = if trace { + all_results .iter() .map(|result| result.accepted_tokens.len()) .collect() } else { Vec::new() }; - for (result, req) in results_a.iter_mut().zip(requests) { + // Normalize every candidate before selecting a winner; a discarded + // suffix must not win the hedge or advance KV, hidden state or counters. + for (result, req) in all_results.iter_mut().zip(&expanded) { spec::truncate_after_terminal( result, &req.stop_policy, &self.model.config().stop_token_ids, ); } - for (slot, (idx, _)) in hedge_spans.iter().enumerate() { - spec::truncate_after_terminal( - &mut results_b[slot], - &requests[*idx].stop_policy, - &self.model.config().stop_token_ids, - ); - } + let (results_a, results_b) = all_results.split_at(requests.len()); + anyhow::ensure!( + results_b.len() == hedge_spans.len(), + "hedge returned {} B results for {} hedge spans", + results_b.len(), + hedge_spans.len() + ); // Per request keep the best-accepting chain; ties keep chain A (no // copies). A later chain of the same request only replaces the // running winner when strictly better, so the final page/hidden @@ -3653,7 +3636,7 @@ impl LocalQwen3Lane { Some(bufs.captured_hidden()), verify_round, )?; - if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + if trace { for idx in 0..requests.len() { let has_hedge = hedge_spans .iter() @@ -3670,14 +3653,14 @@ impl LocalQwen3Lane { .iter() .enumerate() .filter(|(_, (request_idx, _))| *request_idx == idx) - .map(|(slot, _)| raw_b_lengths[slot].to_string()) + .map(|(slot, _)| raw_lengths[requests.len() + slot].to_string()) .collect::>() .join(","); log::debug!( "Qwen3 DFlash hedge detail round={} request={} raw_a={} raw_b_lens={} selected={} selected_len={} matched_draft={}", verify_round, requests[idx].request_id, - raw_a_lengths[idx], + raw_lengths[idx], raw_b_lens, selected, final_results[idx].accepted_tokens.len(), diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 5d309f280..e126b43d8 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -26,8 +26,6 @@ pub(super) fn truncate_after_terminal( policy: &StopPolicy, model_eos: &[u32], ) { - // Keep this helper idempotent: the worker trims before DFlash context is - // recorded, and the executor later asserts the same invariant at commit. let Some(keep) = result.accepted_tokens.iter().position(|&token| { policy .classify(token, |id| model_eos.contains(&id)) @@ -137,27 +135,16 @@ impl Qwen3Executor { req.request_id )); } - } - // The worker must have normalized each span before recording DFlash - // context (both the hedged and plain verify paths truncate at the first - // terminal token). Re-assert that invariant at commit: a broken worker - // fails loudly here instead of being silently truncated a second time. - for (req, req_result) in plan.requests.iter().zip(&result.requests) { - let terminal = |token: u32| { + // Workers normalize before recording context. Reject a broken span + // before any KV commit instead of silently repairing it here. + let terminal_position = req_result.accepted_tokens.iter().position(|&token| { req.stop_policy .classify(token, |id| self.metadata.stop_token_ids.contains(&id)) .is_some() - }; - let terminal_position = req_result - .accepted_tokens - .iter() - .position(|&token| terminal(token)); + }); if let Some(position) = terminal_position && position + 1 != req_result.accepted_tokens.len() { - // Nothing has committed yet: roll back every reservation before - // surfacing the worker defect, mirroring the request-mismatch - // failures above. self.revert_speculative_schedules(&scheduled); return Err(anyhow::anyhow!( "speculative worker returned an untruncated span for {:?}: \ diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index e9ff31b41..eaedf6921 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -54,9 +54,9 @@ pub(crate) fn resolve_step( /// Turn each request's accepted speculative span into a decode effect. A span /// commits 1..=K+1 tokens at once; we walk it in order so a stop token or the -/// max-output budget lands exactly where expected. The executor has already -/// truncated any suffix after a request-terminal token to keep speculative -/// state consistent; the resolver classifies its typed cause here. +/// max-output budget lands exactly where expected. Workers truncate terminal +/// suffixes before recording context, and the executor checks that invariant +/// before KV commit; the resolver classifies the retained trigger here. pub(crate) fn resolve_speculative_outputs( executor: &impl ModelExecutor, active: &[ActiveRequestState], diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index dc1215223..a8eabd6d9 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -72,16 +72,13 @@ pub(crate) struct VerifyPlan<'a> { #[derive(Clone, Debug)] pub(crate) struct VerifyRequestResult { pub request_id: RequestId, - /// Number of draft candidates accepted before the posterior bonus. + /// Number of matched draft candidates retained after terminal truncation. pub matched_draft_tokens: usize, - /// Tokens to commit: the accepted draft prefix followed by the target's - /// posterior token at the first mismatch (or the block-end continuation - /// when every draft is accepted). Always `1..=K + 1` tokens, so a verify - /// step always makes at least one token of progress. The worker normalizes - /// this span after the first request-terminal token before DFlash context is - /// recorded; the executor asserts that invariant at commit rather than - /// re-truncating. The scheduler retains ownership of typed stop-cause - /// emission. + /// Tokens to commit: the accepted draft prefix and the target's posterior + /// token, unless a terminal draft ends the span before the posterior. + /// Always `1..=K + 1` tokens. The worker truncates after the first terminal + /// token before recording DFlash context; the executor checks this at + /// commit, and the scheduler emits the typed stop cause. pub accepted_tokens: Vec, } diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index a117d2b2d..d9a1b879b 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -677,8 +677,17 @@ fn dflash_hedged_midspan_stop_retains_trigger() { baseline_params, GENERATED_TOKENS, )) - .expect_finished() - .tokens; + .expect_finished(); + assert!(matches!( + baseline.terminal, + Terminal::Finished { + reason: FinishReason::Length, + stop_cause: None, + completion_tokens, + .. + } if completion_tokens == GENERATED_TOKENS && completion_tokens == baseline.tokens.len() + )); + let baseline = baseline.tokens; let candidate_stops: Vec<(usize, u32)> = baseline .iter() .enumerate() diff --git a/scripts/qwen3_stop_contract_probe.py b/scripts/qwen3_stop_contract_probe.py index e2bc0edb2..4fec15e49 100644 --- a/scripts/qwen3_stop_contract_probe.py +++ b/scripts/qwen3_stop_contract_probe.py @@ -1,453 +1,472 @@ #!/usr/bin/env python3 -"""Compare the Qwen3 typed stop contract with an un-migrated Qwen3.5 path. +"""Validate the Qwen3 stop contract against a live, un-migrated Qwen3.5 server. -The script talks to already-running OpenAI-compatible PegaInfer servers. It -does not start a server and deliberately contains no model-specific launch -flags. By default each target receives a full-vocabulary explicit stop set; -this guarantees that the first sampled token exercises the explicit-stop path. -Use --stop-token-id when a smaller, known stop set is preferred. +Both servers must already be running. The default explicit stop set covers the +vocabulary, so its first returned token must stop generation. --stop-token-id +selects a known trigger instead. Token IDs, token-ID-formatted logprobs, usage, +and the first terminal position are checked together. SSE is consumed through +[DONE] to HTTP EOF, including content sharing a frame with finish metadata. -Every check validates the wire shape, not merely the presence of a field: +Provide each model's primary EOS ID as resolved by the serving tokenizer; +an unverified EOS must never count as a passing contract check. A healthy legacy +server may fail the new stop semantics, but malformed responses or unavailable +servers are not evidence of a compatibility gap. -- `stop_reason` must be a numeric token ID from the requested stop set (or - absent for a model-EOS finish); -- under a full-vocabulary stop set the first token always matches, so the - explicit-stop runs must report exactly one completion token; -- the triggering token must carry a non-null numeric logprob; -- the exact trigger is verified by requesting `return_token_ids` and comparing - the final returned token ID against `stop_reason`; -- when `--qwen3-eos-token-id` / `--qwen35-eos-token-id` is given, an EOS - finish must end on exactly that trigger ID; -- the streaming case must terminate with `[DONE]` and must not emit content - after the finish event; -- `/v1/models` must actually serve the requested model name. - -Run `--self-check` to exercise the checks against a local mock service that -serves the malformed-response shapes these checks exist to reject; no GPU or -server binary is required. Both targets are reported side by side, and -`--require-legacy-gap` turns the expected adapted-passes/legacy-fails outcome -into an exit-code assertion. - -Example: +Example (Qwen3-4B and Qwen3.5-0.8B): python3 scripts/qwen3_stop_contract_probe.py \ - --qwen3-url http://127.0.0.1:18081 --qwen3-model qwen3-adapted \ - --qwen35-url http://127.0.0.1:18082 --qwen35-model qwen35-legacy \ - --out stop-contract-ab.json + --qwen3-eos-token-id 151645 --qwen35-eos-token-id 248046 \ + --require-legacy-gap --out stop-contract-ab.json + +--self-check tests this probe's assertions using isolated malformed HTTP +responses. It needs no GPU and is not evidence about inference correctness. """ from __future__ import annotations import argparse +import copy import json import math import sys import threading import time from concurrent.futures import ThreadPoolExecutor +from http.client import HTTPException from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path from typing import Any from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen -DEFAULT_PROMPT = "The capital of France is" -DEFAULT_QWEN3_URL = "http://127.0.0.1:18081" -DEFAULT_QWEN35_URL = "http://127.0.0.1:18082" -DEFAULT_QWEN3_MODEL = "qwen3-adapted" -DEFAULT_QWEN35_MODEL = "qwen35-legacy" -DEFAULT_QWEN3_VOCAB = 151_936 -DEFAULT_QWEN35_VOCAB = 248_320 - -def endpoint(base_url: str, suffix: str) -> str: - return base_url.rstrip("/") + suffix +class InvalidResponse(ValueError): + def __init__(self, code: str, detail: str): + super().__init__(f"{code}: {detail}") + self.code = code -def http_json(url: str, payload: dict[str, Any] | None, timeout: float) -> dict[str, Any]: - data = None - headers = {"Accept": "application/json"} - method = "GET" - if payload is not None: - data = json.dumps(payload, separators=(",", ":")).encode("utf-8") - headers["Content-Type"] = "application/json" - method = "POST" - request = Request(url, data=data, headers=headers, method=method) - started = time.perf_counter() - try: - with urlopen(request, timeout=timeout) as response: - raw = response.read() - status = response.status - except HTTPError as error: - raw = error.read() - try: - body: Any = json.loads(raw.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError): - body = raw.decode("utf-8", errors="replace") - return { - "http_status": error.code, - "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), - "error": body, - } - except (TimeoutError, URLError, OSError) as error: - return { - "http_status": None, - "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), - "error": str(error), - } - try: - body = json.loads(raw.decode("utf-8")) - except (UnicodeDecodeError, json.JSONDecodeError) as error: - return { - "http_status": status, - "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), - "error": f"invalid JSON response: {error}", - } - return { - "http_status": status, - "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), - "body": body, - } - +def require(condition: bool, code: str, detail: str) -> None: + if not condition: + raise InvalidResponse(code, detail) -def is_finite_number(value: Any) -> bool: - if isinstance(value, bool) or not isinstance(value, (int, float)): - return False - return math.isfinite(value) +def is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) -def token_ids_from_choice(choice: dict[str, Any]) -> list[int] | None: - ids = choice.get("token_ids") - if isinstance(ids, list) and all(is_int(item) for item in ids): - return ids - return None +def finite_logprob(value: Any) -> bool: + return ( + isinstance(value, (int, float)) + and not isinstance(value, bool) + and -sys.float_info.max <= value <= 0 + and math.isfinite(value) + ) -def logprob_at_last(choice: dict[str, Any], token_count: int) -> float | None: - """Finite logprob of the final emitted token, or None when absent. - The logprob table must cover exactly the emitted tokens: a table shorter - than the token sequence (a missing trigger entry) or a null final entry is - a failure rather than an inherited value from an earlier token. - """ - logprobs = choice.get("logprobs") - if not isinstance(logprobs, dict) or token_count == 0: - return None - token_logprobs = logprobs.get("token_logprobs") - if isinstance(token_logprobs, list): - if len(token_logprobs) != token_count: - return None - value = token_logprobs[-1] - return value if is_finite_number(value) else None - content = logprobs.get("content") - if isinstance(content, list): - if len(content) != token_count: - return None - last = content[-1] - if isinstance(last, dict): - value = last.get("logprob") - return value if is_finite_number(value) else None - return None +def token_ids(value: Any, vocab_size: int) -> list[int]: + require(isinstance(value, list), "token_ids_missing", "expected a token_ids array") + require( + all(is_int(item) and 0 <= item < vocab_size for item in value), + "token_id_range", + "token IDs must be integers within the model vocabulary", + ) + return value -def stream_token_count(choice: dict[str, Any]) -> int: - """Number of tokens a streaming choice emits (token_ids, logprobs, or text).""" - ids = token_ids_from_choice(choice) - if ids is not None: - return len(ids) - logprobs = choice.get("logprobs") - if isinstance(logprobs, dict): - tokens = logprobs.get("tokens") - if isinstance(tokens, list): - return len(tokens) - content = logprobs.get("content") - if isinstance(content, list): - return len(content) - return 0 +def choice_data( + choice: dict[str, Any], vocab_size: int, want_logprobs: bool, *, stream: bool +) -> tuple[list[int], list[float]]: + require(isinstance(choice.get("text"), str), "text_type", "choice.text must be a string") + finish = choice.get("finish_reason") + require(finish in (None, "stop", "length"), "finish_reason", "unexpected finish_reason") + require( + stream or finish is not None, "finish_reason", "a completion requires terminal metadata" + ) + reason = choice.get("stop_reason") + require( + reason is None or is_int(reason), + "stop_reason_type", + "stop_reason must be an integer or null", + ) + require( + reason is None or 0 <= reason < vocab_size, + "stop_reason_range", + "stop_reason is outside the vocabulary", + ) + require( + reason is None or finish == "stop", + "stop_reason_without_stop", + "stop_reason requires a stop finish", + ) + ids = choice.get("token_ids") + lp = choice.get("logprobs") + if stream and ids is None: + # vLLM emits prompt metadata and a pure finish frame without generated IDs. + require( + not choice["text"] + and lp is None + and (finish is not None or choice.get("prompt_token_ids") is not None), + "token_ids_missing", + "content requires returned token IDs", + ) + return [], [] + ids = token_ids(ids, vocab_size) + require( + not (stream and finish is not None and choice["text"] and not ids), + "terminal_text_without_tokens", + "text-only decoder flush must precede finish metadata", + ) + if lp is None: + require( + not (want_logprobs and ids), + "logprobs_missing", + "generated tokens need their own logprobs", + ) + return ids, [] + require(isinstance(lp, dict), "logprobs_shape", "expected completions-style logprobs") + fields = ("tokens", "token_logprobs", "top_logprobs", "text_offset") + require( + all(isinstance(lp.get(key), list) and len(lp[key]) == len(ids) for key in fields), + "logprob_count_mismatch", + "all logprob arrays must cover exactly the returned IDs", + ) + require( + lp["tokens"] == [f"token_id:{item}" for item in ids], + "logprob_token_mismatch", + "logprob token identities differ from returned IDs", + ) + require( + all(finite_logprob(value) for value in lp["token_logprobs"]), + "logprob_value", + "invalid token logprob", + ) + require( + all(is_int(value) and value >= 0 for value in lp["text_offset"]), + "logprob_offset", + "text offsets must be nonnegative integers", + ) + for top in lp["top_logprobs"]: + require(isinstance(top, dict), "top_logprobs_shape", "expected a top-logprob mapping") + for key, value in top.items(): + require( + isinstance(key, str) + and key.startswith("token_id:") + and key[9:].isascii() + and key[9:].isdigit() + and 0 <= int(key[9:]) < vocab_size + and finite_logprob(value), + "top_logprobs_value", + "invalid top-logprob token or value", + ) + return ids, lp["token_logprobs"] -def first_choice(body: dict[str, Any]) -> dict[str, Any]: +def response_choice(body: Any, model: str) -> dict[str, Any]: + require(isinstance(body, dict), "response_shape", "expected a JSON object") + require("error" not in body, "server_error", "server returned an error response") + require( + body.get("model") == model, "model_mismatch", "response model differs from requested model" + ) choices = body.get("choices") - if not isinstance(choices, list) or not choices: - return {} + require( + isinstance(choices, list) and len(choices) == 1, + "choice_count", + "n=1 requires exactly one choice", + ) choice = choices[0] - return choice if isinstance(choice, dict) else {} + require(isinstance(choice, dict), "choice_shape", "expected a choice object") + require( + is_int(choice.get("index")) and choice["index"] == 0, + "choice_index", + "expected choice index 0", + ) + require( + body.get("stop_reason") is None, "stop_reason_location", "stop_reason belongs to the choice" + ) + return choice -def completion_summary(response: dict[str, Any]) -> dict[str, Any]: - body = response.get("body") - if not isinstance(body, dict): - return { - **response, - "finish_reason": None, - "stop_reason": None, - "completion_tokens": None, - "token_ids": None, - "trigger_logprob": None, - } - choice = first_choice(body) - usage = body.get("usage") if isinstance(body.get("usage"), dict) else {} - stop_reason = choice.get("stop_reason", body.get("stop_reason")) - token_ids = token_ids_from_choice(choice) - return { - "http_status": response.get("http_status"), - "elapsed_ms": response.get("elapsed_ms"), - "finish_reason": choice.get("finish_reason"), - "stop_reason": stop_reason, - "completion_tokens": usage.get("completion_tokens"), - "token_ids": token_ids, - "trigger_logprob": logprob_at_last(choice, len(token_ids) if token_ids is not None else 0), - } +def completion_count(usage: Any, ids: list[int]) -> int: + require(isinstance(usage, dict), "usage_missing", "expected completion usage") + fields = ("prompt_tokens", "completion_tokens", "total_tokens") + require( + all(is_int(usage.get(key)) and usage[key] >= 0 for key in fields), + "usage_type", + "usage counts must be nonnegative integers", + ) + require( + usage["total_tokens"] == usage["prompt_tokens"] + usage["completion_tokens"], + "usage_total", + "inconsistent total_tokens", + ) + require( + usage["completion_tokens"] == len(ids), + "completion_count_mismatch", + "usage differs from returned token count", + ) + return usage["completion_tokens"] + + +def read_stream(response: Any, model: str, vocab_size: int, timeout: float) -> dict[str, Any]: + result: dict[str, Any] = {"token_ids": [], "token_logprobs": [], "text": "", "done_seen": False} + finished = False + prompt_seen = False + usage = None + data_lines: list[str] = [] + deadline = time.monotonic() + timeout + + def event(data: str) -> None: + nonlocal finished, prompt_seen, usage + require(not result["done_seen"], "stream_after_done", "received an event after [DONE]") + if data == "[DONE]": + require(finished, "stream_missing_finish", "[DONE] arrived before finish metadata") + require( + usage is not None, + "stream_missing_usage", + "include_usage requires a final usage frame", + ) + result["done_seen"] = True + return + body = json.loads(data) + require(isinstance(body, dict), "response_shape", "expected an SSE JSON object") + require("error" not in body, "server_error", "server returned an SSE error") + require( + body.get("model") == model, + "model_mismatch", + "stream model differs from requested model", + ) + if body.get("choices") == []: + require( + finished and usage is None, + "stream_usage_order", + "usage must occur once after finish", + ) + require( + isinstance(body.get("usage"), dict), "usage_missing", "empty choices require usage" + ) + require( + all( + body.get(key) is None + for key in ("text", "token_ids", "logprobs", "finish_reason", "stop_reason") + ), + "stream_usage_content", + "usage-only frame contains completion content", + ) + usage = body["usage"] + return + require(not finished, "stream_after_finish", "received a choice after finish metadata") + choice = response_choice(body, model) + require(body.get("usage") is None, "stream_usage_order", "unexpected per-delta usage") + prompt = choice.get("prompt_token_ids") + if prompt is not None: + require( + not prompt_seen and not result["token_ids"], + "stream_prompt_order", + "prompt metadata must precede generated tokens", + ) + token_ids(prompt, vocab_size) + prompt_seen = True + ids, values = choice_data(choice, vocab_size, True, stream=True) + result["token_ids"].extend(ids) + result["token_logprobs"].extend(values) + result["text"] += choice["text"] + if choice.get("finish_reason") is not None: + result["finish_reason"] = choice["finish_reason"] + result["stop_reason"] = choice.get("stop_reason") + finished = True + + while True: + require( + time.monotonic() <= deadline, + "stream_timeout", + "stream did not reach HTTP EOF within timeout", + ) + raw = response.readline(1_048_577) + if not raw: + break + require(len(raw) <= 1_048_576, "sse_line_size", "SSE line exceeds 1 MiB") + line = raw.decode("utf-8").rstrip("\r\n") + if not line: + if data_lines: + event("\n".join(data_lines)) + data_lines.clear() + elif line.startswith(":"): + continue + else: + field, _, value = line.partition(":") + if field == "data": + require(not result["done_seen"], "stream_after_done", "received data after [DONE]") + data_lines.append(value.removeprefix(" ")) + require(not data_lines, "sse_incomplete_event", "EOF before the SSE event delimiter") + require(result["done_seen"], "stream_missing_done", "EOF before [DONE]") + result["completion_tokens"] = completion_count(usage, result["token_ids"]) + return result -def completion_payload( +def request_completion( + url: str, model: str, - prompt: str, - max_tokens: int, + vocab_size: int, + args: argparse.Namespace, + *, ignore_eos: bool, - stop_token_ids: list[int] | None, - logprobs: int, - stream: bool, + stop_ids: list[int] | None = None, + logprobs: bool = False, + stream: bool = False, ) -> dict[str, Any]: payload: dict[str, Any] = { "model": model, - "prompt": prompt, + "prompt": args.prompt, "temperature": 0.0, - "max_tokens": max_tokens, + "max_tokens": args.max_tokens, + "n": 1, "ignore_eos": ignore_eos, "stream": stream, "return_token_ids": True, + "return_tokens_as_token_ids": True, } - if stop_token_ids is not None: - payload["stop_token_ids"] = stop_token_ids - if logprobs > 0: - payload["logprobs"] = logprobs - return payload - - -def request_completion( - base_url: str, - model: str, - prompt: str, - max_tokens: int, - timeout: float, - *, - ignore_eos: bool, - stop_token_ids: list[int] | None = None, - logprobs: int = 0, -) -> dict[str, Any]: - payload = completion_payload(model, prompt, max_tokens, ignore_eos, stop_token_ids, logprobs, stream=False) - return completion_summary(http_json(endpoint(base_url, "/v1/completions"), payload, timeout)) - - -def request_stream( - base_url: str, - model: str, - prompt: str, - max_tokens: int, - timeout: float, - *, - ignore_eos: bool, - stop_token_ids: list[int] | None = None, - logprobs: int = 0, -) -> dict[str, Any]: - payload = completion_payload(model, prompt, max_tokens, ignore_eos, stop_token_ids, logprobs, stream=True) - data = json.dumps(payload, separators=(",", ":")).encode("utf-8") + if stop_ids is not None: + payload["stop_token_ids"] = stop_ids + if logprobs or stream: + payload["logprobs"] = 1 + if stream: + payload["stream_options"] = {"include_usage": True} request = Request( - endpoint(base_url, "/v1/completions"), - data=data, - headers={"Accept": "text/event-stream", "Content-Type": "application/json"}, - method="POST", + url.rstrip("/") + "/v1/completions", + data=json.dumps(payload).encode(), + headers={ + "Content-Type": "application/json", + "Accept": "text/event-stream" if stream else "application/json", + }, ) - started = time.perf_counter() - events = 0 - emitted_tokens = 0 - stream_token_ids: list[int] = [] - last_content_logprob: float | None = None - trigger_logprob_valid = False - finish_reason: Any = None - stop_reason: Any = None - done_seen = False - content_after_finish = False - malformed_chunks = 0 - finished = False - status = None - error = None + started = time.monotonic() + result: dict[str, Any] = {"http_status": None, "valid_response": False, "error_codes": []} try: - with urlopen(request, timeout=timeout) as response: - status = response.status - for raw_line in response: - line = raw_line.decode("utf-8", errors="replace").strip() - if not line.startswith("data:"): - continue - data_line = line[5:].strip() - if data_line == "[DONE]": - done_seen = True - break - try: - chunk = json.loads(data_line) - except json.JSONDecodeError: - malformed_chunks += 1 - continue - if not isinstance(chunk, dict): - malformed_chunks += 1 - continue - choices = chunk.get("choices") - if not isinstance(choices, list) or not choices: - # A usage-only frame carries no content and is legal. - if isinstance(chunk.get("usage"), dict): - continue - malformed_chunks += 1 - continue - choice = choices[0] if isinstance(choices[0], dict) else {} - text = choice.get("text") - frame_ids = token_ids_from_choice(choice) - token_count = len(frame_ids) if frame_ids is not None else stream_token_count(choice) - has_content = bool(text) or token_count > 0 - this_finish = choice.get("finish_reason") - this_stop = chunk.get("stop_reason", choice.get("stop_reason")) - terminal_event = this_finish is not None or this_stop is not None - if finished: - if has_content or terminal_event: - content_after_finish = True - continue - if has_content: - # Count content even when it shares a frame with the finish - # metadata; a terminal frame must not hide emitted tokens. - events += 1 - emitted_tokens += token_count if token_count else 1 - if frame_ids is not None: - stream_token_ids.extend(frame_ids) - value = logprob_at_last(choice, token_count) - trigger_logprob_valid = value is not None - if value is not None: - last_content_logprob = value - if terminal_event: - if this_finish is not None: - finish_reason = this_finish - if stop_reason is None and this_stop is not None: - stop_reason = this_stop - finished = True - except HTTPError as exc: - status = exc.code - error = exc.read().decode("utf-8", errors="replace") - except (TimeoutError, URLError, OSError) as exc: - error = str(exc) - return { - "http_status": status, - "elapsed_ms": round((time.perf_counter() - started) * 1000, 2), - "stream_events": events, - "emitted_tokens": emitted_tokens, - "token_ids": stream_token_ids, - "trigger_logprob": last_content_logprob if trigger_logprob_valid else None, - "finish_reason": finish_reason, - "stop_reason": stop_reason, - "completion_tokens": None, - "done_seen": done_seen, - "content_after_finish": content_after_finish, - "malformed_chunks": malformed_chunks, - "error": error, - } - - -def probe_models(base_url: str, timeout: float, expected_model: str) -> dict[str, Any]: - result = http_json(endpoint(base_url, "/v1/models"), None, timeout) - body = result.get("body") - ids: list[str] = [] - if isinstance(body, dict) and isinstance(body.get("data"), list): - ids = [ - item.get("id") - for item in body["data"] - if isinstance(item, dict) and item.get("id") - ] - result["model_ids"] = ids - result["expected_model_present"] = expected_model in ids + with urlopen(request, timeout=args.timeout) as response: + result["http_status"] = response.status + require(response.status == 200, "http_status", "expected HTTP 200") + expected_type = "text/event-stream" if stream else "application/json" + require( + response.headers.get_content_type() == expected_type, + "content_type", + "unexpected response Content-Type", + ) + if stream: + result.update(read_stream(response, model, vocab_size, args.timeout)) + else: + body = json.loads(response.read()) + choice = response_choice(body, model) + ids, values = choice_data(choice, vocab_size, logprobs, stream=False) + result.update( + { + "token_ids": ids, + "token_logprobs": values, + "text": choice["text"], + "finish_reason": choice.get("finish_reason"), + "stop_reason": choice.get("stop_reason"), + "completion_tokens": completion_count(body.get("usage"), ids), + } + ) + result["valid_response"] = True + values = result["token_logprobs"] + result["trigger_logprob"] = values[-1] if values else None + except InvalidResponse as error: + result.update(error=str(error), error_codes=[error.code]) + except (UnicodeDecodeError, json.JSONDecodeError) as error: + result.update(error=f"invalid_json: {error}", error_codes=["invalid_json"]) + except HTTPError as error: + result.update( + http_status=error.code, error=f"http_error: {error}", error_codes=["http_error"] + ) + except (URLError, OSError, HTTPException) as error: + result.update(error=f"transport_error: {error}", error_codes=["transport_error"]) + result["elapsed_ms"] = round((time.monotonic() - started) * 1000, 2) return result -def is_int(value: Any) -> bool: - # JSON `true` parses to Python True, which is also an `int` subclass; - # exclude it so a boolean wire value can never satisfy a numeric check. - return isinstance(value, int) and not isinstance(value, bool) - - -def numeric_stop_reason(result: dict[str, Any]) -> bool: - return is_int(result.get("stop_reason")) - - -def is_stop_finish(result: dict[str, Any]) -> bool: - return result.get("http_status") == 200 and result.get("finish_reason") == "stop" - - -def typed_stop(result: dict[str, Any], stop_ids: list[int], full_vocab: bool) -> bool: - """A typed request stop: the reported stop_reason is the emitted trigger.""" - if not is_stop_finish(result): - return False - if not numeric_stop_reason(result) or result["stop_reason"] not in stop_ids: - return False - ids = result.get("token_ids") - if not isinstance(ids, list) or not ids: - return False - tokens = result.get("completion_tokens") - if not is_int(tokens) or tokens != len(ids) or tokens < 1: - return False - if ids[-1] != result["stop_reason"]: - return False - if full_vocab: - # The first sampled token is in the full-vocabulary stop set, so the - # trigger must be the only completion token. - return len(ids) == 1 - return True - - -def eos_or_typed_stop( +def validate_sequence( result: dict[str, Any], - stop_ids: list[int], - full_vocab: bool, - eos_id: int | None, -) -> bool: - """EOS-enabled explicit-stop case: a model EOS may win over the stop set.""" - if not is_stop_finish(result): - return False - ids = result.get("token_ids") - if not isinstance(ids, list) or not ids: - return False - tokens = result.get("completion_tokens") - if not is_int(tokens) or tokens != len(ids) or tokens < 1: - return False - stop_reason = result.get("stop_reason") - if stop_reason is None: - # A model EOS won. With a configured EOS ID the final token must be - # exactly that trigger; otherwise only the count and terminal shape - # are verifiable. - if eos_id is not None and ids[-1] != eos_id: - return False - return not full_vocab or len(ids) == 1 - if not is_int(stop_reason) or stop_reason not in stop_ids: - return False - if ids[-1] != stop_reason: - return False - return not full_vocab or len(ids) == 1 - - -def length_control(result: dict[str, Any], max_tokens: int) -> bool: - ids = result.get("token_ids") - return ( - result.get("http_status") == 200 - and result.get("finish_reason") == "length" - and result.get("stop_reason") is None - and isinstance(ids, list) - and len(ids) == max_tokens - and is_int(result.get("completion_tokens")) - and result.get("completion_tokens") == max_tokens - ) + stop_ids: set[int], + max_tokens: int, + *, + ignore_eos: bool, + eos_id: int, + require_stop: bool, +) -> list[str]: + if not result["valid_response"]: + return result["error_codes"] + try: + ids = result["token_ids"] + require( + 0 < len(ids) <= max_tokens, + "token_budget_exceeded", + "output must contain 1..max_tokens tokens", + ) + first_stop = next( + ( + index + for index, token in enumerate(ids) + if token in stop_ids or (not ignore_eos and token == eos_id) + ), + None, + ) + if first_stop is None: + require( + result["finish_reason"] == "length" + and result["stop_reason"] is None + and len(ids) == max_tokens, + "length_outcome", + "a sequence without a trigger must exhaust the token budget", + ) + require( + not require_stop, "stop_not_exercised", "the selected stop token was not generated" + ) + else: + require( + first_stop == len(ids) - 1, + "tokens_after_stop", + f"first stop at {first_stop}, but output has {len(ids)} tokens", + ) + require( + result["finish_reason"] == "stop", + "stop_finish_reason", + "the trigger must finish with stop", + ) + expected_reason = None if not ignore_eos and ids[-1] == eos_id else ids[-1] + require( + result["stop_reason"] == expected_reason, + "stop_reason_mismatch", + f"expected stop_reason={expected_reason!r}", + ) + except InvalidResponse as error: + return [error.code] + return [] -def build_stop_ids(vocab_size: int, selected_id: int | None) -> list[int]: - if selected_id is not None: - if selected_id < 0 or selected_id >= vocab_size: - raise ValueError(f"--stop-token-id must be in [0, {vocab_size}), got {selected_id}") - return [selected_id] - return list(range(vocab_size)) +def probe_models(url: str, model: str, timeout: float) -> dict[str, Any]: + try: + with urlopen(url.rstrip("/") + "/v1/models", timeout=timeout) as response: + require(response.status == 200, "http_status", "expected HTTP 200") + body = json.loads(response.read()) + require( + isinstance(body, dict) and isinstance(body.get("data"), list), + "models_shape", + "expected a model list", + ) + require( + any(isinstance(item, dict) and item.get("id") == model for item in body["data"]), + "model_missing", + "requested model is not served", + ) + return {"expected_model_present": True, "error_codes": []} + except InvalidResponse as error: + return {"expected_model_present": False, "error_codes": [error.code], "error": str(error)} + except (UnicodeDecodeError, json.JSONDecodeError, URLError, OSError, HTTPException) as error: + return { + "expected_model_present": False, + "error_codes": ["models_unavailable"], + "error": str(error), + } def run_target( @@ -456,235 +475,185 @@ def run_target( model: str, vocab_size: int, args: argparse.Namespace, - eos_token_id: int | None = None, + eos_token_id: int, ) -> dict[str, Any]: - stop_ids = build_stop_ids(vocab_size, args.stop_token_id) - full_vocab = args.stop_token_id is None - model_probe = probe_models(url, args.timeout, model) + require(vocab_size > 0, "vocab_size", "vocabulary size must be positive") + require( + is_int(eos_token_id) and 0 <= eos_token_id < vocab_size, + "eos_id", + "supply the model's primary EOS ID within its vocabulary", + ) + selected = args.stop_token_id + require( + selected is None or 0 <= selected < vocab_size, + "stop_id", + "selected stop ID is outside the vocabulary", + ) + stops = list(range(vocab_size)) if selected is None else [selected] + stop_set = set(stops) + models = probe_models(url, model, args.timeout) + cases: dict[str, Any] = {} - def call(ignore_eos: bool, ids: list[int] | None, lp: int = 0) -> dict[str, Any]: - return request_completion( + def call( + label: str, + *, + explicit: bool, + ignore_eos: bool = True, + lp: bool = False, + stream: bool = False, + descending: bool = False, + ) -> dict[str, Any]: + ids = list(reversed(stops)) if descending else stops + result = request_completion( url, model, - args.prompt, - args.max_tokens, - args.timeout, + vocab_size, + args, ignore_eos=ignore_eos, - stop_token_ids=ids, + stop_ids=ids if explicit else None, logprobs=lp, + stream=stream, ) - - cases: dict[str, Any] = {} - cases["control"] = call(True, None) - cases["explicit_stop_ignore_eos"] = call(True, stop_ids) - cases["explicit_stop_eos_enabled"] = call(False, stop_ids) - cases["stop_ascending"] = call(True, stop_ids) - cases["stop_descending"] = call(True, list(reversed(stop_ids))) - cases["trigger_logprob"] = call(True, stop_ids, lp=1) - cases["streaming"] = request_stream( - url, - model, - args.prompt, - args.max_tokens, - args.timeout, - ignore_eos=True, - stop_token_ids=stop_ids, - logprobs=1, - ) - - mixed_jobs: list[tuple[str, bool]] = [("control", False)] * 3 + [("explicit", True)] * 3 - - def mixed_call(job: tuple[str, bool]) -> dict[str, Any]: - _, explicit = job - return request_completion( - url, - model, - args.prompt, + result["validation_errors"] = validate_sequence( + result, + stop_set if explicit else set(), args.max_tokens, - args.timeout, - ignore_eos=True, - stop_token_ids=stop_ids if explicit else None, + ignore_eos=ignore_eos, + eos_id=eos_token_id, + require_stop=explicit, ) - - with ThreadPoolExecutor(max_workers=len(mixed_jobs)) as pool: - mixed_results = list(pool.map(mixed_call, mixed_jobs)) - cases["mixed"] = [ - {"kind": kind, "result": result} - for (kind, _), result in zip(mixed_jobs, mixed_results) - ] - - mixed_controls = [item["result"] for item in cases["mixed"] if item["kind"] == "control"] - mixed_explicit = [item["result"] for item in cases["mixed"] if item["kind"] == "explicit"] - streaming = cases["streaming"] - checks = { - "model_present": bool(model_probe.get("expected_model_present")), - "baseline_control": length_control(cases["control"], args.max_tokens), - "explicit_stop_ignore_eos": typed_stop(cases["explicit_stop_ignore_eos"], stop_ids, full_vocab), - "explicit_stop_eos_enabled": eos_or_typed_stop( - cases["explicit_stop_eos_enabled"], stop_ids, full_vocab, eos_token_id - ), - "stop_set_order_invariant": ( - typed_stop(cases["stop_ascending"], stop_ids, full_vocab) - and typed_stop(cases["stop_descending"], stop_ids, full_vocab) - and cases["stop_ascending"].get("stop_reason") - == cases["stop_descending"].get("stop_reason") - and cases["stop_ascending"].get("completion_tokens") - == cases["stop_descending"].get("completion_tokens") - ), - "trigger_logprob_preserved": ( - typed_stop(cases["trigger_logprob"], stop_ids, full_vocab) - and is_finite_number(cases["trigger_logprob"].get("trigger_logprob")) - ), - "stream_reports_typed_stop": ( - streaming.get("http_status") == 200 - and streaming.get("finish_reason") == "stop" - and numeric_stop_reason(streaming) - and streaming.get("stop_reason") in stop_ids - and streaming.get("stream_events", 0) > 0 - and streaming.get("done_seen") is True - and not streaming.get("content_after_finish") - and streaming.get("malformed_chunks", 0) == 0 - and is_finite_number(streaming.get("trigger_logprob")) - and isinstance(streaming.get("token_ids"), list) - and len(streaming["token_ids"]) > 0 - and streaming["token_ids"][-1] == streaming.get("stop_reason") - and streaming.get("emitted_tokens") == len(streaming["token_ids"]) - and ( - streaming.get("emitted_tokens") == 1 - if full_vocab - else (is_int(streaming.get("emitted_tokens")) and streaming["emitted_tokens"] >= 1) + return {"name": label, **result} + + for label, options in ( + ("control", {"explicit": False}), + ("explicit_stop_ignore_eos", {"explicit": True}), + ("explicit_stop_eos_enabled", {"explicit": True, "ignore_eos": False}), + ("stop_descending", {"explicit": True, "descending": True}), + ("trigger_logprob", {"explicit": True, "lp": True}), + ("streaming", {"explicit": True, "stream": True}), + ): + cases[label] = call(label, **options) + with ThreadPoolExecutor(max_workers=6) as pool: + cases["mixed"] = list( + pool.map( + lambda explicit: call("explicit" if explicit else "control", explicit=explicit), + [False, True] * 3, ) - ), - "mixed_controls_pass_3_of_3": sum(length_control(item, args.max_tokens) for item in mixed_controls) == 3, - "mixed_explicit_stops_pass_3_of_3": sum(typed_stop(item, stop_ids, full_vocab) for item in mixed_explicit) == 3, + ) + ascending = cases["explicit_stop_ignore_eos"] + descending = cases["stop_descending"] + order_errors = ascending["validation_errors"] + descending["validation_errors"] + if not order_errors and any( + ascending[key] != descending[key] for key in ("token_ids", "finish_reason", "stop_reason") + ): + order_errors = ["stop_set_order_changed_output"] + failures = { + "model_present": models["error_codes"], + "baseline_control": cases["control"]["validation_errors"], + "explicit_stop_ignore_eos": ascending["validation_errors"], + "explicit_stop_eos_enabled": cases["explicit_stop_eos_enabled"]["validation_errors"], + "stop_set_order_invariant": order_errors, + "trigger_logprob_preserved": cases["trigger_logprob"]["validation_errors"], + "stream_reports_typed_stop": cases["streaming"]["validation_errors"], + "mixed_controls_pass_3_of_3": [ + code + for item in cases["mixed"] + if item["name"] == "control" + for code in item["validation_errors"] + ], + "mixed_explicit_stops_pass_3_of_3": [ + code + for item in cases["mixed"] + if item["name"] == "explicit" + for code in item["validation_errors"] + ], } + checks = {key: not errors for key, errors in failures.items()} + responses = [value for key, value in cases.items() if key != "mixed"] + cases["mixed"] + healthy = ( + checks["model_present"] + and checks["baseline_control"] + and checks["mixed_controls_pass_3_of_3"] + and all(item["valid_response"] for item in responses) + ) + explicit = [item for item in responses if item["name"] != "control"] + ignored_stops = [ + item["valid_response"] + and item.get("finish_reason") == "length" + and item.get("stop_reason") is None + and item.get("completion_tokens") == args.max_tokens + and any(token in stop_set for token in item.get("token_ids", [])) + for item in explicit + ] + explicit_stop_gap = ( + healthy + and any(ignored_stops) + and all( + not item["validation_errors"] or ignored + for item, ignored in zip(explicit, ignored_stops) + ) + and "stop_set_order_changed_output" not in order_errors + ) return { "name": name, "url": url, "model": model, "vocab_size": vocab_size, - "stop_set": "single" if args.stop_token_id is not None else "full-vocabulary", - "stop_set_size": len(stop_ids), - "model_probe": model_probe, + "eos_token_id": eos_token_id, + "stop_set_size": len(stops), + "model_probe": models, "cases": cases, "checks": checks, + "failures": failures, + "healthy": healthy, + "explicit_stop_gap": explicit_stop_gap, "new_contract_passed": all(checks.values()), } def print_target(target: dict[str, Any]) -> None: - print(f"\n{target['name']} ({target['model']})") - print("case finish stop_reason tokens http") - for name, result in target["cases"].items(): - if name == "mixed": - continue - if name == "streaming": - print( - f"{name:28} {str(result.get('finish_reason')):7} " - f"{str(result.get('stop_reason')):12} {'n/a':>7} " - f"{str(result.get('http_status'))} events={result.get('stream_events')} " - f"done={result.get('done_seen')} tail={result.get('content_after_finish')}" - ) - continue - print( - f"{name:28} {str(result.get('finish_reason')):7} " - f"{str(result.get('stop_reason')):12} {str(result.get('completion_tokens')):>7} " - f"{str(result.get('http_status'))}" - ) - print("checks:") + print(f"\n{target['name']}: {target['url']} ({target['model']})") for name, passed in target["checks"].items(): - print(f" [{'PASS' if passed else 'FAIL'}] {name}") - print(f"overall new contract: {'PASS' if target['new_contract_passed'] else 'FAIL'}") - - -def print_comparison(qwen3: dict[str, Any], qwen35: dict[str, Any]) -> None: - print("\ncontract checks (adapted vs legacy):") - print(f"{'check':28} {'qwen3':7} {'qwen35':7}") - for name in qwen3["checks"]: - left = "PASS" if qwen3["checks"][name] else "FAIL" - right = "PASS" if qwen35["checks"][name] else "FAIL" - print(f"{name:28} {left:7} {right:7}") - print( - f"{'overall new contract':28} " - f"{'PASS' if qwen3['new_contract_passed'] else 'FAIL':7} " - f"{'PASS' if qwen35['new_contract_passed'] else 'FAIL':7}" - ) - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--qwen3-url", default=DEFAULT_QWEN3_URL) - parser.add_argument("--qwen3-model", default=DEFAULT_QWEN3_MODEL) - parser.add_argument("--qwen3-vocab-size", type=int, default=DEFAULT_QWEN3_VOCAB) - parser.add_argument("--qwen35-url", default=DEFAULT_QWEN35_URL) - parser.add_argument("--qwen35-model", default=DEFAULT_QWEN35_MODEL) - parser.add_argument("--qwen35-vocab-size", type=int, default=DEFAULT_QWEN35_VOCAB) - parser.add_argument( - "--qwen3-eos-token-id", - type=int, - default=None, - help="Verify an EOS finish ends on exactly this token ID (optional)", - ) - parser.add_argument( - "--qwen35-eos-token-id", - type=int, - default=None, - help="Verify an EOS finish ends on exactly this token ID (optional)", - ) - parser.add_argument("--prompt", default=DEFAULT_PROMPT) - parser.add_argument("--max-tokens", type=int, default=8) - parser.add_argument("--timeout", type=float, default=300.0) - parser.add_argument( - "--stop-token-id", - type=int, - help="Use one explicit stop ID instead of the default full-vocabulary set", - ) - parser.add_argument("--out", type=Path, help="Write the complete result as JSON") - parser.add_argument( - "--strict-both", - action="store_true", - help="Return failure unless both targets satisfy every new-contract check", - ) - parser.add_argument( - "--self-check", - action="store_true", - help="Run the checks against a local mock service that serves malformed " - "responses; verify each is rejected. Requires no server or GPU.", - ) - parser.add_argument( - "--require-legacy-gap", - action="store_true", - help="Return failure unless the adapted target passes every check and the " - "legacy target fails at least one (the expected A/B outcome).", - ) - return parser.parse_args() + detail = "" if passed else ": " + ", ".join(sorted(set(target["failures"][name]))) + print(f" [{'PASS' if passed else 'FAIL'}] {name}{detail}") + print(f"healthy={target['healthy']}, new_contract_passed={target['new_contract_passed']}") -SELF_CHECK_MODEL = "qwen3-adapted" -SELF_CHECK_STOP_ID = 12095 -SELF_CHECK_EOS_ID = 151645 +# These fixtures exercise the probe, not the inference engine. Each negative +# case changes one response property and must fail for its named diagnostic. +SELF_MODEL = "probe-self-check" +SELF_STOP = 17 +SELF_EOS = 31 +SELF_VOCAB = 32 -_MODE = "valid" -_MODE_LOCK = threading.Lock() - -def set_mock_mode(mode: str) -> None: - global _MODE - with _MODE_LOCK: - _MODE = mode - - -def mock_mode() -> str: - with _MODE_LOCK: - return _MODE +def fixture_choice( + ids: list[int], finish: str | None, reason: int | None, lp: bool +) -> dict[str, Any]: + choice: dict[str, Any] = { + "index": 0, + "text": "", + "token_ids": ids, + "finish_reason": finish, + "stop_reason": reason, + "logprobs": None, + } + if lp: + choice["logprobs"] = { + "tokens": [f"token_id:{item}" for item in ids], + "token_logprobs": [-0.5] * len(ids), + "top_logprobs": [{f"token_id:{item}": -0.5} for item in ids], + "text_offset": list(range(len(ids))), + } + return choice -class MockHandler(BaseHTTPRequestHandler): +class FixtureHandler(BaseHTTPRequestHandler): def log_message(self, *args: Any) -> None: pass - def _json(self, payload: dict[str, Any]) -> None: - raw = json.dumps(payload).encode("utf-8") + def send_json(self, body: dict[str, Any]) -> None: + raw = json.dumps(body).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(raw))) @@ -692,353 +661,320 @@ def _json(self, payload: dict[str, Any]) -> None: self.wfile.write(raw) def do_GET(self) -> None: - if self.path.rstrip("/").endswith("/v1/models"): - if mock_mode() == "wrong_model": - self._json({"data": [{"id": "different-model"}]}) - else: - self._json({"data": [{"id": SELF_CHECK_MODEL}]}) - return - self.send_error(404) + self.send_json({"data": [{"id": SELF_MODEL}]}) def do_POST(self) -> None: - length = int(self.headers.get("Content-Length", "0")) - try: - body = json.loads(self.rfile.read(length) or b"{}") - except json.JSONDecodeError: - self.send_error(400) + payload = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + if not ( + payload.get("return_token_ids") + and payload.get("return_tokens_as_token_ids") + and payload.get("n") == 1 + ): + self.send_error(400, "probe must request exact token identities") return - if not self.path.rstrip("/").endswith("/v1/completions"): - self.send_error(404) + mode = self.server.mode + explicit = "stop_token_ids" in payload + lp = bool(payload.get("logprobs")) + ids = [SELF_STOP] if explicit else [7] * payload["max_tokens"] + finish, reason = ("stop", SELF_STOP) if explicit else ("length", None) + if mode in ("valid_prefix", "missing_final_logprob", "cross_frame_missing_logprob"): + ids = [7, SELF_STOP] + if mode == "earlier_stop": + ids = [SELF_STOP, 7, SELF_STOP] + if mode == "over_budget": + ids = [7] * payload["max_tokens"] + [SELF_STOP] + if mode == "legacy" and explicit: + ids, finish, reason = [7] * payload["max_tokens"], "length", None + if mode == "eos": + ids, reason = [SELF_EOS], None + if mode == "wrong_eos_priority": + ids, reason = [SELF_EOS], SELF_EOS + if mode == "wrong_trigger" and explicit: + reason = 19 + if mode == "fake_eos": + reason = None + if mode == "string_reason": + reason = str(SELF_STOP) + if mode == "boolean_reason": + reason = True + if mode == "invalid_id": + ids = [-1] + choice = fixture_choice(ids, finish, reason, lp) + if mode == "wrong_logprob_token": + choice["logprobs"]["tokens"][-1] = "token_id:19" + if mode == "missing_final_logprob": + choice["logprobs"]["token_logprobs"].pop() + if mode == "null_logprob": + choice["logprobs"]["token_logprobs"][-1] = None + if mode == "missing_ids": + choice.pop("token_ids") + if mode == "choice_index": + choice["index"] = 1 + if mode == "missing_finish": + choice["finish_reason"] = None + choice["stop_reason"] = None + if mode == "server_error": + self.send_json({"model": SELF_MODEL, "error": {"message": "injected failure"}}) return - mode = mock_mode() - if body.get("stream"): - self._stream(mode) - else: - self._completion(body, mode) - - def _completion(self, body: dict[str, Any], mode: str) -> None: - explicit = body.get("stop_token_ids") is not None - max_tokens = body.get("max_tokens", 8) - if not explicit: - control_choice: dict[str, Any] = { - "text": " word", - "index": 0, - "finish_reason": "length", - "logprobs": None, - } - if mode != "missing_token_ids": - control_choice["token_ids"] = list(range(max_tokens)) - self._json( - { - "choices": [control_choice], - "usage": { - "prompt_tokens": 5, - "completion_tokens": max_tokens, - "total_tokens": 5 + max_tokens, - }, - } - ) + model = "wrong-model" if mode == "wrong_model" else SELF_MODEL + usage = {"prompt_tokens": 1, "completion_tokens": len(ids), "total_tokens": 1 + len(ids)} + if mode == "wrong_count": + usage["completion_tokens"] += 1 + usage["total_tokens"] += 1 + body = {"model": model, "choices": [choice], "usage": usage} + if mode == "extra_choice": + body["choices"].append(fixture_choice([19], None, None, lp)) + if not payload["stream"]: + self.send_json(body) return - if mode == "eos_win" and not body.get("ignore_eos"): - self._json( - { - "choices": [ - { - "text": "", - "index": 0, - "finish_reason": "stop", - "logprobs": None, - "stop_reason": None, - "token_ids": [SELF_CHECK_EOS_ID], - } - ], - "usage": { - "prompt_tokens": 5, - "completion_tokens": 1, - "total_tokens": 6, - }, - } - ) + if payload.get("stream_options") != {"include_usage": True}: + self.send_error(400, "probe must request stream usage") return - stop_reason: Any = SELF_CHECK_STOP_ID if mode != "string_stop_reason" else "oops" - tokens = max_tokens if mode == "extra_tokens" else 1 - token_ids = list(range(tokens)) if mode == "extra_tokens" else [SELF_CHECK_STOP_ID] - if mode == "missing_trigger_logprob": - tokens = 2 - token_ids = [999, SELF_CHECK_STOP_ID] - if mode == "wrong_trigger_id": - stop_reason = 17 - token_ids = [SELF_CHECK_STOP_ID] - choice: dict[str, Any] = { - "text": " stop", - "index": 0, - "finish_reason": "stop", - "logprobs": None, + prompt = {"model": model, "choices": [{"index": 0, "text": "", "prompt_token_ids": [1]}]} + terminal = { + "model": model, + "choices": [{"index": 0, "text": "", "finish_reason": finish, "stop_reason": reason}], } - if mode != "missing_token_ids": - choice["token_ids"] = token_ids - if (body.get("logprobs") or 0) > 0: - if mode == "null_logprobs": - choice["logprobs"] = {"content": [{"token": "", "logprob": None}]} - elif mode == "missing_trigger_logprob": - # One logprob entry short of the two emitted tokens. - choice["logprobs"] = {"tokens": [""], "token_logprobs": [-0.5]} - else: - choice["logprobs"] = {"content": [{"token": "", "logprob": -0.53125}]} - choice["stop_reason"] = stop_reason - self._json( - { - "choices": [choice], - "usage": { - "prompt_tokens": 5, - "completion_tokens": tokens, - "total_tokens": 5 + tokens, - }, - } - ) - - def _stream(self, mode: str) -> None: - stop_reason: Any = SELF_CHECK_STOP_ID if mode != "string_stop_reason" else "oops" + delta = copy.deepcopy(body) + delta.pop("usage") + delta["choices"][0]["finish_reason"] = None + delta["choices"][0]["stop_reason"] = None + flush = {"model": model, "choices": [fixture_choice([], None, None, True)]} + flush["choices"][0]["text"] = "decoded text" + frames = [prompt, delta, flush, terminal] + if mode == "terminal_token": + frames = [prompt, {"model": model, "choices": [choice]}] + if mode == "terminal_extra_token": + terminal["choices"] = [fixture_choice([SELF_STOP], finish, reason, True)] + usage["completion_tokens"] += 1 + usage["total_tokens"] += 1 + if mode == "hidden_terminal_logprob": + hidden = fixture_choice([19], finish, reason, True) + hidden["token_ids"] = [] + terminal["choices"] = [hidden] + if mode == "cross_frame_missing_logprob": + delta["choices"] = [fixture_choice([7], None, None, True)] + terminal["choices"] = [fixture_choice([SELF_STOP], finish, reason, False)] + if mode == "terminal_text_only": + terminal["choices"] = [fixture_choice([], finish, reason, True)] + terminal["choices"][0]["text"] = "extra text" + if mode == "after_finish": + frames.append({"model": model, "choices": [fixture_choice([19], None, None, True)]}) + if mode != "missing_usage": + frames.append({"model": model, "choices": [], "usage": usage}) + if mode == "usage_content": + frames[-1]["token_ids"] = [19] + if mode == "sse_error": + frames.insert(-1, {"model": model, "error": {"message": "injected failure"}}) + wire = "".join(f"data: {json.dumps(frame)}\n\n" for frame in frames) + if mode != "missing_done": + wire += "data: [DONE]\n\n" + if mode == "after_done": + wire += f"data: {json.dumps(delta)}\n\n" self.send_response(200) self.send_header("Content-Type", "text/event-stream") - self.send_header("Cache-Control", "no-cache") self.end_headers() - - def emit(chunk: dict[str, Any]) -> None: - self.wfile.write(f"data: {json.dumps(chunk)}\n\n".encode()) - - first_choice: dict[str, Any] = { - "text": "", - "index": 0, - "finish_reason": None, - "token_ids": [SELF_CHECK_STOP_ID], - } - if mode != "stream_missing_trigger_logprob": - first_choice["logprobs"] = { - "tokens": [" stop"], - "token_logprobs": [-0.53125], - "top_logprobs": [], - } - emit({"id": "cmpl-1", "choices": [first_choice]}) - if mode in ("terminal_frame_extra_token", "terminal_frame_extra_token_full_vocab"): - # The finish frame smuggles a second token next to the metadata. - smuggled = [17] if mode == "terminal_frame_extra_token" else [SELF_CHECK_STOP_ID] - emit( - { - "id": "cmpl-1", - "choices": [ - { - "text": "", - "index": 0, - "finish_reason": "stop", - "logprobs": None, - "token_ids": smuggled, - } - ], - "stop_reason": stop_reason, - } - ) - else: - emit( - { - "id": "cmpl-1", - "choices": [{"text": "", "index": 0, "finish_reason": "stop", "logprobs": None}], - "stop_reason": stop_reason, - } - ) - if mode == "stream_tail": - emit( - { - "id": "cmpl-1", - "choices": [ - { - "text": "", - "index": 0, - "finish_reason": None, - "logprobs": { - "tokens": [" tail"], - "token_logprobs": [-1.0], - "top_logprobs": [], - }, - "token_ids": [17], - } - ], - } - ) - if mode != "stream_missing_done": - self.wfile.write(b"data: [DONE]\n\n") + try: + self.wfile.write(wire.encode()) + except (BrokenPipeError, ConnectionResetError): + pass def run_self_check() -> int: - server = ThreadingHTTPServer(("127.0.0.1", 0), MockHandler) - port = server.server_address[1] - thread = threading.Thread(target=server.serve_forever, daemon=True) - thread.start() - base_url = f"http://127.0.0.1:{port}" - - class Args: - prompt = DEFAULT_PROMPT - max_tokens = 8 - timeout = 10.0 - stop_token_id = None - + args = argparse.Namespace(prompt="probe", max_tokens=8, timeout=2.0, stop_token_id=None) cases = [ - ("valid", None, None, True, "valid service passes every check"), - ( - "eos_win", - SELF_CHECK_STOP_ID, - SELF_CHECK_EOS_ID, - True, - "an EOS finish must end on the configured EOS trigger", - ), - ("string_stop_reason", SELF_CHECK_STOP_ID, None, False, "string stop_reason must be rejected"), - ("null_logprobs", SELF_CHECK_STOP_ID, None, False, "null trigger logprob must be rejected"), - ("stream_tail", SELF_CHECK_STOP_ID, None, False, "content after the finish event must be rejected"), - ("wrong_model", SELF_CHECK_STOP_ID, None, False, "model-name mismatch must be rejected"), - ( - "extra_tokens", - None, - None, - False, - "extra completion tokens under a full-vocabulary stop set must be rejected", - ), - ( - "missing_trigger_logprob", - SELF_CHECK_STOP_ID, - None, - False, - "a missing final logprob must not inherit the previous token's value", - ), - ( - "wrong_trigger_id", - None, - None, - False, - "stop_reason must match the actual final token ID", - ), - ( - "terminal_frame_extra_token", - SELF_CHECK_STOP_ID, - None, - False, - "tokens sharing a frame with finish metadata must be counted", - ), - ( - "terminal_frame_extra_token_full_vocab", - None, - None, - False, - "a duplicated trigger in the finish frame must still be counted", - ), - ( - "missing_token_ids", - None, - None, - False, - "responses without token IDs cannot satisfy the trigger checks", - ), - ( - "stream_missing_done", - SELF_CHECK_STOP_ID, - None, - False, - "a stream that never sends [DONE] must be rejected", - ), - ( - "stream_missing_trigger_logprob", - SELF_CHECK_STOP_ID, - None, - False, - "the trigger frame must carry its own finite logprob", - ), + ("valid", False, False, True, None), + ("valid", True, False, True, None), + ("terminal_token", True, False, True, None), + ("valid_prefix", False, True, True, None), + ("eos", False, False, False, None), + ("earlier_stop", False, True, True, "tokens_after_stop"), + ("earlier_stop", True, True, True, "tokens_after_stop"), + ("over_budget", False, True, True, "token_budget_exceeded"), + ("over_budget", True, True, True, "token_budget_exceeded"), + ("wrong_trigger", False, False, True, "stop_reason_mismatch"), + ("wrong_trigger", True, False, True, "stop_reason_mismatch"), + ("fake_eos", False, False, False, "stop_reason_mismatch"), + ("wrong_eos_priority", False, False, False, "stop_reason_mismatch"), + ("string_reason", False, False, True, "stop_reason_type"), + ("boolean_reason", False, False, True, "stop_reason_type"), + ("missing_final_logprob", False, True, True, "logprob_count_mismatch"), + ("cross_frame_missing_logprob", True, True, True, "logprobs_missing"), + ("null_logprob", False, False, True, "logprob_value"), + ("wrong_logprob_token", False, False, True, "logprob_token_mismatch"), + ("wrong_logprob_token", True, False, True, "logprob_token_mismatch"), + ("terminal_extra_token", True, False, True, "tokens_after_stop"), + ("hidden_terminal_logprob", True, False, True, "logprob_count_mismatch"), + ("terminal_text_only", True, False, True, "terminal_text_without_tokens"), + ("after_finish", True, False, True, "stream_after_finish"), + ("usage_content", True, False, True, "stream_usage_content"), + ("server_error", False, False, True, "server_error"), + ("sse_error", True, False, True, "server_error"), + ("missing_finish", False, False, True, "finish_reason"), + ("after_done", True, False, True, "stream_after_done"), + ("missing_usage", True, False, True, "stream_missing_usage"), + ("missing_done", True, False, True, "stream_missing_done"), + ("wrong_count", False, False, True, "completion_count_mismatch"), + ("wrong_count", True, False, True, "completion_count_mismatch"), + ("extra_choice", False, False, True, "choice_count"), + ("extra_choice", True, False, True, "choice_count"), + ("choice_index", True, False, True, "choice_index"), + ("missing_ids", False, False, True, "token_ids_missing"), + ("missing_ids", True, False, True, "token_ids_missing"), + ("invalid_id", False, False, True, "token_id_range"), + ("wrong_model", False, False, True, "model_mismatch"), + ("wrong_model", True, False, True, "model_mismatch"), ] failures = 0 - for mode, stop_id, eos_id, expect_pass, label in cases: - set_mock_mode(mode) - args = Args() - args.stop_token_id = stop_id - target = run_target("qwen3_adapted", base_url, SELF_CHECK_MODEL, DEFAULT_QWEN3_VOCAB, args, eos_id) - passed = target["new_contract_passed"] - ok = passed == expect_pass - if not ok: - failures += 1 - for check_name, check_value in target["checks"].items(): - if not check_value: - print(f" unexpected check state: {check_name}=FAIL") - print( - f"[{'PASS' if ok else 'FAIL'}] self-check {mode}: " - f"overall={'PASS' if passed else 'FAIL'}, expected={'PASS' if expect_pass else 'FAIL'} ({label})" + with ThreadingHTTPServer(("127.0.0.1", 0), FixtureHandler) as server: + thread = threading.Thread(target=server.serve_forever, daemon=True) + thread.start() + url = f"http://127.0.0.1:{server.server_port}" + try: + for mode, stream, single, ignore_eos, expected in cases: + server.mode = mode + stops = [SELF_STOP] if single else list(range(SELF_VOCAB)) + result = request_completion( + url, + SELF_MODEL, + SELF_VOCAB, + args, + ignore_eos=ignore_eos, + stop_ids=stops, + logprobs=True, + stream=stream, + ) + actual = validate_sequence( + result, + set(stops), + args.max_tokens, + ignore_eos=ignore_eos, + eos_id=SELF_EOS, + require_stop=True, + ) + ok = actual == ([] if expected is None else [expected]) + failures += not ok + print( + f"[{'PASS' if ok else 'FAIL'}] {mode} ({'SSE' if stream else 'JSON'}): {actual or 'valid'}, expected={expected or 'valid'}" + ) + for mode, expected_pass, expected_health, expected_gap in ( + ("valid", True, True, False), + ("legacy", False, True, True), + ("wrong_trigger", False, True, False), + ("wrong_model", False, False, False), + ): + server.mode = mode + target = run_target("self-check", url, SELF_MODEL, SELF_VOCAB, args, SELF_EOS) + ok = ( + target["new_contract_passed"] == expected_pass + and target["healthy"] == expected_health + and target["explicit_stop_gap"] == expected_gap + ) + failures += not ok + print( + f"[{'PASS' if ok else 'FAIL'}] full probe {mode}: pass={target['new_contract_passed']}, healthy={target['healthy']}" + ) + finally: + server.shutdown() + thread.join() + return int(failures > 0) + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + for prefix, port, model, vocab in ( + ("qwen3", 18081, "qwen3-adapted", 151936), + ("qwen35", 18082, "qwen35-legacy", 248320), + ): + parser.add_argument(f"--{prefix}-url", default=f"http://127.0.0.1:{port}") + parser.add_argument(f"--{prefix}-model", default=model) + parser.add_argument(f"--{prefix}-vocab-size", type=int, default=vocab) + parser.add_argument( + f"--{prefix}-eos-token-id", + type=int, + help="Primary EOS ID resolved by the serving tokenizer (required for live validation)", ) - server.shutdown() - thread.join(timeout=5.0) - return 1 if failures else 0 + parser.add_argument("--prompt", default="The capital of France is") + parser.add_argument("--max-tokens", type=int, default=8) + parser.add_argument( + "--timeout", + type=float, + default=300.0, + help="HTTP read timeout and SSE deadline in seconds, including the tail after [DONE]", + ) + parser.add_argument( + "--stop-token-id", type=int, help="Use one known trigger instead of all vocabulary IDs" + ) + parser.add_argument( + "--out", type=Path, help="Write results and exact failure diagnostics as JSON" + ) + parser.add_argument( + "--strict-both", + action="store_true", + help="Require both healthy targets to satisfy the new contract", + ) + parser.add_argument( + "--require-legacy-gap", + action="store_true", + help="Require a healthy legacy target with a semantic stop-contract failure", + ) + parser.add_argument( + "--self-check", + action="store_true", + help="Check isolated malformed HTTP responses and exact rejection reasons; no GPU required", + ) + args = parser.parse_args() + if not args.self_check: + if args.qwen3_eos_token_id is None or args.qwen35_eos_token_id is None: + parser.error( + "live validation requires both --qwen3-eos-token-id and --qwen35-eos-token-id" + ) + if args.max_tokens <= 0 or not math.isfinite(args.timeout) or args.timeout <= 0: + parser.error("--max-tokens and --timeout must be positive and finite") + if args.strict_both and args.require_legacy_gap: + parser.error("--strict-both conflicts with --require-legacy-gap") + return args def main() -> int: args = parse_args() if args.self_check: return run_self_check() - if args.max_tokens <= 0: - print("--max-tokens must be positive", file=sys.stderr) - return 2 try: - qwen3 = run_target( - "qwen3_adapted", - args.qwen3_url, - args.qwen3_model, - args.qwen3_vocab_size, - args, - args.qwen3_eos_token_id, - ) - qwen35 = run_target( - "qwen35_legacy", - args.qwen35_url, - args.qwen35_model, - args.qwen35_vocab_size, - args, - args.qwen35_eos_token_id, - ) - except ValueError as error: + targets = { + prefix: run_target( + prefix, + getattr(args, f"{prefix}_url"), + getattr(args, f"{prefix}_model"), + getattr(args, f"{prefix}_vocab_size"), + args, + getattr(args, f"{prefix}_eos_token_id"), + ) + for prefix in ("qwen3", "qwen35") + } + except InvalidResponse as error: print(str(error), file=sys.stderr) return 2 - - comparison = { - "qwen3_new_contract_passed": qwen3["new_contract_passed"], - "qwen35_new_contract_passed": qwen35["new_contract_passed"], - "legacy_gap_observed": ( - qwen3["new_contract_passed"] and not qwen35["new_contract_passed"] - ), - } + adapted, legacy = targets["qwen3"], targets["qwen35"] + gap = adapted["new_contract_passed"] and legacy["explicit_stop_gap"] report = { - "schema_version": 2, - "config": { - "prompt": args.prompt, - "max_tokens": args.max_tokens, - "stop_mode": "single" if args.stop_token_id is not None else "full-vocabulary", - "stop_token_id": args.stop_token_id, - "qwen3_eos_token_id": args.qwen3_eos_token_id, - "qwen35_eos_token_id": args.qwen35_eos_token_id, - }, - "targets": {"qwen3_adapted": qwen3, "qwen35_legacy": qwen35}, - "comparison": comparison, + "schema_version": 3, + "config": {key: value for key, value in vars(args).items() if key != "out"}, + "targets": targets, + "legacy_gap_observed": gap, } - print_target(qwen3) - print_target(qwen35) - print_comparison(qwen3, qwen35) - print("\ncomparison:") - print(json.dumps(comparison, indent=2)) + for target in targets.values(): + print_target(target) + print(f"\nlegacy_gap_observed={gap}") if args.out: args.out.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") print(f"wrote {args.out}") - - if not qwen3["new_contract_passed"]: - return 1 - if args.strict_both and not qwen35["new_contract_passed"]: - return 1 - if args.require_legacy_gap and not comparison["legacy_gap_observed"]: - print("expected adapted-passes/legacy-fails gap was not observed", file=sys.stderr) - return 1 - return 0 + return int( + not adapted["new_contract_passed"] + or not legacy["healthy"] + or (args.strict_both and not legacy["new_contract_passed"]) + or (args.require_legacy_gap and not gap) + ) if __name__ == "__main__":