diff --git a/docs/benchmarks/qwen3-8b-pd-vs-mix-h200.md b/docs/benchmarks/qwen3-8b-pd-vs-mix-h200.md index 4145ccce8..10d235909 100644 --- a/docs/benchmarks/qwen3-8b-pd-vs-mix-h200.md +++ b/docs/benchmarks/qwen3-8b-pd-vs-mix-h200.md @@ -17,7 +17,6 @@ vllm-bench \ --multi-turn --multi-turn-num-turns 5 \ --random-input-len 4096 --per-turn-input-len 1024 --random-output-len 128 \ --num-prompts 20 --multi-turn-concurrency 10 \ - --extra-body '{"min_tokens":1}' \ --temperature 0 ``` diff --git a/docs/index.md b/docs/index.md index cbe3d91a5..2e64df7e6 100644 --- a/docs/index.md +++ b/docs/index.md @@ -171,6 +171,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `subsystems/frontend/frontend-architecture.md` | `pegainfer-frontend` owns everything north of the model schedulers. Two contract generations coexist: the step contract (qwen3 + pegainfer-sim migrated) and the legacy `EngineHandle`/`TokenEvent` path (other five lines). Next: migrate glm52, then delete the legacy contract. | | `subsystems/frontend/simulated-inference-engine.md` | CPU-only simulated model crate on the step contract (`SimScheduler` → `LaunchedEngine::Stepped`) for vLLM/OpenAI frontend and `vllm bench serve` validation without CUDA or weights. | | `subsystems/frontend/sim-step-contract.md` | Cut `pegainfer-sim` from the legacy `EngineHandle`/`TokenEvent` path onto the step contract. | +| `subsystems/frontend/stop-token-policy.md` | Shared EOS/explicit-stop contract: preserve the triggering token and typed stop cause across migrated schedulers; retain the legacy sentinel only as a compatibility fallback. | | `subsystems/frontend/sim-high-concurrency-bench.md` | Same-session A/B vs main: feat TPOT ~30–180× better, TTFT worse and linear in C; E2EL/throughput win at c=64 and c=1024. | | `subsystems/frontend/cpu-profiling-baseline.md` | Frontend CPU profiling baseline using `pegainfer-sim` with fixed TTFT=5ms/TPOT=12ms: 200 req / concurrency=16 shows ~150ms TTFT overhead (no dominant hotspot), heap allocation ~10%, stream polling ~7.5%, IPC ~1%; reproducible benchmark command and perf evidence documented. | | `subsystems/frontend/startup-time.md` | Qwen3-4B warm startup-to-ready: frontend tokenizer load runs concurrently with the engine load (HTTP still binds only after the engine registers); mmap teardown is paid at the end of load since #377; pinned-staging upload (2026-07) cuts warm ready 5.22s → 4.66s on sm_89, and the remaining floor is the engine's own post-load startup work. | diff --git a/docs/models/kimi-k2/sampling.md b/docs/models/kimi-k2/sampling.md index cc958b320..36827cb2d 100644 --- a/docs/models/kimi-k2/sampling.md +++ b/docs/models/kimi-k2/sampling.md @@ -1,8 +1,8 @@ # Kimi-K2 sampling: param surface and design -**TL;DR**: temperature/top_k/top_p are honored on TP1/DP8 via one batched FlashInfer pass (greedy rows keep the in-graph argmax, zero perf cost); TP8 rejects non-greedy explicitly; everything else on the OpenAI surface is documented below — nothing is silently ignored anymore (#237). +**TL;DR**: temperature/top_k/top_p are honored on TP1/DP8 via one batched FlashInfer pass (greedy rows keep the in-graph argmax, zero perf cost); TP8 rejects non-greedy explicitly; unsupported sampling fields are rejected or explicitly documented, and request stop-token IDs are carried through the shared stop policy (#237). -Last touched: 2026-06 +Last touched: 2026-08 ## Param surface (`/v1/completions`) @@ -20,10 +20,14 @@ scheduler/worker. | `top_p` = 0 / out of range | **rejected** (HTTP 500, see below) | rejected | engine (`lifecycle.rs validate_sampling_params`) | | `top_k` ≥ 1 | **honored** (`top_k=1` routes greedy) | rejected if non-greedy | engine | | `top_k` = 0 | all tokens (disabled) | — | frontend maps 0 → -1; protocol type is `u32`, negatives don't parse | -| `seed` | **accepted, ignored** — engine seed is fixed at 42, per-request seed is dropped at `convert_sampling` | same | frontend | +| `seed` | greedy requests are **accepted, ignored**; non-greedy per-request seeds are rejected until row-local seed wiring lands | same | frontend | | `logprobs` | honored; for sampled rows the logprob follows the **sampled** token (reported rank is a placeholder, see PR #96) | honored (greedy only) | engine | | `max_tokens`, `echo`, `stop` (EOS) | honored | honored | engine / frontend | -| `min_p`, `frequency_penalty`, `presence_penalty`, `repetition_penalty`, `logit_bias`, `min_tokens`, `prompt_logprobs`, custom `stop_token_ids` | **accepted, ignored** — dropped at `convert_sampling`, never reach the engine | same | frontend (all models, not kimi-specific) | +| `min_p` | **honored** when in `[0, 1)`; out-of-range values are rejected | same | frontend / engine | +| `frequency_penalty`, `presence_penalty`, `repetition_penalty` | defaults are accepted; non-default values are rejected because no matching sampler path exists | same | shared frontend wire validation | +| `logit_bias`, `prompt_logprobs` | **accepted, ignored** — no engine-side implementation yet | same | frontend (all models, not kimi-specific) | +| `stop_token_ids` | **honored** independently of EOS; the matching token is preserved and reported as the stop cause | same | shared `StopPolicy` | +| `min_tokens` | **rejected** — the current scheduler contracts do not carry the threshold needed to mask EOS/stop IDs | same | shared frontend wire validation | Rejection UX pitfall: an engine-side rejection surfaces as a generic HTTP 500 (`"Internal server error"`). The real message ("top_p must be in (0, 1]…", diff --git a/docs/models/qwen3/model-crate.md b/docs/models/qwen3/model-crate.md index 20a8b3024..84f8a67b8 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`. At that point EOS still finished without emitting the stop token; the current typed stop contract emits the trigger token and its `StopCause` together. Length-limited output likewise retains the sampled final token before `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 dc84f0636..5ffec1180 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -93,7 +93,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). The typed step contract carries the trigger token and `StopCause` together; only the legacy bridge retains synthetic sentinel fallback for producers that do not yet report a 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/docs/subsystems/frontend/stop-token-policy.md b/docs/subsystems/frontend/stop-token-policy.md new file mode 100644 index 000000000..6f4a02e68 --- /dev/null +++ b/docs/subsystems/frontend/stop-token-policy.md @@ -0,0 +1,177 @@ +# Stop-Token Policy Contract + +> **TL;DR:** Keep EOS policy, explicit request stop IDs, generated tokens, and +> the concrete stop cause separate from wire parsing through every migrated +> scheduler; preserve the trigger token and drop only speculative suffixes. +> +> **Last touched:** 2026-08 + +## Preparation + +- **Read:** + - `docs/index.md` — frontend has two contract generations; Qwen3 and the + simulator use the step contract while the remaining model lines use the + legacy event path. + - `docs/subsystems/frontend/frontend-architecture.md` — migration must keep + the legacy bridge compatible until each model has its own lifecycle tests. + - `pegainfer-frontend/src/engine/stop.rs` — shared `StopPolicy` and + `StopCause` implementation. +- **Plan:** + 1. Audit every scheduler's prefill, decode, speculative, and P/D terminal + paths for policy propagation and trigger-token ordering. + 2. Add CPU contract coverage where a multi-token or handoff path can lose the + trigger or miscount completion tokens. + 3. Run formatting, frontend/model checks, and GPU tests where the server + toolchain supports them; record environmental blockers separately. + +## Contract + +`StopPolicy` has two independent inputs: + +- `eos`: `ModelDefault`, an explicit primary EOS ID, or `Ignore`; +- `token_ids`: explicit request stop IDs, active regardless of `eos`. + +`StopPolicy::classify` gives EOS precedence when the same ID appears in both +sets. A scheduler emits the sampled token first, then emits `Finished` with: + +- `StopCause::Eos(id)` for EOS stops (wire `stop_reason` remains absent); +- `StopCause::Token(id)` for explicit request stops (wire `stop_reason` is the + actual token ID); +- `None` for length stops. + +The completion count is incremented exactly once for every emitted token, +including the trigger. For speculative spans, only the prefix through the +first terminal token is committed; later candidates are discarded. + +## Migration Matrix + +| Model/path | Policy carried | Prefill | Decode/speculative | Terminal evidence | +| --- | --- | --- | --- | --- | +| Qwen3 step contract | yes | migrated | migrated, including speculative verify suffix truncation | `StopCause` | +| Qwen3.5 legacy scheduler | yes | migrated | migrated | `StopCause` | +| DeepSeek-V2-Lite | yes | migrated | migrated | `StopCause` | +| Kimi-K2 TP/DP | yes | migrated | migrated | `StopCause` | +| GLM5.2 | yes | migrated | decode, DSpark/MTP, and native P/D handoff migrated | `StopCause` / native handoff cause | +| K3 | yes | migrated | scheduler span truncation migrated; CUDA build gate pending | `StopCause` | +| Legacy bridge | compatibility | N/A | consumes typed cause when present; sentinel only for old `None` producers | typed cause or fallback | + +The legacy bridge's `StopCause::None` sentinel branch is intentionally retained +until every old producer is migrated. It must never run when a real +`StopCause` is present. + +`min_tokens` is a separate, still-unimplemented sampling contract: vLLM +requires it to mask EOS and explicit stop IDs until the requested completion +count, but the shared request types do not carry that threshold. Both bridge +generations therefore reject a non-zero value at the common wire-validation +boundary instead of letting legacy models silently ignore it. Implementing the +masking semantics is a follow-up that must add the field to the scheduler +contract and every sampler path. + +## Execution Log + +### Shared and model migration + +- Added the shared `StopPolicy` / `StopCause` types and threaded them through + the request, ledger, step, and event contracts. +- Updated wire conversion so `ignore_eos` does not disable explicit + `stop_token_ids`. +- Updated Qwen3, Qwen3.5, DeepSeek-V2-Lite, Kimi-K2, GLM5.2, K3, and Qwen3 + speculative paths to emit the trigger token before terminal metadata. +- Added native GLM5.2 P/D serialization of the typed stop cause and preserved + the anchor/cause distinction. +- Unified `min_tokens` handling at the shared vLLM wire boundary; legacy and + stepped bridges now fail closed with the same error until sampler masking is + implemented. + +### K3 multi-token contract tests + +- Extended `pegainfer-k3/src/scheduler/tests.rs` with a scripted + `decode_many` fixture. +- Added tests for an explicit stop in the middle of a span, a length cap in the + middle of a span, and EOS precedence over an overlapping explicit stop ID. +- The tests assert emitted IDs, suffix removal, `StopCause`, completion count, + and slot release. + +### Verification + +- `cargo fmt --all` — pass (run from the Linux login shell). +- `git diff --check` — pass. +- `cargo test --release -p pegainfer-frontend --lib` — **77 passed, 0 failed**. + This includes both legacy and stepped `min_tokens` rejection tests, typed + explicit-stop mapping, EOS handling, and the shared wire-policy tests. +- `cargo test --release -p pegainfer-qwen3 --lib` — **93 passed, 0 failed**. + This covers prefill/decode, speculative-span truncation, EOS/explicit-stop + precedence, and request cleanup. +- `cargo test --release -p pegainfer-deepseek-v2-lite --lib` — + **5 passed, 0 failed**. +- `cargo test --release -p pegainfer-sim --tests -- --test-threads=1` — + **22 passed, 0 failed** (7 unit, 12 frontend HTTP, 3 tool-call round-trip). + The serial test flag is required because several simulator tests bind a + fixed localhost port. +- `cargo build --release -p pegainfer-server --bin pegainfer` — pass when the + user-local `protoc-31.1` and CUDA 12.8 library paths are selected (see the + command below). The default login environment links the obsolete system + CUDA 10.1 libraries and is not a valid build environment for this tree. +- `cargo test --release -p pegainfer-k3 --lib scheduler::tests` — blocked before + Rust test compilation by the server CUDA toolchain: the installed headers do + not define `CUmemFabricHandle`, `CU_MEM_HANDLE_TYPE_FABRIC`, or + `CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED`; TileLang is also absent. + This is a mainline K3 environment prerequisite, not a stop-policy compiler + error. + +### Qwen3 HTTP smoke + +- Server: Qwen3-0.6B at `/home/ricardo.zheng/models/Qwen3/Qwen3-0.6B`, served + as `qwen3-0.6b` on `127.0.0.1:18080`; `/v1/models` returned HTTP 200 and the + expected model metadata. +- `ignore_eos=true`, `max_tokens=8`: HTTP 200, `finish_reason=length`, + `stop_reason=null`, `completion_tokens=8`. This proves ignored EOS does not + terminate the request early. +- `stop_token_ids=[0..151935]`, `ignore_eos=true`: HTTP 200, + `finish_reason=stop`, `stop_reason=12095`, `completion_tokens=1`. The first + generated token was preserved and reported as the actual explicit stop ID. +- `min_tokens=1`: HTTP 500 with the standard OpenAI error envelope. The server + log contains the detailed rejection, and the request is rejected before + scheduler submission by the shared wire validator. + +## Remaining Risks + +- K3's real DSpark executor performs speculative KV work before the scheduler + applies request stop policy. Terminal slot release resets that state, so no + suffix reaches the next request, but this is extra work rather than early + executor-side truncation. +- The legacy sentinel fallback must be removed only after each remaining + producer has a resolver and real HTTP lifecycle gate. +- A full GPU K3 test should be rerun after the server upgrades CUDA headers and + installs TileLang. + +## Next Step + +HTTP smoke is complete. Keep the legacy sentinel fallback as a separate +migration after maintainer feedback; it is not required for the typed-cause +paths covered here. + +## Debrief + +- **Outcome:** The shared stop-policy migration is implemented and verified at + unit, model-crate, simulator HTTP, and real Qwen3 HTTP levels. EOS, explicit + stop IDs, trigger-token preservation, completion counts, speculative suffix + truncation, and fail-closed `min_tokens` behavior are covered. +- **Environment caveats:** The default login shell selects obsolete CUDA 10.1 + libraries and `protoc 3.6.1`; the successful build used the user-local + `protoc-31.1` and CUDA 12.8 paths recorded below. K3 GPU validation remains + blocked by missing CUDA fabric headers and TileLang. +- **Follow-up:** Before opening a PR, perform the final diff review, decide + which local-only `.codex` artifacts stay untracked, then stage only the + intended source and documentation files. Do not remove the legacy sentinel + fallback in this change. + +Linux build environment used for the successful checks: + +```bash +export PROTOC=/database/ricardo.zheng/.local/opt/protoc-31.1/bin/protoc +export CUDA_HOME=/usr/local/cuda-12.8 +export LIBRARY_PATH=/usr/local/cuda-12.8/lib64:/usr/local/cuda-12.8/targets/x86_64-linux/lib:/usr/lib/x86_64-linux-gnu +export LD_LIBRARY_PATH=/usr/local/cuda-12.8/lib64:/usr/local/cuda-12.8/targets/x86_64-linux/lib:/usr/lib/x86_64-linux-gnu +export RUSTFLAGS=-Lnative=/usr/local/cuda-12.8/lib64 +``` diff --git a/pegainfer-deepseek-v2-lite/src/scheduler.rs b/pegainfer-deepseek-v2-lite/src/scheduler.rs index e26f3a4d6..8bc35b2a8 100644 --- a/pegainfer-deepseek-v2-lite/src/scheduler.rs +++ b/pegainfer-deepseek-v2-lite/src/scheduler.rs @@ -20,6 +20,7 @@ use grouping::take_decode_position_groups; use log::info; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::SubmittedRequest; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; @@ -51,6 +52,7 @@ struct PendingRequest { queued_at_unix_s: Option, prompt_tokens: Vec, params: SamplingParams, + stop_policy: StopPolicy, max_tokens: usize, lora_adapter: Option, token_tx: TokenSink, @@ -65,18 +67,13 @@ struct ActiveRequestState { max_tokens: usize, generated: usize, last_token: u32, - finish_policy: FinishPolicy, + stop_policy: StopPolicy, + model_eos_token_id: u32, cache: DecodeCache, stats: GenerationStats, trace: RequestTrace, } -#[derive(Clone, Copy)] -struct FinishPolicy { - eos_token_id: u32, - ignore_eos: bool, -} - struct AdmissionBatch { admitted: Vec, rejected: Vec<(PendingRequest, String)>, @@ -186,6 +183,7 @@ impl MixedRequestScheduler { ); let _ = pending.token_tx.send(TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: pending.prompt_tokens.len(), completion_tokens: 0, }); @@ -279,10 +277,8 @@ impl MixedRequestScheduler { max_tokens: pending.max_tokens, generated: 0, last_token: next, - finish_policy: FinishPolicy { - eos_token_id: self.generator.config().eos_token_id, - ignore_eos: pending.params.ignore_eos, - }, + stop_policy: pending.stop_policy, + model_eos_token_id: self.generator.config().eos_token_id, cache, stats, trace: RequestTrace::new( @@ -535,6 +531,7 @@ impl From for PendingRequest { queued_at_unix_s: req.queued_at_unix_s, prompt_tokens: req.prompt_tokens, params: req.params, + stop_policy: req.stop_policy, max_tokens: req.max_tokens, lora_adapter: req.lora_adapter, token_tx: req.token_tx, @@ -568,20 +565,9 @@ impl ActiveRequestState { pending_queue_size_at_terminal: usize, ) -> bool { self.last_token = token; - if !self.finish_policy.ignore_eos && token == self.finish_policy.eos_token_id { - self.log_http_trace( - FinishReason::Stop, - None, - active_set_size_at_terminal, - pending_queue_size_at_terminal, - ); - let _ = self.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: self.prompt_len, - completion_tokens: self.generated, - }); - return true; - } + let stop_cause = self + .stop_policy + .classify(token, |token_id| token_id == self.model_eos_token_id); let first_emit_at = self .trace @@ -610,6 +596,22 @@ impl ActiveRequestState { } self.generated += 1; + if let Some(stop_cause) = stop_cause { + self.log_http_trace( + FinishReason::Stop, + None, + active_set_size_at_terminal, + pending_queue_size_at_terminal, + ); + let _ = self.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), + prompt_tokens: self.prompt_len, + completion_tokens: self.generated, + }); + return true; + } + if self.generated == self.max_tokens { self.log_http_trace( FinishReason::Length, @@ -619,6 +621,7 @@ impl ActiveRequestState { ); let _ = self.token_tx.send(TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: self.prompt_len, completion_tokens: self.generated, }); diff --git a/pegainfer-deepseek-v2-lite/src/scheduler/tests.rs b/pegainfer-deepseek-v2-lite/src/scheduler/tests.rs index 5e3f001c5..46f702e44 100644 --- a/pegainfer-deepseek-v2-lite/src/scheduler/tests.rs +++ b/pegainfer-deepseek-v2-lite/src/scheduler/tests.rs @@ -1,8 +1,11 @@ use std::sync::Arc; use std::sync::atomic::AtomicU8; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::RequestAbortReason; use pegainfer_frontend::engine::RequestTag; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use tokio::sync::mpsc; @@ -27,6 +30,7 @@ fn request( queued_at_unix_s: None, prompt_tokens: vec![1; prompt_len], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, token_tx, @@ -60,10 +64,8 @@ fn active_state( max_tokens: 8, generated, last_token, - finish_policy: FinishPolicy { - eos_token_id: config.eos_token_id, - ignore_eos: false, - }, + stop_policy: StopPolicy::default(), + model_eos_token_id: config.eos_token_id, cache: DecodeCache::new(config), stats: GenerationStats::default(), trace: trace(), @@ -187,6 +189,7 @@ fn terminal_admission_events_keep_scheduler_contract() { assert!(send_prompt_echo(&zero)); let _ = zero.token_tx.send(TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: zero.prompt_tokens.len(), completion_tokens: 0, }); @@ -203,6 +206,7 @@ fn terminal_admission_events_keep_scheduler_contract() { recv_event(&mut zero_rx), TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, completion_tokens: 0, .. } @@ -303,16 +307,22 @@ fn eos_retirement_is_independent_per_request() { assert!(stop_state.emit_token_or_finish(config.eos_token_id, 2, 0)); assert!(!live_state.emit_token_or_finish(12, 2, 0)); + match recv_event(&mut rx_stop) { + TokenEvent::Token { id, .. } => assert_eq!(id, config.eos_token_id), + _ => panic!("EOS request should emit its triggering token"), + } match recv_event(&mut rx_stop) { TokenEvent::Finished { finish_reason, + stop_cause, completion_tokens, .. } => { assert_eq!(finish_reason, FinishReason::Stop); - assert_eq!(completion_tokens, 1); + assert_eq!(stop_cause, Some(StopCause::Eos(config.eos_token_id))); + assert_eq!(completion_tokens, 2); } - _ => panic!("EOS request should finish without emitting EOS"), + _ => panic!("EOS request should finish after emitting EOS"), } match recv_event(&mut rx_live) { TokenEvent::Token { id, .. } => assert_eq!(id, 12), @@ -343,16 +353,22 @@ fn batch_decoded_tokens_retire_eos_independently() { assert_eq!(survivors[0].0, 1); assert_eq!(survivors[0].1.request_id.as_deref(), Some("live")); assert_eq!(survivors[0].1.generated, 2); + match recv_event(&mut rx_stop) { + TokenEvent::Token { id, .. } => assert_eq!(id, config.eos_token_id), + _ => panic!("EOS row should emit its triggering token"), + } match recv_event(&mut rx_stop) { TokenEvent::Finished { finish_reason, + stop_cause, completion_tokens, .. } => { assert_eq!(finish_reason, FinishReason::Stop); - assert_eq!(completion_tokens, 1); + assert_eq!(stop_cause, Some(StopCause::Eos(config.eos_token_id))); + assert_eq!(completion_tokens, 2); } - _ => panic!("EOS row should finish without emitting EOS"), + _ => panic!("EOS row should finish after emitting EOS"), } match recv_event(&mut rx_live) { TokenEvent::Token { id, .. } => assert_eq!(id, 12), @@ -360,6 +376,78 @@ fn batch_decoded_tokens_retire_eos_independently() { } } +#[test] +fn ignored_eos_does_not_disable_explicit_stop_tokens() { + let config = test_lite_config(); + let (token_tx, mut token_rx) = TokenSink::standalone(); + let mut state = active_state("explicit-stop", token_tx, 3, 1, 10, &config); + state.stop_policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![42], + }; + state.max_tokens = 2; + + assert!(state.emit_token_or_finish(42, 0, 0)); + + assert!(matches!( + recv_event(&mut token_rx), + TokenEvent::Token { id: 42, .. } + )); + assert!(matches!( + recv_event(&mut token_rx), + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(42)), + completion_tokens: 2, + .. + } + )); +} + +#[test] +fn ignored_model_eos_remains_a_normal_generated_token() { + let config = test_lite_config(); + let (token_tx, mut token_rx) = TokenSink::standalone(); + let mut state = active_state("ignore-eos", token_tx, 3, 1, 10, &config); + state.stop_policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![], + }; + + assert!(!state.emit_token_or_finish(config.eos_token_id, 1, 0)); + + assert!(matches!( + recv_event(&mut token_rx), + TokenEvent::Token { id, .. } if id == config.eos_token_id + )); + assert_eq!(state.generated, 2); + assert!(token_rx.try_recv().is_err()); +} + +#[test] +fn eos_wins_when_it_is_also_an_explicit_stop_token() { + let config = test_lite_config(); + let (token_tx, mut token_rx) = TokenSink::standalone(); + let mut state = active_state("overlap", token_tx, 3, 1, 10, &config); + state.stop_policy = StopPolicy { + eos: EosPolicy::ModelDefault, + token_ids: vec![config.eos_token_id], + }; + + assert!(state.emit_token_or_finish(config.eos_token_id, 0, 0)); + assert!(matches!( + recv_event(&mut token_rx), + TokenEvent::Token { .. } + )); + assert!(matches!( + recv_event(&mut token_rx), + TokenEvent::Finished { + stop_cause: Some(StopCause::Eos(id)), + .. + } if id == config.eos_token_id + )); +} + #[test] fn cancelled_token_sink_retires_request() { let config = test_lite_config(); diff --git a/pegainfer-deepseek-v2-lite/tests/e2e_ep2.rs b/pegainfer-deepseek-v2-lite/tests/e2e_ep2.rs index b5234a085..1665ccb49 100644 --- a/pegainfer-deepseek-v2-lite/tests/e2e_ep2.rs +++ b/pegainfer-deepseek-v2-lite/tests/e2e_ep2.rs @@ -18,8 +18,10 @@ use anyhow::Result; use anyhow::ensure; use pegainfer_deepseek_v2_lite::DeepSeekV2LiteEp2Generator; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::engine::TokenStreamReceiver; @@ -48,6 +50,17 @@ const DSV2_LITE_MOE_LAYERS: usize = 26; const E2E_JSON_OUT_ENV: &str = "PEGAINFER_DSV2_LITE_E2E_JSON_OUT"; const E2E_CASE_SET_ENV: &str = "PEGAINFER_DSV2_LITE_E2E_CASE_SET"; +fn request_stop_policy(ignore_eos: bool) -> StopPolicy { + StopPolicy { + eos: if ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + token_ids: vec![], + } +} + #[derive(Debug, Deserialize)] struct CaseSet { cases: Vec, @@ -634,6 +647,7 @@ fn run_mixed_serving_generation(model_path: &Path, model_path_label: &str) -> Re ignore_eos, ..SamplingParams::default() }, + stop_policy: request_stop_policy(ignore_eos), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -721,6 +735,7 @@ fn run_mixed_serving_position_fallback( ignore_eos, ..SamplingParams::default() }, + stop_policy: request_stop_policy(ignore_eos), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -789,6 +804,7 @@ fn run_mixed_serving_rejection_isolation( data_parallel_rank: None, prompt_tokens: vec![1, 2, 3], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 4, lora_adapter: None, kv_transfer_params: None, @@ -805,6 +821,7 @@ fn run_mixed_serving_rejection_isolation( data_parallel_rank: None, prompt_tokens: valid_prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 6, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-frontend/examples/echo-server.rs b/pegainfer-frontend/examples/echo-server.rs index 1396f24ec..f8c36f41c 100644 --- a/pegainfer-frontend/examples/echo-server.rs +++ b/pegainfer-frontend/examples/echo-server.rs @@ -150,7 +150,7 @@ impl Scheduler for EchoScheduler { } else { FinishReason::Stop }; - ledger.finish(running.id, reason); + ledger.finish(running.id, reason, None); } } self.running = still_running; diff --git a/pegainfer-frontend/src/engine/driver.rs b/pegainfer-frontend/src/engine/driver.rs index 827adeb07..b15166b45 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; @@ -141,7 +142,7 @@ mod tests { let next = ledger.completion_tokens(id) as u32; ledger.push_tokens(id, &[next], &[]); if ledger.completion_tokens(id) >= max_tokens { - ledger.finish(id, FinishReason::Length); + ledger.finish(id, FinishReason::Length, None); } else { still_running.push((id, max_tokens)); } @@ -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/event.rs b/pegainfer-frontend/src/engine/event.rs index 4802c5d37..6477fdadf 100644 --- a/pegainfer-frontend/src/engine/event.rs +++ b/pegainfer-frontend/src/engine/event.rs @@ -1,5 +1,7 @@ use std::sync::Arc; +use super::stop::StopCause; + #[derive(Clone, Debug, PartialEq)] pub struct TokenLogprob { pub logprob: f32, @@ -36,6 +38,7 @@ pub enum TokenEvent { KvTransfer { params: serde_json::Value }, Finished { finish_reason: FinishReason, + stop_cause: Option, prompt_tokens: usize, completion_tokens: usize, }, diff --git a/pegainfer-frontend/src/engine/handle.rs b/pegainfer-frontend/src/engine/handle.rs index f6347790f..f82d8dd1a 100644 --- a/pegainfer-frontend/src/engine/handle.rs +++ b/pegainfer-frontend/src/engine/handle.rs @@ -343,6 +343,7 @@ mod tests { data_parallel_rank: rank, prompt_tokens: vec![1], params: SamplingParams::default(), + stop_policy: crate::engine::StopPolicy::default(), max_tokens: 1, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-frontend/src/engine/ledger.rs b/pegainfer-frontend/src/engine/ledger.rs index 34b4f16d6..4dcddf554 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 @@ -238,13 +239,14 @@ impl RequestLedger { // ── Terminal transitions ──────────────────────────────────────────── /// Finish the request. Token counts come from the ledger's tally. - pub fn finish(&mut self, id: RequestId, reason: FinishReason) { + pub fn finish(&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, }); @@ -278,7 +280,12 @@ impl RequestLedger { /// visibility). Closes the account; the request's buffered update for /// 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 { + pub fn defer_finish( + &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 +296,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 +382,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 +391,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 +412,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(id, FinishReason::Stop, Some(StopCause::Token(11))); backend.ledger.commit_step(); let mut steps = handle_steps(handle); @@ -419,6 +431,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 +492,7 @@ 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(id, FinishReason::Length, None); backend.ledger.commit_step(); let mut steps = handle_steps(handle); @@ -498,6 +511,7 @@ mod tests { update.terminal, Some(Terminal::Finished { reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 2, completion_tokens: 1, }) @@ -570,7 +584,7 @@ mod tests { let envelope = backend.submissions.try_recv().expect("envelope"); let id = backend.ledger.register(envelope).id; backend.ledger.admit(id); - backend.ledger.finish(id, FinishReason::Stop); + backend.ledger.finish(id, FinishReason::Stop, None); backend.ledger.push_tokens(id, &[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/request.rs b/pegainfer-frontend/src/engine/request.rs index e1f98ba5d..978a90890 100644 --- a/pegainfer-frontend/src/engine/request.rs +++ b/pegainfer-frontend/src/engine/request.rs @@ -1,4 +1,5 @@ use super::sink::TokenSink; +use super::stop::StopPolicy; use crate::sampler::SamplingParams; pub struct GenerateRequest { @@ -18,6 +19,7 @@ pub struct GenerateRequest { pub data_parallel_rank: Option, pub prompt_tokens: Vec, pub params: SamplingParams, + pub stop_policy: StopPolicy, pub max_tokens: usize, pub lora_adapter: Option, /// Opaque router/P-D metadata from the request's diff --git a/pegainfer-frontend/src/engine/step.rs b/pegainfer-frontend/src/engine/step.rs index 2cf242f8c..66f0369c8 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 @@ -228,6 +231,10 @@ impl fmt::Display for RejectReason { pub enum Terminal { Finished { reason: FinishReason, + /// Present for token-driven stop finishes. The triggering token remains + /// in `RequestUpdate.tokens`, with its real logprob in the matching + /// `RequestUpdate.logprobs` entry. + stop_cause: Option, prompt_tokens: usize, completion_tokens: usize, }, diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs new file mode 100644 index 000000000..064088ba6 --- /dev/null +++ b/pegainfer-frontend/src/engine/stop.rs @@ -0,0 +1,104 @@ +/// How a request treats end-of-sequence tokens. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum EosPolicy { + /// Do not stop on model EOS tokens. + Ignore, + /// Use the model executor's configured EOS set. + #[default] + ModelDefault, + /// Stop only on this protocol-provided primary EOS token. + Token(u32), +} + +/// Request-scoped token stopping policy. +/// +/// EOS is kept separate from caller stop tokens because the vLLM protocol +/// reports them differently: EOS has no 'stop_reason', while a request stop +/// reports the actual matching token ID. +#[derive(Clone, Debug, Default, Eq, PartialEq)] +pub struct StopPolicy { + pub eos: EosPolicy, + pub token_ids: Vec, +} + +impl StopPolicy { + /// Classify a token using vLLM's priority: EOS first, then the request's + /// explicit stop-token set. + #[must_use] + pub fn classify( + &self, + token_id: u32, + is_model_eos: impl FnOnce(u32) -> bool, + ) -> Option { + let is_eos = match self.eos { + EosPolicy::Ignore => false, + EosPolicy::ModelDefault => is_model_eos(token_id), + EosPolicy::Token(eos_token_id) => token_id == eos_token_id, + }; + + if is_eos { + Some(StopCause::Eos(token_id)) + } else if self.token_ids.contains(&token_id) { + Some(StopCause::Token(token_id)) + } else { + None + } + } +} + +/// The token-level cause of a normal stop finish. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StopCause { + /// A primary or model-default EOS token. + Eos(u32), + /// A token from the request's explicit stop-token set. + Token(u32), +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_default_classifies_model_eos() { + let policy = StopPolicy::default(); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Eos(99)) + ); + } + + #[test] + fn ignored_eos_does_not_disable_an_explicit_stop() { + let policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![99], + }; + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Token(99)) + ); + } + + #[test] + fn eos_wins_when_the_same_id_is_also_an_explicit_stop() { + let policy = StopPolicy { + eos: EosPolicy::Token(99), + token_ids: vec![99], + }; + + assert_eq!(policy.classify(99, |_| false), Some(StopCause::Eos(99))); + } + + #[test] + fn unmatched_token_does_not_stop() { + let policy = StopPolicy { + eos: EosPolicy::Token(99), + token_ids: vec![42], + }; + + assert_eq!(policy.classify(7, |_| false), None); + } +} diff --git a/pegainfer-frontend/src/vllm/bridge.rs b/pegainfer-frontend/src/vllm/bridge.rs index 0744e75d2..dc302ff8b 100644 --- a/pegainfer-frontend/src/vllm/bridge.rs +++ b/pegainfer-frontend/src/vllm/bridge.rs @@ -57,7 +57,9 @@ use crate::engine::RequestAbortReason; use crate::engine::RequestTag; use crate::engine::SchedulerMetrics; use crate::engine::SpecDecodeCounters; +use crate::engine::StopCause; use crate::engine::TokenEvent; +use crate::engine::TokenLogprob; use crate::engine::TokenSink; use crate::engine::TokenStreamReceiver; use crate::vllm::wire::convert_finish_reason; @@ -274,7 +276,7 @@ impl LocalEngineBridge { output_tx, request_id, EngineCoreFinishReason::Error, - None, + Some(StopReason::Text(unsupported)), None, None, )?; @@ -306,6 +308,13 @@ impl LocalEngineBridge { sampling_params.eos_token_id, &sampling_params.stop_token_ids, ); + let mut stop_candidate_ids = sampling_params.stop_token_ids.clone(); + if let Some(eos_token_id) = sampling_params.eos_token_id + && !stop_candidate_ids.contains(&eos_token_id) + { + stop_candidate_ids.push(eos_token_id); + } + let stop_policy = crate::vllm::wire::convert_stop_policy(&sampling_params); let tag: RequestTag = Arc::from(request_id.as_str()); let abort_reason = Arc::new(AtomicU8::new(RequestAbortReason::None as u8)); @@ -333,6 +342,7 @@ impl LocalEngineBridge { data_parallel_rank: Some(self.engine_index as usize), prompt_tokens, params: convert_sampling(&sampling_params), + stop_policy, max_tokens: sampling_params.max_tokens as usize, lora_adapter, kv_transfer_params, @@ -344,7 +354,12 @@ impl LocalEngineBridge { streams.insert( tag, - RequestStreamState::new(abort_reason, trace_root, stop_sentinel_id), + RequestStreamState::new( + abort_reason, + trace_root, + stop_sentinel_id, + stop_candidate_ids, + ), ); Ok(()) } @@ -364,6 +379,12 @@ struct RequestStreamState { /// The vLLM text decoder removes the final token from a stop-finished /// output. Keep an EOS or explicit stop token as that removable sentinel. stop_sentinel_id: Option, + /// Potential terminal ids are held until the matching `Finished` event. + /// Legacy producers send the token and terminal in separate channel sends; + /// this prevents a wakeup between those sends from leaking the trigger in + /// a non-terminal output. + stop_candidate_ids: Vec, + pending_stop_token: Option<(u32, Option)>, abort_reason: Arc, has_emitted_tokens: bool, /// Request-lifetime root span (submit → finish). The scheduler opens @@ -378,18 +399,29 @@ struct RequestStreamState { } impl RequestStreamState { - fn new(abort_reason: Arc, trace_root: Span, stop_sentinel_id: Option) -> Self { + fn new( + abort_reason: Arc, + trace_root: Span, + stop_sentinel_id: Option, + stop_candidate_ids: Vec, + ) -> Self { Self { first_token_events: None, first_token_prefill_stats: None, kv_transfer_params: None, stop_sentinel_id, + stop_candidate_ids, + pending_stop_token: None, abort_reason, has_emitted_tokens: false, trace_root, } } + fn is_stop_candidate(&self, token_id: u32) -> bool { + self.stop_candidate_ids.contains(&token_id) + } + fn abort(&self, reason: RequestAbortReason) { reason.store(&self.abort_reason); } @@ -469,6 +501,24 @@ fn dispatch_burst( /// output goes first. A lone `Scheduled` (no token, no terminal) yields no /// output — its metadata waits in `state` for the first real output. Returns /// `(output, terminated)`. +fn append_token( + token_ids: &mut Vec, + positions: &mut Vec, + has_logprobs: &mut bool, + id: u32, + logprob: Option, +) { + token_ids.push(id); + if let Some(position) = to_wire_position_logprobs(id, logprob) { + *has_logprobs = true; + positions.push(position); + } else { + positions.push(PositionLogprobs { + entries: Vec::new(), + }); + } +} + fn reduce_request( request_id: &str, state: &mut RequestStreamState, @@ -511,14 +561,25 @@ fn reduce_request( }); } TokenEvent::Token { id, logprob } => { - token_ids.push(id); - if let Some(position) = to_wire_position_logprobs(id, logprob) { - has_logprobs = true; - positions.push(position); + if let Some((pending_id, pending_logprob)) = state.pending_stop_token.take() { + append_token( + &mut token_ids, + &mut positions, + &mut has_logprobs, + pending_id, + pending_logprob, + ); + } + if state.is_stop_candidate(id) { + state.pending_stop_token = Some((id, logprob)); } else { - positions.push(PositionLogprobs { - entries: Vec::new(), - }); + append_token( + &mut token_ids, + &mut positions, + &mut has_logprobs, + id, + logprob, + ); } } TokenEvent::PromptTokens { .. } => { @@ -528,25 +589,62 @@ fn reduce_request( state.kv_transfer_params = Some(params); } TokenEvent::Finished { - finish_reason: fr, .. + finish_reason: fr, + stop_cause, + .. } => { - // PegaInfer suppresses EOS before emitting TokenEvents, while - // vLLM's text decoder expects the terminal Stop output to - // contain EOS and unconditionally removes its final token. - // Without this protocol token, a speculative step that commits - // [visible token, EOS] loses the visible token at the frontend. - if fr == FinishReason::Stop - && let Some(stop_sentinel_id) = state.stop_sentinel_id - { - token_ids.push(stop_sentinel_id); - positions.push(PositionLogprobs { - entries: Vec::new(), - }); + let had_pending_stop = state + .pending_stop_token + .take() + .map(|(id, logprob)| { + append_token( + &mut token_ids, + &mut positions, + &mut has_logprobs, + id, + logprob, + ); + }) + .is_some(); + if fr == FinishReason::Stop { + match stop_cause { + // A migrated scheduler already emitted the triggering + // token with its real logprob. Tell vLLM which token it + // is instead of guessing from the request's policy. + Some(StopCause::Token(token_id)) => { + stop_reason = Some(StopReason::TokenId(token_id)); + } + // EOS has no wire stop_reason. The emitted EOS token + // remains the terminal sentinel for vLLM's decoder. + Some(StopCause::Eos(_)) => {} + // Older schedulers suppress their terminal token and + // cannot report its cause. Preserve the old fallback + // until every producer has migrated to StopCause. + None => { + if !had_pending_stop + && let Some(stop_sentinel_id) = state.stop_sentinel_id + { + token_ids.push(stop_sentinel_id); + positions.push(PositionLogprobs { + entries: Vec::new(), + }); + } + } + } } finish_reason = Some(convert_finish_reason(fr)); terminated = true; } TokenEvent::Error { message, .. } => { + if let Some((id, logprob)) = state.pending_stop_token.take() { + append_token( + &mut token_ids, + &mut positions, + &mut has_logprobs, + id, + logprob, + ); + } warn!("request {request_id} failed: {message}"); finish_reason = Some(EngineCoreFinishReason::Error); stop_reason = Some(StopReason::Text(message)); @@ -555,6 +653,15 @@ fn reduce_request( TokenEvent::Rejected { message, .. } => { // Rejected means the request could not be admitted, not that it // completed cleanly. + if let Some((id, logprob)) = state.pending_stop_token.take() { + append_token( + &mut token_ids, + &mut positions, + &mut has_logprobs, + id, + logprob, + ); + } warn!("request {request_id} rejected: {message}"); finish_reason = Some(EngineCoreFinishReason::Error); stop_reason = Some(StopReason::Text(message)); diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 663484106..86b995780 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -44,8 +44,6 @@ use super::scheduler_stats_from; use super::send_outputs; use super::send_terminal_output; use super::send_utility_response; -use super::stop_sentinel_id; -use crate::engine::FinishReason; use crate::engine::KvCapacity; use crate::engine::Request; use crate::engine::RequestControl; @@ -53,9 +51,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; @@ -328,11 +328,12 @@ impl SteppedEngineBridge { output_tx, request_id, EngineCoreFinishReason::Error, - None, + Some(StopReason::Text(unsupported)), None, None, ); } + let lora_adapter = match lora_adapter_from_sampling_params(&sampling_params) { Ok(adapter) => adapter, Err(error) => { @@ -354,10 +355,6 @@ impl SteppedEngineBridge { .as_ref() .and_then(|args| args.get("kv_transfer_params")) .cloned(); - let stop_sentinel_id = stop_sentinel_id( - sampling_params.eos_token_id, - &sampling_params.stop_token_ids, - ); // Open the request's root span before submit so its context travels // into the scheduler as the parent of the queue/prefill/decode spans; @@ -369,9 +366,11 @@ impl SteppedEngineBridge { Span::noop() }; let trace_parent = SpanContext::from_span(&trace_root); + let control = self.scheduler.submit(Request { prompt_tokens, 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, @@ -384,7 +383,7 @@ impl SteppedEngineBridge { names.insert(request_id.clone(), control.id()); streams.insert( control.id(), - SteppedStream::new(request_id, control, trace_root, stop_sentinel_id), + SteppedStream::new(request_id, control, trace_root), ); Ok(()) } @@ -404,9 +403,6 @@ 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. - stop_sentinel_id: Option, /// Request-lifetime root span; held only for its `Drop`, which closes the /// trace when the stream state is removed. #[allow(dead_code)] @@ -414,12 +410,7 @@ struct SteppedStream { } impl SteppedStream { - fn new( - request_id: String, - control: RequestControl, - trace_root: Span, - stop_sentinel_id: Option, - ) -> Self { + fn new(request_id: String, control: RequestControl, trace_root: Span) -> Self { Self { request_id, control, @@ -428,7 +419,6 @@ impl SteppedStream { cached_tokens: 0, prefill_stats_sent: false, kv_transfer_params: None, - stop_sentinel_id, trace_root, } } @@ -485,7 +475,7 @@ fn reduce_update( // before submission), matching the legacy bridge. Wiring it up means mapping // PromptEcho into EngineCoreOutput's prompt_logprobs fields. - let mut token_ids = update.tokens; + let token_ids = update.tokens; let mut has_logprobs = false; let mut positions: Vec = Vec::with_capacity(token_ids.len()); for (i, &id) in token_ids.iter().enumerate() { @@ -505,18 +495,13 @@ 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. - if reason == FinishReason::Stop - && let Some(stop_sentinel_id) = state.stop_sentinel_id - { - token_ids.push(stop_sentinel_id); - positions.push(PositionLogprobs { - entries: Vec::new(), - }); + Some(Terminal::Finished { + reason, stop_cause, .. + }) => { + 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; } @@ -579,3 +564,138 @@ 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::FinishReason; + use crate::engine::TokenLogprob; + + #[test] + fn min_tokens_is_rejected_before_scheduler_submission() { + let (scheduler, backend) = crate::engine::scheduler_pair(); + let bridge = SteppedEngineBridge { + input_address: String::new(), + output_address: String::new(), + scheduler, + kv_capacity: None, + max_model_len: 4096, + engine_index: 3, + data_parallel_size: 1, + }; + let mut sampling_params = EngineCoreSamplingParams::for_test(); + sampling_params.min_tokens = 1; + let request = EngineCoreRequest { + request_id: "min-tokens".to_string(), + prompt_token_ids: Some(vec![1, 2]), + sampling_params: Some(sampling_params), + ..EngineCoreRequest::default() + }; + let mut streams = HashMap::new(); + let mut names = HashMap::new(); + let (output_tx, mut output_rx) = tokio::sync::mpsc::unbounded_channel(); + + bridge + .start_request(request, &mut streams, &mut names, &output_tx) + .expect("reject min_tokens request"); + + assert!(matches!( + backend.submissions.try_recv(), + Err(crossbeam_channel::TryRecvError::Empty) + )); + assert!(streams.is_empty()); + assert!(names.is_empty()); + + let batch = match output_rx.try_recv().expect("rejection output") { + EngineCoreOutputs::RequestBatch(batch) => batch, + other => panic!("expected request batch, got {other:?}"), + }; + assert_eq!(batch.outputs.len(), 1); + let output = &batch.outputs[0]; + assert_eq!(output.request_id, "min-tokens"); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Error)); + assert_eq!( + output.stop_reason, + Some(StopReason::Text( + "min_tokens=1 is not supported by current engine contracts".to_string() + )) + ); + assert!( + batch + .finished_requests + .as_ref() + .is_some_and(|ids| ids.contains("min-tokens")) + ); + assert!(output_rx.try_recv().is_err()); + } + + #[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()); + + 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()); + + 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); + } +} diff --git a/pegainfer-frontend/src/vllm/bridge/tests.rs b/pegainfer-frontend/src/vllm/bridge/tests.rs index 0926609f7..0b98f7f43 100644 --- a/pegainfer-frontend/src/vllm/bridge/tests.rs +++ b/pegainfer-frontend/src/vllm/bridge/tests.rs @@ -9,11 +9,68 @@ use std::sync::atomic::Ordering; +use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; + use super::*; use crate::engine::FinishReason; use crate::engine::RequestAbortReason; +use crate::engine::StopCause; use crate::engine::TokenLogprob; +#[test] +fn legacy_bridge_rejects_min_tokens_before_scheduler_submission() { + let (submit_tx, mut submit_rx) = mpsc::unbounded_channel(); + let bridge = LocalEngineBridge { + input_address: String::new(), + output_address: String::new(), + handle: EngineHandle::new(submit_tx), + max_model_len: 4096, + engine_index: 2, + data_parallel_size: 1, + metrics_watch: None, + }; + let mut sampling_params = EngineCoreSamplingParams::for_test(); + sampling_params.min_tokens = 1; + let request = EngineCoreRequest { + request_id: "legacy-min-tokens".to_string(), + prompt_token_ids: Some(vec![1, 2]), + sampling_params: Some(sampling_params), + ..EngineCoreRequest::default() + }; + let (event_tx, _event_rx) = mpsc::unbounded_channel(); + let (output_tx, mut output_rx) = mpsc::unbounded_channel(); + let mut streams = HashMap::new(); + + bridge + .start_request(request, &event_tx, &output_tx, &mut streams) + .expect("reject min_tokens request"); + + assert!(submit_rx.try_recv().is_err()); + assert!(streams.is_empty()); + + let batch = match output_rx.try_recv().expect("rejection output") { + EngineCoreOutputs::RequestBatch(batch) => batch, + other => panic!("expected request batch, got {other:?}"), + }; + assert_eq!(batch.outputs.len(), 1); + let output = &batch.outputs[0]; + assert_eq!(output.request_id, "legacy-min-tokens"); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Error)); + assert_eq!( + output.stop_reason, + Some(StopReason::Text( + "min_tokens=1 is not supported by current engine contracts".to_string() + )) + ); + assert!( + batch + .finished_requests + .as_ref() + .is_some_and(|ids| ids.contains("legacy-min-tokens")) + ); + assert!(output_rx.try_recv().is_err()); +} + /// Test harness that exercises the bridge's demux path directly: register /// requests, emit tagged events onto the shared channel, drain one ready /// burst at a time, and inspect the coalesced outputs — the same flow the @@ -41,14 +98,23 @@ impl Demux { /// Register a request as `start_request` does and return its abort reason. fn add(&mut self, id: &str) -> Arc { - self.add_with_stop_sentinel(id, None) + self.add_with_stop_candidates(id, None, Vec::new()) } fn add_with_eos(&mut self, id: &str, eos_token_id: Option) -> Arc { - self.add_with_stop_sentinel(id, eos_token_id) + self.add_with_stop_candidates(id, eos_token_id, eos_token_id.into_iter().collect()) } fn add_with_stop_sentinel(&mut self, id: &str, stop_sentinel_id: Option) -> Arc { + self.add_with_stop_candidates(id, stop_sentinel_id, stop_sentinel_id.into_iter().collect()) + } + + fn add_with_stop_candidates( + &mut self, + id: &str, + stop_sentinel_id: Option, + stop_candidate_ids: Vec, + ) -> Arc { let tag: RequestTag = Arc::from(id); let abort_reason = Arc::new(AtomicU8::new(RequestAbortReason::None as u8)); self.streams.insert( @@ -57,6 +123,7 @@ impl Demux { Arc::clone(&abort_reason), fastrace::Span::noop(), stop_sentinel_id, + stop_candidate_ids, ), ); abort_reason @@ -135,6 +202,7 @@ fn token_and_finish_in_one_burst_coalesce() { "req-1", TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 16, completion_tokens: 2, }, @@ -186,6 +254,7 @@ fn stop_output_appends_eos_for_vllm_decoder() { "req-stop-eos", TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: None, prompt_tokens: 16, completion_tokens: 3, }, @@ -215,6 +284,7 @@ fn explicit_stop_token_is_used_as_sentinel_when_eos_is_absent() { "req-stop-token", TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: None, prompt_tokens: 16, completion_tokens: 2, }, @@ -227,6 +297,145 @@ fn explicit_stop_token_is_used_as_sentinel_when_eos_is_absent() { assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); } +/// A migrated scheduler sends the actual explicit stop token and identifies +/// it in the terminal event. The bridge must preserve both its logprob and +/// real token ID instead of adding a guessed sentinel from the request. +#[test] +fn explicit_stop_cause_uses_actual_token_without_sentinel() { + let mut d = Demux::new(); + d.add_with_stop_candidates("req-explicit-stop", Some(2), vec![2, 43]); + d.emit( + "req-explicit-stop", + TokenEvent::Token { + id: 11, + logprob: None, + }, + ); + d.emit( + "req-explicit-stop", + TokenEvent::Token { + id: 43, + logprob: Some(TokenLogprob { + logprob: -0.25, + top_logprobs: vec![(43, -0.25), (44, -1.0)], + }), + }, + ); + d.emit( + "req-explicit-stop", + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(43)), + prompt_tokens: 16, + completion_tokens: 2, + }, + ); + assert!(d.drain()); + + let batch = d.next_output().expect("terminal output"); + let output = &batch.outputs[0]; + 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.as_ref().expect("stop-token logprob") { + MaybeWireLogprobs::Direct(direct) => direct, + MaybeWireLogprobs::Wire(_) => panic!("expected direct batched 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); +} + +/// Legacy schedulers publish the trigger and terminal through two channel +/// sends. If the bridge wakes between them, it must hold the candidate until +/// the terminal arrives instead of exposing it as a non-terminal token. +#[test] +fn explicit_stop_cause_coalesces_across_split_bursts() { + let mut d = Demux::new(); + d.add_with_stop_candidates("req-split-stop", Some(2), vec![2, 43]); + d.emit( + "req-split-stop", + TokenEvent::Token { + id: 43, + logprob: Some(TokenLogprob { + logprob: -0.25, + top_logprobs: vec![(43, -0.25), (44, -1.0)], + }), + }, + ); + + assert!(d.drain()); + assert!( + d.next_output().is_none(), + "the possible trigger must wait for its terminal event" + ); + + d.emit( + "req-split-stop", + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(43)), + prompt_tokens: 16, + completion_tokens: 1, + }, + ); + assert!(d.drain()); + + let batch = d.next_output().expect("terminal output"); + let output = &batch.outputs[0]; + assert_eq!(output.new_token_ids, vec![43]); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); + assert_eq!(output.stop_reason, Some(StopReason::TokenId(43))); + + let direct = match output.new_logprobs.as_ref().expect("stop-token logprob") { + MaybeWireLogprobs::Direct(direct) => direct, + MaybeWireLogprobs::Wire(_) => panic!("expected direct batched logprobs"), + }; + assert_eq!(direct.positions.len(), 1); + assert_eq!(direct.positions[0].entries[0].token_id, 43); + assert!((direct.positions[0].entries[0].logprob + 0.25).abs() < f32::EPSILON); + assert!(!d.streams.contains_key("req-split-stop")); +} + +/// EOS stops carry no wire stop reason. Because a migrated scheduler already +/// sent EOS, the bridge must not append a second synthetic sentinel either. +#[test] +fn eos_stop_cause_has_no_wire_stop_reason_or_duplicate_sentinel() { + let mut d = Demux::new(); + d.add_with_eos("req-eos-stop", Some(2)); + d.emit( + "req-eos-stop", + TokenEvent::Token { + id: 11, + logprob: None, + }, + ); + d.emit( + "req-eos-stop", + TokenEvent::Token { + id: 2, + logprob: None, + }, + ); + d.emit( + "req-eos-stop", + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(2)), + prompt_tokens: 16, + completion_tokens: 2, + }, + ); + assert!(d.drain()); + + let batch = d.next_output().expect("terminal output"); + let output = &batch.outputs[0]; + assert_eq!(output.new_token_ids, vec![11, 2]); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); + assert_eq!(output.stop_reason, None); +} + /// A lone `Scheduled` (no token yet) emits nothing; its metadata waits in the /// stream state across bursts and flushes onto the first real output. This is /// the reason `RequestStreamState` holds `first_token_*` between bursts. @@ -287,6 +496,7 @@ fn lone_kv_transfer_defers_until_terminal_output() { "req-handoff", TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: None, prompt_tokens: 4, completion_tokens: 1, }, @@ -378,6 +588,7 @@ fn stop_on_prefill_terminal_output_carries_prefill_stats() { "req-stop", TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: None, prompt_tokens: 16, completion_tokens: 0, }, diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index fbc5a9f74..e229c0f0e 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,9 +46,10 @@ 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. - let ignore_eos = params.eos_token_id.is_none() && params.stop_token_ids.is_empty(); + // models with a real EOS. Explicit stop tokens travel independently in + // StopPolicy, so they must not re-enable model EOS after vLLM lowered + // `ignore_eos=true` to a missing `_eos_token_id`. + let ignore_eos = params.eos_token_id.is_none(); if params.temperature <= 0.0 { return SamplingParams { temperature: 0.0, @@ -76,6 +79,15 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar } } +pub(crate) fn convert_stop_policy(params: &EngineCoreSamplingParams) -> StopPolicy { + StopPolicy { + eos: params + .eos_token_id + .map_or(EosPolicy::Ignore, EosPolicy::Token), + token_ids: params.stop_token_ids.clone(), + } +} + /// Reject request parameters the engine would otherwise silently ignore. /// Returns the offending description; `None` means the request is servable. /// @@ -84,6 +96,12 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar /// carrying 1.0000001 wants a penalty and must be rejected, not rounded away. #[allow(clippy::float_cmp)] pub(crate) fn unsupported_request_params(params: &EngineCoreSamplingParams) -> Option { + 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)); } @@ -191,11 +209,46 @@ 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. + // Explicit client stop tokens do not re-enable EOS after the frontend + // dropped _eos_token_id for ignore_eos=true. params.eos_token_id = None; params.stop_token_ids = vec![42]; - assert!(!convert_sampling(¶ms).ignore_eos); + 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), + StopPolicy { + eos: EosPolicy::Token(99), + token_ids: vec![11], + } + ); + } + + #[test] + fn all_stop_tokens_do_not_reenable_ignored_eos() { + let mut params = EngineCoreSamplingParams::for_test(); + + // `_all_stop_token_ids` is the model's masking set. It must not turn + // EOS detection back on after vLLM lowers `ignore_eos=true` to a + // missing `_eos_token_id`. + params.eos_token_id = None; + params.stop_token_ids = vec![42]; + params.all_stop_token_ids = BTreeSet::from([99]); + + assert_eq!( + convert_stop_policy(¶ms), + StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![42], + } + ); } #[test] @@ -246,6 +299,17 @@ mod tests { assert!(unsupported_request_params(¶ms).is_some()); } + #[test] + fn unsupported_sampling_rejects_min_tokens_until_masking_is_implemented() { + let mut params = EngineCoreSamplingParams::for_test(); + params.min_tokens = 1; + + assert_eq!( + unsupported_request_params(¶ms), + Some("min_tokens=1 is not supported by current engine contracts".to_string()) + ); + } + #[test] fn nonempty_ec_transfer_params_is_refused_not_dropped() { let mut params = EngineCoreSamplingParams::for_test(); diff --git a/pegainfer-gemma4/src/engine.rs b/pegainfer-gemma4/src/engine.rs index f76a23bcf..2c738a7ff 100644 --- a/pegainfer-gemma4/src/engine.rs +++ b/pegainfer-gemma4/src/engine.rs @@ -195,6 +195,54 @@ impl GenerationPolicy { } } +/// Emit one sampled token and, when it is terminal, the matching finish +/// event. Token delivery comes first so the frontend can retain the exact +/// trigger token and its logprob before applying protocol stop semantics. +fn emit_sampled_token( + request: &GenerateRequest, + policy: &GenerationPolicy, + token_id: u32, + logprob: Option, + prompt_tokens: usize, + completion_tokens: usize, +) -> bool { + if request + .token_tx + .send(TokenEvent::Token { + id: token_id, + logprob, + }) + .is_err() + { + return false; + } + + if let Some(stop_cause) = request + .stop_policy + .classify(token_id, |id| policy.eos.contains(&id)) + { + let _ = request.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), + prompt_tokens, + completion_tokens, + }); + return false; + } + + if completion_tokens >= request.max_tokens { + let _ = request.token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + stop_cause: None, + prompt_tokens, + completion_tokens, + }); + return false; + } + + true +} + fn token_ids(value: &serde_json::Value) -> Result> { fn one(value: &serde_json::Value) -> Result { let raw = value @@ -532,18 +580,17 @@ struct Walker { } /// One newcomer row of a mixed step's logits head, ahead of the active -/// rows: its sampling params, its logprob request (0 = none), and whether -/// its pick may stop it — a mid-walk segment's row is sampled and -/// discarded, so it never stops. +/// rows: its sampling params and logprob request (0 = none). A mid-walk +/// segment's row is sampled and discarded; its request policy is applied +/// only when the final sampled row is actually emitted. struct HeadRow<'a> { params: &'a pegainfer_frontend::sampler::SamplingParams, logprobs: usize, - ignore_eos: bool, } /// Suppress, sample and score one mixed step's logits — `head` describes /// rows `0..head.len()`, the active rows follow — then deliver the active -/// rows' events. Returns every row's pick, logprob and stop flag; `Err` +/// rows' events. Returns every row's pick and logprob; `Err` /// carries the failure message after the active batch has been failed. #[allow(clippy::too_many_arguments, clippy::type_complexity)] fn mixed_head_flow( @@ -556,7 +603,7 @@ fn mixed_head_flow( head: &[HeadRow<'_>], active: &mut Vec, logits: &mut HiddenStates, -) -> Result<(Vec, Vec>, Vec), String> { +) -> Result<(Vec, Vec>), String> { let k = head.len(); let fail_batch = |active: &mut Vec, what: &str, err: &anyhow::Error| { log::error!("{what} failed: {err:#}"); @@ -592,18 +639,11 @@ fn mixed_head_flow( } } }; - let mut stops = vec![false; active.len() + k]; - for (j, h) in head.iter().enumerate() { - stops[j] = !h.ignore_eos && policy.eos.contains(&picked[j]); - } - for (row, entry) in active.iter().enumerate() { - stops[row + k] = !entry.request.params.ignore_eos && policy.eos.contains(&picked[row + k]); - } let mut lp_requests: Vec = Vec::new(); lp_requests.extend( head.iter() .enumerate() - .filter(|(j, h)| h.logprobs > 0 && !stops[*j]) + .filter(|(_, h)| h.logprobs > 0) .map(|(j, h)| LogprobRequest { row: j, picked: picked[j], @@ -614,7 +654,7 @@ fn mixed_head_flow( active .iter() .enumerate() - .filter(|(row, entry)| entry.request.logprobs > 0 && !stops[row + k]) + .filter(|(_, entry)| entry.request.logprobs > 0) .map(|(row, entry)| LogprobRequest { row: row + k, picked: picked[row + k], @@ -637,8 +677,8 @@ fn mixed_head_flow( } // Active rows: the decode-round event flow, `k` logits rows up. - emit_decode_rows(active, &picked, &stops, &mut logprobs, k); - Ok((picked, logprobs, stops)) + emit_decode_rows(active, &picked, &mut logprobs, k, policy); + Ok((picked, logprobs)) } struct Active { @@ -1219,21 +1259,6 @@ impl EngineState { Ok(tokens) => tokens[0], Err(err) => return fail(format!("{err:#}")), }; - let finish = |reason: FinishReason, completion_tokens: usize| { - let _ = sink.send(TokenEvent::Finished { - finish_reason: reason, - prompt_tokens, - completion_tokens, - }); - Admitted::Done - }; - // The stop token retires the request without being emitted: the - // frontend appends its own sentinel for a terminal Stop and drops the - // last id, so an engine that emits EOS costs the client its final - // visible token. - if !request.params.ignore_eos && self.policy.eos.contains(&next) { - return finish(FinishReason::Stop, 0); - } let logprob = if request.logprobs > 0 { match pegainfer_sample::token_logprobs_batch( &self.ctx, @@ -1250,12 +1275,9 @@ impl EngineState { } else { None }; - if sink.send(TokenEvent::Token { id: next, logprob }).is_err() { + if !emit_sampled_token(&request, &self.policy, next, logprob, prompt_tokens, 1) { return Admitted::Done; } - if request.max_tokens <= 1 { - return finish(FinishReason::Length, 1); - } Admitted::Active(Box::new(Active { request, kv, @@ -1499,7 +1521,6 @@ impl EngineState { let head = [HeadRow { params: &w.request.params, logprobs: w.request.logprobs, - ignore_eos: w.request.params.ignore_eos, }]; match mixed_head_flow( &self.ctx, @@ -1512,7 +1533,7 @@ impl EngineState { active, &mut logits, ) { - Ok((picked, mut lps, _)) => { + Ok((picked, mut lps)) => { w.first = Some((picked[0], lps[0].take())); } Err(message) => { @@ -1614,11 +1635,6 @@ impl EngineState { HeadRow { params: &w.request.params, logprobs: if last { w.request.logprobs } else { 0 }, - ignore_eos: if last { - w.request.params.ignore_eos - } else { - true - }, } }) .collect(); @@ -1635,7 +1651,7 @@ impl EngineState { ) }; match flow { - Ok((picked, mut lps, _)) => { + Ok((picked, mut lps)) => { let mut si = 0usize; for (w, t) in walkers.iter_mut().zip(&takes) { if let Some((take, last)) = *t { @@ -1684,29 +1700,7 @@ impl EngineState { cache.insert(entry, resumed); } } - // The stop token retires the request without being emitted, the - // same contract as every other admission path. - if !request.params.ignore_eos && self.policy.eos.contains(&next) { - let _ = request.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens, - completion_tokens: 0, - }); - return; - } - if request - .token_tx - .send(TokenEvent::Token { id: next, logprob }) - .is_err() - { - return; - } - if request.max_tokens <= 1 { - let _ = request.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens, - completion_tokens: 1, - }); + if !emit_sampled_token(&request, &self.policy, next, logprob, prompt_tokens, 1) { return; } active.push(Active { @@ -1817,13 +1811,12 @@ impl EngineState { } } } - let (picked, mut logprobs, stops) = { + let (picked, mut logprobs) = { let head: Vec> = newcomers .iter() .map(|(request, _, _)| HeadRow { params: &request.params, logprobs: request.logprobs, - ignore_eos: request.params.ignore_eos, }) .collect(); match mixed_head_flow( @@ -1845,30 +1838,15 @@ impl EngineState { // The newcomers: their first tokens are logits rows `0..k`. for (j, (request, kv, _)) in newcomers.into_iter().enumerate() { let prompt_tokens = request.prompt_tokens.len(); - let finish = |reason: FinishReason, completion_tokens: usize| { - let _ = request.token_tx.send(TokenEvent::Finished { - finish_reason: reason, - prompt_tokens, - completion_tokens, - }); - }; - if stops[j] { - finish(FinishReason::Stop, 0); - continue; - } let next = picked[j]; - if request - .token_tx - .send(TokenEvent::Token { - id: next, - logprob: logprobs[j].take(), - }) - .is_err() - { - continue; - } - if request.max_tokens <= 1 { - finish(FinishReason::Length, 1); + if !emit_sampled_token( + &request, + &self.policy, + next, + logprobs[j].take(), + prompt_tokens, + 1, + ) { continue; } active.push(Active { @@ -1935,14 +1913,10 @@ impl EngineState { Err(err) => return fail_batch(active, "batched sampling", &err), } }; - let mut stops = vec![false; active.len()]; - for (row, entry) in active.iter().enumerate() { - stops[row] = !entry.request.params.ignore_eos && self.policy.eos.contains(&picked[row]); - } let lp_requests: Vec = active .iter() .enumerate() - .filter(|(row, entry)| entry.request.logprobs > 0 && !stops[*row]) + .filter(|(_, entry)| entry.request.logprobs > 0) .map(|(row, entry)| LogprobRequest { row, picked: picked[row], @@ -1961,53 +1935,34 @@ impl EngineState { } } - emit_decode_rows(active, &picked, &stops, &mut logprobs, 0); + emit_decode_rows(active, &picked, &mut logprobs, 0, &self.policy); } } /// Deliver one decode step's outcome to every active row and retire the /// finished ones — the event flow both the pure decode round and the mixed /// admission share; `row_base` is the row's offset into the step's logits -/// (the mixed step's row 0 is the newcomer). A stop token retires the -/// request without being emitted; a send failure retires a cancelled one. +/// (the mixed step's row 0 is the newcomer). A send failure or terminal +/// token retires the request after the sampled token has been delivered. fn emit_decode_rows( active: &mut Vec, picked: &[u32], - stops: &[bool], logprobs: &mut [Option], row_base: usize, + policy: &GenerationPolicy, ) { let mut retire: Vec = Vec::new(); for (row, entry) in active.iter_mut().enumerate() { - if stops[row + row_base] { - let _ = entry.request.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: entry.prompt_tokens, - completion_tokens: entry.emitted, - }); - retire.push(row); - continue; - } let token = picked[row + row_base]; entry.emitted += 1; - if entry - .request - .token_tx - .send(TokenEvent::Token { - id: token, - logprob: logprobs[row + row_base].take(), - }) - .is_err() - { - retire.push(row); - continue; - } - if entry.emitted >= entry.request.max_tokens { - let _ = entry.request.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: entry.prompt_tokens, - completion_tokens: entry.emitted, - }); + if !emit_sampled_token( + &entry.request, + policy, + token, + logprobs[row + row_base].take(), + entry.prompt_tokens, + entry.emitted, + ) { retire.push(row); continue; } @@ -2018,6 +1973,136 @@ fn emit_decode_rows( } } +#[cfg(test)] +mod stop_contract_tests { + use pegainfer_frontend::engine::EosPolicy; + use pegainfer_frontend::engine::StopCause; + use pegainfer_frontend::engine::StopPolicy; + use pegainfer_frontend::engine::TokenSink; + use pegainfer_frontend::sampler::SamplingParams; + + use super::*; + + fn request( + stop_policy: StopPolicy, + max_tokens: usize, + ) -> ( + GenerateRequest, + pegainfer_frontend::engine::TokenStreamReceiver, + ) { + let (token_tx, token_rx) = TokenSink::standalone(); + ( + GenerateRequest { + trace_parent: None, + request_id: None, + queued_at_unix_s: None, + data_parallel_rank: None, + prompt_tokens: vec![1, 2], + params: SamplingParams::default(), + stop_policy, + max_tokens, + lora_adapter: None, + kv_transfer_params: None, + token_tx, + logprobs: 1, + echo: false, + }, + token_rx, + ) + } + + fn policy() -> GenerationPolicy { + GenerationPolicy { + eos: vec![99], + suppress: Vec::new(), + } + } + + #[test] + fn model_eos_is_emitted_before_its_stop_finish() { + let (request, mut rx) = request(StopPolicy::default(), 8); + let logprob = TokenLogprob { + logprob: -0.25, + top_logprobs: vec![(99, -0.25)], + }; + + assert!(!emit_sampled_token( + &request, + &policy(), + 99, + Some(logprob.clone()), + 2, + 1, + )); + + let Ok((_, TokenEvent::Token { id, logprob: seen })) = rx.try_recv() else { + panic!("expected trigger Token event"); + }; + assert_eq!(id, 99); + assert_eq!(seen, Some(logprob)); + let Ok(( + _, + TokenEvent::Finished { + finish_reason, + stop_cause, + completion_tokens, + .. + }, + )) = rx.try_recv() + else { + panic!("expected terminal Finished event"); + }; + assert_eq!(finish_reason, FinishReason::Stop); + assert_eq!(stop_cause, Some(StopCause::Eos(99))); + assert_eq!(completion_tokens, 1); + } + + #[test] + fn ignored_eos_still_honors_an_explicit_stop_token() { + let (request, mut rx) = request( + StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![99], + }, + 8, + ); + + assert!(!emit_sampled_token(&request, &policy(), 99, None, 2, 1)); + assert!(matches!( + rx.try_recv(), + Ok((_, TokenEvent::Token { id: 99, .. })) + )); + let Ok((_, TokenEvent::Finished { stop_cause, .. })) = rx.try_recv() else { + panic!("expected terminal Finished event"); + }; + assert_eq!(stop_cause, Some(StopCause::Token(99))); + } + + #[test] + fn unmatched_token_at_the_budget_finishes_by_length() { + let (request, mut rx) = request(StopPolicy::default(), 1); + + assert!(!emit_sampled_token(&request, &policy(), 7, None, 2, 1)); + assert!(matches!( + rx.try_recv(), + Ok((_, TokenEvent::Token { id: 7, .. })) + )); + let Ok(( + _, + TokenEvent::Finished { + finish_reason, + stop_cause, + .. + }, + )) = rx.try_recv() + else { + panic!("expected terminal Finished event"); + }; + assert_eq!(finish_reason, FinishReason::Length); + assert_eq!(stop_cause, None); + } +} + #[cfg(test)] mod gate { use super::*; @@ -2058,8 +2143,10 @@ mod lane_tests { use std::path::Path; use pegainfer_frontend::engine::EngineLoadOptions; + use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; + use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::engine::TokenStreamReceiver; @@ -2289,6 +2376,10 @@ mod lane_tests { ignore_eos: true, ..pegainfer_frontend::sampler::SamplingParams::default() }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-glm52/src/oracle/freerun_step.rs b/pegainfer-glm52/src/oracle/freerun_step.rs index 2722f9586..392f83335 100644 --- a/pegainfer-glm52/src/oracle/freerun_step.rs +++ b/pegainfer-glm52/src/oracle/freerun_step.rs @@ -27,7 +27,9 @@ use std::path::PathBuf; use anyhow::Context; use anyhow::Result; use pegainfer_frontend::engine::EngineHandle; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; use pegainfer_sample::SamplingParams; @@ -95,6 +97,10 @@ fn run_request( data_parallel_rank: Some(rank), prompt_tokens: PATHOLOGICAL_PROMPT[..prompt_len].to_vec(), params, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-glm52/src/oracle/mtp_production.rs b/pegainfer-glm52/src/oracle/mtp_production.rs index b9c94de54..eb5cbb5fb 100644 --- a/pegainfer-glm52/src/oracle/mtp_production.rs +++ b/pegainfer-glm52/src/oracle/mtp_production.rs @@ -4,7 +4,9 @@ use std::path::PathBuf; use anyhow::Context; use anyhow::Result; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; use pegainfer_sample::SamplingParams; @@ -60,6 +62,10 @@ fn native_mtp_uses_final_normalized_target_hidden() -> Result<()> { ignore_eos: true, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, max_tokens: 256, lora_adapter: None, kv_transfer_params: None, @@ -119,6 +125,10 @@ fn native_mtp_uses_final_normalized_target_hidden() -> Result<()> { ignore_eos: true, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, max_tokens: output_lengths[rank], lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-glm52/src/scheduler/admission.rs b/pegainfer-glm52/src/scheduler/admission.rs index a203005cf..c55788212 100644 --- a/pegainfer-glm52/src/scheduler/admission.rs +++ b/pegainfer-glm52/src/scheduler/admission.rs @@ -13,6 +13,7 @@ use std::sync::Arc; use anyhow::Context as _; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::unix_now_s; use pegainfer_kv_store::BlockPool; @@ -155,30 +156,44 @@ pub(super) fn admit_from_queue( ) -> anyhow::Result<()> { // Zero-capacity natives finish at intake, before the slot and budget // gates — a saturated rank must never delay a reply that needs no - // capacity. Two forms: P consumed EOS (anchor None → Stop), and the + // capacity. Two forms: P consumed a typed stop token, and the // replayed anchor exhausting max_tokens (→ Length); neither restores KV. pending.retain(|entry| { let Resolved::Native { req, handoff, .. } = entry else { return true; }; - let (anchor, finish_reason) = match handoff.anchor_token_id { - None => (None, FinishReason::Stop), - Some(anchor) if req.max_tokens == 1 => (Some(anchor), FinishReason::Length), - Some(_) => return true, - }; + let (terminal_token, finish_reason, stop_cause) = + match (handoff.stop_cause, handoff.anchor_token_id) { + (Some(cause), _) => { + let cause = StopCause::from(cause); + let token = match cause { + StopCause::Eos(id) | StopCause::Token(id) => id, + }; + (Some(token), FinishReason::Stop, Some(cause)) + } + // Compatibility for a v5 envelope that omitted the optional + // cause. Its fingerprint prevents mixing with older peers, + // but retaining the fallback makes malformed metadata fail in + // the same non-allocating way as the old EOS marker. + (None, None) => (None, FinishReason::Stop, None), + (None, Some(anchor)) if req.max_tokens == 1 => { + (Some(anchor), FinishReason::Length, None) + } + (None, Some(_)) => return true, + }; let prompt_tokens = req.prompt_tokens.len(); let _ = req.token_tx.send(TokenEvent::Scheduled { queued_at_unix_s: req.queued_at_unix_s.unwrap_or_else(unix_now_s), scheduled_at_unix_s: unix_now_s(), prompt_tokens, - cached_tokens: anchor.map_or(0, |_| handoff.committed_len), + cached_tokens: terminal_token.map_or(0, |_| handoff.committed_len), }); - // P sampled the anchor but never sent it; it reaches the client here. - if let Some(anchor) = anchor { + // P sampled the terminal/anchor token; replay it to this request once. + if let Some(token) = terminal_token { if req .token_tx .send(TokenEvent::Token { - id: anchor, + id: token, logprob: None, }) .is_err() @@ -188,6 +203,7 @@ pub(super) fn admit_from_queue( } let _ = req.token_tx.send(TokenEvent::Finished { finish_reason, + stop_cause, prompt_tokens, completion_tokens: 1, }); @@ -384,7 +400,7 @@ pub(super) fn admit_from_queue( let state = Glm52SlotState::new( req.prompt_tokens.clone(), req.max_tokens, - req.params.ignore_eos, + req.stop_policy.clone(), cached_tokens, ); if drafter_enabled { @@ -517,7 +533,7 @@ fn admit_native( let mut state = Glm52SlotState::new( req.prompt_tokens.clone(), req.max_tokens, - req.params.ignore_eos, + req.stop_policy.clone(), handoff.committed_len, ); state.seed_native_pd_replayed_anchor(); @@ -679,6 +695,7 @@ mod tests { fingerprint: offload::handoff_fingerprint(), committed_len: PAGE, anchor_token_id: Some(11), + stop_cause: None, draft_tokens: vec![1, 2], }; let req = request(vec![10; PAGE], SamplingParams::default(), PAGE); diff --git a/pegainfer-glm52/src/scheduler/contract_tests.rs b/pegainfer-glm52/src/scheduler/contract_tests.rs index 2fab1fcec..3833cc5ac 100644 --- a/pegainfer-glm52/src/scheduler/contract_tests.rs +++ b/pegainfer-glm52/src/scheduler/contract_tests.rs @@ -68,7 +68,12 @@ fn load_snapshot_reports_the_ranks_own_state() { let mut slots: RankSlots = std::array::from_fn(|_| None); let req = request(vec![10, 11], SamplingParams::default(), 4); - let state = Glm52SlotState::new(req.prompt_tokens.clone(), req.max_tokens, true, 0); + let state = Glm52SlotState::new( + req.prompt_tokens.clone(), + req.max_tokens, + super::testkit::stop_policy(true), + 0, + ); let mut kv = pool.new_request(req.prompt_tokens.clone(), req.max_tokens, None); kv.schedule_prefill(1, &pool).expect("one live KV block"); slots[0] = Some(ActiveRequest { @@ -420,6 +425,7 @@ fn resolved_native( fingerprint: super::offload::handoff_fingerprint(), committed_len: committed.len(), anchor_token_id: Some(anchor), + stop_cause: None, draft_tokens: vec![anchor; crate::mtp::GLM52_MTP_DRAFTS], }; let mut req = request(committed, SamplingParams::default(), max_tokens); @@ -613,6 +619,7 @@ fn suppressed_eos_finishes_at_admission_without_a_slot() { fingerprint: super::offload::handoff_fingerprint(), committed_len: 3, anchor_token_id: None, + stop_cause: Some(super::offload::NativeStopCause::Eos(7)), draft_tokens: Vec::new(), }; let mut req = request(vec![10, 11, 12], SamplingParams::default(), 8); @@ -646,12 +653,20 @@ fn suppressed_eos_finishes_at_admission_without_a_slot() { rx.try_recv(), Ok((_, pegainfer_frontend::engine::TokenEvent::Scheduled { .. })) )); + assert!(matches!( + rx.try_recv(), + Ok(( + _, + pegainfer_frontend::engine::TokenEvent::Token { id: 7, .. } + )) + )); assert!(matches!( rx.try_recv(), Ok(( _, pegainfer_frontend::engine::TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Eos(7)), completion_tokens: 1, .. } @@ -673,6 +688,7 @@ fn suppressed_eos_finishes_behind_a_budget_stalled_front() { fingerprint: super::offload::handoff_fingerprint(), committed_len: 3, anchor_token_id: None, + stop_cause: Some(super::offload::NativeStopCause::Eos(7)), draft_tokens: Vec::new(), }; let mut req = request(vec![10, 11, 12], SamplingParams::default(), 8); @@ -713,12 +729,20 @@ fn suppressed_eos_finishes_behind_a_budget_stalled_front() { rx.try_recv(), Ok((_, pegainfer_frontend::engine::TokenEvent::Scheduled { .. })) )); + assert!(matches!( + rx.try_recv(), + Ok(( + _, + pegainfer_frontend::engine::TokenEvent::Token { id: 7, .. } + )) + )); assert!(matches!( rx.try_recv(), Ok(( _, pegainfer_frontend::engine::TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Eos(7)), completion_tokens: 1, .. } @@ -736,6 +760,7 @@ fn anchor_exhausting_max_tokens_finishes_as_length() { fingerprint: super::offload::handoff_fingerprint(), committed_len: 3, anchor_token_id: Some(70_001), + stop_cause: None, draft_tokens: Vec::new(), }; let mut req = request(vec![10, 11, 12], SamplingParams::default(), 1); @@ -852,7 +877,12 @@ fn drive_request( with_drafts: bool, ) -> Result<(), String> { let prompt: Vec = (0..prompt_len as u32).map(|t| 10_000 + t).collect(); - let mut state = Glm52SlotState::new(prompt.clone(), max_tokens, true, 0); + let mut state = Glm52SlotState::new( + prompt.clone(), + max_tokens, + super::testkit::stop_policy(true), + 0, + ); let mut kv = pool.new_request(prompt, max_tokens, None); let mut fresh = 60_000u32; loop { @@ -958,7 +988,7 @@ fn eos_truncated_speculative_apply_stays_in_contract() { // the release must both stay clean. let pool = Arc::new(BlockPool::new(PAGE, 16)); let prompt: Vec = (0..70).collect(); - let mut state = Glm52SlotState::new(prompt.clone(), 32, false, 0); + let mut state = Glm52SlotState::new(prompt.clone(), 32, super::testkit::stop_policy(false), 0); let mut kv = pool.new_request(prompt, 32, None); loop { if !state.mid_prefill() { @@ -988,14 +1018,19 @@ fn eos_truncated_speculative_apply_stays_in_contract() { committed, emit, finish, + stop_cause, .. } = outcome else { panic!("verify span must commit"); }; assert_eq!(committed, vec![21, 7], "truncated to the consumed run"); - assert_eq!(emit, 1, "the suppressed EOS is consumed, not emitted"); + assert_eq!(emit, 2, "the triggering EOS is retained exactly once"); assert_eq!(finish, Some(FinishReason::Stop)); + assert_eq!( + stop_cause, + Some(pegainfer_frontend::engine::StopCause::Eos(7)) + ); kv.apply_speculative(&committed, &pool) .expect("apply_speculative with the truncated run"); kv.release().expect("release"); diff --git a/pegainfer-glm52/src/scheduler/mod.rs b/pegainfer-glm52/src/scheduler/mod.rs index b58a40ce2..0e11e954c 100644 --- a/pegainfer-glm52/src/scheduler/mod.rs +++ b/pegainfer-glm52/src/scheduler/mod.rs @@ -668,8 +668,9 @@ impl Glm52Engine { self.runtime.spawn(async move { let resolved = match native { Some(handoff) => match handoff.anchor_token_id { - // P consumed EOS: nothing to restore or decode — the - // anchored finish happens at admission, no resolve runs. + // P consumed a token-driven stop: nothing to restore or + // decode — the typed finish happens at admission, no + // resolve runs. None => offload::Resolved::Native { req, prefix: KvPrefix::none(), @@ -1046,6 +1047,7 @@ impl Glm52Engine { committed, emit, finish, + stop_cause, context_rows, } => { // A dropped receiver (client disconnect) frees the @@ -1071,6 +1073,7 @@ impl Glm52Engine { { let _ = active.req.token_tx.send(TokenEvent::Finished { finish_reason, + stop_cause, prompt_tokens, completion_tokens: active.state.completion_tokens(), }); @@ -1401,6 +1404,7 @@ impl Glm52Engine { committed, emit, finish, + stop_cause, .. } => { active.kv.apply_prefill(committed[0], pool)?; @@ -1440,7 +1444,9 @@ impl Glm52Engine { let handoff = offload::NativeMtpHandoff { fingerprint: offload::handoff_fingerprint(), committed_len, - anchor_token_id: (emit == 1).then(|| committed[0]), + anchor_token_id: (stop_cause.is_none() && emit == 1) + .then(|| committed[0]), + stop_cause: stop_cause.map(Into::into), draft_tokens: drafts.to_vec(), }; let _ = active.req.token_tx.send(TokenEvent::KvTransfer { @@ -1456,6 +1462,7 @@ impl Glm52Engine { { let _ = active.req.token_tx.send(TokenEvent::Finished { finish_reason, + stop_cause, prompt_tokens, completion_tokens: active.state.completion_tokens(), }); diff --git a/pegainfer-glm52/src/scheduler/offload.rs b/pegainfer-glm52/src/scheduler/offload.rs index f1c8624a0..b697f5cfd 100644 --- a/pegainfer-glm52/src/scheduler/offload.rs +++ b/pegainfer-glm52/src/scheduler/offload.rs @@ -11,6 +11,7 @@ use anyhow::Context as _; use pegainfer_frontend::engine::GenerateRequest; use pegainfer_frontend::engine::KvPrefix; +use pegainfer_frontend::engine::StopCause; use pegainfer_kv_store::CacheScope; use pegainfer_kv_store::CancelProbe; use pegainfer_kv_store::KvStore; @@ -21,17 +22,44 @@ use serde::Serialize; use super::PAGE; -/// The handoff envelope (v3), one struct for both sides. Boundary rule: +/// The handoff envelope (v5), one struct for both sides. Boundary rule: /// content shareable in the radix is a pure prompt function and travels /// the KV data plane; anchor-dependent continuation state travels here. +#[derive(Clone, Copy, Debug, Serialize, Deserialize)] +pub(super) enum NativeStopCause { + Eos(u32), + Token(u32), +} + +impl From for NativeStopCause { + fn from(value: StopCause) -> Self { + match value { + StopCause::Eos(id) => Self::Eos(id), + StopCause::Token(id) => Self::Token(id), + } + } +} + +impl From for StopCause { + fn from(value: NativeStopCause) -> Self { + match value { + NativeStopCause::Eos(id) => Self::Eos(id), + NativeStopCause::Token(id) => Self::Token(id), + } + } +} + #[derive(Clone, Debug, Serialize, Deserialize)] pub(super) struct NativeMtpHandoff { /// Capability manifest (see [`handoff_fingerprint`]); a mismatch is an /// intake rejection carrying both strings. pub(super) fingerprint: String, pub(super) committed_len: usize, - /// P's first sampled token; `None` = EOS — nothing to restore or decode. + /// P's first sampled non-terminal token; `None` means no decode slot is + /// required (the typed terminal cause is carried separately). pub(super) anchor_token_id: Option, + #[serde(default)] + pub(super) stop_cause: Option, pub(super) draft_tokens: Vec, } @@ -47,7 +75,7 @@ pub(super) struct PegaInferPdEnvelope { /// whole per-block byte layout. pub(super) fn handoff_fingerprint() -> String { format!( - "glm52-native-mtp/4/page:{}/salt:{}/drafts:{}", + "glm52-native-mtp/5/page:{}/salt:{}/drafts:{}", crate::model::GLM52_KV_PAGE_STRIDE, super::native_mtp_cache_salt(), crate::mtp::glm52_mtp_draft_len(), @@ -75,7 +103,7 @@ pub(super) fn native_mtp_handoff( ); anyhow::ensure!( req.prompt_tokens.len() == handoff.committed_len, - "native P/D v3 expects the original prompt: committed_len {}, got {} prompt tokens", + "native P/D v5 expects the original prompt: committed_len {}, got {} prompt tokens", handoff.committed_len, req.prompt_tokens.len() ); diff --git a/pegainfer-glm52/src/scheduler/plan.rs b/pegainfer-glm52/src/scheduler/plan.rs index 77e2ab0c8..cb3665673 100644 --- a/pegainfer-glm52/src/scheduler/plan.rs +++ b/pegainfer-glm52/src/scheduler/plan.rs @@ -304,7 +304,12 @@ mod tests { /// prompt token is already fed, so `feed_want() == 1`). fn decoding_rank(params: pegainfer_sample::SamplingParams) -> RankSlots { let req = request(vec![10], params, 8); - let mut state = Glm52SlotState::new(req.prompt_tokens.clone(), req.max_tokens, false, 0); + let mut state = Glm52SlotState::new( + req.prompt_tokens.clone(), + req.max_tokens, + super::super::testkit::stop_policy(false), + 0, + ); assert!(matches!( state.advance_span(&[20], &[]), Glm52StepOutcome::Commit { .. } diff --git a/pegainfer-glm52/src/scheduler/slot.rs b/pegainfer-glm52/src/scheduler/slot.rs index f61049c42..33c2a41c7 100644 --- a/pegainfer-glm52/src/scheduler/slot.rs +++ b/pegainfer-glm52/src/scheduler/slot.rs @@ -4,6 +4,8 @@ //! ([`Glm52SlotState::advance_span`]). use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use crate::dspark::GLM52_DSPARK_DRAFTS; use crate::dspark::accept_prefix_match; @@ -85,15 +87,17 @@ pub(super) enum Glm52StepOutcome { /// Mid-prefill: the model's outputs are discarded, keep feeding the prompt. Prefilling, /// Commit the span's agreed tokens: `committed` is the consumed run - /// (what advances the request's KV bookkeeping — a suppressed EOS is its - /// last entry), of which the leading `emit` tokens are sent to the - /// client, then finish if `finish` is set. Plain decode commits exactly - /// one token; a verify span commits the accepted draft prefix plus the + /// (what advances the request's KV bookkeeping), of which the leading + /// `emit` tokens are sent to the client, then finish if `finish` is set. + /// A token-driven finish includes its trigger in both `committed` and + /// `emit`, with the typed cause retained separately. Plain decode commits + /// exactly one token; a verify span commits the accepted draft prefix plus the /// model's correction or bonus token (1..=span tokens). Commit { committed: Vec, emit: usize, finish: Option, + stop_cause: Option, /// Leading span rows whose tokens are now committed context for the /// draft lane: all rows of a prompt span; anchor + accepted drafts of /// a verify span (rejected rows' captured hidden is dead). @@ -109,10 +113,10 @@ pub(super) enum Glm52StepOutcome { pub(super) struct Glm52SlotState { prompt: Vec, max_tokens: usize, - ignore_eos: bool, + stop_policy: StopPolicy, /// Prompt tokens already fed to the model. fed: usize, - /// Generated tokens (a suppressed EOS counts). + /// Generated tokens, including a token that triggers termination. completion: usize, /// Sampling steps already consumed outside this D request's client-visible /// completion budget. Manual native-P/D uses one: P sampled the forwarded @@ -155,7 +159,7 @@ impl Glm52SlotState { pub(super) fn new( prompt: Vec, max_tokens: usize, - ignore_eos: bool, + stop_policy: StopPolicy, cached_tokens: usize, ) -> Self { debug_assert!(cached_tokens < prompt.len()); @@ -163,7 +167,7 @@ impl Glm52SlotState { fed: cached_tokens, prompt, max_tokens, - ignore_eos, + stop_policy, completion: 0, sampling_offset: 0, last_token: 0, @@ -391,28 +395,32 @@ impl Glm52SlotState { let mut committed = committed; let mut emit = 0usize; let mut finish = None; + let mut stop_cause = None; for &token in &committed { self.completion += 1; - if !self.ignore_eos && eos_token_ids.contains(&token) { + emit += 1; + self.last_token = token; + if let Some(cause) = self + .stop_policy + .classify(token, |id| eos_token_ids.contains(&id)) + { finish = Some(FinishReason::Stop); + stop_cause = Some(cause); break; } - emit += 1; - self.last_token = token; if self.completion >= self.max_tokens { finish = Some(FinishReason::Length); break; } } - // Truncate to the consumed run (a suppressed EOS is consumed but not - // emitted) so the caller's KV bookkeeping advances by exactly the - // tokens this state accounted for. - let consumed = emit + usize::from(matches!(finish, Some(FinishReason::Stop))); - committed.truncate(consumed); + // Truncate after the first terminal token so speculative suffixes are + // neither committed to KV nor exposed to the client. + committed.truncate(emit); Glm52StepOutcome::Commit { committed, emit, finish, + stop_cause, context_rows, } } @@ -450,6 +458,7 @@ mod tests { use crate::scheduler::testkit::EOS; use crate::scheduler::testkit::commit; use crate::scheduler::testkit::state; + use crate::scheduler::testkit::stop_policy; #[test] fn prefill_rides_decode_then_emits() { @@ -567,11 +576,11 @@ mod tests { } #[test] - fn eos_is_suppressed_and_counts_toward_completion() { + fn eos_is_emitted_and_counts_toward_completion() { let mut state = state(vec![10], 4, false); assert_eq!( state.advance_span(&[7], EOS), - commit(&[7], 0, Some(FinishReason::Stop), 1) + commit(&[7], 1, Some(FinishReason::Stop), 1) ); assert_eq!(state.completion_tokens(), 1); } @@ -605,7 +614,55 @@ mod tests { let mut state = state(vec![10], 1, false); assert_eq!( state.advance_span(&[7], EOS), - commit(&[7], 0, Some(FinishReason::Stop), 1) + commit(&[7], 1, Some(FinishReason::Stop), 1) + ); + } + + #[test] + fn ignored_eos_does_not_disable_an_explicit_stop_token() { + let mut state = Glm52SlotState::new( + vec![10], + 4, + StopPolicy { + eos: pegainfer_frontend::engine::EosPolicy::Ignore, + token_ids: vec![42], + }, + 0, + ); + + assert_eq!( + state.advance_span(&[42], EOS), + Glm52StepOutcome::Commit { + committed: vec![42], + emit: 1, + finish: Some(FinishReason::Stop), + stop_cause: Some(StopCause::Token(42)), + context_rows: 1, + } + ); + } + + #[test] + fn eos_precedes_an_overlapping_explicit_stop_token() { + let mut state = Glm52SlotState::new( + vec![10], + 4, + StopPolicy { + eos: pegainfer_frontend::engine::EosPolicy::Token(7), + token_ids: vec![7], + }, + 0, + ); + + assert_eq!( + state.advance_span(&[7], EOS), + Glm52StepOutcome::Commit { + committed: vec![7], + emit: 1, + finish: Some(FinishReason::Stop), + stop_cause: Some(StopCause::Eos(7)), + context_rows: 1, + } ); } @@ -666,7 +723,7 @@ mod tests { #[test] fn native_pd_starts_by_verifying_the_forwarded_anchor() { - let mut state = Glm52SlotState::new(vec![10, 11, 20], 8, false, 2); + let mut state = Glm52SlotState::new(vec![10, 11, 20], 8, stop_policy(false), 2); state.seed_native_pd_anchor(); state.set_drafts(vec![21, 22, 99, 98, 97], GLM52_MTP_DRAFTS); @@ -700,7 +757,7 @@ mod tests { #[test] fn manual_native_pd_advances_rng_without_spending_client_budget() { - let mut state = Glm52SlotState::new(vec![10, 11, 12, 42], 4, true, 3); + let mut state = Glm52SlotState::new(vec![10, 11, 12, 42], 4, stop_policy(true), 3); state.seed_native_pd_anchor(); assert_eq!( @@ -720,7 +777,7 @@ mod tests { #[test] fn native_pd_replayed_anchor_counts_against_the_router_budget() { - let mut state = Glm52SlotState::new(vec![10, 11, 20], 8, false, 2); + let mut state = Glm52SlotState::new(vec![10, 11, 20], 8, stop_policy(false), 2); state.seed_native_pd_replayed_anchor(); state.set_drafts(vec![21, 22, 99, 98, 97], GLM52_MTP_DRAFTS); @@ -770,7 +827,7 @@ mod tests { fn eos_inside_the_committed_run_truncates_and_finishes() { let mut state = state(vec![10], 32, false); assert_eq!(state.advance_span(&[20], EOS), commit(&[20], 1, None, 1)); - // Draft 2 is the EOS token (7): accepted, counted, suppressed; the + // Draft 2 is the EOS token (7): accepted, counted, emitted; the // rest of the committed run is dropped. state.set_drafts( vec![21, 7, 23, 24, 25, 26, 27], @@ -779,7 +836,7 @@ mod tests { let outputs = [21, 7, 23, 24]; assert_eq!( state.advance_span(&outputs, EOS), - commit(&[21, 7], 1, Some(FinishReason::Stop), 4) + commit(&[21, 7], 2, Some(FinishReason::Stop), 4) ); assert_eq!(state.completion_tokens(), 3); } @@ -898,7 +955,7 @@ mod tests { // 3 blocks of prompt with the first 2 cache-hit: feeding starts at // position 128 and only the suffix is ever fed. let prompt: Vec = (0..192).collect(); - let s = Glm52SlotState::new(prompt, 8, false, 128); + let s = Glm52SlotState::new(prompt, 8, stop_policy(false), 128); assert_eq!(s.feed_want(), 64); assert_eq!(s.next_input_at(0).position, 128); assert_eq!(s.next_input_at(0).token, 128); diff --git a/pegainfer-glm52/src/scheduler/testkit.rs b/pegainfer-glm52/src/scheduler/testkit.rs index b461f3808..ec78b0be1 100644 --- a/pegainfer-glm52/src/scheduler/testkit.rs +++ b/pegainfer-glm52/src/scheduler/testkit.rs @@ -1,7 +1,10 @@ //! Shared fixtures for the scheduler module tests. +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_kv_store::BlockPool; use pegainfer_kv_store::RequestKv; @@ -12,7 +15,18 @@ use super::slot::Glm52StepOutcome; pub(super) const EOS: &[u32] = &[7]; pub(super) fn state(prompt: Vec, max_tokens: usize, ignore_eos: bool) -> Glm52SlotState { - Glm52SlotState::new(prompt, max_tokens, ignore_eos, 0) + Glm52SlotState::new(prompt, max_tokens, stop_policy(ignore_eos), 0) +} + +pub(super) fn stop_policy(ignore_eos: bool) -> StopPolicy { + StopPolicy { + eos: if ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + token_ids: Vec::new(), + } } /// A standalone `RequestKv` for tests that never schedule KV (the pool @@ -28,10 +42,13 @@ pub(super) fn commit( finish: Option, context_rows: usize, ) -> Glm52StepOutcome { + let stop_cause = matches!(finish, Some(FinishReason::Stop)) + .then(|| StopCause::Eos(*committed.last().expect("a stop commit has a trigger token"))); Glm52StepOutcome::Commit { committed: committed.to_vec(), emit, finish, + stop_cause, context_rows, } } @@ -42,6 +59,7 @@ pub(super) fn request( max_tokens: usize, ) -> GenerateRequest { let (token_tx, _token_rx) = pegainfer_frontend::engine::TokenSink::standalone(); + let stop_policy = stop_policy(params.ignore_eos); GenerateRequest { trace_parent: None, request_id: None, @@ -49,6 +67,7 @@ pub(super) fn request( data_parallel_rank: None, prompt_tokens: prompt, params, + stop_policy, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-k3/src/scheduler/mod.rs b/pegainfer-k3/src/scheduler/mod.rs index 9c851c650..94728f0e8 100644 --- a/pegainfer-k3/src/scheduler/mod.rs +++ b/pegainfer-k3/src/scheduler/mod.rs @@ -33,6 +33,8 @@ use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestLedger; use pegainfer_frontend::engine::Scheduler; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::spawn_scheduler; pub use self::executor::DecodeSlot; @@ -44,8 +46,9 @@ pub use self::executor::StepExecutor; /// Scheduler facts that come from the model line rather than the executor. #[derive(Clone, Debug, Default)] pub struct K3SchedulerConfig { - /// Token ids that end a stream with [`FinishReason::Stop`]. Requests that - /// set `ignore_eos` opt out. + /// Model-default token ids that can end a stream with + /// [`FinishReason::Stop`]. A request's [`StopPolicy`] decides whether the + /// model EOS set is active and which explicit stop ids also apply. pub eos_token_ids: Vec, /// KV pool capacity to advertise, or `None` from an engine that does not /// own a pool yet. Injected rather than asked of the executor so the @@ -91,11 +94,16 @@ where /// Answer a request that just reached its end: silence if the frontend /// abandoned it since the scheduler last looked, the finish otherwise. Abort /// can land at any moment, so the check belongs on every finish path. -fn finish_or_retire(id: RequestId, reason: FinishReason, ledger: &mut RequestLedger) { +fn finish_or_retire( + id: RequestId, + reason: FinishReason, + stop_cause: Option, + ledger: &mut RequestLedger, +) { if ledger.is_aborted(id) { ledger.retire(id); } else { - ledger.finish(id, reason); + ledger.finish(id, reason, stop_cause); } } @@ -106,7 +114,7 @@ struct RunningRequest { /// This request's most recent committed token — next step's input. last_token: u32, max_tokens: usize, - ignore_eos: bool, + stop_policy: StopPolicy, } pub struct K3Scheduler { @@ -156,9 +164,10 @@ impl K3Scheduler { ) } - /// Whether `token` ends this request's stream by end-of-sequence. - fn is_stop_token(&self, token: u32, ignore_eos: bool) -> bool { - !ignore_eos && self.eos_token_ids.contains(&token) + /// Classify a generated token under the request's independent EOS and + /// explicit stop-token policy. EOS wins when an id appears in both sets. + fn stop_cause(&self, token: u32, policy: &StopPolicy) -> Option { + policy.classify(token, |id| self.eos_token_ids.contains(&id)) } /// Fill free slots from the queue: retire what the frontend abandoned, @@ -180,7 +189,7 @@ impl K3Scheduler { ledger.admit(id); if pending.request.max_tokens == 0 { // Nothing to generate: answer without occupying a slot. - finish_or_retire(id, FinishReason::Length, ledger); + finish_or_retire(id, FinishReason::Length, None, ledger); continue; } let slot = self @@ -207,18 +216,21 @@ impl K3Scheduler { slot, last_token: first, max_tokens: pending.request.max_tokens, - ignore_eos: pending.request.params.ignore_eos, + stop_policy: pending.request.stop_policy, }; - if self.is_stop_token(first, state.ignore_eos) { - // The stop token itself is not part of the completion. + if let Some(stop_cause) = self.stop_cause(first, &state.stop_policy) { + // Keep the triggering token in the contract update. The vLLM + // bridge uses it for usage/logprob accounting and maps only + // explicit request stops to a wire stop reason. + ledger.push_tokens(id, &[first], &[]); self.release_slot(slot); - finish_or_retire(id, FinishReason::Stop, ledger); + finish_or_retire(id, FinishReason::Stop, Some(stop_cause), ledger); continue; } ledger.push_tokens(id, &[first], &[]); if ledger.completion_tokens(id) >= state.max_tokens { self.release_slot(slot); - finish_or_retire(id, FinishReason::Length, ledger); + finish_or_retire(id, FinishReason::Length, None, ledger); continue; } self.running.push(state); @@ -274,17 +286,20 @@ impl K3Scheduler { // stop/length cut are computed-but-dead, like a rejected draft. let already = ledger.completion_tokens(state.id); let mut kept: Vec = Vec::with_capacity(committed.len()); - let mut finished = None; + let mut finished: Option<(FinishReason, Option)> = None; for &token in &committed { - if self.is_stop_token(token, state.ignore_eos) { - // The stop token itself is not part of the completion. - finished = Some(FinishReason::Stop); + if let Some(stop_cause) = self.stop_cause(token, &state.stop_policy) { + // The trigger is a real completion token. Suffix tokens + // from the speculative round are computed-but-dead. + kept.push(token); + state.last_token = token; + finished = Some((FinishReason::Stop, Some(stop_cause))); break; } kept.push(token); state.last_token = token; if already + kept.len() >= state.max_tokens { - finished = Some(FinishReason::Length); + finished = Some((FinishReason::Length, None)); break; } } @@ -292,9 +307,9 @@ impl K3Scheduler { ledger.push_tokens(state.id, &kept, &[]); } match finished { - Some(reason) => { + Some((reason, stop_cause)) => { self.release_slot(state.slot); - finish_or_retire(state.id, reason, ledger); + finish_or_retire(state.id, reason, stop_cause, ledger); } None => still_running.push(state), } diff --git a/pegainfer-k3/src/scheduler/tests.rs b/pegainfer-k3/src/scheduler/tests.rs index 2d20d9e1e..a160095a4 100644 --- a/pegainfer-k3/src/scheduler/tests.rs +++ b/pegainfer-k3/src/scheduler/tests.rs @@ -13,6 +13,7 @@ use std::time::Instant; use anyhow::Result; use pegainfer_frontend::engine::Engine; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::LiveScheduler; use pegainfer_frontend::engine::RejectReason; @@ -20,6 +21,8 @@ use pegainfer_frontend::engine::Request; use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestUpdate; use pegainfer_frontend::engine::StepReceiver; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; @@ -47,6 +50,10 @@ struct FakeExecutor { /// Emit [`EOS_TOKEN`] instead of the scripted token at this step index /// (0 = at prefill). eos_at: Option, + /// Optional multi-token spans returned by a speculative executor. The + /// scheduler must classify each span in order and discard the suffix + /// after the first terminal token. + decode_spans: VecDeque>, fail_next_prefill: bool, fail_next_decode: bool, decode_delay: Duration, @@ -61,6 +68,7 @@ impl FakeExecutor { steps: HashMap::new(), live: HashSet::new(), eos_at: None, + decode_spans: VecDeque::new(), fail_next_prefill: false, fail_next_decode: false, decode_delay: Duration::ZERO, @@ -78,6 +86,11 @@ impl FakeExecutor { self } + fn with_decode_spans(mut self, spans: impl IntoIterator>) -> Self { + self.decode_spans = spans.into_iter().collect(); + self + } + fn with_one_prefill_failure(mut self) -> Self { self.fail_next_prefill = true; self @@ -152,6 +165,31 @@ impl StepExecutor for FakeExecutor { .collect()) } + fn decode_many(&mut self, batch: &[DecodeSlot]) -> Result>> { + if self.decode_spans.is_empty() { + return self + .decode(batch) + .map(|tokens| tokens.into_iter().map(|token| vec![token]).collect()); + } + assert_eq!(batch.len(), 1, "scripted spans are single-request fixtures"); + let entry = batch[0]; + assert!( + self.live.contains(&entry.slot), + "slot {} decoded after release", + entry.slot + ); + let span = self + .decode_spans + .pop_front() + .expect("scripted span queue unexpectedly empty"); + assert!(!span.is_empty(), "a scripted decode span must not be empty"); + self.steps + .entry(entry.slot) + .and_modify(|step| *step += span.len() as u32) + .or_insert(span.len() as u32); + Ok(vec![span]) + } + fn release(&mut self, slot: SlotId) { assert!(self.live.remove(&slot), "slot {slot} released while free"); self.steps.remove(&slot); @@ -165,6 +203,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, @@ -299,6 +338,7 @@ fn admitted_request_streams_its_tokens_and_finishes_at_max_tokens() { terminal, Terminal::Finished { reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 4, completion_tokens: 3, } @@ -345,7 +385,7 @@ fn a_zero_length_completion_finishes_without_taking_a_slot() { } #[test] -fn eos_finishes_with_stop_and_is_not_streamed() { +fn eos_finishes_with_stop_and_retains_the_triggering_token() { let released = Arc::new(Mutex::new(Vec::new())); let executor = FakeExecutor::new(4, released).with_eos_at(2); let (partition, mut steps) = launch(executor); @@ -354,14 +394,78 @@ fn eos_finishes_with_stop_and_is_not_streamed() { let (tokens, terminal) = steps.collect_terminal(control.id()); assert_eq!( tokens, - vec![10, 11], - "the stop token itself is not part of the completion" + vec![10, 11, EOS_TOKEN], + "the stop token stays in the engine update for wire accounting" + ); + assert!( + matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(EOS_TOKEN)), + completion_tokens: 3, + .. + } + ), + "{terminal:?}" + ); +} + +#[test] +fn prefill_eos_finishes_and_releases_the_slot_for_reuse() { + let released = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(1, Arc::clone(&released)).with_eos_at(0); + let (partition, mut steps) = launch(executor); + + for _ in 0..2 { + let control = partition.handle.submit(request(4, 64)); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!(tokens, vec![EOS_TOKEN]); + + assert!( + matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(EOS_TOKEN)), + completion_tokens: 1, + .. + } + ), + "{terminal:?}" + ); + } + + assert_eq!( + *released.lock().expect("released log"), + vec![0, 0], + "the only slot must be reusable after a prefill stop" ); +} + +#[test] +fn explicit_stop_token_finishes_independently_of_eos() { + let released = Arc::new(Mutex::new(Vec::new())); + let (partition, mut steps) = launch(FakeExecutor::new(4, released)); + + let mut req = request(4, 64); + req.stop_policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![11], + }; + + let control = partition.handle.submit(req); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!(tokens, vec![10, 11]); + assert!( matches!( terminal, Terminal::Finished { reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(11)), completion_tokens: 2, .. } @@ -370,6 +474,34 @@ fn eos_finishes_with_stop_and_is_not_streamed() { ); } +#[test] +fn ignored_eos_is_streamed_and_length_finishes_the_request() { + let released = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(4, released).with_eos_at(2); + let (partition, mut steps) = launch(executor); + + let mut req = request(4, 3); + req.stop_policy.eos = EosPolicy::Ignore; + + let control = partition.handle.submit(req); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!(tokens, vec![10, 11, EOS_TOKEN]); + + assert!( + matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Length, + stop_cause: None, + completion_tokens: 3, + .. + } + ), + "{terminal:?}" + ); +} + #[test] fn oversized_request_is_rejected_without_taking_a_slot() { let released = Arc::new(Mutex::new(Vec::new())); @@ -539,3 +671,101 @@ fn requests_beyond_the_slot_budget_wait_instead_of_being_refused() { ); } } + +#[test] +fn speculative_span_keeps_the_first_explicit_stop_and_discards_suffix() { + let released = Arc::new(Mutex::new(Vec::new())); + let executor = + FakeExecutor::new(1, Arc::clone(&released)).with_decode_spans([vec![20, 42, 77]]); + let (partition, mut steps) = launch(executor); + + let mut req = request(4, 64); + req.stop_policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![42], + }; + let control = partition.handle.submit(req); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!(tokens, vec![10, 20, 42]); + assert!( + matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(42)), + completion_tokens: 3, + .. + } + ), + "{terminal:?}" + ); + assert!( + wait_until(Duration::from_secs(1), || released + .lock() + .expect("released log") + .contains(&0)), + "a speculative stop must release its slot" + ); +} + +#[test] +fn speculative_span_stops_at_max_tokens_without_a_stop_cause() { + let released = Arc::new(Mutex::new(Vec::new())); + let executor = + FakeExecutor::new(1, Arc::clone(&released)).with_decode_spans([vec![20, 21, 22]]); + let (partition, mut steps) = launch(executor); + + let control = partition.handle.submit(request(4, 2)); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!(tokens, vec![10, 20]); + assert!( + matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Length, + stop_cause: None, + completion_tokens: 2, + .. + } + ), + "{terminal:?}" + ); + assert!( + wait_until(Duration::from_secs(1), || released + .lock() + .expect("released log") + .contains(&0)), + "a length-truncated speculative span must release its slot" + ); +} + +#[test] +fn speculative_eos_wins_over_an_overlapping_explicit_stop() { + let released = Arc::new(Mutex::new(Vec::new())); + let executor = FakeExecutor::new(1, released).with_decode_spans([vec![42, 77]]); + let (partition, mut steps) = launch(executor); + + let mut req = request(4, 64); + req.stop_policy = StopPolicy { + eos: EosPolicy::Token(42), + token_ids: vec![42], + }; + let control = partition.handle.submit(req); + let (tokens, terminal) = steps.collect_terminal(control.id()); + + assert_eq!(tokens, vec![10, 42]); + assert!( + matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(42)), + completion_tokens: 2, + .. + } + ), + "{terminal:?}" + ); +} diff --git a/pegainfer-kimi-k2/src/batch_decode_trace.rs b/pegainfer-kimi-k2/src/batch_decode_trace.rs index fec41a657..fa9a19ab4 100644 --- a/pegainfer-kimi-k2/src/batch_decode_trace.rs +++ b/pegainfer-kimi-k2/src/batch_decode_trace.rs @@ -5,8 +5,12 @@ use pegainfer_core::ops::call_trace; #[cfg(feature = "kernel-call-trace")] use pegainfer_frontend::engine::EngineLoadOptions; #[cfg(feature = "kernel-call-trace")] +use pegainfer_frontend::engine::EosPolicy; +#[cfg(feature = "kernel-call-trace")] use pegainfer_frontend::engine::GenerateRequest; #[cfg(feature = "kernel-call-trace")] +use pegainfer_frontend::engine::StopPolicy; +#[cfg(feature = "kernel-call-trace")] use pegainfer_frontend::engine::TokenEvent; #[cfg(feature = "kernel-call-trace")] use pegainfer_frontend::engine::TokenSink; @@ -152,6 +156,10 @@ pub fn trace_runtime_decode_kernel_calls( seed: None, ignore_eos: true, }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, max_tokens: 2, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-kimi-k2/src/runner/scheduler.rs b/pegainfer-kimi-k2/src/runner/scheduler.rs index f5ed9d1af..4aebe9fbb 100644 --- a/pegainfer-kimi-k2/src/runner/scheduler.rs +++ b/pegainfer-kimi-k2/src/runner/scheduler.rs @@ -14,8 +14,10 @@ use lifecycle::validate_kv_capacity; use log::error; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::SubmittedRequest; use pegainfer_frontend::engine::TokenEvent; +use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::TokenSink; use pegainfer_kv_cache::BlockPool; use pegainfer_kv_cache::RequestKv; @@ -43,8 +45,55 @@ fn row_options(req: &GenerateRequest) -> KimiRowOptions { } } +/// Deliver a sampled token before its terminal event so the protocol layer +/// receives the exact trigger token and logprob. Stop classification stays on +/// the host and does not change the GPU sampling row or CUDA graph inputs. +fn emit_sampled_token( + token_tx: &TokenSink, + stop_policy: &StopPolicy, + model_eos: &[u32], + token_id: u32, + logprob: Option, + prompt_tokens: usize, + completion_tokens: usize, + max_tokens: usize, +) -> bool { + if token_tx + .send(TokenEvent::Token { + id: token_id, + logprob, + }) + .is_err() + { + return false; + } + + if let Some(stop_cause) = stop_policy.classify(token_id, |id| model_eos.contains(&id)) { + let _ = token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), + prompt_tokens, + completion_tokens, + }); + return false; + } + + if completion_tokens >= max_tokens { + let _ = token_tx.send(TokenEvent::Finished { + finish_reason: FinishReason::Length, + stop_cause: None, + prompt_tokens, + completion_tokens, + }); + return false; + } + + true +} + struct ActiveKimiRequest { token_tx: TokenSink, + stop_policy: StopPolicy, prompt_len: usize, completion_tokens: usize, max_tokens: usize, @@ -328,34 +377,16 @@ impl KimiK2Scheduler { retire.push(idx); continue; } - // EOS outranks the length limit; the stop token itself is not - // emitted (same contract as the Qwen schedulers). - if !req.options.sampling.ignore_eos && self.stop_token_ids.contains(&token_id) { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_len, - completion_tokens: req.completion_tokens, - }); - retire.push(idx); - continue; - } - if req - .token_tx - .send(TokenEvent::Token { - id: token_id, - logprob: report.logprob, - }) - .is_err() - { - retire.push(idx); - continue; - } - if req.completion_tokens >= req.max_tokens { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_len, - completion_tokens: req.completion_tokens, - }); + if !emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &self.stop_token_ids, + token_id, + report.logprob, + req.prompt_len, + req.completion_tokens, + req.max_tokens, + ) { retire.push(idx); } else { req.last_token = token_id; @@ -441,22 +472,16 @@ impl KimiK2Scheduler { }); return None; } - if !req.params.ignore_eos && self.stop_token_ids.contains(&token_id) { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - return None; - } - if req - .token_tx - .send(TokenEvent::Token { - id: token_id, - logprob: report.logprob, - }) - .is_err() - { + if !emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &self.stop_token_ids, + token_id, + report.logprob, + req.prompt_tokens.len(), + 1, + req.max_tokens, + ) { return None; } token_id @@ -476,17 +501,10 @@ impl KimiK2Scheduler { } }; let completion_tokens = completion_tokens + 1; - if completion_tokens >= req.max_tokens { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens, - }); - return None; - } let options = row_options(&req); Some(ActiveKimiRequest { token_tx: req.token_tx, + stop_policy: req.stop_policy, prompt_len: req.prompt_tokens.len(), completion_tokens, max_tokens: req.max_tokens, @@ -576,36 +594,23 @@ impl KimiK2Scheduler { }); continue; } - if !req.params.ignore_eos && self.stop_token_ids.contains(&token_id) { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - continue; - } - if req - .token_tx - .send(TokenEvent::Token { - id: token_id, - logprob: report.logprob, - }) - .is_err() - { - continue; - } let completion_tokens = 1usize; - if completion_tokens >= req.max_tokens { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens, - }); + if !emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &self.stop_token_ids, + token_id, + report.logprob, + req.prompt_tokens.len(), + completion_tokens, + req.max_tokens, + ) { continue; } let options = row_options(&req); active.push(ActiveKimiRequest { token_tx: req.token_tx, + stop_policy: req.stop_policy, prompt_len: req.prompt_tokens.len(), completion_tokens, max_tokens: req.max_tokens, @@ -624,6 +629,8 @@ mod tests { use std::sync::Arc; use std::sync::Mutex; + use pegainfer_frontend::engine::EosPolicy; + use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::sampler::SamplingParams; use super::*; @@ -775,6 +782,7 @@ mod tests { data_parallel_rank: None, prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -785,6 +793,75 @@ mod tests { (req, token_rx) } + #[test] + fn sampled_model_eos_is_emitted_before_stop_finish() { + let (req, mut token_rx) = request_with_channel(vec![11, 22], 4); + let logprob = TokenLogprob { + logprob: -0.5, + top_logprobs: vec![(99, -0.5)], + }; + + assert!(!emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &[99], + 99, + Some(logprob.clone()), + 2, + 1, + req.max_tokens, + )); + + let Ok((_, TokenEvent::Token { id, logprob: seen })) = token_rx.try_recv() else { + panic!("expected trigger Token event"); + }; + assert_eq!(id, 99); + assert_eq!(seen, Some(logprob)); + let Ok(( + _, + TokenEvent::Finished { + finish_reason, + stop_cause, + completion_tokens, + .. + }, + )) = token_rx.try_recv() + else { + panic!("expected terminal Finished event"); + }; + assert_eq!(finish_reason, FinishReason::Stop); + assert_eq!(stop_cause, Some(StopCause::Eos(99))); + assert_eq!(completion_tokens, 1); + } + + #[test] + fn ignored_eos_still_honors_request_stop_token() { + let (mut req, mut token_rx) = request_with_channel(vec![11, 22], 4); + req.stop_policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![99], + }; + + assert!(!emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &[99], + 99, + None, + 2, + 1, + req.max_tokens, + )); + assert!(matches!( + token_rx.try_recv(), + Ok((_, TokenEvent::Token { id: 99, .. })) + )); + let Ok((_, TokenEvent::Finished { stop_cause, .. })) = token_rx.try_recv() else { + panic!("expected terminal Finished event"); + }; + assert_eq!(stop_cause, Some(StopCause::Token(99))); + } + #[test] fn mixed_prompt_batch_routes_single_token_requests_to_decode() { let calls = Arc::new(Mutex::new(Vec::new())); diff --git a/pegainfer-kimi-k2/src/runner/scheduler/dp.rs b/pegainfer-kimi-k2/src/runner/scheduler/dp.rs index 33ce71e87..536e43597 100644 --- a/pegainfer-kimi-k2/src/runner/scheduler/dp.rs +++ b/pegainfer-kimi-k2/src/runner/scheduler/dp.rs @@ -5,6 +5,7 @@ use crossbeam_channel::bounded; use log::error; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::SubmittedRequest; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; @@ -13,6 +14,7 @@ use pegainfer_kv_cache::RequestKv; use rand::rngs::StdRng; use tokio::sync::mpsc; +use super::emit_sampled_token; use super::lifecycle::preflight_prefill_candidate; use super::lifecycle::request_lifetime_blocks; use super::lifecycle::send_scheduled; @@ -58,6 +60,7 @@ pub(in crate::runner) struct DpRankState { struct RequestState { token_tx: TokenSink, + stop_policy: StopPolicy, prompt_len: usize, completion_tokens: usize, max_tokens: usize, @@ -512,38 +515,24 @@ impl DpCoordinator { }); return; } - if !req.params.ignore_eos && self.stop_token_ids.contains(&last_token) { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: prompt_len, - completion_tokens: 0, - }); - return; - } - if req - .token_tx - .send(TokenEvent::Token { - id: last_token, - logprob: owner_report.logprob, - }) - .is_err() - { - return; - } - let completion_tokens = 1; - if completion_tokens >= req.max_tokens { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: prompt_len, - completion_tokens, - }); + if !emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &self.stop_token_ids, + last_token, + owner_report.logprob, + prompt_len, + completion_tokens, + req.max_tokens, + ) { return; } let options = row_options(&req); self.ranks[dp_rank].slots[slot] = Some(RequestState { token_tx: req.token_tx, + stop_policy: req.stop_policy, prompt_len, completion_tokens, max_tokens: req.max_tokens, @@ -647,38 +636,24 @@ impl DpCoordinator { }); return; } - if !req.params.ignore_eos && self.stop_token_ids.contains(&token_id) { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - return; - } - if req - .token_tx - .send(TokenEvent::Token { - id: token_id, - logprob: report.logprob.clone(), - }) - .is_err() - { - return; - } - let completion_tokens = 1; - if completion_tokens >= req.max_tokens { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens, - }); + if !emit_sampled_token( + &req.token_tx, + &req.stop_policy, + &self.stop_token_ids, + token_id, + report.logprob.clone(), + req.prompt_tokens.len(), + completion_tokens, + req.max_tokens, + ) { return; } let options = row_options(&req); self.ranks[dp_rank].slots[slot] = Some(RequestState { token_tx: req.token_tx, + stop_policy: req.stop_policy, prompt_len: req.prompt_tokens.len(), completion_tokens, max_tokens: req.max_tokens, @@ -908,36 +883,16 @@ impl DpRankState { return; } - // EOS outranks the length limit; the stop token itself is not emitted - // (same contract as the Qwen schedulers). - if !req.options.sampling.ignore_eos && stop_token_ids.contains(&token_id) { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_len, - completion_tokens: req.completion_tokens, - }); - self.slots[slot_idx] = None; - return; - } - - if req - .token_tx - .send(TokenEvent::Token { - id: token_id, - logprob: report.logprob.clone(), - }) - .is_err() - { - self.slots[slot_idx] = None; - return; - } - - if req.completion_tokens >= req.max_tokens { - let _ = req.token_tx.send(TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_len, - completion_tokens: req.completion_tokens, - }); + if !emit_sampled_token( + &req.token_tx, + &req.stop_policy, + stop_token_ids, + token_id, + report.logprob.clone(), + req.prompt_len, + req.completion_tokens, + req.max_tokens, + ) { self.slots[slot_idx] = None; } else { req.last_token = token_id; @@ -1095,6 +1050,8 @@ fn rank_forward_loop( #[cfg(test)] mod tests { + use pegainfer_frontend::engine::EosPolicy; + use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::sampler::SamplingParams; use super::*; @@ -1108,6 +1065,7 @@ mod tests { data_parallel_rank: None, prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -1155,6 +1113,7 @@ mod tests { let (token_tx, _token_rx) = TokenSink::standalone(); RequestState { token_tx, + stop_policy: StopPolicy::default(), prompt_len, completion_tokens, max_tokens, @@ -1323,6 +1282,7 @@ mod tests { kv.schedule_decode(&pool).expect("decode block"); rank.slots[0] = Some(RequestState { token_tx, + stop_policy: StopPolicy::default(), prompt_len: 4, completion_tokens: 1, max_tokens: 16, @@ -1334,10 +1294,15 @@ mod tests { rank.process_decode_report(0, &dummy_report(163_586), &[163_586], &pool); assert!(rank.slots[0].is_none()); + let Ok((_, TokenEvent::Token { id, .. })) = token_rx.try_recv() else { + panic!("expected trigger Token event"); + }; + assert_eq!(id, 163_586); let Ok(( _, TokenEvent::Finished { finish_reason, + stop_cause, completion_tokens, .. }, @@ -1346,13 +1311,13 @@ mod tests { panic!("expected Finished event"); }; assert_eq!(finish_reason, FinishReason::Stop); + assert_eq!(stop_cause, Some(StopCause::Eos(163_586))); assert_eq!(completion_tokens, 2); - // The stop token itself is not emitted. assert!(token_rx.try_recv().is_err()); } #[test] - fn decode_report_honors_ignore_eos() { + fn decode_report_honors_request_eos_policy() { let pool = test_pool(); let mut rank = DpRankState { slots: (0..MAX_BATCH_PER_DP).map(|_| None).collect(), @@ -1362,6 +1327,10 @@ mod tests { kv.schedule_decode(&pool).expect("decode block"); rank.slots[0] = Some(RequestState { token_tx, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, prompt_len: 4, completion_tokens: 1, max_tokens: 16, @@ -1385,6 +1354,42 @@ mod tests { assert_eq!(id, 163_586); } + #[test] + fn decode_report_stops_on_explicit_token_while_eos_is_ignored() { + let pool = test_pool(); + let mut rank = DpRankState { + slots: (0..MAX_BATCH_PER_DP).map(|_| None).collect(), + }; + let (token_tx, mut token_rx) = TokenSink::standalone(); + let mut kv = dummy_kv(&pool, 4, 1, 16, 7); + kv.schedule_decode(&pool).expect("decode block"); + rank.slots[0] = Some(RequestState { + token_tx, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![42], + }, + prompt_len: 4, + completion_tokens: 1, + max_tokens: 16, + last_token: 7, + options: KimiRowOptions::default(), + kv, + }); + + rank.process_decode_report(0, &dummy_report(42), &[163_586], &pool); + + assert!(rank.slots[0].is_none()); + assert!(matches!( + token_rx.try_recv(), + Ok((_, TokenEvent::Token { id: 42, .. })) + )); + let Ok((_, TokenEvent::Finished { stop_cause, .. })) = token_rx.try_recv() else { + panic!("expected terminal Finished event"); + }; + assert_eq!(stop_cause, Some(StopCause::Token(42))); + } + #[test] fn rank_kv_budget_reserves_remaining_lifetime_of_active_requests() { let mut coordinator = test_coordinator(1); diff --git a/pegainfer-kimi-k2/src/runner/scheduler/lifecycle.rs b/pegainfer-kimi-k2/src/runner/scheduler/lifecycle.rs index dbfaf35ce..e559990e4 100644 --- a/pegainfer-kimi-k2/src/runner/scheduler/lifecycle.rs +++ b/pegainfer-kimi-k2/src/runner/scheduler/lifecycle.rs @@ -75,6 +75,7 @@ pub(in crate::runner) fn preflight_prefill_candidate( UnschedulableVerdict::Finish => { let _ = req.token_tx.send(TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: req.prompt_tokens.len(), completion_tokens: 0, }); diff --git a/pegainfer-kimi-k2/tests/vllm_golden_gate.rs b/pegainfer-kimi-k2/tests/vllm_golden_gate.rs index b33a93d80..347abf24a 100644 --- a/pegainfer-kimi-k2/tests/vllm_golden_gate.rs +++ b/pegainfer-kimi-k2/tests/vllm_golden_gate.rs @@ -57,7 +57,9 @@ use std::time::Instant; use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::EpBackend; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::TokenSink; @@ -345,6 +347,10 @@ fn submit( seed: None, ignore_eos: true, }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: Vec::new(), + }, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 488f261f0..282471d2f 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,29 @@ 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. +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); + result.matched_draft_tokens = result.matched_draft_tokens.min(keep); +} + impl Qwen3Executor { pub(super) fn execute_speculative_verify_impl( &mut self, @@ -25,6 +47,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(), @@ -80,7 +108,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 +136,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 +225,60 @@ 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); + } + + #[test] + fn posterior_stop_does_not_reduce_the_matched_draft_count() { + let policy = StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![99], + }; + let mut result = VerifyRequestResult { + request_id: RequestId::new(3), + matched_draft_tokens: 2, + accepted_tokens: vec![5, 6, 99], + }; + + truncate_after_terminal(&mut result, &policy, &[7]); + + assert_eq!(result.accepted_tokens, vec![5, 6, 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 3954e3644..e24ab7c13 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; @@ -342,7 +343,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) { @@ -366,32 +367,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 @@ -405,14 +385,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, @@ -435,7 +415,7 @@ impl Qwen3Scheduler { req.generated_count = completion_tokens; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id, tokens, completion_tokens, @@ -459,10 +439,11 @@ impl Qwen3Scheduler { } } } - DecodeEffect::EmitManyAndFinish { + DecodeEffect::FinishMany { request_id, tokens, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -476,7 +457,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); @@ -506,32 +487,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); @@ -559,12 +527,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(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(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 faace5079..a4bbda206 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); @@ -173,27 +175,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 { @@ -244,6 +254,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::*; @@ -255,6 +266,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, @@ -274,6 +286,7 @@ mod tests { max_tokens, prompt_len: 10, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: 0, } } @@ -286,12 +299,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..1219fe52e 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -1,4 +1,5 @@ use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; use super::ActiveRequestState; use super::PendingRequest; @@ -13,6 +14,24 @@ use crate::executor::ModelExecutor; use crate::executor::PrefillRequestResult; use crate::speculative::VerifyRequestResult; +fn stop_cause( + executor: &impl ModelExecutor, + req: &ActiveRequestState, + token: u32, +) -> Option { + req.stop_policy + .classify(token, |token_id| executor.is_stop_token(token_id)) +} + +fn pending_stop_cause( + executor: &impl ModelExecutor, + req: &PendingRequest, + token: u32, +) -> Option { + req.stop_policy + .classify(token, |token_id| executor.is_stop_token(token_id)) +} + pub(crate) fn resolve_step( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -60,26 +79,30 @@ pub(crate) fn resolve_speculative_outputs( .expect("speculative request_id must exist in active set"); let mut emitted = Vec::new(); let mut completion_tokens = req.generated_count; + for &token in &result.accepted_tokens { completion_tokens += 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(token); - if is_eos { - return DecodeEffect::EmitManyAndFinish { + emitted.push(token); + + if let Some(stop_cause) = stop_cause(executor, req, token) { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }; } - emitted.push(token); + if completion_tokens >= req.max_tokens { - return DecodeEffect::EmitManyAndFinish { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Length, + stop_cause: None, }; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id: result.request_id, tokens: emitted, completion_tokens, @@ -126,10 +149,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 +166,7 @@ fn resolve_prefill_outputs( token: result.first_token, logprob: result.first_token_logprob, finish_reason: FinishReason::Length, + stop_cause: None, }); continue; } @@ -154,6 +181,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 +205,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..2da3e6d4b 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,7 +563,7 @@ fn speculative_full_span_accept_continues() { "completion = prior generated + span len" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), } } @@ -564,41 +576,41 @@ fn speculative_stop_token_midspan_finishes_and_suppresses_eos() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, + stop_cause, .. }, ] => { - assert_eq!( - tokens, - &vec![10, 11], - "EOS itself is suppressed from emission" - ); + assert_eq!(tokens, &vec![10, 11, SPEC_EOS],); assert!(matches!(finish_reason, FinishReason::Stop)); + assert_eq!(*stop_cause, Some(StopCause::Eos(SPEC_EOS))); } - _ => panic!("expected EmitManyAndFinish(Stop)"), + _ => panic!("expected FinishMany(Stop)"), } } #[test] -fn speculative_stop_token_at_span_start_emits_nothing() { +fn speculative_stop_token_at_span_start_retains_the_token() { let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); let active = [spec_active(1, 5, 100, false)]; let results = [spec_result(1, vec![SPEC_EOS, 11, 12])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, + stop_cause, .. }, ] => { - assert!(tokens.is_empty(), "stop at position 0 emits no tokens"); + assert_eq!(tokens, &vec![SPEC_EOS]); assert!(matches!(finish_reason, FinishReason::Stop)); + assert_eq!(*stop_cause, Some(StopCause::Eos(SPEC_EOS))); } - _ => panic!("expected EmitManyAndFinish(Stop)"), + _ => panic!("expected FinishMany(Stop)"), } } @@ -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 b5ed57378..f18cdde90 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -24,6 +24,7 @@ //! target distribution; acceptance only decides how many ride one step. use anyhow::Result; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use crate::executor::RequestId; @@ -56,6 +57,9 @@ impl VerifyStepItem { #[derive(Clone, Copy)] pub(crate) struct VerifyPlan<'a> { pub requests: &'a [VerifyStepItem], + /// Request-local stop policies in the same order as `requests`. They stay + /// executor-side and are not copied into the worker command or GPU batch. + pub stop_policies: &'a [StopPolicy], /// Engine step seed for the verify rows' sampler pass (same contract as /// decode: fresh per step; seeded rows re-mix their own request seed). pub sample_seed: u64, @@ -69,8 +73,9 @@ pub(crate) struct VerifyRequestResult { /// Tokens to commit: the accepted draft prefix followed by the target's /// posterior token at the first mismatch (or the block-end continuation /// when every draft is accepted). Always `1..=K + 1` tokens, so a verify - /// step always makes at least one token of progress. The scheduler still - /// owns stop-token suppression before client emission. + /// step always makes at least one token of progress. Before KV commit the + /// executor truncates this span after the first request-terminal token; + /// the scheduler retains ownership of typed stop-cause emission. pub accepted_tokens: Vec, } diff --git a/pegainfer-qwen3/tests/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index dcbe38ebe..8c54ed21e 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; @@ -36,9 +37,18 @@ pub(crate) fn request( params: SamplingParams, max_tokens: usize, ) -> Request { + let eos = if params.ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }; Request { prompt_tokens, params, + stop_policy: pegainfer_frontend::engine::StopPolicy { + eos, + token_ids: Vec::new(), + }, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler.rs index c2552c23c..84c8076f0 100644 --- a/pegainfer-qwen35/src/scheduler.rs +++ b/pegainfer-qwen35/src/scheduler.rs @@ -31,6 +31,7 @@ use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest as SchedulerRequest; use pegainfer_frontend::engine::KvCapacity; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::SubmittedRequest; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenLogprob; @@ -85,6 +86,7 @@ struct ActiveRequest35 { max_tokens: usize, prompt_len: usize, params: SamplingParams, + stop_policy: StopPolicy, /// Number of top logprobs to return (0 = disabled). logprobs: usize, } @@ -2073,10 +2075,12 @@ fn dispatch_decode_tokens( let req = &mut active[i]; req.generated_count += 1; - let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); + let stop_cause = req + .stop_policy + .classify(token, |token_id| backend.is_stop_token(token_id)); let at_limit = req.generated_count >= req.max_tokens; - if is_eos { + if let Some(stop_cause) = stop_cause { debug!( "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", req.request_id, @@ -2084,15 +2088,21 @@ fn dispatch_decode_tokens( req.generated_count, FinishReason::Stop ); - let event = TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }; + let events = vec![ + TokenEvent::Token { id: token, logprob }, + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), + prompt_tokens: req.prompt_len, + completion_tokens: req.generated_count, + }, + ]; if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(vec![event]))); + to_retire.push((i, Retirement::Completion(events))); } else { - let _ = req.token_tx.send(event); + for event in events { + let _ = req.token_tx.send(event); + } to_retire.push((i, Retirement::CleanupOnly)); } } else if at_limit { @@ -2107,6 +2117,7 @@ fn dispatch_decode_tokens( TokenEvent::Token { id: token, logprob }, TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: req.prompt_len, completion_tokens: req.generated_count, }, @@ -2392,21 +2403,31 @@ fn promote_or_requeue( let first_token = artifact.token; let logprob = artifact.logprob; - if !req.params.ignore_eos && backend.is_stop_token(first_token) { + let stop_cause = req + .stop_policy + .classify(first_token, |token_id| backend.is_stop_token(token_id)); + if let Some(stop_cause) = stop_cause { debug!( "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", req.request_id, prompt_len, - 0, + 1, FinishReason::Stop ); let candidate = CompletionCandidate { request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: prompt_len, - completion_tokens: 0, - }], + final_events: vec![ + TokenEvent::Token { + id: first_token, + logprob, + }, + TokenEvent::Finished { + finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), + prompt_tokens: prompt_len, + completion_tokens: 1, + }, + ], }; if let Err(err) = backend .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) @@ -2439,6 +2460,7 @@ fn promote_or_requeue( }, TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: prompt_len, completion_tokens: 1, }, @@ -2494,6 +2516,7 @@ fn promote_or_requeue( max_tokens: req.max_tokens, prompt_len, params: req.params, + stop_policy: req.stop_policy, logprobs: req.logprobs, }); } diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 39a521517..6cc621b75 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -5,7 +5,10 @@ use std::time::Duration; use std::time::Instant; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::EpBackend; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use super::*; @@ -29,6 +32,10 @@ fn test_request_with_shape( ignore_eos: true, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![], + }, max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -50,6 +57,7 @@ fn active_request(request_id: u64, label: &str, token_tx: TokenSink) -> ActiveRe max_tokens: 8, prompt_len: 1, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: 0, } } @@ -427,7 +435,7 @@ fn prune_drop_failure_preserves_pending_for_terminal_fanout() { fn decode_eos_waits_for_drop_before_finished() { let (token_tx, token_rx) = TokenSink::standalone(); let mut request = active_request(30, "decode-eos", token_tx); - request.params.ignore_eos = false; + request.stop_policy = StopPolicy::default(); let mut active = vec![request]; let mut backend = LifecycleTestBackend::new(Some(9), token_rx); @@ -436,10 +444,15 @@ fn decode_eos_waits_for_drop_before_finished() { assert!(active.is_empty()); assert_eq!(backend.active_drops, vec![RequestId::new(30)]); let mut token_rx = backend.observer.take().unwrap(); + assert!(matches!( + next_event(&mut token_rx, "decode EOS token"), + TokenEvent::Token { id: 9, .. } + )); assert!(matches!( next_event(&mut token_rx, "decode EOS"), TokenEvent::Finished { finish_reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(9)), .. } )); @@ -684,6 +697,7 @@ fn terminal_shutdown_closes_drains_and_errors_every_owner_once() { request: active_request(42, "candidate", candidate_tx), final_events: vec![TokenEvent::Finished { finish_reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 1, completion_tokens: 2, }], @@ -808,6 +822,7 @@ fn send_rejection_reports_kv_lifetime_request_tokens() { data_parallel_rank: None, prompt_tokens: vec![1; 16], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 65, lora_adapter: None, kv_transfer_params: None, @@ -948,6 +963,7 @@ fn send_rejection_reports_context_window_limit() { data_parallel_rank: None, prompt_tokens: vec![1; 16], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 17, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen35/tests/chunked_prefill.rs b/pegainfer-qwen35/tests/chunked_prefill.rs index dee13c8b2..bd9a846c7 100644 --- a/pegainfer-qwen35/tests/chunked_prefill.rs +++ b/pegainfer-qwen35/tests/chunked_prefill.rs @@ -10,8 +10,10 @@ use std::path::Path; use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::sampler::SamplingParams; @@ -51,6 +53,10 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec) -> (Vec, Finish ignore_eos: true, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![], + }, max_tokens: GENERATED_TOKENS, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen35/tests/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index 9886399bb..dc01cfeff 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -9,9 +9,11 @@ use std::time::Instant; use log::info; use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::FinishReason; use pegainfer_frontend::engine::GenerateRequest; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::TokenSink; @@ -128,6 +130,7 @@ fn generate_tokens_with_logprobs( data_parallel_rank: None, prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -159,6 +162,10 @@ fn submit_repeated_token_request( ignore_eos: true, ..SamplingParams::default() }, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![], + }, max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -345,6 +352,7 @@ fn expect_context_window_rejection(handle: &EngineHandle, max_context_tokens: us data_parallel_rank: None, prompt_tokens: vec![1; max_context_tokens], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 1, lora_adapter: None, kv_transfer_params: None, @@ -558,6 +566,7 @@ fn run_full_scheduler_e2e( data_parallel_rank: None, prompt_tokens, params: concurrent_params(case_idx), + stop_policy: StopPolicy::default(), max_tokens: case.max_new_tokens, lora_adapter: None, kv_transfer_params: None, @@ -600,6 +609,7 @@ fn run_full_scheduler_e2e( data_parallel_rank: None, prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 8, lora_adapter: None, kv_transfer_params: None, @@ -643,6 +653,7 @@ fn run_full_scheduler_e2e( data_parallel_rank: None, prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 10, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen35/tests/sampling_behavior.rs b/pegainfer-qwen35/tests/sampling_behavior.rs index 4649ec186..e1b561ff4 100644 --- a/pegainfer-qwen35/tests/sampling_behavior.rs +++ b/pegainfer-qwen35/tests/sampling_behavior.rs @@ -12,7 +12,9 @@ use std::path::Path; use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::GenerateRequest; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenEvent; use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::sampler::SamplingParams; @@ -37,6 +39,10 @@ fn generate(handle: &EngineHandle, prompt_tokens: Vec, params: SamplingPara data_parallel_rank: None, prompt_tokens, params, + stop_policy: StopPolicy { + eos: EosPolicy::Ignore, + token_ids: vec![], + }, max_tokens: GENERATED_TOKENS, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-sim/src/lib.rs b/pegainfer-sim/src/lib.rs index bba3357c2..86c765707 100644 --- a/pegainfer-sim/src/lib.rs +++ b/pegainfer-sim/src/lib.rs @@ -12,6 +12,7 @@ use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestLedger; use pegainfer_frontend::engine::Scheduler; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::spawn_scheduler; @@ -129,6 +130,7 @@ struct RunningRequest { next_token_at: Instant, finish_reason: FinishReason, logprobs: usize, + stop_policy: StopPolicy, } impl SimScheduler { @@ -177,7 +179,7 @@ impl Scheduler for SimScheduler { planned_completion(&self.config, &request.prompt_tokens, request.max_tokens); ledger.admit(id); if pending.is_empty() { - ledger.finish(id, finish_reason); + ledger.finish(id, finish_reason, None); continue; } self.running.push(RunningRequest { @@ -186,6 +188,7 @@ impl Scheduler for SimScheduler { next_token_at: Instant::now() + self.config.ttft(prompt_len), finish_reason, logprobs: request.logprobs, + stop_policy: request.stop_policy, }); } @@ -201,7 +204,7 @@ impl Scheduler for SimScheduler { continue; } let Some(token) = running.pending.pop() else { - ledger.finish(running.id, running.finish_reason); + ledger.finish(running.id, running.finish_reason, None); continue; }; let logprob = (running.logprobs > 0).then_some(TokenLogprob { @@ -213,8 +216,13 @@ impl Scheduler for SimScheduler { None => Vec::new(), }; ledger.push_tokens(running.id, &[token], &logprobs); - if running.pending.is_empty() { - ledger.finish(running.id, running.finish_reason); + // The simulator has no model EOS set; explicit request stop IDs + // still exercise the same token-before-terminal contract as a + // real scheduler. + if let Some(stop_cause) = running.stop_policy.classify(token, |_| false) { + ledger.finish(running.id, FinishReason::Stop, Some(stop_cause)); + } else if running.pending.is_empty() { + ledger.finish(running.id, running.finish_reason, None); } else { running.next_token_at = Instant::now() + self.config.tpot(); still_running.push(running); @@ -277,15 +285,27 @@ fn duration_from_ms(ms: f64) -> Duration { #[cfg(test)] mod tests { use pegainfer_frontend::engine::Request; + use pegainfer_frontend::engine::StopCause; + use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; use super::*; fn request(prompt_tokens: Vec, max_tokens: usize, logprobs: usize) -> Request { + request_with_policy(prompt_tokens, max_tokens, logprobs, StopPolicy::default()) + } + + fn request_with_policy( + prompt_tokens: Vec, + max_tokens: usize, + logprobs: usize, + stop_policy: StopPolicy, + ) -> Request { Request { prompt_tokens, params: SamplingParams::default(), + stop_policy, max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -368,6 +388,34 @@ mod tests { )); } + #[test] + fn explicit_stop_keeps_trigger_token_and_reports_cause() { + let config = SimulatedEngineConfig::new(0.0, 100.0, 0.0, 0) + .unwrap() + .with_scripted_completion(vec![11, 42, 99]); + let request = request_with_policy( + vec![7], + 8, + 0, + StopPolicy { + eos: pegainfer_frontend::engine::EosPolicy::Ignore, + token_ids: vec![42], + }, + ); + let (tokens, _, terminal) = collect_completion(&config, request); + + assert_eq!(tokens, [11, 42]); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(42)), + completion_tokens: 2, + .. + } + )); + } + #[test] fn config_rejects_invalid_timing_values() { assert!(SimulatedEngineConfig::new(-1.0, 100.0, 12.0, 0).is_err()); @@ -391,6 +439,7 @@ mod tests { reason: FinishReason::Length, prompt_tokens: 2, completion_tokens: 3, + .. } )); }