From 83c128e09fcad2281aed2f5214f0e79338745b1d Mon Sep 17 00:00:00 2001 From: xiaguan <751080330@qq.com> Date: Fri, 28 Aug 2026 04:39:58 -0400 Subject: [PATCH] feat(qwen35): launch only on the step contract Cut Qwen3.5 off EngineHandle/TokenEvent so ModelLine::launch returns LaunchedEngine::Stepped and Qwen35Scheduler owns submit/step/metrics. Co-authored-by: Cursor --- docs/index.md | 7 +- docs/models/qwen35/load-snapshot.md | 20 +- docs/models/qwen35/model-crate.md | 4 +- docs/models/qwen35/step-contract-migration.md | 143 + .../frontend/frontend-architecture.md | 6 +- .../subsystems/frontend/prometheus-metrics.md | 4 +- pegainfer-qwen35/src/lib.rs | 12 +- pegainfer-qwen35/src/model_line.rs | 7 +- pegainfer-qwen35/src/scheduler.rs | 2624 ----------------- pegainfer-qwen35/src/scheduler/backend.rs | 1152 ++++++++ pegainfer-qwen35/src/scheduler/mod.rs | 1135 +++++++ pegainfer-qwen35/src/scheduler/tests.rs | 1085 ++----- pegainfer-qwen35/tests/chunked_prefill.rs | 85 +- pegainfer-qwen35/tests/common/harness.rs | 271 ++ pegainfer-qwen35/tests/common/mod.rs | 2 + pegainfer-qwen35/tests/e2e_scheduler.rs | 569 ++-- pegainfer-qwen35/tests/sampling_behavior.rs | 81 +- pegainfer-qwen35/tests/serving_tp2.rs | 201 +- 18 files changed, 3437 insertions(+), 3971 deletions(-) create mode 100644 docs/models/qwen35/step-contract-migration.md delete mode 100644 pegainfer-qwen35/src/scheduler.rs create mode 100644 pegainfer-qwen35/src/scheduler/backend.rs create mode 100644 pegainfer-qwen35/src/scheduler/mod.rs create mode 100644 pegainfer-qwen35/tests/common/harness.rs diff --git a/docs/index.md b/docs/index.md index 63a5950b1..08dfc6859 100644 --- a/docs/index.md +++ b/docs/index.md @@ -46,18 +46,19 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | | `models/qwen35/roadmap.md` | Qwen3.5 dense roadmap v2 (#654): core correctness/admission/chunked-prefill/sampling/step-tail gates are landed; current 4B HTTP boundary is the retained #469 RTX 5090 sweep, which completed with zero failed requests but trails vLLM at high concurrency. Next: HTTP gap attribution, mixed-load ITL (#470), lifecycle recovery (#471), joint-state prefix reuse (#257), and design-first TP (#446). | -| `models/qwen35/load-snapshot.md` | Qwen3.5 publishes logical running, current pending, and KV load after drain/cancellation pruning and before admission through the shared single-GPU/TP scheduler backend. | +| `models/qwen35/load-snapshot.md` | Qwen3.5 publishes logical running, waiting, and KV load after each `step` (same cadence as Qwen3) from the shared single-GPU/TP backend. | | `models/qwen35/prefix-cache.md` | Qwen3.5 prefix-cache design: a hit is valid only when full-attention KV and a complete recurrent/conv snapshot exist at the same 256-token boundary; the first version uses a fixed-budget GPU snapshot pool with joint lookup, pinning, and LRU eviction. | | `models/qwen35/kv-admission.md` | Issue #254 complete: Qwen3.5 now uses full-lifetime KV admission, deferred pressure handling, impossible-request rejection, explicit error semantics, direct rejection-event coverage, RTX 5090 e2e, and real HTTP pressure/post-pressure validation. | | `models/qwen35/optimization.md` | Hybrid 24 linear + 8 full attn optimization ledger. Decode-tuning refresh fuses MLP gate/up and tunes decode cublasLt buckets, improving direct TPOT by 2-3%; vLLM still leads 1024/256 HTTP decode. | | `models/qwen35/accuracy.md` | Qwen3.5 HF bf16 logits goldens, size-keyed (0.8b/2b/4b/9b/27b all committed), through `past_key_values`: short replay covers sequential graph, bucket-straddling batched graph, and slot-compaction; long replay covers 4097/8192-token prompts; full GSM8K 8-shot now matches the HF baseline within 0.15 percentage points. | -| `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (Triton AOT is the only Python build dependency); root loads it through `EngineHandle`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | +| `models/qwen35/model-crate.md` | `pegainfer-qwen35` owns Qwen3.5 model/scheduler/recurrent ops/tests/benches; feature-gated behind `qwen35` (Triton AOT is the only Python build dependency); root loads it through `Engine` / `LaunchedEngine::Stepped`. Build/check/clippy, root bench sanity check, historical Qwen3.5 e2e, and scheduler e2e records live here. | | `models/qwen35/batched-step-tail.md` | Qwen3.5 issue #353 implementation record: final prefill tail is batched, decode/unified sample from batched logits, host full-vocab copies are logprobs-only, HF + scheduler e2e pass, and final serving A/B supports only the first-token/short-output TTFT claim. | | `models/qwen35/tp-design.md` | Qwen3.5 TP design: Phase 1 is eager dense TP on Qwen3's controller/worker runtime; validate TP2 first, fail closed for indivisible degrees and TP+CUDA Graph, shard dense full-attention/MLP, and leave sharded linear/GDR state to follow-up. | | `models/qwen35/tp-implementation.md` | Qwen3.5 TP Phase 1 and P2A are complete: TP2 has start-gated eager unified prefill+decode, strict ID-aligned artifacts, fail-closed lifecycle recovery, and pre-load ordinal validation; P2B GDR state sharding is next. | | `models/qwen35/mixed-load-itl-470.md` | Issue #470: full cold `--max-batch 8/bg=4` matrix on RTX 4090 (24/24 valid) + starvation negative control. Qwen3.5 is not immune; chunking bounds max/per-step stall but raises p99 at low QPS (~14→~80–92ms) and pulls p99/max back from the prefill wall to the chunk wall at high load; `qps·prefill_s≳1` is a throughput wall (chunking can't fix it, and ON's +15% TTFT can trip it earlier). The old "p99 immunity" was a slot-starvation artifact. | | `models/qwen35/adaptive-scheduler-policy.md` | Issue #727 adaptive scheduler policy record: default `off`, opt-in `auto`, hard `--max-prefill-tokens` cap, TP `auto` rejection, and pre-review whole-prefill benchmark tradeoff retained as non-default evidence. | | `models/qwen35/unified-prefill-overlap.md` | Issue #715 implementation record: opt-in single-GPU shared-SM overlap keeps one prefill chunk in flight while active decode continues; default serial policy and unsupported-combination guards remain explicit. | +| `models/qwen35/step-contract-migration.md` | Qwen3.5 launches only as `LaunchedEngine::Stepped`; `Qwen35Scheduler` implements the step contract in `scheduler/`. Legacy `EngineHandle`/`TokenEvent` is gone from this crate. | ## models / gemma4 @@ -170,7 +171,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | Path | TL;DR | | --- | --- | -| `subsystems/frontend/frontend-architecture.md` | `pegainfer-frontend` owns everything north of the model schedulers. Two contract generations coexist: the step contract (qwen3 + pegainfer-sim migrated) and the legacy `EngineHandle`/`TokenEvent` path (other five lines). Next: migrate glm52, then delete the legacy contract. | +| `subsystems/frontend/frontend-architecture.md` | `pegainfer-frontend` owns everything north of the model schedulers. Two contract generations coexist: the step contract (qwen3 + qwen35 + pegainfer-sim migrated) and the legacy `EngineHandle`/`TokenEvent` path (glm52/kimi/dsv2/gemma4). Next: migrate glm52, then delete the legacy contract. | | `subsystems/frontend/simulated-inference-engine.md` | CPU-only simulated model crate on the step contract (`SimScheduler` → `LaunchedEngine::Stepped`) for vLLM/OpenAI frontend and `vllm bench serve` validation without CUDA or weights. | | `subsystems/frontend/sim-step-contract.md` | Cut `pegainfer-sim` from the legacy `EngineHandle`/`TokenEvent` path onto the step contract. | | `subsystems/frontend/sim-high-concurrency-bench.md` | Same-session A/B vs main: feat TPOT ~30–180× better, TTFT worse and linear in C; E2EL/throughput win at c=64 and c=1024. | diff --git a/docs/models/qwen35/load-snapshot.md b/docs/models/qwen35/load-snapshot.md index dcbad0ab7..c1abcfca3 100644 --- a/docs/models/qwen35/load-snapshot.md +++ b/docs/models/qwen35/load-snapshot.md @@ -1,6 +1,6 @@ # Qwen3.5 Scheduler LoadSnapshot -> **TL;DR:** Qwen3.5 publishes one logical post-drain/post-prune `LoadSnapshot` stream from its shared single-GPU/TP scheduler: running counts active and prefilling requests, waiting counts all current pending work, and KV usage is request-page capacity minus available pages. +> **TL;DR:** Qwen3.5 publishes one logical `SchedulerMetrics` snapshot from its shared single-GPU/TP backend after each `step` (same cadence as Qwen3): running counts active, prefilling, and in-flight overlap prefill; waiting counts deferred/queued work; KV usage is request-page capacity minus available pages. > > **Last touched:** 2026-08 @@ -25,28 +25,28 @@ The data path reuses the existing frontend contract: ```text -Qwen3.5 SchedulerBackend - -> LoadSnapshot watch - -> EngineHandle - -> LocalEngineBridge +Qwen3.5 Qwen35Backend + -> Scheduler::metrics after prune (inside step) + -> driver publishes once per iteration + -> SchedulerHandle::metrics / SteppedEngineBridge -> SchedulerStats -> /metrics ``` -Both Qwen3.5 execution modes own one logical request stream, so single-GPU and TP each attach one `EngineHandle::with_load_watch` receiver. The frontend bridge, metric names, labels, and scheduler-stat conversion remain unchanged. +Both Qwen3.5 execution modes own one logical request stream, so single-GPU and TP each expose one scheduler. The frontend bridge, metric names, labels, and scheduler-stat conversion remain unchanged. -Each scheduler tick first merges deferred work with every submission currently available, then prunes closed pending, active, and prefilling requests before publishing. The fixed boundary is `drain -> prune -> publish load -> admission -> plan`. If the idle scheduler wakes through `blocking_recv()`, it drains, prunes, and publishes again before admission so work closed before admission never consumes a slot or appears in the snapshot. +The driver publishes *after* `step()` returns, same as Qwen3. `step` prunes aborted work before admission, then admits and executes; `metrics()` reads the queues at the end of that step. In-flight overlap prefill counts as running so an overlap wait inside `step` is never published as idle. Snapshot accounting is: | Metric field | Existing Qwen3.5 state | | --- | --- | -| `num_running_reqs` | `active.len() + prefilling.len()` | -| `num_waiting_reqs` | the merged pending queue: prior deferred work plus newly drained submissions | +| `num_running_reqs` | `active.len() + prefilling.len() + inflight_prefill` | +| `num_waiting_reqs` | deferred/queued work not yet admitted | | `kv_used_blocks` | request KV capacity minus currently available request pages | | `kv_total_blocks` | backend request KV capacity, excluding the CUDA Graph padding page | -Publication reads the scheduler's queues and KV allocator after closed resident state has gone through its normal retirement path. The snapshot therefore describes the state used by the following admission decision: cancelled residents no longer count as running or hold capacity, while live pending requests count as waiting even if they were submitted during the current tick. +Publication reads the scheduler's queues and KV allocator after aborted resident state has gone through its normal retirement path. Live pending requests count as waiting even if they were submitted during the current driver drain. The live gate uses `scripts/bench_http_serving.py` to create overlapping HTTP traffic and a 100 ms `curl /metrics` sampler to retain the three labeled gauges. diff --git a/docs/models/qwen35/model-crate.md b/docs/models/qwen35/model-crate.md index 6d8d4d6a7..f2190d496 100644 --- a/docs/models/qwen35/model-crate.md +++ b/docs/models/qwen35/model-crate.md @@ -1,8 +1,8 @@ # Qwen3.5-4B Model Crate **Created**: 2026-05-05 -**TL;DR**: `pegainfer-qwen35` now owns Qwen3.5 config, weights, prefill/decode/unified forward, recurrent state, scheduler, recurrent op wrappers, scheduler integration tests, and Qwen3.5 op benches. The whole crate is behind the `qwen35` feature (`--features qwen35` on `pegainfer-server`) because its GDR prefill kernels are Triton AOT-generated — this keeps the default Qwen3 build Python-free. Root `pegainfer` loads Qwen3.5 through `pegainfer_qwen35::start_engine(...)` / generic `EngineHandle`; root no longer exposes `pegainfer::model::Qwen35Model` or `pegainfer::scheduler_qwen35`. The original exact-text e2e/regen tests described in this migration record were later retired by the HF logits gate in `docs/models/qwen35/accuracy.md`. -**Last touched**: 2026-07 +**TL;DR**: `pegainfer-qwen35` now owns Qwen3.5 config, weights, prefill/decode/unified forward, recurrent state, scheduler, recurrent op wrappers, scheduler integration tests, and Qwen3.5 op benches. The whole crate is behind the `qwen35` feature (`--features qwen35` on `pegainfer-server`) because its GDR prefill kernels are Triton AOT-generated — this keeps the default Qwen3 build Python-free. Root `pegainfer` loads Qwen3.5 through `pegainfer_qwen35::start_engine(...)` / `Engine` (`LaunchedEngine::Stepped`); root no longer exposes `pegainfer::model::Qwen35Model` or `pegainfer::scheduler_qwen35`. The original exact-text e2e/regen tests described in this migration record were later retired by the HF logits gate in `docs/models/qwen35/accuracy.md`. +**Last touched**: 2026-08 ## Feature gate (2026-06) diff --git a/docs/models/qwen35/step-contract-migration.md b/docs/models/qwen35/step-contract-migration.md new file mode 100644 index 000000000..d0c1755f3 --- /dev/null +++ b/docs/models/qwen35/step-contract-migration.md @@ -0,0 +1,143 @@ +# Qwen3.5 onto the step contract + +> **TL;DR:** `pegainfer-qwen35` launches only as `LaunchedEngine::Stepped`. `Qwen35Scheduler` implements the step contract in `scheduler/`; the legacy `EngineHandle` / `TokenEvent` path is gone. GPU gates passed on RTX 5070 Ti against `/data/models/Qwen3.5-4B`. +> +> **Last touched:** 2026-08 + +## Preparation + +- **Read**: + - `docs/index.md` — Qwen3.5 is still listed as a legacy-handle line; frontend architecture says glm52 is next, then qwen35. User asked for qwen35 now, and to drop the old path. + - `docs/subsystems/frontend/frontend-architecture.md` — onboarding checklist: implement `Scheduler`, `spawn_scheduler`, return `Stepped`. Echo-server / K3 `scheduler/mod.rs` / Qwen3 `frontend_adapter.rs` are the references. Legacy modules stay in the frontend crate until every line migrates. + - `docs/subsystems/frontend/sim-step-contract.md` — sim already cut over with no `EngineHandle` leftover; tests drive `StepOutputs` / `Terminal`. + - `docs/conventions/migration-defense.md` — every old defensive structure needs an heir (inherit / replace / argue impossible). + - `docs/lessons/exact-match-gate-thread-cublas.md` — scheduler thread must rebind CUDA context and init thread-local cuBLAS. Today's `bind_model_thread` must land on the contract driver thread, not the load thread. + - `docs/models/qwen35/load-snapshot.md` — today's snapshot is drain → prune → publish → admission. The step driver publishes *after* `step()`, same as Qwen3 already does. + - `docs/models/qwen35/model-crate.md` — crate still documents `start_engine` / `EngineHandle` as the root-facing API. + - `docs/models/qwen35/unified-prefill-overlap.md` — Shared-SM overlap: at most one inflight prefill; when decode retires first, wait on the CUDA event instead of parking on submit. + - `pegainfer-k3/src/model_line.rs` + `pegainfer-k3/src/scheduler/mod.rs` — target shape: `submit` parks, `step` writes the ledger, `start_with_executors` → `Engine`, `launch` → `Stepped`. Fake-executor contract tests in `scheduler/tests.rs`. + - `pegainfer-qwen3/src/frontend_adapter.rs` + `tests/common/harness.rs` — closest *behavioral* reference for chunked prefill / overlap / abort-as-flag. GPU tests use a step-stream harness, not `TokenSink`. We are **not** copying Qwen3's extra `frontend_adapter.rs` layer. + - `pegainfer-qwen35/src/model_line.rs`, `lib.rs`, `scheduler.rs` — `launch` currently maps to `LaunchedEngine::Handle`. `scheduler_loop` owns submit drain, TokenSink dispatch, idle `blocking_recv`, and a watch metrics publisher. +- **Relevant history**: + - Qwen3 and `pegainfer-sim` already migrated; glm52/qwen35/kimi/dsv2/gemma4 still launch through `Handle`. + - Qwen3 kept `scheduler.rs` as contract-free mechanics and put the ledger writes in `frontend_adapter.rs`. K3 folded both into `scheduler/`. Qwen3.5 already has `scheduler/plan.rs` as the mechanics split, so K3's layout is the closer fit. + - #830-style risk: `prune_closed_requests` (channel-close = cancel), `terminal_scheduler_shutdown` (fail every in-flight request), `bind_model_thread`, and the overlap "don't park while inflight" rule are the defenses that must not vanish. +- **Plan**: + 1. Reshape `pegainfer-qwen35` like K3: + - `scheduler/mod.rs` implements `Scheduler` (`submit` / `step` / `metrics`) and `start_with_backend` → `Engine`. + - Keep `scheduler/plan.rs` as admission / chunk / KV-budget mechanics. + - Split the current 2.6k-line `scheduler.rs` into `scheduler/{mod,backend}.rs` (single-GPU + TP backends stay behind one enum; GPU execute/overlap stay south of the ledger). + - `model_line::launch` returns `LaunchedEngine::Stepped` only. + 2. Delete the old path from this crate — no `EngineHandle`, no `TokenSink` / `TokenEvent` / `GenerateRequest`, no `start_engine*`, no private `scheduler_loop` / submit channel / watch publisher. Tests and `runtime` re-exports follow the new `Engine`. + 3. Wire replacements (defense table in execution): + - Token dispatch → `RequestLedger::push_tokens` / `finish` / `fail` / `reject` / `retire`. + - Channel-close cancel → `ledger.is_aborted` on every touch (same as K3 `finish_or_retire`). + - Idle `blocking_recv` → contract driver spin (overlap wait on the CUDA event stays *inside* `step`, matching today's "don't park while inflight"). + - Load watch → `Scheduler::metrics` after prune; driver publishes once per iteration. + - Fatal shutdown → `step` returns `Err`; driver `fail_all`s open accounts (drop-bomb covers the rest). + - `bind_model_thread` + cuBLAS init on first `step` of the driver thread. + 4. Rewrite tests onto the contract: + - Protocol: fake executor in `scheduler/tests.rs` (K3 pattern) — abort, reject, stop/length, overlap idle, fatal fail-all. + - GPU: copy Qwen3's `EngineHarness` into `tests/common/`; rewrite `e2e_scheduler`, `sampling_behavior`, `chunked_prefill`, `serving_tp2` off `EngineHandle`. + - Keep plan.rs unit tests; they do not speak the wire. + 5. Docs: this file's execution log; `frontend-architecture.md` (qwen35 is stepped); `model-crate.md` / `load-snapshot.md` TL;DRs if they would mislead. + 6. Verify: `cargo fmt` / clippy `-p pegainfer-qwen35 --features qwen35 --all-targets -- -D warnings`; lib + scheduler unit tests; GPU gates `e2e_scheduler` / `sampling_behavior` / `chunked_prefill` when a card is available. +- **Risks / open questions**: + - Metrics cadence moves from pre-admission to post-`step` (Qwen3 already lives with this). Load-snapshot tests must assert the new boundary, not the old one. + - Shared-SM overlap currently parks on a CUDA event when decode empties first. That wait belongs inside `step`, not as a driver idle. Easy to get wrong and stall admission. + - cuBLAS thread-local handles: skipping `bind_model_thread` on the driver thread reproduces the historical gibberish bug. + - Frontend crate still carries the legacy contract for glm52/kimi/dsv2/gemma4. This task does not delete those modules. + +## Execution Log + +### GPU tests `sampling_behavior` + `chunked_prefill` onto EngineHarness +- Copied `pegainfer-qwen3/tests/common/harness.rs` into `pegainfer-qwen35/tests/common/harness.rs` (identical). `tests/common/mod.rs` already exports `harness`. +- Rewrote `tests/sampling_behavior.rs` and `tests/chunked_prefill.rs` onto `EngineHarness` / `Request` / `Terminal`. Deleted `EngineHandle`, `GenerateRequest`, `TokenSink`, `TokenEvent`, `TokenStreamReceiver`. Sampling params stay on `Request.params`. `ignore_eos` still forces a length finish, so `outcome.tokens` is exactly `max_tokens` (no stop/EOS token to drop). +- Coverage kept: greedy determinism, `top_k=1` / tiny `top_p` collapse, hot-temp diversity, chunked vs unchunked greedy token match + `FinishReason::Length`. +- Not in these two files before, so not added here: mixed-batch (lives in `e2e_scheduler.rs`, untouched), abort/cancel, `Terminal::Rejected`, timeouts. +- Result: files formatted; compile/GPU run wait on `start_engine` returning `Engine`. + +### GPU test `serving_tp2` onto EngineHarness +- Rewrote `tests/serving_tp2.rs` off `LaunchedEngine::Handle` / HTTP `vllm::serve`. Live generate now: `start_engine_with_capacity` → `EngineHarness::new` → `submit` / `expect_finished`. CUDA Graph + TP still asserts `start_engine_with_capacity(...).err()` contains `"eager execution only"` (now a non-ignored pre-load test; no weights/GPU). +- Coverage kept: greedy `ignore_eos` length stop, completion-token count, finite logprobs, two concurrent prompts, TP2 + CUDA Graph fail-closed before load. +- Dropped: OpenAI `/v1/models` + `/v1/completions` HTTP (streaming SSE `[DONE]`, usage JSON, health wait). That was this file's unique serving smoke; in-process harness does not replace the frontend wire. +- Removed unused `reqwest` / `tokio-util` from `pegainfer-qwen35` dev-dependencies (HTTP client was only used here). `lib.rs` launch-validation tests already match `.err()` and do not name `EngineHandle`. +- Result: files formatted; compile waits on `start_engine_with_capacity` returning `Engine`. + +### HTTP TP2 serving gate inherited onto Stepped +- Restored `qwen35_tp2_serves_openai_completions_over_http` from `HEAD:pegainfer-qwen35/tests/serving_tp2.rs`. The generate harness and CUDA Graph + TP pre-load reject tests stay; the HTTP gate was not abolished. +- `spawn_ready_server` still `spawn_blocking`s `start_engine_with_capacity`, then hands the engine to `pegainfer_frontend::vllm::serve(std::future::ready(Ok(LaunchedEngine::Stepped(engine))), …)` — same model path / port / `CancellationToken` args as HEAD, `Handle` swapped for `Stepped`. No `EngineHandle` / `TokenSink`. +- Assertions inherited: `/v1/models` advertises `qwen35-tp2-serving-smoke`; non-streaming `/v1/completions` (non-empty text, `finish_reason=length`, usage JSON, finite logprobs); SSE streaming including `data: [DONE]` and a choice payload; concurrent completions; `/health` wait before the first request; 30s shutdown. `#[ignore]` text and `PEGAINFER_TEST_FRONTEND_MODEL_PATH` fallback (`common::model_fixture::frontend_model_path_or_skip`) match HEAD. CUDA Graph fail-closed stays the standalone pre-load test (not re-nested in HTTP). +- Restored `reqwest` (`features = ["json"]`) and `tokio-util` in `pegainfer-qwen35` dev-dependencies, same as HEAD. +- Result: files formatted; HTTP compile waits on `start_engine_with_capacity` returning `Engine`. + +### Crate-internal scheduler unit tests onto the step contract + +Rewrote `pegainfer-qwen35/src/scheduler/tests.rs` off TokenSink / TokenEvent. No dual path in the test file. Production `scheduler.rs` still owns the legacy loop (rewritten in parallel); tests now call `echo_refusal` / `contract_reject_reason` / `prefill_drop_expectation` / `logical_load_counts` helpers added there. `plan.rs` tests were not touched. GPU tests under `pegainfer-qwen35/tests/` were not touched by this slice. + +`cargo test --release -p pegainfer-qwen35 --features qwen35 --lib -- scheduler::tests`: 7 passed, 1 ignored (`tp2_scheduler_runs_forced_mixed_steps`). `start_tp_with_capacity` still returns `EngineHandle`; the Length-terminal assertions live in `assert_forced_mixed_steps(Engine)` and become live once start returns `Engine`. + +#### Defense table (crate-internal tests) + +| Old defense | Failure mode | Heir | +| --- | --- | --- | +| `send_rejection` KvBudget TokenEvent message | client does not see lifetime KV tokens | **inherit**: `contract_reject_reason` → `RejectReason::KvBudget { worst_case_tokens }` Display (`max_request_tokens=80`) | +| `send_rejection` ContextLength TokenEvent message | client does not see window + requested length | **inherit**: `contract_reject_reason` → `RejectReason::ContextLength` Display | +| `reject_unsupported_echo` + `UNSUPPORTED_ECHO_MESSAGE` | echo request reaches backend admission | **replace**: `echo_refusal` → `RejectReason::EchoPrefillTokens { limit: 0 }`; assert contract Display, not the old TokenEvent string | +| `tp_engine_rejects_cuda_graph_before_model_load` | TP+CUDA Graph starts loading weights | **inherit**: `start_engine_with_capacity` still fail-closes before load (`eager execution only`); Ok type is irrelevant because the test asserts `Err` | +| `tp2_scheduler_runs_forced_mixed_steps` TokenSink collection | TP mixed decode+chunk-prefill does not finish at Length | **replace**: ignored GPU test launches via `start_tp_with_capacity`; step-stream assertions in `assert_forced_mixed_steps(Engine)` (K3 StepCollector). TokenSink collection abolished | +| `closed_pending_work_is_pruned_before_admission` | cancelled queued request is admitted | **replace**: cannot mint `RequestLedger` here. Heir is K3 `aborted_request_retires_silently_and_frees_its_slot` (`ledger.is_aborted` on admit). Comment in the test file | +| `closed_resident_work_is_absent_from_post_prune_load` | cancelled resident still in load snapshot | **replace**: abort prune needs a ledger (K3). Load formula heir: `logical_load_counts` still counts inflight as running (`overlap_wait_policy_is_inside_step`) | +| `closed_resident_frees_capacity_for_same_tick_admission` | cancelled decode slot not reused this tick | **replace**: same K3 abort+admit path; not unit-testable here without a ledger | +| `closed_materialized_prefill_requires_existing_worker_state` | drop missing TP worker state / drop unmaterialized state | **inherit**: `prefill_drop_expectation(cursor)` — `MustBeAbsent` at cursor 0, `MustExist` after | +| `prune_drop_failure_preserves_pending_for_terminal_fanout` | prune fatal starves pending of an error | **abolish**: `step` returns `Err`; driver `fail_all` (`pegainfer-frontend/src/engine/driver.rs` `fatal_step_fails_in_flight_requests_with_the_error`) | +| `decode_eos_waits_for_drop_before_finished` | client sees Finished before KV drop | **abolish**: one `RequestUpdate` committed after the whole step; drop always happens before the client sees the terminal | +| `decode_length_waits_for_drop_before_token_and_finished` | Token then Finished before drop | **abolish**: same structure | +| `non_tp_decode_preserves_publish_before_retire_order` | TokenEvent order vs retire | **abolish**: same structure | +| `decode_completion_drop_failure_publishes_only_terminal_error` | drop fail publishes Token then Error | **abolish**: `step` `Err` → driver `fail_all`; no TokenEvent fan-out | +| `immediate_prefill_completion_waits_for_drop` | prefill Length before drop | **abolish**: same one-update-per-step structure | +| `immediate_prefill_drop_failure_publishes_only_terminal_error` | remaining scheduled prefill not failed | **abolish**: driver `fail_all` covers every open account | +| `terminal_shutdown_closes_drains_and_errors_every_owner_once` | shutdown misses an owner / double Error / load not zeroed | **abolish**: frontend `drive()` `fail_all` + drop-bomb; already tested in `pegainfer-frontend/src/engine/driver.rs` | +| `FatalSchedulerError.transient` + TokenEvent::Error fan-out | fatal loses in-flight requests | **abolish**: the ledger holds an account for every unanswered request; `fail_all` writes them off | +| `should_block_on_submit` after last decode retires | scheduler parks on submit while inflight prefill is the only work | **replace**: wait is inside `step` (overlap_wait). Untestable without GPU; `overlap_wait_policy_is_inside_step` asserts `logical_load_counts` still counts inflight as running so the driver cannot see idle | + +Admission/chunk tests that never spoke TokenSink stay in `scheduler/plan.rs`. + +### Production cutover: `Qwen35Scheduler` on the step contract + +No dual path. `scheduler.rs` is now `scheduler/{mod,backend,plan,tests}.rs`. `Qwen35Scheduler` implements `submit` / `step` / `metrics`. `start_with_capacity` / `start_with_capacity_and_policy` / `start_tp_with_capacity` return `Engine` via `spawn_scheduler` (`qwen35-scheduler` / `qwen35-scheduler-tp`). `model_line::launch` returns `LaunchedEngine::Stepped` only. `start_engine*` keep their names but return `Engine`. GPU tests and HTTP TP2 were already on `Engine` / `Stepped`; they compile against this cutover. `assert_forced_mixed_steps` is now called from the ignored TP2 lib test. + +`bind_model_thread` + `CublasThreadGuard` run on first `step` of the driver thread (single-GPU only; stored on the scheduler). TP workers still bind themselves. + +#### Defense table (loop / TokenSink / bind / overlap / metrics / shutdown) + +| Old defense | Failure mode | Heir | +| --- | --- | --- | +| Own `scheduler_loop` + tokio submit + idle `blocking_recv` | second loop disagrees with the driver; parks while GPU work is in flight | **replace**: contract `drive()` drains submit, calls `step`, publishes metrics, commits. Idle is `spin_loop` in the driver | +| `TokenSink` / `TokenEvent` dispatch | per-request send, ordering by convention, send-fail = cancel | **replace**: `ledger.admit` / `reject` / `push_tokens` / `finish` / `fail` / `retire`. Stop token is not pushed (bridge appends EOS for usage) | +| `token_tx.is_closed` prune (`prune_closed_requests`) | cancelled work is admitted or stays in load | **replace**: `ledger.is_aborted` on every touch; `finish_or_retire` on every finish path | +| `should_block_on_submit` (`owned_work_empty && !inflight`) | parks on submit while inflight prefill is the only work | **replace**: if inflight and active empty, `step` waits on the CUDA event (`overlap_wait`); if active nonempty, decode while polling. Never park on submit. Inflight counts as running so the driver cannot see idle | +| Watch `LoadSnapshot` publish before admission | cancelled residents appear as running; waiting misses same-tick submits | **replace**: `Scheduler::metrics` after `step` (same cadence as Qwen3). Inflight in running; deferred/queued in waiting. Driver publishes once per iteration | +| `FatalSchedulerError` + TokenEvent Error fan-out + `terminal_scheduler_shutdown` | fatal misses an in-flight owner / double Error | **abolish**: `step` `Err` → driver `fail_all` + drop-bomb. Single-GPU execute failures `ledger.fail` the touched requests and keep serving (`Ok`) | +| `completion_requires_drop_ack` TokenEvent-vs-drop order | client sees Finished before TP KV drop | **abolish**: drop backend state, then `ledger.finish`; one `RequestUpdate` committed after the whole step — structurally drop-before-visible-terminal | +| `bind_model_thread` before `scheduler_loop` on a private thread | first `step` of `spawn_scheduler` runs without CUDA/cuBLAS binding → gibberish logits | **inherit**: first `step` of the driver thread calls `bind_model_thread`; `CublasThreadGuard` lives on `Qwen35Scheduler` for the thread lifetime. Single-GPU only (TP workers bind themselves) | + +### GPU gates `e2e_scheduler` / `sampling_behavior` / `chunked_prefill` + +- Weights are not under the repo `models/` (that directory does not exist). Fixture is `/data/models/Qwen3.5-4B` (`config.json` `model_type=qwen3_5`, `text_config.model_type=qwen3_5_text`; ~9.3 GiB safetensors). Duplicate copy at `/data/openclaw-data/workspace/models/Qwen3.5-4B` not used. +- `nvidia-smi` before run: RTX 5070 Ti, 34 MiB / 16303 MiB, 0% util, no processes. One GPU, so TP2 ignored. +- Command: + +``` +PEGAINFER_TEST_MODEL_PATH=/data/models/Qwen3.5-4B cargo test --release -p pegainfer-qwen35 --features qwen35 --test e2e_scheduler --test sampling_behavior --test chunked_prefill -- --test-threads=1 +``` + +- First run: `chunked_prefill` passed; `test_e2e_qwen35_scheduler` OOM'd at load (`from_safetensors_with_options` sizes recurrent state to MAX_BATCH=64: 6288 MB + 3248 MB scratch vs 7227 MB free). Pre-existing vs HEAD, but the test comment already wanted 8 slots for 16GB. Aligned TP1 onto `start_engine_with_capacity(..., 8)` like TP2 in the same file. Not a bind_model_thread / EngineHandle leftover. +- Second run: **4 passed, 0 failed, 1 ignored** (`test_e2e_qwen35_scheduler_tp2`). `chunked_prefill` 1/1, `e2e_scheduler` 2/2 (+1 ignored), `sampling_behavior` 1/1. No gibberish / cuBLAS first-step failure. + +## Debrief + +- **Outcome**: `pegainfer-qwen35` production launch is on the step contract only. `Qwen35Scheduler` lives in `scheduler/`, `launch` returns `LaunchedEngine::Stepped`, `start_engine*` return `Engine`. Lib tests: 90 passed, 7 ignored, 0 failed. Clippy `-D warnings --all-targets` passed. GPU gates on RTX 5070 Ti + `/data/models/Qwen3.5-4B`: 4 passed, 0 failed, 1 ignored (TP2). +- **Pitfalls encountered**: `bind_model_thread` used to run before the private loop and `start_*` waited on it. The contract driver thread is a different spawn, so binding had to move to first `step` or the historical cuBLAS gibberish bug returns. Overlap wait must stay *inside* `step`; returning idle to the driver while inflight would spin-loop instead of waiting on the CUDA event. `test_e2e_qwen35_scheduler` still loaded via `from_safetensors_with_options` (64 slots) then `start_with_capacity(8)` — that OOM'd on 16GB before any step ran; the 8-slot comment only applied after load. +- **Lessons learned**: K3's crate shape (Scheduler in `scheduler/`, ledger writes next to start) fits Qwen3.5 better than Qwen3's extra `frontend_adapter.rs`, because `plan.rs` was already the mechanics split. TokenEvent drop-vs-publish tests are structurally gone; keep the failure-mode comments. Repo `models/` is absent on this machine; the Qwen3.5-4B fixture lives at `/data/models/Qwen3.5-4B`. GPU tests must load through `start_engine*` with an explicit small `max_batch` so recurrent-state reservation matches the 16GB budget. +- **Follow-ups**: `serving_tp2` / `tp2_scheduler_runs_forced_mixed_steps` still skipped — only one GPU. A fake-backend protocol suite (abort/reject/stop) still cannot mint `RequestLedger` from this crate. diff --git a/docs/subsystems/frontend/frontend-architecture.md b/docs/subsystems/frontend/frontend-architecture.md index dc84f0636..55e6fdc17 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 + typestate request handles + a contract-owned polling driver — Qwen3 and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/qwen35/kimi-k2/deepseek-v2-lite/gemma4 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 + typestate request handles + a contract-owned polling driver — Qwen3, Qwen3.5, and `pegainfer-sim` are migrated) and the **legacy handle contract** (`EngineHandle` + `TokenEvent` per-request events — glm52/kimi-k2/deepseek-v2-lite/gemma4 still launch through it). **Next step: migrate glm52, then delete the legacy contract.** Last touched: 2026-08 @@ -56,7 +56,7 @@ The contract's own invariants are tested in `emitter.rs`/`driver.rs` tests; the ## The legacy handle contract (migration pending) -`request.rs`/`event.rs`/`sink.rs`/`kv.rs`/`handle.rs` still carry the previous generation: `launch -> EngineHandle`, per-request `TokenSink` events (`Scheduled … Token* … terminal` by convention), send-failure-as-cancellation. glm52, qwen35, kimi-k2, deepseek-v2-lite, and gemma4 launch through it (`LaunchedEngine::Handle`), and the vllm stack keeps both bridge paths (`bridge.rs` for handles, `bridge/stepped.rs` for step engines). KV-prefix resolution (`KvPrefix`, `submit_resolved`) currently exists only on the legacy path; fold it into the step contract when the first offload-capable line migrates. +`request.rs`/`event.rs`/`sink.rs`/`kv.rs`/`handle.rs` still carry the previous generation: `launch -> EngineHandle`, per-request `TokenSink` events (`Scheduled … Token* … terminal` by convention), send-failure-as-cancellation. glm52, kimi-k2, deepseek-v2-lite, and gemma4 launch through it (`LaunchedEngine::Handle`), and the vllm stack keeps both bridge paths (`bridge.rs` for handles, `bridge/stepped.rs` for step engines). KV-prefix resolution (`KvPrefix`, `submit_resolved`) currently exists only on the legacy path; fold it into the step contract when the first offload-capable line migrates. ## Crate layout @@ -111,4 +111,4 @@ All six lines are onboarded. Adding a model line = write `model_line.rs` in the ## Next step -Migrate glm52 onto the step contract (second pilot; brings P/D and EP multi-scheduler requirements), then qwen35/kimi-k2/deepseek-v2-lite, then delete the legacy contract modules and `LaunchedEngine::Handle`. +Migrate glm52 onto the step contract (P/D and EP multi-scheduler requirements), then kimi-k2/deepseek-v2-lite/gemma4, then delete the legacy contract modules and `LaunchedEngine::Handle`. diff --git a/docs/subsystems/frontend/prometheus-metrics.md b/docs/subsystems/frontend/prometheus-metrics.md index e17df01c8..0085b1b3c 100644 --- a/docs/subsystems/frontend/prometheus-metrics.md +++ b/docs/subsystems/frontend/prometheus-metrics.md @@ -2,7 +2,7 @@ **TL;DR:** `/metrics` exposes request histograms for every model and engine gauges for schedulers that publish `LoadSnapshot`: Qwen3 and Qwen3.5 use one logical engine, while GLM5.2 EP8/DP8 uses eight rank-local engines and GLM5.2 TP8 uses one logical engine. The bridge forwards each partition's stats under the same identity the vLLM frontend uses for least-load routing. -Last touched: 2026-07 +Last touched: 2026-08 ## How the numbers flow @@ -11,7 +11,7 @@ Two independent paths feed the upstream Prometheus registry (`vllm-metrics`, ser 1. **Per-request path (works for every model crate).** The bridge stamps each request's first output with `Queued`/`Scheduled` timestamps and `PrefillStats` (prompt/computed/cached token split). The upstream `RequestMetricsTracker` turns those into `time_to_first_token_seconds`, `inter_token_latency_seconds`, `request_queue_time_seconds`, `prompt_tokens_total`, `generation_tokens_total`, `request_success_total`, `prompt_tokens_by_source_total`, … unconditionally — `disable_log_stats` only gates the periodic *text* logger, not Prometheus. 2. **Engine-gauge path (needs one `LoadSnapshot` watch per scheduler partition).** The scheduler publishes `LoadSnapshot { kv_used_blocks, kv_total_blocks, num_running_reqs, num_waiting_reqs }` at scheduler boundaries; one bridge identity per partition forwards its snapshot as a stats-only `RequestBatchOutputs`. The enclosing `engine_index` is both the routing identity and the Prometheus `engine` label. Watches coalesce to ≤1 message per scheduler step, and the scheduler's final idle publish settles the gauges back to 0. -For a single-partition model, `EngineHandle::with_load_watch` keeps the original one-engine contract. Qwen3.5 uses that contract for both its single-GPU backend and its TP backend because both execute one logical request stream through one scheduler. A partitioned scheduler uses `with_load_watches`, and the frontend launch declares the same engine count; a mismatch fails startup. GLM5.2 EP8 therefore registers engines 0–7, each bound to its own pending queue and KV pool. TP8 registers only engine 0 because its eight workers mirror one logical request stream. +One logical request stream is one engine. Qwen3 and Qwen3.5 publish `Scheduler::metrics` after each `step` (single-GPU and TP share that one stream); the stepped bridge pulls the snapshot once per driver iteration. Legacy-handle lines still attach `EngineHandle::with_load_watch`. A partitioned scheduler uses one watch per partition, and the frontend launch declares the same engine count; a mismatch fails startup. GLM5.2 EP8 therefore registers engines 0–7, each bound to its own pending queue and KV pool. TP8 registers only engine 0 because its eight workers mirror one logical request stream. Measured cost is noise in both covered configurations: diff --git a/pegainfer-qwen35/src/lib.rs b/pegainfer-qwen35/src/lib.rs index 0638d222c..96df209b5 100644 --- a/pegainfer-qwen35/src/lib.rs +++ b/pegainfer-qwen35/src/lib.rs @@ -30,7 +30,7 @@ use std::path::Path; use anyhow::Result; use anyhow::anyhow; pub(crate) use config::probe_config_json; -use pegainfer_frontend::engine::EngineHandle; +use pegainfer_frontend::engine::Engine; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::EpBackend; pub use scheduler::DEFAULT_MAX_PREFILL_TOKENS; @@ -101,7 +101,7 @@ pub fn start_engine( options: EngineLoadOptions, max_batch: usize, max_prefill_tokens: usize, -) -> Result { +) -> Result { start_engine_with_capacity_and_policy( model_path, options, @@ -140,7 +140,7 @@ pub fn launch_with_options_policy_and_overlap( options: Qwen35LaunchOptions, scheduler_policy: Qwen35SchedulerPolicy, decode_overlap: Qwen35DecodeOverlap, -) -> Result { +) -> Result { let device_ordinals = options.device_ordinals()?; start_engine_with_capacity_policy_and_overlap( model_path, @@ -163,7 +163,7 @@ pub fn start_engine_with_capacity( options: EngineLoadOptions, max_batch: usize, max_prefill_tokens: usize, -) -> Result { +) -> Result { start_engine_with_capacity_and_policy( model_path, options, @@ -179,7 +179,7 @@ pub(crate) fn start_engine_with_capacity_and_policy( max_batch: usize, max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, -) -> Result { +) -> Result { start_engine_with_capacity_policy_and_overlap( model_path, options, @@ -197,7 +197,7 @@ pub fn start_engine_with_capacity_policy_and_overlap( max_prefill_tokens: usize, scheduler_policy: Qwen35SchedulerPolicy, decode_overlap: Qwen35DecodeOverlap, -) -> Result { +) -> Result { anyhow::ensure!( (1..=MAX_DECODE_BATCH).contains(&max_batch), "Qwen3.5 max_batch must be in 1..={MAX_DECODE_BATCH}, got {max_batch}" diff --git a/pegainfer-qwen35/src/model_line.rs b/pegainfer-qwen35/src/model_line.rs index e6a47194f..88332dbcf 100644 --- a/pegainfer-qwen35/src/model_line.rs +++ b/pegainfer-qwen35/src/model_line.rs @@ -9,6 +9,7 @@ use pegainfer_frontend::model_line::CliDecodeOverlap; use pegainfer_frontend::model_line::CliError; use pegainfer_frontend::model_line::LaunchContext; use pegainfer_frontend::model_line::ModelLine; +use pegainfer_frontend::model_line::ServePlan; use crate::Qwen35DecodeOverlap; use crate::Qwen35LaunchOptions; @@ -144,6 +145,10 @@ impl ModelLine for Qwen35Line { Ok(()) } + fn serve_plan(&self, _ctx: &LaunchContext<'_>) -> Result { + Ok(ServePlan::default()) + } + fn launch(&self, ctx: &LaunchContext<'_>) -> anyhow::Result { let cli = cli(ctx); crate::launch_with_options_policy_and_overlap( @@ -162,7 +167,7 @@ impl ModelLine for Qwen35Line { resolve_decode_overlap(ctx.shared.decode_overlap) .map_err(|error| anyhow::anyhow!(error.to_string()))?, ) - .map(LaunchedEngine::Handle) + .map(LaunchedEngine::Stepped) } } diff --git a/pegainfer-qwen35/src/scheduler.rs b/pegainfer-qwen35/src/scheduler.rs deleted file mode 100644 index c2552c23c..000000000 --- a/pegainfer-qwen35/src/scheduler.rs +++ /dev/null @@ -1,2624 +0,0 @@ -//! Scheduler for Qwen3.5: dedicated GPU thread that batches concurrent requests. -//! -//! Mirrors the Qwen3 scheduler but manages: -//! - `RecurrentState` alongside `KvState` (linear attention layers) -//! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) - -mod plan; - -use std::collections::HashMap; -use std::collections::HashSet; -use std::collections::VecDeque; -use std::sync::Arc; -use std::sync::OnceLock; -use std::sync::mpsc as std_mpsc; -use std::thread; -use std::time::Instant; -use std::time::SystemTime; -use std::time::UNIX_EPOCH; - -use anyhow::Result; -use cudarc::driver::CudaEvent; -use cudarc::driver::CudaStream; -use cudarc::driver::sys; -use log::debug; -use log::info; -use log::warn; -use pegainfer_core::kv_pool::KvState; -use pegainfer_core::tensor::HiddenStates; -use pegainfer_frontend::engine::EngineHandle as SchedulerHandle; -use pegainfer_frontend::engine::FinishReason; -use pegainfer_frontend::engine::GenerateRequest as SchedulerRequest; -use pegainfer_frontend::engine::KvCapacity; -use pegainfer_frontend::engine::SchedulerMetrics; -use pegainfer_frontend::engine::SubmittedRequest; -use pegainfer_frontend::engine::TokenEvent; -use pegainfer_frontend::engine::TokenLogprob; -use pegainfer_frontend::engine::TokenSink; -use pegainfer_frontend::engine::panic_message; -use pegainfer_frontend::sampler::SamplingParams; -use rand::SeedableRng; -use rand::rngs::StdRng; -use tokio::sync::mpsc; -use tokio::sync::watch; - -use self::plan::ActiveDecodeState; -use self::plan::ActiveKvBudget; -use self::plan::ExecutionPlan; -use self::plan::PrefillKvBudget; -use self::plan::PrefillQueueState; -use self::plan::RejectReason; -use self::plan::admit_pending_requests; -use self::plan::choose_prefill_budget; -use self::plan::compaction_after_retire; -use self::plan::max_kv_tokens; -use self::plan::plan_prefill_chunks; -use self::plan::prefilling_future_pages; -use self::plan::slot_for_new_request; -use crate::Qwen35DecodeOverlap; -use crate::Qwen35SchedulerPolicy; -use crate::batch_decode_graph::BatchDecodeGraphState; -use crate::executor::DecodeRequestResult; -use crate::executor::DecodeResult; -use crate::executor::PrefillRequestResult; -use crate::executor::PrefillResult; -use crate::executor::RequestId; -use crate::logprobs::snapshot_requested_logprobs; -use crate::recurrent_state::RecurrentState; -use crate::tp_executor::DropExpectation; -use crate::tp_executor::Qwen35TpExecutor; -use crate::tp_executor::TpDecodeStepItem; -use crate::tp_executor::TpPrefillChunkItem; -use crate::tp_executor::TpUnifiedPlan; -use crate::weights::Qwen35Model; - -// ── Internal types ────────────────────────────────────────────────────── - -/// An in-flight request being decoded. Recurrent state lives in the -/// `BatchDecodeGraphState` at `graph_slot_idx` — NOT owned here. -struct ActiveRequest35 { - request_id: Option, - token_tx: TokenSink, - backend_state: ActiveBackendState, - last_token: u32, - generated_count: usize, - max_tokens: usize, - prompt_len: usize, - params: SamplingParams, - /// Number of top logprobs to return (0 = disabled). - logprobs: usize, -} - -/// A request whose prompt is being prefilled across multiple scheduler steps. -/// It owns its growing KV and recurrent state until the prompt is exhausted, -/// at which point it is promoted into the decode batch. -struct PrefillingRequest35 { - req: SchedulerRequest, - backend_state: PrefillBackendState, - /// Prompt tokens prefilled so far. - cursor: usize, - /// Tokens to prefill in the step currently scheduled (set by `take_prefill_chunks`). - step_chunk: usize, -} - -enum ActiveBackendState { - Single { - kv: KvState, - /// Index into `BatchDecodeGraphState.slot_states`. - graph_slot_idx: usize, - }, - Tp { - request_id: RequestId, - }, -} - -enum PrefillBackendState { - Single { kv: KvState, rec: RecurrentState }, - Tp { request_id: RequestId }, -} - -struct TerminalRequest { - token_tx: TokenSink, - prompt_tokens: usize, - completion_tokens: usize, -} - -impl TerminalRequest { - fn send_error(self, message: &str) { - let _ = self.token_tx.send(TokenEvent::Error { - message: message.to_string(), - prompt_tokens: self.prompt_tokens, - completion_tokens: self.completion_tokens, - }); - } -} - -impl From for TerminalRequest { - fn from(req: SchedulerRequest) -> Self { - Self { - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - token_tx: req.token_tx, - } - } -} - -impl From for TerminalRequest { - fn from(req: ActiveRequest35) -> Self { - Self { - token_tx: req.token_tx, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - } - } -} - -impl From for TerminalRequest { - fn from(req: PrefillingRequest35) -> Self { - req.req.into() - } -} - -struct PrefillCompletionRequest { - req: SchedulerRequest, - backend_state: PrefillBackendState, -} - -trait CompletionRequest { - fn token_tx(&self) -> &TokenSink; - fn into_terminal(self) -> TerminalRequest; -} - -impl CompletionRequest for ActiveRequest35 { - fn token_tx(&self) -> &TokenSink { - &self.token_tx - } - - fn into_terminal(self) -> TerminalRequest { - self.into() - } -} - -impl CompletionRequest for PrefillCompletionRequest { - fn token_tx(&self) -> &TokenSink { - &self.req.token_tx - } - - fn into_terminal(self) -> TerminalRequest { - self.req.into() - } -} - -struct CompletionCandidate { - request: R, - final_events: Vec, -} - -impl CompletionCandidate { - fn commit(self) { - for event in self.final_events { - let _ = self.request.token_tx().send(event); - } - } - - fn into_terminal(self) -> TerminalRequest { - self.request.into_terminal() - } -} - -struct FatalSchedulerError { - message: String, - transient: Vec, -} - -#[derive(Clone, Debug, PartialEq)] -struct PrefillArtifact { - token: u32, - logprob: Option, -} - -#[derive(Clone, Debug, PartialEq)] -struct DecodeArtifact { - token: u32, - logprob: Option, -} - -struct AlignedUnifiedArtifacts { - prefill: Vec>, - decode: Vec, -} - -enum PrefillStepArtifacts { - Single { - tokens: Vec, - logprobs: Vec>, - }, - Tp(Vec>), -} - -impl PrefillStepArtifacts { - fn final_artifact(&self, idx: usize) -> PrefillArtifact { - match self { - Self::Single { tokens, logprobs } => PrefillArtifact { - token: tokens[idx], - logprob: logprobs[idx].clone(), - }, - Self::Tp(artifacts) => artifacts[idx] - .clone() - .expect("validated TP final-prefill row must contain an artifact"), - } - } -} - -impl FatalSchedulerError { - fn new(message: impl Into) -> Self { - Self { - message: message.into(), - transient: Vec::new(), - } - } - - fn with_request(mut self, request: impl Into) -> Self { - self.transient.push(request.into()); - self - } - - fn with_requests(mut self, requests: I) -> Self - where - I: IntoIterator, - R: Into, - { - self.transient.extend(requests.into_iter().map(Into::into)); - self - } -} - -pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; - -/// Env-gated per-step ITL diagnostics (issue #470). When `PEGAINFER_ITL_DEBUG` -/// is set, the scheduler emits one `ITL_STEP` line per executed step, tagging -/// the plan kind, the *actual* prefill-chunk token count associated with the -/// action, the active decode width, and the CPU wall-time. This lets the -/// mixed-load bench separate serial Unified stalls from overlap launch, -/// decode, completion, and wait actions instead of relying on the coarse -/// `[submit, last-token]` injection window. Off by default: no cost on the -/// normal bench path. -fn itl_debug_enabled() -> bool { - static ENABLED: OnceLock = OnceLock::new(); - *ENABLED.get_or_init(|| std::env::var_os("PEGAINFER_ITL_DEBUG").is_some()) -} - -/// Monotonic microseconds since the first ITL step, so `ITL_STEP` timestamps -/// are correlatable within one process run (paired with wall-clock epoch us). -fn itl_debug_mono_us() -> u128 { - static ORIGIN: OnceLock = OnceLock::new(); - ORIGIN.get_or_init(Instant::now).elapsed().as_micros() -} - -fn log_itl_step( - step_start: Option, - plan: &str, - prefill_tokens: usize, - prefill_reqs: usize, - decode_n: usize, -) { - let Some(step_start) = step_start else { - return; - }; - let dur_us = step_start.elapsed().as_micros(); - let epoch_us = SystemTime::now() - .duration_since(UNIX_EPOCH) - .map_or(0, |d| d.as_micros()); - info!( - "ITL_STEP mono_us={} epoch_us={} plan={} prefill_tok={} prefill_reqs={} decode_n={} dur_us={}", - itl_debug_mono_us(), - epoch_us, - plan, - prefill_tokens, - prefill_reqs, - decode_n, - dur_us - ); -} - -// ── Entry point ───────────────────────────────────────────────────────── - -pub fn start_with_capacity( - model: Qwen35Model, - seed: u64, - max_batch: usize, - max_prefill_tokens: usize, -) -> Result { - start_with_capacity_and_policy( - model, - seed, - max_batch, - max_prefill_tokens, - Qwen35SchedulerPolicy::Off, - Qwen35DecodeOverlap::Off, - ) -} - -pub(crate) fn start_with_capacity_and_policy( - model: Qwen35Model, - seed: u64, - max_batch: usize, - max_prefill_tokens: usize, - scheduler_policy: Qwen35SchedulerPolicy, - decode_overlap: Qwen35DecodeOverlap, -) -> Result { - assert!( - max_prefill_tokens > 0, - "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" - ); - // Static instance cap for the vLLM bridge's max_model_len. Live admission - // still uses the current page budget inside the scheduler loop. - let total_blocks = model.kv_pool().capacity_pages().saturating_sub(1); - let kv_total_blocks = total_blocks as u64; - let block_size = model.kv_pool().layout().page_size; - let servable = servable_len( - model.config().max_position_embeddings, - total_blocks, - block_size, - ); - let backend = SingleGpuBackend::new(model, max_batch, decode_overlap)?; - - let (submit_tx, submit_rx) = mpsc::unbounded_channel(); - let (startup_tx, startup_rx) = std_mpsc::channel(); - let (load_tx, load_rx) = watch::channel(SchedulerMetrics { - kv_total_blocks, - ..SchedulerMetrics::default() - }); - - let join_handle = thread::Builder::new() - .name("scheduler-qwen35".into()) - .spawn(move || match bind_model_thread(backend.model()) { - Ok(_guard) => { - let _ = startup_tx.send(Ok(())); - scheduler_loop( - SchedulerBackend::Single(backend), - submit_rx, - seed, - max_prefill_tokens, - scheduler_policy, - load_tx, - ); - } - Err(err) => { - let _ = startup_tx.send(Err(err)); - } - }) - .expect("failed to spawn Qwen3.5 scheduler thread"); - - let Ok(startup) = startup_rx.recv() else { - let panic_note = match join_handle.join() { - Err(panic) => format!(" (thread panicked: {})", panic_message(panic.as_ref())), - Ok(()) => String::new(), - }; - anyhow::bail!("Qwen3.5 scheduler exited during startup{panic_note}"); - }; - if let Err(err) = startup { - let _ = join_handle.join(); - return Err(err); - } - Ok( - SchedulerHandle::new_with_join_handle(submit_tx, join_handle) - .with_servable_len(servable) - .with_kv_capacity(KvCapacity { - total_blocks, - block_size, - }) - .with_metrics_watch(load_rx), - ) -} - -pub(crate) fn start_tp_with_capacity( - model_path: &str, - seed: u64, - device_ordinals: &[usize], - max_batch: usize, - max_prefill_tokens: usize, -) -> Result { - assert!( - max_prefill_tokens > 0, - "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" - ); - let backend = - TpSchedulerBackend::new(model_path, device_ordinals, max_batch, max_prefill_tokens)?; - let servable = servable_len( - backend.max_position_embeddings(), - backend.capacity_pages_for_requests(), - backend.page_size(), - ); - let kv_capacity = KvCapacity { - total_blocks: backend.capacity_pages_for_requests(), - block_size: backend.page_size(), - }; - - let (submit_tx, submit_rx) = mpsc::unbounded_channel(); - let (load_tx, load_rx) = watch::channel(SchedulerMetrics { - kv_total_blocks: kv_capacity.total_blocks as u64, - ..SchedulerMetrics::default() - }); - let join_handle = thread::Builder::new() - .name("scheduler-qwen35-tp".into()) - .spawn(move || { - scheduler_loop( - SchedulerBackend::Tp(backend), - submit_rx, - seed, - max_prefill_tokens, - Qwen35SchedulerPolicy::Off, - load_tx, - ); - }) - .expect("failed to spawn Qwen3.5 TP scheduler thread"); - - Ok( - SchedulerHandle::new_with_join_handle(submit_tx, join_handle) - .with_servable_len(servable) - .with_kv_capacity(kv_capacity) - .with_metrics_watch(load_rx), - ) -} - -struct SingleGpuBackend { - model: Qwen35Model, - graph_state: BatchDecodeGraphState, - prefill_stream: Option>, -} - -// One instance per scheduler; the size asymmetry costs nothing here. -#[allow(clippy::large_enum_variant)] -enum SchedulerBackend { - Single(SingleGpuBackend), - Tp(TpSchedulerBackend), -} - -struct AsyncPrefillOutput { - logits: Option, - done: CudaEvent, - stream: Arc, - completed: bool, -} - -impl AsyncPrefillOutput { - fn is_ready(&mut self) -> bool { - match unsafe { sys::cuEventQuery(self.done.cu_event()) } { - sys::CUresult::CUDA_SUCCESS => { - self.completed = true; - true - } - sys::CUresult::CUDA_ERROR_NOT_READY => false, - err => fatal_cuda_lifecycle(&format!( - "query Qwen3.5 async prefill event failed: {err:?}" - )), - } - } - - fn into_logits(mut self) -> HiddenStates { - if !self.completed { - if let Err(err) = self.done.synchronize() { - fatal_cuda_lifecycle(&format!("wait for Qwen3.5 async prefill failed: {err}")); - } - self.completed = true; - } - self.logits - .take() - .expect("async prefill logits must be consumed exactly once") - } -} - -impl Drop for AsyncPrefillOutput { - fn drop(&mut self) { - if self.completed { - return; - } - if let Err(err) = self.stream.synchronize() { - fatal_cuda_lifecycle(&format!( - "drain Qwen3.5 async prefill during cleanup failed: {err}" - )); - } - } -} - -fn fatal_cuda_lifecycle(message: &str) -> ! { - log::error!("FATAL: {message}; aborting before CUDA-referenced state is released"); - std::process::abort(); -} - -struct TpSchedulerBackend { - executor: Qwen35TpExecutor, - next_request_id: u64, -} - -impl SingleGpuBackend { - fn new( - model: Qwen35Model, - max_batch: usize, - decode_overlap: Qwen35DecodeOverlap, - ) -> Result { - anyhow::ensure!(max_batch > 0, "Qwen3.5 max_batch must be > 0"); - let graph_capacity = crate::batch_decode_graph::bucket_for(max_batch); - let graph_state = model.create_batch_decode_graph_state_with_capacity(graph_capacity)?; - let prefill_stream = match decode_overlap { - Qwen35DecodeOverlap::Off => None, - Qwen35DecodeOverlap::SharedSm => Some( - model - .device_ctx() - .ctx - .new_stream() - .map_err(|err| anyhow::anyhow!("create Qwen3.5 prefill stream: {err}"))?, - ), - }; - Ok(Self { - model, - graph_state, - prefill_stream, - }) - } - - fn model(&self) -> &Qwen35Model { - &self.model - } - - fn max_batch(&self) -> usize { - // #470: admit the requested `--max-batch`, which may sit below the loaded - // graph bucket (e.g. 5 on bucket 8); never exceed the physical slots. - self.model - .decode_admission_batch - .min(self.graph_state.slot_states.len()) - .max(1) - } - - fn page_size(&self) -> usize { - self.model.kv_pool().layout().page_size - } - - fn available_pages(&self) -> usize { - self.model.kv_pool().available_pages() - } - - fn capacity_pages_for_requests(&self) -> usize { - self.model.kv_pool().capacity_pages().saturating_sub(1) - } - - fn max_position_embeddings(&self) -> usize { - self.model.config().max_position_embeddings - } - - fn alloc_kv(&self) -> KvState { - self.model.alloc_kv() - } - - fn alloc_recurrent(&self) -> Result { - RecurrentState::new(self.model.device_ctx(), self.model.config()) - } - - fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { - anyhow::bail!("single-GPU prefill received TP chunk state"); - }; - let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - self.model - .batch_prefill_logits(&window_refs, kvs, &mut rec_refs) - } - - fn overlap_enabled(&self) -> bool { - self.prefill_stream.is_some() - } - - fn launch_async_prefill(&mut self, chunk: &mut ScheduledChunk) -> Result { - let prefill_stream = self - .prefill_stream - .clone() - .ok_or_else(|| anyhow::anyhow!("Qwen3.5 decode overlap is disabled"))?; - - // Request KV/recurrent state was allocated on the model stream. Order - // those producers before the prefill stream without blocking the host. - prefill_stream - .join(&self.model.device_ctx().stream) - .map_err(|err| anyhow::anyhow!("join Qwen3.5 prefill stream: {err}"))?; - - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { - anyhow::bail!("single-GPU async prefill received TP chunk state"); - }; - let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - let logits = match self.model.batch_prefill_logits_on_stream( - Arc::clone(&prefill_stream), - &window_refs, - kvs, - &mut rec_refs, - ) { - Ok(logits) => logits, - Err(err) => { - if let Err(sync_err) = prefill_stream.synchronize() { - fatal_cuda_lifecycle(&format!( - "Qwen3.5 async prefill failed ({err}); stream drain failed: {sync_err}" - )); - } - return Err(err); - } - }; - let done = match prefill_stream.record_event(None) { - Ok(done) => done, - Err(err) => { - if let Err(sync_err) = prefill_stream.synchronize() { - fatal_cuda_lifecycle(&format!( - "record Qwen3.5 async prefill event failed ({err}); stream drain failed: {sync_err}" - )); - } - return Err(anyhow::anyhow!("record Qwen3.5 async prefill event: {err}")); - } - }; - Ok(AsyncPrefillOutput { - logits: Some(logits), - done, - stream: prefill_stream, - completed: false, - }) - } - - fn unified_step( - &mut self, - chunk: &mut ScheduledChunk, - active: &mut [ActiveRequest35], - ) -> Result { - let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); - let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { - anyhow::bail!("single-GPU unified step received TP chunk state"); - }; - let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); - let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); - let mut decode_kv_refs: Vec<&mut KvState> = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU unified step received TP active state") - } - }) - .collect(); - self.model.unified_step( - &window_refs, - kvs, - &mut rec_refs, - &decode_tokens, - &mut decode_kv_refs, - &mut self.graph_state, - ) - } - - fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { - let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); - let mut kv_refs: Vec<&mut KvState> = active - .iter_mut() - .map(|r| match &mut r.backend_state { - ActiveBackendState::Single { kv, .. } => kv, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU decode received TP active state") - } - }) - .collect(); - self.model - .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) - } - - fn sample_prefill_logits( - &mut self, - pending: &[SchedulerRequest], - logits: &HiddenStates, - sample_seed: u64, - ) -> Result<(Vec, Vec>)> { - debug_assert_eq!( - logits.seq_len, - pending.len(), - "Qwen3.5 prefill logits rows must preserve pending request order" - ); - let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect(); - let cpu_logits = - snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?; - let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); - let tokens = self.model.select_tokens_from_logits_varied( - logits, - &mut self.graph_state.buffers, - ¶ms_refs, - sample_seed, - )?; - - let logprobs = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - pegainfer_sample::token_logprob_from_row( - &logits_f32, - tokens[i], - pending[i].logprobs, - ) - }) - }) - .collect(); - Ok((tokens, logprobs)) - } - - fn sample_decode_logits( - &mut self, - active: &[ActiveRequest35], - sample_seed: u64, - ) -> Result<(Vec, Vec>)> { - let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); - let cpu_logits = snapshot_requested_logprobs( - self.model.device_ctx(), - &self.graph_state.buffers.logits, - &requested_logprobs, - )?; - let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); - let tokens = self.model.select_tokens_batch_varied( - &mut self.graph_state.buffers, - ¶ms_refs, - sample_seed, - )?; - - let logprobs = cpu_logits - .into_iter() - .enumerate() - .map(|(i, logits_opt)| { - logits_opt.and_then(|logits_f32| { - pegainfer_sample::token_logprob_from_row( - &logits_f32, - tokens[i], - active[i].logprobs, - ) - }) - }) - .collect(); - Ok((tokens, logprobs)) - } - - fn is_stop_token(&self, token: u32) -> bool { - self.model.is_stop_token(token) - } - - fn copy_recurrent_to_slot( - &mut self, - recurrent: &RecurrentState, - slot_idx: usize, - ) -> Result<()> { - self.graph_state - .copy_state_to_slot(self.model.device_ctx(), recurrent, slot_idx) - } - - fn compact_slot(&mut self, active: &mut [ActiveRequest35], compaction: plan::SlotCompaction) { - let src_slot = match active[compaction.moved_to].backend_state { - ActiveBackendState::Single { graph_slot_idx, .. } => graph_slot_idx, - ActiveBackendState::Tp { .. } => { - panic!("single-GPU slot compaction received TP active state") - } - }; - debug_assert_eq!(src_slot, compaction.moved_from); - - let ctx = self.model.device_ctx(); - let src = &self.graph_state.slot_states[compaction.moved_from]; - for layer_idx in 0..src.layers.len() { - let (src_part, dst_part) = if compaction.moved_to < compaction.moved_from { - let (left, right) = self - .graph_state - .slot_states - .split_at_mut(compaction.moved_from); - ( - &right[0].layers[layer_idx], - &mut left[compaction.moved_to].layers[layer_idx], - ) - } else { - unreachable!("idx < active.len() <= last"); - }; - - ctx.stream - .memcpy_dtod(&src_part.state, &mut dst_part.state) - .expect("compact slot state copy failed"); - ctx.stream - .memcpy_dtod(&src_part.conv_state.data, &mut dst_part.conv_state.data) - .expect("compact slot conv_state copy failed"); - } - self.graph_state.slot_states[compaction.moved_to].seq_len = - self.graph_state.slot_states[compaction.moved_from].seq_len; - - match &mut active[compaction.moved_to].backend_state { - ActiveBackendState::Single { graph_slot_idx, .. } => { - *graph_slot_idx = compaction.moved_to; - } - ActiveBackendState::Tp { .. } => { - panic!("single-GPU slot compaction received TP active state") - } - } - } -} - -impl TpSchedulerBackend { - fn new( - model_path: &str, - device_ordinals: &[usize], - max_batch: usize, - max_prefill_tokens: usize, - ) -> Result { - let executor = Qwen35TpExecutor::from_runtime_with_limits( - model_path, - false, - device_ordinals, - max_batch, - max_prefill_tokens, - )?; - Ok(Self { - executor, - next_request_id: 1, - }) - } - - fn alloc_request_id(&mut self) -> RequestId { - let id = RequestId::new(self.next_request_id); - self.next_request_id = self.next_request_id.wrapping_add(1).max(1); - id - } - - fn max_batch(&self) -> usize { - self.executor.max_batch() - } - - fn page_size(&self) -> usize { - self.executor.page_size() - } - - fn capacity_pages_for_requests(&self) -> usize { - self.executor.capacity_pages_for_requests() - } - - fn max_position_embeddings(&self) -> usize { - self.executor.max_position_embeddings() - } - - fn is_stop_token(&self, token: u32) -> bool { - self.executor.is_stop_token(token) - } - - fn available_pages( - &self, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - ) -> usize { - let page_size = self.page_size(); - let active_pages: usize = active - .iter() - .map(|req| pages_needed(current_active_tokens(req), page_size)) - .sum(); - let prefilling_pages: usize = prefilling - .iter() - .map(|req| pages_needed(req.cursor, page_size)) - .sum(); - self.capacity_pages_for_requests() - .saturating_sub(active_pages.saturating_add(prefilling_pages)) - } - - fn execute_prefill_chunk( - &self, - chunk: &ScheduledChunk, - sample_seed: u64, - ) -> Result>> { - let items = tp_prefill_items(chunk)?; - let result = self - .executor - .execute_prefill_chunks_with_seed(&items, sample_seed)?; - align_prefill_results(chunk, &result) - .map_err(|err| self.executor.poison_artifact_contract("prefill", &err)) - } - - fn execute_decode( - &self, - active: &[ActiveRequest35], - sample_seed: u64, - ) -> Result> { - let items = tp_decode_items(active)?; - let result = self.executor.execute_decode_items(&items, sample_seed)?; - align_decode_results(active, &result) - .map_err(|err| self.executor.poison_artifact_contract("decode", &err)) - } - - fn execute_unified( - &self, - chunk: &ScheduledChunk, - active: &[ActiveRequest35], - decode_sample_seed: u64, - prefill_sample_seed: u64, - ) -> Result { - let plan = TpUnifiedPlan { - prefill: tp_prefill_items(chunk)?, - decode: tp_decode_items(active)?, - prefill_sample_seed, - decode_sample_seed, - }; - let result = self.executor.execute_unified(&plan)?; - let prefill = align_prefill_results(chunk, &result.prefill).map_err(|err| { - self.executor - .poison_artifact_contract("unified prefill", &err) - })?; - let decode = align_decode_results(active, &result.decode).map_err(|err| { - self.executor - .poison_artifact_contract("unified decode", &err) - })?; - Ok(AlignedUnifiedArtifacts { prefill, decode }) - } - - fn drop_request(&self, request_id: RequestId, expectation: DropExpectation) -> Result<()> { - self.executor.drop_request(request_id, expectation) - } -} - -impl SchedulerBackend { - fn max_batch(&self) -> usize { - match self { - Self::Single(backend) => backend.max_batch(), - Self::Tp(backend) => backend.max_batch(), - } - } - - fn page_size(&self) -> usize { - match self { - Self::Single(backend) => backend.page_size(), - Self::Tp(backend) => backend.page_size(), - } - } - - fn available_pages( - &self, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - ) -> usize { - match self { - Self::Single(backend) => backend.available_pages(), - Self::Tp(backend) => backend.available_pages(active, prefilling), - } - } - - fn capacity_pages_for_requests(&self) -> usize { - match self { - Self::Single(backend) => backend.capacity_pages_for_requests(), - Self::Tp(backend) => backend.capacity_pages_for_requests(), - } - } - - fn max_position_embeddings(&self) -> usize { - match self { - Self::Single(backend) => backend.max_position_embeddings(), - Self::Tp(backend) => backend.max_position_embeddings(), - } - } - - fn alloc_prefill_state(&mut self) -> Result { - match self { - Self::Single(backend) => Ok(PrefillBackendState::Single { - kv: backend.alloc_kv(), - rec: backend.alloc_recurrent()?, - }), - Self::Tp(backend) => Ok(PrefillBackendState::Tp { - request_id: backend.alloc_request_id(), - }), - } - } - - fn is_stop_token(&self, token: u32) -> bool { - match self { - Self::Single(backend) => backend.is_stop_token(token), - Self::Tp(backend) => backend.is_stop_token(token), - } - } -} - -fn current_active_tokens(req: &ActiveRequest35) -> usize { - req.prompt_len - .saturating_add(req.generated_count.saturating_sub(1)) -} - -fn pages_needed(token_count: usize, page_size: usize) -> usize { - token_count.div_ceil(page_size) -} - -fn tp_prefill_items(chunk: &ScheduledChunk) -> Result> { - let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { - anyhow::bail!("TP prefill received single-GPU chunk state"); - }; - anyhow::ensure!( - chunk.reqs.len() == request_ids.len() - && chunk.reqs.len() == chunk.windows.len() - && chunk.reqs.len() == chunk.ends.len(), - "Qwen3.5 TP scheduled prefill vectors are misaligned" - ); - Ok(chunk - .reqs - .iter() - .zip(request_ids) - .zip(&chunk.windows) - .zip(&chunk.ends) - .map(|(((req, request_id), window), end)| { - TpPrefillChunkItem::new_with_sampling( - *request_id, - window.clone(), - req.logprobs, - req.params, - *end == req.prompt_tokens.len(), - ) - }) - .collect()) -} - -fn tp_decode_items(active: &[ActiveRequest35]) -> Result> { - active - .iter() - .map(|req| { - let ActiveBackendState::Tp { request_id } = &req.backend_state else { - anyhow::bail!("TP decode received single-GPU active state"); - }; - Ok(TpDecodeStepItem::new( - *request_id, - req.last_token, - req.logprobs, - req.params, - )) - }) - .collect() -} - -fn align_prefill_results( - chunk: &ScheduledChunk, - result: &PrefillResult, -) -> Result>> { - let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { - anyhow::bail!("align_prefill_results requires TP chunk state"); - }; - anyhow::ensure!( - request_ids.len() == chunk.reqs.len() && chunk.ends.len() == chunk.reqs.len(), - "Qwen3.5 TP prefill alignment vectors are misaligned" - ); - let expected: HashSet = request_ids - .iter() - .zip(&chunk.reqs) - .zip(&chunk.ends) - .filter_map(|((&request_id, req), &end)| { - (end == req.prompt_tokens.len()).then_some(request_id) - }) - .collect(); - let mut by_id = HashMap::with_capacity(result.requests.len()); - for PrefillRequestResult { - request_id, - first_token, - first_token_logprob, - } in &result.requests - { - anyhow::ensure!( - expected.contains(request_id), - "Qwen3.5 TP prefill returned unknown or non-final request id {}", - request_id.get() - ); - let artifact = PrefillArtifact { - token: *first_token, - logprob: first_token_logprob.clone(), - }; - anyhow::ensure!( - by_id.insert(*request_id, artifact).is_none(), - "Qwen3.5 TP prefill returned duplicate request id {}", - request_id.get() - ); - } - anyhow::ensure!( - by_id.len() == expected.len(), - "Qwen3.5 TP prefill result is missing final request IDs" - ); - - request_ids - .iter() - .zip(&chunk.reqs) - .zip(&chunk.ends) - .map(|((&request_id, req), &end)| { - if end == req.prompt_tokens.len() { - by_id.remove(&request_id).map(Some).ok_or_else(|| { - anyhow::anyhow!( - "Qwen3.5 TP prefill result is missing final request id {}", - request_id.get() - ) - }) - } else { - Ok(None) - } - }) - .collect() -} - -fn align_decode_results( - active: &[ActiveRequest35], - result: &DecodeResult, -) -> Result> { - let expected: Vec = active - .iter() - .map(|active_req| { - let ActiveBackendState::Tp { request_id } = active_req.backend_state else { - anyhow::bail!("align_decode_results requires TP active state"); - }; - Ok(request_id) - }) - .collect::>()?; - let expected_set: HashSet<_> = expected.iter().copied().collect(); - anyhow::ensure!( - expected_set.len() == expected.len(), - "Qwen3.5 TP active decode IDs contain duplicates" - ); - let mut by_id = HashMap::with_capacity(result.requests.len()); - for DecodeRequestResult { - request_id, - token, - logprob, - } in &result.requests - { - anyhow::ensure!( - expected_set.contains(request_id), - "Qwen3.5 TP decode returned unknown request id {}", - request_id.get() - ); - let artifact = DecodeArtifact { - token: *token, - logprob: logprob.clone(), - }; - anyhow::ensure!( - by_id.insert(*request_id, artifact).is_none(), - "Qwen3.5 TP decode returned duplicate request id {}", - request_id.get() - ); - } - expected - .into_iter() - .map(|request_id| { - by_id.remove(&request_id).ok_or_else(|| { - anyhow::anyhow!( - "Qwen3.5 TP decode result is missing request id {}", - request_id.get() - ) - }) - }) - .collect() -} - -fn split_decode_artifacts(artifacts: &[DecodeArtifact]) -> (Vec, Vec>) { - artifacts - .iter() - .map(|artifact| (artifact.token, artifact.logprob.clone())) - .unzip() -} - -fn servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { - max_context - .min(max_pages.saturating_mul(page_size)) - .try_into() - .unwrap_or(u32::MAX) -} - -struct CublasThreadGuard; - -impl Drop for CublasThreadGuard { - fn drop(&mut self) { - unsafe { - crate::ffi::cublas_destroy(); - } - } -} - -fn bind_model_thread(model: &Qwen35Model) -> Result { - let ctx = model.device_ctx(); - unsafe { - let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); - if err != 0 { - return Err(anyhow::anyhow!( - "Failed to set CUDA device {} on Qwen3.5 scheduler thread: cudaError={}", - ctx.device_ordinal, - err - )); - } - } - ctx.ctx.bind_to_thread().map_err(|e| { - anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 scheduler thread: {e}") - })?; - unsafe { - crate::ffi::cublas_init(); - } - model.tune_decode_gemm_algos()?; - Ok(CublasThreadGuard) -} - -// ── Main loop ─────────────────────────────────────────────────────────── - -fn publish_load( - load_tx: &watch::Sender, - backend: &SchedulerBackend, - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - inflight_prefill_reqs: usize, - num_waiting_reqs: usize, -) { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - let (num_running_reqs, num_waiting_reqs) = - logical_load_counts(active, prefilling, inflight_prefill_reqs, num_waiting_reqs); - load_tx.send_replace(SchedulerMetrics { - kv_used_blocks: kv_total_blocks - .saturating_sub(backend.available_pages(active, prefilling) as u64), - kv_total_blocks, - num_running_reqs, - num_waiting_reqs, - spec_decode: None, - }); -} - -fn logical_load_counts( - active: &[ActiveRequest35], - prefilling: &[PrefillingRequest35], - inflight_prefill_reqs: usize, - num_waiting_reqs: usize, -) -> (u64, u64) { - ( - (active.len() + prefilling.len() + inflight_prefill_reqs) as u64, - num_waiting_reqs as u64, - ) -} - -fn should_block_on_submit(owned_work_empty: bool, inflight_prefill: bool) -> bool { - owned_work_empty && !inflight_prefill -} - -fn terminal_scheduler_shutdown( - submit_rx: &mut mpsc::UnboundedReceiver, - load_tx: &watch::Sender, - kv_total_blocks: u64, - active: Vec, - prefilling: Vec, - pending: Vec, - deferred: Vec, - inflight_prefill: Option, - failure: FatalSchedulerError, -) { - submit_rx.close(); - - let mut requests = failure.transient; - requests.extend(active.into_iter().map(Into::into)); - requests.extend(prefilling.into_iter().map(Into::into)); - requests.extend(pending.into_iter().map(Into::into)); - requests.extend(deferred.into_iter().map(Into::into)); - if let Some(InflightPrefill { output, chunk, .. }) = inflight_prefill { - // The stream must drain before the chunk's KV/recurrent/conv state is - // released or transferred into terminal request ownership. - drop(output); - requests.extend(chunk.reqs.into_iter().map(Into::into)); - } - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - requests.push(req.into()); - } - - warn!( - "Qwen3.5 TP scheduler terminating after replica failure: {}", - failure.message - ); - for request in requests { - request.send_error(&failure.message); - } - load_tx.send_replace(SchedulerMetrics { - kv_used_blocks: 0, - kv_total_blocks, - num_running_reqs: 0, - num_waiting_reqs: 0, - spec_decode: None, - }); -} - -fn prune_closed_requests( - backend: &mut B, - active: &mut Vec, - prefilling: &mut Vec, - pending: &mut Vec, -) -> std::result::Result<(), FatalSchedulerError> -where - B: DecodeDispatchBackend + PrefillPromoteBackend, -{ - pending.retain(|req| !req.token_tx.is_closed()); - - for idx in (0..active.len()).rev() { - if active[idx].token_tx.is_closed() { - debug!( - "request pruned before scheduling: request_id={:?} phase=decode tokens_generated={}", - active[idx].request_id, active[idx].generated_count - ); - let removed = backend.take_active_request(active, idx); - if let Err(err) = backend.drop_active_state(&removed.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); - } - } - } - - for idx in (0..prefilling.len()).rev() { - if prefilling[idx].req.token_tx.is_closed() { - let removed = prefilling.remove(idx); - debug!( - "request pruned before scheduling: request_id={:?} phase=prefill cursor={}", - removed.req.request_id, removed.cursor - ); - let expectation = if removed.cursor == 0 { - DropExpectation::MustBeAbsent - } else { - DropExpectation::MustExist - }; - if let Err(err) = backend.drop_prefill_state(&removed.backend_state, expectation) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(removed)); - } - } - } - Ok(()) -} - -const UNSUPPORTED_ECHO_MESSAGE: &str = "echo=true is unsupported by the Qwen3.5 serving contract"; - -fn reject_unsupported_echo(pending: &mut Vec) { - pending.retain(|req| { - if !req.echo { - return true; - } - let _ = req.token_tx.send(TokenEvent::Rejected { - message: UNSUPPORTED_ECHO_MESSAGE.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - false - }); -} - -#[allow(clippy::needless_pass_by_value)] -fn scheduler_loop( - mut backend: SchedulerBackend, - mut submit_rx: mpsc::UnboundedReceiver, - seed: u64, - prefill_budget: usize, - scheduler_policy: Qwen35SchedulerPolicy, - load_tx: watch::Sender, -) { - let mut rng = StdRng::seed_from_u64(seed); - let mut active: Vec = Vec::new(); - let mut deferred: Vec = Vec::new(); - let mut prefilling: Vec = Vec::new(); - let mut inflight_prefill: Option = None; - let max_batch = backend.max_batch(); - - info!("scheduler ready (max_batch={})", max_batch); - - loop { - if inflight_prefill - .as_mut() - .is_some_and(|prefill| prefill.output.is_ready()) - { - let (prefill_tokens, prefill_reqs) = - inflight_prefill.as_ref().map_or((0, 0), |prefill| { - ( - prefill.chunk.windows.iter().map(Vec::len).sum(), - prefill.chunk.reqs.len(), - ) - }); - let decode_n = active.len(); - let step_start = itl_debug_enabled().then(Instant::now); - let finish_result = finish_async_prefill( - &mut backend, - &mut active, - &mut prefilling, - inflight_prefill - .take() - .expect("ready async prefill must still be present"), - ); - log_itl_step( - step_start, - "overlap_complete", - prefill_tokens, - prefill_reqs, - decode_n, - ); - if let Err(failure) = finish_result { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - } - - // 1. Merge deferred work with every submission currently available. - let mut pending = std::mem::take(&mut deferred); - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - pending.push(req); - } - - // 2. Remove closed work before metrics, admission, or planning. Active - // and prefilling cleanup goes through the backend's normal retirement - // paths so graph slots and TP request state are released consistently. - if let Err(failure) = - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) - { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - pending, - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - reject_unsupported_echo(&mut pending); - - // 3. Publish the settled post-prune state. Requests accepted from the - // channel are waiting until admission below; closed requests never - // appear in this snapshot or consume its KV/slot accounting. - publish_load( - &load_tx, - &backend, - &active, - &prefilling, - inflight_prefill - .as_ref() - .map_or(0, |prefill| prefill.chunk.reqs.len()), - pending.len(), - ); - - // 4. Nothing in flight and nothing pending: the idle snapshot above is - // already visible, so block until work arrives. Drain and prune again - // after wakeup because the first request may already be closed and more - // submissions may have raced with the blocking receive. - if should_block_on_submit( - active.is_empty() && prefilling.is_empty() && pending.is_empty(), - inflight_prefill.is_some(), - ) { - if let Some((req, _kv_prefix)) = submit_rx.blocking_recv() { - pending.push(req); - } else { - info!("scheduler: all handles dropped, exiting"); - return; - } - while let Ok((req, _kv_prefix)) = submit_rx.try_recv() { - pending.push(req); - } - if let Err(failure) = - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) - { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - pending, - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - reject_unsupported_echo(&mut pending); - publish_load(&load_tx, &backend, &active, &prefilling, 0, pending.len()); - if pending.is_empty() { - continue; - } - } - - // One async prefill owns its scheduled request state. Do not admit or - // launch a second chunk until it resolves. Active decode keeps moving; - // if it retires first, wait on the event instead of blocking on submit. - if inflight_prefill.is_some() { - deferred = pending; - let itl_step_start = itl_debug_enabled().then(Instant::now); - let (itl_prefill_tokens, itl_prefill_reqs) = - inflight_prefill.as_ref().map_or((0, 0), |prefill| { - ( - prefill.chunk.windows.iter().map(Vec::len).sum(), - prefill.chunk.reqs.len(), - ) - }); - let itl_decode_n = active.len(); - let (itl_plan_kind, step_result) = if active.is_empty() { - let result = finish_async_prefill( - &mut backend, - &mut active, - &mut prefilling, - inflight_prefill - .take() - .expect("async prefill must be present before blocking wait"), - ); - ("overlap_wait", result) - } else { - let result = decode_step(&mut backend, &mut active, &mut rng); - ("overlap_decode", result) - }; - log_itl_step( - itl_step_start, - itl_plan_kind, - itl_prefill_tokens, - itl_prefill_reqs, - itl_decode_n, - ); - if let Err(failure) = step_result { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - continue; - } - - // 5. Admit new prompts. In-flight prefills reserve their promotion slot - // and future KV growth, so shrink the slot/page budgets accordingly - let active_budget: Vec = active - .iter() - .map(|req| ActiveKvBudget { - prompt_len: req.prompt_len, - generated_count: req.generated_count, - max_tokens: req.max_tokens, - }) - .collect(); - let page_size = backend.page_size(); - let prefilling_budget: Vec = prefilling - .iter() - .map(|p| PrefillKvBudget { - current_tokens: p.cursor, - prompt_len: p.req.prompt_tokens.len(), - max_tokens: p.req.max_tokens, - }) - .collect(); - let page_budget = backend - .available_pages(&active, &prefilling) - .saturating_sub(prefilling_future_pages(&prefilling_budget, page_size)); - let decode_batching_slot = max_batch.saturating_sub(prefilling.len()); - let admission = admit_pending_requests( - pending, - &active_budget, - decode_batching_slot, - page_size, - page_budget, - // KvPool capacity includes the CUDA Graph padding page reserved at - // construction, so a real request can use at most the remaining pages. - backend.capacity_pages_for_requests(), - backend.max_position_embeddings(), - |req| req.prompt_tokens.len(), - |req| req.max_tokens, - ); - for (rejected, reason) in &admission.rejected { - send_rejection(rejected, *reason); - } - - // 6. Move freshly admitted prompts into the chunked-prefill queue. - for req in admission.pending { - debug!( - "request admitted: request_id={:?} prompt_len={} max_tokens={}", - req.request_id, - req.prompt_tokens.len(), - req.max_tokens - ); - match backend.alloc_prefill_state() { - Ok(backend_state) => prefilling.push(PrefillingRequest35 { - backend_state, - cursor: 0, - step_chunk: 0, - req, - }), - Err(e) => { - warn!("failed to allocate recurrent state for new request: {e}"); - let _ = req.token_tx.send(TokenEvent::Error { - message: e.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - } - } - } - - deferred = admission.deferred; - - // 7. Choose this tick's prefill budget, take that chunk off the front of - // the queue, then dispatch by plan. Auto can return 0 for a short - // decode-priority tick; the next iteration reconsiders the same FIFO - // prefill without reordering it. - let active_decode: Vec = active - .iter() - .map(|req| ActiveDecodeState { - generated_count: req.generated_count, - max_tokens: req.max_tokens, - }) - .collect(); - let prefill_queue: Vec = prefilling - .iter() - .map(|req| PrefillQueueState { - remaining_tokens: req.req.prompt_tokens.len().saturating_sub(req.cursor), - }) - .collect(); - let step_prefill_budget = choose_prefill_budget( - scheduler_policy, - prefill_budget, - &active_decode, - &prefill_queue, - ); - let scheduled = take_prefill_chunks(&mut prefilling, step_prefill_budget); - // ITL diagnostics (#470): capture the *actual* prefill-chunk token count - // and the frozen decode width for this step before the plan consumes the - // scheduled set. Off unless PEGAINFER_ITL_DEBUG is set. - let itl_debug = itl_debug_enabled(); - let itl_prefill_tokens: usize = scheduled.iter().map(|p| p.step_chunk).sum(); - let itl_prefill_reqs = scheduled.len(); - let itl_decode_n = active.len(); - let plan = plan::build_next_plan(!active.is_empty(), scheduled); - if let Some(plan) = plan { - let itl_plan_kind = match &plan { - ExecutionPlan::Unified { .. } if matches!(&backend, SchedulerBackend::Single(single) if single.overlap_enabled()) => { - "overlap_launch" - } - ExecutionPlan::Unified { .. } => "unified", - ExecutionPlan::Prefill { .. } => "prefill", - ExecutionPlan::Decode => "decode", - }; - let itl_step_start = itl_debug.then(Instant::now); - let step_result = match plan { - ExecutionPlan::Unified { pending } => { - if matches!(&backend, SchedulerBackend::Single(single) if single.overlap_enabled()) - { - launch_overlap_step( - &mut backend, - &mut active, - pending, - &mut inflight_prefill, - &mut rng, - ) - } else { - unified_step_sched( - &mut backend, - &mut active, - pending, - &mut prefilling, - &mut rng, - ) - } - } - ExecutionPlan::Prefill { pending } => prefill_batch( - &mut backend, - &mut active, - pending, - &mut prefilling, - &mut rng, - ), - ExecutionPlan::Decode => decode_step(&mut backend, &mut active, &mut rng), - }; - log_itl_step( - itl_step_start, - itl_plan_kind, - itl_prefill_tokens, - itl_prefill_reqs, - itl_decode_n, - ); - if let Err(failure) = step_result { - let kv_total_blocks = backend.capacity_pages_for_requests() as u64; - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - kv_total_blocks, - active, - prefilling, - Vec::new(), - deferred, - inflight_prefill.take(), - failure, - ); - return; - } - } - } -} - -fn send_rejection(req: &SchedulerRequest, reason: RejectReason) { - let message = match reason { - RejectReason::ContextLength { limit } => format!( - "request exceeds this model's maximum context length of {limit} tokens: requested {} (prompt={} + max_tokens={})", - req.prompt_tokens.len().saturating_add(req.max_tokens), - req.prompt_tokens.len(), - req.max_tokens - ), - RejectReason::KvBudget => { - let max_request_tokens = max_kv_tokens(req.prompt_tokens.len(), req.max_tokens); - format!( - "request requires more KV pages than this model instance can provide: prompt_tokens={}, max_request_tokens={max_request_tokens}", - req.prompt_tokens.len() - ) - } - }; - let _ = req.token_tx.send(TokenEvent::Rejected { - message, - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); -} - -// ── Batch prefill ─────────────────────────────────────────────────────── - -fn prefill_batch( - backend: &mut SchedulerBackend, - active: &mut Vec, - scheduled: Vec, - prefilling: &mut Vec, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - let mut chunk = ScheduledChunk::from(scheduled); - let sample_seed = rand::RngExt::random(rng); - let artifacts = match backend { - SchedulerBackend::Single(single) => { - // Scope the borrows of `chunk` to the executor call so the error path can - // move `chunk` into `fail_chunk`. - let logits = match single.batch_prefill_logits(&mut chunk) { - Ok(v) => v, - Err(e) => { - warn!("batch prefill failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return Ok(()); - } - }; - let prefill_sample_seed = rand::RngExt::random(rng); - match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { - Ok((tokens, logprobs)) => PrefillStepArtifacts::Single { tokens, logprobs }, - Err(e) => { - warn!("prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return Ok(()); - } - } - } - SchedulerBackend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { - Ok(v) => PrefillStepArtifacts::Tp(v), - Err(e) => { - warn!("TP prefill chunk failed: {e}"); - return Err(FatalSchedulerError::new(e.to_string()).with_requests(chunk.reqs)); - } - }, - }; - - promote_or_requeue(backend, active, prefilling, chunk, &artifacts) -} - -fn launch_overlap_step( - backend: &mut SchedulerBackend, - active: &mut Vec, - scheduled: Vec, - inflight_prefill: &mut Option, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - debug_assert!(inflight_prefill.is_none()); - let mut chunk = ScheduledChunk::from(scheduled); - let decode_seed = rand::RngExt::random(rng); - let prefill_seed = rand::RngExt::random(rng); - let output = match backend { - SchedulerBackend::Single(single) => single.launch_async_prefill(&mut chunk), - SchedulerBackend::Tp(_) => unreachable!("Qwen3.5 TP cannot launch async prefill"), - }; - match output { - Ok(output) => { - *inflight_prefill = Some(InflightPrefill { - chunk, - output, - sample_seed: prefill_seed, - }); - } - Err(err) => { - warn!("async prefill launch failed: {err}"); - fail_chunk(chunk, &err.to_string()); - } - } - decode_step_with_seed(backend, active, decode_seed) -} - -fn finish_async_prefill( - backend: &mut SchedulerBackend, - active: &mut Vec, - prefilling: &mut Vec, - inflight: InflightPrefill, -) -> std::result::Result<(), FatalSchedulerError> { - let InflightPrefill { - chunk, - output, - sample_seed, - } = inflight; - let logits = output.into_logits(); - let SchedulerBackend::Single(single) = backend else { - unreachable!("Qwen3.5 TP cannot finish async prefill"); - }; - let (tokens, logprobs) = match single.sample_prefill_logits(&chunk.reqs, &logits, sample_seed) { - Ok(result) => result, - Err(err) => { - warn!("async prefill sampling failed: {err}"); - fail_chunk(chunk, &err.to_string()); - return Ok(()); - } - }; - let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; - promote_or_requeue(single, active, prefilling, chunk, &artifacts) -} - -// ── Unified step (prefill chunk + decode in one forward pass) ────────────── - -fn unified_step_sched( - backend: &mut SchedulerBackend, - active: &mut Vec, - scheduled: Vec, - prefilling: &mut Vec, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - let mut chunk = ScheduledChunk::from(scheduled); - if matches!(backend, SchedulerBackend::Tp(_)) { - // Preserve the established scheduler RNG order: decode seed first, - // prefill seed second. Workers execute the forwards in the opposite - // (prefill-then-decode) order using these preselected seeds. - let decode_sample_seed = rand::RngExt::random(rng); - let prefill_sample_seed = rand::RngExt::random(rng); - let result = { - let SchedulerBackend::Tp(tp) = backend else { - unreachable!() - }; - tp.execute_unified(&chunk, active, decode_sample_seed, prefill_sample_seed) - }; - let artifacts = match result { - Ok(artifacts) => artifacts, - Err(err) => { - warn!("TP unified step failed: {err}"); - return Err(FatalSchedulerError::new(err.to_string()).with_requests(chunk.reqs)); - } - }; - - let (decode_tokens, decode_logprobs) = split_decode_artifacts(&artifacts.decode); - if let Err(failure) = - dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs) - { - return Err(failure.with_requests(chunk.reqs)); - } - - let prefill = PrefillStepArtifacts::Tp(artifacts.prefill); - return promote_or_requeue(backend, active, prefilling, chunk, &prefill); - } - - let SchedulerBackend::Single(backend) = backend else { - unreachable!() - }; - // Scope the borrows of `chunk` / `active` to the executor call so the error - // and decode-processing paths can use them afterwards. - let result = backend.unified_step(&mut chunk, active); - let output = match result { - Ok(v) => v, - Err(e) => { - warn!("unified step failed: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - fail_chunk(chunk, &message); - return Ok(()); - } - }; - let decode_seed = rand::RngExt::random(rng); - let prefill_seed = rand::RngExt::random(rng); - - // Process decode results FIRST (it may retire requests and free graph slots - // that promotion then fills densely). - if output.decoded { - process_decode_logits(backend, active, decode_seed)?; - } - - let prefill_logits = output - .prefill_logits - .as_ref() - .expect("scheduled prefill chunk must return prefill logits"); - let (tokens, logprobs) = - match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, prefill_seed) { - Ok(v) => v, - Err(e) => { - warn!("unified prefill sampling failed: {e}"); - fail_chunk(chunk, &e.to_string()); - return Ok(()); - } - }; - let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; - promote_or_requeue(backend, active, prefilling, chunk, &prefill) -} - -// ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── - -fn decode_step( - backend: &mut SchedulerBackend, - active: &mut Vec, - rng: &mut StdRng, -) -> std::result::Result<(), FatalSchedulerError> { - // Preserve the historical scheduler RNG sequence: TP consumes the first - // seed, while single-GPU decode consumed a second seed inside sampling. - let first_seed = rand::RngExt::random(rng); - let sample_seed = if matches!(backend, SchedulerBackend::Single(_)) { - rand::RngExt::random(rng) - } else { - first_seed - }; - decode_step_with_seed(backend, active, sample_seed) -} - -fn decode_step_with_seed( - backend: &mut SchedulerBackend, - active: &mut Vec, - sample_seed: u64, -) -> std::result::Result<(), FatalSchedulerError> { - let (tokens, logprobs_vec) = match backend { - SchedulerBackend::Single(single) => { - if let Err(e) = single.decode_graph(active) { - warn!("batch_decode_graph error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return Ok(()); - } - // Snapshot logits to CPU BEFORE sampling (sampling may modify bufs.logits) - match single.sample_decode_logits(active, sample_seed) { - Ok(v) => v, - Err(e) => { - warn!("decode sampling/logprobs error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return Ok(()); - } - } - } - SchedulerBackend::Tp(tp) => match tp.execute_decode(active, sample_seed) { - Ok(v) => split_decode_artifacts(&v), - Err(e) => { - warn!("TP eager decode error: {e}"); - return Err(FatalSchedulerError::new(e.to_string())); - } - }, - }; - - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) -} - -/// Process decode logits from unified step: sample, extract logprobs, dispatch. -fn process_decode_logits( - backend: &mut SingleGpuBackend, - active: &mut Vec, - sample_seed: u64, -) -> std::result::Result<(), FatalSchedulerError> { - let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, sample_seed) { - Ok(v) => v, - Err(e) => { - warn!("decode sampling/logprobs error: {e}"); - let message = e.to_string(); - for req in active.drain(..) { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.clone(), - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }); - } - return Ok(()); - } - }; - - dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec) -} - -/// Dispatch sampled decode tokens: send events, check EOS/limits, retire finished. -/// -/// `tokens` and `logprobs` are indexed by original position in `active`. -/// Retirements collected first, then compacted in reverse order. -fn dispatch_decode_tokens( - backend: &mut impl DecodeDispatchBackend, - active: &mut Vec, - tokens: &[u32], - logprobs: &[Option], -) -> std::result::Result<(), FatalSchedulerError> { - enum Retirement { - Completion(Vec), - CleanupOnly, - Disconnected, - } - - let n = active.len(); - let mut to_retire = Vec::new(); - - for i in 0..n { - let token = tokens[i]; - let logprob = logprobs[i].clone(); - let req = &mut active[i]; - req.generated_count += 1; - - let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); - let at_limit = req.generated_count >= req.max_tokens; - - if is_eos { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - req.prompt_len, - req.generated_count, - FinishReason::Stop - ); - let event = TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }; - if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(vec![event]))); - } else { - let _ = req.token_tx.send(event); - to_retire.push((i, Retirement::CleanupOnly)); - } - } else if at_limit { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - req.prompt_len, - req.generated_count, - FinishReason::Length - ); - let events = vec![ - TokenEvent::Token { id: token, logprob }, - TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: req.prompt_len, - completion_tokens: req.generated_count, - }, - ]; - if backend.completion_requires_drop_ack() { - to_retire.push((i, Retirement::Completion(events))); - } else { - for event in events { - let _ = req.token_tx.send(event); - } - to_retire.push((i, Retirement::CleanupOnly)); - } - } else if req - .token_tx - .send(TokenEvent::Token { id: token, logprob }) - .is_err() - { - debug!( - "request dropped: client disconnected: request_id={:?} tokens_generated={}", - req.request_id, req.generated_count - ); - to_retire.push((i, Retirement::Disconnected)); - } else { - req.last_token = token; - } - } - - // Remove in reverse order so compact_slot indices stay valid - for (i, retirement) in to_retire.into_iter().rev() { - let request = backend.take_active_request(active, i); - match retirement { - Retirement::Completion(final_events) => { - let candidate = CompletionCandidate { - request, - final_events, - }; - if let Err(err) = backend.drop_active_state(&candidate.request.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()) - .with_request(candidate.into_terminal())); - } - candidate.commit(); - } - Retirement::CleanupOnly | Retirement::Disconnected => { - if let Err(err) = backend.drop_active_state(&request.backend_state) { - return Err(FatalSchedulerError::new(err.to_string()).with_request(request)); - } - } - } - } - Ok(()) -} - -trait DecodeDispatchBackend { - fn is_stop_token(&self, token: u32) -> bool; - fn completion_requires_drop_ack(&self) -> bool; - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35; - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()>; -} - -impl DecodeDispatchBackend for SingleGpuBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn completion_requires_drop_ack(&self) -> bool { - false - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - compact_single_slot(self, active, idx) - } - - fn drop_active_state(&mut self, _state: &ActiveBackendState) -> Result<()> { - Ok(()) - } -} - -impl DecodeDispatchBackend for SchedulerBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn completion_requires_drop_ack(&self) -> bool { - matches!(self, SchedulerBackend::Tp(_)) - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - match self { - SchedulerBackend::Single(backend) => compact_single_slot(backend, active, idx), - SchedulerBackend::Tp(_) => active.swap_remove(idx), - } - } - - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - match (self, state) { - (SchedulerBackend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), - (SchedulerBackend::Tp(backend), ActiveBackendState::Tp { request_id }) => { - backend.drop_request(*request_id, DropExpectation::MustExist) - } - _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), - } - } -} - -/// Remove single-GPU request at `idx` via swap_remove and compact graph slots. -/// -/// After swap_remove, the element that was at `active.len()-1` (before remove) -/// now sits at `idx`. Its graph slot must be copied into the vacated slot so -/// that slots 0..active.len() remain dense. -fn compact_single_slot( - backend: &mut SingleGpuBackend, - active: &mut Vec, - idx: usize, -) -> ActiveRequest35 { - let compaction = compaction_after_retire(active.len(), idx); - let removed = active.swap_remove(idx); - - if let Some(compaction) = compaction { - backend.compact_slot(active, compaction); - } - removed -} - -// ── Chunked-prefill helpers ──────────────────────────────────────────────── - -/// Step's scheduled prefill set -struct ScheduledChunk { - reqs: Vec, - backend_state: ScheduledChunkBackendState, - /// Prompt cursor after this step's chunk - ends: Vec, - /// This step's chunked token slice per request - windows: Vec>, -} - -struct InflightPrefill { - // Fields drop in declaration order. Drain the stream before request state - // can return KV pages or release recurrent/convolution buffers on unwind. - output: AsyncPrefillOutput, - chunk: ScheduledChunk, - sample_seed: u64, -} - -enum ScheduledChunkBackendState { - Single { - kvs: Vec, - recs: Vec, - }, - Tp { - request_ids: Vec, - }, -} - -impl From> for ScheduledChunk { - fn from(scheduled: Vec) -> Self { - let n = scheduled.len(); - let is_tp = scheduled - .first() - .is_some_and(|p| matches!(p.backend_state, PrefillBackendState::Tp { .. })); - let mut chunk = ScheduledChunk { - reqs: Vec::with_capacity(n), - backend_state: if is_tp { - ScheduledChunkBackendState::Tp { - request_ids: Vec::with_capacity(n), - } - } else { - ScheduledChunkBackendState::Single { - kvs: Vec::with_capacity(n), - recs: Vec::with_capacity(n), - } - }, - ends: Vec::with_capacity(n), - windows: Vec::with_capacity(n), - }; - for p in scheduled { - let end = p.cursor + p.step_chunk; - chunk - .windows - .push(p.req.prompt_tokens[p.cursor..end].to_vec()); - chunk.ends.push(end); - chunk.reqs.push(p.req); - match (&mut chunk.backend_state, p.backend_state) { - ( - ScheduledChunkBackendState::Single { kvs, recs }, - PrefillBackendState::Single { kv, rec }, - ) => { - kvs.push(kv); - recs.push(rec); - } - ( - ScheduledChunkBackendState::Tp { request_ids }, - PrefillBackendState::Tp { request_id }, - ) => request_ids.push(request_id), - _ => unreachable!("mixed Qwen3.5 scheduler backend states in one chunk"), - } - } - chunk - } -} - -/// Pull this step's prefill set off the FRONT of `prefilling`, capping the -/// step's total forwarded prompt tokens at `prefill_budget`. -fn take_prefill_chunks( - prefilling: &mut Vec, - prefill_budget: usize, -) -> Vec { - let remaining: Vec = prefilling - .iter() - .map(|p| p.req.prompt_tokens.len() - p.cursor) - .collect(); - let chunks = plan_prefill_chunks(&remaining, prefill_budget); - let mut scheduled: Vec = prefilling.drain(0..chunks.len()).collect(); - for (p, chunk) in scheduled.iter_mut().zip(&chunks) { - p.step_chunk = *chunk; - } - scheduled -} - -/// Report a forward/sampling failure to every request in the failed chunk. -fn fail_chunk(chunk: ScheduledChunk, message: &str) { - for req in chunk.reqs { - let _ = req.token_tx.send(TokenEvent::Error { - message: message.to_string(), - prompt_tokens: req.prompt_tokens.len(), - completion_tokens: 0, - }); - } -} - -/// For each request in the just-prefilled chunk: if its prompt is now exhausted, -/// sample its first token, emit events, and move it into the decode batch; -/// otherwise re-queue it (with an advanced cursor) at the FRONT of `prefilling`. -/// `artifacts` are indexed by request order in `chunk`. -fn promote_or_requeue( - backend: &mut impl PrefillPromoteBackend, - active: &mut Vec, - prefilling: &mut Vec, - chunk: ScheduledChunk, - artifacts: &PrefillStepArtifacts, -) -> std::result::Result<(), FatalSchedulerError> { - let ScheduledChunk { - reqs, - backend_state, - ends, - .. - } = chunk; - let mut still_prefilling: Vec = Vec::new(); - let backend_states = split_scheduled_backend_state(backend_state); - let mut entries: VecDeque<_> = reqs - .into_iter() - .zip(backend_states) - .zip(ends) - .enumerate() - .map(|(i, ((req, backend_state), end))| (i, req, backend_state, end)) - .collect(); - - while let Some((i, req, backend_state, end)) = entries.pop_front() { - // Not finished: re-queue with the advanced cursor - if end < req.prompt_tokens.len() { - still_prefilling.push(PrefillingRequest35 { - req, - backend_state, - cursor: end, - step_chunk: 0, - }); - continue; - } - - let prompt_len = req.prompt_tokens.len(); - let artifact = artifacts.final_artifact(i); - let first_token = artifact.token; - let logprob = artifact.logprob; - - if !req.params.ignore_eos && backend.is_stop_token(first_token) { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 0, - FinishReason::Stop - ); - let candidate = CompletionCandidate { - request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![TokenEvent::Finished { - finish_reason: FinishReason::Stop, - prompt_tokens: prompt_len, - completion_tokens: 0, - }], - }; - if let Err(err) = backend - .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - candidate.into_terminal(), - still_prefilling, - entries, - )); - } - candidate.commit(); - continue; - } - - if req.max_tokens <= 1 { - debug!( - "request finished: request_id={:?} prompt_tokens={} completion_tokens={} finish_reason={:?}", - req.request_id, - prompt_len, - 1, - FinishReason::Length - ); - let candidate = CompletionCandidate { - request: PrefillCompletionRequest { req, backend_state }, - final_events: vec![ - TokenEvent::Token { - id: first_token, - logprob, - }, - TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: prompt_len, - completion_tokens: 1, - }, - ], - }; - if let Err(err) = backend - .drop_prefill_state(&candidate.request.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - candidate.into_terminal(), - still_prefilling, - entries, - )); - } - candidate.commit(); - continue; - } - - if req - .token_tx - .send(TokenEvent::Token { - id: first_token, - logprob, - }) - .is_err() - { - debug!( - "request dropped: client disconnected: request_id={:?} tokens_generated={}", - req.request_id, 0 - ); - let removed = PrefillCompletionRequest { req, backend_state }; - if let Err(err) = - backend.drop_prefill_state(&removed.backend_state, DropExpectation::MustExist) - { - return Err(prefill_lifecycle_failure( - err.to_string(), - removed.into_terminal(), - still_prefilling, - entries, - )); - } - continue; - } - - let active_backend_state = backend.promote_prefill_state(active.len(), backend_state); - active.push(ActiveRequest35 { - request_id: req.request_id, - token_tx: req.token_tx, - backend_state: active_backend_state, - last_token: first_token, - generated_count: 1, - max_tokens: req.max_tokens, - prompt_len, - params: req.params, - logprobs: req.logprobs, - }); - } - - prefilling.splice(0..0, still_prefilling); - Ok(()) -} - -fn prefill_lifecycle_failure( - message: String, - current: TerminalRequest, - still_prefilling: Vec, - remaining: VecDeque<(usize, SchedulerRequest, PrefillBackendState, usize)>, -) -> FatalSchedulerError { - FatalSchedulerError::new(message) - .with_request(current) - .with_requests(still_prefilling) - .with_requests(remaining.into_iter().map(|(_, req, _, _)| req)) -} - -trait PrefillPromoteBackend { - fn is_stop_token(&self, token: u32) -> bool; - fn promote_prefill_state( - &mut self, - active_len: usize, - state: PrefillBackendState, - ) -> ActiveBackendState; - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()>; -} - -impl PrefillPromoteBackend for SingleGpuBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn promote_prefill_state( - &mut self, - active_len: usize, - state: PrefillBackendState, - ) -> ActiveBackendState { - let PrefillBackendState::Single { kv, rec } = state else { - panic!("single-GPU promotion received TP prefill state"); - }; - let slot_idx = slot_for_new_request(active_len, self.max_batch()) - .expect("admission must reserve a graph slot"); - self.copy_recurrent_to_slot(&rec, slot_idx) - .expect("copy recurrent state to slot failed"); - ActiveBackendState::Single { - kv, - graph_slot_idx: slot_idx, - } - } - - fn drop_prefill_state( - &mut self, - _state: &PrefillBackendState, - _expectation: DropExpectation, - ) -> Result<()> { - Ok(()) - } -} - -impl PrefillPromoteBackend for SchedulerBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.is_stop_token(token) - } - - fn promote_prefill_state( - &mut self, - active_len: usize, - state: PrefillBackendState, - ) -> ActiveBackendState { - match (self, state) { - (SchedulerBackend::Single(single), PrefillBackendState::Single { kv, rec }) => { - let slot_idx = slot_for_new_request(active_len, single.max_batch()) - .expect("admission must reserve a graph slot"); - single - .copy_recurrent_to_slot(&rec, slot_idx) - .expect("copy recurrent state to slot failed"); - ActiveBackendState::Single { - kv, - graph_slot_idx: slot_idx, - } - } - (SchedulerBackend::Tp(_), PrefillBackendState::Tp { request_id }) => { - ActiveBackendState::Tp { request_id } - } - _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), - } - } - - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()> { - match (self, state) { - (SchedulerBackend::Single(_), PrefillBackendState::Single { .. }) => Ok(()), - (SchedulerBackend::Tp(backend), PrefillBackendState::Tp { request_id }) => { - backend.drop_request(*request_id, expectation) - } - _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during prefill drop"), - } - } -} - -fn split_scheduled_backend_state( - backend_state: ScheduledChunkBackendState, -) -> Vec { - match backend_state { - ScheduledChunkBackendState::Single { kvs, recs } => kvs - .into_iter() - .zip(recs) - .map(|(kv, rec)| PrefillBackendState::Single { kv, rec }) - .collect(), - ScheduledChunkBackendState::Tp { request_ids } => request_ids - .into_iter() - .map(|request_id| PrefillBackendState::Tp { request_id }) - .collect(), - } -} - -#[cfg(test)] -mod tests; diff --git a/pegainfer-qwen35/src/scheduler/backend.rs b/pegainfer-qwen35/src/scheduler/backend.rs new file mode 100644 index 000000000..174bb6f9f --- /dev/null +++ b/pegainfer-qwen35/src/scheduler/backend.rs @@ -0,0 +1,1152 @@ +//! Single-GPU and TP execution backends south of the ledger. +//! +//! Named [`Qwen35Backend`] so it does not collide with the frontend's +//! `engine::SchedulerBackend` (the driver-side wiring). + +use std::collections::HashMap; +use std::collections::HashSet; +use std::sync::Arc; + +use anyhow::Result; +use cudarc::driver::CudaEvent; +use cudarc::driver::CudaStream; +use cudarc::driver::sys; +use pegainfer_core::kv_pool::KvState; +use pegainfer_core::tensor::HiddenStates; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::RequestId as EngineRequestId; +use pegainfer_frontend::engine::TokenLogprob; +use pegainfer_frontend::sampler::SamplingParams; + +use super::plan; +use super::plan::slot_for_new_request; +use crate::Qwen35DecodeOverlap; +use crate::batch_decode_graph::BatchDecodeGraphState; +use crate::executor::DecodeRequestResult; +use crate::executor::DecodeResult; +use crate::executor::PrefillRequestResult; +use crate::executor::PrefillResult; +use crate::executor::RequestId as BackendRequestId; +use crate::logprobs::snapshot_requested_logprobs; +use crate::recurrent_state::RecurrentState; +use crate::tp_executor::DropExpectation; +use crate::tp_executor::Qwen35TpExecutor; +use crate::tp_executor::TpDecodeStepItem; +use crate::tp_executor::TpPrefillChunkItem; +use crate::tp_executor::TpUnifiedPlan; +use crate::weights::Qwen35Model; + +/// An in-flight request being decoded. Recurrent state lives in the +/// `BatchDecodeGraphState` at `graph_slot_idx` — NOT owned here. +pub(super) struct ActiveRequest35 { + pub(super) id: EngineRequestId, + pub(super) backend_state: ActiveBackendState, + pub(super) last_token: u32, + pub(super) generated_count: usize, + pub(super) max_tokens: usize, + pub(super) prompt_len: usize, + pub(super) params: SamplingParams, + /// Number of top logprobs to return (0 = disabled). + pub(super) logprobs: usize, +} + +/// A request whose prompt is being prefilled across multiple scheduler steps. +pub(super) struct PrefillingRequest35 { + pub(super) id: EngineRequestId, + pub(super) request: Request, + pub(super) backend_state: PrefillBackendState, + /// Prompt tokens prefilled so far. + pub(super) cursor: usize, + /// Tokens to prefill in the step currently scheduled (set by `take_prefill_chunks`). + pub(super) step_chunk: usize, +} + +pub(super) enum ActiveBackendState { + Single { + kv: KvState, + /// Index into `BatchDecodeGraphState.slot_states`. + graph_slot_idx: usize, + }, + Tp { + request_id: BackendRequestId, + }, +} + +pub(super) enum PrefillBackendState { + Single { kv: KvState, rec: RecurrentState }, + Tp { request_id: BackendRequestId }, +} + +#[derive(Clone, Debug, PartialEq)] +pub(super) struct PrefillArtifact { + pub(super) token: u32, + pub(super) logprob: Option, +} + +#[derive(Clone, Debug, PartialEq)] +pub(super) struct DecodeArtifact { + pub(super) token: u32, + pub(super) logprob: Option, +} + +pub(super) struct AlignedUnifiedArtifacts { + pub(super) prefill: Vec>, + pub(super) decode: Vec, +} + +pub(super) enum PrefillStepArtifacts { + Single { + tokens: Vec, + logprobs: Vec>, + }, + Tp(Vec>), +} + +impl PrefillStepArtifacts { + pub(super) fn final_artifact(&self, idx: usize) -> PrefillArtifact { + match self { + Self::Single { tokens, logprobs } => PrefillArtifact { + token: tokens[idx], + logprob: logprobs[idx].clone(), + }, + Self::Tp(artifacts) => artifacts[idx] + .clone() + .expect("validated TP final-prefill row must contain an artifact"), + } + } +} + +pub(super) struct SingleGpuBackend { + model: Qwen35Model, + graph_state: BatchDecodeGraphState, + prefill_stream: Option>, +} + +// One instance per scheduler; the size asymmetry costs nothing here. +#[allow(clippy::large_enum_variant)] +pub(super) enum Qwen35Backend { + Single(SingleGpuBackend), + Tp(TpSchedulerBackend), +} + +pub(super) struct AsyncPrefillOutput { + logits: Option, + done: CudaEvent, + stream: Arc, + completed: bool, +} + +impl AsyncPrefillOutput { + pub(super) fn is_ready(&mut self) -> bool { + match unsafe { sys::cuEventQuery(self.done.cu_event()) } { + sys::CUresult::CUDA_SUCCESS => { + self.completed = true; + true + } + sys::CUresult::CUDA_ERROR_NOT_READY => false, + err => fatal_cuda_lifecycle(&format!( + "query Qwen3.5 async prefill event failed: {err:?}" + )), + } + } + + pub(super) fn into_logits(mut self) -> HiddenStates { + if !self.completed { + if let Err(err) = self.done.synchronize() { + fatal_cuda_lifecycle(&format!("wait for Qwen3.5 async prefill failed: {err}")); + } + self.completed = true; + } + self.logits + .take() + .expect("async prefill logits must be consumed exactly once") + } +} + +impl Drop for AsyncPrefillOutput { + fn drop(&mut self) { + if self.completed { + return; + } + if let Err(err) = self.stream.synchronize() { + fatal_cuda_lifecycle(&format!( + "drain Qwen3.5 async prefill during cleanup failed: {err}" + )); + } + } +} + +fn fatal_cuda_lifecycle(message: &str) -> ! { + log::error!("FATAL: {message}; aborting before CUDA-referenced state is released"); + std::process::abort(); +} + +pub(super) struct TpSchedulerBackend { + executor: Qwen35TpExecutor, + next_request_id: u64, +} + +impl SingleGpuBackend { + pub(super) fn new( + model: Qwen35Model, + max_batch: usize, + decode_overlap: Qwen35DecodeOverlap, + ) -> Result { + anyhow::ensure!(max_batch > 0, "Qwen3.5 max_batch must be > 0"); + let graph_capacity = crate::batch_decode_graph::bucket_for(max_batch); + let graph_state = model.create_batch_decode_graph_state_with_capacity(graph_capacity)?; + let prefill_stream = match decode_overlap { + Qwen35DecodeOverlap::Off => None, + Qwen35DecodeOverlap::SharedSm => Some( + model + .device_ctx() + .ctx + .new_stream() + .map_err(|err| anyhow::anyhow!("create Qwen3.5 prefill stream: {err}"))?, + ), + }; + Ok(Self { + model, + graph_state, + prefill_stream, + }) + } + + pub(super) fn model(&self) -> &Qwen35Model { + &self.model + } + + fn max_batch(&self) -> usize { + // #470: admit the requested `--max-batch`, which may sit below the loaded + // graph bucket (e.g. 5 on bucket 8); never exceed the physical slots. + self.model + .decode_admission_batch + .min(self.graph_state.slot_states.len()) + .max(1) + } + + fn page_size(&self) -> usize { + self.model.kv_pool().layout().page_size + } + + fn available_pages(&self) -> usize { + self.model.kv_pool().available_pages() + } + + fn capacity_pages_for_requests(&self) -> usize { + self.model.kv_pool().capacity_pages().saturating_sub(1) + } + + fn max_position_embeddings(&self) -> usize { + self.model.config().max_position_embeddings + } + + fn alloc_kv(&self) -> KvState { + self.model.alloc_kv() + } + + fn alloc_recurrent(&self) -> Result { + RecurrentState::new(self.model.device_ctx(), self.model.config()) + } + + pub(super) fn batch_prefill_logits(&self, chunk: &mut ScheduledChunk) -> Result { + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU prefill received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + self.model + .batch_prefill_logits(&window_refs, kvs, &mut rec_refs) + } + + pub(super) fn overlap_enabled(&self) -> bool { + self.prefill_stream.is_some() + } + + pub(super) fn launch_async_prefill( + &mut self, + chunk: &mut ScheduledChunk, + ) -> Result { + let prefill_stream = self + .prefill_stream + .clone() + .ok_or_else(|| anyhow::anyhow!("Qwen3.5 decode overlap is disabled"))?; + + // Request KV/recurrent state was allocated on the model stream. Order + // those producers before the prefill stream without blocking the host. + prefill_stream + .join(&self.model.device_ctx().stream) + .map_err(|err| anyhow::anyhow!("join Qwen3.5 prefill stream: {err}"))?; + + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU async prefill received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + let logits = match self.model.batch_prefill_logits_on_stream( + Arc::clone(&prefill_stream), + &window_refs, + kvs, + &mut rec_refs, + ) { + Ok(logits) => logits, + Err(err) => { + if let Err(sync_err) = prefill_stream.synchronize() { + fatal_cuda_lifecycle(&format!( + "Qwen3.5 async prefill failed ({err}); stream drain failed: {sync_err}" + )); + } + return Err(err); + } + }; + let done = match prefill_stream.record_event(None) { + Ok(done) => done, + Err(err) => { + if let Err(sync_err) = prefill_stream.synchronize() { + fatal_cuda_lifecycle(&format!( + "record Qwen3.5 async prefill event failed ({err}); stream drain failed: {sync_err}" + )); + } + return Err(anyhow::anyhow!("record Qwen3.5 async prefill event: {err}")); + } + }; + Ok(AsyncPrefillOutput { + logits: Some(logits), + done, + stream: prefill_stream, + completed: false, + }) + } + + pub(super) fn unified_step( + &mut self, + chunk: &mut ScheduledChunk, + active: &mut [ActiveRequest35], + ) -> Result { + let window_refs: Vec<&[u32]> = chunk.windows.iter().map(Vec::as_slice).collect(); + let ScheduledChunkBackendState::Single { kvs, recs } = &mut chunk.backend_state else { + anyhow::bail!("single-GPU unified step received TP chunk state"); + }; + let mut rec_refs: Vec<&mut RecurrentState> = recs.iter_mut().collect(); + let decode_tokens: Vec = active.iter().map(|r| r.last_token).collect(); + let mut decode_kv_refs: Vec<&mut KvState> = active + .iter_mut() + .map(|r| match &mut r.backend_state { + ActiveBackendState::Single { kv, .. } => kv, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU unified step received TP active state") + } + }) + .collect(); + self.model.unified_step( + &window_refs, + kvs, + &mut rec_refs, + &decode_tokens, + &mut decode_kv_refs, + &mut self.graph_state, + ) + } + + pub(super) fn decode_graph(&mut self, active: &mut [ActiveRequest35]) -> Result<()> { + let token_ids: Vec = active.iter().map(|r| r.last_token).collect(); + let mut kv_refs: Vec<&mut KvState> = active + .iter_mut() + .map(|r| match &mut r.backend_state { + ActiveBackendState::Single { kv, .. } => kv, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU decode received TP active state") + } + }) + .collect(); + self.model + .batch_decode_graph(&token_ids, &mut kv_refs, &mut self.graph_state) + } + + pub(super) fn sample_prefill_logits( + &mut self, + pending: &[Request], + logits: &HiddenStates, + sample_seed: u64, + ) -> Result<(Vec, Vec>)> { + debug_assert_eq!( + logits.seq_len, + pending.len(), + "Qwen3.5 prefill logits rows must preserve pending request order" + ); + let requested_logprobs: Vec = pending.iter().map(|r| r.logprobs).collect(); + let cpu_logits = + snapshot_requested_logprobs(self.model.device_ctx(), logits, &requested_logprobs)?; + let params_refs: Vec<&SamplingParams> = pending.iter().map(|r| &r.params).collect(); + let tokens = self.model.select_tokens_from_logits_varied( + logits, + &mut self.graph_state.buffers, + ¶ms_refs, + sample_seed, + )?; + + let logprobs = cpu_logits + .into_iter() + .enumerate() + .map(|(i, logits_opt)| { + logits_opt.and_then(|logits_f32| { + pegainfer_sample::token_logprob_from_row( + &logits_f32, + tokens[i], + pending[i].logprobs, + ) + }) + }) + .collect(); + Ok((tokens, logprobs)) + } + + pub(super) fn sample_decode_logits( + &mut self, + active: &[ActiveRequest35], + sample_seed: u64, + ) -> Result<(Vec, Vec>)> { + let requested_logprobs: Vec = active.iter().map(|r| r.logprobs).collect(); + let cpu_logits = snapshot_requested_logprobs( + self.model.device_ctx(), + &self.graph_state.buffers.logits, + &requested_logprobs, + )?; + let params_refs: Vec<&SamplingParams> = active.iter().map(|r| &r.params).collect(); + let tokens = self.model.select_tokens_batch_varied( + &mut self.graph_state.buffers, + ¶ms_refs, + sample_seed, + )?; + + let logprobs = cpu_logits + .into_iter() + .enumerate() + .map(|(i, logits_opt)| { + logits_opt.and_then(|logits_f32| { + pegainfer_sample::token_logprob_from_row( + &logits_f32, + tokens[i], + active[i].logprobs, + ) + }) + }) + .collect(); + Ok((tokens, logprobs)) + } + + fn is_stop_token(&self, token: u32) -> bool { + self.model.is_stop_token(token) + } + + fn copy_recurrent_to_slot( + &mut self, + recurrent: &RecurrentState, + slot_idx: usize, + ) -> Result<()> { + self.graph_state + .copy_state_to_slot(self.model.device_ctx(), recurrent, slot_idx) + } + + fn compact_slot(&mut self, active: &mut [ActiveRequest35], compaction: plan::SlotCompaction) { + let src_slot = match active[compaction.moved_to].backend_state { + ActiveBackendState::Single { graph_slot_idx, .. } => graph_slot_idx, + ActiveBackendState::Tp { .. } => { + panic!("single-GPU slot compaction received TP active state") + } + }; + debug_assert_eq!(src_slot, compaction.moved_from); + + let ctx = self.model.device_ctx(); + let src = &self.graph_state.slot_states[compaction.moved_from]; + for layer_idx in 0..src.layers.len() { + let (src_part, dst_part) = if compaction.moved_to < compaction.moved_from { + let (left, right) = self + .graph_state + .slot_states + .split_at_mut(compaction.moved_from); + ( + &right[0].layers[layer_idx], + &mut left[compaction.moved_to].layers[layer_idx], + ) + } else { + unreachable!("idx < active.len() <= last"); + }; + + ctx.stream + .memcpy_dtod(&src_part.state, &mut dst_part.state) + .expect("compact slot state copy failed"); + ctx.stream + .memcpy_dtod(&src_part.conv_state.data, &mut dst_part.conv_state.data) + .expect("compact slot conv_state copy failed"); + } + self.graph_state.slot_states[compaction.moved_to].seq_len = + self.graph_state.slot_states[compaction.moved_from].seq_len; + + match &mut active[compaction.moved_to].backend_state { + ActiveBackendState::Single { graph_slot_idx, .. } => { + *graph_slot_idx = compaction.moved_to; + } + ActiveBackendState::Tp { .. } => { + panic!("single-GPU slot compaction received TP active state") + } + } + } +} + +impl TpSchedulerBackend { + pub(super) fn new( + model_path: &str, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, + ) -> Result { + let executor = Qwen35TpExecutor::from_runtime_with_limits( + model_path, + false, + device_ordinals, + max_batch, + max_prefill_tokens, + )?; + Ok(Self { + executor, + next_request_id: 1, + }) + } + + fn alloc_request_id(&mut self) -> BackendRequestId { + let id = BackendRequestId::new(self.next_request_id); + self.next_request_id = self.next_request_id.wrapping_add(1).max(1); + id + } + + fn max_batch(&self) -> usize { + self.executor.max_batch() + } + + pub(super) fn page_size(&self) -> usize { + self.executor.page_size() + } + + pub(super) fn capacity_pages_for_requests(&self) -> usize { + self.executor.capacity_pages_for_requests() + } + + pub(super) fn max_position_embeddings(&self) -> usize { + self.executor.max_position_embeddings() + } + + fn is_stop_token(&self, token: u32) -> bool { + self.executor.is_stop_token(token) + } + + fn available_pages( + &self, + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + ) -> usize { + let page_size = self.page_size(); + let active_pages: usize = active + .iter() + .map(|req| pages_needed(current_active_tokens(req), page_size)) + .sum(); + let prefilling_pages: usize = prefilling + .iter() + .map(|req| pages_needed(req.cursor, page_size)) + .sum(); + self.capacity_pages_for_requests() + .saturating_sub(active_pages.saturating_add(prefilling_pages)) + } + + pub(super) fn execute_prefill_chunk( + &self, + chunk: &ScheduledChunk, + sample_seed: u64, + ) -> Result>> { + let items = tp_prefill_items(chunk)?; + let result = self + .executor + .execute_prefill_chunks_with_seed(&items, sample_seed)?; + align_prefill_results(chunk, &result) + .map_err(|err| self.executor.poison_artifact_contract("prefill", &err)) + } + + pub(super) fn execute_decode( + &self, + active: &[ActiveRequest35], + sample_seed: u64, + ) -> Result> { + let items = tp_decode_items(active)?; + let result = self.executor.execute_decode_items(&items, sample_seed)?; + align_decode_results(active, &result) + .map_err(|err| self.executor.poison_artifact_contract("decode", &err)) + } + + pub(super) fn execute_unified( + &self, + chunk: &ScheduledChunk, + active: &[ActiveRequest35], + decode_sample_seed: u64, + prefill_sample_seed: u64, + ) -> Result { + let plan = TpUnifiedPlan { + prefill: tp_prefill_items(chunk)?, + decode: tp_decode_items(active)?, + prefill_sample_seed, + decode_sample_seed, + }; + let result = self.executor.execute_unified(&plan)?; + let prefill = align_prefill_results(chunk, &result.prefill).map_err(|err| { + self.executor + .poison_artifact_contract("unified prefill", &err) + })?; + let decode = align_decode_results(active, &result.decode).map_err(|err| { + self.executor + .poison_artifact_contract("unified decode", &err) + })?; + Ok(AlignedUnifiedArtifacts { prefill, decode }) + } + + fn drop_request( + &self, + request_id: BackendRequestId, + expectation: DropExpectation, + ) -> Result<()> { + self.executor.drop_request(request_id, expectation) + } +} + +impl Qwen35Backend { + pub(super) fn max_batch(&self) -> usize { + match self { + Self::Single(backend) => backend.max_batch(), + Self::Tp(backend) => backend.max_batch(), + } + } + + pub(super) fn page_size(&self) -> usize { + match self { + Self::Single(backend) => backend.page_size(), + Self::Tp(backend) => backend.page_size(), + } + } + + pub(super) fn available_pages( + &self, + active: &[ActiveRequest35], + prefilling: &[PrefillingRequest35], + ) -> usize { + match self { + Self::Single(backend) => backend.available_pages(), + Self::Tp(backend) => backend.available_pages(active, prefilling), + } + } + + pub(super) fn capacity_pages_for_requests(&self) -> usize { + match self { + Self::Single(backend) => backend.capacity_pages_for_requests(), + Self::Tp(backend) => backend.capacity_pages_for_requests(), + } + } + + pub(super) fn max_position_embeddings(&self) -> usize { + match self { + Self::Single(backend) => backend.max_position_embeddings(), + Self::Tp(backend) => backend.max_position_embeddings(), + } + } + + pub(super) fn alloc_prefill_state(&mut self) -> Result { + match self { + Self::Single(backend) => Ok(PrefillBackendState::Single { + kv: backend.alloc_kv(), + rec: backend.alloc_recurrent()?, + }), + Self::Tp(backend) => Ok(PrefillBackendState::Tp { + request_id: backend.alloc_request_id(), + }), + } + } + + pub(super) fn is_stop_token(&self, token: u32) -> bool { + match self { + Self::Single(backend) => backend.is_stop_token(token), + Self::Tp(backend) => backend.is_stop_token(token), + } + } +} + +fn current_active_tokens(req: &ActiveRequest35) -> usize { + req.prompt_len + .saturating_add(req.generated_count.saturating_sub(1)) +} + +fn pages_needed(token_count: usize, page_size: usize) -> usize { + token_count.div_ceil(page_size) +} + +fn tp_prefill_items(chunk: &ScheduledChunk) -> Result> { + let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { + anyhow::bail!("TP prefill received single-GPU chunk state"); + }; + anyhow::ensure!( + chunk.reqs.len() == request_ids.len() + && chunk.reqs.len() == chunk.windows.len() + && chunk.reqs.len() == chunk.ends.len(), + "Qwen3.5 TP scheduled prefill vectors are misaligned" + ); + Ok(chunk + .reqs + .iter() + .zip(request_ids) + .zip(&chunk.windows) + .zip(&chunk.ends) + .map(|(((req, request_id), window), end)| { + TpPrefillChunkItem::new_with_sampling( + *request_id, + window.clone(), + req.logprobs, + req.params, + *end == req.prompt_tokens.len(), + ) + }) + .collect()) +} + +fn tp_decode_items(active: &[ActiveRequest35]) -> Result> { + active + .iter() + .map(|req| { + let ActiveBackendState::Tp { request_id } = &req.backend_state else { + anyhow::bail!("TP decode received single-GPU active state"); + }; + Ok(TpDecodeStepItem::new( + *request_id, + req.last_token, + req.logprobs, + req.params, + )) + }) + .collect() +} + +fn align_prefill_results( + chunk: &ScheduledChunk, + result: &PrefillResult, +) -> Result>> { + let ScheduledChunkBackendState::Tp { request_ids } = &chunk.backend_state else { + anyhow::bail!("align_prefill_results requires TP chunk state"); + }; + anyhow::ensure!( + request_ids.len() == chunk.reqs.len() && chunk.ends.len() == chunk.reqs.len(), + "Qwen3.5 TP prefill alignment vectors are misaligned" + ); + let expected: HashSet = request_ids + .iter() + .zip(&chunk.reqs) + .zip(&chunk.ends) + .filter_map(|((&request_id, req), &end)| { + (end == req.prompt_tokens.len()).then_some(request_id) + }) + .collect(); + let mut by_id = HashMap::with_capacity(result.requests.len()); + for PrefillRequestResult { + request_id, + first_token, + first_token_logprob, + } in &result.requests + { + anyhow::ensure!( + expected.contains(request_id), + "Qwen3.5 TP prefill returned unknown or non-final request id {}", + request_id.get() + ); + let artifact = PrefillArtifact { + token: *first_token, + logprob: first_token_logprob.clone(), + }; + anyhow::ensure!( + by_id.insert(*request_id, artifact).is_none(), + "Qwen3.5 TP prefill returned duplicate request id {}", + request_id.get() + ); + } + anyhow::ensure!( + by_id.len() == expected.len(), + "Qwen3.5 TP prefill result is missing final request IDs" + ); + + request_ids + .iter() + .zip(&chunk.reqs) + .zip(&chunk.ends) + .map(|((&request_id, req), &end)| { + if end == req.prompt_tokens.len() { + by_id.remove(&request_id).map(Some).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP prefill result is missing final request id {}", + request_id.get() + ) + }) + } else { + Ok(None) + } + }) + .collect() +} + +fn align_decode_results( + active: &[ActiveRequest35], + result: &DecodeResult, +) -> Result> { + let expected: Vec = active + .iter() + .map(|active_req| { + let ActiveBackendState::Tp { request_id } = active_req.backend_state else { + anyhow::bail!("align_decode_results requires TP active state"); + }; + Ok(request_id) + }) + .collect::>()?; + let expected_set: HashSet<_> = expected.iter().copied().collect(); + anyhow::ensure!( + expected_set.len() == expected.len(), + "Qwen3.5 TP active decode IDs contain duplicates" + ); + let mut by_id = HashMap::with_capacity(result.requests.len()); + for DecodeRequestResult { + request_id, + token, + logprob, + } in &result.requests + { + anyhow::ensure!( + expected_set.contains(request_id), + "Qwen3.5 TP decode returned unknown request id {}", + request_id.get() + ); + let artifact = DecodeArtifact { + token: *token, + logprob: logprob.clone(), + }; + anyhow::ensure!( + by_id.insert(*request_id, artifact).is_none(), + "Qwen3.5 TP decode returned duplicate request id {}", + request_id.get() + ); + } + expected + .into_iter() + .map(|request_id| { + by_id.remove(&request_id).ok_or_else(|| { + anyhow::anyhow!( + "Qwen3.5 TP decode result is missing request id {}", + request_id.get() + ) + }) + }) + .collect() +} + +pub(super) fn split_decode_artifacts( + artifacts: &[DecodeArtifact], +) -> (Vec, Vec>) { + artifacts + .iter() + .map(|artifact| (artifact.token, artifact.logprob.clone())) + .unzip() +} + +pub(super) fn servable_len(max_context: usize, max_pages: usize, page_size: usize) -> u32 { + max_context + .min(max_pages.saturating_mul(page_size)) + .try_into() + .unwrap_or(u32::MAX) +} + +pub(super) struct CublasThreadGuard; + +impl Drop for CublasThreadGuard { + fn drop(&mut self) { + unsafe { + crate::ffi::cublas_destroy(); + } + } +} + +/// Bind the CUDA context and init thread-local cuBLAS on the scheduler +/// (driver) thread. Must run on first `step`, not the load thread. +pub(super) fn bind_model_thread(model: &Qwen35Model) -> Result { + let ctx = model.device_ctx(); + unsafe { + let err = crate::ffi::cuda_set_device(ctx.device_ordinal as i32); + if err != 0 { + return Err(anyhow::anyhow!( + "Failed to set CUDA device {} on Qwen3.5 scheduler thread: cudaError={}", + ctx.device_ordinal, + err + )); + } + } + ctx.ctx.bind_to_thread().map_err(|e| { + anyhow::anyhow!("Failed to bind CUDA context to Qwen3.5 scheduler thread: {e}") + })?; + unsafe { + crate::ffi::cublas_init(); + } + model.tune_decode_gemm_algos()?; + Ok(CublasThreadGuard) +} + +pub(super) trait DecodeDispatchBackend { + fn is_stop_token(&self, token: u32) -> bool; + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35; + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()>; +} + +impl DecodeDispatchBackend for SingleGpuBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + compact_single_slot(self, active, idx) + } + + fn drop_active_state(&mut self, _state: &ActiveBackendState) -> Result<()> { + Ok(()) + } +} + +impl DecodeDispatchBackend for Qwen35Backend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn take_active_request( + &mut self, + active: &mut Vec, + idx: usize, + ) -> ActiveRequest35 { + match self { + Qwen35Backend::Single(backend) => compact_single_slot(backend, active, idx), + Qwen35Backend::Tp(_) => active.swap_remove(idx), + } + } + + fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { + match (self, state) { + (Qwen35Backend::Single(_), ActiveBackendState::Single { .. }) => Ok(()), + (Qwen35Backend::Tp(backend), ActiveBackendState::Tp { request_id }) => { + backend.drop_request(*request_id, DropExpectation::MustExist) + } + _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during retirement"), + } + } +} + +/// Remove single-GPU request at `idx` via swap_remove and compact graph slots. +fn compact_single_slot( + backend: &mut SingleGpuBackend, + active: &mut Vec, + idx: usize, +) -> ActiveRequest35 { + let compaction = plan::compaction_after_retire(active.len(), idx); + let removed = active.swap_remove(idx); + + if let Some(compaction) = compaction { + backend.compact_slot(active, compaction); + } + removed +} + +pub(super) struct ScheduledChunk { + pub(super) ids: Vec, + pub(super) reqs: Vec, + pub(super) backend_state: ScheduledChunkBackendState, + /// Prompt cursor after this step's chunk + pub(super) ends: Vec, + /// This step's chunked token slice per request + pub(super) windows: Vec>, +} + +pub(super) struct InflightPrefill { + // Fields drop in declaration order. Drain the stream before request state + // can return KV pages or release recurrent/convolution buffers on unwind. + pub(super) output: AsyncPrefillOutput, + pub(super) chunk: ScheduledChunk, + pub(super) sample_seed: u64, +} + +pub(super) enum ScheduledChunkBackendState { + Single { + kvs: Vec, + recs: Vec, + }, + Tp { + request_ids: Vec, + }, +} + +impl From> for ScheduledChunk { + fn from(scheduled: Vec) -> Self { + let n = scheduled.len(); + let is_tp = scheduled + .first() + .is_some_and(|p| matches!(p.backend_state, PrefillBackendState::Tp { .. })); + let mut chunk = ScheduledChunk { + ids: Vec::with_capacity(n), + reqs: Vec::with_capacity(n), + backend_state: if is_tp { + ScheduledChunkBackendState::Tp { + request_ids: Vec::with_capacity(n), + } + } else { + ScheduledChunkBackendState::Single { + kvs: Vec::with_capacity(n), + recs: Vec::with_capacity(n), + } + }, + ends: Vec::with_capacity(n), + windows: Vec::with_capacity(n), + }; + for p in scheduled { + let end = p.cursor + p.step_chunk; + chunk + .windows + .push(p.request.prompt_tokens[p.cursor..end].to_vec()); + chunk.ends.push(end); + chunk.ids.push(p.id); + chunk.reqs.push(p.request); + match (&mut chunk.backend_state, p.backend_state) { + ( + ScheduledChunkBackendState::Single { kvs, recs }, + PrefillBackendState::Single { kv, rec }, + ) => { + kvs.push(kv); + recs.push(rec); + } + ( + ScheduledChunkBackendState::Tp { request_ids }, + PrefillBackendState::Tp { request_id }, + ) => request_ids.push(request_id), + _ => unreachable!("mixed Qwen3.5 scheduler backend states in one chunk"), + } + } + chunk + } +} + +pub(super) trait PrefillPromoteBackend { + fn is_stop_token(&self, token: u32) -> bool; + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState; + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()>; +} + +impl PrefillPromoteBackend for SingleGpuBackend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState { + let PrefillBackendState::Single { kv, rec } = state else { + panic!("single-GPU promotion received TP prefill state"); + }; + let slot_idx = slot_for_new_request(active_len, self.max_batch()) + .expect("admission must reserve a graph slot"); + self.copy_recurrent_to_slot(&rec, slot_idx) + .expect("copy recurrent state to slot failed"); + ActiveBackendState::Single { + kv, + graph_slot_idx: slot_idx, + } + } + + fn drop_prefill_state( + &mut self, + _state: &PrefillBackendState, + _expectation: DropExpectation, + ) -> Result<()> { + Ok(()) + } +} + +impl PrefillPromoteBackend for Qwen35Backend { + fn is_stop_token(&self, token: u32) -> bool { + self.is_stop_token(token) + } + + fn promote_prefill_state( + &mut self, + active_len: usize, + state: PrefillBackendState, + ) -> ActiveBackendState { + match (self, state) { + (Qwen35Backend::Single(single), PrefillBackendState::Single { kv, rec }) => { + let slot_idx = slot_for_new_request(active_len, single.max_batch()) + .expect("admission must reserve a graph slot"); + single + .copy_recurrent_to_slot(&rec, slot_idx) + .expect("copy recurrent state to slot failed"); + ActiveBackendState::Single { + kv, + graph_slot_idx: slot_idx, + } + } + (Qwen35Backend::Tp(_), PrefillBackendState::Tp { request_id }) => { + ActiveBackendState::Tp { request_id } + } + _ => panic!("mismatched Qwen3.5 scheduler backend state during promotion"), + } + } + + fn drop_prefill_state( + &mut self, + state: &PrefillBackendState, + expectation: DropExpectation, + ) -> Result<()> { + match (self, state) { + (Qwen35Backend::Single(_), PrefillBackendState::Single { .. }) => Ok(()), + (Qwen35Backend::Tp(backend), PrefillBackendState::Tp { request_id }) => { + backend.drop_request(*request_id, expectation) + } + _ => anyhow::bail!("mismatched Qwen3.5 scheduler backend state during prefill drop"), + } + } +} + +pub(super) fn split_scheduled_backend_state( + backend_state: ScheduledChunkBackendState, +) -> Vec { + match backend_state { + ScheduledChunkBackendState::Single { kvs, recs } => kvs + .into_iter() + .zip(recs) + .map(|(kv, rec)| PrefillBackendState::Single { kv, rec }) + .collect(), + ScheduledChunkBackendState::Tp { request_ids } => request_ids + .into_iter() + .map(|request_id| PrefillBackendState::Tp { request_id }) + .collect(), + } +} diff --git a/pegainfer-qwen35/src/scheduler/mod.rs b/pegainfer-qwen35/src/scheduler/mod.rs new file mode 100644 index 000000000..d03ae5e81 --- /dev/null +++ b/pegainfer-qwen35/src/scheduler/mod.rs @@ -0,0 +1,1135 @@ +//! Scheduler for Qwen3.5: a [`Scheduler`] on the contract driver thread. +//! +//! Mirrors the Qwen3 scheduler but manages: +//! - `RecurrentState` alongside `KvState` (linear attention layers) +//! - `BatchDecodeGraphState` for CUDA Graph batch decode (stable-address slots) +//! +//! Shape matches K3: this module implements `submit` / `step` / `metrics`, +//! writes the ledger, and returns [`Engine`] from `start_*`. GPU execute and +//! overlap live in [`backend`]. + +mod backend; +mod plan; + +#[cfg(test)] +mod tests; + +use std::collections::VecDeque; +use std::sync::OnceLock; +use std::time::Instant; +use std::time::SystemTime; +use std::time::UNIX_EPOCH; + +use anyhow::Result; +use log::debug; +use log::info; +use log::warn; +use pegainfer_frontend::engine::Engine; +use pegainfer_frontend::engine::EngineInfo; +use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::KvCapacity; +use pegainfer_frontend::engine::QueuedRequest; +use pegainfer_frontend::engine::RejectReason as ContractRejectReason; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::RequestId as EngineRequestId; +use pegainfer_frontend::engine::RequestLedger; +use pegainfer_frontend::engine::Scheduler; +use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::TokenLogprob; +use pegainfer_frontend::engine::spawn_scheduler; +use rand::SeedableRng; +use rand::rngs::StdRng; + +use self::backend::ActiveRequest35; +use self::backend::CublasThreadGuard; +use self::backend::DecodeDispatchBackend; +use self::backend::InflightPrefill; +use self::backend::PrefillPromoteBackend; +use self::backend::PrefillStepArtifacts; +use self::backend::PrefillingRequest35; +use self::backend::Qwen35Backend; +use self::backend::ScheduledChunk; +use self::backend::SingleGpuBackend; +use self::backend::TpSchedulerBackend; +use self::backend::bind_model_thread; +use self::backend::servable_len; +use self::backend::split_decode_artifacts; +use self::backend::split_scheduled_backend_state; +use self::plan::ActiveDecodeState; +use self::plan::ActiveKvBudget; +use self::plan::ExecutionPlan; +use self::plan::PrefillKvBudget; +use self::plan::PrefillQueueState; +use self::plan::RejectReason; +use self::plan::admit_pending_requests; +use self::plan::choose_prefill_budget; +use self::plan::max_kv_tokens; +use self::plan::plan_prefill_chunks; +use self::plan::prefilling_future_pages; +use crate::Qwen35DecodeOverlap; +use crate::Qwen35SchedulerPolicy; +use crate::tp_executor::DropExpectation; +use crate::weights::Qwen35Model; + +struct PendingRequest { + id: EngineRequestId, + request: Request, +} + +pub const DEFAULT_MAX_PREFILL_TOKENS: usize = 1024; + +/// Env-gated per-step ITL diagnostics (issue #470). When `PEGAINFER_ITL_DEBUG` +/// is set, the scheduler emits one `ITL_STEP` line per executed step, tagging +/// the plan kind, the *actual* prefill-chunk token count associated with the +/// action, the active decode width, and the CPU wall-time. Off by default. +fn itl_debug_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| std::env::var_os("PEGAINFER_ITL_DEBUG").is_some()) +} + +fn itl_debug_mono_us() -> u128 { + static ORIGIN: OnceLock = OnceLock::new(); + ORIGIN.get_or_init(Instant::now).elapsed().as_micros() +} + +fn log_itl_step( + step_start: Option, + plan: &str, + prefill_tokens: usize, + prefill_reqs: usize, + decode_n: usize, +) { + let Some(step_start) = step_start else { + return; + }; + let dur_us = step_start.elapsed().as_micros(); + let epoch_us = SystemTime::now() + .duration_since(UNIX_EPOCH) + .map_or(0, |d| d.as_micros()); + info!( + "ITL_STEP mono_us={} epoch_us={} plan={} prefill_tok={} prefill_reqs={} decode_n={} dur_us={}", + itl_debug_mono_us(), + epoch_us, + plan, + prefill_tokens, + prefill_reqs, + decode_n, + dur_us + ); +} + +// ── Engine assembly ───────────────────────────────────────────────────── + +pub fn start_with_capacity( + model: Qwen35Model, + seed: u64, + max_batch: usize, + max_prefill_tokens: usize, +) -> Result { + start_with_capacity_and_policy( + model, + seed, + max_batch, + max_prefill_tokens, + Qwen35SchedulerPolicy::Off, + Qwen35DecodeOverlap::Off, + ) +} + +pub(crate) fn start_with_capacity_and_policy( + model: Qwen35Model, + seed: u64, + max_batch: usize, + max_prefill_tokens: usize, + scheduler_policy: Qwen35SchedulerPolicy, + decode_overlap: Qwen35DecodeOverlap, +) -> Result { + assert!( + max_prefill_tokens > 0, + "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" + ); + let total_blocks = model.kv_pool().capacity_pages().saturating_sub(1); + let block_size = model.kv_pool().layout().page_size; + let servable = servable_len( + model.config().max_position_embeddings, + total_blocks, + block_size, + ); + let backend = SingleGpuBackend::new(model, max_batch, decode_overlap)?; + let scheduler = Qwen35Scheduler::new( + Qwen35Backend::Single(backend), + seed, + max_prefill_tokens, + scheduler_policy, + ); + Ok(Engine { + schedulers: vec![spawn_scheduler("qwen35-scheduler", scheduler)], + info: EngineInfo { + kv_capacity: Some(KvCapacity { + total_blocks, + block_size, + }), + servable_len: Some(servable), + }, + lora: None, + }) +} + +pub(crate) fn start_tp_with_capacity( + model_path: &str, + seed: u64, + device_ordinals: &[usize], + max_batch: usize, + max_prefill_tokens: usize, +) -> Result { + assert!( + max_prefill_tokens > 0, + "max_prefill_tokens must be positive: a zero budget can never schedule a prefill chunk" + ); + let backend = + TpSchedulerBackend::new(model_path, device_ordinals, max_batch, max_prefill_tokens)?; + let servable = servable_len( + backend.max_position_embeddings(), + backend.capacity_pages_for_requests(), + backend.page_size(), + ); + let kv_capacity = KvCapacity { + total_blocks: backend.capacity_pages_for_requests(), + block_size: backend.page_size(), + }; + let scheduler = Qwen35Scheduler::new( + Qwen35Backend::Tp(backend), + seed, + max_prefill_tokens, + Qwen35SchedulerPolicy::Off, + ); + Ok(Engine { + schedulers: vec![spawn_scheduler("qwen35-scheduler-tp", scheduler)], + info: EngineInfo { + kv_capacity: Some(kv_capacity), + servable_len: Some(servable), + }, + lora: None, + }) +} + +// ── The Scheduler implementation ──────────────────────────────────────── + +fn finish_or_retire(id: EngineRequestId, reason: FinishReason, ledger: &mut RequestLedger) { + if ledger.is_aborted(id) { + ledger.retire(id); + } else { + ledger.finish(id, reason); + } +} + +struct Qwen35Scheduler { + backend: Qwen35Backend, + rng: StdRng, + prefill_budget: usize, + scheduler_policy: Qwen35SchedulerPolicy, + pending: Vec, + active: Vec, + prefilling: Vec, + inflight_prefill: Option, + /// Lives for the driver-thread lifetime. Single-GPU only; TP workers bind + /// themselves. + cublas_guard: Option, + ready_logged: bool, +} + +impl Qwen35Scheduler { + fn new( + backend: Qwen35Backend, + seed: u64, + prefill_budget: usize, + scheduler_policy: Qwen35SchedulerPolicy, + ) -> Self { + Self { + backend, + rng: StdRng::seed_from_u64(seed), + prefill_budget, + scheduler_policy, + pending: Vec::new(), + active: Vec::new(), + prefilling: Vec::new(), + inflight_prefill: None, + cublas_guard: None, + ready_logged: false, + } + } + + /// Bind CUDA + init thread-local cuBLAS on the first `step` of the + /// driver thread (single-GPU only). Stored on `self` so the guard + /// outlives every later step. + fn bind_if_needed(&mut self) -> Result<()> { + if self.cublas_guard.is_some() { + return Ok(()); + } + if let Qwen35Backend::Single(single) = &self.backend { + let guard = bind_model_thread(single.model())?; + self.cublas_guard = Some(guard); + } + if !self.ready_logged { + info!("scheduler ready (max_batch={})", self.backend.max_batch()); + self.ready_logged = true; + } + Ok(()) + } + + fn prune_aborted(&mut self, ledger: &mut RequestLedger) -> Result<()> { + let mut index = 0; + while index < self.pending.len() { + if ledger.is_aborted(self.pending[index].id) { + let removed = self.pending.remove(index); + ledger.retire(removed.id); + } else { + index += 1; + } + } + + for idx in (0..self.active.len()).rev() { + if ledger.is_aborted(self.active[idx].id) { + let removed = self.backend.take_active_request(&mut self.active, idx); + self.backend.drop_active_state(&removed.backend_state)?; + ledger.retire(removed.id); + } + } + + for idx in (0..self.prefilling.len()).rev() { + if ledger.is_aborted(self.prefilling[idx].id) { + let removed = self.prefilling.remove(idx); + let expectation = prefill_drop_expectation(removed.cursor); + self.backend + .drop_prefill_state(&removed.backend_state, expectation)?; + ledger.retire(removed.id); + } + } + Ok(()) + } + + fn reject_echo(&mut self, ledger: &mut RequestLedger) { + let mut index = 0; + while index < self.pending.len() { + let Some(reason) = echo_refusal(&self.pending[index].request) else { + index += 1; + continue; + }; + let removed = self.pending.remove(index); + if ledger.is_aborted(removed.id) { + ledger.retire(removed.id); + } else { + ledger.reject(removed.id, reason); + } + } + } + + fn admit_pending(&mut self, ledger: &mut RequestLedger) { + let pending = std::mem::take(&mut self.pending); + let active_budget: Vec = self + .active + .iter() + .map(|req| ActiveKvBudget { + prompt_len: req.prompt_len, + generated_count: req.generated_count, + max_tokens: req.max_tokens, + }) + .collect(); + let page_size = self.backend.page_size(); + let prefilling_budget: Vec = self + .prefilling + .iter() + .map(|p| PrefillKvBudget { + current_tokens: p.cursor, + prompt_len: p.request.prompt_tokens.len(), + max_tokens: p.request.max_tokens, + }) + .collect(); + let page_budget = self + .backend + .available_pages(&self.active, &self.prefilling) + .saturating_sub(prefilling_future_pages(&prefilling_budget, page_size)); + let decode_batching_slot = self + .backend + .max_batch() + .saturating_sub(self.prefilling.len()); + let admission = admit_pending_requests( + pending, + &active_budget, + decode_batching_slot, + page_size, + page_budget, + self.backend.capacity_pages_for_requests(), + self.backend.max_position_embeddings(), + |req| req.request.prompt_tokens.len(), + |req| req.request.max_tokens, + ); + for (rejected, reason) in admission.rejected { + if ledger.is_aborted(rejected.id) { + ledger.retire(rejected.id); + } else { + ledger.reject( + rejected.id, + contract_reject_reason( + rejected.request.prompt_tokens.len(), + rejected.request.max_tokens, + reason, + ), + ); + } + } + for req in admission.pending { + if ledger.is_aborted(req.id) { + ledger.retire(req.id); + continue; + } + ledger.admit(req.id); + debug!( + "request admitted: request_id={} prompt_len={} max_tokens={}", + req.id, + req.request.prompt_tokens.len(), + req.request.max_tokens + ); + match self.backend.alloc_prefill_state() { + Ok(backend_state) => self.prefilling.push(PrefillingRequest35 { + id: req.id, + request: req.request, + backend_state, + cursor: 0, + step_chunk: 0, + }), + Err(error) => { + warn!("failed to allocate recurrent state for new request: {error}"); + ledger.fail(req.id, error.to_string()); + } + } + } + self.pending = admission.deferred; + } +} + +impl Scheduler for Qwen35Scheduler { + fn submit(&mut self, request: QueuedRequest) { + self.pending.push(PendingRequest { + id: request.id, + request: request.request, + }); + } + + fn step(&mut self, ledger: &mut RequestLedger) -> Result<()> { + self.bind_if_needed()?; + + if self + .inflight_prefill + .as_mut() + .is_some_and(|prefill| prefill.output.is_ready()) + { + let (prefill_tokens, prefill_reqs) = + self.inflight_prefill.as_ref().map_or((0, 0), |prefill| { + ( + prefill.chunk.windows.iter().map(Vec::len).sum(), + prefill.chunk.reqs.len(), + ) + }); + let decode_n = self.active.len(); + let step_start = itl_debug_enabled().then(Instant::now); + finish_async_prefill( + &mut self.backend, + &mut self.active, + &mut self.prefilling, + self.inflight_prefill + .take() + .expect("ready async prefill must still be present"), + ledger, + )?; + log_itl_step( + step_start, + "overlap_complete", + prefill_tokens, + prefill_reqs, + decode_n, + ); + } + + self.prune_aborted(ledger)?; + self.reject_echo(ledger); + + // One async prefill owns its scheduled request state. Do not admit or + // launch a second chunk until it resolves. Active decode keeps moving; + // if it retires first, wait on the event inside `step` instead of + // returning idle to the driver. + if self.inflight_prefill.is_some() { + let itl_step_start = itl_debug_enabled().then(Instant::now); + let (itl_prefill_tokens, itl_prefill_reqs) = + self.inflight_prefill.as_ref().map_or((0, 0), |prefill| { + ( + prefill.chunk.windows.iter().map(Vec::len).sum(), + prefill.chunk.reqs.len(), + ) + }); + let itl_decode_n = self.active.len(); + let itl_plan_kind = if self.active.is_empty() { + finish_async_prefill( + &mut self.backend, + &mut self.active, + &mut self.prefilling, + self.inflight_prefill + .take() + .expect("async prefill must be present before blocking wait"), + ledger, + )?; + "overlap_wait" + } else { + decode_step(&mut self.backend, &mut self.active, &mut self.rng, ledger)?; + "overlap_decode" + }; + log_itl_step( + itl_step_start, + itl_plan_kind, + itl_prefill_tokens, + itl_prefill_reqs, + itl_decode_n, + ); + return Ok(()); + } + + self.admit_pending(ledger); + + let active_decode: Vec = self + .active + .iter() + .map(|req| ActiveDecodeState { + generated_count: req.generated_count, + max_tokens: req.max_tokens, + }) + .collect(); + let prefill_queue: Vec = self + .prefilling + .iter() + .map(|req| PrefillQueueState { + remaining_tokens: req.request.prompt_tokens.len().saturating_sub(req.cursor), + }) + .collect(); + let step_prefill_budget = choose_prefill_budget( + self.scheduler_policy, + self.prefill_budget, + &active_decode, + &prefill_queue, + ); + let scheduled = take_prefill_chunks(&mut self.prefilling, step_prefill_budget); + let itl_debug = itl_debug_enabled(); + let itl_prefill_tokens: usize = scheduled.iter().map(|p| p.step_chunk).sum(); + let itl_prefill_reqs = scheduled.len(); + let itl_decode_n = self.active.len(); + let Some(plan) = plan::build_next_plan(!self.active.is_empty(), scheduled) else { + return Ok(()); + }; + let itl_plan_kind = match &plan { + ExecutionPlan::Unified { .. } if matches!(&self.backend, Qwen35Backend::Single(single) if single.overlap_enabled()) => { + "overlap_launch" + } + ExecutionPlan::Unified { .. } => "unified", + ExecutionPlan::Prefill { .. } => "prefill", + ExecutionPlan::Decode => "decode", + }; + let itl_step_start = itl_debug.then(Instant::now); + match plan { + ExecutionPlan::Unified { pending } => { + if matches!(&self.backend, Qwen35Backend::Single(single) if single.overlap_enabled()) + { + launch_overlap_step( + &mut self.backend, + &mut self.active, + pending, + &mut self.inflight_prefill, + &mut self.rng, + ledger, + )?; + } else { + unified_step_sched( + &mut self.backend, + &mut self.active, + pending, + &mut self.prefilling, + &mut self.rng, + ledger, + )?; + } + } + ExecutionPlan::Prefill { pending } => { + prefill_batch( + &mut self.backend, + &mut self.active, + pending, + &mut self.prefilling, + &mut self.rng, + ledger, + )?; + } + ExecutionPlan::Decode => { + decode_step(&mut self.backend, &mut self.active, &mut self.rng, ledger)?; + } + } + log_itl_step( + itl_step_start, + itl_plan_kind, + itl_prefill_tokens, + itl_prefill_reqs, + itl_decode_n, + ); + Ok(()) + } + + fn metrics(&self) -> SchedulerMetrics { + let kv_total_blocks = self.backend.capacity_pages_for_requests() as u64; + let (num_running_reqs, num_waiting_reqs) = logical_load_counts( + self.active.len(), + self.prefilling.len(), + self.inflight_prefill + .as_ref() + .map_or(0, |prefill| prefill.chunk.reqs.len()), + self.pending.len(), + ); + SchedulerMetrics { + kv_used_blocks: + kv_total_blocks.saturating_sub( + self.backend.available_pages(&self.active, &self.prefilling) as u64, + ), + kv_total_blocks, + num_running_reqs, + num_waiting_reqs, + spec_decode: None, + } + } +} + +/// Running = decode + resident prefill + in-flight overlap prefill. Waiting +/// is the queued/deferred count the caller already computed. Inflight must +/// stay in the running tally so an overlap wait inside `step` is not +/// published as idle. +fn logical_load_counts( + active_len: usize, + prefilling_len: usize, + inflight_prefill_reqs: usize, + num_waiting_reqs: usize, +) -> (u64, u64) { + ( + (active_len + prefilling_len + inflight_prefill_reqs) as u64, + num_waiting_reqs as u64, + ) +} + +/// Cursor 0 never materialized worker state; a later cursor did. TP drop +/// must match that, or a cancelled first chunk looks like a missing request +/// on the worker and a cancelled mid-prompt chunk looks like a leak. +fn prefill_drop_expectation(cursor: usize) -> DropExpectation { + if cursor == 0 { + DropExpectation::MustBeAbsent + } else { + DropExpectation::MustExist + } +} + +/// Echo is unsupported on Qwen3.5: refuse before KV admission with a zero +/// prefill bound so the contract `Display` is the client message. +fn echo_refusal(request: &Request) -> Option { + request + .echo + .then_some(ContractRejectReason::EchoPrefillTokens { + prompt_tokens: request.prompt_tokens.len(), + limit: 0, + }) +} + +/// Widen a plan-level admission verdict into the contract's typed reason. +fn contract_reject_reason( + prompt_tokens: usize, + max_tokens: usize, + reason: RejectReason, +) -> ContractRejectReason { + match reason { + RejectReason::ContextLength { limit } => ContractRejectReason::ContextLength { + prompt_tokens, + max_tokens, + limit, + }, + RejectReason::KvBudget => ContractRejectReason::KvBudget { + prompt_tokens, + worst_case_tokens: max_kv_tokens(prompt_tokens, max_tokens), + }, + } +} + +fn fail_ids( + ids: impl IntoIterator, + ledger: &mut RequestLedger, + message: &str, +) { + for id in ids { + if ledger.is_aborted(id) { + ledger.retire(id); + } else { + ledger.fail(id, message); + } + } +} + +fn fail_active( + backend: &mut impl DecodeDispatchBackend, + active: &mut Vec, + ledger: &mut RequestLedger, + message: &str, +) -> Result<()> { + while !active.is_empty() { + let req = backend.take_active_request(active, active.len() - 1); + backend.drop_active_state(&req.backend_state)?; + if ledger.is_aborted(req.id) { + ledger.retire(req.id); + } else { + ledger.fail(req.id, message); + } + } + Ok(()) +} + +// ── Batch prefill ─────────────────────────────────────────────────────── + +fn prefill_batch( + backend: &mut Qwen35Backend, + active: &mut Vec, + scheduled: Vec, + prefilling: &mut Vec, + rng: &mut StdRng, + ledger: &mut RequestLedger, +) -> Result<()> { + let mut chunk = ScheduledChunk::from(scheduled); + let sample_seed = rand::RngExt::random(rng); + let artifacts = match backend { + Qwen35Backend::Single(single) => { + let logits = match single.batch_prefill_logits(&mut chunk) { + Ok(v) => v, + Err(e) => { + warn!("batch prefill failed: {e}"); + fail_ids(chunk.ids, ledger, &e.to_string()); + return Ok(()); + } + }; + let prefill_sample_seed = rand::RngExt::random(rng); + match single.sample_prefill_logits(&chunk.reqs, &logits, prefill_sample_seed) { + Ok((tokens, logprobs)) => PrefillStepArtifacts::Single { tokens, logprobs }, + Err(e) => { + warn!("prefill sampling failed: {e}"); + fail_ids(chunk.ids, ledger, &e.to_string()); + return Ok(()); + } + } + } + Qwen35Backend::Tp(tp) => match tp.execute_prefill_chunk(&chunk, sample_seed) { + Ok(v) => PrefillStepArtifacts::Tp(v), + Err(e) => { + warn!("TP prefill chunk failed: {e}"); + return Err(e); + } + }, + }; + + promote_or_requeue(backend, active, prefilling, chunk, &artifacts, ledger) +} + +fn launch_overlap_step( + backend: &mut Qwen35Backend, + active: &mut Vec, + scheduled: Vec, + inflight_prefill: &mut Option, + rng: &mut StdRng, + ledger: &mut RequestLedger, +) -> Result<()> { + debug_assert!(inflight_prefill.is_none()); + let mut chunk = ScheduledChunk::from(scheduled); + let decode_seed = rand::RngExt::random(rng); + let prefill_seed = rand::RngExt::random(rng); + let output = match backend { + Qwen35Backend::Single(single) => single.launch_async_prefill(&mut chunk), + Qwen35Backend::Tp(_) => unreachable!("Qwen3.5 TP cannot launch async prefill"), + }; + match output { + Ok(output) => { + *inflight_prefill = Some(InflightPrefill { + chunk, + output, + sample_seed: prefill_seed, + }); + } + Err(err) => { + warn!("async prefill launch failed: {err}"); + fail_ids(chunk.ids, ledger, &err.to_string()); + } + } + decode_step_with_seed(backend, active, decode_seed, ledger) +} + +fn finish_async_prefill( + backend: &mut Qwen35Backend, + active: &mut Vec, + prefilling: &mut Vec, + inflight: InflightPrefill, + ledger: &mut RequestLedger, +) -> Result<()> { + let InflightPrefill { + chunk, + output, + sample_seed, + } = inflight; + let logits = output.into_logits(); + let Qwen35Backend::Single(single) = backend else { + unreachable!("Qwen3.5 TP cannot finish async prefill"); + }; + let (tokens, logprobs) = match single.sample_prefill_logits(&chunk.reqs, &logits, sample_seed) { + Ok(result) => result, + Err(err) => { + warn!("async prefill sampling failed: {err}"); + fail_ids(chunk.ids, ledger, &err.to_string()); + return Ok(()); + } + }; + let artifacts = PrefillStepArtifacts::Single { tokens, logprobs }; + promote_or_requeue(single, active, prefilling, chunk, &artifacts, ledger) +} + +// ── Unified step (prefill chunk + decode in one forward pass) ────────────── + +fn unified_step_sched( + backend: &mut Qwen35Backend, + active: &mut Vec, + scheduled: Vec, + prefilling: &mut Vec, + rng: &mut StdRng, + ledger: &mut RequestLedger, +) -> Result<()> { + let mut chunk = ScheduledChunk::from(scheduled); + if matches!(backend, Qwen35Backend::Tp(_)) { + let decode_sample_seed = rand::RngExt::random(rng); + let prefill_sample_seed = rand::RngExt::random(rng); + let result = { + let Qwen35Backend::Tp(tp) = backend else { + unreachable!() + }; + tp.execute_unified(&chunk, active, decode_sample_seed, prefill_sample_seed) + }; + let artifacts = match result { + Ok(artifacts) => artifacts, + Err(err) => { + warn!("TP unified step failed: {err}"); + return Err(err); + } + }; + + let (decode_tokens, decode_logprobs) = split_decode_artifacts(&artifacts.decode); + dispatch_decode_tokens(backend, active, &decode_tokens, &decode_logprobs, ledger)?; + + let prefill = PrefillStepArtifacts::Tp(artifacts.prefill); + return promote_or_requeue(backend, active, prefilling, chunk, &prefill, ledger); + } + + let Qwen35Backend::Single(backend) = backend else { + unreachable!() + }; + let result = backend.unified_step(&mut chunk, active); + let output = match result { + Ok(v) => v, + Err(e) => { + warn!("unified step failed: {e}"); + let message = e.to_string(); + fail_active(backend, active, ledger, &message)?; + fail_ids(chunk.ids, ledger, &message); + return Ok(()); + } + }; + let decode_seed = rand::RngExt::random(rng); + let prefill_seed = rand::RngExt::random(rng); + + if output.decoded { + process_decode_logits(backend, active, decode_seed, ledger)?; + } + + let prefill_logits = output + .prefill_logits + .as_ref() + .expect("scheduled prefill chunk must return prefill logits"); + let (tokens, logprobs) = + match backend.sample_prefill_logits(&chunk.reqs, prefill_logits, prefill_seed) { + Ok(v) => v, + Err(e) => { + warn!("unified prefill sampling failed: {e}"); + fail_ids(chunk.ids, ledger, &e.to_string()); + return Ok(()); + } + }; + let prefill = PrefillStepArtifacts::Single { tokens, logprobs }; + promote_or_requeue(backend, active, prefilling, chunk, &prefill, ledger) +} + +// ── Decode step (pure decode, CUDA Graph enabled) ────────────────────── + +fn decode_step( + backend: &mut Qwen35Backend, + active: &mut Vec, + rng: &mut StdRng, + ledger: &mut RequestLedger, +) -> Result<()> { + let first_seed = rand::RngExt::random(rng); + let sample_seed = if matches!(backend, Qwen35Backend::Single(_)) { + rand::RngExt::random(rng) + } else { + first_seed + }; + decode_step_with_seed(backend, active, sample_seed, ledger) +} + +fn decode_step_with_seed( + backend: &mut Qwen35Backend, + active: &mut Vec, + sample_seed: u64, + ledger: &mut RequestLedger, +) -> Result<()> { + let (tokens, logprobs_vec) = match backend { + Qwen35Backend::Single(single) => { + if let Err(e) = single.decode_graph(active) { + warn!("batch_decode_graph error: {e}"); + fail_active(single, active, ledger, &e.to_string())?; + return Ok(()); + } + match single.sample_decode_logits(active, sample_seed) { + Ok(v) => v, + Err(e) => { + warn!("decode sampling/logprobs error: {e}"); + fail_active(single, active, ledger, &e.to_string())?; + return Ok(()); + } + } + } + Qwen35Backend::Tp(tp) => match tp.execute_decode(active, sample_seed) { + Ok(v) => split_decode_artifacts(&v), + Err(e) => { + warn!("TP eager decode error: {e}"); + return Err(e); + } + }, + }; + + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec, ledger) +} + +fn process_decode_logits( + backend: &mut SingleGpuBackend, + active: &mut Vec, + sample_seed: u64, + ledger: &mut RequestLedger, +) -> Result<()> { + let (tokens, logprobs_vec) = match backend.sample_decode_logits(active, sample_seed) { + Ok(v) => v, + Err(e) => { + warn!("decode sampling/logprobs error: {e}"); + fail_active(backend, active, ledger, &e.to_string())?; + return Ok(()); + } + }; + + dispatch_decode_tokens(backend, active, &tokens, &logprobs_vec, ledger) +} + +enum Retirement { + Stop, + Length { + token: u32, + logprob: Option, + }, + Aborted, +} + +/// Dispatch sampled decode tokens. Stop tokens are not pushed (the bridge +/// appends EOS for usage). Drop backend state before the ledger terminal so +/// TP drop-then-finish is structurally drop-before-visible-terminal. +fn dispatch_decode_tokens( + backend: &mut impl DecodeDispatchBackend, + active: &mut Vec, + tokens: &[u32], + logprobs: &[Option], + ledger: &mut RequestLedger, +) -> Result<()> { + let n = active.len(); + let mut to_retire = Vec::new(); + + for i in 0..n { + let token = tokens[i]; + let logprob = logprobs[i].clone(); + let req = &mut active[i]; + if ledger.is_aborted(req.id) { + to_retire.push((i, Retirement::Aborted)); + continue; + } + req.generated_count += 1; + + let is_eos = !req.params.ignore_eos && backend.is_stop_token(token); + let at_limit = req.generated_count >= req.max_tokens; + + if is_eos { + debug!( + "request finished: request_id={} prompt_tokens={} finish_reason={:?}", + req.id, + req.prompt_len, + FinishReason::Stop + ); + to_retire.push((i, Retirement::Stop)); + } else if at_limit { + debug!( + "request finished: request_id={} prompt_tokens={} finish_reason={:?}", + req.id, + req.prompt_len, + FinishReason::Length + ); + to_retire.push((i, Retirement::Length { token, logprob })); + } else { + req.last_token = token; + ledger.push_tokens(req.id, &[token], &[logprob]); + } + } + + for (i, retirement) in to_retire.into_iter().rev() { + match retirement { + Retirement::Stop => { + let request = backend.take_active_request(active, i); + backend.drop_active_state(&request.backend_state)?; + finish_or_retire(request.id, FinishReason::Stop, ledger); + } + Retirement::Length { token, logprob } => { + let request = backend.take_active_request(active, i); + backend.drop_active_state(&request.backend_state)?; + if !ledger.is_aborted(request.id) { + ledger.push_tokens(request.id, &[token], &[logprob]); + } + finish_or_retire(request.id, FinishReason::Length, ledger); + } + Retirement::Aborted => { + let request = backend.take_active_request(active, i); + backend.drop_active_state(&request.backend_state)?; + ledger.retire(request.id); + } + } + } + Ok(()) +} + +// ── Chunked-prefill helpers ──────────────────────────────────────────────── + +fn take_prefill_chunks( + prefilling: &mut Vec, + prefill_budget: usize, +) -> Vec { + let remaining: Vec = prefilling + .iter() + .map(|p| p.request.prompt_tokens.len() - p.cursor) + .collect(); + let chunks = plan_prefill_chunks(&remaining, prefill_budget); + let mut scheduled: Vec = prefilling.drain(0..chunks.len()).collect(); + for (p, chunk) in scheduled.iter_mut().zip(&chunks) { + p.step_chunk = *chunk; + } + scheduled +} + +/// For each request in the just-prefilled chunk: if its prompt is now exhausted, +/// sample its first token and move it into the decode batch; otherwise re-queue +/// it (with an advanced cursor) at the FRONT of `prefilling`. +fn promote_or_requeue( + backend: &mut impl PrefillPromoteBackend, + active: &mut Vec, + prefilling: &mut Vec, + chunk: ScheduledChunk, + artifacts: &PrefillStepArtifacts, + ledger: &mut RequestLedger, +) -> Result<()> { + let ScheduledChunk { + ids, + reqs, + backend_state, + ends, + .. + } = chunk; + let mut still_prefilling: Vec = Vec::new(); + let backend_states = split_scheduled_backend_state(backend_state); + let mut entries: VecDeque<_> = ids + .into_iter() + .zip(reqs) + .zip(backend_states) + .zip(ends) + .enumerate() + .map(|(i, (((id, request), backend_state), end))| (i, id, request, backend_state, end)) + .collect(); + + while let Some((i, id, request, backend_state, end)) = entries.pop_front() { + if ledger.is_aborted(id) { + let expectation = prefill_drop_expectation(end); + backend.drop_prefill_state(&backend_state, expectation)?; + ledger.retire(id); + continue; + } + + if end < request.prompt_tokens.len() { + still_prefilling.push(PrefillingRequest35 { + id, + request, + backend_state, + cursor: end, + step_chunk: 0, + }); + continue; + } + + let prompt_len = request.prompt_tokens.len(); + let artifact = artifacts.final_artifact(i); + let first_token = artifact.token; + let logprob = artifact.logprob; + + if !request.params.ignore_eos && backend.is_stop_token(first_token) { + debug!( + "request finished: request_id={} prompt_tokens={} finish_reason={:?}", + id, + prompt_len, + FinishReason::Stop + ); + backend.drop_prefill_state(&backend_state, DropExpectation::MustExist)?; + finish_or_retire(id, FinishReason::Stop, ledger); + continue; + } + + if request.max_tokens <= 1 { + debug!( + "request finished: request_id={} prompt_tokens={} finish_reason={:?}", + id, + prompt_len, + FinishReason::Length + ); + backend.drop_prefill_state(&backend_state, DropExpectation::MustExist)?; + ledger.push_tokens(id, &[first_token], &[logprob]); + finish_or_retire(id, FinishReason::Length, ledger); + continue; + } + + let active_backend_state = backend.promote_prefill_state(active.len(), backend_state); + ledger.push_tokens(id, &[first_token], &[logprob]); + active.push(ActiveRequest35 { + id, + backend_state: active_backend_state, + last_token: first_token, + generated_count: 1, + max_tokens: request.max_tokens, + prompt_len, + params: request.params, + logprobs: request.logprobs, + }); + } + + prefilling.splice(0..0, still_prefilling); + Ok(()) +} diff --git a/pegainfer-qwen35/src/scheduler/tests.rs b/pegainfer-qwen35/src/scheduler/tests.rs index 39a521517..7ceeac71a 100644 --- a/pegainfer-qwen35/src/scheduler/tests.rs +++ b/pegainfer-qwen35/src/scheduler/tests.rs @@ -1,29 +1,50 @@ +//! Contract-facing scheduler unit tests. No TokenSink, no GenerateRequest, +//! no EngineHandle, no watch load_tx, no tokio submit_rx. +//! +//! Protocol tests that need a live ledger go through `spawn_scheduler` once a +//! fake backend exists (K3 `scheduler/tests.rs`). `RequestLedger::new` is +//! `pub(crate)` in the frontend crate, so this model crate cannot mint one. +//! Abort prune therefore cannot be driven here: see K3 +//! `aborted_request_retires_silently_and_frees_its_slot`. +//! +//! Abolished (structure makes the old failure mode impossible — migration-defense): +//! - `terminal_scheduler_shutdown` fan-out: heir is frontend `drive()` `fail_all` +//! (`pegainfer-frontend/src/engine/driver.rs`). +//! - `completion_requires_drop_ack` / `publish_before_retire` TokenEvent-vs-drop +//! ordering: one `RequestUpdate` is committed after the whole step, so drop +//! always happens before the client sees the terminal. +//! - `FatalSchedulerError.transient` + `TokenEvent::Error` fan-out: `step` +//! returns `Err`, driver `fail_all`s. + use std::path::Path; -use std::sync::Arc; -use std::sync::Barrier; -use std::time::Duration; -use std::time::Instant; +use pegainfer_frontend::engine::Engine; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::EpBackend; - -use super::*; - -fn test_request(request_id: &str, token_tx: TokenSink) -> SchedulerRequest { - test_request_with_shape(request_id, token_tx, vec![1], 1) -} - -fn test_request_with_shape( - request_id: &str, - token_tx: TokenSink, - prompt_tokens: Vec, - max_tokens: usize, -) -> SchedulerRequest { - SchedulerRequest { - trace_parent: None, - request_id: Some(request_id.to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, +use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::LiveScheduler; +use pegainfer_frontend::engine::RejectReason; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::RequestId; +use pegainfer_frontend::engine::RequestUpdate; +use pegainfer_frontend::engine::StepReceiver; +use pegainfer_frontend::engine::Terminal; +use pegainfer_frontend::sampler::SamplingParams; + +use super::DropExpectation; +use super::contract_reject_reason; +use super::echo_refusal; +use super::logical_load_counts; +use super::plan::RejectReason as AdmissionReject; +use super::prefill_drop_expectation; +use super::start_tp_with_capacity; + +fn request(prompt_len: usize, max_tokens: usize) -> Request { + request_with_prompt(vec![7; prompt_len], max_tokens) +} + +fn request_with_prompt(prompt_tokens: Vec, max_tokens: usize) -> Request { + Request { prompt_tokens, params: SamplingParams { ignore_eos: true, @@ -32,852 +53,228 @@ fn test_request_with_shape( max_tokens, lora_adapter: None, kv_transfer_params: None, - token_tx, logprobs: 0, echo: false, + trace_parent: None, + client_label: None, } } -fn active_request(request_id: u64, label: &str, token_tx: TokenSink) -> ActiveRequest35 { - ActiveRequest35 { - request_id: Some(label.to_string()), - token_tx, - backend_state: ActiveBackendState::Tp { - request_id: RequestId::new(request_id), - }, - last_token: 1, - generated_count: 1, - max_tokens: 8, - prompt_len: 1, - params: SamplingParams::default(), - logprobs: 0, - } -} - -fn prefilling_request(request_id: u64, label: &str, token_tx: TokenSink) -> PrefillingRequest35 { - PrefillingRequest35 { - req: test_request(label, token_tx), - backend_state: PrefillBackendState::Tp { - request_id: RequestId::new(request_id), - }, - cursor: 0, - step_chunk: 0, - } -} - -#[derive(Default)] -struct PruneTestBackend { - retired_active: Vec, - dropped_prefilling: Vec<(RequestId, DropExpectation)>, -} - -impl DecodeDispatchBackend for PruneTestBackend { - fn is_stop_token(&self, _token: u32) -> bool { - false - } - - fn completion_requires_drop_ack(&self) -> bool { - true - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - active.swap_remove(idx) - } - - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id } = state else { - panic!("prune test expected TP active state"); - }; - self.retired_active.push(*request_id); - Ok(()) - } -} - -impl PrefillPromoteBackend for PruneTestBackend { - fn is_stop_token(&self, _token: u32) -> bool { - false - } - - fn promote_prefill_state( - &mut self, - _active_len: usize, - _state: PrefillBackendState, - ) -> ActiveBackendState { - panic!("prune test must not promote prefill state") - } - - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()> { - let PrefillBackendState::Tp { request_id } = state else { - panic!("prune test expected TP prefill state"); - }; - self.dropped_prefilling.push((*request_id, expectation)); - Ok(()) - } -} - -struct LifecycleTestBackend { - stop_token: Option, - active_completion_requires_drop_ack: bool, - fail_active_drop: bool, - fail_prefill_drop: bool, - active_drops: Vec, - active_events_before_drop: Vec, - prefill_drops: Vec<(RequestId, DropExpectation)>, - observer: Option, +/// Demultiplex the step stream per request. Same shape as K3 / Qwen3 protocol +/// tests; used by the ignored TP2 GPU case. +struct StepCollector { + steps: StepReceiver, + buffered: std::collections::HashMap>, } -impl LifecycleTestBackend { - fn new( - stop_token: Option, - observer: pegainfer_frontend::engine::TokenStreamReceiver, - ) -> Self { +impl StepCollector { + fn new(steps: StepReceiver) -> Self { Self { - stop_token, - active_completion_requires_drop_ack: true, - fail_active_drop: false, - fail_prefill_drop: false, - active_drops: Vec::new(), - active_events_before_drop: Vec::new(), - prefill_drops: Vec::new(), - observer: Some(observer), + steps, + buffered: std::collections::HashMap::new(), } } - fn assert_no_completion_published(&mut self) { - let Some(observer) = &mut self.observer else { - return; - }; - assert!(matches!( - observer.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - } -} - -impl DecodeDispatchBackend for LifecycleTestBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.stop_token == Some(token) - } - - fn completion_requires_drop_ack(&self) -> bool { - self.active_completion_requires_drop_ack - } - - fn take_active_request( - &mut self, - active: &mut Vec, - idx: usize, - ) -> ActiveRequest35 { - active.swap_remove(idx) - } - - fn drop_active_state(&mut self, state: &ActiveBackendState) -> Result<()> { - let ActiveBackendState::Tp { request_id } = state else { - panic!("lifecycle test expected TP active state"); - }; - if self.active_completion_requires_drop_ack { - self.assert_no_completion_published(); - } else if let Some(observer) = &mut self.observer { - while let Ok((_, event)) = observer.try_recv() { - self.active_events_before_drop.push(event); + fn next_for(&mut self, id: RequestId) -> RequestUpdate { + loop { + if let Some(update) = self + .buffered + .get_mut(&id) + .and_then(std::collections::VecDeque::pop_front) + { + return update; + } + let step = self + .steps + .blocking_recv() + .expect("step stream closed while awaiting an update"); + for update in step.updates { + self.buffered + .entry(update.id) + .or_default() + .push_back(update); } } - self.active_drops.push(*request_id); - anyhow::ensure!(!self.fail_active_drop, "injected active drop failure"); - Ok(()) } -} -impl PrefillPromoteBackend for LifecycleTestBackend { - fn is_stop_token(&self, token: u32) -> bool { - self.stop_token == Some(token) - } - - fn promote_prefill_state( - &mut self, - _active_len: usize, - _state: PrefillBackendState, - ) -> ActiveBackendState { - panic!("completion lifecycle test must not promote prefill state") - } - - fn drop_prefill_state( - &mut self, - state: &PrefillBackendState, - expectation: DropExpectation, - ) -> Result<()> { - let PrefillBackendState::Tp { request_id } = state else { - panic!("lifecycle test expected TP prefill state"); - }; - self.assert_no_completion_published(); - self.prefill_drops.push((*request_id, expectation)); - anyhow::ensure!(!self.fail_prefill_drop, "injected prefill drop failure"); - Ok(()) + fn collect_terminal(&mut self, id: RequestId) -> (Vec, Terminal) { + let mut tokens = Vec::new(); + loop { + let update = self.next_for(id); + tokens.extend_from_slice(&update.tokens); + if let Some(terminal) = update.terminal { + return (tokens, terminal); + } + } } } -fn next_event( - rx: &mut pegainfer_frontend::engine::TokenStreamReceiver, - description: &str, -) -> TokenEvent { - rx.blocking_recv() - .unwrap_or_else(|| panic!("{description} channel closed before event")) - .1 +fn partition(mut engine: Engine) -> (LiveScheduler, StepCollector) { + let mut scheduler = engine.schedulers.remove(0); + let steps = scheduler + .handle + .take_steps() + .expect("a fresh scheduler yields its step stream once"); + (scheduler, StepCollector::new(steps)) } -fn assert_no_more_events(rx: &mut pegainfer_frontend::engine::TokenStreamReceiver) { - assert!( - rx.try_recv().is_err(), - "request received more than one terminal event" - ); -} +/// Heir of the TokenSink mixed-step collection. +fn assert_forced_mixed_steps(engine: Engine) { + let (partition, mut steps) = partition(engine); -#[test] -fn closed_pending_work_is_pruned_before_admission() { - let (closed_sink, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let (open_sink, _open_rx) = TokenSink::standalone(); - let mut pending = vec![ - test_request("closed", closed_sink), - test_request("open", open_sink), - ]; - let mut active = Vec::new(); - let mut prefilling = Vec::new(); - let mut backend = PruneTestBackend::default(); + let decode = partition + .handle + .submit(request_with_prompt(vec![151_646], 8)); + let prefill = partition + .handle + .submit(request_with_prompt(vec![151_646, 9707], 2)); + let (decode_tokens, decode_finish) = steps.collect_terminal(decode.id()); + let (prefill_tokens, prefill_finish) = steps.collect_terminal(prefill.id()); + assert_eq!(decode_tokens.len(), 8); assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - assert_eq!(pending.len(), 1); - assert_eq!(pending[0].request_id.as_deref(), Some("open")); - let admission = admit_pending_requests( - pending, - &[], - 1, - 16, - 8, - 8, - 128, - |req| req.prompt_tokens.len(), - |req| req.max_tokens, + matches!( + decode_finish, + Terminal::Finished { + reason: FinishReason::Length, + completion_tokens: 8, + .. + } + ), + "{decode_finish:?}" ); - assert_eq!(admission.pending.len(), 1); - assert!(admission.deferred.is_empty()); - assert!(admission.rejected.is_empty()); -} - -#[test] -fn closed_resident_work_is_absent_from_post_prune_load() { - let (closed_active_sink, closed_active_rx) = TokenSink::standalone(); - drop(closed_active_rx); - let (open_active_sink, _open_active_rx) = TokenSink::standalone(); - let (closed_prefill_sink, closed_prefill_rx) = TokenSink::standalone(); - drop(closed_prefill_rx); - let (pending_sink, _pending_rx) = TokenSink::standalone(); - let mut active = vec![ - active_request(10, "active-closed", closed_active_sink), - active_request(11, "active-open", open_active_sink), - ]; - let mut prefilling = vec![prefilling_request( - 12, - "prefill-closed", - closed_prefill_sink, - )]; - let mut pending = vec![test_request("pending-open", pending_sink)]; - let mut backend = PruneTestBackend::default(); - + assert_eq!(prefill_tokens.len(), 2); assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - assert_eq!(active.len(), 1); - assert_eq!(active[0].request_id.as_deref(), Some("active-open")); - assert!(prefilling.is_empty()); - assert_eq!( - logical_load_counts(&active, &prefilling, 0, pending.len()), - (1, 1) - ); - assert_eq!(backend.retired_active, vec![RequestId::new(10)]); - assert_eq!( - backend.dropped_prefilling, - vec![(RequestId::new(12), DropExpectation::MustBeAbsent)] + matches!( + prefill_finish, + Terminal::Finished { + reason: FinishReason::Length, + completion_tokens: 2, + .. + } + ), + "{prefill_finish:?}" ); } -#[test] -fn closed_resident_frees_capacity_for_same_tick_admission() { - let (closed_sink, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let (pending_sink, _pending_rx) = TokenSink::standalone(); - let mut active = vec![active_request(20, "resident-closed", closed_sink)]; - let mut prefilling = Vec::new(); - let mut pending = vec![test_request("replacement", pending_sink)]; - let mut backend = PruneTestBackend::default(); +// ── Rejection Display (same failure mode, contract heir) ──────────────── +#[test] +fn send_rejection_reports_kv_lifetime_request_tokens() { + let reason = contract_reject_reason(16, 65, AdmissionReject::KvBudget); + match reason { + RejectReason::KvBudget { + prompt_tokens, + worst_case_tokens, + } => { + assert_eq!(prompt_tokens, 16); + assert_eq!(worst_case_tokens, 80); + } + other => panic!("expected KvBudget, got {other:?}"), + } + let message = reason.to_string(); assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - let active_budget: Vec = active - .iter() - .map(|req| ActiveKvBudget { - prompt_len: req.prompt_len, - generated_count: req.generated_count, - max_tokens: req.max_tokens, - }) - .collect(); - let admission = admit_pending_requests( - pending, - &active_budget, - 1usize.saturating_sub(prefilling.len()), - 16, - 8, - 8, - 128, - |req| req.prompt_tokens.len(), - |req| req.max_tokens, + message.contains("max_request_tokens=80"), + "rejection should report the full lifetime KV request: {message}" ); - - assert!(active.is_empty()); - assert_eq!(backend.retired_active, vec![RequestId::new(20)]); - assert_eq!(admission.pending.len(), 1); - assert_eq!( - admission.pending[0].request_id.as_deref(), - Some("replacement") - ); - assert!(admission.deferred.is_empty()); } #[test] -fn closed_materialized_prefill_requires_existing_worker_state() { - let (closed_sink, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let mut prefilling = vec![PrefillingRequest35 { - cursor: 1, - ..prefilling_request(21, "prefill-materialized", closed_sink) - }]; - let mut active = Vec::new(); - let mut pending = Vec::new(); - let mut backend = PruneTestBackend::default(); - +fn send_rejection_reports_context_window_limit() { + let reason = contract_reject_reason(16, 17, AdmissionReject::ContextLength { limit: 32 }); + match reason { + RejectReason::ContextLength { + prompt_tokens, + max_tokens, + limit, + } => { + assert_eq!((prompt_tokens, max_tokens, limit), (16, 17, 32)); + } + other => panic!("expected ContextLength, got {other:?}"), + } + let message = reason.to_string(); assert!( - prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending).is_ok() - ); - - assert!(prefilling.is_empty()); - assert_eq!( - backend.dropped_prefilling, - vec![(RequestId::new(21), DropExpectation::MustExist)] + message.contains("maximum context length of 32 tokens"), + "rejection should report the context-window limit: {message}" ); -} - -#[test] -fn prune_drop_failure_preserves_pending_for_terminal_fanout() { - let (closed_tx, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - let (pending_tx, mut pending_rx) = TokenSink::standalone(); - let (_observer_tx, observer_rx) = TokenSink::standalone(); - let mut active = vec![active_request(22, "closed-active", closed_tx)]; - let mut prefilling = Vec::new(); - let mut pending = vec![test_request("live-pending", pending_tx)]; - let mut backend = LifecycleTestBackend::new(None, observer_rx); - backend.observer = None; - backend.fail_active_drop = true; - - let failure = - match prune_closed_requests(&mut backend, &mut active, &mut prefilling, &mut pending) { - Ok(()) => panic!("injected prune drop should fail"), - Err(failure) => failure, - }; - assert!(active.is_empty()); - assert_eq!(failure.transient.len(), 1); - assert_eq!(pending.len(), 1); - - let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - active, - prefilling, - pending, - Vec::new(), - None, - failure, + assert!( + message.contains("requested 33"), + "rejection should report prompt + max_tokens: {message}" ); - - assert!(matches!( - next_event(&mut pending_rx, "pending after prune failure"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut pending_rx); } #[test] -fn decode_eos_waits_for_drop_before_finished() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(30, "decode-eos", token_tx); - request.params.ignore_eos = false; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(Some(9), token_rx); - - assert!(dispatch_decode_tokens(&mut backend, &mut active, &[9], &[None]).is_ok()); - - assert!(active.is_empty()); - assert_eq!(backend.active_drops, vec![RequestId::new(30)]); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - next_event(&mut token_rx, "decode EOS"), - TokenEvent::Finished { - finish_reason: FinishReason::Stop, - .. - } - )); - assert_no_more_events(&mut token_rx); -} - -#[test] -fn decode_length_waits_for_drop_before_token_and_finished() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(31, "decode-length", token_tx); - request.params.ignore_eos = true; - request.max_tokens = 2; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(None, token_rx); - - assert!(dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None]).is_ok()); +fn echo_request_is_rejected_before_backend_admission() { + let mut echo = request(3, 4); + echo.echo = true; + let regular = request(1, 1); - assert!(active.is_empty()); - assert_eq!(backend.active_drops, vec![RequestId::new(31)]); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - next_event(&mut token_rx, "decode length token"), - TokenEvent::Token { id: 7, .. } - )); - assert!(matches!( - next_event(&mut token_rx, "decode length finish"), - TokenEvent::Finished { - finish_reason: FinishReason::Length, - .. + let reason = echo_refusal(&echo).expect("echo must be refused before KV admission"); + match reason { + RejectReason::EchoPrefillTokens { + prompt_tokens, + limit, + } => { + assert_eq!(prompt_tokens, 3); + assert_eq!( + limit, 0, + "echo is unsupported, not merely over a prefill cap" + ); } - )); - assert_no_more_events(&mut token_rx); -} - -#[test] -fn non_tp_decode_preserves_publish_before_retire_order() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(36, "single-order", token_tx); - request.params.ignore_eos = true; - request.max_tokens = 2; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(None, token_rx); - backend.active_completion_requires_drop_ack = false; - - assert!(dispatch_decode_tokens(&mut backend, &mut active, &[8], &[None]).is_ok()); - - assert!(active.is_empty()); - assert_eq!(backend.active_drops, vec![RequestId::new(36)]); - assert_eq!(backend.active_events_before_drop.len(), 2); - assert!(matches!( - &backend.active_events_before_drop[0], - TokenEvent::Token { id: 8, .. } - )); - assert!(matches!( - &backend.active_events_before_drop[1], - TokenEvent::Finished { - finish_reason: FinishReason::Length, - .. + other => panic!("expected EchoPrefillTokens, got {other:?}"), + } + assert_eq!( + reason.to_string(), + RejectReason::EchoPrefillTokens { + prompt_tokens: 3, + limit: 0, } - )); -} - -#[test] -fn decode_completion_drop_failure_publishes_only_terminal_error() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = active_request(32, "decode-drop-failure", token_tx); - request.params.ignore_eos = true; - request.max_tokens = 2; - let mut active = vec![request]; - let mut backend = LifecycleTestBackend::new(None, token_rx); - backend.fail_active_drop = true; - - let failure = match dispatch_decode_tokens(&mut backend, &mut active, &[7], &[None]) { - Ok(()) => panic!("injected active drop should fail"), - Err(failure) => failure, - }; - assert!(active.is_empty()); - assert_eq!(failure.transient.len(), 1); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - token_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - - let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - Vec::new(), - Vec::new(), - Vec::new(), - Vec::new(), - None, - failure, + .to_string() ); - - assert!(matches!( - next_event(&mut token_rx, "failed decode completion"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut token_rx); -} - -#[test] -fn immediate_prefill_completion_waits_for_drop() { - let (token_tx, token_rx) = TokenSink::standalone(); - let mut request = prefilling_request(33, "prefill-length", token_tx); - request.req.max_tokens = 1; - request.req.params.ignore_eos = true; - request.step_chunk = 1; - let chunk = ScheduledChunk::from(vec![request]); - let mut active = Vec::new(); - let mut prefilling = Vec::new(); - let mut backend = LifecycleTestBackend::new(None, token_rx); - assert!( - promote_or_requeue( - &mut backend, - &mut active, - &mut prefilling, - chunk, - &PrefillStepArtifacts::Single { - tokens: vec![11], - logprobs: vec![None], - }, - ) - .is_ok() - ); - - assert!(active.is_empty()); - assert!(prefilling.is_empty()); - assert_eq!( - backend.prefill_drops, - vec![(RequestId::new(33), DropExpectation::MustExist)] + echo_refusal(®ular).is_none(), + "only echo requests are ineligible for backend admission" ); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - next_event(&mut token_rx, "prefill length token"), - TokenEvent::Token { id: 11, .. } - )); - assert!(matches!( - next_event(&mut token_rx, "prefill length finish"), - TokenEvent::Finished { - finish_reason: FinishReason::Length, - .. - } - )); - assert_no_more_events(&mut token_rx); } -#[test] -fn immediate_prefill_drop_failure_publishes_only_terminal_error() { - let (token_tx, token_rx) = TokenSink::standalone(); - let (remaining_tx, mut remaining_rx) = TokenSink::standalone(); - let mut request = prefilling_request(34, "prefill-drop-failure", token_tx); - request.req.max_tokens = 1; - request.req.params.ignore_eos = true; - request.step_chunk = 1; - let mut remaining = prefilling_request(35, "remaining-scheduled", remaining_tx); - remaining.req.max_tokens = 1; - remaining.req.params.ignore_eos = true; - remaining.step_chunk = 1; - let chunk = ScheduledChunk::from(vec![request, remaining]); - let mut active = Vec::new(); - let mut prefilling = Vec::new(); - let mut backend = LifecycleTestBackend::new(None, token_rx); - backend.fail_prefill_drop = true; - - let failure = match promote_or_requeue( - &mut backend, - &mut active, - &mut prefilling, - chunk, - &PrefillStepArtifacts::Single { - tokens: vec![12, 13], - logprobs: vec![None, None], - }, - ) { - Ok(()) => panic!("injected prefill drop should fail"), - Err(failure) => failure, - }; - assert_eq!(failure.transient.len(), 2); - let mut token_rx = backend.observer.take().unwrap(); - assert!(matches!( - token_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - assert!(matches!( - remaining_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); - - let (_submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let (load_tx, _load_rx) = watch::channel(SchedulerMetrics::default()); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - active, - prefilling, - Vec::new(), - Vec::new(), - None, - failure, - ); - - assert!(matches!( - next_event(&mut token_rx, "failed prefill completion"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut token_rx); - assert!(matches!( - next_event(&mut remaining_rx, "remaining scheduled prefill"), - TokenEvent::Error { .. } - )); - assert_no_more_events(&mut remaining_rx); -} +// ── Drop expectation (pure data; abort prune needs a ledger) ──────────── +// +// `prune_aborted` will take `&mut RequestLedger`, which this crate cannot +// mint. Queued/active/prefilling abort-retire is covered by K3 +// `aborted_request_retires_silently_and_frees_its_slot`. What remains +// testable without a ledger is the TP drop expectation on cursor. #[test] -fn terminal_shutdown_closes_drains_and_errors_every_owner_once() { - let (active_tx, active_rx) = TokenSink::standalone(); - let (prefill_tx, prefill_rx) = TokenSink::standalone(); - let (pending_tx, pending_rx) = TokenSink::standalone(); - let (deferred_tx, deferred_rx) = TokenSink::standalone(); - let (candidate_tx, candidate_rx) = TokenSink::standalone(); - let (scheduled_tx, scheduled_rx) = TokenSink::standalone(); - let (queued_tx, queued_rx) = TokenSink::standalone(); - let (after_close_tx, mut after_close_rx) = TokenSink::standalone(); - let (closed_tx, closed_rx) = TokenSink::standalone(); - drop(closed_rx); - - let active = vec![active_request(40, "active", active_tx)]; - let prefilling = vec![prefilling_request(41, "prefilling", prefill_tx)]; - let pending = vec![ - test_request("duplicate-external-id", pending_tx), - test_request("closed-sink", closed_tx), - ]; - let deferred = vec![test_request("duplicate-external-id", deferred_tx)]; - let candidate = CompletionCandidate { - request: active_request(42, "candidate", candidate_tx), - final_events: vec![TokenEvent::Finished { - finish_reason: FinishReason::Length, - prompt_tokens: 1, - completion_tokens: 2, - }], - }; - let failure = FatalSchedulerError::new("injected TP replica failure") - .with_request(candidate.into_terminal()) - .with_request(test_request("scheduled", scheduled_tx)); - - let queued = test_request("queued-before-close", queued_tx); - let after_close = test_request("queued-after-close", after_close_tx); - let (submit_tx, mut submit_rx) = mpsc::unbounded_channel(); - let before_send = Arc::new(Barrier::new(2)); - let sent_before_close = Arc::new(Barrier::new(2)); - let after_receiver_close = Arc::new(Barrier::new(2)); - let sender = { - let before_send = Arc::clone(&before_send); - let sent_before_close = Arc::clone(&sent_before_close); - let after_receiver_close = Arc::clone(&after_receiver_close); - std::thread::spawn(move || { - before_send.wait(); - submit_tx - .send((queued, pegainfer_frontend::engine::KvPrefix::none())) - .expect("close-before request should be accepted"); - sent_before_close.wait(); - after_receiver_close.wait(); - submit_tx - .send((after_close, pegainfer_frontend::engine::KvPrefix::none())) - .is_err() - }) - }; - - before_send.wait(); - sent_before_close.wait(); - let (load_tx, load_rx) = watch::channel(SchedulerMetrics { - kv_used_blocks: 9, - kv_total_blocks: 64, - num_running_reqs: 9, - num_waiting_reqs: 9, - spec_decode: None, - }); - terminal_scheduler_shutdown( - &mut submit_rx, - &load_tx, - 64, - active, - prefilling, - pending, - deferred, - None, - failure, - ); - after_receiver_close.wait(); - assert!(sender.join().expect("submit race thread panicked")); - - let mut receivers = vec![ - ("active", active_rx), - ("prefilling", prefill_rx), - ("pending", pending_rx), - ("deferred", deferred_rx), - ("candidate", candidate_rx), - ("scheduled", scheduled_rx), - ("queued", queued_rx), - ]; - for (owner, rx) in &mut receivers { - match next_event(rx, owner) { - TokenEvent::Error { message, .. } => { - assert_eq!(message, "injected TP replica failure"); - } - other => panic!("{owner} received non-error terminal event: {other:?}"), - } - assert_no_more_events(rx); - } - assert!(after_close_rx.try_recv().is_err()); - - let snapshot = *load_rx.borrow(); - assert_eq!(snapshot.kv_used_blocks, 0); - assert_eq!(snapshot.kv_total_blocks, 64); - assert_eq!(snapshot.num_running_reqs, 0); - assert_eq!(snapshot.num_waiting_reqs, 0); -} - -fn collect_finished_with_timeout( - token_rx: &mut pegainfer_frontend::engine::TokenStreamReceiver, - description: &str, -) -> (usize, FinishReason) { - let deadline = Instant::now() + Duration::from_secs(30); - let mut token_count = 0; - loop { - match token_rx.try_recv() { - Ok((_, TokenEvent::Token { .. })) => token_count += 1, - Ok((_, TokenEvent::Finished { finish_reason, .. })) => { - return (token_count, finish_reason); - } - Ok((_, TokenEvent::Error { message, .. })) => { - panic!("{description} failed: {message}") - } - Ok((_, TokenEvent::Rejected { message, .. })) => { - panic!("{description} was rejected: {message}") - } - Ok((_, _)) => {} - Err(tokio::sync::mpsc::error::TryRecvError::Empty) => { - assert!( - Instant::now() < deadline, - "timed out waiting for {description}" - ); - std::thread::sleep(Duration::from_millis(1)); - } - Err(tokio::sync::mpsc::error::TryRecvError::Disconnected) => { - panic!("{description} channel disconnected before Finished") - } - } - } +fn unmaterialized_prefill_drop_expects_absent_worker_state() { + assert_eq!(prefill_drop_expectation(0), DropExpectation::MustBeAbsent); } #[test] -fn send_rejection_reports_kv_lifetime_request_tokens() { - let (token_tx, mut token_rx) = TokenSink::standalone(); - let req = SchedulerRequest { - trace_parent: None, - request_id: Some("too-large".to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens: vec![1; 16], - params: SamplingParams::default(), - max_tokens: 65, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }; - - send_rejection(&req, RejectReason::KvBudget); - - match token_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens, - }) => { - assert_eq!(prompt_tokens, 16); - assert_eq!(completion_tokens, 0); - assert!( - message.contains("max_request_tokens=80"), - "rejection should report the full lifetime KV request" - ); - } - _ => panic!("expected rejection event"), - } +fn materialized_prefill_drop_requires_existing_worker_state() { + assert_eq!(prefill_drop_expectation(1), DropExpectation::MustExist); + assert_eq!(prefill_drop_expectation(16), DropExpectation::MustExist); } -#[test] -fn echo_request_is_rejected_before_backend_admission() { - let (echo_tx, mut echo_rx) = TokenSink::standalone(); - let (regular_tx, mut regular_rx) = TokenSink::standalone(); - let mut echo = test_request_with_shape("unsupported-echo", echo_tx, vec![1, 2, 3], 4); - echo.echo = true; - let regular = test_request("regular", regular_tx); - let mut pending = vec![echo, regular]; - - reject_unsupported_echo(&mut pending); - - assert_eq!(pending.len(), 1); - assert_eq!(pending[0].request_id.as_deref(), Some("regular")); - assert!( - !pending[0].echo, - "only requests eligible for backend admission may remain" - ); - match echo_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens, - }) => { - assert_eq!(message, UNSUPPORTED_ECHO_MESSAGE); - assert_eq!(prompt_tokens, 3); - assert_eq!(completion_tokens, 0); - } - event => panic!("expected unsupported echo rejection, got {event:?}"), - } - assert!(matches!( - regular_rx.try_recv(), - Err(tokio::sync::mpsc::error::TryRecvError::Empty) - )); -} +// ── Overlap wait lives inside step, not as driver idle ────────────────── #[test] -fn inflight_prefill_waits_instead_of_parking_after_last_decode_retires() { - assert!( - !should_block_on_submit(true, true), - "an in-flight prefill must keep the scheduler off submit_rx.blocking_recv()" +fn overlap_wait_policy_is_inside_step() { + // When inflight prefill is set and active is empty, `step` waits on the + // CUDA event (overlap_wait). It must not return idle to the driver until + // that prefill is finished — the old `should_block_on_submit` gate + // (`owned_work_empty && !inflight`) is deleted; the wait is inside step. + // GPU-untestable here; the reachable heir is that inflight still counts + // as running so metrics cannot look drained mid-wait. + assert_eq!( + logical_load_counts(0, 0, 1, 0), + (1, 0), + "inflight prefill is running work, not an idle scheduler" ); + assert_eq!(logical_load_counts(1, 0, 1, 2), (2, 2)); } +// ── Launch validation ─────────────────────────────────────────────────── + #[test] fn tp_engine_rejects_cuda_graph_before_model_load() { let err = match crate::start_engine_with_capacity( @@ -906,75 +303,7 @@ fn tp2_scheduler_runs_forced_mixed_steps() { else { return; }; - let handle = + let engine = start_tp_with_capacity(&model_path, 42, &[0, 1], 2, 1).expect("start TP2 scheduler"); - let (decode_tx, mut decode_rx) = TokenSink::standalone(); - let (prefill_tx, mut prefill_rx) = TokenSink::standalone(); - - handle - .submit(test_request_with_shape( - "mixed-active", - decode_tx, - vec![151_646], - 8, - )) - .expect("submit request that becomes active first"); - handle - .submit(test_request_with_shape( - "mixed-prefill", - prefill_tx, - vec![151_646, 9707], - 2, - )) - .expect("submit request that remains chunk-prefilling"); - - let (decode_tokens, decode_finish) = - collect_finished_with_timeout(&mut decode_rx, "mixed active request"); - let (prefill_tokens, prefill_finish) = - collect_finished_with_timeout(&mut prefill_rx, "mixed prefill request"); - assert_eq!(decode_tokens, 8); - assert_eq!(decode_finish, FinishReason::Length); - assert_eq!(prefill_tokens, 2); - assert_eq!(prefill_finish, FinishReason::Length); -} - -#[test] -fn send_rejection_reports_context_window_limit() { - let (token_tx, mut token_rx) = TokenSink::standalone(); - let req = SchedulerRequest { - trace_parent: None, - request_id: Some("too-long".to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens: vec![1; 16], - params: SamplingParams::default(), - max_tokens: 17, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }; - - send_rejection(&req, RejectReason::ContextLength { limit: 32 }); - - match token_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens, - }) => { - assert_eq!(prompt_tokens, 16); - assert_eq!(completion_tokens, 0); - assert!( - message.contains("maximum context length of 32 tokens"), - "rejection should report the context-window limit" - ); - assert!( - message.contains("requested 33"), - "rejection should report prompt + max_tokens" - ); - } - _ => panic!("expected rejection event"), - } + assert_forced_mixed_steps(engine); } diff --git a/pegainfer-qwen35/tests/chunked_prefill.rs b/pegainfer-qwen35/tests/chunked_prefill.rs index dee13c8b2..9a5b848a6 100644 --- a/pegainfer-qwen35/tests/chunked_prefill.rs +++ b/pegainfer-qwen35/tests/chunked_prefill.rs @@ -8,73 +8,52 @@ use std::path::Path; -use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::FinishReason; -use pegainfer_frontend::engine::GenerateRequest; -use pegainfer_frontend::engine::TokenEvent; -use pegainfer_frontend::engine::TokenSink; +use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::sampler::SamplingParams; mod common; +use common::harness::EngineHarness; + const CHUNK_BUDGET: usize = 16; const BASELINE_PREFILL_BUDGET: usize = 1 << 20; const MAX_BATCH: usize = 2; const GENERATED_TOKENS: usize = 8; -fn start_engine(model_path: &str, max_prefill_tokens: usize) -> EngineHandle { - pegainfer_qwen35::start_engine( - Path::new(model_path), - EngineLoadOptions { - enable_cuda_graph: true, - device_ordinals: vec![0], - seed: 42, - ..EngineLoadOptions::default() - }, - MAX_BATCH, - max_prefill_tokens, +fn start_engine(model_path: &str, max_prefill_tokens: usize) -> EngineHarness { + EngineHarness::new( + pegainfer_qwen35::start_engine( + Path::new(model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0], + seed: 42, + ..EngineLoadOptions::default() + }, + MAX_BATCH, + max_prefill_tokens, + ) + .expect("failed to start Qwen3.5 engine"), ) - .expect("failed to start Qwen3.5 engine") } -fn generate(handle: &EngineHandle, prompt_tokens: Vec) -> (Vec, FinishReason) { - let (token_tx, mut rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, +fn generate(harness: &EngineHarness, prompt_tokens: Vec) -> (Vec, FinishReason) { + let outcome = harness + .submit(common::harness::request( prompt_tokens, - params: SamplingParams { + SamplingParams { ignore_eos: true, ..SamplingParams::default() }, - max_tokens: GENERATED_TOKENS, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }) - .expect("submit failed"); - - let mut tokens = Vec::new(); - loop { - match rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Token { id, .. }) => tokens.push(id), - Some( - TokenEvent::Scheduled { .. } - | TokenEvent::PromptTokens { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - Some(TokenEvent::Finished { finish_reason, .. }) => return (tokens, finish_reason), - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), - None => panic!("scheduler channel closed without Finished"), - } - } + GENERATED_TOKENS, + )) + .expect_finished(); + let Terminal::Finished { reason, .. } = outcome.terminal else { + unreachable!("expect_finished returned a non-Finished terminal"); + }; + (outcome.tokens, reason) } #[test] @@ -102,8 +81,8 @@ fn chunked_prefill_matches_unchunked_prefill_for_resumed_paged_kv() { ); let (baseline_tokens, baseline_finish) = { - let handle = start_engine(&model_path, BASELINE_PREFILL_BUDGET); - generate(&handle, prompt_tokens.clone()) + let harness = start_engine(&model_path, BASELINE_PREFILL_BUDGET); + generate(&harness, prompt_tokens.clone()) }; assert_eq!( baseline_finish, @@ -112,8 +91,8 @@ fn chunked_prefill_matches_unchunked_prefill_for_resumed_paged_kv() { ); let (chunked_tokens, chunked_finish) = { - let handle = start_engine(&model_path, CHUNK_BUDGET); - generate(&handle, prompt_tokens) + let harness = start_engine(&model_path, CHUNK_BUDGET); + generate(&harness, prompt_tokens) }; assert_eq!( chunked_finish, diff --git a/pegainfer-qwen35/tests/common/harness.rs b/pegainfer-qwen35/tests/common/harness.rs new file mode 100644 index 000000000..f26c76daa --- /dev/null +++ b/pegainfer-qwen35/tests/common/harness.rs @@ -0,0 +1,271 @@ +//! Blocking test harness over the engine contract. +//! +//! Integration tests predate the step-batched wire protocol and want the old +//! per-request stream ergonomics; this pump restores them the same way the +//! production bridge does — one thread demultiplexes `StepOutputs` into +//! per-request queues. Everything here is synchronous: GPU tests drive +//! requests from plain test threads. + +// Each integration test crate compiles this module independently and uses a +// different slice of it. +#![allow(dead_code)] + +use std::collections::HashMap; +use std::collections::VecDeque; +use std::sync::Arc; +use std::sync::Condvar; +use std::sync::Mutex; + +use pegainfer_frontend::engine::Engine; +use pegainfer_frontend::engine::EngineInfo; +use pegainfer_frontend::engine::LoraClient; +use pegainfer_frontend::engine::PromptEcho; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::RequestControl; +use pegainfer_frontend::engine::RequestId; +use pegainfer_frontend::engine::RequestUpdate; +use pegainfer_frontend::engine::SchedulerHandle; +use pegainfer_frontend::engine::SchedulerMetrics; +use pegainfer_frontend::engine::Terminal; +use pegainfer_frontend::engine::TokenLogprob; +use pegainfer_frontend::sampler::SamplingParams; + +/// A contract request with test defaults; adjust fields on the result for +/// echo/logprobs/LoRA variants. +pub(crate) fn request( + prompt_tokens: Vec, + params: SamplingParams, + max_tokens: usize, +) -> Request { + Request { + prompt_tokens, + params, + max_tokens, + lora_adapter: None, + kv_transfer_params: None, + logprobs: 0, + echo: false, + trace_parent: None, + client_label: None, + } +} + +pub(crate) struct EngineHarness { + handle: Option, + lora: Option, + scheduler_join: Option>, + pump_join: Option>, + inbox: Arc, + pub(crate) info: EngineInfo, +} + +struct Inbox { + state: Mutex, + cv: Condvar, +} + +struct InboxState { + updates: HashMap>, + closed: bool, +} + +impl EngineHarness { + pub(crate) fn new(mut engine: Engine) -> Self { + assert_eq!( + engine.schedulers.len(), + 1, + "test harness drives single-scheduler engines" + ); + let mut scheduler = engine.schedulers.remove(0); + let mut steps = scheduler + .handle + .take_steps() + .expect("a fresh scheduler yields its step stream once"); + let inbox = Arc::new(Inbox { + state: Mutex::new(InboxState { + updates: HashMap::new(), + closed: false, + }), + cv: Condvar::new(), + }); + let pump_inbox = Arc::clone(&inbox); + let pump_join = std::thread::spawn(move || { + while let Some(step) = steps.blocking_recv() { + let mut state = pump_inbox.state.lock().unwrap(); + for update in step.updates { + state + .updates + .entry(update.id) + .or_default() + .push_back(update); + } + drop(state); + pump_inbox.cv.notify_all(); + } + pump_inbox.state.lock().unwrap().closed = true; + pump_inbox.cv.notify_all(); + }); + Self { + handle: Some(scheduler.handle), + lora: engine.lora, + scheduler_join: Some(scheduler.join), + pump_join: Some(pump_join), + inbox, + info: engine.info, + } + } + + pub(crate) fn submit(&self, request: Request) -> RequestStream { + let control = self + .handle + .as_ref() + .expect("harness handle lives until drop") + .submit(request); + RequestStream { + id: control.id(), + control, + inbox: Arc::clone(&self.inbox), + } + } + + /// Latest scheduler occupancy snapshot. Pull-only — the step driver + /// publishes into a cell, not a watch channel. + pub(crate) fn metrics(&self) -> SchedulerMetrics { + self.handle + .as_ref() + .expect("harness handle lives until drop") + .metrics() + } + + /// The engine's LoRA client; panics when the engine serves no adapter + /// control — the `Option` on `Engine::lora` is the capability. + pub(crate) fn lora_client(&self) -> LoraClient { + self.lora.clone().expect("engine exposes LoRA control") + } + + /// Submit one request and return its generated token ids, panicking on any + /// non-`Finished` terminal — the shape most GPU tests want. + pub(crate) fn generate( + &self, + prompt_tokens: Vec, + params: SamplingParams, + max_tokens: usize, + ) -> Vec { + self.submit(request(prompt_tokens, params, max_tokens)) + .expect_finished() + .tokens + } +} + +impl Drop for EngineHarness { + fn drop(&mut self) { + // Closing the submission channel lets the scheduler drain and exit; + // the step stream then closes and the pump follows. + drop(self.handle.take()); + if let Some(join) = self.scheduler_join.take() { + let _ = join.join(); + } + if let Some(join) = self.pump_join.take() { + let _ = join.join(); + } + } +} + +/// One request's demultiplexed view of the step stream. `Send`, so tests may +/// drive concurrent requests from separate threads. +pub(crate) struct RequestStream { + id: RequestId, + pub(crate) control: RequestControl, + inbox: Arc, +} + +impl RequestStream { + pub(crate) fn id(&self) -> RequestId { + self.id + } + + /// Next update for this request; `None` once the engine is gone. + pub(crate) fn recv(&mut self) -> Option { + let mut state = self.inbox.state.lock().unwrap(); + loop { + if let Some(update) = state + .updates + .get_mut(&self.id) + .and_then(VecDeque::pop_front) + { + return Some(update); + } + if state.closed { + return None; + } + state = self.inbox.cv.wait(state).unwrap(); + } + } + + /// Non-blocking poll. `None` means nothing is queued yet (or the engine + /// is gone); overlap tests use this to drain without waiting. + pub(crate) fn try_recv(&mut self) -> Option { + let mut state = self.inbox.state.lock().unwrap(); + state + .updates + .get_mut(&self.id) + .and_then(VecDeque::pop_front) + } + + /// Fold the stream to its terminal. Panics if the engine dies first — a + /// vanished stream is a test failure, not an outcome. + pub(crate) fn outcome(mut self) -> Outcome { + let mut outcome = Outcome { + tokens: Vec::new(), + logprobs: Vec::new(), + cached_tokens: None, + prompt_echo: None, + terminal: Terminal::Failed { + message: String::new(), + prompt_tokens: 0, + completion_tokens: 0, + }, + }; + loop { + let update = self + .recv() + .expect("engine closed the stream without a terminal"); + outcome.tokens.extend_from_slice(&update.tokens); + outcome.logprobs.extend(update.logprobs); + if update.cached_tokens.is_some() { + assert!( + outcome.cached_tokens.is_none(), + "cached_tokens must be reported at most once per request" + ); + outcome.cached_tokens = update.cached_tokens; + } + if update.prompt_echo.is_some() { + outcome.prompt_echo = update.prompt_echo; + } + if let Some(terminal) = update.terminal { + outcome.terminal = terminal; + return outcome; + } + } + } + + /// The common success path: fold to the terminal and require `Finished`. + pub(crate) fn expect_finished(self) -> Outcome { + let outcome = self.outcome(); + match &outcome.terminal { + Terminal::Finished { .. } => outcome, + Terminal::Rejected { reason, .. } => panic!("generation rejected: {reason}"), + Terminal::Failed { message, .. } => panic!("generation failed: {message}"), + } + } +} + +/// A finished request, folded: everything the old per-request event stream +/// delivered incrementally. +pub(crate) struct Outcome { + pub(crate) tokens: Vec, + pub(crate) logprobs: Vec>, + pub(crate) cached_tokens: Option, + pub(crate) prompt_echo: Option, + pub(crate) terminal: Terminal, +} diff --git a/pegainfer-qwen35/tests/common/mod.rs b/pegainfer-qwen35/tests/common/mod.rs index c56f1cec2..e0b0f8dd3 100644 --- a/pegainfer-qwen35/tests/common/mod.rs +++ b/pegainfer-qwen35/tests/common/mod.rs @@ -1,3 +1,5 @@ +pub(crate) mod harness; + use std::sync::Arc; use vllm_text::Error; diff --git a/pegainfer-qwen35/tests/e2e_scheduler.rs b/pegainfer-qwen35/tests/e2e_scheduler.rs index 9886399bb..a67228b71 100644 --- a/pegainfer-qwen35/tests/e2e_scheduler.rs +++ b/pegainfer-qwen35/tests/e2e_scheduler.rs @@ -1,26 +1,29 @@ //! E2E scheduler integration test for Qwen3.5-4B. //! //! Tests the Qwen3.5 reduced-capacity scheduler path (batch prefill + -//! CUDA Graph decode) with sequential, concurrent, and consumer-drop requests. +//! CUDA Graph decode) with sequential, concurrent, and client-abort requests. use std::collections::HashSet; use std::path::Path; +use std::sync::Arc; +use std::time::Duration; use std::time::Instant; use log::info; -use pegainfer_frontend::engine::EngineHandle; +use pegainfer_frontend::engine::Engine; use pegainfer_frontend::engine::EngineLoadOptions; use pegainfer_frontend::engine::FinishReason; -use pegainfer_frontend::engine::GenerateRequest; -use pegainfer_frontend::engine::SchedulerMetrics; -use pegainfer_frontend::engine::TokenEvent; +use pegainfer_frontend::engine::RejectReason; +use pegainfer_frontend::engine::RequestUpdate; +use pegainfer_frontend::engine::Terminal; use pegainfer_frontend::engine::TokenLogprob; -use pegainfer_frontend::engine::TokenSink; -use pegainfer_frontend::engine::TokenStreamReceiver; use pegainfer_frontend::sampler::SamplingParams; use vllm_text::tokenizer::DynTokenizer; mod common; +use common::harness::EngineHarness; +use common::harness::RequestStream; + const CASES: &[TestCase] = &[ TestCase { name: "tell_story", @@ -101,133 +104,114 @@ struct GenerationResult { } fn generate_tokens( - handle: &EngineHandle, + engine: &EngineHarness, tokenizer: &DynTokenizer, prompt: &str, max_tokens: usize, ) -> (Vec, FinishReason) { - let result = generate_tokens_with_logprobs(handle, tokenizer, prompt, max_tokens, 0); + let result = generate_tokens_with_logprobs(engine, tokenizer, prompt, max_tokens, 0); (result.tokens, result.finish_reason) } fn generate_tokens_with_logprobs( - handle: &EngineHandle, + engine: &EngineHarness, tokenizer: &DynTokenizer, prompt: &str, max_tokens: usize, logprobs: usize, ) -> GenerationResult { let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); - let (token_tx, mut token_rx) = TokenSink::standalone(); - - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams::default(), - max_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs, - echo: false, - }) - .expect("submit failed"); - - collect_generation(&mut token_rx, prompt, logprobs) + let mut request = + common::harness::request(prompt_tokens, SamplingParams::default(), max_tokens); + request.logprobs = logprobs; + collect_generation(engine.submit(request), prompt, logprobs) } fn submit_repeated_token_request( - handle: &EngineHandle, + engine: &EngineHarness, request_id: &str, token: u32, prompt_len: usize, max_tokens: usize, -) -> TokenStreamReceiver { - let (token_tx, token_rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: Some(request_id.to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens: vec![token; prompt_len], - params: SamplingParams { - ignore_eos: true, - ..SamplingParams::default() - }, - max_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }) - .unwrap_or_else(|err| panic!("submit {request_id}: {err}")); - token_rx +) -> RequestStream { + let mut request = common::harness::request( + vec![token; prompt_len], + SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens, + ); + request.client_label = Some(Arc::from(request_id)); + engine.submit(request) } -fn wait_for_first_token(rx: &mut TokenStreamReceiver, request_id: &str) { - let deadline = Instant::now() + std::time::Duration::from_secs(30); +fn wait_for_first_token(stream: &mut RequestStream, request_id: &str) { + let deadline = Instant::now() + Duration::from_secs(30); loop { - match recv_event_before(rx, request_id, deadline) { - Some(TokenEvent::Token { .. }) => return, - Some(TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. }) => {} - Some(event) => panic!("{request_id} emitted {event:?} before its first token"), + match recv_update_before(stream, request_id, deadline) { + Some(update) if !update.tokens.is_empty() => return, + Some(update) if update.terminal.is_some() => { + panic!( + "{request_id} emitted {:?} before its first token", + update.terminal + ) + } + Some(_) => {} None => panic!("scheduler closed before {request_id} emitted a token"), } } } -fn recv_event_before( - rx: &mut TokenStreamReceiver, +fn recv_update_before( + stream: &mut RequestStream, request_id: &str, deadline: Instant, -) -> Option { +) -> Option { loop { - if let Ok((_, event)) = rx.try_recv() { - return Some(event); + if let Some(update) = stream.try_recv() { + return Some(update); } assert!( Instant::now() < deadline, "timed out waiting for {request_id} scheduler event" ); - std::thread::sleep(std::time::Duration::from_millis(1)); + std::thread::sleep(Duration::from_millis(1)); } } -fn assert_no_generated_event(rx: &mut TokenStreamReceiver, request_id: &str) { - while let Ok((_, event)) = rx.try_recv() { - match event { - TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. } => {} - event => panic!("{request_id} emitted {event:?} before the overlap bound"), - } +fn assert_no_generated_event(stream: &mut RequestStream, request_id: &str) { + while let Some(update) = stream.try_recv() { + assert!( + update.tokens.is_empty(), + "{request_id} emitted tokens {:?} before the overlap bound", + update.tokens + ); + assert!( + update.terminal.is_none(), + "{request_id} emitted {:?} before the overlap bound", + update.terminal + ); } } -fn drain_tokens(rx: &mut TokenStreamReceiver, request_id: &str) -> usize { +fn drain_tokens(stream: &mut RequestStream, request_id: &str) -> usize { let mut tokens = 0; - while let Ok((_, event)) = rx.try_recv() { - match event { - TokenEvent::Token { .. } => tokens += 1, - TokenEvent::PromptTokens { .. } | TokenEvent::Scheduled { .. } => {} - event => panic!("{request_id} emitted {event:?} while it must remain active"), - } + while let Some(update) = stream.try_recv() { + assert!( + update.terminal.is_none(), + "{request_id} emitted {:?} while it must remain active", + update.terminal + ); + tokens += update.tokens.len(); } tokens } -fn wait_for_running_requests( - load: &mut tokio::sync::watch::Receiver, - expected: u64, - timeout: std::time::Duration, -) { +fn wait_for_running_requests(engine: &EngineHarness, expected: u64, timeout: Duration) { let deadline = Instant::now() + timeout; loop { - let snapshot = *load.borrow_and_update(); + let snapshot = engine.metrics(); if snapshot.num_running_reqs == expected { return; } @@ -235,29 +219,25 @@ fn wait_for_running_requests( Instant::now() < deadline, "timed out waiting for {expected} running requests; last snapshot: {snapshot:?}" ); - std::thread::sleep(std::time::Duration::from_millis(1)); + std::thread::sleep(Duration::from_millis(1)); } } -fn collect_generation( - token_rx: &mut TokenStreamReceiver, - name: &str, - logprobs: usize, -) -> GenerationResult { - collect_generation_until(token_rx, name, logprobs, None) +fn collect_generation(stream: RequestStream, name: &str, logprobs: usize) -> GenerationResult { + collect_generation_until(stream, name, logprobs, None) } fn collect_generation_with_timeout( - token_rx: &mut TokenStreamReceiver, + stream: RequestStream, name: &str, logprobs: usize, - timeout: std::time::Duration, + timeout: Duration, ) -> GenerationResult { - collect_generation_until(token_rx, name, logprobs, Some(Instant::now() + timeout)) + collect_generation_until(stream, name, logprobs, Some(Instant::now() + timeout)) } fn collect_generation_until( - token_rx: &mut TokenStreamReceiver, + mut stream: RequestStream, name: &str, logprobs: usize, deadline: Option, @@ -265,62 +245,88 @@ fn collect_generation_until( let mut tokens = Vec::new(); let mut token_logprobs = Vec::new(); loop { - let event = match deadline { - Some(deadline) => recv_event_before(token_rx, name, deadline), - None => token_rx.blocking_recv().map(|(_, event)| event), + let update = match deadline { + Some(deadline) => recv_update_before(&mut stream, name, deadline), + None => stream.recv(), }; - match event { - Some(TokenEvent::Token { id, logprob }) => { - if logprobs == 0 { - assert!( - logprob.is_none(), - "{name}: logprobs=0 should not return token logprobs" - ); - } else { - let lp = logprob - .as_ref() - .unwrap_or_else(|| panic!("{name}: logprobs={logprobs} returned None")); - assert!( - lp.logprob.is_finite(), - "{name}: sampled token logprob must be finite" - ); - assert_eq!( - lp.top_logprobs.len(), - logprobs, - "{name}: top_logprobs length should match the request" - ); - assert!( - lp.top_logprobs.iter().all(|&(_, v)| v.is_finite()), - "{name}: top_logprobs must be finite" - ); - assert_eq!( - lp.top_logprobs.first().map(|&(token, _)| token), - Some(id), - "{name}: greedy sampled token should match top-1 logprob row" - ); + match update { + Some(update) => { + apply_update(name, logprobs, &update, &mut tokens, &mut token_logprobs); + if let Some(terminal) = update.terminal { + return generation_from_terminal(name, tokens, token_logprobs, terminal); } - tokens.push(id); - token_logprobs.push(logprob); - } - Some( - TokenEvent::PromptTokens { .. } - | TokenEvent::Scheduled { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - Some(TokenEvent::Finished { finish_reason, .. }) => { - return GenerationResult { - tokens, - logprobs: token_logprobs, - finish_reason, - }; } - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), None => panic!("{name}: scheduler channel closed without Finished"), } } } +fn apply_update( + name: &str, + requested_logprobs: usize, + update: &RequestUpdate, + tokens: &mut Vec, + token_logprobs: &mut Vec>, +) { + if requested_logprobs == 0 { + assert!( + update.logprobs.iter().all(Option::is_none), + "{name}: logprobs=0 should not return token logprobs" + ); + } else if !update.tokens.is_empty() { + assert_eq!( + update.logprobs.len(), + update.tokens.len(), + "{name}: logprobs must be parallel to tokens" + ); + for (id, logprob) in update.tokens.iter().copied().zip(update.logprobs.iter()) { + let lp = logprob + .as_ref() + .unwrap_or_else(|| panic!("{name}: logprobs={requested_logprobs} returned None")); + assert!( + lp.logprob.is_finite(), + "{name}: sampled token logprob must be finite" + ); + assert_eq!( + lp.top_logprobs.len(), + requested_logprobs, + "{name}: top_logprobs length should match the request" + ); + assert!( + lp.top_logprobs.iter().all(|&(_, v)| v.is_finite()), + "{name}: top_logprobs must be finite" + ); + assert_eq!( + lp.top_logprobs.first().map(|&(token, _)| token), + Some(id), + "{name}: greedy sampled token should match top-1 logprob row" + ); + } + } + tokens.extend_from_slice(&update.tokens); + token_logprobs.extend(update.logprobs.iter().cloned()); +} + +fn generation_from_terminal( + name: &str, + tokens: Vec, + logprobs: Vec>, + terminal: Terminal, +) -> GenerationResult { + match terminal { + Terminal::Finished { + reason: finish_reason, + .. + } => GenerationResult { + tokens, + logprobs, + finish_reason, + }, + Terminal::Rejected { reason, .. } => panic!("{name}: generation rejected: {reason}"), + Terminal::Failed { message, .. } => panic!("{name}: generation failed: {message}"), + } +} + fn concurrent_params(case_idx: usize) -> SamplingParams { if case_idx.is_multiple_of(2) { SamplingParams::default() @@ -334,45 +340,37 @@ fn concurrent_params(case_idx: usize) -> SamplingParams { } } -fn expect_context_window_rejection(handle: &EngineHandle, max_context_tokens: usize) { - let (token_tx, mut token_rx) = TokenSink::standalone(); - - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: Some("over-context-window".to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens: vec![1; max_context_tokens], - params: SamplingParams::default(), - max_tokens: 1, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }) - .expect("submit over-context request"); - - match token_rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Rejected { - message, - prompt_tokens, - completion_tokens, - }) => { +fn expect_context_window_rejection(engine: &EngineHarness, max_context_tokens: usize) { + let mut request = + common::harness::request(vec![1; max_context_tokens], SamplingParams::default(), 1); + request.client_label = Some(Arc::from("over-context-window")); + + match engine.submit(request).outcome().terminal { + Terminal::Rejected { + reason: + RejectReason::ContextLength { + prompt_tokens, + max_tokens, + limit, + }, + prompt_tokens: reported_prompt, + } => { + assert_eq!(reported_prompt, max_context_tokens); assert_eq!(prompt_tokens, max_context_tokens); - assert_eq!(completion_tokens, 0); - assert!( - message.contains("maximum context length"), - "expected context-window rejection, got: {message}" - ); - assert!( - message.contains(&(max_context_tokens + 1).to_string()), - "rejection should report prompt + max_tokens, got: {message}" + assert_eq!(max_tokens, 1); + assert_eq!(limit, max_context_tokens); + assert_eq!( + prompt_tokens.saturating_add(max_tokens), + max_context_tokens + 1 ); } - Some(_) => panic!("expected context-window rejection"), - None => panic!("scheduler channel closed without rejection"), + Terminal::Rejected { reason, .. } => { + panic!("expected context-window rejection, got: {reason}") + } + Terminal::Failed { message, .. } => { + panic!("oversized prompt errored instead of clean rejection: {message}") + } + Terminal::Finished { .. } => panic!("expected context-window rejection"), } } @@ -449,7 +447,7 @@ fn assert_no_model_wide_collapse(collapses: &[(&str, Collapse)]) { } fn run_full_scheduler_e2e( - handle: &EngineHandle, + engine: &EngineHarness, tokenizer: &DynTokenizer, max_context_tokens: usize, label: &str, @@ -458,7 +456,7 @@ fn run_full_scheduler_e2e( // ── 0. Static context-window rejection ───────────────────────────── info!("=== Phase 0: Context-window rejection ==="); - expect_context_window_rejection(handle, max_context_tokens); + expect_context_window_rejection(engine, max_context_tokens); info!(" PASS: over-context request rejected before prefill"); // ── 1. logprobs must not change greedy tokens ───────────────────── @@ -466,9 +464,9 @@ fn run_full_scheduler_e2e( for case in CASES.iter().take(3) { let max_tokens = case.max_new_tokens.min(16); let no_logprobs = - generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, 0); + generate_tokens_with_logprobs(engine, tokenizer, case.prompt, max_tokens, 0); let with_logprobs = - generate_tokens_with_logprobs(handle, tokenizer, case.prompt, max_tokens, 1); + generate_tokens_with_logprobs(engine, tokenizer, case.prompt, max_tokens, 1); assert_eq!(no_logprobs.finish_reason, with_logprobs.finish_reason); assert_eq!( no_logprobs.tokens, with_logprobs.tokens, @@ -503,7 +501,7 @@ fn run_full_scheduler_e2e( info!("--- {:?} ---", case.name); let start = Instant::now(); let (tokens, finish_reason) = - generate_tokens(handle, tokenizer, case.prompt, case.max_new_tokens); + generate_tokens(engine, tokenizer, case.prompt, case.max_new_tokens); let elapsed = start.elapsed(); let text = tokenizer.decode(&tokens, true).expect("decode failed"); @@ -530,7 +528,7 @@ fn run_full_scheduler_e2e( // ── 3. Multi-request (scheduler state reuse) ──────────────────────── info!("=== Phase 3: Multi-request ==="); for case in CASES { - let (tokens, _) = generate_tokens(handle, tokenizer, case.prompt, case.max_new_tokens); + let (tokens, _) = generate_tokens(engine, tokenizer, case.prompt, case.max_new_tokens); let text = tokenizer.decode(&tokens, true).expect("decode failed"); assert!( !text.is_empty(), @@ -543,35 +541,23 @@ fn run_full_scheduler_e2e( // ── 4. Concurrent requests ────────────────────────────────────────── info!("=== Phase 4: Concurrent requests ==="); { - let mut receivers: Vec<(String, usize, TokenStreamReceiver)> = Vec::new(); + let mut streams: Vec<(String, usize, RequestStream)> = Vec::new(); // Submit all cases concurrently, alternating greedy and sampled rows so // batch decode covers the mixed token-selection path from #284. for (case_idx, case) in CASES.iter().enumerate() { let prompt_tokens = tokenizer.encode(case.prompt, false).expect("encode failed"); - let (token_tx, token_rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: concurrent_params(case_idx), - max_tokens: case.max_new_tokens, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }) - .expect("submit failed"); - receivers.push((case.name.to_string(), 0, token_rx)); + let mut request = common::harness::request( + prompt_tokens, + concurrent_params(case_idx), + case.max_new_tokens, + ); + request.client_label = Some(Arc::from(case.name)); + streams.push((case.name.to_string(), 0, engine.submit(request))); } - // Collect all results - for (name, logprobs, mut rx) in receivers { - let result = collect_generation(&mut rx, &name, logprobs); + for (name, logprobs, stream) in streams { + let result = collect_generation(stream, &name, logprobs); let text = tokenizer .decode(&result.tokens, true) .expect("decode failed"); @@ -587,32 +573,18 @@ fn run_full_scheduler_e2e( ("mixed_no_logprobs", CASES[0].prompt, 0usize), ("mixed_with_logprobs", CASES[1].prompt, 1usize), ]; - let mut receivers: Vec<(&str, usize, TokenStreamReceiver)> = Vec::new(); + let mut streams: Vec<(&str, usize, RequestStream)> = Vec::new(); for (name, prompt, logprobs) in mixed { let prompt_tokens = tokenizer.encode(prompt, false).expect("encode failed"); - let (token_tx, token_rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: Some(name.to_string()), - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams::default(), - max_tokens: 8, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs, - echo: false, - }) - .expect("submit failed"); - receivers.push((name, logprobs, token_rx)); + let mut request = common::harness::request(prompt_tokens, SamplingParams::default(), 8); + request.logprobs = logprobs; + request.client_label = Some(Arc::from(name)); + streams.push((name, logprobs, engine.submit(request))); } - for (name, logprobs, mut rx) in receivers { - let result = collect_generation(&mut rx, name, logprobs); + for (name, logprobs, stream) in streams { + let result = collect_generation(stream, name, logprobs); assert!(!result.tokens.is_empty(), "{name}: produced no tokens"); if logprobs == 0 { assert!( @@ -629,37 +601,26 @@ fn run_full_scheduler_e2e( } } - // ── 5. Consumer drop safety ───────────────────────────────────────── - info!("=== Phase 5: Consumer drop ==="); + // ── 5. Client abort safety ────────────────────────────────────────── + info!("=== Phase 5: Client abort ==="); { let prompt_tokens = tokenizer.encode("Hello", false).expect("encode failed"); - let (token_tx, rx) = TokenSink::standalone(); - drop(rx); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, - prompt_tokens, - params: SamplingParams::default(), - max_tokens: 10, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }) - .expect("submit failed"); - std::thread::sleep(std::time::Duration::from_millis(500)); - info!(" PASS: consumer drop handled"); + let stream = engine.submit(common::harness::request( + prompt_tokens, + SamplingParams::default(), + 10, + )); + stream.control.abort(); + drop(stream); + std::thread::sleep(Duration::from_millis(500)); + info!(" PASS: client abort handled"); } // Verify scheduler survives - let (tokens, _) = generate_tokens(handle, tokenizer, "Hello", 5); + let (tokens, _) = generate_tokens(engine, tokenizer, "Hello", 5); let text = tokenizer.decode(&tokens, true).expect("decode failed"); - assert!(!text.is_empty(), "scheduler dead after consumer drop"); - info!(" PASS: scheduler survived consumer drop"); + assert!(!text.is_empty(), "scheduler dead after client abort"); + info!(" PASS: scheduler survived client abort"); info!("All Qwen3.5 scheduler tests passed for {label}!"); } @@ -672,22 +633,27 @@ fn test_e2e_qwen35_scheduler() { info!("Loading Qwen3.5 model for scheduler test..."); let start = Instant::now(); - let model = - pegainfer_qwen35::runtime::Qwen35Model::from_safetensors_with_options(&model_path, true) - .expect("Failed to load model"); let tokenizer = common::load_tokenizer(&model_path); - // Use reduced batch capacity (8) to fit on 16GB GPUs alongside the model. - let handle = pegainfer_qwen35::runtime::start_with_capacity( - model, - 42, + // Load through `start_engine_with_capacity` so recurrent-state reservation + // matches the intended 8-slot 16GB budget. `from_safetensors_with_options` + // still sizes to MAX_BATCH=64 and OOMs on a 16GB card before start. + let engine: Engine = pegainfer_qwen35::start_engine_with_capacity( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0], + seed: 42, + ..EngineLoadOptions::default() + }, 8, pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, ) .expect("Failed to start Qwen3.5 scheduler"); + let engine = EngineHarness::new(engine); info!("scheduler loaded in {:.2?}", start.elapsed()); let max_context_tokens = max_position_embeddings(&model_path); - run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP1"); + run_full_scheduler_e2e(&engine, &tokenizer, max_context_tokens, "TP1"); } #[test] @@ -706,7 +672,7 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { .expect("test prompt must contain a token"); let off_reference_tokens = { - let off_handle = pegainfer_qwen35::start_engine_with_capacity_policy_and_overlap( + let off_engine: Engine = pegainfer_qwen35::start_engine_with_capacity_policy_and_overlap( Path::new(&model_path), EngineLoadOptions { enable_cuda_graph: true, @@ -720,18 +686,19 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { pegainfer_qwen35::Qwen35DecodeOverlap::Off, ) .expect("Failed to start Qwen3.5 default-Off scheduler"); - let mut off_rx = submit_repeated_token_request( - &off_handle, + let off_engine = EngineHarness::new(off_engine); + let off_stream = submit_repeated_token_request( + &off_engine, "overlap-off-reference", seed_token, 8192, 2, ); let off = collect_generation_with_timeout( - &mut off_rx, + off_stream, "overlap-off-reference", 0, - std::time::Duration::from_secs(30), + Duration::from_secs(30), ); assert_eq!( off.tokens.len(), @@ -741,7 +708,7 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { off.tokens }; - let handle = pegainfer_qwen35::start_engine_with_capacity_policy_and_overlap( + let engine: Engine = pegainfer_qwen35::start_engine_with_capacity_policy_and_overlap( Path::new(&model_path), EngineLoadOptions { enable_cuda_graph: true, @@ -755,29 +722,28 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { pegainfer_qwen35::Qwen35DecodeOverlap::SharedSm, ) .expect("Failed to start Qwen3.5 shared-SM scheduler"); - let mut load = handle - .metrics_watch() - .expect("scheduler must expose metrics"); - - let mut active_rx = - submit_repeated_token_request(&handle, "overlap-last-decoder", seed_token, 512, 128); - wait_for_first_token(&mut active_rx, "overlap-last-decoder"); - let _ = drain_tokens(&mut active_rx, "overlap-last-decoder"); - let mut prefill_rx = - submit_repeated_token_request(&handle, "overlap-inflight-prefill", seed_token, 8192, 2); - - wait_for_running_requests(&mut load, 2, std::time::Duration::from_secs(10)); - let _ = drain_tokens(&mut active_rx, "overlap-last-decoder"); + let engine = EngineHarness::new(engine); + + let mut active = + submit_repeated_token_request(&engine, "overlap-last-decoder", seed_token, 512, 128); + wait_for_first_token(&mut active, "overlap-last-decoder"); + let _ = drain_tokens(&mut active, "overlap-last-decoder"); + let mut prefill = + submit_repeated_token_request(&engine, "overlap-inflight-prefill", seed_token, 8192, 2); + + wait_for_running_requests(&engine, 2, Duration::from_secs(10)); + let _ = drain_tokens(&mut active, "overlap-last-decoder"); for _ in 0..2 { - wait_for_first_token(&mut active_rx, "overlap-last-decoder"); - assert_no_generated_event(&mut prefill_rx, "overlap-inflight-prefill"); + wait_for_first_token(&mut active, "overlap-last-decoder"); + assert_no_generated_event(&mut prefill, "overlap-inflight-prefill"); } - drop(active_rx); + active.control.abort(); + drop(active); let prefill = collect_generation_with_timeout( - &mut prefill_rx, + prefill, "overlap-inflight-prefill", 0, - std::time::Duration::from_secs(30), + Duration::from_secs(30), ); assert_eq!( prefill.tokens.len(), @@ -789,7 +755,7 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { "Shared-SM overlapped prefill must match the greedy default-Off reference" ); - let (tokens, finish_reason) = generate_tokens(&handle, &tokenizer, "Hello again", 2); + let (tokens, finish_reason) = generate_tokens(&engine, &tokenizer, "Hello again", 2); assert_eq!( tokens.len(), 2, @@ -797,24 +763,26 @@ fn test_e2e_qwen35_shared_sm_last_decoder() { ); assert_eq!(finish_reason, FinishReason::Length); - let mut shutdown_active_rx = - submit_repeated_token_request(&handle, "overlap-shutdown-decoder", seed_token, 512, 128); - wait_for_first_token(&mut shutdown_active_rx, "overlap-shutdown-decoder"); - let _ = drain_tokens(&mut shutdown_active_rx, "overlap-shutdown-decoder"); - let mut shutdown_prefill_rx = - submit_repeated_token_request(&handle, "overlap-shutdown-prefill", seed_token, 8192, 2); - wait_for_running_requests(&mut load, 2, std::time::Duration::from_secs(10)); - assert_no_generated_event(&mut shutdown_prefill_rx, "overlap-shutdown-prefill"); - drop(shutdown_active_rx); - drop(shutdown_prefill_rx); + let mut shutdown_active = + submit_repeated_token_request(&engine, "overlap-shutdown-decoder", seed_token, 512, 128); + wait_for_first_token(&mut shutdown_active, "overlap-shutdown-decoder"); + let _ = drain_tokens(&mut shutdown_active, "overlap-shutdown-decoder"); + let mut shutdown_prefill = + submit_repeated_token_request(&engine, "overlap-shutdown-prefill", seed_token, 8192, 2); + wait_for_running_requests(&engine, 2, Duration::from_secs(10)); + assert_no_generated_event(&mut shutdown_prefill, "overlap-shutdown-prefill"); + shutdown_active.control.abort(); + shutdown_prefill.control.abort(); + drop(shutdown_active); + drop(shutdown_prefill); let (done_tx, done_rx) = std::sync::mpsc::channel(); let shutdown = std::thread::spawn(move || { - drop(handle); + drop(engine); let _ = done_tx.send(()); }); done_rx - .recv_timeout(std::time::Duration::from_secs(30)) + .recv_timeout(Duration::from_secs(30)) .expect("dropping the last handle must drain in-flight prefill and return"); shutdown.join().expect("scheduler shutdown thread panicked"); } @@ -830,7 +798,7 @@ fn test_e2e_qwen35_scheduler_tp2() { let start = Instant::now(); let tokenizer = common::load_tokenizer(&model_path); // TP Phase 1 is eager-only; CUDA Graph must stay disabled for multi-device startup. - let handle = pegainfer_qwen35::start_engine_with_capacity( + let engine: Engine = pegainfer_qwen35::start_engine_with_capacity( Path::new(&model_path), EngineLoadOptions { enable_cuda_graph: false, @@ -842,8 +810,9 @@ fn test_e2e_qwen35_scheduler_tp2() { pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, ) .expect("Failed to start Qwen3.5 TP2 scheduler"); + let engine = EngineHarness::new(engine); info!("TP2 scheduler loaded in {:.2?}", start.elapsed()); let max_context_tokens = max_position_embeddings(&model_path); - run_full_scheduler_e2e(&handle, &tokenizer, max_context_tokens, "TP2"); + run_full_scheduler_e2e(&engine, &tokenizer, max_context_tokens, "TP2"); } diff --git a/pegainfer-qwen35/tests/sampling_behavior.rs b/pegainfer-qwen35/tests/sampling_behavior.rs index 4649ec186..e0880373c 100644 --- a/pegainfer-qwen35/tests/sampling_behavior.rs +++ b/pegainfer-qwen35/tests/sampling_behavior.rs @@ -10,15 +10,13 @@ use std::path::Path; -use pegainfer_frontend::engine::EngineHandle; use pegainfer_frontend::engine::EngineLoadOptions; -use pegainfer_frontend::engine::GenerateRequest; -use pegainfer_frontend::engine::TokenEvent; -use pegainfer_frontend::engine::TokenSink; use pegainfer_frontend::sampler::SamplingParams; mod common; +use common::harness::EngineHarness; + const GENERATED_TOKENS: usize = 32; fn params(mut params: SamplingParams) -> SamplingParams { @@ -27,40 +25,15 @@ fn params(mut params: SamplingParams) -> SamplingParams { } /// Submit one request and collect the generated token ids until `Finished`. -fn generate(handle: &EngineHandle, prompt_tokens: Vec, params: SamplingParams) -> Vec { - let (token_tx, mut rx) = TokenSink::standalone(); - handle - .submit(GenerateRequest { - trace_parent: None, - request_id: None, - queued_at_unix_s: None, - data_parallel_rank: None, +fn generate(harness: &EngineHarness, prompt_tokens: Vec, params: SamplingParams) -> Vec { + harness + .submit(common::harness::request( prompt_tokens, params, - max_tokens: GENERATED_TOKENS, - lora_adapter: None, - kv_transfer_params: None, - token_tx, - logprobs: 0, - echo: false, - }) - .expect("submit failed"); - - let mut tokens = Vec::new(); - loop { - match rx.blocking_recv().map(|(_, event)| event) { - Some(TokenEvent::Token { id, .. }) => tokens.push(id), - Some( - TokenEvent::Scheduled { .. } - | TokenEvent::PromptTokens { .. } - | TokenEvent::KvTransfer { .. }, - ) => {} - Some(TokenEvent::Finished { .. }) => return tokens, - Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), - Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), - None => panic!("scheduler channel closed without Finished"), - } - } + GENERATED_TOKENS, + )) + .expect_finished() + .tokens } #[test] @@ -70,18 +43,20 @@ fn sampling_params_steer_the_qwen35_sampler() { return; }; - let handle = pegainfer_qwen35::start_engine( - Path::new(&model_path), - EngineLoadOptions { - enable_cuda_graph: true, - device_ordinals: vec![0], - seed: 42, - ..EngineLoadOptions::default() - }, - 4, - pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, - ) - .expect("failed to start Qwen3.5 engine"); + let harness = EngineHarness::new( + pegainfer_qwen35::start_engine( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0], + seed: 42, + ..EngineLoadOptions::default() + }, + 4, + pegainfer_qwen35::DEFAULT_MAX_PREFILL_TOKENS, + ) + .expect("failed to start Qwen3.5 engine"), + ); let tokenizer = common::load_tokenizer(&model_path); let prompt = "Here is a short story about a dragon. Once upon a time"; @@ -89,17 +64,17 @@ fn sampling_params_steer_the_qwen35_sampler() { let greedy_params = params(SamplingParams::default()); - let greedy = generate(&handle, prompt_tokens.clone(), greedy_params); + let greedy = generate(&harness, prompt_tokens.clone(), greedy_params); assert_eq!( greedy.len(), GENERATED_TOKENS, "ignore_eos should force a full 32-token generation" ); - let greedy_again = generate(&handle, prompt_tokens.clone(), greedy_params); + let greedy_again = generate(&harness, prompt_tokens.clone(), greedy_params); assert_eq!(greedy, greedy_again, "greedy decode must be deterministic"); let top_k_one = generate( - &handle, + &harness, prompt_tokens.clone(), params(SamplingParams { temperature: 0.8, @@ -110,7 +85,7 @@ fn sampling_params_steer_the_qwen35_sampler() { assert_eq!(top_k_one, greedy, "top_k=1 must collapse to greedy"); let top_p_tiny = generate( - &handle, + &harness, prompt_tokens.clone(), params(SamplingParams { temperature: 1.0, @@ -127,7 +102,7 @@ fn sampling_params_steer_the_qwen35_sampler() { ..SamplingParams::default() }); let runs: Vec> = (0..4) - .map(|_| generate(&handle, prompt_tokens.clone(), hot)) + .map(|_| generate(&harness, prompt_tokens.clone(), hot)) .collect(); assert!( runs.iter().any(|run| *run != runs[0]), diff --git a/pegainfer-qwen35/tests/serving_tp2.rs b/pegainfer-qwen35/tests/serving_tp2.rs index 2ed223e49..ed42c3e6a 100644 --- a/pegainfer-qwen35/tests/serving_tp2.rs +++ b/pegainfer-qwen35/tests/serving_tp2.rs @@ -1,3 +1,12 @@ +//! TP2 eager generate smoke on the step-contract Engine, plus OpenAI HTTP +//! serving through `vllm::serve(LaunchedEngine::Stepped)`. +//! +//! CUDA Graph + TP is fail-closed before load. Live generate coverage is +//! in-process via [`common::harness::EngineHarness`]. The HTTP gate drives +//! `/v1/models` and `/v1/completions` (streaming and concurrent) against a +//! real frontend. + +use std::mem::ManuallyDrop; use std::net::TcpListener; use std::path::Path; use std::path::PathBuf; @@ -8,6 +17,11 @@ use anyhow::Result; use anyhow::anyhow; use anyhow::bail; use pegainfer_frontend::engine::EngineLoadOptions; +use pegainfer_frontend::engine::FinishReason; +use pegainfer_frontend::engine::LaunchedEngine; +use pegainfer_frontend::engine::Request; +use pegainfer_frontend::engine::Terminal; +use pegainfer_frontend::sampler::SamplingParams; use reqwest::Client; use serde_json::Value; use serde_json::json; @@ -16,6 +30,13 @@ use tokio_util::sync::CancellationToken; mod common; +use common::harness::EngineHarness; +use common::harness::Outcome; + +const PROMPT: &[u32] = &[151_644, 872, 198, 9707, 151_645, 198, 151_644, 77091, 198]; +const ALTERNATE_PROMPT: &[u32] = &[ + 151_644, 872, 198, 3838, 374, 220, 17, 489, 220, 17, 30, 151_645, 198, 151_644, 77091, 198, +]; const MODEL_NAME: &str = "qwen35-tp2-serving-smoke"; const HTTP_TIMEOUT: Duration = Duration::from_secs(120); @@ -35,6 +56,78 @@ impl Qwen35Tp2Server { } } +#[test] +fn tp2_rejects_cuda_graph_before_loading_model() { + let err = pegainfer_qwen35::start_engine_with_capacity( + Path::new("unused-model-path"), + EngineLoadOptions { + enable_cuda_graph: true, + device_ordinals: vec![0, 1], + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + 1, + ) + .err() + .expect("TP2 + CUDA Graph must fail before serving requests") + .to_string(); + assert!( + err.contains("eager execution only"), + "unexpected TP2 + CUDA Graph startup error: {err}" + ); +} + +#[test] +#[ignore = "requires two CUDA devices, CUDA-12 NCCL, and Qwen3.5 weights"] +fn qwen35_tp2_generates_greedy_logprobs_and_concurrent() { + let Some(model_path) = + common::model_path_or_skip("qwen35_tp2_generates_greedy_logprobs_and_concurrent") + else { + return; + }; + let gpus = cuda_device_count(); + if gpus < 2 { + eprintln!( + "SKIP qwen35_tp2_generates_greedy_logprobs_and_concurrent: needs >=2 GPUs, have {gpus}" + ); + return; + } + + let engine = pegainfer_qwen35::start_engine_with_capacity( + Path::new(&model_path), + EngineLoadOptions { + enable_cuda_graph: false, + device_ordinals: common::tp2_device_ordinals(), + seed: 42, + ..EngineLoadOptions::default() + }, + 8, + 1, + ) + .expect("start Qwen3.5 TP2 engine"); + // Dropping the harness joins the scheduler thread; on a panic the engine may + // be wedged and the join would hang, so panics leak it — only the happy path + // drops. + let harness = ManuallyDrop::new(EngineHarness::new(engine)); + + let outcome = harness + .submit(completion_request(PROMPT.to_vec(), 5, 1)) + .expect_finished(); + assert_length_outcome(&outcome, 5); + assert_logprobs(&outcome); + + let first = harness.submit(completion_request(PROMPT.to_vec(), 3, 0)); + let second = harness.submit(completion_request(ALTERNATE_PROMPT.to_vec(), 3, 1)); + let first = first.expect_finished(); + let second = second.expect_finished(); + assert_length_outcome(&first, 3); + assert_length_outcome(&second, 3); + assert_logprobs(&second); + + drop(ManuallyDrop::into_inner(harness)); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 4)] #[ignore = "requires two CUDA devices, CUDA-12 NCCL, Qwen3.5 weights, and real HTTP frontend startup"] async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { @@ -51,7 +144,6 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { return Ok(()); }; let frontend_model_path = PathBuf::from(frontend_model_path); - let invalid_graph_model_path = engine_model_path.clone(); let server = spawn_ready_server(engine_model_path, frontend_model_path, 1).await?; let client = test_client()?; @@ -59,11 +151,6 @@ async fn qwen35_tp2_serves_openai_completions_over_http() -> Result<()> { assert_non_streaming_completion(&client, &server.base_url).await?; assert_streaming_completion(&client, &server.base_url).await?; assert_concurrent_completions(&client, &server.base_url).await?; - assert_invalid_cuda_graph_tp_startup_fails( - invalid_graph_model_path - .to_str() - .context("Qwen3.5 engine fixture path is not valid UTF-8")?, - )?; server.shutdown().await } @@ -74,7 +161,7 @@ async fn spawn_ready_server( max_prefill_tokens: usize, ) -> Result { let device_ordinals = common::tp2_device_ordinals(); - let handle = tokio::task::spawn_blocking(move || { + let engine = tokio::task::spawn_blocking(move || { pegainfer_qwen35::start_engine_with_capacity( &engine_model_path, EngineLoadOptions { @@ -96,9 +183,7 @@ async fn spawn_ready_server( let server_shutdown = shutdown.clone(); let mut task = tokio::spawn(async move { pegainfer_frontend::vllm::serve( - std::future::ready(Ok(pegainfer_frontend::engine::LaunchedEngine::Handle( - handle, - ))), + std::future::ready(Ok(LaunchedEngine::Stepped(engine))), &frontend_model_path, vec![MODEL_NAME.to_string()], port, @@ -166,7 +251,7 @@ async fn assert_non_streaming_completion(client: &Client, base_url: &str) -> Res bail!("expected length finish_reason for ignore_eos request, got {completion}"); } assert_usage(&completion, 5)?; - assert_logprobs(&completion)?; + assert_http_logprobs(&completion)?; Ok(()) } @@ -195,28 +280,7 @@ async fn assert_concurrent_completions(client: &Client, base_url: &str) -> Resul let (first, second) = tokio::try_join!(first, second)?; assert_usage(&first, 3)?; assert_usage(&second, 3)?; - assert_logprobs(&second)?; - Ok(()) -} - -fn assert_invalid_cuda_graph_tp_startup_fails(model_path: &str) -> Result<()> { - let Err(error) = pegainfer_qwen35::start_engine_with_capacity( - Path::new(model_path), - EngineLoadOptions { - enable_cuda_graph: true, - device_ordinals: common::tp2_device_ordinals(), - seed: 42, - ..EngineLoadOptions::default() - }, - 8, - 1, - ) else { - bail!("TP2 + CUDA Graph must fail before serving requests"); - }; - let message = error.to_string(); - if !message.contains("eager execution only") { - bail!("unexpected TP2 + CUDA Graph startup error: {message}"); - } + assert_http_logprobs(&second)?; Ok(()) } @@ -247,7 +311,7 @@ async fn post_completion_stream(client: &Client, base_url: &str, body: Value) -> fn completion_body(stream: bool, max_tokens: usize, logprobs: usize) -> Value { let mut body = json!({ "model": MODEL_NAME, - "prompt": [151_644, 872, 198, 9707, 151_645, 198, 151_644, 77091, 198], + "prompt": PROMPT, "max_tokens": max_tokens, "temperature": 0.0, "ignore_eos": true, @@ -262,7 +326,7 @@ fn completion_body(stream: bool, max_tokens: usize, logprobs: usize) -> Value { fn alternate_completion_body(stream: bool, max_tokens: usize, logprobs: usize) -> Value { let mut body = json!({ "model": MODEL_NAME, - "prompt": [151_644, 872, 198, 3838, 374, 220, 17, 489, 220, 17, 30, 151_645, 198, 151_644, 77091, 198], + "prompt": ALTERNATE_PROMPT, "max_tokens": max_tokens, "temperature": 0.0, "ignore_eos": true, @@ -294,7 +358,7 @@ fn assert_usage(completion: &Value, expected_completion_tokens: usize) -> Result Ok(()) } -fn assert_logprobs(completion: &Value) -> Result<()> { +fn assert_http_logprobs(completion: &Value) -> Result<()> { let logprobs = &completion["choices"][0]["logprobs"]; if logprobs.is_null() { bail!("completion requested logprobs but response has null logprobs: {completion}"); @@ -344,3 +408,68 @@ fn reserve_loopback_port() -> Result { .context("failed to reserve loopback port for Qwen3.5 TP2 serving test")?; Ok(listener.local_addr()?.port()) } + +fn completion_request(prompt_tokens: Vec, max_tokens: usize, logprobs: usize) -> Request { + let mut request = common::harness::request( + prompt_tokens, + SamplingParams { + ignore_eos: true, + ..SamplingParams::default() + }, + max_tokens, + ); + request.logprobs = logprobs; + request +} + +fn assert_length_outcome(outcome: &Outcome, expected_completion_tokens: usize) { + assert_eq!( + outcome.tokens.len(), + expected_completion_tokens, + "expected {expected_completion_tokens} generated tokens, got {}", + outcome.tokens.len() + ); + let Terminal::Finished { + reason, + prompt_tokens, + completion_tokens, + } = &outcome.terminal + else { + panic!("expected Finished, got {:?}", outcome.terminal); + }; + assert_eq!( + *reason, + FinishReason::Length, + "expected length stop for ignore_eos request" + ); + assert_eq!( + *completion_tokens, expected_completion_tokens, + "expected {expected_completion_tokens} completion tokens, got {completion_tokens}" + ); + assert!(*prompt_tokens > 0, "engine reported zero prompt tokens"); +} + +fn assert_logprobs(outcome: &Outcome) { + assert!( + !outcome.logprobs.is_empty(), + "request asked for logprobs but received none" + ); + assert_eq!( + outcome.logprobs.len(), + outcome.tokens.len(), + "logprobs length must match generated tokens" + ); + assert!( + outcome.logprobs.iter().all(|value| { + value + .as_ref() + .is_some_and(|logprob| logprob.logprob.is_finite()) + }), + "logprobs contain missing or non-finite values: {:?}", + outcome.logprobs + ); +} + +fn cuda_device_count() -> usize { + cudarc::driver::CudaContext::device_count().map_or(0, |n| n.max(0) as usize) +}