diff --git a/docs/index.md b/docs/index.md index 89198a07..ffd72ffd 100644 --- a/docs/index.md +++ b/docs/index.md @@ -31,6 +31,7 @@ Organized by domain (model line / subsystem / playbook / lesson) instead of by l | `models/qwen3/model-crate.md` | `openinfer-qwen3-4b` owns Qwen3 config/weights/executor/scheduler/tests/kernel plan; root sees generic `EngineHandle`; split-K decode gated on `padded_bs<=32` (64-token `Tuned` floor, cap 64 chunks; `Pin`/`PerToken` fixed 160-token split), 4k/64 serving TPOT p50 `6.46ms` on RTX 5090. | | `models/qwen3/prefix-cache.md` | Prefix caching on by default for Qwen3-4B: full-block kvbm radix matching at the executor, suffix-only prefill. Repeated ~1900-token prompt TTFT 141.8 → 16.3ms p50 (8.7×); warm TTFT ≈ TPOT + ~5ms setup. Includes the RoPE scalar-path corruption fix and the drain-the-stream TTFT measurement pitfall. | | `models/qwen3/dflash-speculative-decoding.md` | DFlash speculative decoding behind `--dflash-draft-model-path`, modelled as an optimistic transaction (propose K → verify K+1 span → accept longest argmax prefix + 1 bonus → commit/roll back KV). Lossless up to bf16 tie-flips (bit-identical multi-token accepts; lm-eval gsm8k strict-match identical spec on/off). Single-stream decode 1.82× on 5070 Ti, 1.56× on 5090. Concurrent throughput fixed by batching the draft forward, then a piecewise verify CUDA Graph (dense ops captured, attention eager) closed single-stream: 5090 greedy c1 274 ≈ vLLM 278, c8 1525 > 1240, c16 1834 ≈ 1846 — all batch sizes now ≥ vLLM. Accept measured equal (9.1% vs 8.85%, same drafter); draft-side piecewise graph tracked next. Proposer trait deferred to EAGLE. | +| `models/qwen3/ngram-speculative.md` | Draft-model-free (n-gram / prompt-lookup) speculative decoding for Qwen3-4B behind `--ngram-speculative`, built on the same #436 verify/accept/KV transaction as DFlash — only *propose* differs (host-side longest-suffix lookup over the request's own token history, no worker forward, no draft model). Greedy verification is **lossless** (bit-for-bit identical to plain greedy, modulo bf16 tie flips), GPU-validated 3/3 in `tests/ngram_speculative_gate.rs`. Method state collapsed into one `SpeculativeMethod{ Dflash \| Ngram }` enum; full proposer trait deferred to EAGLE (#325). Single-GPU greedy only; mutually exclusive with DFlash / LoRA / KV-offload / decode-overlap; forces prefix cache off. | | `models/qwen3/accuracy-gate.md` | Qwen3-4B instance of the logits golden gate (`tests/hf_golden_gate.rs`): 48 teacher-forced sequences / 816 positions vs a stored HF bf16 golden, replayed over bs=1 / batched eager / CUDA-graph. Strict guards: regret check + mean ≤ 0.06 + p99 ≤ 0.20; absolute max printed but not asserted (coverage-unstable). Methodology in `subsystems/correctness/`. | | `models/qwen3/kernels-crate.md` | Phase 1 split implemented and 5090-verified: Qwen3-4B kernel surface lives in `openinfer-kernels`; release build, test-target compile, accuracy gate, and bench snapshot pass. | | `models/qwen3/tp-design.md` | Qwen3 tensor-parallel design: `TP=2` milestone scope plus the controller/worker broadcast execution model, request identity, and coarse-grained step protocol for future TP/MoE work. | diff --git a/docs/models/qwen3/ngram-speculative.md b/docs/models/qwen3/ngram-speculative.md new file mode 100644 index 00000000..7ac4eebc --- /dev/null +++ b/docs/models/qwen3/ngram-speculative.md @@ -0,0 +1,200 @@ +# Qwen3-4B n-gram speculative decoding (design) + +**TL;DR**: Draft-model-free speculative decoding for Qwen3-4B. An n-gram +(prompt-lookup) proposer suggests up to `K` continuation tokens from the +request's running context; the target model verifies them in one forward pass +and commits the longest greedy-agreed prefix plus one model token. Greedy +verification is **lossless** — output is bit-for-bit identical to plain greedy +decode, just fewer forward passes on repetitive / structured text. It is a second +*propose* implementation on the shared speculative core (#436/#442): verify, +greedy acceptance, and the KV transaction are reused unchanged; only *propose* +differs. An **acceptance gate** closes speculation when recent draft acceptance +is too low to pay for the verify forward, so low-acceptance traffic (prose) falls +back to plain decode instead of regressing — see *Serving A/B*. + +Last touched: 2026-06 · Enabled with `--ngram-speculative` (off by default). +Proposer, host-side propose, the shared verify/accept/KV transaction, scheduler +wiring, and the acceptance gate are landed and GPU-validated lossless. + +## Shape: one transaction, only *propose* varies + +Speculative decoding is a single optimistic transaction; the only method-specific +part is where the drafts come from. The loaded method is one enum on the +executor (`openinfer-qwen3-4b/src/executor.rs`): + +```rust +enum SpeculativeMethod { + Dflash(DFlashMeta), // worker-side draft-model forward, captures hidden state + Ngram(NgramRuntime), // host-side prompt-lookup, no draft model +} +// Qwen3Executor::speculative_method: Option — single source of truth +``` + +This is the pre-trait consolidation #436 deferred. A full `SpeculativeProposer` +trait object is **not** introduced: DFlash's propose is a worker RPC (needs +`&mut executor` + lane), so it can't be a self-contained proposer without borrow +gymnastics. Enum dispatch is the honest shape until a third method (EAGLE, #325) +makes a shared trait pay off. + +Per step (`scheduler/plan.rs`, `ExecutionPlan::SpeculativeDecode`): + +1. **Propose** — produce up to `K` candidate tokens per request. +2. **Verify** — one target forward over each request's `K + 1` span + `[current, draft_1..draft_K]`, argmax at every position. +3. **Accept** — `accept_greedy` keeps the longest prefix matching the target + argmax, then appends the target's own token at the first mismatch (always + commits `1..=K + 1` tokens). +4. **Commit / roll back** — accepted KV committed; the unused draft tail is + LIFO-released by kvbm (`RequestKv::apply_speculative`). + +Because the draft↔verify boundary is a pure token span, `accept_greedy` / +`build_verify_results`, the verify transaction +(`executor/spec.rs::execute_speculative_verify_impl`), and the +`schedule/apply/revert_speculative` KV lifecycle are **reused unchanged** from +the DFlash core. Only two things are method-specific: where drafts come from, and +whether verify captures hidden states. + +## What n-gram adds + +- **Proposer** (`ngram.rs`, `NgramProposer`) — stateless longest-suffix + prompt-lookup over a request's own token history: tries the longest configured + suffix first, falls back to shorter. Unit-tested in isolation. +- **`NgramRuntime`** (`ngram.rs`) — owns the method's entire state: the proposer, + the acceptance gate (`NgramGate`), and the per-request running context map. The + executor drives it through intent-level calls — `seed`, `append_committed`, + `drop_request`, `is_ready`, `propose_step`, `record_verified` — instead of + maintaining three parallel fields across every execution path. +- **Host-side propose** — `Qwen3Executor::execute_speculative_draft` calls + `NgramRuntime::propose_step` on the executor thread: no worker forward, no draft + model, no readiness handshake. (DFlash instead dispatches + `StepCommand::SpeculativeDraft` to the lane.) +- **No-capture verify** — `LocalQwen3Lane::execute_ngram_verify` reuses + `batch_prefill_into` with **zero capture layers** (a token-only proposer keeps + no model state to seed), then argmax + `build_verify_results`. The worker routes + verify by `lane.dflash.is_some()`; the methods are mutually exclusive. +- **Running context** (`NgramRuntime.ctx`) — per-request token context, seeded at + prefill (`prompt + first_token`), appended at every commit path (verify / plain + decode / fused unified decode, plus a defensive overlap branch), dropped on + retire. Always ends with the request's dangling token so the next draft + continues from it. + +## Enabling it + +`--ngram-speculative` (server flag) → `Qwen3LaunchOptions.ngram_speculative` → +`scheduler::start_qwen3` → `Qwen3Executor::load_ngram_drafter(NgramConfig::from_env())`. + +- Single-GPU greedy only; **requires `--tp-size=1`**; mutually exclusive with + `--dflash-draft-model-path`, `--enable-lora`, `--kv-offload`, and + `--decode-overlap` (validated in `lib.rs::launch` and `server/src/main.rs`). +- Forces the prefix cache off — the verify forward writes each request's own + speculative KV span, so cross-request prefix reuse is unsafe here. + +Proposer knobs (`NgramConfig::from_env`): + +- `OPENINFER_QWEN3_NGRAM_TOKENS=K` — draft count (default 4). +- `OPENINFER_QWEN3_NGRAM_MAX_NGRAM=N` — longest matched suffix (default 3). +- `OPENINFER_QWEN3_NGRAM_ACCEPT_THRESHOLD=f` — acceptance gate threshold (mean + accepted draft tokens/step below which speculation falls back to plain decode; + default 0.3, `0` disables the gate). See *Serving A/B*. + +**Per-request eligibility.** `should_speculative_decode` takes the speculative +path only when *every* active request is greedy (`SamplingParams::is_greedy()`), +asks for no decode logprobs (`logprobs == 0`), and is non-LoRA — a single +ineligible request falls the whole batch back to plain decode for that step. So +enabling speculation never changes a sampled request's output or strips requested +logprobs. A miss (length-1 verify span) already degrades to a plain single-token +decode in `build_verify_results`, so it needs no special fallback. + +## Why it is lossless + +`accept_greedy` keeps only the prefix where `draft[i] == target_argmax[i]`, then +appends one of the target's own tokens, so the committed sequence is exactly what +plain greedy decode would have produced one token at a time. This is a built-in +oracle: **spec-on must equal spec-off** token-for-token under greedy params +(modulo the benign prefill-vs-decode bf16 tie flip). The gate only changes +*whether* we speculate, never the committed tokens, so it preserves this. + +GPU-validated by `tests/ngram_speculative_gate.rs` (3 tests, real Qwen3-4B): +greedy-matches-plain (bs=1), concurrent heterogeneous (bs>1 verify-span path), +and mixed greedy+sampling. The proposer and gate also have isolated unit tests in +`ngram.rs`. + +## Serving A/B (`vllm bench serve`) and the acceptance gate + +Greedy speculation can only win when enough drafts are accepted to offset the +verify forward (which costs ~`K + 1`× a plain decode). Single RTX 4090, Qwen3-4B, +`temperature=0`, `--ignore-eos`, `--no-prefix-cache`, 12 prompts × 128 output +tokens, sonnet (550/128, prefix 200) and random (550/128). `Δ` is n-gram vs plain +decode on the **same** engine; vLLM (0.21 / 0.23, ngram `K=4`, +`prompt_lookup_max=3`) is the reference. + +| dataset / conc | acceptance | vLLM 0.23 Δ | openinfer (no gate) Δ | **openinfer (gate) Δ** | +| --- | --- | --- | --- | --- | +| sonnet c1 | ~5% | −5.8% | −26.2% | **+0.6%** | +| sonnet c4 | ~7% | −5.8% | −40.0% | **−2.2%** | +| random c1 | ~33% | +19.3% | −10.3% | **+21.8%** | +| random c4 | ~37% | +17.0% | −33.8% | **−20.0%** | + +Reading this: + +- **Prose (sonnet) acceptance is ~5–7%** — prompt-lookup drafts are almost never + the model's greedy continuation, so *every* engine loses here (vLLM included). + Without the gate the loss is severe (−40% at c4) because each step pays the + verify forward (a `K+1`-wide prefill-kernel pass) plus, when only some requests + matched, a second plain-decode pass — both wasted. The gate detects the low + acceptance and falls back to plain decode, recovering the regression to ≈0 and + **beating vLLM**, which keeps speculating and eats its −6%. +- **random c1 matches vLLM** (+21.8% vs +19.3%): low concurrency leaves the GPU + compute-idle, so the gate stays open and the ~33% acceptance is a free win. +- **random c4 is the open gap**: vLLM **+17%** vs openinfer **−20%** at the *same* + ~37% acceptance (proven by the matching c1 numbers). This is purely the verify + path — see below. The gate can't help: acceptance is high, so it correctly + stays open; the loss is in *how* we verify, not *whether* we should. + +**Root cause of the random-c4 gap (nsys, `--cuda-graph-trace=node`).** At c4 the +GPU is already ~92% busy on plain decode. Under n-gram the same 22 s window shows +attention running through the **prefill** kernel (`BatchPrefill`, ~23 µs) instead +of the decode kernel (`BatchDecode`, ~15 µs), ~989 verify forwards **plus** ~887 +decode forwards (the batch splits two ways every mixed step), and ~43% more GEMM. +vLLM instead runs **one fused forward**: draft tokens are extra rows in the same +batch (a missed request is just `query_len=1`), verified by one rejection step — +no second pass, no prefill-kernel penalty, no ragged-span CUDA-graph thrash. +Folding our verify into that single unified forward is the way to turn random-c4 +positive; tracked as the follow-up below, not in this PR. + +## The acceptance gate (`NgramGate`) + +`ngram.rs`, owned by `NgramRuntime`. An engine-wide EWMA of accepted draft tokens +per drafted step. `should_draft` opens while in warmup, while the EWMA clears +`accept_threshold`, or on a periodic probe; otherwise `propose_step` produces no +drafts for the step and the existing draft/undraft partition in `plan.rs` routes +the whole batch to plain decode (no verify forward). The probe (every +`PROBE_INTERVAL` steps) re-opens it so a shift into repetitive text is picked back +up. `record_verified` (called from the verify transaction in `spec.rs`) feeds the +step's mean acceptance back in. Lossless either way. Unit-tested (`gate_*` in +`ngram.rs`); the engine-level losslessness gate still passes with the gate live. + +## Remaining work / follow-up + +1. **Single fused verify forward** (the random-c4 fix, separate PR): fold the + verify tokens into one unified batched forward (FlashInfer varlen) with a GPU + rejection step — a missed request becomes a `query_len=1` row in the same + batch — instead of the current `batch_prefill_into`-based verify plus the + second plain-decode pass for undrafted requests. The *Serving A/B* above + quantifies the prize: closing the ~37-point random-c4 gap to vLLM + (+17% vs −20%). + +## Scope / deferred + +Greedy, single-GPU, no LoRA. Deferred: sampling (non-greedy) acceptance (needs +draft distributions + a rejection sampler, not argmax), the unified fused verify +forward above, and a general `SpeculativeProposer` trait (revisit when EAGLE +#325 lands a third method). + +## Prior art + +- vLLM V1 n-gram / prompt-lookup spec decode (`NgramProposer` in the runner, GPU + rejection sampler, scheduler reserves KV for `k` draft tokens). Its single + fused forward is what the follow-up above borrows. +- kvbm `SchedulableSequence` speculative lifecycle (`schedule_speculative` / + `apply_speculative`, LIFO block release), reused unchanged here. diff --git a/openinfer-dynamo-backend/src/engine.rs b/openinfer-dynamo-backend/src/engine.rs index 1e1f8123..6affcafc 100644 --- a/openinfer-dynamo-backend/src/engine.rs +++ b/openinfer-dynamo-backend/src/engine.rs @@ -153,6 +153,7 @@ impl OpeninferBackend { // Speculative decoding is a standalone-server knob; the Dynamo // worker never drafts. dflash_draft_model_path: None, + ngram_speculative: false, enable_kv_events, }; diff --git a/openinfer-qwen3-4b/src/executor.rs b/openinfer-qwen3-4b/src/executor.rs index 694c1522..b1132081 100644 --- a/openinfer-qwen3-4b/src/executor.rs +++ b/openinfer-qwen3-4b/src/executor.rs @@ -501,7 +501,11 @@ fn execute_step_on_lane( // each span position) and captures the target hidden states (at the // DFlash layers) to seed the next draft — all into reused, // pointer-stable scratch (`VerifyGraphBuffers`). - let result = lane.execute_dflash_verify(requests, kv_views)?; + let result = if lane.dflash.is_some() { + lane.execute_dflash_verify(requests, kv_views)? + } else { + lane.execute_ngram_verify(requests, kv_views)? + }; Ok(WorkerStepOutcome::SpeculativeVerify(result)) } StepCommand::SpeculativeDraft { requests } => Ok(WorkerStepOutcome::SpeculativeDraft( @@ -793,6 +797,18 @@ struct Qwen3ExecutorMetadata { config: Config, } +/// The loaded speculative-decode method. Exactly one is ever active; the +/// *propose* step is the only method-specific part (DFlash runs a draft-model +/// forward on the worker; n-gram scans the request's token context host-side), +/// while verify, greedy acceptance, and the KV transaction are shared. +enum SpeculativeMethod { + /// DFlash draft model: worker-side draft forward with hidden-state capture. + Dflash(DFlashMeta), + /// Host-side n-gram / prompt-lookup proposer; needs no draft model. The + /// runtime owns the proposer, the acceptance gate, and per-request context. + Ngram(crate::ngram::NgramRuntime), +} + pub struct Qwen3Executor { metadata: Qwen3ExecutorMetadata, kv_mgr: KvCacheManager, @@ -826,9 +842,12 @@ pub struct Qwen3Executor { /// In-flight async prefill state. Populated by the SplitConcurrent step, /// consumed by `poll_async_prefill`. async_prefill: Option, - /// DFlash draft metadata; `Some` once a draft model is loaded into the - /// primary lane. Speculative decoding is enabled iff this is set. - speculative: Option, + /// The active speculative-decode method and its loaded state — DFlash draft + /// metadata or the n-gram proposer — or `None` when speculation is off. + /// Exactly one method is ever loaded; this is the single source of truth for + /// whether speculation is on and which kind. Only the *propose* step varies + /// between methods; verify, accept, and the KV transaction are shared. + speculative_method: Option, /// Requests whose DFlash context is captured and ready to draft. A request /// enters this set when its prompt finishes prefilling with captured target /// context, and leaves on retire or a plain (non-speculative) decode. @@ -952,7 +971,7 @@ impl Qwen3Executor { l1_retention_disabled: false, overlap: None, async_prefill: None, - speculative: None, + speculative_method: None, dflash_ready_requests: HashSet::new(), kv_events, }) @@ -1181,7 +1200,7 @@ impl Qwen3Executor { l1_retention_disabled: false, overlap: None, async_prefill: None, - speculative: None, + speculative_method: None, dflash_ready_requests: HashSet::new(), // KV events are single-GPU only (asserted above); never wired here. kv_events: None, @@ -1266,6 +1285,30 @@ impl Qwen3Executor { } } + /// The loaded DFlash draft metadata, or `None` under no/other method. + fn dflash_meta(&self) -> Option<&DFlashMeta> { + match &self.speculative_method { + Some(SpeculativeMethod::Dflash(meta)) => Some(meta), + _ => None, + } + } + + /// The n-gram runtime (proposer + gate + per-request context), or `None` + /// under no/other method. + fn ngram_runtime(&self) -> Option<&crate::ngram::NgramRuntime> { + match &self.speculative_method { + Some(SpeculativeMethod::Ngram(runtime)) => Some(runtime), + _ => None, + } + } + + fn ngram_runtime_mut(&mut self) -> Option<&mut crate::ngram::NgramRuntime> { + match &mut self.speculative_method { + Some(SpeculativeMethod::Ngram(runtime)) => Some(runtime), + _ => None, + } + } + /// Enable speculative decoding by loading a DFlash draft model into the /// primary lane. /// @@ -1289,7 +1332,34 @@ impl Qwen3Executor { meta.block_size ); self.prefix_cache_enabled = false; - self.speculative = Some(meta); + self.speculative_method = Some(SpeculativeMethod::Dflash(meta)); + Ok(()) + } + + /// Enable host-side n-gram (prompt-lookup) speculative decoding. Like DFlash + /// it is single-GPU greedy-only and forces the prefix cache off (the verify + /// forward writes each request's own speculative KV span); unlike DFlash it + /// loads no draft model — drafts come from scanning the request's own token + /// context on the executor thread. + pub fn load_ngram_drafter(&mut self, config: crate::ngram::NgramConfig) -> Result<()> { + anyhow::ensure!( + self.workers.is_empty(), + "speculative decoding requires the single-GPU path (got {} extra ranks)", + self.workers.len() + ); + anyhow::ensure!( + self.offload.is_none(), + "speculative decoding is not supported together with KV offload" + ); + anyhow::ensure!( + self.speculative_method.is_none(), + "n-gram speculative decoding cannot be combined with a DFlash draft model" + ); + log::info!("Qwen3 n-gram speculative decoding enabled: {config}"); + self.prefix_cache_enabled = false; + self.speculative_method = Some(SpeculativeMethod::Ngram(crate::ngram::NgramRuntime::new( + config, + ))); Ok(()) } @@ -1607,12 +1677,18 @@ impl ModelExecutor for Qwen3Executor { fn max_context_tokens(&self) -> usize { let target = self.metadata.config.max_position_embeddings; - match &self.speculative { + match self.dflash_meta() { // The draft's fixed-width in-fill block writes `block_size` positions // past the committed length each step, so a request may use at most // `draft.max_pos - block_size` tokens before the draft cache would // overflow. Reject the rest at admission instead of crashing mid-prefill. Some(meta) => target.min(meta.max_position_embeddings.saturating_sub(meta.block_size)), + // n-gram (or no speculation): full window. n-gram has no separate + // draft cache, and its verify span is clamped to the request's + // remaining output budget (see `build_speculative_verify_items`), so + // the furthest KV position it writes is `prompt + generated + span <= + // prompt + max_tokens <= max_position_embeddings` — never past the + // window. No DFlash-style block_size reduction is needed. None => target, } } @@ -1662,7 +1738,10 @@ impl ModelExecutor for Qwen3Executor { } } self.saved_cursor.remove(&request_id); - if self.speculative.is_some() { + if let Some(runtime) = self.ngram_runtime_mut() { + runtime.drop_request(request_id); + } + if self.dflash_meta().is_some() { self.dflash_ready_requests.remove(&request_id); self.primary.drop_dflash_request(request_id)?; } @@ -1849,10 +1928,30 @@ impl ModelExecutor for Qwen3Executor { for req_result in &result.requests { self.apply_prefill_result(req_result)?; } + // Seed the n-gram running context for each freshly-completed prompt: + // the full prompt followed by the first sampled token (the request's + // current dangling token). Later committed tokens are appended at verify + // and decode time so the context always ends with the dangling token. + if let Some(runtime) = self.ngram_runtime_mut() { + for req_result in &result.requests { + if req_result.completed { + if let Some(req) = plan + .requests + .iter() + .find(|r| r.request_id == req_result.request_id) + { + let mut context = Vec::with_capacity(req.prompt_tokens.len() + 1); + context.extend_from_slice(&req.prompt_tokens); + context.push(req_result.first_token); + runtime.seed(req_result.request_id, context); + } + } + } + } // A request becomes draft-ready once its prompt is fully prefilled with // captured target context. Partial chunks stay pending; ineligible // requests drop any stale worker state. - if self.speculative.is_some() { + if self.dflash_meta().is_some() { for req_result in &result.requests { let captured = result .dflash_context_captured_requests @@ -1923,9 +2022,16 @@ impl ModelExecutor for Qwen3Executor { .expect("request must exist after decode"); rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; } + // A plain decode commits one token per request; keep the n-gram context + // in sync so it always ends with the request's dangling token. + if let Some(runtime) = self.ngram_runtime_mut() { + for req_result in &result.requests { + runtime.append_committed(req_result.request_id, &[req_result.token]); + } + } // A plain decode advances the sequence outside the speculative path, so // any captured draft context is now stale — drop it. - if self.speculative.is_some() { + if self.dflash_meta().is_some() { for req_result in &result.requests { if self.dflash_ready_requests.remove(&req_result.request_id) { self.primary.drop_dflash_request(req_result.request_id)?; @@ -1941,6 +2047,14 @@ impl ModelExecutor for Qwen3Executor { } fn execute_speculative_draft(&mut self, plan: DraftPlan<'_>) -> Result { + // N-gram drafts are produced host-side on the executor thread (no worker + // forward): the runtime scans each request's own token context for a + // prompt-lookup match, gated by recent acceptance. DFlash instead runs a + // draft-model forward on the lane. + if let Some(runtime) = self.ngram_runtime_mut() { + let requests = runtime.propose_step(plan.requests)?; + return Ok(DraftResult { requests }); + } self.execute_speculative_draft_impl(plan) } @@ -1949,11 +2063,17 @@ impl ModelExecutor for Qwen3Executor { } fn speculative_enabled(&self) -> bool { - self.speculative.is_some() + self.speculative_method.is_some() } fn speculative_request_ready(&self, request_id: RequestId) -> bool { - self.dflash_ready_requests.contains(&request_id) + if let Some(runtime) = self.ngram_runtime() { + // Ready as soon as the prompt is prefilled and its context seeded; + // n-gram needs no captured hidden state. + runtime.is_ready(request_id) + } else { + self.dflash_ready_requests.contains(&request_id) + } } fn execute_unified(&mut self, plan: UnifiedPlan<'_>) -> Result { @@ -2022,12 +2142,23 @@ impl ModelExecutor for Qwen3Executor { .expect("request must exist after unified decode"); rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; } + // A unified step commits one token per active request outside the + // speculative path; keep the n-gram context in sync (mirroring + // execute_decode) so it stays anchored on the dangling token when + // the batch returns to pure-greedy n-gram speculation. Reachable + // whenever a non-greedy request shares the batch and pushes the + // greedy ones onto the fused decode path. + if let Some(runtime) = self.ngram_runtime_mut() { + for req_result in &result.decode_requests { + runtime.append_committed(req_result.request_id, &[req_result.token]); + } + } // A plain decode via the fused unified step advances the sequence // outside the speculative path, so any captured draft context is // now stale — drop it, mirroring execute_decode. (Eligible pending // are routed to a dedicated prefill step, so unified prefills never // need DFlash mark-ready here.) - if self.speculative.is_some() { + if self.dflash_meta().is_some() { for req_result in &result.decode_requests { if self.dflash_ready_requests.remove(&req_result.request_id) { self.primary.drop_dflash_request(req_result.request_id)?; @@ -2055,6 +2186,15 @@ impl ModelExecutor for Qwen3Executor { .expect("request must exist after split decode"); rkv.apply_decode(req_result.token, self.kv_mgr.pool())?; } + // Defensive n-gram context sync (mirrors the Unified branch). The + // overlap path is unreachable under n-gram today (speculation forces + // decode-overlap off), but keep the context anchored if that + // mutual-exclusion is ever relaxed. + if let Some(runtime) = self.ngram_runtime_mut() { + for req_result in &decode_result.requests { + runtime.append_committed(req_result.request_id, &[req_result.token]); + } + } for req_result in &decode_result.requests { self.save_sealed_blocks(req_result.request_id); } @@ -2660,6 +2800,65 @@ impl LocalQwen3Lane { result } + /// Verify forward for host-side proposers (n-gram): one target forward over + /// each request's `K + 1` span with its speculative KV view, returning the + /// greedy argmax per position. Unlike [`Self::execute_dflash_verify`] it + /// captures no hidden states — a token-only proposer keeps no model state to + /// seed — so the verify graph runs with zero capture layers. + fn execute_ngram_verify( + &mut self, + requests: &[VerifyStepItem], + kv_views: &[KvView], + ) -> Result { + let span = requests + .iter() + .map(|req| req.as_slice().len()) + .max() + .unwrap_or(1) + .max(1); + // (Re)allocate the verify scratch if it is missing or too narrow for this + // step's widest span. The widest span grows from 1 (no match) up to K + 1 + // as matches lengthen, so this may reallocate a few times early on before + // settling at the K + 1 ceiling — never per step in steady state. + let needs_alloc = self + .verify_bufs + .as_ref() + .map_or(true, |bufs| bufs.span() < span); + if needs_alloc { + let max_batch = *BATCH_BUCKETS.last().unwrap(); + self.verify_bufs = Some(VerifyGraphBuffers::new( + &self.model, + max_batch, + span, + 0, + self.total_blocks, + )?); + } + + let mut bufs = self.verify_bufs.take().expect("verify buffers just set"); + let result = (|| -> Result { + let spans: Vec<&[u32]> = requests.iter().map(VerifyStepItem::as_slice).collect(); + self.model.batch_prefill_into( + &spans, + kv_views, + self.kv_buffer.buffer(), + &self.layout, + &[], + &mut bufs, + )?; + let total_tokens: usize = requests.iter().map(|req| req.as_slice().len()).sum(); + let greedy = SamplingParams::default(); + let params: Vec<&SamplingParams> = vec![&greedy; total_tokens]; + let target_tokens = self.select_step_tokens(bufs.all_logits(), ¶ms, 0)?; + let request_results = build_verify_results(requests, &target_tokens)?; + Ok(VerifyResult { + requests: request_results, + }) + })(); + self.verify_bufs = Some(bufs); + result + } + fn execute_decode( &mut self, token_ids: &[u32], diff --git a/openinfer-qwen3-4b/src/executor/spec.rs b/openinfer-qwen3-4b/src/executor/spec.rs index 0ec42881..9c469f3a 100644 --- a/openinfer-qwen3-4b/src/executor/spec.rs +++ b/openinfer-qwen3-4b/src/executor/spec.rs @@ -17,8 +17,8 @@ impl Qwen3Executor { plan: VerifyPlan<'_>, ) -> Result { anyhow::ensure!( - self.speculative.is_some(), - "speculative verification requested but no draft model is loaded" + self.speculative_method.is_some(), + "speculative verification requested but no speculative method is enabled" ); for req in plan.requests { anyhow::ensure!( @@ -30,11 +30,13 @@ impl Qwen3Executor { req.params.is_greedy(), "speculative verification currently supports greedy sampling only" ); - anyhow::ensure!( - self.dflash_ready_requests.contains(&req.request_id), - "speculative verification requested before DFlash state is ready for {:?}", - req.request_id - ); + if self.dflash_meta().is_some() { + anyhow::ensure!( + self.dflash_ready_requests.contains(&req.request_id), + "speculative verification requested before DFlash state is ready for {:?}", + req.request_id + ); + } anyhow::ensure!( self.request_kvs.contains_key(&req.request_id), "missing RequestKv for {:?}", @@ -134,6 +136,22 @@ impl Qwen3Executor { self.save_sealed_blocks(req_result.request_id); } + // Keep the n-gram running context in sync: append every committed token + // so the next draft scans the full history and continues from the new + // dangling token. DFlash keeps its own per-request hidden state instead. + // This verify forward ran because the gate was open, so also fold its + // mean accepted-draft-tokens into the gate (`result.requests` is the + // drafted subset, hence non-empty here). + if let Some(runtime) = self.ngram_runtime_mut() { + for req_result in &result.requests { + runtime.append_committed(req_result.request_id, &req_result.accepted_tokens); + } + if !result.requests.is_empty() { + let accepted: usize = result.requests.iter().map(|r| r.matched_draft_tokens).sum(); + runtime.record_verified(accepted as f32 / result.requests.len() as f32); + } + } + Ok(result) } @@ -142,7 +160,7 @@ impl Qwen3Executor { plan: DraftPlan<'_>, ) -> Result { anyhow::ensure!( - self.speculative.is_some(), + self.dflash_meta().is_some(), "speculative draft requested but no draft model is loaded" ); for req in plan.requests { diff --git a/openinfer-qwen3-4b/src/lib.rs b/openinfer-qwen3-4b/src/lib.rs index e7eacba8..56d6dfa3 100644 --- a/openinfer-qwen3-4b/src/lib.rs +++ b/openinfer-qwen3-4b/src/lib.rs @@ -10,6 +10,7 @@ mod executor; pub(crate) mod green_ctx; pub mod kernel_bench; mod lora; +mod ngram; mod prefill; mod scheduler; mod speculative; @@ -177,6 +178,10 @@ pub struct Qwen3LaunchOptions { /// `Some` enables DFlash speculative decoding with this drafter model. /// Single-GPU only and mutually exclusive with LoRA and KV offload. pub dflash_draft_model_path: Option, + /// Enables host-side n-gram (prompt-lookup) speculative decoding. Needs no + /// draft model; single-GPU greedy only and mutually exclusive with DFlash, + /// LoRA, KV offload, and decode overlap. + pub ngram_speculative: bool, /// Publish KV block store/remove events for an out-of-band cache-aware /// router (e.g. a Dynamo KV router). Off for plain single-machine serving; /// single-GPU + base-model only (rejected with LoRA or tensor parallel). @@ -226,6 +231,22 @@ pub fn launch(model_path: &Path, options: Qwen3LaunchOptions) -> Result Result Result, + ngram_speculative: bool, enable_kv_events: bool, ) -> Result { let EngineLoadOptions { @@ -330,6 +354,7 @@ pub fn start_engine_with_offload( memory_options, decode_overlap, dflash_draft_model_path, + ngram_speculative, enable_kv_events, ) } diff --git a/openinfer-qwen3-4b/src/ngram.rs b/openinfer-qwen3-4b/src/ngram.rs new file mode 100644 index 00000000..45487927 --- /dev/null +++ b/openinfer-qwen3-4b/src/ngram.rs @@ -0,0 +1,479 @@ +//! N-gram (prompt-lookup) speculative proposer for Qwen3. +//! +//! This proposes candidate continuation tokens **without a draft model**: it +//! looks for the most recent earlier occurrence of the current token suffix in +//! the running context (prompt + tokens generated so far) and proposes the +//! tokens that followed that occurrence. The target model then verifies these +//! candidates in a single forward pass and accepts the longest matching prefix. +//! +//! This is effective for repetitive / structured text (code, quoting, JSON, +//! long-form copying) at zero extra model cost. When no n-gram match is found +//! it proposes nothing and decoding falls back to the normal one-token step. +//! +//! Only the proposer is implemented here; verification and KV rollback are +//! handled by the decode path and are intentionally out of scope for this +//! module so the lookup logic can be unit-tested in isolation. + +use std::collections::HashMap; + +use anyhow::Result; + +use crate::executor::RequestId; +use crate::speculative::{DraftRequestResult, DraftStepItem}; + +/// Configuration for [`NgramProposer`]. +#[derive(Clone, Copy, Debug, PartialEq)] +pub struct NgramConfig { + /// Largest suffix length to match. Longer suffixes are tried first because + /// they are more specific (a longer match is a stronger predictor). + pub max_ngram: usize, + /// Smallest suffix length to match before giving up. Must be `>= 1`. + pub min_ngram: usize, + /// Maximum number of speculative tokens proposed per step. + pub num_speculative: usize, + /// Mean accepted-draft-tokens-per-step below which speculation is judged not + /// worth its verify forward, so the engine falls back to plain decode (see + /// [`NgramGate`]). On non-repetitive text (e.g. prose) prompt-lookup drafts + /// are almost never accepted, and at higher concurrency the wasted verify + /// compute is a net throughput loss; gating recovers it. `0.0` disables the + /// gate (always speculate). + pub accept_threshold: f32, +} + +impl Default for NgramConfig { + fn default() -> Self { + Self { + max_ngram: 3, + min_ngram: 1, + num_speculative: 4, + accept_threshold: 0.3, + } + } +} + +impl std::fmt::Display for NgramConfig { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "K={}, max_ngram={}, accept_threshold={}", + self.num_speculative, self.max_ngram, self.accept_threshold + ) + } +} + +impl NgramConfig { + /// Read the n-gram-specific knobs from the environment, falling back to + /// [`Default`] for anything unset. The generic on/off switch lives in + /// [`crate::speculative::SpeculativeConfig::from_env`]; this only owns the + /// proposer's own parameters: + /// + /// * `OPENINFER_QWEN3_NGRAM_TOKENS` = draft count `K`. + /// * `OPENINFER_QWEN3_NGRAM_MAX_NGRAM` = longest suffix to match. + /// * `OPENINFER_QWEN3_NGRAM_ACCEPT_THRESHOLD` = gate threshold (`0` disables). + #[must_use] + pub fn from_env() -> Self { + let defaults = Self::default(); + Self { + max_ngram: env_usize("OPENINFER_QWEN3_NGRAM_MAX_NGRAM").unwrap_or(defaults.max_ngram), + min_ngram: defaults.min_ngram, + num_speculative: env_usize("OPENINFER_QWEN3_NGRAM_TOKENS") + .unwrap_or(defaults.num_speculative), + accept_threshold: env_f32("OPENINFER_QWEN3_NGRAM_ACCEPT_THRESHOLD") + .unwrap_or(defaults.accept_threshold), + } + } +} + +fn env_usize(name: &str) -> Option { + std::env::var(name).ok().and_then(|v| v.parse().ok()) +} + +fn env_f32(name: &str) -> Option { + std::env::var(name).ok().and_then(|v| v.parse().ok()) +} + +/// Stateless n-gram / prompt-lookup proposer. +/// +/// The proposer keeps no history of its own; each call scans the supplied +/// context, so the caller is free to reuse a single instance across requests. +#[derive(Clone, Copy, Debug)] +pub(crate) struct NgramProposer { + config: NgramConfig, +} + +impl NgramProposer { + /// Create a proposer. `min_ngram` is clamped to at least 1 and to at most + /// `max_ngram`, and `max_ngram` to at least 1, so the config is always + /// usable regardless of caller input. + #[must_use] + pub(crate) fn new(config: NgramConfig) -> Self { + let max_ngram = config.max_ngram.max(1); + let min_ngram = config.min_ngram.clamp(1, max_ngram); + Self { + config: NgramConfig { + max_ngram, + min_ngram, + ..config + }, + } + } + + /// The (clamped) configuration this proposer was built with. + #[must_use] + pub(crate) fn config(&self) -> NgramConfig { + self.config + } + + /// Propose up to `num_speculative` continuation tokens for `context` + /// (the full token sequence so far: prompt followed by generated tokens). + /// + /// Returns an empty `Vec` when no usable n-gram match exists or when + /// `num_speculative == 0`. Tries the longest configured suffix first and + /// falls back to shorter suffixes. + /// + /// Cost: each call is a linear reverse scan of `context` per suffix length, + /// so a miss is `O(context_len * max_ngram)` on the executor thread, every + /// speculative step. Fine for the target case (repetitive / bounded + /// contexts); a very long context could erode the speedup it buys. If that + /// becomes a concern, bound the look-back window (as vLLM's prompt-lookup + /// does) rather than scanning the whole history. + #[must_use] + pub(crate) fn propose(&self, context: &[u32]) -> Vec { + if self.config.num_speculative == 0 { + return Vec::new(); + } + let len = context.len(); + // A match plus at least one following token needs `n + 1` tokens. + let max_n = self.config.max_ngram.min(len.saturating_sub(1)); + for n in (self.config.min_ngram..=max_n).rev() { + let suffix = &context[len - n..]; + if let Some(start) = latest_earlier_match(context, suffix) { + let pred_start = start + n; + let end = (pred_start + self.config.num_speculative).min(len); + debug_assert!(pred_start < end, "match must leave a following token"); + return context[pred_start..end].to_vec(); + } + } + Vec::new() + } +} + +/// Find the latest start index `i < len - n` such that `context[i..i + n]` +/// equals `suffix` (the trailing `n` tokens). "Latest" prefers the most recent +/// context. Guaranteed to leave at least one following token at `i + n`. +fn latest_earlier_match(context: &[u32], suffix: &[u32]) -> Option { + let n = suffix.len(); + let len = context.len(); + if n == 0 || len < n + 1 { + return None; + } + // Candidate starts run from `len - n - 1` (immediately before the trailing + // suffix) down to 0; the trailing occurrence itself (start `len - n`) is + // excluded so we never "predict" from the suffix we are matching. + (0..len - n).rev().find(|&i| &context[i..i + n] == suffix) +} + +/// Engine-wide switch that decides whether n-gram speculation is currently +/// paying for itself. +/// +/// Speculation is only a win when enough drafted tokens are accepted to offset +/// the verify forward (which costs ~`K + 1`× a plain decode). On repetitive text +/// acceptance is high and this stays open; on prose it collapses to near zero, so +/// the gate closes and decoding falls back to plain — recovering the throughput +/// the wasted verify would otherwise burn, which matters most at the higher batch +/// sizes where the GPU has no spare compute to hide it. +/// +/// The estimate is a single EWMA of accepted draft tokens per step, updated only +/// when we actually drafted. While closed the gate **probes** once every +/// `PROBE_INTERVAL` steps so a shift into repetitive text re-opens it. An initial +/// warmup keeps it open until there is enough data to judge. +#[derive(Clone, Copy, Debug)] +pub(crate) struct NgramGate { + /// EWMA of accepted draft tokens per drafted step. + accept_ewma: f32, + /// Drafted steps observed so far (capped at `WARMUP_STEPS`). + warmup_left: u32, + /// Steps elapsed since the gate last allowed a draft (drives probing). + since_draft: u32, +} + +impl NgramGate { + /// Smoothing for the acceptance EWMA: low enough to react within a handful + /// of steps when a request switches between repetitive and prose regions. + const EWMA_ALPHA: f32 = 0.2; + /// Always speculate for this many drafted steps before trusting the EWMA. + const WARMUP_STEPS: u32 = 8; + /// While closed, allow one probing draft every this many steps. + const PROBE_INTERVAL: u32 = 32; + + #[must_use] + pub(crate) fn new() -> Self { + Self { + accept_ewma: 0.0, + warmup_left: Self::WARMUP_STEPS, + since_draft: 0, + } + } + + /// Whether this step should draft. Open during warmup, while recent + /// acceptance clears the threshold, or on a periodic probe. A non-positive + /// threshold disables gating (always open). + #[must_use] + pub(crate) fn should_draft(&self, threshold: f32) -> bool { + threshold <= 0.0 + || self.warmup_left > 0 + || self.accept_ewma >= threshold + || self.since_draft >= Self::PROBE_INTERVAL + } + + /// Record the outcome of a step that drafted: `accepted` is the mean number + /// of draft tokens verify took across the step's drafted requests. Resets the + /// probe clock and folds the count into the EWMA. + pub(crate) fn record_drafted(&mut self, accepted: f32) { + self.since_draft = 0; + self.warmup_left = self.warmup_left.saturating_sub(1); + self.accept_ewma += Self::EWMA_ALPHA * (accepted - self.accept_ewma); + } + + /// Record a step that did not draft (gate closed, or no match found), so the + /// probe clock advances toward the next re-check. + pub(crate) fn record_skipped(&mut self) { + self.since_draft = self.since_draft.saturating_add(1); + } +} + +/// Owns all per-engine n-gram speculative state: the stateless [`NgramProposer`], +/// the [`NgramGate`], and each request's running token context (prompt + every +/// committed token, always ending on the dangling token the next draft continues +/// from). This is the single home for the method's state so the executor drives +/// it through intent-level calls (`seed` / `append_committed` / `drop_request` / +/// `propose_step` / `record_verified`) instead of threading three fields through +/// every execution path. Held inside [`crate::executor::SpeculativeMethod::Ngram`]. +pub(crate) struct NgramRuntime { + proposer: NgramProposer, + gate: NgramGate, + /// Per-request running context. Present once a request's prompt is prefilled + /// (see [`Self::seed`]); the proposer scans it for prompt-lookup matches. + ctx: HashMap>, +} + +impl NgramRuntime { + #[must_use] + pub(crate) fn new(config: NgramConfig) -> Self { + Self { + proposer: NgramProposer::new(config), + gate: NgramGate::new(), + ctx: HashMap::new(), + } + } + + /// Seed a freshly-prefilled request's context (`prompt + first_token`). Later + /// commits extend it via [`Self::append_committed`]. + pub(crate) fn seed(&mut self, request_id: RequestId, context: Vec) { + self.ctx.insert(request_id, context); + } + + /// Append committed tokens to a request's context, keeping it anchored on the + /// new dangling token. No-op for an unseeded request (defensive). + pub(crate) fn append_committed(&mut self, request_id: RequestId, tokens: &[u32]) { + if let Some(context) = self.ctx.get_mut(&request_id) { + context.extend_from_slice(tokens); + } + } + + /// Forget a retired request's context. + pub(crate) fn drop_request(&mut self, request_id: RequestId) { + self.ctx.remove(&request_id); + } + + /// Whether a request can speculate: its prompt is prefilled and context + /// seeded (n-gram keeps no captured hidden state, so that is all it needs). + #[must_use] + pub(crate) fn is_ready(&self, request_id: RequestId) -> bool { + self.ctx.contains_key(&request_id) + } + + /// Build the verify spans for one speculative step. The gate is consulted + /// once: when closed, no request drafts, so every span is just + /// `[current_token]` and the scheduler routes the whole batch to plain decode. + /// Advances the gate's probe clock when the step produced no draft. + pub(crate) fn propose_step( + &mut self, + items: &[DraftStepItem], + ) -> Result> { + let drafting = self + .gate + .should_draft(self.proposer.config().accept_threshold); + let mut any_drafted = false; + let mut results = Vec::with_capacity(items.len()); + for item in items { + anyhow::ensure!( + item.params.is_greedy(), + "n-gram speculative draft currently supports greedy sampling only" + ); + let drafts = if drafting { + let context = self.ctx.get(&item.request_id).ok_or_else(|| { + anyhow::anyhow!("missing n-gram context for {:?}", item.request_id) + })?; + self.proposer.propose(context) + } else { + Vec::new() + }; + any_drafted |= !drafts.is_empty(); + // Verify span = [current dangling token, draft_1, …, draft_K]. + let mut token_ids = Vec::with_capacity(1 + drafts.len()); + token_ids.push(item.current_token); + token_ids.extend(drafts); + results.push(DraftRequestResult { + request_id: item.request_id, + token_ids, + }); + } + // No draft produced (gate closed, or open but every request missed) means + // no verify forward runs this step, so advance the probe clock. A step + // that drafted updates the acceptance estimate via `record_verified`. + if !any_drafted { + self.gate.record_skipped(); + } + Ok(results) + } + + /// Fold a verify forward's mean accepted-draft-tokens into the gate. Called + /// only when a verify ran, i.e. the gate was open and at least one request + /// drafted. + pub(crate) fn record_verified(&mut self, mean_accepted: f32) { + self.gate.record_drafted(mean_accepted); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn proposer(max_ngram: usize, min_ngram: usize, k: usize) -> NgramProposer { + NgramProposer::new(NgramConfig { + max_ngram, + min_ngram, + num_speculative: k, + ..NgramConfig::default() + }) + } + + #[test] + fn proposes_continuation_after_recent_match() { + // ...1 2 3 1 2 3 1 2 ; suffix [1,2] last matched at index 3 -> follow [3,1,2] + let ctx = [1u32, 2, 3, 1, 2, 3, 1, 2]; + let got = proposer(2, 1, 4).propose(&ctx); + assert_eq!(got, vec![3, 1, 2]); + } + + #[test] + fn prefers_longest_suffix_match() { + // Both [9] (len1) and [4,9] (len2) recur, but the len-2 match is more + // specific and must win. + let ctx = [4u32, 9, 7, 1, 9, 0, 4, 9]; + // suffix len2 = [4,9] earlier at index 0 -> follow [7,1,...] + let got = proposer(2, 1, 2).propose(&ctx); + assert_eq!(got, vec![7, 1]); + } + + #[test] + fn falls_back_to_shorter_suffix() { + // No repeated 2-gram suffix, but the last token 5 occurred earlier. + let ctx = [5u32, 8, 2, 7, 5]; + // suffix len2 = [7,5] has no earlier match; len1 = [5] matches at 0 -> [8] + let got = proposer(2, 1, 3).propose(&ctx); + assert_eq!(got, vec![8, 2, 7]); + } + + #[test] + fn returns_empty_when_no_match() { + let ctx = [1u32, 2, 3, 4, 5]; + assert!(proposer(3, 1, 4).propose(&ctx).is_empty()); + } + + #[test] + fn caps_proposal_at_num_speculative() { + let ctx = [1u32, 2, 1, 2, 1, 2, 1, 2]; + let got = proposer(2, 1, 1).propose(&ctx); + assert_eq!(got.len(), 1); + } + + #[test] + fn handles_short_context() { + assert!(proposer(3, 1, 4).propose(&[]).is_empty()); + assert!(proposer(3, 1, 4).propose(&[7]).is_empty()); + } + + #[test] + fn zero_speculative_proposes_nothing() { + let ctx = [1u32, 2, 1, 2]; + assert!(proposer(2, 1, 0).propose(&ctx).is_empty()); + } + + #[test] + fn new_clamps_invalid_config() { + // min_ngram > max_ngram and zero max are clamped, not panicked. + let p = NgramProposer::new(NgramConfig { + max_ngram: 0, + min_ngram: 5, + num_speculative: 2, + ..NgramConfig::default() + }); + // With max_ngram clamped to 1, a single-token recurrence still works: + // suffix [3] recurs at index 0, so the follower [1, 3] is proposed. + let ctx = [3u32, 1, 3]; + assert_eq!(p.propose(&ctx), vec![1, 3]); + } + + #[test] + fn gate_disabled_with_nonpositive_threshold() { + let mut gate = NgramGate::new(); + // Drive acceptance to zero; a zero threshold must keep speculating. + for _ in 0..50 { + gate.record_drafted(0.0); + } + assert!(gate.should_draft(0.0)); + } + + #[test] + fn gate_closes_after_warmup_on_low_acceptance() { + let mut gate = NgramGate::new(); + // Warmup keeps it open even while acceptance reads zero. + for _ in 0..NgramGate::WARMUP_STEPS { + assert!(gate.should_draft(0.3)); + gate.record_drafted(0.0); + } + // Warmup spent + EWMA below threshold -> closed. + assert!(!gate.should_draft(0.3)); + } + + #[test] + fn gate_stays_open_on_high_acceptance() { + let mut gate = NgramGate::new(); + for _ in 0..20 { + gate.record_drafted(3.0); + } + assert!(gate.should_draft(0.3)); + } + + #[test] + fn gate_probes_after_cooldown_then_recovers() { + let mut gate = NgramGate::new(); + for _ in 0..NgramGate::WARMUP_STEPS { + gate.record_drafted(0.0); + } + assert!( + !gate.should_draft(0.3), + "closed after low-acceptance warmup" + ); + // Skipping for PROBE_INTERVAL steps re-opens it for one probe. + for _ in 0..NgramGate::PROBE_INTERVAL { + gate.record_skipped(); + } + assert!(gate.should_draft(0.3), "probe re-opens the gate"); + // A probe that finds repetitive text (high acceptance) keeps it open. + gate.record_drafted(4.0); + assert!(gate.should_draft(0.3)); + } +} diff --git a/openinfer-qwen3-4b/src/scheduler.rs b/openinfer-qwen3-4b/src/scheduler.rs index fb679522..4183725f 100644 --- a/openinfer-qwen3-4b/src/scheduler.rs +++ b/openinfer-qwen3-4b/src/scheduler.rs @@ -150,6 +150,7 @@ pub(crate) fn start_qwen3( memory_options: Qwen3MemoryOptions, decode_overlap: crate::DecodeOverlap, dflash_draft_model_path: Option<&str>, + ngram_speculative: bool, enable_kv_events: bool, ) -> Result { let mut executor = Qwen3Executor::from_runtime_with_lora_options( @@ -171,6 +172,11 @@ pub(crate) fn start_qwen3( // was already reserved during profiling from the draft path passed above. if let Some(draft_path) = dflash_draft_model_path { executor.load_dflash_draft_model(draft_path)?; + } else if ngram_speculative { + // Host-side n-gram drafter: no draft model to load, just enable the + // proposer and force the prefix cache off (same as DFlash) so verify + // writes each request's own speculative KV span. + executor.load_ngram_drafter(crate::ngram::NgramConfig::from_env())?; } Ok(start_with_executor(executor, seed, max_prefill_tokens)) diff --git a/openinfer-qwen3-4b/src/scheduler/plan.rs b/openinfer-qwen3-4b/src/scheduler/plan.rs index 19b1bf7e..e295ece3 100644 --- a/openinfer-qwen3-4b/src/scheduler/plan.rs +++ b/openinfer-qwen3-4b/src/scheduler/plan.rs @@ -6,7 +6,8 @@ use crate::executor::{ PrefillStepItem, UnifiedPlan, UnifiedResult, }; use crate::speculative::{ - DraftPlan, DraftRequestResult, DraftStepItem, VerifyPlan, VerifyResult, VerifyStepItem, + DraftPlan, DraftRequestResult, DraftStepItem, VerifyPlan, VerifyRequestResult, VerifyResult, + VerifyStepItem, }; use super::{ActiveRequestState, PendingRequest}; @@ -113,20 +114,65 @@ pub(super) fn execute_plan( Ok(ExecutionArtifacts::Decode { result }) } ExecutionPlan::SpeculativeDecode => { - // Two executor calls per step: draft proposes K tokens per request, - // then a single target forward verifies the K+1 span. Both index by - // request_id; sorting keeps draft and verify results aligned. + // Propose first, then route each request by whether it actually got a + // draft. Drafted requests (verify span > 1) take the speculative + // verify forward. Requests with no draft (n-gram miss, or K = 0) would + // otherwise pay a full verify forward for a single token — swapping the + // optimized decode kernel for the prefill-style verify kernel, which + // both costs more and perturbs near-tie greedy picks by the + // prefill-vs-decode bf16 gap. Route those through the plain decode path + // instead: same kernel and cost as baseline, committing one token each. + // All-miss steps (common on non-repetitive text) thus degrade exactly + // to plain decode. All index by request_id; sorting keeps results + // aligned for resolve. let draft_requests = build_speculative_draft_items(active); let mut draft = executor.execute_speculative_draft(DraftPlan { requests: &draft_requests, })?; draft.requests.sort_by_key(|result| result.request_id); let verify_requests = build_speculative_verify_items(active, &draft.requests); - let mut verify = executor.execute_speculative_verify(VerifyPlan { - requests: &verify_requests, - })?; - verify.requests.sort_by_key(|result| result.request_id); - Ok(ExecutionArtifacts::SpeculativeDecode { verify }) + let (drafted, undrafted): (Vec, Vec) = verify_requests + .into_iter() + .partition(|item| item.token_ids.len() > 1); + + let mut request_results: Vec = Vec::with_capacity(active.len()); + if !drafted.is_empty() { + let mut verify = + executor.execute_speculative_verify(VerifyPlan { requests: &drafted })?; + request_results.append(&mut verify.requests); + } + if !undrafted.is_empty() { + // logprobs are off and LoRA is absent on the speculative path + // (should_speculative_decode requires both), so a minimal decode + // item is faithful to these requests. + let decode_items: Vec = undrafted + .iter() + .map(|item| DecodeStepItem { + request_id: item.request_id, + token_id: item.token_ids[0], + params: item.params, + logprobs: 0, + lora_adapter: None, + }) + .collect(); + let decode = executor.execute_decode(DecodePlan { + requests: &decode_items, + sample_seed: rand::RngExt::random(rng), + })?; + for result in decode.requests { + request_results.push(VerifyRequestResult { + request_id: result.request_id, + matched_draft_tokens: 0, + accepted_tokens: vec![result.token], + }); + } + } + request_results.sort_by_key(|result| result.request_id); + Ok(ExecutionArtifacts::SpeculativeDecode { + verify: VerifyResult { + requests: request_results, + }, + }) } ExecutionPlan::Unified { pending } => { let scheduled_at_unix_s = openinfer_core::engine::unix_now_s(); diff --git a/openinfer-qwen3-4b/src/verify_graph.rs b/openinfer-qwen3-4b/src/verify_graph.rs index 25c4e626..0e384477 100644 --- a/openinfer-qwen3-4b/src/verify_graph.rs +++ b/openinfer-qwen3-4b/src/verify_graph.rs @@ -157,6 +157,11 @@ impl VerifyGraphBuffers { pub(crate) fn captured_hidden(&self) -> &HiddenStates { &self.captured_hidden } + + /// The fixed per-request span width these buffers were allocated for. + pub(crate) fn span(&self) -> usize { + self.span + } } impl Qwen3Model { diff --git a/openinfer-qwen3-4b/tests/batch_invariance_reject.rs b/openinfer-qwen3-4b/tests/batch_invariance_reject.rs index 01f04546..26ee842b 100644 --- a/openinfer-qwen3-4b/tests/batch_invariance_reject.rs +++ b/openinfer-qwen3-4b/tests/batch_invariance_reject.rs @@ -29,6 +29,7 @@ fn batch_invariant_rejects_decode_overlap() { true, None, false, + false, ) .err() .expect("--batch-invariant + --decode-overlap must be rejected"); diff --git a/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs b/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs index 20de5aa6..b28d57fb 100644 --- a/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs +++ b/openinfer-qwen3-4b/tests/dflash_speculative_gate.rs @@ -116,6 +116,7 @@ fn launch_options(draft: Option) -> Qwen3LaunchOptions { decode_overlap: DecodeOverlap::Off, batch_invariant: false, dflash_draft_model_path: draft, + ngram_speculative: false, enable_kv_events: false, } } diff --git a/openinfer-qwen3-4b/tests/dflash_speculative_perf.rs b/openinfer-qwen3-4b/tests/dflash_speculative_perf.rs index 9b31fbf8..fea3fb0b 100644 --- a/openinfer-qwen3-4b/tests/dflash_speculative_perf.rs +++ b/openinfer-qwen3-4b/tests/dflash_speculative_perf.rs @@ -66,6 +66,7 @@ fn launch_options(draft: Option) -> Qwen3LaunchOptions { decode_overlap: DecodeOverlap::Off, batch_invariant: false, dflash_draft_model_path: draft, + ngram_speculative: false, enable_kv_events: false, } } diff --git a/openinfer-qwen3-4b/tests/ngram_speculative_gate.rs b/openinfer-qwen3-4b/tests/ngram_speculative_gate.rs new file mode 100644 index 00000000..7ef4f087 --- /dev/null +++ b/openinfer-qwen3-4b/tests/ngram_speculative_gate.rs @@ -0,0 +1,649 @@ +//! N-gram (prompt-lookup) speculative-decoding losslessness gate. +//! +//! Greedy speculative decoding must be *lossless*: every draft is verified by a +//! target forward and only the matching-argmax prefix (plus one bonus) is +//! committed, so the accepted tokens are the target model's own greedy +//! continuation. The catch is pure numerics — the verify path runs the +//! *prefill* attention kernel over the `K + 1` span while a plain decode runs +//! the *decode* kernel, and the two differ by ~1 bf16 ULP. On a near-tie that +//! flips one argmax, and from there two greedy runs fan out completely. +//! +//! So an exact `spec == baseline` token match is the wrong gate: it false-fails +//! on a benign tie flip. We use the same *regret* test as the DFlash gate (and +//! `hf_golden_gate`): at the first position the two sequences disagree (where +//! they still share an identical context, so the comparison is valid) we ask how +//! far below the argmax the speculative pick sits — measured *in the prefill +//! kernel's own distribution*, because that is the kernel the verify path runs. +//! A re-prefill of the shared context (`prefill_next`) gives that reference. +//! Within `MARGIN_TOL` of the prefill argmax ⇒ a benign numerical tie. Clearly +//! worse (or outside the prefill top-K) ⇒ the verify/accept logic chose a token +//! the forward never favored — a real bug. +//! +//! Unlike the DFlash gate this needs no draft model: n-gram drafts come from +//! scanning each request's own token context, so the only knob is +//! `--ngram-speculative`. Prompts are repetitive/structured so the proposer +//! actually fires (its suffix recurs earlier) and the verify path is exercised; +//! losslessness holds regardless, but a no-op proposer would make the gate +//! vacuous. +//! +//! Requires a CUDA GPU and Qwen3-4B weights; skips cleanly when absent (set +//! `OPENINFER_TEST_MODEL_PATH` to the weights to run it). + +use std::path::Path; +use std::time::Duration; + +use openinfer_core::engine::{EngineHandle, GenerateRequest, TokenEvent, TokenSink}; +use openinfer_core::sampler::SamplingParams; +use openinfer_qwen3_4b::{ + DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES, DEFAULT_MAX_PREFILL_TOKENS, DecodeOverlap, + Qwen3LaunchOptions, Qwen3MemoryOptions, Qwen3OffloadOptions, +}; +use vllm_text::tokenizer::DynTokenizer; + +mod common; + +const MODEL_PATH: &str = concat!(env!("CARGO_MANIFEST_DIR"), "/../models/Qwen3-4B"); +const GENERATED_TOKENS: usize = 64; +/// Top-K logprobs requested from the baseline; wide enough that the speculative +/// pick is in the set on any real tie (a pick outside the top-K is itself a red +/// flag the gate should catch). +const LOGPROBS: usize = 20; +/// Max acceptable regret: how far below the baseline's argmax (in the baseline's +/// own logprobs) the speculative pick may sit at the divergence point. ~3 bf16 +/// ULP at typical logit magnitudes — mirrors `hf_golden_gate`'s `MARGIN_TOL`. +const MARGIN_TOL: f32 = 0.20; + +/// Both tests launch a Qwen3-4B engine, and two at once overflow a 16 GB card. +/// Cargo runs tests in one binary concurrently, so serialize the engine-holding +/// bodies — only one engine is ever resident on the GPU. +static GPU: std::sync::Mutex<()> = std::sync::Mutex::new(()); + +fn target_path_or_skip() -> Option { + match std::env::var("OPENINFER_TEST_MODEL_PATH") { + Ok(path) => Some(path), + Err(_) if Path::new(MODEL_PATH).join("config.json").exists() => { + Some(MODEL_PATH.to_string()) + } + Err(_) => { + eprintln!( + "skipping ngram gate: {MODEL_PATH}/config.json missing; set OPENINFER_TEST_MODEL_PATH" + ); + None + } + } +} + +/// `ngram = true` switches on host-side n-gram speculation (no draft model). The +/// baseline (`false`) still forces the prefix cache off so both runs take the +/// same cold prefill path. +fn launch_options(ngram: bool) -> Qwen3LaunchOptions { + Qwen3LaunchOptions { + device_ordinal: 0, + tp_size: 1, + cuda_graph: true, + offload: Qwen3OffloadOptions::disabled(), + no_prefix_cache: true, + max_prefill_tokens: DEFAULT_MAX_PREFILL_TOKENS, + memory: Qwen3MemoryOptions::new(0.85, DEFAULT_KV_CACHE_MEMORY_MARGIN_BYTES) + .validate() + .expect("valid memory options"), + lora: None, + decode_overlap: DecodeOverlap::Off, + batch_invariant: false, + dflash_draft_model_path: None, + ngram_speculative: ngram, + enable_kv_events: false, + } +} + +/// One decoded position: the chosen token and (when requested) the top-K +/// `(token, logprob)` distribution that produced it. +struct Step { + id: u32, + top_logprobs: Vec<(u32, f32)>, +} + +/// Submit one greedy request and collect the decoded steps until `Finished`. +fn generate( + handle: &EngineHandle, + prompt_tokens: Vec, + logprobs: usize, + max_tokens: usize, +) -> Vec { + let (token_tx, mut rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens, + params: SamplingParams::default(), + max_tokens, + lora_adapter: None, + token_tx, + logprobs, + echo: false, + }) + .expect("submit failed"); + + let mut steps = Vec::new(); + loop { + match rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, logprob }) => steps.push(Step { + id, + top_logprobs: logprob.map(|lp| lp.top_logprobs).unwrap_or_default(), + }), + Some(TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. }) => {} + Some(TokenEvent::Finished { .. }) => return steps, + Some(TokenEvent::Error { message, .. }) => panic!("generation failed: {message}"), + Some(TokenEvent::Rejected { message, .. }) => panic!("generation rejected: {message}"), + None => panic!("scheduler channel closed without Finished"), + } + } +} + +/// Submit several greedy requests at once, then collect each one's steps. They +/// run concurrently in the one engine — the scheduler batches them — so with +/// heterogeneous `max_tokens` the verify spans differ across a batch, exercising +/// the real bs>1 verify path. Each tuple is `(prompt_tokens, max_tokens)`; +/// logprobs are off so the speculative path stays active. Returns one step list +/// per request, in submission order. +fn generate_concurrent(handle: &EngineHandle, requests: Vec<(Vec, usize)>) -> Vec> { + let receivers: Vec<_> = requests + .into_iter() + .map(|(prompt_tokens, max_tokens)| { + let (token_tx, rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens, + params: SamplingParams::default(), + max_tokens, + lora_adapter: None, + token_tx, + logprobs: 0, + echo: false, + }) + .expect("submit failed"); + rx + }) + .collect(); + + receivers + .into_iter() + .map(|mut rx| { + let mut steps = Vec::new(); + loop { + match rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, logprob }) => steps.push(Step { + id, + top_logprobs: logprob.map(|lp| lp.top_logprobs).unwrap_or_default(), + }), + Some(TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. }) => {} + Some(TokenEvent::Finished { .. }) => break, + Some(TokenEvent::Error { message, .. }) => { + panic!("generation failed: {message}") + } + Some(TokenEvent::Rejected { message, .. }) => { + panic!("generation rejected: {message}") + } + None => panic!("scheduler channel closed without Finished"), + } + } + steps + }) + .collect() +} + +/// Prefill `context` (echo) and return the next-token distribution the *prefill* +/// kernel produces — the kernel the speculative verify path also uses. This is +/// the reference the spec pick should match (vs the plain-decode baseline, whose +/// kernel resolves bifurcation ties to the other side). +fn prefill_next(handle: &EngineHandle, context: Vec, logprobs: usize) -> Step { + let (token_tx, mut rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens: context, + params: SamplingParams::default(), + max_tokens: 1, + lora_adapter: None, + token_tx, + logprobs, + echo: true, + }) + .expect("submit failed"); + + loop { + match rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, logprob }) => { + return Step { + id, + top_logprobs: logprob.map(|lp| lp.top_logprobs).unwrap_or_default(), + }; + } + Some(TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. }) => {} + Some(TokenEvent::Finished { .. }) => panic!("echo prefill finished without a token"), + Some(TokenEvent::Error { message, .. }) => panic!("prefill failed: {message}"), + Some(TokenEvent::Rejected { message, .. }) => panic!("prefill rejected: {message}"), + None => panic!("scheduler channel closed without a token"), + } + } +} + +/// Compare one prompt's speculative `spec` steps against its plain-greedy `base`, +/// tolerating only the benign prefill-vs-decode kernel-gap tie flip (the spec +/// pick sits within `MARGIN_TOL` of the prefill kernel's own argmax, measured in +/// the prefill distribution the verify path actually runs). `Ok(())` ⇒ lossless +/// or a benign tie; `Err(diagnostic)` ⇒ a real spec bug. `handle` must be the +/// live speculative engine — at a divergence it re-prefills the shared context +/// (`prompt_tokens` + the matched prefix) to read that prefill-kernel reference. +fn check_lossless( + handle: &EngineHandle, + tokenizer: &DynTokenizer, + i: usize, + prompt: &str, + prompt_tokens: &[u32], + base: &[Step], + spec: &[Step], +) -> Result<(), String> { + let matched = base + .iter() + .zip(spec) + .take_while(|(b, s)| b.id == s.id) + .count(); + + // A strict-prefix relationship (one run stops early, or runs on past the + // other) must fail the gate even though every overlapping token matched: + // greedy speculation is lossless only if the two streams are token-for-token + // *and* length identical. Reject length mismatches before the divergence + // diagnostic below, which indexes `spec[matched]` and would otherwise panic + // out of bounds when `spec` is the shorter run. + if base.len() != spec.len() { + return Err(format!( + "prompt {i} ({prompt:?}): length mismatch — baseline produced {} tokens, \ + speculative produced {} tokens ({matched} prefix tokens matched); greedy \ + speculation must be token-for-token identical, not merely a shared prefix", + base.len(), + spec.len(), + )); + } + + if matched == base.len() { + eprintln!( + "prompt {i} ({prompt:?}): {matched}/{} tokens identical (100% lossless)", + base.len() + ); + return Ok(()); + } + + let spec_id = spec[matched].id; + let decode_argmax = base[matched].top_logprobs[0].0; + + { + let lo = matched.saturating_sub(2); + let hi = (matched + 3).min(base.len()).min(spec.len()); + let base_ids: Vec = base[..hi].iter().map(|s| s.id).collect(); + eprintln!(" [diag] prompt {i} matched={matched}"); + eprintln!( + " [diag] context+gen base ids {:?} = {:?}", + &base_ids, + tokenizer.decode(&base_ids, false).unwrap_or_default() + ); + eprintln!( + " [diag] base[{lo}..{hi}] = {:?}", + base[lo..hi] + .iter() + .map(|s| (s.id, tokenizer.decode(&[s.id], false).unwrap_or_default())) + .collect::>() + ); + eprintln!( + " [diag] spec[{lo}..{hi}] = {:?}", + spec[lo..hi] + .iter() + .map(|s| (s.id, tokenizer.decode(&[s.id], false).unwrap_or_default())) + .collect::>() + ); + } + + let mut context = prompt_tokens.to_vec(); + context.extend(base[..matched].iter().map(|s| s.id)); + let prefill_ref = prefill_next(handle, context, LOGPROBS); + + if prefill_ref.id == spec_id { + let decode_lp = base[matched] + .top_logprobs + .iter() + .find(|(t, _)| *t == spec_id) + .map(|(_, lp)| base[matched].top_logprobs[0].1 - lp); + eprintln!( + "prompt {i} ({prompt:?}): kernel-gap flip at token {matched} — verify(prefill)→{spec_id}, \ + decode→{decode_argmax}; spec matches prefill greedy (decode-margin {:?}). Not a spec bug.", + decode_lp + ); + return Ok(()); + } + + let prefill_regret = prefill_ref + .top_logprobs + .iter() + .find(|(t, _)| *t == spec_id) + .map(|(_, lp)| prefill_ref.top_logprobs[0].1 - lp); + + if let Some(regret) = prefill_regret { + if regret <= MARGIN_TOL { + eprintln!( + "prompt {i} ({prompt:?}): tie flip at token {matched} — \ + verify(prefill)→{}, spec→{spec_id}, decode→{decode_argmax}; \ + spec pick is #2 in the prefill distribution (regret {regret:.3} ≤ {MARGIN_TOL}). \ + Not a spec bug.", + prefill_ref.id, + ); + return Ok(()); + } + } + + let decode_regret = base[matched] + .top_logprobs + .iter() + .find(|(t, _)| *t == spec_id) + .map(|(_, lp)| base[matched].top_logprobs[0].1 - lp); + Err(format!( + "prompt {i}: at token {matched} spec chose {spec_id} but prefill greedy says {} and \ + decode greedy says {decode_argmax} (spec regret in prefill dist: {prefill_regret:?} > \ + {MARGIN_TOL}; in decode dist: {decode_regret:?}) — real spec bug", + prefill_ref.id, + )) +} + +/// bs=1 losslessness: each prompt's greedy n-gram speculative decode must match +/// its plain-greedy baseline token-for-token (modulo the benign bf16 tie flip). +#[test] +fn ngram_speculative_greedy_matches_plain_greedy() { + let Some(model_path) = target_path_or_skip() else { + return; + }; + let _gpu = GPU.lock().unwrap_or_else(|p| p.into_inner()); + + // Repetitive / structured prompts so the n-gram proposer fires often (its + // suffix recurs earlier in the context) and the verify path is exercised. + let prompts = [ + "1 2 3 4 1 2 3 4 1 2 3 4 1 2 3 4", + "def fibonacci(n):\n if n < 2:\n return n\n return fibonacci(n - 1) + fibonacci(", + "The quick brown fox. The quick brown fox. The quick brown fox.", + "Q: What is 17 multiplied by 23? A: Let's think step by step.", + "{\"a\": 1, \"b\": 2, \"a\": 1, \"b\": 2, \"a\": 1, \"b\": 2,", + ]; + + let tokenizer = common::load_tokenizer(&model_path); + let encoded: Vec> = prompts + .iter() + .map(|p| tokenizer.encode(p, false).expect("encode failed")) + .collect(); + + // 1. Baseline: plain greedy decode (speculative off), with logprobs so the + // regret check has the reference distribution at the divergence point. + let baseline: Vec> = { + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(false)) + .expect("failed to start baseline engine"); + let out = encoded + .iter() + .map(|t| generate(&handle, t.clone(), LOGPROBS, GENERATED_TOKENS)) + .collect(); + drop(handle); + // Free the target before the speculative engine loads the same weights. + std::thread::sleep(Duration::from_secs(2)); + out + }; + + // 2. Speculative: n-gram propose + verify (logprobs off ⇒ spec path active). + // Keep the engine alive through analysis: at a divergence it re-prefills + // the shared context to read the prefill-kernel reference. + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(true)) + .expect("failed to start speculative engine"); + + let mut failures = Vec::new(); + for (i, &prompt) in prompts.iter().enumerate() { + let spec = generate(&handle, encoded[i].clone(), 0, GENERATED_TOKENS); + if let Err(failure) = check_lossless( + &handle, + &tokenizer, + i, + prompt, + &encoded[i], + &baseline[i], + &spec, + ) { + failures.push(failure); + } + } + + drop(handle); + + assert!( + failures.is_empty(), + "n-gram speculative greedy decode is not lossless:\n{}", + failures.join("\n") + ); +} + +/// Concurrent, heterogeneous-`max_tokens` losslessness for the bs>1 verify path: +/// several greedy requests at staggered budgets run in one engine so the in-flight +/// batch mixes full and near-budget (truncated) verify spans. Each must stay +/// lossless vs its own plain-greedy baseline. A per-request context mix-up (the +/// proposer scanning the wrong request's history) or a bs>1 verify-span indexing +/// bug would surface here as a real (non-tie) divergence. +#[test] +fn ngram_concurrent_heterogeneous_is_lossless() { + let Some(model_path) = target_path_or_skip() else { + return; + }; + let _gpu = GPU.lock().unwrap_or_else(|p| p.into_inner()); + + let cases: [(&str, usize); 4] = [ + ( + "def fibonacci(n):\n if n < 2:\n return n\n return fibonacci(", + 64, + ), + ("1 2 3 4 1 2 3 4 1 2 3 4 1 2 3 4", 24), + ( + "Q: What is 17 multiplied by 23? A: Let's think step by step.", + 48, + ), + ( + "The quick brown fox. The quick brown fox. The quick brown fox.", + 40, + ), + ]; + + let tokenizer = common::load_tokenizer(&model_path); + let encoded: Vec> = cases + .iter() + .map(|(p, _)| tokenizer.encode(p, false).expect("encode failed")) + .collect(); + + // 1. Baselines: each prompt's plain-greedy decode (spec off) at ITS budget, + // with logprobs for the regret reference. Sequential, one engine. + let baselines: Vec> = { + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(false)) + .expect("failed to start baseline engine"); + let out = encoded + .iter() + .zip(&cases) + .map(|(t, (_, max_tokens))| generate(&handle, t.clone(), LOGPROBS, *max_tokens)) + .collect(); + drop(handle); + std::thread::sleep(Duration::from_secs(2)); + out + }; + + // 2. Speculative engine: submit all four at once so they form real batches. + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(true)) + .expect("failed to start speculative engine"); + let specs = generate_concurrent( + &handle, + encoded + .iter() + .zip(&cases) + .map(|(t, (_, max_tokens))| (t.clone(), *max_tokens)) + .collect(), + ); + + let mut failures = Vec::new(); + for (i, (prompt, _)) in cases.iter().enumerate() { + if let Err(failure) = check_lossless( + &handle, + &tokenizer, + i, + prompt, + &encoded[i], + &baselines[i], + &specs[i], + ) { + failures.push(failure); + } + } + drop(handle); + + assert!( + failures.is_empty(), + "concurrent heterogeneous n-gram speculative decode is not lossless:\n{}", + failures.join("\n") + ); +} + +/// Mixed greedy + sampling traffic must not break the greedy request. +/// +/// A non-greedy request sharing the batch flips `should_speculative_decode` to +/// false (it requires *every* active request to be greedy), so the co-resident +/// greedy request decodes outside the speculative path — via plain decode, or +/// via the fused unified step when a non-greedy request is pending while the +/// greedy one is active. Both paths must keep `ngram_ctx` anchored on the +/// dangling token; a missed sync there silently degrades n-gram acceptance once +/// the batch returns to pure greedy (it is not a correctness bug — the verify +/// seed comes from the real KV, not `ngram_ctx`). This guards that mixed-traffic +/// class: the greedy request must stay token-identical to its solo plain-greedy +/// baseline. The sampling request's output is non-deterministic, so only the +/// greedy request is checked; it is there purely to hold the batch non-greedy. +/// +/// Note: deterministically landing the greedy request on the *unified* branch +/// (vs plain decode) depends on scheduler arrival timing the public API can't +/// pin down; the unified context-sync itself mirrors the proven `execute_decode` +/// path, so this test guards the broader property rather than that one branch. +#[test] +fn ngram_mixed_greedy_and_sampling_is_lossless() { + let Some(model_path) = target_path_or_skip() else { + return; + }; + let _gpu = GPU.lock().unwrap_or_else(|p| p.into_inner()); + + let greedy_prompt = "def fibonacci(n):\n if n < 2:\n return n\n return fibonacci("; + const GREEDY_MAX: usize = 48; + let tokenizer = common::load_tokenizer(&model_path); + let greedy_tokens = tokenizer + .encode(greedy_prompt, false) + .expect("encode failed"); + let sampling_tokens = tokenizer + .encode("Tell me a long and winding story about a dragon.", false) + .expect("encode failed"); + + // 1. Baseline: the greedy prompt's solo plain-greedy decode (spec off) with + // logprobs for the regret reference. + let baseline = { + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(false)) + .expect("failed to start baseline engine"); + let out = generate(&handle, greedy_tokens.clone(), LOGPROBS, GREEDY_MAX); + drop(handle); + std::thread::sleep(Duration::from_secs(2)); + out + }; + + // 2. Speculative engine (n-gram on). Submit a long *sampling* request and the + // greedy request concurrently so they co-reside: the sampling request keeps + // the batch non-greedy, forcing the greedy one off the speculative path. + let handle = openinfer_qwen3_4b::launch(Path::new(&model_path), launch_options(true)) + .expect("failed to start speculative engine"); + + // Non-greedy (temperature >= eps, top_k != 1) so the batch is held non-greedy; + // start from Default to stay robust to added fields. + let sampling_params = SamplingParams { + temperature: 0.8, + top_p: 0.95, + ..SamplingParams::default() + }; + let (sampling_tx, mut sampling_rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens: sampling_tokens, + params: sampling_params, + max_tokens: 96, + lora_adapter: None, + token_tx: sampling_tx, + logprobs: 0, + echo: false, + }) + .expect("submit sampling failed"); + let (greedy_tx, mut greedy_rx) = TokenSink::standalone(); + handle + .submit(GenerateRequest { + request_id: None, + queued_at_unix_s: None, + prompt_tokens: greedy_tokens.clone(), + params: SamplingParams::default(), + max_tokens: GREEDY_MAX, + lora_adapter: None, + token_tx: greedy_tx, + logprobs: 0, + echo: false, + }) + .expect("submit greedy failed"); + + // Drain the greedy request's steps; the sampling request just runs alongside. + let mut greedy_steps = Vec::new(); + loop { + match greedy_rx.blocking_recv().map(|(_, event)| event) { + Some(TokenEvent::Token { id, logprob }) => greedy_steps.push(Step { + id, + top_logprobs: logprob.map(|lp| lp.top_logprobs).unwrap_or_default(), + }), + Some(TokenEvent::Scheduled { .. } | TokenEvent::PromptTokens { .. }) => {} + Some(TokenEvent::Finished { .. }) => break, + Some(TokenEvent::Error { message, .. }) => { + panic!("greedy generation failed: {message}") + } + Some(TokenEvent::Rejected { message, .. }) => { + panic!("greedy generation rejected: {message}") + } + None => panic!("scheduler channel closed without Finished"), + } + } + // Drain the sampling request to completion so the engine tears down cleanly. + while let Some((_, event)) = sampling_rx.blocking_recv() { + if matches!( + event, + TokenEvent::Finished { .. } | TokenEvent::Error { .. } | TokenEvent::Rejected { .. } + ) { + break; + } + } + + let result = check_lossless( + &handle, + &tokenizer, + 0, + greedy_prompt, + &greedy_tokens, + &baseline, + &greedy_steps, + ); + drop(handle); + + assert!( + result.is_ok(), + "greedy n-gram request diverged from plain greedy under mixed greedy+sampling traffic:\n{}", + result.unwrap_err() + ); +} diff --git a/openinfer-qwen3-4b/tests/scheduler_robustness.rs b/openinfer-qwen3-4b/tests/scheduler_robustness.rs index fdaec356..cde5098c 100644 --- a/openinfer-qwen3-4b/tests/scheduler_robustness.rs +++ b/openinfer-qwen3-4b/tests/scheduler_robustness.rs @@ -106,6 +106,7 @@ fn scheduler_survives_consumer_drop() { true, None, false, + false, ) .expect("failed to start engine"); assert_eq!( diff --git a/openinfer-server/src/bin/bench_serving/main.rs b/openinfer-server/src/bin/bench_serving/main.rs index 63c5f090..8b3af320 100644 --- a/openinfer-server/src/bin/bench_serving/main.rs +++ b/openinfer-server/src/bin/bench_serving/main.rs @@ -193,6 +193,12 @@ fn main() -> Result<()> { openinfer_qwen3_4b::Qwen3OffloadOptions::disabled(), false, max_prefill_tokens, + openinfer_qwen3_4b::Qwen3MemoryOptions::default(), + openinfer_qwen3_4b::DecodeOverlap::Off, + false, + None, + false, + false, )?; finish(handle, cli.cuda_graph) } diff --git a/openinfer-server/src/config.rs b/openinfer-server/src/config.rs index c03ef693..e395225a 100644 --- a/openinfer-server/src/config.rs +++ b/openinfer-server/src/config.rs @@ -95,6 +95,13 @@ pub(crate) struct Args { #[arg(long = "dflash-draft-model-path")] pub dflash_draft_model_path: Option, + /// Enable Qwen3 host-side n-gram (prompt-lookup) speculative decoding. Needs + /// no draft model; single-GPU greedy only, and forces the prefix cache off. + /// Mutually exclusive with --dflash-draft-model-path, --enable-lora, + /// --kv-offload, and --decode-overlap. + #[arg(long = "ngram-speculative", default_value_t = false)] + pub ngram_speculative: bool, + /// Cap on total prompt tokens forwarded in one scheduler step (chunked /// prefill), for both the Qwen3 and Qwen3.5 schedulers. Prefill activation /// scratch scales with the step's prompt tokens, so this bounds peak VRAM diff --git a/openinfer-server/src/main.rs b/openinfer-server/src/main.rs index eceab2d6..a56cf77e 100644 --- a/openinfer-server/src/main.rs +++ b/openinfer-server/src/main.rs @@ -119,6 +119,10 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result openinfer_deepseek_v4::launch( @@ -182,6 +186,10 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result None, }; + anyhow::ensure!( + !args.ngram_speculative || args.tp_size == 1, + "--ngram-speculative currently requires --tp-size=1" + ); openinfer_qwen3_4b::launch( &args.model_path, Qwen3LaunchOptions { @@ -200,6 +208,7 @@ fn load_engine(args: &Args, model_type: ModelType) -> anyhow::Result