diff --git a/docs/models/qwen3/model-crate.md b/docs/models/qwen3/model-crate.md index 20a8b3024..187a3339a 100644 --- a/docs/models/qwen3/model-crate.md +++ b/docs/models/qwen3/model-crate.md @@ -1,8 +1,8 @@ # Qwen3-4B Model Crate **Created**: 2026-05-03 -**Last touched**: 2026-07 -**TL;DR**: `crates/pegainfer-qwen3` now owns Qwen3 config, weights, execution, scheduler, tests, benches, and kernel plan. Root `pegainfer` loads Qwen3 through a generic `EngineHandle` and no longer contains `Qwen3Model`, `Qwen3Executor`, `ModelRuntimeConfig`, root Qwen3 tests, or `src/model/qwen3/*`. The old `ModelForward` path has been removed; decode length-limit now emits the final token before `Finished`. Long-context `bs=1` TPOT was traced to non-partition FlashInfer paged decode under-filling the GPU; Qwen3 runtime gates FlashInfer split-K decode on `padded_bs<=32` (the `seq_len>=1024` gate was dropped in #437) with a 64-token chunk floor (`Tuned` capped at 64 chunks; opt-in `--batch-invariant` pins a fixed 160-token split), cutting 4k/64 serving steady TPOT from about `11.7ms` to `6.46ms` on RTX 5090. Qwen3 now keeps a single model-crate bench entry: `qwen3_kernel_snapshot`, a JSON snapshot runner with warm/cold-L2 latency, default-on CUPTI counters, and compare. Correctness/truth is intentionally out of this snapshot for now. +**Last touched**: 2026-09 +**TL;DR**: `pegainfer-qwen3` owns Qwen3 config, weights, execution, scheduling, tests, benches, and kernel metadata. The model line exposes Qwen3 through the stepped `LaunchedEngine` contract; frontend code does not depend on Qwen3 executor internals. The old `ModelForward` path has been removed; decode length-limit now emits the final token before `Finished`. Long-context `bs=1` TPOT was traced to non-partition FlashInfer paged decode under-filling the GPU; Qwen3 runtime gates FlashInfer split-K decode on `padded_bs<=32` (the `seq_len>=1024` gate was dropped in #437) with a 64-token chunk floor (`Tuned` capped at 64 chunks; opt-in `--batch-invariant` pins a fixed 160-token split), cutting 4k/64 serving steady TPOT from about `11.7ms` to `6.46ms` on RTX 5090. Qwen3 now keeps a single model-crate bench entry: `qwen3_kernel_snapshot`, a JSON snapshot runner with warm/cold-L2 latency, default-on CUPTI counters, and compare. Correctness/truth is intentionally out of this snapshot for now. ## Determinism scope @@ -10,112 +10,35 @@ It also **requires `--no-prefix-cache` and rejects `--kv-offload`**. A prefix hit advances a prompt's `chunk_start` off the request-local chunk grid, re-cutting its prefill chunks according to prior traffic. `--no-prefix-cache` closes that — except under offload, where it only disables HBM retention and deliberately leaves prefix *matching* on (host-tier reuse is the point of that mode), so that pair is rejected outright. -## Preparation - -- **Read**: - - `docs/index.md` - identified the kernels/core crate split and per-model boundary docs. - - `docs/models/qwen3/kernels-crate.md` - Qwen3 kernel source/build ownership and human kernel index already live in `pegainfer-kernels`; model-owned DAG metadata should live with the model crate. - - `docs/subsystems/kernels/pegainfer-kernels-boundary.md` - records the per-model engine direction and says root should be reusable frontend/control-plane infrastructure, not a universal model abstraction. - - `src/main.rs`, `src/lib.rs`, `src/server_engine.rs`, `src/scheduler.rs`, `src/model_executor.rs`, `src/model/qwen3/*`, `src/bin/bench_serving.rs`, and Qwen3 tests - mapped what root currently knows about Qwen3. -- **Relevant history**: - - The earlier shared-runtime work (now consolidated into `docs/subsystems/runtime/runtime.md`) was a useful simplification, but the next boundary should not make `ModelForward` the long-term universal engine API. -- **Plan**: - 1. Define the model crate/root interface before moving code. - 2. Move the generic text-generation handle/request/event types into `pegainfer-core` so root and model crates can communicate without model crates depending on root. - 3. Create `crates/pegainfer-qwen3` and move Qwen3 config, weights, forward paths, decode buffers, `Qwen3Executor`, Qwen3 scheduler internals, Qwen3 correctness tests, and Qwen3-specific benches into it. - 4. Keep root `pegainfer` as frontend plus model registry. The registry can know crate names, but `main`, `vllm_frontend`, and generic benchmark code should only see `EngineHandle`, `ModelInfo`, and tokenizer path. - 5. Add a model-owned `kernel_plan.rs` in the Qwen3 crate as the LLM/human index from model DAG phases to reusable kernels. Do not add a hand-maintained public TOML in `pegainfer-kernels`. - 6. Verify locally with format/metadata, then on the CUDA validation host with release build, clippy, Qwen3 crate e2e, and root `bench_serving snapshot`. Keep microbench timing in Criterion benches instead of duplicating it as a test. -- **Risks / open questions**: - - If the scheduler stays in root, root still knows Qwen3's execution shape. To meet the stated goal, the Qwen3 scheduler should move into the Qwen3 crate and expose only a generic handle. - - `bench_serving` previously had a direct `ModelForward` path for Qwen3 and a scheduler path for Qwen3.5. It needed to become generic over `EngineHandle`, while Qwen3 crate-local benches should use the model executor phase API. - - Qwen3.5 remains in root for this phase. The registry may temporarily wrap root-local Qwen3.5, but new Qwen3 code should not depend on that temporary shape. - -## Interface Proposal - -The root-visible interface should be request/response oriented, not prefill/decode oriented. - -```rust -// pegainfer-core -pub struct EngineLoadOptions { - pub enable_cuda_graph: bool, - pub device_ordinals: Vec, - pub seed: u64, -} - -pub struct ModelInfo { - pub id: &'static str, - pub display_name: String, - pub max_model_len: Option, -} - -pub struct GenerateRequest { - pub prompt_tokens: Vec, - pub params: SamplingParams, - pub max_tokens: usize, - pub token_tx: tokio::sync::mpsc::UnboundedSender, - pub logprobs: usize, - pub echo: bool, -} - -pub enum TokenEvent { - Token { id: u32, logprob: Option }, - PromptTokens { ids: Vec, logprobs: Vec> }, - Finished { finish_reason: FinishReason, prompt_tokens: usize, completion_tokens: usize }, -} - -#[derive(Clone)] -pub struct EngineHandle { - submit_tx: tokio::sync::mpsc::UnboundedSender, -} -``` +## Stop contract -```rust -// pegainfer-qwen3 -pub fn start_engine( - model_path: &std::path::Path, - options: EngineLoadOptions, -) -> anyhow::Result; -``` +The stepped Qwen3 scheduler keeps the request's EOS policy and explicit +`stop_token_ids` independent in `StopPolicy`. When a generated token triggers +either policy, that token is retained exactly once in the internal +`RequestUpdate`; decode results carry its real logprob when one is available, +and the terminal carries a typed `StopCause`. `StopCause::Eos` maps to +`finish_reason=stop` without a wire +`stop_reason`; an explicit request stop maps to the actual token ID. +`ignore_eos=true` disables only model EOS and does not disable explicit request +stop IDs. A length finish has no token-level stop cause. + +## Runtime boundary + +- `pegainfer-qwen3/src/model_line.rs` is the model-line dispatch seam. Its + `launch` implementation returns `LaunchedEngine::Stepped` and keeps model + loading, scheduler state, and CUDA execution inside this crate. +- `pegainfer-qwen3::runtime` is the intentional low-level surface for model- + local tools and benches. The frontend only consumes the generic step + contract (`Request`, `StepOutputs`, `RequestLedger`, and `StopPolicy`). +- The stop contract above is part of the model-line boundary: Qwen3 retains + the trigger token and reports the typed cause; it is not a legacy + per-request event protocol. + +## Execution log -`Qwen3Model`, `BatchDecodeBuffers`, and `KvState` should not be root-facing APIs. The deliberate low-level escape hatch is `pegainfer_qwen3::runtime`, which exposes `Qwen3Executor` plus prefill/decode/unified plan types. That is the production phase boundary used by the scheduler and by model-local benches; root should still use `start_engine`. - -## Execution Log - -### Step 1: Add generic engine API to core -- Added `pegainfer_core::engine` with: - - `EngineLoadOptions` - - `ModelInfo` - - `TokenLogprob` - - `FinishReason` - - `GenerateRequest` - - `TokenEvent` - - `EngineHandle` -- Root `server_engine` now re-exports `FinishReason` and `TokenLogprob` for compatibility. -- Root `scheduler.rs` is reduced to compatibility re-exports for `SchedulerHandle`, `SchedulerRequest`, and `TokenEvent`. - -### Step 2: Extract Qwen3 crate -- Added `crates/pegainfer-qwen3`. -- Moved Qwen3-owned code into the crate: - - config/weights/forward/prefill/decode/unified forward - - batch decode buffers - - `Qwen3Executor` - - Qwen3 scheduler internals - - Qwen3 e2e and paged-attention correctness tests - - Qwen3 regression data generator - - Qwen3 prefill Criterion bench -- Added `kernel_plan.rs` as the model-owned kernel routing index. It is typed Rust metadata, not a hand-maintained public TOML. - -### Step 3: Remove root Qwen3 execution knowledge -- Root no longer has: - - `src/model/qwen3.rs` - - `src/model/qwen3/*` - - `src/model_executor.rs` - - Qwen3 root tests: `tests/e2e.rs`, `tests/paged_attention.rs`, `tests/bench_prefill.rs` -- Root `main.rs` starts Qwen3 through `pegainfer_qwen3::start_engine(...)`. -- Root `vllm_frontend.rs` accepts a generic `EngineHandle`. -- Root `bench_serving` uses the same generic scheduler bench path for Qwen3 instead of constructing `Qwen3Model` directly. -- Checked root with `rg` and confirmed no hits for `Qwen3Model`, `Qwen3Executor`, `ModelRuntimeConfig`, `model_executor`, `src/model/qwen3`, or stale "Qwen3 continuous" comments under root source/tests/benches/README. +The entries below are historical validation notes. They describe files and +commands that existed at the time of each measurement; current ownership and +interfaces are defined by the runtime boundary above. ### Step 4: Link and validation fixes - Added explicit `stdc++` link output in `pegainfer-kernels` build script. Once Qwen3 became an independent crate with its own tests, the FlashInfer C++ CUDA objects needed the C++ runtime linked for test binaries as well as root binaries. @@ -142,7 +65,7 @@ pub fn start_engine( - Rejected a bench-only support API and also rejected using `ModelForward` as the benchmark entry. - Added an explicit `runtime` module that re-exports the scheduler's real `Qwen3Executor` phase API: `PrefillPlan`, `DecodePlan`, `UnifiedPlan`, request items, and result types. - Removed top-level public `Qwen3Model`, `ModelRuntimeConfig`, and `Qwen3State` re-exports. External low-level tools must opt into `runtime`; root continues to use `start_engine`. -- Replaced `crates/pegainfer-qwen3/benches/qwen3_prefill.rs` with `benches/qwen3_runtime.rs`. It measures executor prefill TTFT over `128`, `512`, `1024`, `2048`, `4096`, and `10000` token prompts, plus executor decode TPOT for batch sizes `1`, `2`, `4`, `8`, `16`, and `32` at a `1024` token context. +- Replaced `pegainfer-qwen3/benches/qwen3_prefill.rs` with `benches/qwen3_runtime.rs`. It measures executor prefill TTFT over `128`, `512`, `1024`, `2048`, `4096`, and `10000` token prompts, plus executor decode TPOT for batch sizes `1`, `2`, `4`, `8`, `16`, and `32` at a `1024` token context. - Updated `tests/paged_attention.rs` to use the same executor phase API: prefill once to create KV state, then decode through `execute_decode`. - Verification after the cleanup: - Local `cargo fmt --all --check` and `cargo metadata --no-deps --format-version 1` pass. @@ -157,7 +80,7 @@ pub fn start_engine( ### Step 7: Retire ModelForward and Fix Length Limit - Deleted `pegainfer_core::model::{ModelForward, GenerationState}` and removed the root `src/model.rs` re-export. - Deleted the Qwen3 `forward.rs` compatibility path. Qwen3 tests that used it now build their baselines from `batch_prefill(bs=1)` plus `batch_decode(bs=1)`, so they exercise the same phase APIs as production. -- Fixed Qwen3 decode length-limit handling by adding `DecodeEffect::EmitAndFinish`. EOS behavior is unchanged: EOS finishes without emitting the stop token. Length limit now emits the sampled final token, then sends `Finished { finish_reason: Length }`. +- Fixed Qwen3 decode length-limit handling by adding `DecodeEffect::EmitAndFinish`. The stepped stop contract now retains an EOS trigger in the internal token update and records `StopCause::Eos`; the bridge omits EOS from wire `stop_reason`. Length limit emits the sampled final token, then sends `Finished { finish_reason: Length }`. - Regenerated `test_data/Qwen3-4B.json` because every length-limited golden output now includes the final requested token. - Re-ran `bench_serving snapshot` on the CUDA validation host and pulled back `bench_snapshots/rtx-5090/qwen3-4b.json`; `decode_heavy (1024,256)` now records `generated_tokens min=max=avg=256`. - Performance stayed within noise on RTX 5090: @@ -169,7 +92,7 @@ pub fn start_engine( - CUDA host `PEGAINFER_CUDA_SM=120 PEGAINFER_TEST_MODEL_PATH= cargo test --release -p pegainfer-qwen3 --test e2e -- --nocapture` passes. ### Step 8: Decode Context-Length Sweep and Compile Audit -- Added `crates/pegainfer-qwen3/src/bin/qwen3_decode_context.rs` as a production-path fixed-context decode probe. It prefills a fresh request to a selected context length, then measures or profiles real `Qwen3Executor::execute_decode`; the optional `cudaProfilerStart/Stop` range only exists for profiler capture and does not run in normal serving. +- Added `pegainfer-qwen3/src/bin/qwen3_decode_context.rs` as a production-path fixed-context decode probe. It prefills a fresh request to a selected context length, then measures or profiles real `Qwen3Executor::execute_decode`; the optional `cudaProfilerStart/Stop` range only exists for profiler capture and does not run in normal serving. - GPU fixed-context command: - `PEGAINFER_CUDA_SM=120 target/release/qwen3_decode_context --model-path --iters 10 --contexts 128,512,1024,2048,4096,8192,10000` - Result on RTX 5090: @@ -195,13 +118,13 @@ pub fn start_engine( - H2D traffic in the profiled decode range was only about `20-23us/step`, so metadata dirty caching is good runtime hygiene but cannot explain a multi-ms TPOT gap. - Compile audit on the same validation worktree: - GPU reports compute capability `12.0`; default toolkit is CUDA `12.9` (`nvcc V12.9.86`), driver `575.57.08`. - - `crates/pegainfer-kernels/build.rs` emits `-O3 -gencode arch=compute_120,code=sm_120 -gencode arch=compute_120,code=compute_120 --compiler-options -fPIC`; FlashInfer translation units add `--std=c++17` and the FlashInfer include path. + - `pegainfer-kernels/build.rs` emits `-O3 -gencode arch=compute_120,code=sm_120 -gencode arch=compute_120,code=compute_120 --compiler-options -fPIC`; FlashInfer translation units add `--std=c++17` and the FlashInfer include path. - `cuobjdump -lelf` confirms both `libkernels_cuda.a` and `target/release/pegainfer` contain `sm_120.cubin`. `compute_120` PTX fallback is also embedded, but the matching SASS is present, so this is not PTX-JIT-only execution. - CUDA `13.1` is installed and can build the same code into `sm_120` cubins, but the current driver/runtime combination cannot run it (`cudaError=35` after linking `libcudart.so.13`). Until the driver is upgraded, CUDA `12.9` is the latest runnable toolkit on this box. - Interpretation: the compile target is correct. The `bs=1` long-context slope is the known non-partition FlashInfer paged decode issue: grid shape is effectively `(batch_size, num_kv_heads) = (1, 8)`, so only 8 CTAs scan the whole KV context. At `ctx=4096`, Qwen3-4B attention reads about `604MB` (`576MiB`) of K/V per token; the measured attention time is about `5.7ms`, or roughly `105GB/s` effective aggregate bandwidth, far below the RTX 5090 memory system because the kernel under-fills the GPU. The next real fix is partition-KV/split-K decode for `bs=1` or low-batch, not build-flag tuning. ### Step 9: Pure Paged Decode Attention Bench -- Added `crates/pegainfer-qwen3/benches/qwen3_attention.rs`. +- Added `pegainfer-qwen3/benches/qwen3_attention.rs`. - The bench does not load Qwen3 weights. It constructs synthetic non-zero Q and paged KV buffers using Qwen3-4B attention shape: `num_qo_heads=32`, `num_kv_heads=8`, `head_dim=128`, `page_size=16`, one layer. - The bench calls the FlashInfer paged decode FFI directly and uses CUDA events around the kernel launches. It measures decode attention only; it excludes QKV projection, KV append, O projection, MLP, scheduler, tokenizer, and host-side serving overhead. - Added `paged_attention_decode_split_kv_cuda` as a reusable kernel entry for FlashInfer partition-KV/split-K decode. Runtime dispatch still uses the existing non-partition path; this step only exposes and benchmarks the candidate operator. @@ -315,7 +238,7 @@ Result: Interpretation: split-K removes the long-context attention slope for the low-batch case. The remaining `~6.8-7.1ms` TPOT is now dominated by the non-attention decode body: GEMMs/GEMVs, MLP, norms, logits, sampling, and graph replay overhead. Next optimization work should not keep pushing paged attention first; it should re-profile the post-split decode step and pick the new largest kernel family. ### Step 11: Attention Theoretical Bandwidth Estimate -- Updated `crates/pegainfer-qwen3/benches/qwen3_attention.rs` to print a one-time theoretical bandwidth report before Criterion runs. +- Updated `pegainfer-qwen3/benches/qwen3_attention.rs` to print a one-time theoretical bandwidth report before Criterion runs. - The report queries CUDA Driver attributes: - `CU_DEVICE_ATTRIBUTE_MEMORY_CLOCK_RATE` - `CU_DEVICE_ATTRIBUTE_GLOBAL_MEMORY_BUS_WIDTH` @@ -350,8 +273,8 @@ Batch sweep sanity rows at `kv_len=1024`: Interpretation: the estimate is good enough to prove the original `bs=1` non-partition path was badly under-filling memory bandwidth. It is not good enough to make final hardware-utilization claims because single-layer KV working sets fit in the RTX 5090's `96MiB` L2; the `bs16` non-partition row exceeding `100%` of DRAM peak is the warning sign. The next measurement step should use CUPTI Profiler or NCU counters for `dram__bytes_*`, `lts__t_bytes.*`, and `*_pct_of_peak_sustained_elapsed`. ### Step 12: CUPTI Counters and Split-K Retune -- Added `crates/pegainfer-cupti`, a small CUPTI Range Profiler wrapper used by the attention bench. It profiles only the attention launch range and lets the bench clear L2 before `cuptiRangeProfilerStart`, so cache-clear traffic is excluded from the measured range. -- Extended `crates/pegainfer-qwen3/benches/qwen3_attention.rs`: +- Added `pegainfer-cupti`, a small CUPTI Range Profiler wrapper used by the attention bench. It profiles only the attention launch range and lets the bench clear L2 before `cuptiRangeProfilerStart`, so cache-clear traffic is excluded from the measured range. +- Extended `pegainfer-qwen3/benches/qwen3_attention.rs`: - `PEGAINFER_QWEN3_ATTENTION_CUPTI=1` prints cold-L2 CUPTI rows for `gpu__time_duration.sum`, `dram__bytes.sum`, `dram__bytes_op_read.sum`, `dram__bytes_op_write.sum`, and `lts__t_bytes.sum`. - `PEGAINFER_QWEN3_ATTENTION_SPLITK_SWEEP=1` sweeps split-K chunk sizes and max chunk slots. - `PEGAINFER_QWEN3_ATTENTION_REPORT_ONLY=1` prints reports without running Criterion samples. @@ -445,8 +368,8 @@ Verification: Note: an initial remote e2e run failed because the remote `test_data/Qwen3-4B.json` was stale and expected the pre length-limit baseline. Syncing the tracked baseline fixed it; this was not a split-K numerical drift. ### Step 13: Kernel Snapshot MVP -- Extracted the Qwen3 paged decode attention case construction into `crates/pegainfer-qwen3/src/kernel_bench.rs`. -- Added `crates/pegainfer-qwen3/benches/qwen3_kernel_snapshot.rs` as a deterministic `harness=false` runner. +- Extracted the Qwen3 paged decode attention case construction into `pegainfer-qwen3/src/kernel_bench.rs`. +- Added `pegainfer-qwen3/benches/qwen3_kernel_snapshot.rs` as a deterministic `harness=false` runner. - Removed the temporary correctness envelope from the snapshot runner. We do not have a settled truth source for this layer yet, so correctness belongs in a separate design rather than a misleading "non-partition equals truth" field. - CUPTI is default-on in the snapshot runner. `--no-cupti` is available only for latency-only smoke runs. @@ -522,8 +445,8 @@ The SM counters are intentionally minimal. `sm__throughput.avg.pct_of_peak_susta ### Step 14: Consolidate Bench Entry Points - Deleted the retired Criterion benches: - - `crates/pegainfer-qwen3/benches/qwen3_runtime.rs` - - `crates/pegainfer-qwen3/benches/qwen3_attention.rs` + - `pegainfer-qwen3/benches/qwen3_runtime.rs` + - `pegainfer-qwen3/benches/qwen3_attention.rs` - Removed their `[[bench]]` entries and the Qwen3 crate-local `criterion` dev dependency. - Qwen3 now has exactly one model-crate bench entry: `qwen3_kernel_snapshot`. - Rationale: the human CSV report, split-K tuning sweep, and machine-readable JSON runner were duplicating case construction, metric selection, and interpretation. Kernel maintenance should have one durable artifact first; optional human views should be generated from snapshot data rather than maintained as separate benches. diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index 26008934a..38b482ad2 100644 --- a/docs/subsystems/frontend/frontend-architecture.md +++ b/docs/subsystems/frontend/frontend-architecture.md @@ -1,6 +1,6 @@ # Frontend architecture: pegainfer-frontend and the engine boundary -**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). **Next step: migrate glm52, then delete the legacy contract.** +**TL;DR:** `pegainfer-frontend` owns everything north of the model schedulers: the engine contract, the vLLM protocol stack, and the `ModelLine` dispatch trait. The contract now has two generations living side by side: the **step contract** (`StepOutputs` wire + `RequestLedger` lifecycle + a contract-owned polling driver — Qwen3, Gemma 4, K3 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite still launch through it). **Next step: migrate glm52, then delete the legacy contract.** Last touched: 2026-09 @@ -12,14 +12,14 @@ An engine is a set of schedulers, each a `Scheduler` implementation driven by th ``` pegainfer-frontend/src/engine/ -├── step.rs # the wire: RequestId, Request, QueuedRequest, -│ # StepOutputs { Vec }, -│ # RequestUpdate { scheduled, tokens, logprobs, cached_tokens, -│ # prompt_echo, kv_transfer, terminal }, Terminal +├── step.rs # the wire: RequestId, Request, StepOutputs { Vec }, +│ # Request { ..., stop_policy }, RequestUpdate { scheduled, tokens, +│ # logprobs, cached_tokens, prompt_echo, kv_transfer, terminal }, +│ # Terminal { ..., stop_cause } ├── request_lifecycle.rs # submission envelope, abort control and step sender plumbing; │ # DeferredFinish remains available for P/D handoff -├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, -│ # prompt/completion tallies, one merged update per touched id +├── ledger.rs # RequestLedger: admit/reject/push/finish/fail/retire, +│ # prompt/completion tallies, one merged update per touched id ├── wiring.rs # scheduler_pair, SchedulerHandle (submit/take_steps/load), │ # Engine { schedulers, info, lora }, LiveScheduler, │ # EngineInfo, LaunchedEngine { Handle | Stepped } @@ -32,10 +32,11 @@ pegainfer-frontend/src/engine/ Design decisions worth knowing before touching it: - **Step-batched wire.** One message per scheduler step, not one channel per request: the scheduler's natural output unit is the step batch, and per-request channels were tried and rejected (the scheduler for-loop over N channels was the bottleneck). The protocol stack demuxes. -- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. The ledger merges every write for an id into that one record before the driver commits the step. -- **Ledger lifecycle.** `RequestLedger` owns one account for every unanswered submission. A scheduler receives only `QueuedRequest { id, request }` and writes every lifecycle transition by id. The ledger rejects touches after closure, derives completion counts from `push_tokens`, and writes off every open account when an engine-fatal `step` error ends the driver. -- **Ledger as single writer.** Schedulers never touch the step channel; they call ledger methods. `admit` stamps `ScheduledInfo` from the registered prompt length, `push_tokens` tallies completions, and `commit_step` publishes the merged statement once per driver iteration. -- **Polling driver, scheduler-owned park.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish metrics, commit. An idle iteration ends in a `spin_loop` hint. Gemma 4 has one deliberate park inside this policy: while async prefill is the only remaining work, its scheduler drains the lane and joins it rather than hot-polling the completion (a drain failure is engine-fatal); when decode or queued work exists it keeps polling the lane without blocking. +- **Flat `RequestUpdate`.** All facts a step produced for one request travel in one struct, so intra-request ordering is structure, not convention. This is what makes `defer_finish` safe: a P/D prefill executor can withhold a request's `Finished` until its KV saves are peer-visible and send it later from any thread — the deferred message carries the request's entire buffered update, so late delivery cannot reorder. +- **Independent stop policy and cause.** `Request.stop_policy` keeps model EOS handling separate from explicit request stop IDs. A token-driven finish retains the triggering token in `RequestUpdate` and carries `Terminal::Finished.stop_cause`; decode paths include its real logprob when available, so the stepped bridge can report the actual explicit stop ID without reconstructing it. `ignore_eos` affects only model EOS. Legacy producers may leave the cause empty while they are migrated individually. +- **Ledger lifecycle.** Schedulers carry plain `RequestId`s and mutate `RequestLedger`; `RequestEnvelope` answers submissions dropped before registration, while `DeferredFinish` preserves a complete buffered update when a terminal is delivered after a P/D handoff. `RequestControl` is the frontend abort flag. The ledger remains the single writer for admission, token accounting, terminal transitions, and step publication. +- **Pure polling driver.** `spawn_scheduler` owns the serve loop: drain submissions, `Scheduler::step`, publish load, commit. No idle/park distinction — the scheduler owns the GPU and spinning on it costs nothing anyone else could use; async KV I/O (prefetch, decode-overlap prefill) is naturally absorbed by polling. An idle iteration ends in a `spin_loop` hint (relaxes the core's issue slots, no latency cost — busy iterations never pause). The loop exits when the frontend drops the handle and the queue drains. +- **Gemma 4 async prefill exception.** While asynchronous prefill is the only remaining work, Gemma 4 drains and joins that lane rather than hot-polling its completion; decode or queued work keeps the normal polling path. - **Abort is a flag, not channel teardown.** `SchedulerHandle::submit` returns a `RequestControl`; the frontend flips its boolean abort flag and the scheduler retires the request silently on its next touch (no terminal — the frontend already dropped its state for that id). - **Channels:** the submit channel is crossbeam (sync consumer on the scheduler thread), steps are tokio mpsc (async consumer in the bridge); load is a shared cell read via `SchedulerHandle::load()` — pull-only by design, "notify me on load change" is deliberately unrepresentable (the driver busy-polls, so a subscription edge would fire per spin). All channels unbounded on purpose — admission control is the scheduler's job, expressed as `Rejected`, never as backpressure on submit. - **Control plane lives outside the contract.** `Scheduler` has no control method and the contract carries no control channel. A capability like LoRA is a private channel the model crate mints *before* `spawn_scheduler` — the scheduler closes over the receiver, the `LoraClient` sender surfaces as `Engine.lora: Option`, and the `Option` *is* the capability (no `bool` flag, no registry until a second capability exists). The vocabulary (`LoraControl`, `LoraClient`) is still defined in the frontend crate because the frontend must speak it without holding model structs; only the wiring is the model's business. @@ -91,7 +92,7 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Protocol stacks -**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a `Finished{Stop}` appends the stop sentinel token, which is how usage keeps counting the suppressed EOS). +**`vllm` (current default, fleet-proven).** Impersonates a vLLM EngineCore process over in-process ZMQ/msgpack because upstream `vllm-server` assumes the engine is a separate process. HTTP routes, OpenAI types, tokenizer, chat templates, Prometheus live in the external `vllm-server`/`vllm-metrics`/`vllm-text` crates. `SteppedEngineBridge` translates each `RequestUpdate` 1:1 into an EngineCore output (wall-clock timestamps are reconstructed from the contract's `Instant`s via a per-bridge unix anchor; a typed `StopCause` reports the actual request stop token, while the synthetic sentinel remains only as a compatibility fallback for producers that provide no cause). **`dynamo` (planned second stack).** dynamo's `lib/llm` in-process path removes the wire protocol entirely (`EngineConfig::InProcessTokens` + `run_input`). The step contract was shaped so this stack can consume `StepOutputs` directly without impersonation overhead. Decision gate: prototype, A/B against the vllm stack, let TTFT/step-overhead numbers pick the default. diff --git a/pegainfer-frontend/src/engine/driver.rs b/pegainfer-frontend/src/engine/driver.rs index 827adeb07..c8c4e97bf 100644 --- a/pegainfer-frontend/src/engine/driver.rs +++ b/pegainfer-frontend/src/engine/driver.rs @@ -108,6 +108,7 @@ mod tests { use super::super::step::Request; use super::super::step::RequestId; use super::super::step::Terminal; + use super::super::stop::StopPolicy; use super::*; use crate::engine::FinishReason; @@ -163,6 +164,7 @@ mod tests { Request { prompt_tokens: vec![1, 2], params: crate::sampler::SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -197,6 +199,7 @@ mod tests { terminal, Some(Terminal::Finished { reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 2, completion_tokens: 3, }) diff --git a/pegainfer-frontend/src/engine/ledger.rs b/pegainfer-frontend/src/engine/ledger.rs index 34b4f16d6..08806c001 100644 --- a/pegainfer-frontend/src/engine/ledger.rs +++ b/pegainfer-frontend/src/engine/ledger.rs @@ -38,6 +38,7 @@ use super::step::RequestUpdate; use super::step::ScheduledInfo; use super::step::StepOutputs; use super::step::Terminal; +use super::stop::StopCause; /// One open account: the request's admission facts and running tally. The /// payload is not here — it went to the scheduler at `submit`; the account is @@ -239,12 +240,23 @@ impl RequestLedger { /// Finish the request. Token counts come from the ledger's tally. pub fn finish(&mut self, id: RequestId, reason: FinishReason) { + self.finish_with_cause(id, reason, None); + } + + /// Finish a request while preserving a typed token-level stop cause. + pub fn finish_with_cause( + &mut self, + id: RequestId, + reason: FinishReason, + stop_cause: Option, + ) { let account = self.close(id); let AccountState::Active { completion_tokens } = account.state else { panic!("finish on {id} before admission"); }; self.statement.entry(id).terminal = Some(Terminal::Finished { reason, + stop_cause, prompt_tokens: account.prompt_len, completion_tokens, }); @@ -279,6 +291,16 @@ impl RequestLedger { /// this step — tokens included — folds into the returned message, so late /// delivery cannot reorder against the step stream. pub fn defer_finish(&mut self, id: RequestId, reason: FinishReason) -> DeferredFinish { + self.defer_finish_with_cause(id, reason, None) + } + + /// Defer a finish while preserving a typed token-level stop cause. + pub fn defer_finish_with_cause( + &mut self, + id: RequestId, + reason: FinishReason, + stop_cause: Option, + ) -> DeferredFinish { let account = self.close(id); let AccountState::Active { completion_tokens } = account.state else { panic!("defer_finish on {id} before admission"); @@ -289,6 +311,7 @@ impl RequestLedger { .unwrap_or_else(|| RequestUpdate::empty(id)); update.terminal = Some(Terminal::Finished { reason, + stop_cause, prompt_tokens: account.prompt_len, completion_tokens, }); @@ -374,6 +397,7 @@ mod tests { use super::super::request_lifecycle::StepReceiver; use super::super::step::Request; use super::super::step::Terminal; + use super::super::stop::StopPolicy; use super::super::wiring::SchedulerHandle; use super::super::wiring::scheduler_pair; use super::*; @@ -382,6 +406,7 @@ mod tests { Request { prompt_tokens: prompt, params: crate::sampler::SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 8, lora_adapter: None, kv_transfer_params: None, @@ -402,7 +427,9 @@ mod tests { backend.ledger.admit(id); backend.ledger.push_tokens(id, &[10, 11], &[]); backend.ledger.set_cached_tokens(id, 2); - backend.ledger.finish(id, FinishReason::Stop); + backend + .ledger + .finish_with_cause(id, FinishReason::Stop, Some(StopCause::Token(11))); backend.ledger.commit_step(); let mut steps = handle_steps(handle); @@ -419,6 +446,7 @@ mod tests { update.terminal, Some(Terminal::Finished { reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(11)), prompt_tokens: 3, completion_tokens: 2, }) @@ -479,7 +507,9 @@ mod tests { let id = backend.ledger.register(envelope).id; backend.ledger.admit(id); backend.ledger.push_tokens(id, &[7], &[]); - let deferred = backend.ledger.defer_finish(id, FinishReason::Length); + let deferred = backend + .ledger + .defer_finish_with_cause(id, FinishReason::Length, None); backend.ledger.commit_step(); let mut steps = handle_steps(handle); @@ -498,6 +528,7 @@ mod tests { update.terminal, Some(Terminal::Finished { reason: FinishReason::Length, + stop_cause: None, prompt_tokens: 2, completion_tokens: 1, }) diff --git a/pegainfer-frontend/src/engine/mod.rs b/pegainfer-frontend/src/engine/mod.rs index 2c69a3d63..a2a8382fd 100644 --- a/pegainfer-frontend/src/engine/mod.rs +++ b/pegainfer-frontend/src/engine/mod.rs @@ -39,6 +39,7 @@ mod request; mod request_lifecycle; mod sink; mod step; +mod stop; mod wiring; pub use control::*; @@ -52,4 +53,5 @@ pub use request::*; pub use request_lifecycle::*; pub use sink::*; pub use step::*; +pub use stop::*; pub use wiring::*; diff --git a/pegainfer-frontend/src/engine/step.rs b/pegainfer-frontend/src/engine/step.rs index 50fa3df58..01d164382 100644 --- a/pegainfer-frontend/src/engine/step.rs +++ b/pegainfer-frontend/src/engine/step.rs @@ -14,6 +14,8 @@ use std::time::Instant; use super::event::FinishReason; use super::event::TokenLogprob; +use super::stop::StopCause; +use super::stop::StopPolicy; /// In-process routing id for one generate request, minted by /// [`super::SchedulerHandle::submit`] from a per-scheduler counter. `Copy` and @@ -49,6 +51,7 @@ impl std::fmt::Display for RequestId { pub struct Request { pub prompt_tokens: Vec, pub params: crate::sampler::SamplingParams, + pub stop_policy: StopPolicy, pub max_tokens: usize, pub lora_adapter: Option, /// Opaque router/P-D metadata from the request's @@ -233,6 +236,10 @@ impl fmt::Display for RejectReason { pub enum Terminal { Finished { reason: FinishReason, + /// Present for token-driven stop finishes. The triggering token remains + /// in `RequestUpdate.tokens`, with its real logprob in the matching + /// `RequestUpdate.logprobs` entry. + stop_cause: Option, prompt_tokens: usize, completion_tokens: usize, }, diff --git a/pegainfer-frontend/src/engine/stop.rs b/pegainfer-frontend/src/engine/stop.rs new file mode 100644 index 000000000..83c62b3fc --- /dev/null +++ b/pegainfer-frontend/src/engine/stop.rs @@ -0,0 +1,121 @@ +use std::sync::Arc; + +/// How a request treats end-of-sequence tokens. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum EosPolicy { + /// Do not stop on model EOS tokens. + Ignore, + /// Use the model executor's configured EOS set. + #[default] + ModelDefault, +} + +/// 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 { + eos: EosPolicy, + /// Sorted and deduplicated explicit stop IDs. + /// + /// Requests clone this policy while building a plan and sending it to + /// worker ranks. Keeping the normalized set behind an `Arc` makes those + /// clones cheap and lets classification use binary search for large stop + /// sets without regressing the common one-ID case. + token_ids: Arc<[u32]>, +} + +impl StopPolicy { + /// Build a policy from wire-provided stop IDs. + /// + /// Normalization happens once at the request boundary. Internal copies can + /// then share the immutable slice instead of repeatedly sorting, deduping, + /// or cloning the caller's vector. + #[must_use] + pub fn new(eos: EosPolicy, mut token_ids: Vec) -> Self { + token_ids.sort_unstable(); + token_ids.dedup(); + Self { + eos, + token_ids: token_ids.into(), + } + } + + /// Classify a token using vLLM's priority: EOS first, then the request's + /// explicit stop-token set. + #[must_use] + pub fn classify( + &self, + token_id: u32, + is_model_eos: impl FnOnce(u32) -> bool, + ) -> Option { + let is_eos = match self.eos { + EosPolicy::Ignore => false, + EosPolicy::ModelDefault => is_model_eos(token_id), + }; + + if is_eos { + Some(StopCause::Eos(token_id)) + } else if self.token_ids.binary_search(&token_id).is_ok() { + Some(StopCause::Token(token_id)) + } else { + None + } + } +} + +/// The token-level cause of a normal stop finish. +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +pub enum StopCause { + /// A primary or model-default EOS token. + Eos(u32), + /// A token from the request's explicit stop-token set. + Token(u32), +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn model_default_classifies_model_eos() { + let policy = StopPolicy::default(); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Eos(99)) + ); + } + + #[test] + fn ignored_eos_does_not_disable_an_explicit_stop() { + let policy = StopPolicy::new(EosPolicy::Ignore, vec![99]); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Token(99)) + ); + } + + #[test] + fn normalizes_unsorted_duplicate_stop_ids() { + let policy = StopPolicy::new(EosPolicy::Ignore, vec![7, 3, 7, 1]); + + assert_eq!(policy.classify(1, |_| false), Some(StopCause::Token(1))); + assert_eq!(policy.classify(3, |_| false), Some(StopCause::Token(3))); + assert_eq!(policy.classify(7, |_| false), Some(StopCause::Token(7))); + assert!(policy.classify(8, |_| false).is_none()); + } + + #[test] + fn model_eos_has_priority_over_explicit_stop() { + let policy = StopPolicy::new(EosPolicy::ModelDefault, vec![99]); + + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(StopCause::Eos(99)) + ); + } +} diff --git a/pegainfer-frontend/src/vllm/bridge/stepped.rs b/pegainfer-frontend/src/vllm/bridge/stepped.rs index 0a1d037cc..731fbecfe 100644 --- a/pegainfer-frontend/src/vllm/bridge/stepped.rs +++ b/pegainfer-frontend/src/vllm/bridge/stepped.rs @@ -55,9 +55,11 @@ use crate::engine::RequestId; use crate::engine::RequestUpdate; use crate::engine::SchedulerHandle; use crate::engine::StepOutputs; +use crate::engine::StopCause; use crate::engine::Terminal; use crate::vllm::wire::convert_finish_reason; use crate::vllm::wire::convert_sampling; +use crate::vllm::wire::convert_stop_policy; use crate::vllm::wire::lora_adapter_from_sampling_params; use crate::vllm::wire::requested_logprobs; use crate::vllm::wire::to_wire_position_logprobs; @@ -362,6 +364,7 @@ impl SteppedEngineBridge { None, ); } + let lora_adapter = match lora_adapter_from_sampling_params(&sampling_params) { Ok(adapter) => adapter, Err(error) => { @@ -383,6 +386,10 @@ impl SteppedEngineBridge { .as_ref() .and_then(|args| args.get("kv_transfer_params")) .cloned(); + // Older stepped model producers still suppress their terminal token + // and report only `FinishReason::Stop`. Keep the legacy sentinel for + // that producer shape; typed stop causes carry the real token and do + // not need a synthetic suffix. let stop_sentinel_id = stop_sentinel_id( sampling_params.eos_token_id, &sampling_params.stop_token_ids, @@ -398,9 +405,14 @@ impl SteppedEngineBridge { Span::noop() }; let trace_parent = SpanContext::from_span(&trace_root); + let control = self.scheduler.submit(Request { prompt_tokens, + // Keep the legacy SamplingParams lowering unchanged for stepped + // producers that have not migrated to StopPolicy. Qwen3 uses the + // independent policy below for stop classification. params: convert_sampling(&sampling_params), + stop_policy: convert_stop_policy(&sampling_params), max_tokens: sampling_params.max_tokens as usize, lora_adapter, kv_transfer_params, @@ -435,8 +447,9 @@ struct SteppedStream { /// P/D handoff metadata can arrive in an update with no token or /// terminal, so retain it until the next output carries it to the router. kv_transfer_params: Option, - /// The vLLM text decoder removes the final token from a stop-finished - /// output. Keep an EOS or explicit stop token as that removable sentinel. + /// Compatibility sentinel for stepped producers that predate typed + /// [`StopCause`]. New producers must include their triggering token in the + /// update and therefore bypass this fallback. stop_sentinel_id: Option, /// Request-lifetime root span; held only for its `Drop`, which closes the /// trace when the stream state is removed. @@ -539,11 +552,11 @@ fn reduce_update( let mut terminated = false; match update.terminal { None => {} - Some(Terminal::Finished { reason, .. }) => { - // PegaInfer suppresses EOS before emitting tokens, while vLLM's - // text decoder expects the terminal Stop output to contain EOS - // and unconditionally removes its final token. + Some(Terminal::Finished { + reason, stop_cause, .. + }) => { if reason == FinishReason::Stop + && stop_cause.is_none() && let Some(stop_sentinel_id) = state.stop_sentinel_id { token_ids.push(stop_sentinel_id); @@ -551,6 +564,10 @@ fn reduce_update( entries: Vec::new(), }); } + if let Some(StopCause::Token(token_id)) = stop_cause { + stop_reason = Some(StopReason::TokenId(token_id)); + } + finish_reason = Some(convert_finish_reason(reason)); terminated = true; } @@ -616,14 +633,19 @@ impl UnixAnchor { #[cfg(test)] mod tests { + use std::sync::atomic::AtomicBool; + use super::*; use crate::engine::RejectReason; + use crate::engine::StopPolicy; + use crate::engine::TokenLogprob; use crate::engine::scheduler_pair; fn request() -> Request { Request { prompt_tokens: vec![1, 2], params: crate::sampler::SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 1, lora_adapter: None, kv_transfer_params: None, @@ -662,4 +684,93 @@ mod tests { "a request refused while queued did no prefill" ); } + + #[test] + fn request_stop_maps_the_actual_token_and_preserves_its_logprob() { + let id = RequestId::new(7); + let control = RequestControl::new(id, Arc::new(AtomicBool::new(false))); + let mut state = SteppedStream::new("request-7".to_string(), control, Span::noop(), None); + + let mut update = RequestUpdate::empty(id); + update.tokens = vec![11, 43]; + update.logprobs = vec![ + None, + Some(TokenLogprob { + logprob: -0.25, + top_logprobs: vec![(43, -0.25), (44, -1.0)], + }), + ]; + update.terminal = Some(Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(43)), + prompt_tokens: 16, + completion_tokens: 2, + }); + + let (output, terminated) = reduce_update(&mut state, update, &UnixAnchor::now()); + let output = output.expect("terminal output"); + + assert!(terminated); + assert_eq!(output.new_token_ids, vec![11, 43]); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); + assert_eq!(output.stop_reason, Some(StopReason::TokenId(43))); + + let direct = match output.new_logprobs.expect("stop-token logprob") { + MaybeWireLogprobs::Direct(direct) => direct, + MaybeWireLogprobs::Wire(_) => panic!("expected direct logprobs"), + }; + + assert_eq!(direct.positions.len(), 2); + assert_eq!(direct.positions[1].entries[0].token_id, 43); + assert!((direct.positions[1].entries[0].logprob + 0.25).abs() < f32::EPSILON); + } + + #[test] + fn model_eos_has_no_wire_stop_reason() { + let id = RequestId::new(8); + let control = RequestControl::new(id, Arc::new(AtomicBool::new(false))); + let mut state = SteppedStream::new("request-8".to_string(), control, Span::noop(), None); + + let mut update = RequestUpdate::empty(id); + update.tokens = vec![2]; + update.logprobs = vec![None]; + update.terminal = Some(Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Eos(2)), + prompt_tokens: 16, + completion_tokens: 1, + }); + + let (output, terminated) = reduce_update(&mut state, update, &UnixAnchor::now()); + let output = output.expect("terminal output"); + + assert!(terminated); + assert_eq!(output.new_token_ids, vec![2]); + assert_eq!(output.finish_reason, Some(EngineCoreFinishReason::Stop)); + assert_eq!(output.stop_reason, None); + } + + #[test] + fn legacy_stop_without_typed_cause_keeps_the_wire_sentinel() { + let id = RequestId::new(9); + let control = RequestControl::new(id, Arc::new(AtomicBool::new(false))); + let mut state = + SteppedStream::new("request-9".to_string(), control, Span::noop(), Some(99)); + + let mut update = RequestUpdate::empty(id); + update.tokens = vec![11]; + update.terminal = Some(Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: None, + prompt_tokens: 16, + completion_tokens: 1, + }); + + let (output, terminated) = reduce_update(&mut state, update, &UnixAnchor::now()); + let output = output.expect("terminal output"); + + assert!(terminated); + assert_eq!(output.new_token_ids, vec![11, 99]); + assert_eq!(output.stop_reason, None); + } } diff --git a/pegainfer-frontend/src/vllm/wire.rs b/pegainfer-frontend/src/vllm/wire.rs index fbc5a9f74..3c22727be 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; @@ -42,10 +44,10 @@ pub(crate) fn to_wire_position_logprobs( pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingParams { // The vLLM frontend lowers a client `ignore_eos=true` to `_eos_token_id: // None`, but `_all_stop_token_ids` always carries the model EOS set (it - // exists for min_tokens masking, not stop detection). Deriving ignore_eos - // from all_stop_token_ids would therefore void every ignore_eos request on - // models with a real EOS. Only `_eos_token_id` and the client's explicit - // `stop_token_ids` express a stop intent. + // exists for min_tokens masking, not stop detection). This conversion feeds + // the legacy SamplingParams contract, which has no field for explicit + // request stop IDs; keep its historical lowering until each producer is + // migrated to StopPolicy. Qwen3 receives the independent policy below. let ignore_eos = params.eos_token_id.is_none() && params.stop_token_ids.is_empty(); if params.temperature <= 0.0 { return SamplingParams { @@ -76,7 +78,23 @@ pub(crate) fn convert_sampling(params: &EngineCoreSamplingParams) -> SamplingPar } } -/// Reject request parameters the engine would otherwise silently ignore. +pub(crate) fn convert_stop_policy(params: &EngineCoreSamplingParams) -> StopPolicy { + StopPolicy::new( + // Qwen3 owns the complete model EOS set in generation_config. The + // protocol's optional primary ID only tells us whether EOS is active; + // using it as a singleton would miss secondary model EOS IDs. + params + .eos_token_id + .map_or(EosPolicy::Ignore, |_| EosPolicy::ModelDefault), + params.stop_token_ids.clone(), + ) +} + +/// Reject request parameters the frontend cannot represent faithfully. +/// +/// The stepped contract carries explicit request stop IDs independently in +/// [`StopPolicy`]; this helper only validates unrelated sampling/transfer +/// fields that would otherwise be silently ignored. /// Returns the offending description; `None` means the request is servable. /// /// The float comparisons are exact on purpose: they detect "the client sent @@ -84,6 +102,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,13 +215,31 @@ mod tests { params.eos_token_id = Some(163_586); assert!(!convert_sampling(¶ms).ignore_eos); - // Explicit client stop tokens keep EOS detection on even when the - // frontend dropped _eos_token_id. + // The legacy SamplingParams contract keeps EOS active when explicit + // stop IDs are present; the stepped Qwen3 path carries those IDs in + // StopPolicy instead. params.eos_token_id = None; params.stop_token_ids = vec![42]; assert!(!convert_sampling(¶ms).ignore_eos); } + #[test] + fn convert_stop_policy_keeps_eos_and_explicit_stops_independent() { + let mut params = EngineCoreSamplingParams::for_test(); + params.eos_token_id = Some(99); + params.stop_token_ids = vec![11]; + + let policy = convert_stop_policy(¶ms); + assert_eq!( + policy.classify(99, |token_id| token_id == 99), + Some(crate::engine::StopCause::Eos(99)) + ); + assert_eq!( + policy.classify(11, |_| false), + Some(crate::engine::StopCause::Token(11)) + ); + } + #[test] fn convert_sampling_passes_min_p_and_never_seed() { let mut params = EngineCoreSamplingParams::for_test(); @@ -222,6 +264,13 @@ mod tests { params.repetition_penalty = 1.0; assert_eq!(unsupported_request_params(¶ms), None); + params.min_tokens = 1; + assert_eq!( + unsupported_request_params(¶ms).as_deref(), + Some("min_tokens=1 is not supported by current engine contracts") + ); + params.min_tokens = 0; + params.min_p = 0.2; assert_eq!(unsupported_request_params(¶ms), None); params.min_p = 1.5; diff --git a/pegainfer-gemma4/src/engine/lane_tests.rs b/pegainfer-gemma4/src/engine/lane_tests.rs index 1a5e145db..1102ddec6 100644 --- a/pegainfer-gemma4/src/engine/lane_tests.rs +++ b/pegainfer-gemma4/src/engine/lane_tests.rs @@ -55,6 +55,7 @@ impl Harness { ignore_eos: true, ..pegainfer_frontend::sampler::SamplingParams::default() }, + stop_policy: pegainfer_frontend::engine::StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-k3/src/scheduler/tests.rs b/pegainfer-k3/src/scheduler/tests.rs index ac1c0b0b8..afd787841 100644 --- a/pegainfer-k3/src/scheduler/tests.rs +++ b/pegainfer-k3/src/scheduler/tests.rs @@ -23,6 +23,7 @@ use pegainfer_frontend::engine::Request; use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestUpdate; use pegainfer_frontend::engine::StepReceiver; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; @@ -175,6 +176,7 @@ fn request(prompt_len: usize, max_tokens: usize) -> Request { Request { prompt_tokens: vec![7; prompt_len], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -313,6 +315,7 @@ fn admitted_request_streams_its_tokens_and_finishes_at_max_tokens() { reason: FinishReason::Length, prompt_tokens: 4, completion_tokens: 3, + .. } ), "{terminal:?}" diff --git a/pegainfer-qwen3/src/executor.rs b/pegainfer-qwen3/src/executor.rs index 086afc75b..c29da4e7b 100644 --- a/pegainfer-qwen3/src/executor.rs +++ b/pegainfer-qwen3/src/executor.rs @@ -19,6 +19,7 @@ use pegainfer_core::weight_loader::load_shard_info; use pegainfer_frontend::engine::DeferredFinish; use pegainfer_frontend::engine::LoadLoraAdapterRequest; use pegainfer_frontend::engine::SpecDecodeCounters; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::engine::panic_message; @@ -594,6 +595,7 @@ fn execute_step_on_lane( StepCommand::SpeculativeVerify { requests, kv_views, + stop_policies, sample_seed, } => { // One target forward over each request's K+1 draft span with a @@ -602,7 +604,8 @@ fn execute_step_on_lane( // token at each span position) and captures the target hidden states // (at the DFlash layers) to seed the next draft — all into reused, // pointer-stable scratch (`VerifyGraphBuffers`). - let result = lane.execute_dflash_verify(requests, kv_views, *sample_seed)?; + let result = + lane.execute_dflash_verify(requests, kv_views, stop_policies, *sample_seed)?; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -3525,6 +3528,7 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], + stop_policies: &[StopPolicy], capture_layer_ids: &[usize], sample_seed: u64, bufs: &mut VerifyGraphBuffers, @@ -3632,9 +3636,45 @@ impl LocalQwen3Lane { .flat_map(|req| std::iter::repeat_n(&req.params, req.as_slice().len())) .collect(); let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, sample_seed)?; - let all_results = build_verify_results(&expanded, &target_tokens)?; - let (results_a, results_b) = all_results.split_at(requests.len()); - + let mut all_results = build_verify_results(&expanded, &target_tokens)?; + // A terminal token ends the request even when it appears in the middle + // of a speculative span. Normalize every candidate before selecting a + // winner; otherwise a discarded suffix can win the hedge and advance + // the wrong KV/hidden state and acceptance statistics. + let (results_a, results_b) = all_results.split_at_mut(requests.len()); + anyhow::ensure!( + results_b.len() == hedge_spans.len(), + "hedge returned {} B results for {} hedge spans", + results_b.len(), + hedge_spans.len() + ); + let trace = std::env::var_os("PEGAINFER_TEST_LOG").is_some(); + let raw_a_lengths: Vec = if trace { + results_a + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + } else { + Vec::new() + }; + let raw_b_lengths: Vec = if trace { + results_b + .iter() + .map(|result| result.accepted_tokens.len()) + .collect() + } else { + Vec::new() + }; + for (result, policy) in results_a.iter_mut().zip(stop_policies) { + spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + } + for (slot, (idx, _)) in hedge_spans.iter().enumerate() { + spec::truncate_after_terminal( + &mut results_b[slot], + &stop_policies[*idx], + &self.model.config().stop_token_ids, + ); + } // Per request keep the best-accepting chain; ties keep chain A (no // copies). A later chain of the same request only replaces the // running winner when strictly better, so the final page/hidden @@ -3644,6 +3684,11 @@ impl LocalQwen3Lane { let elem = std::mem::size_of::(); let mut final_requests: Vec = requests.to_vec(); let mut final_results: Vec = results_a.to_vec(); + let mut selected_is_b = if trace { + Some(vec![false; requests.len()]) + } else { + None + }; let mut b_wins = 0usize; let mut b_row_offset = a_total_rows; for (slot, (idx, replaced)) in hedge_spans.iter().enumerate() { @@ -3663,12 +3708,15 @@ impl LocalQwen3Lane { cudarc::driver::result::memcpy_dtod_async( dst, src, - span_len * hidden_dim * elem, + res_b.accepted_tokens.len() * hidden_dim * elem, ctx.stream.cu_stream(), ) } .map_err(|e| anyhow::anyhow!("hedge hidden compaction failed: {e}"))?; b_wins += 1; + if let Some(selected) = selected_is_b.as_mut() { + selected[*idx] = true; + } final_requests[*idx] = expanded[requests.len() + slot].clone(); final_results[*idx] = res_b.clone(); } @@ -3692,6 +3740,37 @@ impl LocalQwen3Lane { &final_results, Some(bufs.captured_hidden()), )?; + if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + for idx in 0..requests.len() { + let has_hedge = hedge_spans + .iter() + .any(|(request_idx, _)| *request_idx == idx); + if !has_hedge { + continue; + } + let selected = if selected_is_b.as_ref().is_some_and(|selected| selected[idx]) { + 'B' + } else { + 'A' + }; + let raw_b_lens = hedge_spans + .iter() + .enumerate() + .filter(|(_, (request_idx, _))| *request_idx == idx) + .map(|(slot, _)| raw_b_lengths[slot].to_string()) + .collect::>() + .join(","); + log::debug!( + "Qwen3 DFlash hedge detail request={} raw_a={} raw_b_lens={} selected={} selected_len={} matched_draft={}", + requests[idx].request_id, + raw_a_lengths[idx], + raw_b_lens, + selected, + final_results[idx].accepted_tokens.len(), + final_results[idx].matched_draft_tokens, + ); + } + } Ok(Some(VerifyResult { requests: final_results, })) @@ -3705,8 +3784,15 @@ impl LocalQwen3Lane { &mut self, requests: &[VerifyStepItem], kv_views: &[KvView], + stop_policies: &[StopPolicy], sample_seed: u64, ) -> Result { + anyhow::ensure!( + stop_policies.len() == requests.len(), + "DFlash verify received {} stop policies for {} requests", + stop_policies.len(), + requests.len() + ); let capture_layer_ids = self.dflash_capture_layer_ids().ok_or_else(|| { anyhow::anyhow!("DFlash verify requested but no draft model is loaded") })?; @@ -3756,6 +3842,7 @@ impl LocalQwen3Lane { if let Some(result) = self.try_execute_hedged_verify( requests, kv_views, + stop_policies, &capture_layer_ids, sample_seed, &mut bufs, @@ -3791,8 +3878,13 @@ impl LocalQwen3Lane { .flat_map(|req| std::iter::repeat_n(&req.params, req.as_slice().len())) .collect(); let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, sample_seed)?; - let request_results = build_verify_results(requests, &target_tokens)?; - + let mut request_results = build_verify_results(requests, &target_tokens)?; + // Apply the request policy before recording target hidden states; + // otherwise a suffix discarded by terminal handling would leak + // into the next DFlash draft context and acceptance counters. + for (policy, result) in stop_policies.iter().zip(&mut request_results) { + spec::truncate_after_terminal(result, policy, &self.model.config().stop_token_ids); + } self.record_verify_dflash_context( requests, &request_results, @@ -3905,6 +3997,9 @@ enum StepCommand { SpeculativeVerify { requests: Vec, kv_views: Vec, + /// Request-local stop policies used before DFlash context recording; + /// these are host metadata and never enter the GPU batch. + stop_policies: Vec, sample_seed: u64, }, /// Speculative draft: roll the DFlash draft model forward one block per diff --git a/pegainfer-qwen3/src/executor/spec.rs b/pegainfer-qwen3/src/executor/spec.rs index 488f261f0..1fffd0d35 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,34 @@ use super::WorkerStepOutcome; use crate::speculative::DraftPlan; use crate::speculative::DraftResult; use crate::speculative::VerifyPlan; +use crate::speculative::VerifyRequestResult; use crate::speculative::VerifyResult; +/// Remove accepted tokens after the first request-terminal token before the +/// speculative KV transaction commits. The scheduler classifies the same +/// retained trigger later to produce the typed protocol stop cause. +pub(super) fn truncate_after_terminal( + result: &mut VerifyRequestResult, + policy: &StopPolicy, + model_eos: &[u32], +) { + // Keep this helper idempotent: the worker trims before DFlash context is + // recorded, and the executor repeats the invariant before KV commit. + let Some(keep) = result.accepted_tokens.iter().position(|&token| { + policy + .classify(token, |id| model_eos.contains(&id)) + .is_some() + }) else { + return; + }; + let keep = keep + 1; + result.accepted_tokens.truncate(keep); + // A trigger in the accepted draft prefix is itself a matched draft. If the + // trigger is the posterior token, the original count is already `keep - 1`; + // clamping to the retained prefix handles both cases. + result.matched_draft_tokens = result.matched_draft_tokens.min(keep); +} + impl Qwen3Executor { pub(super) fn execute_speculative_verify_impl( &mut self, @@ -25,6 +52,12 @@ impl Qwen3Executor { self.speculative.is_some(), "speculative verification requested but no draft model is loaded" ); + anyhow::ensure!( + plan.stop_policies.len() == plan.requests.len(), + "speculative verify received {} stop policies for {} requests", + plan.stop_policies.len(), + plan.requests.len() + ); for req in plan.requests { anyhow::ensure!( !req.as_slice().is_empty(), @@ -71,6 +104,7 @@ impl Qwen3Executor { let step = StepCommand::SpeculativeVerify { requests: plan.requests.to_vec(), kv_views, + stop_policies: plan.stop_policies.to_vec(), sample_seed: plan.sample_seed, }; let outcome = match self.run_step(&step) { @@ -80,7 +114,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 +142,13 @@ impl Qwen3Executor { )); } } + // The worker normally applies the request contract before copying + // worker-side state. Recheck the same invariant here before touching + // RequestKv so a terminal suffix is rolled back with its reservation, + // including legacy workers that return an untrimmed span. + for (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 @@ -130,6 +171,13 @@ impl Qwen3Executor { req_result.request_id )); } + if std::env::var_os("PEGAINFER_TEST_LOG").is_some() { + log::debug!( + "Qwen3 DFlash commit request={} accepted_len={}", + req_result.request_id, + req_result.accepted_tokens.len(), + ); + } applied.push(req_result.request_id); } for req_result in &result.requests { @@ -191,3 +239,25 @@ impl Qwen3Executor { } } } + +#[cfg(test)] +mod tests { + use pegainfer_frontend::engine::EosPolicy; + + use super::*; + + #[test] + fn explicit_stop_truncates_the_kv_commit_after_the_trigger() { + let policy = StopPolicy::new(EosPolicy::Ignore, vec![7]); + let mut result = VerifyRequestResult { + request_id: RequestId::new(1), + matched_draft_tokens: 3, + accepted_tokens: vec![5, 7, 8, 9], + }; + + truncate_after_terminal(&mut result, &policy, &[99]); + + assert_eq!(result.accepted_tokens, vec![5, 7]); + assert_eq!(result.matched_draft_tokens, 2); + } +} diff --git a/pegainfer-qwen3/src/frontend_adapter.rs b/pegainfer-qwen3/src/frontend_adapter.rs index e02125ea5..c89d8aec4 100644 --- a/pegainfer-qwen3/src/frontend_adapter.rs +++ b/pegainfer-qwen3/src/frontend_adapter.rs @@ -42,6 +42,7 @@ use pegainfer_frontend::engine::RejectReason; use pegainfer_frontend::engine::RequestLedger; use pegainfer_frontend::engine::Scheduler; use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::engine::spawn_scheduler; use pegainfer_kernels::ops::NumericPolicy; use pegainfer_kernels::ops::numeric_policy; @@ -343,7 +344,7 @@ impl Qwen3Scheduler { // terminal rides the committed step, which the driver ships after // publishing metrics — the finishing batch's send-time stats then // read the drained occupancy instead of racing the publish. - let mut finishes: Vec<(RequestId, FinishReason)> = Vec::new(); + let mut finishes: Vec<(RequestId, FinishReason, Option)> = Vec::new(); for cached in effects.cached { if ledger.is_active(cached.request_id) { @@ -367,32 +368,11 @@ impl Qwen3Scheduler { for effect in effects.decode { match effect { DecodeEffect::Finish { - request_id, - finish_reason, - } => { - let Some(index) = self - .active - .iter() - .position(|req| req.request_id == request_id) - else { - continue; - }; - if ledger.is_active(request_id) { - if ledger.is_aborted(request_id) { - ledger.retire(request_id); - } else { - finishes.push((request_id, finish_reason)); - } - } - self.tracker.finish(request_id); - let _ = self.executor.drop_request(request_id); - to_retire.push(index); - } - DecodeEffect::EmitAndFinish { request_id, token, logprob, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -406,14 +386,14 @@ impl Qwen3Scheduler { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &[token], &[logprob]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); let _ = self.executor.drop_request(request_id); to_retire.push(index); } - DecodeEffect::EmitAndContinue { + DecodeEffect::Continue { request_id, token, logprob, @@ -436,7 +416,7 @@ impl Qwen3Scheduler { req.generated_count = completion_tokens; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id, tokens, completion_tokens, @@ -460,10 +440,11 @@ impl Qwen3Scheduler { } } } - DecodeEffect::EmitManyAndFinish { + DecodeEffect::FinishMany { request_id, tokens, finish_reason, + stop_cause, } => { let Some(index) = self .active @@ -477,7 +458,7 @@ impl Qwen3Scheduler { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &tokens, &[]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); @@ -507,32 +488,19 @@ impl Qwen3Scheduler { continued.push(req); } } - PendingEffect::Finish { - request_id, - finish_reason, - } => { - if ledger.is_active(request_id) { - if ledger.is_aborted(request_id) { - ledger.retire(request_id); - } else { - finishes.push((request_id, finish_reason)); - } - } - self.tracker.finish(request_id); - let _ = self.executor.drop_request(request_id); - } PendingEffect::EmitAndFinish { request_id, token, logprob, finish_reason, + stop_cause, } => { if ledger.is_active(request_id) { if ledger.is_aborted(request_id) { ledger.retire(request_id); } else { ledger.push_tokens(request_id, &[token], &[logprob]); - finishes.push((request_id, finish_reason)); + finishes.push((request_id, finish_reason, stop_cause)); } } self.tracker.finish(request_id); @@ -560,12 +528,14 @@ impl Qwen3Scheduler { if self.executor.withholds_finishes() { let withheld: Vec = finishes .into_iter() - .map(|(request_id, reason)| ledger.defer_finish(request_id, reason)) + .map(|(request_id, reason, stop_cause)| { + ledger.defer_finish_with_cause(request_id, reason, stop_cause) + }) .collect(); self.executor.release_finished_events(withheld); } else { - for (request_id, reason) in finishes { - ledger.finish(request_id, reason); + for (request_id, reason, stop_cause) in finishes { + ledger.finish_with_cause(request_id, reason, stop_cause); } } } diff --git a/pegainfer-qwen3/src/frontend_adapter/tests.rs b/pegainfer-qwen3/src/frontend_adapter/tests.rs index 35441bc7d..1d1b2244d 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,62 @@ fn wait_until(timeout: Duration, mut predicate: impl FnMut() -> bool) -> bool { false } +#[test] +fn request_stop_token_beats_length_during_prefill_when_eos_is_ignored() { + let executor = FakeExecutor::new(4, Arc::new(Mutex::new(Vec::new()))).with_logprobs(); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 1); + req.stop_policy = pegainfer_frontend::engine::StopPolicy::new( + pegainfer_frontend::engine::EosPolicy::Ignore, + vec![100], + ); + + let control = partition.handle.submit(req); + let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); + + assert_eq!(tokens, vec![100]); + assert_eq!(logprobs.len(), 1); + assert!((logprobs[0].as_ref().expect("stop-token logprob").logprob + 0.1).abs() < f32::EPSILON); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(100)), + completion_tokens: 1, + .. + } + )); +} + +#[test] +fn request_stop_token_beats_length_during_decode_when_eos_is_ignored() { + let executor = FakeExecutor::new(4, Arc::new(Mutex::new(Vec::new()))).with_logprobs(); + let (partition, _lora, mut steps) = launch(executor, false); + + let mut req = request(16, 2); + req.stop_policy = pegainfer_frontend::engine::StopPolicy::new( + pegainfer_frontend::engine::EosPolicy::Ignore, + vec![200], + ); + + let control = partition.handle.submit(req); + let (tokens, logprobs, terminal) = steps.collect_terminal_with_logprobs(control.id()); + + assert_eq!(tokens, vec![100, 200]); + assert_eq!(logprobs.len(), 2); + assert!((logprobs[1].as_ref().expect("stop-token logprob").logprob + 0.2).abs() < f32::EPSILON); + assert!(matches!( + terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(pegainfer_frontend::engine::StopCause::Token(200)), + completion_tokens: 2, + .. + } + )); +} + #[test] fn unknown_lora_request_is_rejected_without_blocking_base_request() { let dropped = Arc::new(Mutex::new(Vec::new())); diff --git a/pegainfer-qwen3/src/scheduler.rs b/pegainfer-qwen3/src/scheduler.rs index dcd3da28d..5ba9e0231 100644 --- a/pegainfer-qwen3/src/scheduler.rs +++ b/pegainfer-qwen3/src/scheduler.rs @@ -18,6 +18,7 @@ use std::collections::HashSet; use log::debug; use log::warn; use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use crate::executor::ModelExecutor; @@ -34,6 +35,7 @@ pub(crate) struct ActiveRequestState { pub(crate) max_tokens: usize, pub(crate) prompt_len: usize, pub(crate) params: SamplingParams, + pub(crate) stop_policy: StopPolicy, /// Number of top logprobs to return (0 = disabled). pub(crate) logprobs: usize, } @@ -47,6 +49,7 @@ pub(crate) struct PendingRequest { pub(crate) lora_adapter: Option, pub(crate) prompt_tokens: Vec, pub(crate) params: SamplingParams, + pub(crate) stop_policy: StopPolicy, pub(crate) max_tokens: usize, pub(crate) logprobs: usize, pub(crate) echo: bool, @@ -73,6 +76,7 @@ impl PendingRequest { lora_adapter: req.lora_adapter, prompt_tokens: req.prompt_tokens, params: req.params, + stop_policy: req.stop_policy, max_tokens: req.max_tokens, logprobs: req.logprobs, echo: req.echo, diff --git a/pegainfer-qwen3/src/scheduler/effects.rs b/pegainfer-qwen3/src/scheduler/effects.rs index 8e37a70e3..69d78411c 100644 --- a/pegainfer-qwen3/src/scheduler/effects.rs +++ b/pegainfer-qwen3/src/scheduler/effects.rs @@ -7,6 +7,7 @@ //! resolve logic stay a pure function of executor results. use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; use pegainfer_frontend::engine::TokenLogprob; use super::ActiveRequestState; @@ -27,15 +28,12 @@ pub(crate) struct PromptEchoEffect { } pub(crate) enum PendingEffect { - Finish { - request_id: RequestId, - finish_reason: FinishReason, - }, EmitAndFinish { request_id: RequestId, token: u32, logprob: Option, finish_reason: FinishReason, + stop_cause: Option, }, Promote { state: ActiveRequestState, @@ -49,16 +47,13 @@ pub(crate) enum PendingEffect { pub(crate) enum DecodeEffect { Finish { - request_id: RequestId, - finish_reason: FinishReason, - }, - EmitAndFinish { request_id: RequestId, token: u32, logprob: Option, finish_reason: FinishReason, + stop_cause: Option, }, - EmitAndContinue { + Continue { request_id: RequestId, token: u32, logprob: Option, @@ -68,17 +63,18 @@ pub(crate) enum DecodeEffect { completion_tokens: usize, }, /// Commit several accepted speculative tokens and keep the request running. - EmitManyAndContinue { + ContinueMany { request_id: RequestId, tokens: Vec, completion_tokens: usize, }, /// Commit several accepted speculative tokens, then finish — a stop token or /// the max-output budget was hit partway through the accepted span. - EmitManyAndFinish { + FinishMany { request_id: RequestId, tokens: Vec, finish_reason: FinishReason, + stop_cause: Option, }, } diff --git a/pegainfer-qwen3/src/scheduler/plan.rs b/pegainfer-qwen3/src/scheduler/plan.rs index c8ca7d32d..414eee610 100644 --- a/pegainfer-qwen3/src/scheduler/plan.rs +++ b/pegainfer-qwen3/src/scheduler/plan.rs @@ -119,9 +119,11 @@ pub(crate) fn execute_plan( requests: &draft_requests, })?; draft.requests.sort_by_key(|result| result.request_id); - let verify_requests = build_speculative_verify_items(active, &draft.requests); + let (verify_requests, stop_policies) = + build_speculative_verify_items(active, &draft.requests); let mut verify = executor.execute_speculative_verify(VerifyPlan { requests: &verify_requests, + stop_policies: &stop_policies, sample_seed: rand::RngExt::random(rng), })?; verify.requests.sort_by_key(|result| result.request_id); @@ -183,27 +185,35 @@ fn build_speculative_draft_items(active: &[ActiveRequestState]) -> Vec Vec { - draft_results - .iter() - .map(|draft| { - let active = active - .iter() - .find(|req| req.request_id == draft.request_id) - .expect("draft request_id must exist in active set"); - // Clamp the verify span to the request's remaining output budget so - // a long accepted run can't overshoot max_tokens. - let remaining = active.max_tokens.saturating_sub(active.generated_count); - // A continuing active request always has budget left (resolve emits - // EmitManyAndFinish the moment generated_count hits max_tokens), so - // this is a true invariant, not a runtime condition — don't crash the - // scheduler thread in release on a state we've proven unreachable. - debug_assert!(remaining > 0, "active request must have output budget"); - let mut token_ids = draft.token_ids.clone(); - token_ids.truncate(remaining); - VerifyStepItem::new(draft.request_id, token_ids, active.params) - }) - .collect() +) -> ( + Vec, + Vec, +) { + let mut requests = Vec::with_capacity(draft_results.len()); + let mut stop_policies = Vec::with_capacity(draft_results.len()); + for draft in draft_results { + let active = active + .iter() + .find(|req| req.request_id == draft.request_id) + .expect("draft request_id must exist in active set"); + // Clamp the verify span to the request's remaining output budget so + // a long accepted run can't overshoot max_tokens. + let remaining = active.max_tokens.saturating_sub(active.generated_count); + // A continuing active request always has budget left (resolve emits + // FinishMany the moment generated_count hits max_tokens), so + // this is a true invariant, not a runtime condition — don't crash the + // scheduler thread in release on a state we've proven unreachable. + debug_assert!(remaining > 0, "active request must have output budget"); + let mut token_ids = draft.token_ids.clone(); + token_ids.truncate(remaining); + requests.push(VerifyStepItem::new( + draft.request_id, + token_ids, + active.params, + )); + stop_policies.push(active.stop_policy.clone()); + } + (requests, stop_policies) } fn build_prefill_items(pending: &[PendingRequest], indices: &[usize]) -> Vec { @@ -254,6 +264,7 @@ fn sort_decode_results(results: &mut [crate::executor::DecodeRequestResult]) { #[cfg(test)] mod tests { + use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use super::*; @@ -265,6 +276,7 @@ mod tests { lora_adapter: None, prompt_tokens: vec![1, 2, 3], params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens: 8, logprobs: 0, echo: false, @@ -284,6 +296,7 @@ mod tests { max_tokens, prompt_len: 10, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), logprobs: 0, } } @@ -308,12 +321,13 @@ mod tests { token_ids: (0..16).collect(), }; - let verify = build_speculative_verify_items(&active, &[draft]); + let (verify, stop_policies) = build_speculative_verify_items(&active, &[draft]); assert_eq!(verify.len(), 1); // 32 - 24 = 8 remaining → the 16-token span truncates to 8. assert_eq!(verify[0].as_slice().len(), 8); assert_eq!(verify[0].as_slice(), (0..8).collect::>()); + assert_eq!(stop_policies, vec![StopPolicy::default()]); } // The plan selector is the whole batch-formation policy: what the scheduler diff --git a/pegainfer-qwen3/src/scheduler/resolve.rs b/pegainfer-qwen3/src/scheduler/resolve.rs index 12356210f..840ca989d 100644 --- a/pegainfer-qwen3/src/scheduler/resolve.rs +++ b/pegainfer-qwen3/src/scheduler/resolve.rs @@ -1,4 +1,6 @@ use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use super::ActiveRequestState; use super::PendingRequest; @@ -13,6 +15,14 @@ use crate::executor::ModelExecutor; use crate::executor::PrefillRequestResult; use crate::speculative::VerifyRequestResult; +fn classify_stop( + executor: &impl ModelExecutor, + policy: &StopPolicy, + token: u32, +) -> Option { + policy.classify(token, |token_id| executor.is_stop_token(token_id)) +} + pub(crate) fn resolve_step( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -44,8 +54,9 @@ pub(crate) fn resolve_step( /// Turn each request's accepted speculative span into a decode effect. A span /// commits 1..=K+1 tokens at once; we walk it in order so a stop token or the -/// max-output budget truncates exactly where it lands (the executor already -/// suppressed nothing — stop handling lives here, mirroring single-token decode). +/// max-output budget lands exactly where expected. The executor has already +/// truncated any suffix after a request-terminal token to keep speculative +/// state consistent; the resolver classifies its typed cause here. pub(crate) fn resolve_speculative_outputs( executor: &impl ModelExecutor, active: &[ActiveRequestState], @@ -60,26 +71,30 @@ pub(crate) fn resolve_speculative_outputs( .expect("speculative request_id must exist in active set"); let mut emitted = Vec::new(); let mut completion_tokens = req.generated_count; + for &token in &result.accepted_tokens { completion_tokens += 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(token); - if is_eos { - return DecodeEffect::EmitManyAndFinish { + emitted.push(token); + + if let Some(stop_cause) = classify_stop(executor, &req.stop_policy, token) { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }; } - emitted.push(token); + if completion_tokens >= req.max_tokens { - return DecodeEffect::EmitManyAndFinish { + return DecodeEffect::FinishMany { request_id: result.request_id, tokens: emitted, finish_reason: FinishReason::Length, + stop_cause: None, }; } } - DecodeEffect::EmitManyAndContinue { + DecodeEffect::ContinueMany { request_id: result.request_id, tokens: emitted, completion_tokens, @@ -126,10 +141,13 @@ fn resolve_prefill_outputs( }); } - if !req.params.ignore_eos && executor.is_stop_token(result.first_token) { - effects.pending.push(PendingEffect::Finish { + if let Some(stop_cause) = classify_stop(executor, &req.stop_policy, result.first_token) { + effects.pending.push(PendingEffect::EmitAndFinish { request_id: req.request_id, + token: result.first_token, + logprob: result.first_token_logprob, finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), }); continue; } @@ -140,6 +158,7 @@ fn resolve_prefill_outputs( token: result.first_token, logprob: result.first_token_logprob, finish_reason: FinishReason::Length, + stop_cause: None, }); continue; } @@ -154,6 +173,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 +197,27 @@ fn resolve_decode_outputs( .find(|req| req.request_id == result.request_id) .expect("decode request_id must exist in active set"); let completion_tokens = req.generated_count + 1; - let is_eos = !req.params.ignore_eos && executor.is_stop_token(result.token); + let stop_cause = classify_stop(executor, &req.stop_policy, result.token); let at_limit = completion_tokens >= req.max_tokens; - if is_eos { + + if let Some(stop_cause) = stop_cause { DecodeEffect::Finish { request_id: result.request_id, + token: result.token, + logprob: result.logprob.clone(), finish_reason: FinishReason::Stop, + stop_cause: Some(stop_cause), } } else if at_limit { - DecodeEffect::EmitAndFinish { + DecodeEffect::Finish { request_id: result.request_id, token: result.token, logprob: result.logprob.clone(), finish_reason: FinishReason::Length, + stop_cause: None, } } else { - DecodeEffect::EmitAndContinue { + DecodeEffect::Continue { request_id: result.request_id, token: result.token, logprob: result.logprob.clone(), diff --git a/pegainfer-qwen3/src/scheduler/test_support.rs b/pegainfer-qwen3/src/scheduler/test_support.rs index 16d19323a..cf752b97b 100644 --- a/pegainfer-qwen3/src/scheduler/test_support.rs +++ b/pegainfer-qwen3/src/scheduler/test_support.rs @@ -9,6 +9,7 @@ use std::time::Duration; use anyhow::Result; use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::UnloadLoraAdapterRequest; use pegainfer_frontend::sampler::SamplingParams; @@ -39,6 +40,7 @@ pub(crate) struct FakeExecutor { pub(crate) dropped: Arc>>, pub(crate) prefetch_offers: Arc>>, stop_token: Option, + emit_logprobs: bool, } impl FakeExecutor { @@ -56,6 +58,7 @@ impl FakeExecutor { dropped, prefetch_offers: Arc::new(Mutex::new(Vec::new())), stop_token: None, + emit_logprobs: false, } } @@ -64,6 +67,11 @@ impl FakeExecutor { self } + pub(crate) fn with_logprobs(mut self) -> Self { + self.emit_logprobs = true; + self + } + pub(crate) fn with_decode_failure(mut self) -> Self { self.fail_decode_once = true; self @@ -99,7 +107,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 +238,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(), }) @@ -267,6 +287,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..d67f11ed0 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::new( + if ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + Vec::new(), + ), ..active_state(id, generated_count, max_tokens) } } @@ -537,7 +549,7 @@ fn speculative_full_span_accept_continues() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndContinue { + effects::DecodeEffect::ContinueMany { request_id, tokens, completion_tokens, @@ -551,54 +563,54 @@ fn speculative_full_span_accept_continues() { "completion = prior generated + span len" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), } } #[test] -fn speculative_stop_token_midspan_finishes_and_suppresses_eos() { +fn speculative_stop_token_midspan_finishes_and_retains_the_trigger() { let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); let active = [spec_active(1, 5, 100, false)]; - // EOS lands at span position 2; tokens before it are emitted, EOS is not. + // EOS lands at span position 2; the trigger is retained and the suffix is not. let results = [spec_result(1, vec![10, 11, SPEC_EOS, 13])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, + stop_cause, .. }, ] => { - assert_eq!( - tokens, - &vec![10, 11], - "EOS itself is suppressed from emission" - ); + assert_eq!(tokens, &vec![10, 11, SPEC_EOS],); assert!(matches!(finish_reason, FinishReason::Stop)); + assert_eq!(*stop_cause, Some(StopCause::Eos(SPEC_EOS))); } - _ => panic!("expected EmitManyAndFinish(Stop)"), + _ => panic!("expected FinishMany(Stop)"), } } #[test] -fn speculative_stop_token_at_span_start_emits_nothing() { +fn speculative_stop_token_at_span_start_retains_the_token() { let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); let active = [spec_active(1, 5, 100, false)]; let results = [spec_result(1, vec![SPEC_EOS, 11, 12])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, + stop_cause, .. }, ] => { - assert!(tokens.is_empty(), "stop at position 0 emits no tokens"); + assert_eq!(tokens, &vec![SPEC_EOS]); assert!(matches!(finish_reason, FinishReason::Stop)); + assert_eq!(*stop_cause, Some(StopCause::Eos(SPEC_EOS))); } - _ => panic!("expected EmitManyAndFinish(Stop)"), + _ => panic!("expected FinishMany(Stop)"), } } @@ -611,7 +623,7 @@ fn speculative_max_tokens_truncates_midspan() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { [ - effects::DecodeEffect::EmitManyAndFinish { + effects::DecodeEffect::FinishMany { tokens, finish_reason, .. @@ -624,7 +636,7 @@ fn speculative_max_tokens_truncates_midspan() { ); assert!(matches!(finish_reason, FinishReason::Length)); } - _ => panic!("expected EmitManyAndFinish(Length)"), + _ => panic!("expected FinishMany(Length)"), } } @@ -635,14 +647,43 @@ fn speculative_ignore_eos_does_not_stop() { let results = [spec_result(1, vec![SPEC_EOS, SPEC_EOS])]; let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); match &effects[..] { - [effects::DecodeEffect::EmitManyAndContinue { tokens, .. }] => { + [effects::DecodeEffect::ContinueMany { tokens, .. }] => { assert_eq!( tokens, &vec![SPEC_EOS, SPEC_EOS], "ignore_eos passes stop tokens through" ); } - _ => panic!("expected EmitManyAndContinue"), + _ => panic!("expected ContinueMany"), + } +} + +#[test] +fn speculative_request_stop_truncates_the_span_when_eos_is_ignored() { + let exec = FakeExecutor::new(64, Arc::new(Mutex::new(Vec::new()))).with_stop_token(SPEC_EOS); + + let mut request = spec_active(1, 0, 100, true); + request.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![12]); + + let active = [request]; + let results = [spec_result(1, vec![10, 11, 12, 13])]; + + let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); + + match &effects[..] { + [ + effects::DecodeEffect::FinishMany { + tokens, + finish_reason, + stop_cause, + .. + }, + ] => { + assert_eq!(tokens, &vec![10, 11, 12]); + assert_eq!(*finish_reason, FinishReason::Stop); + assert_eq!(*stop_cause, Some(StopCause::Token(12))); + } + _ => panic!("expected request stop to finish the speculative span"), } } @@ -657,11 +698,11 @@ fn speculative_resolves_each_request_independently() { let effects = resolve::resolve_speculative_outputs(&exec, &active, &results); assert!(matches!( &effects[0], - effects::DecodeEffect::EmitManyAndContinue { request_id, .. } if *request_id == RequestId::new(1) + effects::DecodeEffect::ContinueMany { request_id, .. } if *request_id == RequestId::new(1) )); assert!(matches!( &effects[1], - effects::DecodeEffect::EmitManyAndFinish { request_id, finish_reason: FinishReason::Stop, .. } + effects::DecodeEffect::FinishMany { request_id, finish_reason: FinishReason::Stop, .. } if *request_id == RequestId::new(2) )); } diff --git a/pegainfer-qwen3/src/speculative.rs b/pegainfer-qwen3/src/speculative.rs index bea499611..0614433b4 100644 --- a/pegainfer-qwen3/src/speculative.rs +++ b/pegainfer-qwen3/src/speculative.rs @@ -24,6 +24,7 @@ //! target distribution; acceptance only decides how many ride one step. use anyhow::Result; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::sampler::SamplingParams; use crate::executor::RequestId; @@ -56,6 +57,9 @@ impl VerifyStepItem { #[derive(Clone, Copy)] pub(crate) struct VerifyPlan<'a> { pub requests: &'a [VerifyStepItem], + /// Request-local stop policies in the same order as `requests`. They remain + /// host-side and are not copied into GPU buffers. + pub stop_policies: &'a [StopPolicy], /// Engine step seed for the verify rows' sampler pass (same contract as /// decode: fresh per step; seeded rows re-mix their own request seed). pub sample_seed: u64, @@ -69,8 +73,9 @@ pub(crate) struct VerifyRequestResult { /// Tokens to commit: the accepted draft prefix followed by the target's /// posterior token at the first mismatch (or the block-end continuation /// when every draft is accepted). Always `1..=K + 1` tokens, so a verify - /// step always makes at least one token of progress. The scheduler still - /// owns stop-token suppression before client emission. + /// step always makes at least one token of progress. Before KV commit the + /// executor truncates this span after the first request-terminal token; + /// the scheduler retains ownership of typed stop-cause emission. pub accepted_tokens: Vec, } diff --git a/pegainfer-qwen3/tests/common/harness.rs b/pegainfer-qwen3/tests/common/harness.rs index 0459712ff..461c641f4 100644 --- a/pegainfer-qwen3/tests/common/harness.rs +++ b/pegainfer-qwen3/tests/common/harness.rs @@ -18,6 +18,7 @@ use std::sync::Mutex; use pegainfer_frontend::engine::Engine; use pegainfer_frontend::engine::EngineInfo; +use pegainfer_frontend::engine::EosPolicy; use pegainfer_frontend::engine::LoraClient; use pegainfer_frontend::engine::PromptEcho; use pegainfer_frontend::engine::Request; @@ -25,6 +26,7 @@ use pegainfer_frontend::engine::RequestControl; use pegainfer_frontend::engine::RequestId; use pegainfer_frontend::engine::RequestUpdate; use pegainfer_frontend::engine::SchedulerHandle; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::engine::TokenLogprob; use pegainfer_frontend::sampler::SamplingParams; @@ -36,9 +38,19 @@ pub(crate) fn request( params: SamplingParams, max_tokens: usize, ) -> Request { + let stop_policy = StopPolicy::new( + if params.ignore_eos { + EosPolicy::Ignore + } else { + EosPolicy::ModelDefault + }, + Vec::new(), + ); + Request { prompt_tokens, params, + stop_policy, max_tokens, lora_adapter: None, kv_transfer_params: None, diff --git a/pegainfer-qwen3/tests/dflash_speculative_gate.rs b/pegainfer-qwen3/tests/dflash_speculative_gate.rs index 228e0b478..5c67648a5 100644 --- a/pegainfer-qwen3/tests/dflash_speculative_gate.rs +++ b/pegainfer-qwen3/tests/dflash_speculative_gate.rs @@ -39,11 +39,18 @@ //! `PEGAINFER_TEST_MODEL_PATH` (target) and `PEGAINFER_DFLASH_TEST_MODEL_PATH` //! (drafter); skips cleanly when either is absent. +use std::collections::HashMap; +use std::collections::HashSet; +use std::collections::VecDeque; use std::path::Path; use std::path::PathBuf; use std::process::Command; use std::time::Duration; +use pegainfer_frontend::engine::EosPolicy; +use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::StopCause; +use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; use pegainfer_qwen3::DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES; @@ -620,6 +627,109 @@ fn dflash_concurrent_heterogeneous_is_lossless() { ); } +/// Production hedge regression: an explicit stop in the middle of a verify +/// span must be applied before the hedge winner is selected and committed. +/// The parent gate also checks the request-local worker trace, because the +/// final stream alone is protected by the executor's legacy safety truncation. +#[test] +fn dflash_hedged_midspan_stop_retains_trigger() { + common::harness::init_capture_logging(); + let (Some(model_path), Some(draft_path)) = (target_path_or_skip(), draft_path_or_skip()) else { + return; + }; + if std::env::var_os("PEGAINFER_SPEC_HEDGE").is_none() { + eprintln!( + "skipping hedged mid-span stop gate: run it through hedged_ladder_passes_the_lossless_gates" + ); + return; + } + let _gpu = GPU + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + + let prompt = "Write a short paragraph about a blue"; + let tokenizer = common::load_tokenizer(&model_path); + let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); + let draft_config: serde_json::Value = serde_json::from_str( + &std::fs::read_to_string(Path::new(&draft_path).join("config.json")) + .expect("read draft config"), + ) + .expect("parse draft config"); + let block_size = draft_config["block_size"] + .as_u64() + .expect("draft block_size") as usize; + let min_index = (block_size / 2).max(1); + + let engine = EngineHarness::new( + pegainfer_qwen3::launch( + Path::new(&model_path), + launch_options(Some(PathBuf::from(&draft_path))), + ) + .expect("failed to start speculative engine"), + ); + let baseline_params = SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }; + let baseline = engine + .submit(request( + prompt_tokens.clone(), + baseline_params, + GENERATED_TOKENS, + )) + .expect_finished() + .tokens; + let candidate_stops: Vec<(usize, u32)> = baseline + .iter() + .enumerate() + .skip(min_index) + .take(block_size.saturating_sub(min_index)) + .filter(|(index, token)| !baseline[..*index].contains(token)) + .map(|(index, token)| (index, *token)) + .collect(); + assert!( + !candidate_stops.is_empty(), + "baseline did not produce a unique token inside the first verify span" + ); + + let mut stopped_cases = 0usize; + for (baseline_index, stop_id) in candidate_stops { + let stopped_params = SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }; + let mut stopped = request(prompt_tokens.clone(), stopped_params, GENERATED_TOKENS); + stopped.stop_policy = StopPolicy::new(EosPolicy::Ignore, vec![stop_id]); + let stream = engine.submit(stopped); + let request_id = stream.id(); + eprintln!("hedge stop request={request_id}"); + let outcome = stream.expect_finished(); + + if outcome.tokens.last() != Some(&stop_id) { + continue; + } + assert!(!outcome.tokens[..outcome.tokens.len() - 1].contains(&stop_id)); + assert!(matches!( + outcome.terminal, + Terminal::Finished { + reason: FinishReason::Stop, + stop_cause: Some(StopCause::Token(id)), + completion_tokens, + .. + } if id == stop_id && completion_tokens == outcome.tokens.len() + )); + eprintln!( + "hedge stop candidate baseline_index={baseline_index} token={stop_id} retained_len={}", + outcome.tokens.len() + ); + stopped_cases += 1; + } + assert!( + stopped_cases > 0, + "none of the candidate tokens produced a mid-span explicit stop" + ); +} + /// P2 regression: a request that fits the target context window but lands in the /// draft's `block_size` in-fill headroom (`max_pos - block_size < prompt + /// max_tokens <= max_pos`) must be rejected cleanly at admission. Before the @@ -700,15 +810,15 @@ fn dflash_request_in_draft_headroom_is_rejected_not_panicked() { } } -/// Execution gate for the hedge ladder: re-runs the two losslessness tests +/// Execution gate for the hedge ladder: re-runs the hedge child suites /// above in child processes, since the hedge config is read once per process /// and cannot be toggled in-process. /// -/// Scope: this proves the copy-back and discard branches were ENTERED, not -/// that a later round consumes the winner's state correctly — `check_lossless` -/// returns at the first benign tie flip, so the suffix past it is never -/// compared. A tie is also a non-win (the win test is strictly-greater), so -/// the counters cannot separate a tie from a shorter chain. +/// The stop child additionally checks that an untrimmed B candidate is not +/// selected over the truncated A candidate, and that context append and KV +/// commit use the same retained length. The losslessness children retain their +/// numerical tie tolerance and only require that the configured hedge path +/// actually ran. /// /// Strict token equality against an unhedged run is NOT a valid contract: /// hedged rounds change the verify batch shape, which legally flips bf16 ties, @@ -750,16 +860,17 @@ fn hedged_ladder_passes_the_lossless_gates() { .lock() .unwrap_or_else(std::sync::PoisonError::into_inner); let exe = std::env::current_exe().expect("test binary path"); - // One child per lossless suite, each with a single exact filter, so the - // invocation shape is beyond dispute on any libtest version. + // One child per suite, each with a single exact filter, so the invocation + // shape is beyond dispute on any libtest version. // A hedge-free child would make its lossless pass vacuous, so each child // must show expanded spans in the executor's per-round trace. let mut total_spans = 0usize; - let mut total_wins = 0usize; let mut total_rounds = 0usize; + let mut total_wins = 0usize; for child_test in [ "dflash_speculative_greedy_matches_plain_greedy", "dflash_concurrent_heterogeneous_is_lossless", + "dflash_hedged_midspan_stop_retains_trigger", ] { let output = Command::new(&exe) .args(["--exact", child_test, "--test-threads=1", "--nocapture"]) @@ -777,7 +888,6 @@ fn hedged_ladder_passes_the_lossless_gates() { ); let mut rounds = 0usize; let mut spans = 0usize; - let mut wins = 0usize; for line in child_stderr.lines() { let Some(rest) = line.split("DFlash hedge: ").nth(1) else { continue; @@ -787,23 +897,110 @@ fn hedged_ladder_passes_the_lossless_gates() { .filter(|tok| !tok.is_empty()) .map(|tok| tok.parse::().expect("hedge trace number")); spans += nums.next().expect("span count"); - wins += nums.next().expect("win count"); + let wins = nums.next().expect("win count"); + if child_test != "dflash_hedged_midspan_stop_retains_trigger" { + total_wins += wins; + } rounds += 1; } assert!( rounds > 0 && spans > 0, "child '{child_test}' executed no hedged verify round:\n{child_stderr}" ); - total_rounds += rounds; - total_spans += spans; - total_wins += wins; + if child_test == "dflash_hedged_midspan_stop_retains_trigger" { + let stop_requests: HashSet = child_stderr + .lines() + .filter_map(|line| line.strip_prefix("hedge stop request=")) + .map(str::to_owned) + .collect(); + assert!( + !stop_requests.is_empty(), + "stop child emitted no request marker" + ); + let mut context_by_request: HashMap> = HashMap::new(); + let mut commit_by_request: HashMap> = HashMap::new(); + for line in child_stderr.lines() { + let fields: Vec<_> = line.split_whitespace().collect(); + let request = fields + .iter() + .find_map(|field| field.strip_prefix("request=")); + if let Some(request) = request { + if let Some(appended) = fields + .iter() + .find_map(|field| field.strip_prefix("appended=")) + .and_then(|value| value.parse::().ok()) + { + context_by_request + .entry(request.to_string()) + .or_default() + .push_back(appended); + } + if let Some(accepted_len) = fields + .iter() + .find_map(|field| field.strip_prefix("accepted_len=")) + .and_then(|value| value.parse::().ok()) + { + commit_by_request + .entry(request.to_string()) + .or_default() + .push_back(accepted_len); + } + } + } + let mut worker_side_truncation = false; + for line in child_stderr.lines() { + if !line.contains("Qwen3 DFlash hedge detail ") { + continue; + } + let value = |name: &str| { + line.split_whitespace() + .find_map(|field| field.strip_prefix(name)) + }; + let selected = value("selected="); + let request = value("request="); + let raw_a = value("raw_a=").and_then(|v| v.parse().ok()); + let raw_b_lens = value("raw_b_lens=").map(|v| { + v.split(',') + .filter(|value| !value.is_empty()) + .map(|value| value.parse::().expect("raw B length")) + .collect::>() + }); + let selected_len = value("selected_len=").and_then(|v| v.parse().ok()); + if !request.is_some_and(|id| stop_requests.contains(id)) { + continue; + } + let context_len = request + .and_then(|id| context_by_request.get_mut(id)) + .and_then(VecDeque::pop_front); + let commit_len = request + .and_then(|id| commit_by_request.get_mut(id)) + .and_then(VecDeque::pop_front); + let raw_b_max = raw_b_lens + .as_ref() + .and_then(|lengths| lengths.iter().copied().max()); + if raw_b_max.is_some_and(|raw_b| raw_a.is_some_and(|raw_a| raw_b > raw_a)) + && selected == Some("A") + && selected_len + .is_some_and(|selected| raw_b_max.is_some_and(|raw| selected < raw)) + && selected_len.is_some() + && selected_len == context_len + && selected_len == commit_len + { + worker_side_truncation = true; + break; + } + } + assert!( + worker_side_truncation, + "stop hedge never truncated a candidate before winner/context commit:\n{child_stderr}" + ); + } + if child_test != "dflash_hedged_midspan_stop_retains_trigger" { + total_rounds += rounds; + total_spans += spans; + } } - assert!( - total_wins > 0, - "no hedge chain ever won across {total_rounds} hedged rounds" - ); - assert!( - total_spans > total_wins, - "every hedge span won ({total_wins}/{total_spans}) — the discard path never executed" - ); + assert!(total_rounds > 0 && total_spans > 0); + assert!(total_wins > 0); + assert!(total_spans > total_wins); } diff --git a/pegainfer-sim/src/lib.rs b/pegainfer-sim/src/lib.rs index 1e2f52ebe..d99d0908f 100644 --- a/pegainfer-sim/src/lib.rs +++ b/pegainfer-sim/src/lib.rs @@ -311,6 +311,7 @@ fn duration_from_ms(ms: f64) -> Duration { #[cfg(test)] mod tests { use pegainfer_frontend::engine::Request; + use pegainfer_frontend::engine::StopPolicy; use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; @@ -320,6 +321,7 @@ mod tests { Request { prompt_tokens, params: SamplingParams::default(), + stop_policy: StopPolicy::default(), max_tokens, lora_adapter: None, kv_transfer_params: None, @@ -425,6 +427,7 @@ mod tests { reason: FinishReason::Length, prompt_tokens: 2, completion_tokens: 3, + .. } )); }