diff --git a/docs/models/qwen3/model-crate.md b/docs/models/qwen3/model-crate.md index 2f3468ffa..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`. 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; 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 239e2a61b..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,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`: 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/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..e6c68ab53 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,9 @@ 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 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 new file mode 100644 index 000000000..f7e34afca --- /dev/null +++ b/pegainfer-frontend/src/engine/stop.rs @@ -0,0 +1,133 @@ +use std::sync::Arc; + +/// 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, + /// 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. +/// +/// 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, + /// 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. + 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] + pub fn classify( + &self, + token_id: u32, + is_model_eos: impl FnOnce(u32) -> bool, + ) -> Option { + 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_primary_eos { + Some(StopCause::Eos(token_id)) + } else if self.token_ids.binary_search(&token_id).is_ok() { + 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::new(EosPolicy::Ignore, vec![99]); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Token(99)) + ); + } + + #[test] + fn normalizes_unsorted_duplicate_stop_ids() { + let policy = StopPolicy::new(EosPolicy::Ignore, vec![7, 3, 7, 1]); + + 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]); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + 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/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 48179e9cc..9c4927ea5 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,12 +651,15 @@ 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; use super::*; use crate::engine::PromptEcho; use crate::engine::RejectReason; + use crate::engine::StopPolicy; use crate::engine::TokenLogprob; use crate::engine::scheduler_pair; @@ -647,6 +667,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 +811,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..a1a75e1ed 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; @@ -77,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 a 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 { @@ -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::new( + // 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::Token), + 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 @@ -125,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)); } @@ -233,13 +257,45 @@ 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 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); } + #[test] + 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![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, |_| false), + Some(crate::engine::StopCause::Eos(99)) + ); + // Secondary model EOS IDs must keep their token stop reason. + assert_eq!( + policy.classify(100, |_| true), + 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)) + ); + assert_eq!(policy.classify(100, |_| true), None); + } + #[test] fn convert_sampling_passes_min_p_and_never_seed() { let mut params = EngineCoreSamplingParams::for_test(); @@ -264,6 +320,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-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, 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..195a39686 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -495,6 +495,7 @@ fn execute_step_on_lane( requests, kv_views, 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- @@ -502,7 +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, *sample_seed)?; + let result = + lane.execute_dflash_verify(requests, kv_views, *sample_seed, *verify_round)?; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -942,6 +944,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. @@ -1144,6 +1148,7 @@ impl Qwen3Executor { spec_decode_counters: None, dflash_ready_requests: HashSet::new(), device_ordinal, + verify_round: 0, }) } @@ -1536,6 +1541,7 @@ impl Qwen3Executor { spec_decode_counters: None, dflash_ready_requests: HashSet::new(), device_ordinal: device_ordinals[0], + verify_round: 0, }) } @@ -3427,6 +3433,7 @@ impl LocalQwen3Lane { kv_views: &[KvView], capture_layer_ids: &[usize], sample_seed: u64, + verify_round: u64, bufs: &mut VerifyGraphBuffers, ) -> Result> { let page_size = self.layout.page_size; @@ -3504,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); @@ -3532,9 +3544,32 @@ 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 mut all_results = build_verify_results(&expanded, &target_tokens)?; + let trace = std::env::var_os("PEGAINFER_TEST_LOG").is_some(); + let raw_lengths: Vec = if trace { + all_results + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + } else { + Vec::new() + }; + // 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, + ); + } 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 @@ -3544,6 +3579,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 = 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() { @@ -3563,12 +3603,15 @@ 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(), ) } .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(); } @@ -3591,7 +3634,40 @@ impl LocalQwen3Lane { &final_requests, &final_results, Some(bufs.captured_hidden()), + verify_round, )?; + if trace { + for idx in 0..requests.len() { + let has_hedge = hedge_spans + .iter() + .any(|(request_idx, _)| *request_idx == idx); + if !has_hedge { + continue; + } + 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_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_lengths[idx], + raw_b_lens, + selected, + final_results[idx].accepted_tokens.len(), + final_results[idx].matched_draft_tokens, + ); + } + } Ok(Some(VerifyResult { requests: final_results, })) @@ -3606,6 +3682,7 @@ impl LocalQwen3Lane { requests: &[VerifyStepItem], kv_views: &[KvView], sample_seed: u64, + verify_round: u64, ) -> Result { let capture_layer_ids = self.dflash_capture_layer_ids().ok_or_else(|| { anyhow::anyhow!("DFlash verify requested but no draft model is loaded") @@ -3658,6 +3735,7 @@ impl LocalQwen3Lane { kv_views, &capture_layer_ids, sample_seed, + verify_round, &mut bufs, )? { return Ok(result); @@ -3691,12 +3769,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 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 (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, &request_results, Some(bufs.captured_hidden()), + verify_round, )?; Ok(VerifyResult { requests: request_results, @@ -3805,6 +3893,7 @@ enum StepCommand { requests: Vec, kv_views: 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 488f261f0..e126b43d8 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,13 +15,39 @@ 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], +) { + 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); + // 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); +} + impl Qwen3Executor { pub(super) fn execute_speculative_verify_impl( &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" @@ -72,6 +99,7 @@ impl Qwen3Executor { requests: plan.requests.to_vec(), kv_views, sample_seed: plan.sample_seed, + verify_round, }; let outcome = match self.run_step(&step) { Ok(outcome) => outcome, @@ -107,6 +135,25 @@ impl Qwen3Executor { req.request_id )); } + // 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() + }); + if let Some(position) = terminal_position + && position + 1 != req_result.accepted_tokens.len() + { + 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. @@ -130,6 +177,14 @@ impl Qwen3Executor { req_result.request_id )); } + if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + log::debug!( + "Qwen3 DFlash commit round={} request={} accepted_len={}", + verify_round, + req_result.request_id, + req_result.accepted_tokens.len(), + ); + } applied.push(req_result.request_id); } for req_result in &result.requests { @@ -191,3 +246,25 @@ 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::new(EosPolicy::Ignore, 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); + } +} 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..b4a5be33c 100644 --- a/pegainfer-qwen3/src/frontend_adapter/tests.rs +++ b/pegainfer-qwen3/src/frontend_adapter/tests.rs @@ -79,12 +79,28 @@ 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 (tokens, _, terminal) = self.collect_terminal_with_logprobs(id); + (tokens, terminal) + } + + 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, terminal); + return (tokens, logprobs, terminal); } } } @@ -123,6 +139,62 @@ 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 = 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()); + + 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 = 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()); + + 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 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..c9c2d2d4c 100644 --- a/pegainfer-qwen3/src/scheduler/plan.rs +++ b/pegainfer-qwen3/src/scheduler/plan.rs @@ -183,26 +183,30 @@ fn build_speculative_verify_items( active: &[ActiveRequestState], draft_results: &[DraftRequestResult], ) -> 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() + let mut requests = 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, + active.stop_policy.clone(), + )); + } + requests } fn build_prefill_items(pending: &[PendingRequest], indices: &[usize]) -> Vec { @@ -253,6 +257,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 +269,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 +289,7 @@ mod tests { max_tokens, prompt_len: 10, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: None, } } @@ -313,6 +320,7 @@ mod tests { // 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!(verify[0].stop_policy, 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..eaedf6921 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -1,4 +1,6 @@ use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use super::ActiveRequestState; use super::PendingRequest; @@ -13,6 +15,14 @@ use crate::executor::ModelExecutor; use crate::executor::PrefillRequestResult; use crate::speculative::VerifyRequestResult; +fn classify_stop( + executor: &impl ModelExecutor, + policy: &StopPolicy, + token: u32, +) -> Option { + policy.classify(token, |token_id| executor.is_stop_token(token_id)) +} + pub(crate) fn resolve_step( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -44,8 +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 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. 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], @@ -60,26 +71,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) = classify_stop(executor, &req.stop_policy, 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 +139,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) = classify_stop(executor, &req.stop_policy, 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 +156,7 @@ fn resolve_prefill_outputs( token: result.first_token, logprob: result.first_token_logprob, finish_reason: FinishReason::Length, + stop_cause: None, }); continue; } @@ -152,6 +171,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 +195,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 = classify_stop(executor, &req.stop_policy, 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..76a597f6a 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; @@ -39,6 +40,7 @@ pub(crate) struct FakeExecutor { pub(crate) dropped: Arc>>, pub(crate) prefetch_offers: Arc>>, stop_token: Option, + emit_logprobs: bool, } impl FakeExecutor { @@ -56,6 +58,7 @@ impl FakeExecutor { dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, + emit_logprobs: false, } } @@ -64,6 +67,11 @@ impl FakeExecutor { self } + pub(crate) fn with_logprobs(mut self) -> Self { + self.emit_logprobs = true; + self + } + pub(crate) fn with_decode_failure(mut self) -> Self { self.fail_decode_once = true; self @@ -99,7 +107,14 @@ 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, + rank: 0, + top_logprobs: vec![(token, -0.1)], + } + }), prompt_logprobs: None, cached_tokens: 0, completed, @@ -224,7 +239,14 @@ 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, + rank: 0, + top_logprobs: vec![(token, -0.2)], + } + }), }) .collect(), }) @@ -267,6 +289,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..4ac09de6a 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::new( + if ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + 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,54 +566,54 @@ fn speculative_full_span_accept_continues() { "completion = prior generated + span len" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), } } #[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[..] { [ - 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 = StopPolicy::new(EosPolicy::Ignore, 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..a8eabd6d9 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; @@ -37,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, } } @@ -64,13 +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 scheduler still - /// owns stop-token suppression before client 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, } @@ -257,6 +265,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); @@ -271,6 +280,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); @@ -279,8 +289,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/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index ea27a5abe..f66adbd5e 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; @@ -25,6 +26,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; @@ -36,9 +38,19 @@ pub(crate) fn request( params: SamplingParams, max_tokens: usize, ) -> Request { + let stop_policy = StopPolicy::new( + if params.ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + Vec::new(), + ); + Request { prompt_tokens, params, + stop_policy, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 16e07781b..d9a1b879b 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -39,11 +39,17 @@ //! `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::path::Path; 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 +626,140 @@ 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 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(); + 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"; + 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 baseline_params = SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }; + let baseline = engine + .submit(request( + prompt_tokens.clone(), + baseline_params, + GENERATED_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() + .skip(min_index) + .take(block_size.saturating_sub(min_index)) + .filter(|(index, token)| !baseline[..*index].contains(token)) + .map(|(index, token)| (index, *token)) + .collect(); + assert!( + !candidate_stops.is_empty(), + "baseline did not produce a unique token inside the first verify span" + ); + + 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, + ..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(); + + // 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!( + 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 /// draft's `block_size` in-fill headroom (`max_pos - block_size < prompt + /// max_tokens <= max_pos`) must be rejected cleanly at admission. Before the @@ -700,15 +840,15 @@ 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 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, @@ -750,16 +890,17 @@ 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; + let mut total_wins = 0usize; 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"]) @@ -777,7 +918,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; @@ -787,23 +927,121 @@ 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 wins = nums.next().expect("win count"); + if child_test != "dflash_hedged_midspan_stop_retains_trigger" { + total_wins += wins; + } rounds += 1; } assert!( rounds > 0 && spans > 0, "child '{child_test}' executed no hedged verify round:\n{child_stderr}" ); - total_rounds += rounds; - total_spans += spans; - total_wins += wins; + 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_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 (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() { + if !line.contains("Qwen3 DFlash hedge detail ") { + continue; + } + let value = |name: &str| { + line.split_whitespace() + .find_map(|field| field.strip_prefix(name)) + }; + let selected = value("selected="); + 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(',') + .filter(|value| !value.is_empty()) + .map(|value| value.parse::().expect("raw B length")) + .collect::>() + }); + let selected_len = value("selected_len=") + .expect("selected length") + .parse::() + .expect("numeric selected length"); + if !stop_requests.contains(request) { + continue; + } + 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") + && raw_b_max.is_some_and(|raw| selected_len < raw) + { + worker_side_truncation = true; + } + } + assert!( + worker_side_truncation, + "stop hedge never truncated a candidate before winner/context commit:\n{child_stderr}" + ); + } + if child_test != "dflash_hedged_midspan_stop_retains_trigger" { + total_rounds += rounds; + total_spans += spans; + } } - 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); + assert!(total_wins > 0); + assert!(total_spans > total_wins); } 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, + .. } )); } diff --git a/scripts/qwen3_stop_contract_probe.py b/scripts/qwen3_stop_contract_probe.py new file mode 100644 index 000000000..4fec15e49 --- /dev/null +++ b/scripts/qwen3_stop_contract_probe.py @@ -0,0 +1,981 @@ +#!/usr/bin/env python3 +"""Validate the Qwen3 stop contract against a live, un-migrated Qwen3.5 server. + +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. + +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. + +Example (Qwen3-4B and Qwen3.5-0.8B): + python3 scripts/qwen3_stop_contract_probe.py \ + --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 + + +class InvalidResponse(ValueError): + def __init__(self, code: str, detail: str): + super().__init__(f"{code}: {detail}") + self.code = code + + +def require(condition: bool, code: str, detail: str) -> None: + if not condition: + raise InvalidResponse(code, detail) + + +def is_int(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +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 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 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 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") + require( + isinstance(choices, list) and len(choices) == 1, + "choice_count", + "n=1 requires exactly one choice", + ) + choice = choices[0] + 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_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 request_completion( + url: str, + model: str, + vocab_size: int, + args: argparse.Namespace, + *, + ignore_eos: bool, + stop_ids: list[int] | None = None, + logprobs: bool = False, + stream: bool = False, +) -> dict[str, Any]: + payload: dict[str, Any] = { + "model": model, + "prompt": args.prompt, + "temperature": 0.0, + "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_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( + 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.monotonic() + result: dict[str, Any] = {"http_status": None, "valid_response": False, "error_codes": []} + try: + 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 validate_sequence( + result: dict[str, Any], + 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 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( + name: str, + url: str, + model: str, + vocab_size: int, + args: argparse.Namespace, + eos_token_id: int, +) -> dict[str, Any]: + 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( + 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, + vocab_size, + args, + ignore_eos=ignore_eos, + stop_ids=ids if explicit else None, + logprobs=lp, + stream=stream, + ) + result["validation_errors"] = validate_sequence( + result, + stop_set if explicit else set(), + args.max_tokens, + ignore_eos=ignore_eos, + eos_id=eos_token_id, + require_stop=explicit, + ) + 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, + ) + ) + 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, + "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['url']} ({target['model']})") + for name, passed in target["checks"].items(): + 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']}") + + +# 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 + + +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 FixtureHandler(BaseHTTPRequestHandler): + def log_message(self, *args: Any) -> None: + pass + + 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))) + self.end_headers() + self.wfile.write(raw) + + def do_GET(self) -> None: + self.send_json({"data": [{"id": SELF_MODEL}]}) + + def do_POST(self) -> None: + 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 + 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 + 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 payload.get("stream_options") != {"include_usage": True}: + self.send_error(400, "probe must request stream usage") + return + prompt = {"model": model, "choices": [{"index": 0, "text": "", "prompt_token_ids": [1]}]} + terminal = { + "model": model, + "choices": [{"index": 0, "text": "", "finish_reason": finish, "stop_reason": reason}], + } + 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.end_headers() + try: + self.wfile.write(wire.encode()) + except (BrokenPipeError, ConnectionResetError): + pass + + +def run_self_check() -> int: + args = argparse.Namespace(prompt="probe", max_tokens=8, timeout=2.0, stop_token_id=None) + cases = [ + ("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 + 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)", + ) + 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() + try: + 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 + adapted, legacy = targets["qwen3"], targets["qwen35"] + gap = adapted["new_contract_passed"] and legacy["explicit_stop_gap"] + report = { + "schema_version": 3, + "config": {key: value for key, value in vars(args).items() if key != "out"}, + "targets": targets, + "legacy_gap_observed": gap, + } + 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}") + 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__": + raise SystemExit(main())