PR: feat(qwen3): n-gram speculative decoding

[wjinxu]

Adds lossless n-gram (prompt-lookup) speculative decoding for Qwen3, built on the method-agnostic speculative core that landed in #436 / #442 (merged into this branch). Off by default; enable with --ngram-speculative.

This started as a parallel speculative stack (its own speculative.rs, verify, KV lifecycle). It has been re-based onto #436's shared core: n-gram is now just a second propose implementation behind the same verify / accept / KV transaction — no duplicated controller.

The shape

Speculative decoding is one optimistic transaction; only propose differs by method:

1. Propose — produce up to K candidate tokens.
2. Verify — one target forward over the 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.

The draft↔verify boundary is a pure token span. So accept_greedy / build_verify_results, execute_speculative_verify_impl, and the schedule/apply/revert_speculative KV transaction are reused unchanged. Only two things are method-specific: where the drafts come from, and whether verify captures hidden states.

What n-gram adds

• Proposer (ngram.rs): NgramProposer — longest-suffix prompt-lookup over the request's own token history, no draft model. Tries the longest configured suffix first, falls back to shorter. + NgramConfig (from_env) + 8 unit tests.
• Propose is host-side: execute_speculative_draft runs the lookup on the executor thread — no worker forward, no draft model, no readiness handshake. DFlash instead dispatches StepCommand::SpeculativeDraft to the lane.
• No-capture verify: 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.
• Context tracking (ngram_ctx): per-request running token context, seeded at prefill (prompt + first_token), appended at every commit (verify / plain decode / fused unified decode), dropped on retire. Always ends with the request's dangling token.
• Method state (SpeculativeMethod enum): the two mutually-exclusive executor fields are collapsed into one Option<SpeculativeMethod>{ Dflash | Ngram } — single source of truth, no "assert not both set" smell. This is the pre-trait consolidation feat(qwen3): DFlash speculative decoding #436 deferred. A full SpeculativeProposer trait object is not introduced yet: 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, feat(qwen3): explore speculative decoding #325) makes a shared trait pay off.

Config

• --ngram-speculative — enable. Single-GPU greedy only; mutually exclusive with --dflash-draft-model-path, --enable-lora, --kv-offload, --decode-overlap; requires --tp-size=1. Forces the prefix cache off (verify writes each request's own speculative KV span).
• OPENINFER_QWEN3_NGRAM_TOKENS — draft count K (default 4).
• OPENINFER_QWEN3_NGRAM_MAX_NGRAM — longest suffix to match (default 3).

Losslessness

Greedy speculation is lossless (the accepted tokens are the target's own greedy continuation), which is a built-in oracle: spec-on must equal spec-off token-for-token (modulo the benign prefill-vs-decode bf16 tie flip).

GPU-validated on RTX 4090 (Qwen3-4B), tests/ngram_speculative_gate.rs, 3/3 pass, every prompt 100% token-identical to plain greedy:

• ngram_speculative_greedy_matches_plain_greedy — bs=1, 5 prompts.
• ngram_concurrent_heterogeneous_is_lossless — 4 concurrent, staggered budgets (bs>1 verify-span path).
• ngram_mixed_greedy_and_sampling_is_lossless — a greedy request stays identical while a sampling request holds the batch off the speculative path.

A length-1 verify span (n-gram "no match") already degrades to a plain single-token decode in build_verify_results, so a miss needs no special fallback.

Files

File	Change
openinfer-qwen3-4b/src/ngram.rs	New: NgramProposer + NgramConfig + 8 unit tests
openinfer-qwen3-4b/src/executor.rs	SpeculativeMethod enum; load_ngram_drafter; host-side execute_speculative_draft; execute_ngram_verify; ngram_ctx lifecycle (prefill / decode / unified / drop)
openinfer-qwen3-4b/src/executor/spec.rs	Generalize the verify transaction to be method-agnostic (relax DFlash-only asserts; append accepted tokens to ngram_ctx)
openinfer-qwen3-4b/src/verify_graph.rs	span() accessor for the no-capture verify-buffer sizing
openinfer-qwen3-4b/src/scheduler.rs, src/lib.rs, openinfer-server/src/{config,main}.rs	--ngram-speculative wiring + mutual-exclusion validation
openinfer-qwen3-4b/tests/ngram_speculative_gate.rs	New: engine-level regret-based losslessness gate (3 tests)

Sits on top of main after merging #436 / #442, so the diff is the n-gram addition, not a re-derivation of the shared core.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 2f6d3e45b6

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[wjinxu]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3f27ad7ada

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[wjinxu]

@codex review

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 8c5b5adc4a

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[wjinxu]

@codex review

[chatgpt-codex-connector]

Codex Review: Didn't find any major issues. Breezy!

Reviewed commit: a47adbea29

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[wjinxu]

Codex Review: Didn't find any major issues. Breezy!

Reviewed commit: a47adbea29

ℹ️ About Codex in GitHub

Thank you

[wjinxu]

@xiaguan Rebased onto latest main — the only conflict was the enable_kv_events field added in #457 colliding with this branch's ngram_speculative field in the DFlash test launch_options(); kept both. Builds clean (cargo test --no-run, exit 0), n-gram unit tests + the GPU losslessness gate all pass. Ready for review whenever you have a chance 🙏

[xiaguan]

Thanks for putting this together. The correctness side looks solid from the losslessness gate; the remaining thing I’d like to see before we position this as a performance win is a serving-level A/B with vllm bench serve datasets.

I ran a quick local pass on an RTX 5070 Ti, Qwen3-4B, temperature=0, ignore_eos, no prefix cache, 12 prompts × 128 output tokens:

dataset / mode	plain	n-gram	DFlash
sonnet c1	90.2 tok/s	95.9 tok/s (+6.3%)	140.3 tok/s (+55.6%)
sonnet c4	307.0 tok/s	198.2 tok/s (-35.4%)	353.8 tok/s (+15.2%)
random c1	90.1 tok/s	112.8 tok/s (+25.2%)	not run
random c4	306.6 tok/s	220.1 tok/s (-28.2%)	not run

So the current shape looks promising for some single-stream cases, but the concurrent serving path needs more investigation before we call it a general speedup. Could you add vllm bench serve A/B results to the PR and either address the c4 regression or document the feature as experimental / single-stream-oriented for now?

Here are the exact commands I used. Paths may need adjusting for your machine:

set -euo pipefail

MODEL=/data/models/Qwen3-4B
DRAFT=/data/models/Qwen3-4B-DFlash-b16
PORT=18080
RESULT_DIR=/tmp/openinfer-ngram-pr349
SONNET=/tmp/openinfer-sonnet.txt
mkdir -p "$RESULT_DIR"

curl -fsSL https://raw.githubusercontent.com/vllm-project/vllm/main/benchmarks/sonnet.txt -o "$SONNET"
PATH=/usr/local/cuda-13.3/bin:$PATH cargo build --release

Start one server mode at a time:

# plain
RUST_LOG=warn target/release/openinfer \
  --model-path "$MODEL" \
  --served-model-name qwen3-4b \
  --port "$PORT" \
  --no-prefix-cache \
  --gpu-memory-utilization 0.85

# n-gram
RUST_LOG=warn target/release/openinfer \
  --model-path "$MODEL" \
  --served-model-name qwen3-4b \
  --port "$PORT" \
  --no-prefix-cache \
  --gpu-memory-utilization 0.85 \
  --ngram-speculative

# DFlash
RUST_LOG=warn target/release/openinfer \
  --model-path "$MODEL" \
  --served-model-name qwen3-4b \
  --port "$PORT" \
  --no-prefix-cache \
  --gpu-memory-utilization 0.85 \
  --dflash-draft-model-path "$DRAFT"

Run sonnet for c1/c4:

MODE=plain   # plain | ngram | dflash
CONCURRENCY=1  # then repeat with 4

uv run --with vllm vllm bench serve \
  --backend openai \
  --base-url "http://127.0.0.1:${PORT}" \
  --model qwen3-4b \
  --tokenizer "$MODEL" \
  --dataset-name sonnet \
  --dataset-path "$SONNET" \
  --sonnet-input-len 550 \
  --sonnet-output-len 128 \
  --sonnet-prefix-len 200 \
  --num-prompts 12 \
  --num-warmups 1 \
  --max-concurrency "$CONCURRENCY" \
  --request-rate inf \
  --ignore-eos \
  --temperature 0 \
  --ready-check-timeout-sec 60 \
  --disable-tqdm \
  --metric-percentiles 50,90,99 \
  --save-result \
  --save-detailed \
  --result-dir "$RESULT_DIR" \
  --result-filename "sonnet-${MODE}-c${CONCURRENCY}.json"

Run random for a control workload:

MODE=plain   # plain | ngram
CONCURRENCY=1  # then repeat with 4

uv run --with vllm vllm bench serve \
  --backend openai \
  --base-url "http://127.0.0.1:${PORT}" \
  --model qwen3-4b \
  --tokenizer "$MODEL" \
  --dataset-name random \
  --random-input-len 550 \
  --random-output-len 128 \
  --random-range-ratio 0 \
  --num-prompts 12 \
  --num-warmups 1 \
  --max-concurrency "$CONCURRENCY" \
  --request-rate inf \
  --ignore-eos \
  --temperature 0 \
  --ready-check-timeout-sec 60 \
  --disable-tqdm \
  --metric-percentiles 50,90,99 \
  --save-result \
  --save-detailed \
  --result-dir "$RESULT_DIR" \
  --result-filename "random-${MODE}-c${CONCURRENCY}.json"

[xiaguan]

One more note: I’m fine with keeping the optimization work as a follow-up PR if we want to land the correctness / feature scaffolding first, but then I think we should open an issue to explicitly track the concurrent serving regression.

The next step should be profiling why c4 regresses before changing the implementation. My guess is around the current speculative batch/verify path, but that is only a hypothesis until we have an nsys or per-step timing breakdown for plain vs n-gram at c4.

[wjinxu]

Added serving-level A/B with vllm bench serve (RTX 4090, Qwen3-4B, temperature=0, --ignore-eos, --no-prefix-cache, 12×128), sonnet + random, vs vLLM0.21/0.23.

The regression was from paying the verify forward at near-zero acceptance. Added an acceptance gate (NgramGate): an EWMA of accepted draft tokens/step; below threshold the proposer emits nothing and the batch falls back to plain decode (no verify forward), with a periodic probe to re-open on repetitive text. Lossless — the gate only changes whether we speculate; the losslessness gate still passes 3/3.

n-gram vs plain, after the gate:

dataset / conc	acceptance	vLLM 0.23	openinfer
sonnet c1	~5%	−5.8%	+0.6%
sonnet c4	~7%	−5.8%	−2.2%
random c1	~33%	+19.3%	+21.8%
random c4	~37%	+17.0%	−20.0%

Prose regression is gone (better than vLLM, which keeps speculating); random c1 matches vLLM. random c4 is the remaining gap (same ~37% acceptance,but vLLM's single fused forward beats our two-forward + prefill-kernel verify) — root-caused with nsys, tracked as a follow-up PR

[xiaguan]

Thanks for the update and for adding the serving A/B data.

Two code-quality things I’d like to see cleaned up before merge:

1. The new docs/models/qwen3/ngram-speculative.md still describes an older version of this branch in a few places. It mentions things like tests/ngram_speculative.rs, scheduler/speculative.rs, SpeculativeProposer, and OPENINFER_QWEN3_SPEC, but the current implementation is based on --ngram-speculative, SpeculativeMethod::{Dflash, Ngram}, executor-side n-gram context, and the shared verify/accept/KV transaction. Could you refresh the doc to match the actual code path?

2. The n-gram lifecycle is spread across executor.rs: context seeding, append-on-decode, append-on-unified, append-on-verify, drop cleanup, and gate bookkeeping all live in separate branches. Since executor.rs is already very large, I’d prefer to keep this feature’s state in a small dedicated helper, e.g. an NgramState/NgramRuntime that owns the context map, proposer, and gate. Then the executor can call intent-level methods instead of manually maintaining n-gram state in every execution path.

The feature shape looks reasonable; I mainly want the documentation and ownership boundary to stay maintainable.

[wjinxu]

1. Ownership — all n-gram state now lives in a dedicated NgramRuntime (proposer + gate + per-request context). SpeculativeMethod::Ngram holds it; the two executor fields and the per-path bookkeeping collapse to intent-level calls (seed / append_committed / drop_request / is_ready / propose_step / record_verified). No behaviour change.

2. Docs — rewrote ngram-speculative.md to the shipped path: dropped the stale SpeculativeProposer trait, scheduler/speculative.rs, OPENINFER_QWEN3_SPEC, tests/ngram_speculative.rs, verify_forward; now reflects --ngram-speculative, the SpeculativeMethod { Dflash, Ngram } enum, NgramRuntime, and the shared verify/accept/KV transaction.

Green locally: build + fmt + clippy clean, unit tests 12/12, GPU losslessness gate 3/3.

[xiaguan]

Review: n-gram speculative decoding — thanks!

Great to see this land on top of the #436 shared core rather than as a parallel stack. Reusing verify / accept / the KV transaction and only swapping propose is the right call, and the enum dispatch instead of a premature trait is honest. The GPU losslessness gate (3/3, token-identical) is solid evidence. A few things I'd like to see addressed before this is ready — none of them block the design, mostly scope/defaults and some polish.

Substantive

1. random c4 is −20% in the PR's own A/B, and the fix (fused verify forward) is a follow-up. The doc is transparent about it, which I appreciate — but it does mean --ngram-speculative, in what should be spec decode's sweet spot (high concurrency + moderate acceptance), regresses below plain decode (vLLM is +17% there). Two options I'd be happy with:

   • Land the fused verify forward in this PR, or
   • Mark the flag experimental and document that concurrency > 1 can regress until fused verify lands (so nobody enables it in prod expecting a flat win).

2. NgramGate is engine-wide; the mixed-load case isn't tested. record_verified feeds mean(accepted_draft_tokens) across the drafted requests in a step, so a repetitive request (accept=4) co-resident with prose requests (accept=0) keeps the EWMA healthy and the gate open — the prose requests then keep paying the prefill-kernel verify for ~0 accepted tokens. The A/B tables are all homogeneous datasets. Could you add a 50/50 mixed A/B (repetitive + prose) to show the gate doesn't degrade the prose side? If it does, a per-request gate (or per-request EWMA) is probably the real shape.

3. accept_threshold semantics drift with K. It's compared against "mean accepted draft tokens per step" (absolute, 0..K), so 0.3 means 7.5% acceptance at K=4 but 3.75% at K=8. The doc quotes acceptance as percentages (~5% / ~37%), but the knob is in absolute tokens — users tuning K have to re-tune the threshold. Making it accepted / drafted (ratio, 0..1) would make it K-portable and match the doc's mental model directly. Happy if this is a follow-up, but worth noting now since it touches the gate's public surface.

Minor

4. execute_ngram_verify reallocs verify scratch as the span grows 1 → K+1. K is known at load_ngram_drafter time; pre-allocating at K+1 there avoids the startup alloc stalls on the lane. DFlash already does the equivalent with its fixed block_size.

5. NgramProposer::propose is O(context_len · max_ngram) per step, full scan, no look-back window. Fine for bounded contexts, but it'll erode the win on long generations — vLLM bounds the look-back. The doc mentions this as a future concern; it'll become one, so worth a tracked issue at least.

6. ngram_mixed_greedy_and_sampling_is_lossless can't deterministically hit the unified ngram_ctx sync branch (the test's own comment says so), and commit 1d7b6c98 fixed a bug exactly there. The fact that the IT can't pin it suggests the unified sync depends on scheduler arrival timing — a mild design smell. Not blocking, but worth a comment in the test or a deterministic harness later.

7. propose_step's anyhow::ensure!(item.params.is_greedy()) is unreachable — should_speculative_decode already guarantees the whole batch is greedy. Either drop it or demote to debug_assert! to mark it as an invariant guard.

8. NgramConfig fields are all pub. The type needs to be pub (it's in load_ngram_drafter's signature), but the fields are only constructed by from_env(); pub(crate) would be enough and keeps the contract tight.

Comment density

This is subjective, but I found the commenting heavy. A lot of the prose explains what the code does line-by-line rather than why — e.g. the execute_plan::SpeculativeDecode block, the NgramRuntime method docs, the append_committed sync blocks duplicated across decode/unified/split paths. Where the code already reads as a self-explanatory sequence, letting it speak cuts the surface area reviewers have to hold and makes the why comments (the ones that actually earn their keep, like the bf16 tie-flip rationale in the gate test) stand out more. Not asking for a mass strip — just biasing toward "delete unless it explains intent the code can't."

Performance

The acceptance gate is a nice save for prose, and the nsys root-cause for random-c4 (two forwards + prefill-kernel verify vs vLLM's single fused forward) is exactly the right diagnosis. The fused unified verify forward is the obvious next prize — looking forward to it closing the c4 gap. Happy to review that PR.

Conclusion

Solid foundation and the losslessness story is convincing. I'd just like to see the c4 framing tightened up and a mixed-load A/B before this merges — the rest can be follow-ups.