diff --git a/docs/models/qwen3/model-crate.md b/docs/models/qwen3/model-crate.md index 20a8b3024..03d56720a 100644 --- a/docs/models/qwen3/model-crate.md +++ b/docs/models/qwen3/model-crate.md @@ -1,7 +1,7 @@ # Qwen3-4B Model Crate **Created**: 2026-05-03 -**Last touched**: 2026-07 +**Last touched**: 2026-08 **TL;DR**: `crates/pegainfer-qwen3` now owns Qwen3 config, weights, execution, scheduler, tests, benches, and kernel plan. Root `pegainfer` loads Qwen3 through a generic `EngineHandle` and no longer contains `Qwen3Model`, `Qwen3Executor`, `ModelRuntimeConfig`, root Qwen3 tests, or `src/model/qwen3/*`. The old `ModelForward` path has been removed; decode length-limit now emits the final token before `Finished`. Long-context `bs=1` TPOT was traced to non-partition FlashInfer paged decode under-filling the GPU; Qwen3 runtime gates FlashInfer split-K decode on `padded_bs<=32` (the `seq_len>=1024` gate was dropped in #437) with a 64-token chunk floor (`Tuned` capped at 64 chunks; opt-in `--batch-invariant` pins a fixed 160-token split), cutting 4k/64 serving steady TPOT from about `11.7ms` to `6.46ms` on RTX 5090. Qwen3 now keeps a single model-crate bench entry: `qwen3_kernel_snapshot`, a JSON snapshot runner with warm/cold-L2 latency, default-on CUPTI counters, and compare. Correctness/truth is intentionally out of this snapshot for now. ## Determinism scope @@ -10,6 +10,18 @@ It also **requires `--no-prefix-cache` and rejects `--kv-offload`**. A prefix hit advances a prompt's `chunk_start` off the request-local chunk grid, re-cutting its prefill chunks according to prior traffic. `--no-prefix-cache` closes that — except under offload, where it only disables HBM retention and deliberately leaves prefix *matching* on (host-tier reuse is the point of that mode), so that pair is rejected outright. +## Stop contract + +The stepped Qwen3 scheduler keeps the request's EOS policy and explicit +`stop_token_ids` independent in `StopPolicy`. When a generated token triggers +either policy, that token is retained exactly once in the internal +`RequestUpdate`; decode results carry its real logprob when one is available, +and the terminal carries a typed `StopCause`. `StopCause::Eos` maps to +`finish_reason=stop` without a wire +`stop_reason`; an explicit request stop maps to the actual token ID. +`ignore_eos=true` disables only model EOS and does not disable explicit request +stop IDs. A length finish has no token-level stop cause. + ## Preparation - **Read**: @@ -157,7 +169,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`. The stepped stop contract now retains an EOS trigger in the internal token update and records `StopCause::Eos`; the bridge omits EOS from wire `stop_reason`. Length limit emits the sampled final token, then sends `Finished { finish_reason: Length }`. - 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 26008934a..faa686eb9 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -12,14 +12,19 @@ An engine is a set of schedulers, each a `Scheduler` implementation driven by th ``` pegainfer-frontend/src/engine/ -├── step.rs # the wire: RequestId, Request, QueuedRequest, -│ # StepOutputs { Vec }, -│ # RequestUpdate { scheduled, tokens, logprobs, cached_tokens, -│ # prompt_echo, kv_transfer, terminal }, Terminal -├── request_lifecycle.rs # submission envelope, abort control and step sender plumbing; -│ # DeferredFinish remains available for P/D handoff -├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, -│ # prompt/completion tallies, one merged update per touched id +├── step.rs # the wire: RequestId, Request, StepOutputs { Vec }, +│ # Request { ..., stop_policy }, RequestUpdate { scheduled, tokens, +│ # logprobs, cached_tokens, prompt_echo, kv_transfer, terminal }, +│ # Terminal { ..., stop_cause } +├── request_lifecycle.rs # typestate handles: QueuedRequest ─admit→ +│ # ActiveRequest ─finish/fail/defer→ consumed; every +│ # transition is by-move, a dropped handle emits Failed +│ # (drop bomb); DeferredFinish; RequestControl +├── emitter.rs # StepEmitter: the single writer of the per-step buffer; stamps +│ # timestamps, tallies prompt/completion counts, folds each +│ # request's step into one RequestUpdate; commit_step sends once +├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, +│ # prompt/completion tallies, one merged update per touched id ├── wiring.rs # scheduler_pair, SchedulerHandle (submit/take_steps/load), │ # Engine { schedulers, info, lora }, LiveScheduler, │ # EngineInfo, LaunchedEngine { Handle | Stepped } @@ -32,10 +37,12 @@ pegainfer-frontend/src/engine/ Design decisions worth knowing before touching it: - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. -- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. -- **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. -- **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. -- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. +- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. This is what makes `defer_finish` safe: a P/D prefill executor can withhold a request's `Finished` until its KV saves are peer-visible and send it later from any thread — the deferred message carries the request's entire buffered update, so late delivery cannot reorder. +- **Independent stop policy and cause.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in `RequestUpdate` and carries `Terminal::Finished.stop_cause`; decode paths include its real logprob when available, so the stepped bridge can report the actual explicit stop ID without reconstructing it. `ignore_eos` affects only model EOS. Legacy producers may leave the cause empty while they are migrated individually. +- **Typestate lifecycle.** `QueuedRequest` (queued) and `ActiveRequest` (streaming) are owned tokens; admit/reject/retire/finish/fail consume them, so "terminal exactly once, nothing after it" cannot be miscoded — it does not compile. A handle dropped without a transition emits `Failed` from its `Drop`, which is also how a crashed scheduler answers every in-flight request: the driver drops the scheduler, the handles fall, the terminals ship. +- **Emitter as single writer.** Schedulers never touch the channel; they call `StepEmitter` methods against their handles. The emitter stamps `ScheduledInfo` at admission, tallies token counts (terminal counts derive from the tally, never from model-side arithmetic), and `commit_step` publishes the whole step in one send. +- **Pure polling driver.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish load, commit. No idle/park distinction — the scheduler owns the GPU and spinning on it costs nothing anyone else could use; async KV I/O (prefetch, decode-overlap prefill) is naturally absorbed by polling. An idle iteration ends in a `spin_loop` hint (relaxes the core's issue slots, no latency cost — busy iterations never pause). The loop exits when the frontend drops the handle and the queue drains. +- **Gemma 4 async prefill exception.** While asynchronous prefill is the only remaining work, Gemma 4 drains and joins that lane rather than hot-polling its completion; decode or queued work keeps the normal polling path. - **Abort is a flag, not channel teardown.** `SchedulerHandle::submit` returns a `RequestControl`; the frontend flips its boolean abort flag and the scheduler retires the request silently on its next touch (no terminal — the frontend already dropped its state for that id). - **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::load()` — pull-only by design, "notify me on load change" is deliberately unrepresentable (the driver busy-polls, so a subscription edge would fire per spin). All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit. - **Control plane lives outside the contract.** `Scheduler` has no control method and the contract carries no control channel. A capability like LoRA is a private channel the model crate mints *before* `spawn_scheduler` — the scheduler closes over the receiver, the `LoraClient` sender surfaces as `Engine.lora: Option`, and the `Option` *is* the capability (no `bool` flag, no registry until a second capability exists). The vocabulary (`LoraControl`, `LoraClient`) is still defined in the frontend crate because the frontend must speak it without holding model structs; only the wiring is the model's business. @@ -91,7 +98,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Protocol stacks -**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` appends the stop sentinel token, which is how usage keeps counting the suppressed EOS). +**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a typed `StopCause` reports the actual request stop token, while the synthetic sentinel remains only as a compatibility fallback for producers that provide no cause). **`dynamo` (planned second stack).** dynamo's `lib/llm` in-process path removes the wire protocol entirely (`EngineConfig::InProcessTokens` + `run_input`). The step contract was shaped so this stack can consume `StepOutputs` directly without impersonation overhead. Decision gate: prototype, A/B against the vllm stack, let TTFT/step-overhead numbers pick the default. diff --git a/pegainfer-frontend/src/engine/driver.rs b/pegainfer-frontend/src/engine/driver.rs index 827adeb07..c8c4e97bf 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 34b4f16d6..08806c001 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 50fa3df58..01d164382 100644 --- a/pegainfer-frontend/src/engine/step.rs +++ b/pegainfer-frontend/src/engine/step.rs @@ -14,6 +14,8 @@ use std::time::Instant; use super::event::FinishReason; use super::event::TokenLogprob; +use super::stop::StopCause; +use super::stop::StopPolicy; /// In-process routing id for one generate request, minted by /// [`super::SchedulerHandle::submit`] from a per-scheduler counter. `Copy` and @@ -49,6 +51,7 @@ impl std::fmt::Display for RequestId { pub struct Request { pub prompt_tokens: Vec, pub params: crate::sampler::SamplingParams, + pub stop_policy: StopPolicy, pub max_tokens: usize, pub lora_adapter: Option, /// Opaque router/P-D metadata from the request's @@ -233,6 +236,10 @@ impl fmt::Display for RejectReason { pub enum Terminal { Finished { reason: FinishReason, + /// Present for token-driven stop finishes. The triggering token remains + /// in `RequestUpdate.tokens`, with its real logprob in the matching + /// `RequestUpdate.logprobs` entry. + stop_cause: Option, prompt_tokens: usize, completion_tokens: usize, }, diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs new file mode 100644 index 000000000..099b1372b --- /dev/null +++ b/pegainfer-frontend/src/engine/stop.rs @@ -0,0 +1,94 @@ +/// How a request treats end-of-sequence tokens. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum EosPolicy { + /// Do not stop on model EOS tokens. + Ignore, + /// Use the model executor's configured EOS set. + #[default] + ModelDefault, + /// Stop only on this protocol-provided primary EOS token. + Token(u32), +} + +/// Request-scoped token stopping policy. +/// +/// EOS is kept separate from caller stop tokens because the vLLM protocol +/// reports them differently: EOS has no 'stop_reason', while a request stop +/// reports the actual matching token ID. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct StopPolicy { + pub eos: EosPolicy, + pub token_ids: Vec, +} + +impl StopPolicy { + /// Classify a token using vLLM's priority: EOS first, then the request's + /// explicit stop-token set. + #[must_use] + pub fn classify( + &self, + token_id: u32, + is_model_eos: impl FnOnce(u32) -> bool, + ) -> Option { + let is_eos = match self.eos { + EosPolicy::Ignore => false, + EosPolicy::ModelDefault => is_model_eos(token_id), + EosPolicy::Token(eos_token_id) => token_id == eos_token_id, + }; + + if is_eos { + Some(StopCause::Eos(token_id)) + } else if self.token_ids.contains(&token_id) { + Some(StopCause::Token(token_id)) + } else { + None + } + } +} + +/// The token-level cause of a normal stop finish. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StopCause { + /// A primary or model-default EOS token. + Eos(u32), + /// A token from the request's explicit stop-token set. + Token(u32), +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_default_classifies_model_eos() { + let policy = StopPolicy::default(); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Eos(99)) + ); + } + + #[test] + fn ignored_eos_does_not_disable_an_explicit_stop() { + let policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![99], + }; + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Token(99)) + ); + } + + #[test] + fn eos_wins_when_the_same_id_is_also_an_explicit_stop() { + let policy = StopPolicy { + eos: EosPolicy::Token(99), + token_ids: vec![99], + }; + + assert_eq!(policy.classify(99, |_| false), Some(StopCause::Eos(99))); + } +} diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 0a1d037cc..731fbecfe 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::to_wire_position_logprobs; @@ -362,6 +364,7 @@ impl SteppedEngineBridge { None, ); } + let lora_adapter = match lora_adapter_from_sampling_params(&sampling_params) { Ok(adapter) => adapter, Err(error) => { @@ -383,6 +386,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, @@ -398,9 +405,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, @@ -435,8 +447,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. @@ -539,11 +552,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); @@ -551,6 +564,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; } @@ -616,14 +633,19 @@ impl UnixAnchor { #[cfg(test)] mod tests { + use std::sync::atomic::AtomicBool; + use super::*; use crate::engine::RejectReason; + use crate::engine::StopPolicy; + use crate::engine::TokenLogprob; use crate::engine::scheduler_pair; fn request() -> Request { 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, @@ -662,4 +684,93 @@ mod tests { "a request refused while queued did no prefill" ); } + + #[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 { + 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 fbc5a9f74..6e31d6e53 100644 --- a/pegainfer-frontend/src/vllm/wire.rs +++ b/pegainfer-frontend/src/vllm/wire.rs @@ -5,7 +5,9 @@ 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::StopPolicy; use crate::engine::TokenLogprob; use crate::sampler::SamplingParams; @@ -44,8 +46,8 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar // None`, but `_all_stop_token_ids` always carries the model EOS set (it // exists for min_tokens masking, not stop detection). Deriving ignore_eos // from all_stop_token_ids would therefore void every ignore_eos request on - // models with a real EOS. Only `_eos_token_id` and the client's explicit - // `stop_token_ids` express a stop intent. + // models with a real EOS. Only _eos_token_id and the client's explicit + // stop_token_ids express the legacy scheduler's stop intent. let ignore_eos = params.eos_token_id.is_none() && params.stop_token_ids.is_empty(); if params.temperature <= 0.0 { return SamplingParams { @@ -76,7 +78,23 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar } } -/// Reject request parameters the engine would otherwise silently ignore. +pub(crate) fn convert_stop_policy(params: &EngineCoreSamplingParams) -> StopPolicy { + StopPolicy { + // Qwen3 owns the complete model EOS set in generation_config. The + // protocol's optional primary ID only tells us whether EOS is active; + // using it as a singleton would miss secondary model EOS IDs. + eos: params + .eos_token_id + .map_or(EosPolicy::Ignore, |_| EosPolicy::ModelDefault), + token_ids: params.stop_token_ids.clone(), + } +} + +/// Reject request parameters the frontend cannot represent faithfully. +/// +/// The stepped contract carries explicit request stop IDs independently in +/// [`StopPolicy`]; this helper only validates unrelated sampling/transfer +/// fields that would otherwise be silently ignored. /// Returns the offending description; `None` means the request is servable. /// /// The float comparisons are exact on purpose: they detect "the client sent @@ -191,13 +209,23 @@ mod tests { params.eos_token_id = Some(163_586); assert!(!convert_sampling(¶ms).ignore_eos); - // Explicit client stop tokens keep EOS detection on even when the - // frontend dropped _eos_token_id. + // The legacy scheduler keeps EOS active when an explicit stop token is + // present; the stepped bridge carries explicit stops in StopPolicy. params.eos_token_id = None; params.stop_token_ids = vec![42]; assert!(!convert_sampling(¶ms).ignore_eos); } + #[test] + fn convert_stop_policy_keeps_eos_and_explicit_stops_independent() { + let mut params = EngineCoreSamplingParams::for_test(); + params.eos_token_id = Some(99); + params.stop_token_ids = vec![11]; + + assert_eq!(convert_stop_policy(¶ms).eos, EosPolicy::ModelDefault); + assert_eq!(convert_stop_policy(¶ms).token_ids, vec![11]); + } + #[test] fn convert_sampling_passes_min_p_and_never_seed() { let mut params = EngineCoreSamplingParams::for_test(); diff --git a/pegainfer-gemma4/src/engine/lane_tests.rs b/pegainfer-gemma4/src/engine/lane_tests.rs index 1a5e145db..1102ddec6 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 ac1c0b0b8..afd787841 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 086afc75b..f42ec5959 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -19,6 +19,7 @@ use pegainfer_core::weight_loader::load_shard_info; use pegainfer_frontend::engine::DeferredFinish; use pegainfer_frontend::engine::LoadLoraAdapterRequest; use pegainfer_frontend::engine::SpecDecodeCounters; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::engine::panic_message; @@ -594,6 +595,7 @@ fn execute_step_on_lane( StepCommand::SpeculativeVerify { requests, kv_views, + stop_policies, sample_seed, } => { // One target forward over each request's K+1 draft span with a @@ -602,7 +604,8 @@ fn execute_step_on_lane( // token at each span position) and captures the target hidden states // (at the DFlash layers) to seed the next draft — all into reused, // pointer-stable scratch (`VerifyGraphBuffers`). - let result = lane.execute_dflash_verify(requests, kv_views, *sample_seed)?; + let result = + lane.execute_dflash_verify(requests, kv_views, stop_policies, *sample_seed)?; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -3705,8 +3708,15 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], + stop_policies: &[StopPolicy], sample_seed: u64, ) -> Result { + anyhow::ensure!( + stop_policies.len() == requests.len(), + "DFlash verify received {} stop policies for {} requests", + stop_policies.len(), + requests.len() + ); let capture_layer_ids = self.dflash_capture_layer_ids().ok_or_else(|| { anyhow::anyhow!("DFlash verify requested but no draft model is loaded") })?; @@ -3791,8 +3801,13 @@ impl LocalQwen3Lane { .flat_map(|req| std::iter::repeat_n(&req.params, req.as_slice().len())) .collect(); let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, sample_seed)?; - let request_results = build_verify_results(requests, &target_tokens)?; - + let mut request_results = build_verify_results(requests, &target_tokens)?; + // Apply the request policy before recording target hidden states; + // otherwise a suffix discarded by terminal handling would leak + // into the next DFlash draft context and acceptance counters. + for (policy, result) in stop_policies.iter().zip(&mut request_results) { + spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + } self.record_verify_dflash_context( requests, &request_results, @@ -3905,6 +3920,9 @@ enum StepCommand { SpeculativeVerify { requests: Vec, kv_views: Vec, + /// Request-local stop policies used before DFlash context recording; + /// these are host metadata and never enter the GPU batch. + stop_policies: Vec, sample_seed: u64, }, /// Speculative draft: roll the DFlash draft model forward one block per diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 488f261f0..a55ef96cb 100644 --- a/pegainfer-qwen3/src/executor/spec.rs +++ b/pegainfer-qwen3/src/executor/spec.rs @@ -6,6 +6,7 @@ //! module owns only the KV bookkeeping the executor thread is responsible for. use anyhow::Result; +use pegainfer_frontend::engine::StopPolicy; use super::Qwen3Executor; use super::RequestId; @@ -14,8 +15,31 @@ use super::WorkerStepOutcome; use crate::speculative::DraftPlan; use crate::speculative::DraftResult; use crate::speculative::VerifyPlan; +use crate::speculative::VerifyRequestResult; use crate::speculative::VerifyResult; +/// Remove accepted tokens after the first request-terminal token before the +/// speculative KV transaction commits. The scheduler classifies the same +/// retained trigger later to produce the typed protocol stop cause. +pub(super) fn truncate_after_terminal( + result: &mut VerifyRequestResult, + policy: &StopPolicy, + model_eos: &[u32], +) { + // Keep this helper idempotent: the worker trims before DFlash context is + // recorded, and the executor repeats the invariant before KV commit. + let Some(keep) = result.accepted_tokens.iter().position(|&token| { + policy + .classify(token, |id| model_eos.contains(&id)) + .is_some() + }) else { + return; + }; + let keep = keep + 1; + result.accepted_tokens.truncate(keep); + result.matched_draft_tokens = result.matched_draft_tokens.min(keep); +} + impl Qwen3Executor { pub(super) fn execute_speculative_verify_impl( &mut self, @@ -25,6 +49,12 @@ impl Qwen3Executor { self.speculative.is_some(), "speculative verification requested but no draft model is loaded" ); + anyhow::ensure!( + plan.stop_policies.len() == plan.requests.len(), + "speculative verify received {} stop policies for {} requests", + plan.stop_policies.len(), + plan.requests.len() + ); for req in plan.requests { anyhow::ensure!( !req.as_slice().is_empty(), @@ -71,6 +101,7 @@ impl Qwen3Executor { let step = StepCommand::SpeculativeVerify { requests: plan.requests.to_vec(), kv_views, + stop_policies: plan.stop_policies.to_vec(), sample_seed: plan.sample_seed, }; let outcome = match self.run_step(&step) { @@ -80,7 +111,7 @@ impl Qwen3Executor { return Err(e); } }; - let result = match outcome { + let mut result = match outcome { WorkerStepOutcome::SpeculativeVerify(result) => result, other => { self.revert_speculative_schedules(&scheduled); @@ -108,6 +139,12 @@ impl Qwen3Executor { )); } } + // The worker returns the mathematically accepted span. Apply the + // request contract before touching RequestKv so a terminal token's + // speculative suffix is rolled back with the unused reservation. + for (policy, req_result) in plan.stop_policies.iter().zip(&mut result.requests) { + truncate_after_terminal(req_result, policy, &self.metadata.stop_token_ids); + } // Commit the accepted prefix of each request's KV and free the rest. // On a mid-loop failure, only the not-yet-applied requests roll back @@ -191,3 +228,42 @@ impl Qwen3Executor { } } } + +#[cfg(test)] +mod tests { + use pegainfer_frontend::engine::EosPolicy; + + use super::*; + + #[test] + fn explicit_stop_truncates_the_kv_commit_after_the_trigger() { + let policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![7], + }; + let mut result = VerifyRequestResult { + request_id: RequestId::new(1), + matched_draft_tokens: 3, + accepted_tokens: vec![5, 7, 8, 9], + }; + + truncate_after_terminal(&mut result, &policy, &[99]); + + assert_eq!(result.accepted_tokens, vec![5, 7]); + assert_eq!(result.matched_draft_tokens, 2); + } + + #[test] + fn model_eos_truncates_but_keeps_the_trigger() { + let mut result = VerifyRequestResult { + request_id: RequestId::new(2), + matched_draft_tokens: 3, + accepted_tokens: vec![5, 99, 8, 9], + }; + + truncate_after_terminal(&mut result, &StopPolicy::default(), &[99]); + + assert_eq!(result.accepted_tokens, vec![5, 99]); + assert_eq!(result.matched_draft_tokens, 2); + } +} diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index e02125ea5..c89d8aec4 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -42,6 +42,7 @@ use pegainfer_frontend::engine::RejectReason; use pegainfer_frontend::engine::RequestLedger; use pegainfer_frontend::engine::Scheduler; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::engine::spawn_scheduler; use pegainfer_kernels::ops::NumericPolicy; use pegainfer_kernels::ops::numeric_policy; @@ -343,7 +344,7 @@ impl Qwen3Scheduler { // terminal rides the committed step, which the driver ships after // publishing metrics — the finishing batch's send-time stats then // read the drained occupancy instead of racing the publish. - let mut finishes: Vec<(RequestId, FinishReason)> = Vec::new(); + let mut finishes: Vec<(RequestId, FinishReason, Option)> = Vec::new(); for cached in effects.cached { if ledger.is_active(cached.request_id) { @@ -367,32 +368,11 @@ impl Qwen3Scheduler { for effect in effects.decode { match effect { DecodeEffect::Finish { - request_id, - finish_reason, - } => { - let Some(index) = self - .active - .iter() - .position(|req| req.request_id == request_id) - else { - continue; - }; - if ledger.is_active(request_id) { - if ledger.is_aborted(request_id) { - ledger.retire(request_id); - } else { - finishes.push((request_id, finish_reason)); - } - } - self.tracker.finish(request_id); - let _ = self.executor.drop_request(request_id); - to_retire.push(index); - } - DecodeEffect::EmitAndFinish { request_id, token, logprob, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -406,14 +386,14 @@ impl Qwen3Scheduler { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &[token], &[logprob]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); let _ = self.executor.drop_request(request_id); to_retire.push(index); } - DecodeEffect::EmitAndContinue { + DecodeEffect::Continue { request_id, token, logprob, @@ -436,7 +416,7 @@ impl Qwen3Scheduler { req.generated_count = completion_tokens; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id, tokens, completion_tokens, @@ -460,10 +440,11 @@ impl Qwen3Scheduler { } } } - DecodeEffect::EmitManyAndFinish { + DecodeEffect::FinishMany { request_id, tokens, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -477,7 +458,7 @@ impl Qwen3Scheduler { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &tokens, &[]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); @@ -507,32 +488,19 @@ impl Qwen3Scheduler { continued.push(req); } } - PendingEffect::Finish { - request_id, - finish_reason, - } => { - if ledger.is_active(request_id) { - if ledger.is_aborted(request_id) { - ledger.retire(request_id); - } else { - finishes.push((request_id, finish_reason)); - } - } - self.tracker.finish(request_id); - let _ = self.executor.drop_request(request_id); - } PendingEffect::EmitAndFinish { request_id, token, logprob, finish_reason, + stop_cause, } => { if ledger.is_active(request_id) { if ledger.is_aborted(request_id) { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &[token], &[logprob]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); @@ -560,12 +528,14 @@ impl Qwen3Scheduler { if self.executor.withholds_finishes() { let withheld: Vec = finishes .into_iter() - .map(|(request_id, reason)| ledger.defer_finish(request_id, reason)) + .map(|(request_id, reason, stop_cause)| { + ledger.defer_finish_with_cause(request_id, reason, stop_cause) + }) .collect(); self.executor.release_finished_events(withheld); } else { - for (request_id, reason) in finishes { - ledger.finish(request_id, reason); + for (request_id, reason, stop_cause) in finishes { + ledger.finish_with_cause(request_id, reason, stop_cause); } } } diff --git a/pegainfer-qwen3/src/frontend_adapter/tests.rs b/pegainfer-qwen3/src/frontend_adapter/tests.rs index 35441bc7d..0a89aac2b 100644 --- a/pegainfer-qwen3/src/frontend_adapter/tests.rs +++ b/pegainfer-qwen3/src/frontend_adapter/tests.rs @@ -89,6 +89,28 @@ impl StepCollector { } } + fn collect_terminal_with_logprobs( + &mut self, + id: RequestId, + ) -> ( + Vec, + Vec>, + Terminal, + ) { + let mut tokens = Vec::new(); + let mut logprobs = Vec::new(); + + loop { + let update = self.next_for(id); + tokens.extend_from_slice(&update.tokens); + logprobs.extend(update.logprobs); + + if let Some(terminal) = update.terminal { + return (tokens, logprobs, terminal); + } + } + } + /// Drain the remaining stream (until the scheduler is gone) and return /// every terminal seen for `id`. For asserting silence after an abort. fn drain_terminals_for(&mut self, id: RequestId) -> Vec { @@ -123,6 +145,105 @@ fn wait_until(timeout: Duration, mut predicate: impl FnMut() -> bool) -> bool { false } +#[test] +fn request_stop_token_beats_length_during_prefill_when_eos_is_ignored() { + let executor = FakeExecutor::new(4, Arc::new(Mutex::new(Vec::new()))).with_logprobs(); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 1); + req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; + req.stop_policy.token_ids = vec![100]; + + let control = partition.handle.submit(req); + let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); + + assert_eq!(tokens, vec![100]); + assert_eq!(logprobs.len(), 1); + assert!((logprobs[0].as_ref().expect("stop-token logprob").logprob + 0.1).abs() < f32::EPSILON); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(100)), + completion_tokens: 1, + .. + } + )); +} + +#[test] +fn request_stop_token_beats_length_during_decode_when_eos_is_ignored() { + let executor = FakeExecutor::new(4, Arc::new(Mutex::new(Vec::new()))).with_logprobs(); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 2); + req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; + req.stop_policy.token_ids = vec![200]; + + let control = partition.handle.submit(req); + let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); + + assert_eq!(tokens, vec![100, 200]); + assert_eq!(logprobs.len(), 2); + assert!((logprobs[1].as_ref().expect("stop-token logprob").logprob + 0.2).abs() < f32::EPSILON); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(200)), + completion_tokens: 2, + .. + } + )); +} + +#[test] +fn speculative_request_stop_beats_length_midspan_and_cleans_up() { + let dropped = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(4, Arc::clone(&dropped)) + .with_stop_token(12) + .with_speculative_accepted_tokens(&[10, 11, 12, 13]); + + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 4); + req.stop_policy.eos = pegainfer_frontend::engine::EosPolicy::Ignore; + req.stop_policy.token_ids = vec![12]; + + let control = partition.handle.submit(req); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!( + tokens, + vec![100, 10, 11, 12], + "the trigger is retained and the accepted suffix is discarded" + ); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(12)), + completion_tokens: 4, + .. + } + )); + + assert!( + wait_until(Duration::from_secs(1), || { + dropped.lock().unwrap().contains(&0) + }), + "stopped speculative request state should be dropped" + ); + + assert!( + wait_until(Duration::from_secs(1), || { + let metrics = partition.handle.metrics(); + metrics.num_running_reqs == 0 && metrics.kv_used_blocks == 0 + }), + "stopped speculative request should release scheduler state and KV blocks" + ); +} + #[test] fn unknown_lora_request_is_rejected_without_blocking_base_request() { let dropped = Arc::new(Mutex::new(Vec::new())); diff --git a/pegainfer-qwen3/src/scheduler.rs b/pegainfer-qwen3/src/scheduler.rs index dcd3da28d..5ba9e0231 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, /// Number of top logprobs to return (0 = disabled). pub(crate) logprobs: usize, } @@ -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: usize, pub(crate) echo: bool, @@ -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, echo: req.echo, 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 c8ca7d32d..414eee610 100644 --- a/pegainfer-qwen3/src/scheduler/plan.rs +++ b/pegainfer-qwen3/src/scheduler/plan.rs @@ -119,9 +119,11 @@ pub(crate) fn execute_plan( requests: &draft_requests, })?; draft.requests.sort_by_key(|result| result.request_id); - let verify_requests = build_speculative_verify_items(active, &draft.requests); + let (verify_requests, stop_policies) = + build_speculative_verify_items(active, &draft.requests); let mut verify = executor.execute_speculative_verify(VerifyPlan { requests: &verify_requests, + stop_policies: &stop_policies, sample_seed: rand::RngExt::random(rng), })?; verify.requests.sort_by_key(|result| result.request_id); @@ -183,27 +185,35 @@ fn build_speculative_draft_items(active: &[ActiveRequestState]) -> Vec Vec { - draft_results - .iter() - .map(|draft| { - let active = active - .iter() - .find(|req| req.request_id == draft.request_id) - .expect("draft request_id must exist in active set"); - // Clamp the verify span to the request's remaining output budget so - // a long accepted run can't overshoot max_tokens. - let remaining = active.max_tokens.saturating_sub(active.generated_count); - // A continuing active request always has budget left (resolve emits - // EmitManyAndFinish the moment generated_count hits max_tokens), so - // this is a true invariant, not a runtime condition — don't crash the - // scheduler thread in release on a state we've proven unreachable. - debug_assert!(remaining > 0, "active request must have output budget"); - let mut token_ids = draft.token_ids.clone(); - token_ids.truncate(remaining); - VerifyStepItem::new(draft.request_id, token_ids, active.params) - }) - .collect() +) -> ( + Vec, + Vec, +) { + let mut requests = Vec::with_capacity(draft_results.len()); + let mut stop_policies = Vec::with_capacity(draft_results.len()); + for draft in draft_results { + let active = active + .iter() + .find(|req| req.request_id == draft.request_id) + .expect("draft request_id must exist in active set"); + // Clamp the verify span to the request's remaining output budget so + // a long accepted run can't overshoot max_tokens. + let remaining = active.max_tokens.saturating_sub(active.generated_count); + // A continuing active request always has budget left (resolve emits + // FinishMany the moment generated_count hits max_tokens), so + // this is a true invariant, not a runtime condition — don't crash the + // scheduler thread in release on a state we've proven unreachable. + debug_assert!(remaining > 0, "active request must have output budget"); + let mut token_ids = draft.token_ids.clone(); + token_ids.truncate(remaining); + requests.push(VerifyStepItem::new( + draft.request_id, + token_ids, + active.params, + )); + stop_policies.push(active.stop_policy.clone()); + } + (requests, stop_policies) } fn build_prefill_items(pending: &[PendingRequest], indices: &[usize]) -> Vec { @@ -254,6 +264,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::*; @@ -265,6 +276,7 @@ mod tests { lora_adapter: None, prompt_tokens: vec![1, 2, 3], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 8, logprobs: 0, echo: false, @@ -284,6 +296,7 @@ mod tests { max_tokens, prompt_len: 10, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: 0, } } @@ -308,12 +321,13 @@ mod tests { token_ids: (0..16).collect(), }; - let verify = build_speculative_verify_items(&active, &[draft]); + let (verify, stop_policies) = build_speculative_verify_items(&active, &[draft]); assert_eq!(verify.len(), 1); // 32 - 24 = 8 remaining → the 16-token span truncates to 8. assert_eq!(verify[0].as_slice().len(), 8); assert_eq!(verify[0].as_slice(), (0..8).collect::>()); + assert_eq!(stop_policies, vec![StopPolicy::default()]); } // The plan selector is the whole batch-formation policy: what the scheduler diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 12356210f..4d748bba4 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -1,4 +1,5 @@ use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; use super::ActiveRequestState; use super::PendingRequest; @@ -13,6 +14,24 @@ use crate::executor::ModelExecutor; use crate::executor::PrefillRequestResult; use crate::speculative::VerifyRequestResult; +fn stop_cause( + executor: &impl ModelExecutor, + req: &ActiveRequestState, + token: u32, +) -> Option { + req.stop_policy + .classify(token, |token_id| executor.is_stop_token(token_id)) +} + +fn pending_stop_cause( + executor: &impl ModelExecutor, + req: &PendingRequest, + token: u32, +) -> Option { + req.stop_policy + .classify(token, |token_id| executor.is_stop_token(token_id)) +} + pub(crate) fn resolve_step( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -44,8 +63,9 @@ pub(crate) fn resolve_step( /// Turn each request's accepted speculative span into a decode effect. A span /// commits 1..=K+1 tokens at once; we walk it in order so a stop token or the -/// max-output budget truncates exactly where it lands (the executor already -/// suppressed nothing — stop handling lives here, mirroring single-token decode). +/// max-output budget lands exactly where expected. The executor has already +/// truncated any suffix after a request-terminal token to keep speculative +/// state consistent; the resolver classifies its typed cause here. pub(crate) fn resolve_speculative_outputs( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -60,26 +80,30 @@ pub(crate) fn resolve_speculative_outputs( .expect("speculative request_id must exist in active set"); let mut emitted = Vec::new(); let mut completion_tokens = req.generated_count; + for &token in &result.accepted_tokens { completion_tokens += 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(token); - if is_eos { - return DecodeEffect::EmitManyAndFinish { + emitted.push(token); + + if let Some(stop_cause) = stop_cause(executor, req, token) { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }; } - emitted.push(token); + if completion_tokens >= req.max_tokens { - return DecodeEffect::EmitManyAndFinish { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Length, + stop_cause: None, }; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id: result.request_id, tokens: emitted, completion_tokens, @@ -126,10 +150,13 @@ fn resolve_prefill_outputs( }); } - if !req.params.ignore_eos && executor.is_stop_token(result.first_token) { - effects.pending.push(PendingEffect::Finish { + if let Some(stop_cause) = pending_stop_cause(executor, &req, result.first_token) { + effects.pending.push(PendingEffect::EmitAndFinish { request_id: req.request_id, + token: result.first_token, + logprob: result.first_token_logprob, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }); continue; } @@ -140,6 +167,7 @@ fn resolve_prefill_outputs( token: result.first_token, logprob: result.first_token_logprob, finish_reason: FinishReason::Length, + stop_cause: None, }); continue; } @@ -154,6 +182,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, @@ -177,22 +206,27 @@ fn resolve_decode_outputs( .find(|req| req.request_id == result.request_id) .expect("decode request_id must exist in active set"); let completion_tokens = req.generated_count + 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(result.token); + let stop_cause = stop_cause(executor, req, result.token); let at_limit = completion_tokens >= req.max_tokens; - if is_eos { + + if let Some(stop_cause) = stop_cause { DecodeEffect::Finish { request_id: result.request_id, + token: result.token, + logprob: result.logprob.clone(), finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), } } else if at_limit { - DecodeEffect::EmitAndFinish { + DecodeEffect::Finish { request_id: result.request_id, token: result.token, logprob: result.logprob.clone(), finish_reason: FinishReason::Length, + stop_cause: None, } } else { - DecodeEffect::EmitAndContinue { + DecodeEffect::Continue { request_id: result.request_id, token: result.token, logprob: result.logprob.clone(), diff --git a/pegainfer-qwen3/src/scheduler/test_support.rs b/pegainfer-qwen3/src/scheduler/test_support.rs index 16d19323a..00d99a00a 100644 --- a/pegainfer-qwen3/src/scheduler/test_support.rs +++ b/pegainfer-qwen3/src/scheduler/test_support.rs @@ -9,6 +9,7 @@ use std::time::Duration; use anyhow::Result; use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::sampler::SamplingParams; @@ -23,6 +24,12 @@ use crate::executor::PrefillStepItem; use crate::executor::RequestId; use crate::executor::UnifiedPlan; use crate::executor::UnifiedResult; +use crate::speculative::DraftPlan; +use crate::speculative::DraftRequestResult; +use crate::speculative::DraftResult; +use crate::speculative::VerifyPlan; +use crate::speculative::VerifyRequestResult; +use crate::speculative::VerifyResult; pub(crate) struct FakeExecutor { pub(crate) block_size: usize, @@ -39,6 +46,8 @@ pub(crate) struct FakeExecutor { pub(crate) dropped: Arc>>, pub(crate) prefetch_offers: Arc>>, stop_token: Option, + emit_logprobs: bool, + speculative_accepted_tokens: Option>, } impl FakeExecutor { @@ -56,6 +65,8 @@ impl FakeExecutor { dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, + emit_logprobs: false, + speculative_accepted_tokens: None, } } @@ -64,6 +75,20 @@ impl FakeExecutor { self } + pub(crate) fn with_logprobs(mut self) -> Self { + self.emit_logprobs = true; + self + } + + pub(crate) fn with_speculative_accepted_tokens(mut self, tokens: &[u32]) -> Self { + assert!( + !tokens.is_empty(), + "fake speculative span must make progress" + ); + self.speculative_accepted_tokens = Some(tokens.to_vec()); + self + } + pub(crate) fn with_decode_failure(mut self) -> Self { self.fail_decode_once = true; self @@ -99,7 +124,13 @@ impl FakeExecutor { PrefillRequestResult { request_id: req.request_id, first_token: 100 + req.request_id.raw() as u32, - first_token_logprob: None, + first_token_logprob: self.emit_logprobs.then(|| { + let token = 100 + req.request_id.raw() as u32; + pegainfer_frontend::engine::TokenLogprob { + logprob: -0.1, + top_logprobs: vec![(token, -0.1)], + } + }), prompt_logprobs: None, cached_tokens: 0, completed, @@ -224,7 +255,13 @@ impl ModelExecutor for FakeExecutor { .map(|req| DecodeRequestResult { request_id: req.request_id, token: 200 + req.request_id.raw() as u32, - logprob: None, + logprob: self.emit_logprobs.then(|| { + let token = 200 + req.request_id.raw() as u32; + pegainfer_frontend::engine::TokenLogprob { + logprob: -0.2, + top_logprobs: vec![(token, -0.2)], + } + }), }) .collect(), }) @@ -260,6 +297,75 @@ impl ModelExecutor for FakeExecutor { .collect(), }) } + + fn execute_speculative_draft(&mut self, plan: DraftPlan<'_>) -> Result { + let accepted_tokens = self + .speculative_accepted_tokens + .as_ref() + .ok_or_else(|| anyhow::anyhow!("fake speculative decoding is disabled"))?; + + Ok(DraftResult { + requests: plan + .requests + .iter() + .map(|req| { + let mut token_ids = Vec::with_capacity(accepted_tokens.len()); + token_ids.push(req.current_token); + token_ids.extend( + accepted_tokens + .iter() + .copied() + .take(accepted_tokens.len().saturating_sub(1)), + ); + + DraftRequestResult { + request_id: req.request_id, + token_ids, + } + }) + .collect(), + }) + } + + fn execute_speculative_verify(&mut self, plan: VerifyPlan<'_>) -> Result { + let configured = self + .speculative_accepted_tokens + .clone() + .ok_or_else(|| anyhow::anyhow!("fake speculative decoding is disabled"))?; + + let mut requests = Vec::with_capacity(plan.requests.len()); + + for req in plan.requests { + let span_len = req.as_slice().len(); + anyhow::ensure!(span_len > 0, "fake speculative verify span is empty"); + + let accepted_tokens = configured[..configured.len().min(span_len)].to_vec(); + + let current_tokens = self + .held_tokens + .get(&req.request_id) + .copied() + .ok_or_else(|| anyhow::anyhow!("missing fake request state"))?; + + self.ensure_request_tokens(req.request_id, current_tokens + accepted_tokens.len())?; + + requests.push(VerifyRequestResult { + request_id: req.request_id, + matched_draft_tokens: accepted_tokens.len().saturating_sub(1), + accepted_tokens, + }); + } + + Ok(VerifyResult { requests }) + } + + fn speculative_enabled(&self) -> bool { + self.speculative_accepted_tokens.is_some() + } + + fn speculative_request_ready(&self, request_id: RequestId) -> bool { + self.speculative_accepted_tokens.is_some() && self.held_tokens.contains_key(&request_id) + } } /// A minimal contract request: `prompt_len` filler tokens, default sampling. @@ -267,6 +373,7 @@ pub(crate) fn request(prompt_len: usize, max_tokens: usize) -> Request { Request { prompt_tokens: vec![1; prompt_len], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen3/src/scheduler/tests.rs b/pegainfer-qwen3/src/scheduler/tests.rs index 9d866f343..03f71b876 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: 0, } } @@ -515,6 +519,14 @@ fn spec_active( ignore_eos, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: if ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + token_ids: Vec::new(), + }, ..active_state(id, generated_count, max_tokens) } } @@ -537,7 +549,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, @@ -551,54 +563,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)"), } } @@ -611,7 +623,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, .. @@ -624,7 +636,7 @@ fn speculative_max_tokens_truncates_midspan() { ); assert!(matches!(finish_reason, FinishReason::Length)); } - _ => panic!("expected EmitManyAndFinish(Length)"), + _ => panic!("expected FinishMany(Length)"), } } @@ -635,14 +647,43 @@ fn speculative_ignore_eos_does_not_stop() { let results = [spec_result(1, vec![SPEC_EOS, SPEC_EOS])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { - [effects::DecodeEffect::EmitManyAndContinue { tokens, .. }] => { + [effects::DecodeEffect::ContinueMany { tokens, .. }] => { assert_eq!( tokens, &vec![SPEC_EOS, SPEC_EOS], "ignore_eos passes stop tokens through" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), + } +} + +#[test] +fn speculative_request_stop_truncates_the_span_when_eos_is_ignored() { + let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); + + let mut request = spec_active(1, 0, 100, true); + request.stop_policy.token_ids = vec![12]; + + let active = [request]; + let results = [spec_result(1, vec![10, 11, 12, 13])]; + + let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); + + match &effects[..] { + [ + effects::DecodeEffect::FinishMany { + tokens, + finish_reason, + stop_cause, + .. + }, + ] => { + assert_eq!(tokens, &vec![10, 11, 12]); + assert_eq!(*finish_reason, FinishReason::Stop); + assert_eq!(*stop_cause, Some(StopCause::Token(12))); + } + _ => panic!("expected request stop to finish the speculative span"), } } @@ -657,11 +698,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..0614433b4 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -24,6 +24,7 @@ //! target distribution; acceptance only decides how many ride one step. use anyhow::Result; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use crate::executor::RequestId; @@ -56,6 +57,9 @@ impl VerifyStepItem { #[derive(Clone, Copy)] pub(crate) struct VerifyPlan<'a> { pub requests: &'a [VerifyStepItem], + /// Request-local stop policies in the same order as `requests`. They remain + /// host-side and are not copied into GPU buffers. + pub stop_policies: &'a [StopPolicy], /// Engine step seed for the verify rows' sampler pass (same contract as /// decode: fresh per step; seeded rows re-mix their own request seed). pub sample_seed: u64, @@ -69,8 +73,9 @@ pub(crate) struct VerifyRequestResult { /// Tokens to commit: the accepted draft prefix followed by the target's /// posterior token at the first mismatch (or the block-end continuation /// when every draft is accepted). Always `1..=K + 1` tokens, so a verify - /// step always makes at least one token of progress. The scheduler still - /// owns stop-token suppression before client emission. + /// step always makes at least one token of progress. Before KV commit the + /// executor truncates this span after the first request-terminal token; + /// the scheduler retains ownership of typed stop-cause emission. pub accepted_tokens: Vec, } diff --git a/pegainfer-qwen3/tests/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index 0459712ff..735d2e7ba 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 { + eos: if params.ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + ..StopPolicy::default() + }; + Request { prompt_tokens, params, + stop_policy, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-sim/src/lib.rs b/pegainfer-sim/src/lib.rs index 1e2f52ebe..d99d0908f 100644 --- a/pegainfer-sim/src/lib.rs +++ b/pegainfer-sim/src/lib.rs @@ -311,6 +311,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; @@ -320,6 +321,7 @@ mod tests { Request { prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -425,6 +427,7 @@ mod tests { reason: FinishReason::Length, prompt_tokens: 2, completion_tokens: 3, + .. } )); }