feat(qwen3): rejection sampling for speculative decoding (non-greedy requests)

[xiaguan]

Current state

Qwen3 speculative decoding (DFlash) is greedy-only. Acceptance is argmax equality, not distribution-correct rejection sampling:

• accept_greedy in openinfer-qwen3/src/speculative.rs compares draft tokens against the target model's argmax at each position. The module doc explicitly notes: "Sampling-correct acceptance would need the target model's full distribution … it is left until a sampling method actually lands."
• The executor hard-asserts greedy in openinfer-qwen3/src/executor/spec.rs (both draft and verify paths): speculative verification currently supports greedy sampling only.
• The scheduler gates the speculative path via should_speculative_decode in openinfer-qwen3/src/scheduler/plan.rs, requiring params.is_greedy() for every active request. A single sampled request silently drops the whole batch back to plain decode.

Net effect: any request with temperature > 0 gets zero speedup from speculative decoding.

What's needed

Standard rejection sampling (Leviathan et al. 2023 / vLLM-style):

1. Verify forward must expose the target distribution (post-temperature/top-p) at each of the K + 1 positions, not just argmax.
2. Draft side must record its proposal probabilities q(x) per proposed token.
3. Acceptance: accept draft token x with probability min(1, p(x) / q(x)); on rejection, resample from the normalized residual max(0, p − q) and stop.
4. Bonus token at full acceptance is sampled from the target distribution at position K + 1.
5. Relax the is_greedy() gates in the scheduler and executor; keep per-request SamplingParams respected within a batch.

Correctness bar: the output distribution must be provably identical to non-speculative sampling (the rejection-sampling guarantee), gated by a statistical test against plain decode at matching seeds/params.

Notes

• Batch-level all-or-nothing gating in should_speculative_decode could also be revisited (mixed greedy/sampled batches), but that's a follow-up — per-request rejection sampling is the prerequisite.

[n-WN]

Claiming this — it's the natural continuation of the sampling-parity line (#490/#491 gave select_batch per-row params and the per-request philox-seed groundwork, which the statistical gate here will lean on: fixed-seed replay makes "provably identical to non-speculative sampling" testable as same-seed distribution comparison, not just aggregate stats).

Rough attack order, smallest correct slices:

1. Verify-side distribution exposure: the verify forward already computes logits at all K+1 positions; expose the post-temperature/top-k/top-p probabilities (the same renorm pipeline gpu_sample_batch_into uses, so filtered-out tokens are exact zeros — that matters for the residual). No behavior change, greedy path untouched.
2. Draft-side q(x): record the draft's proposal probability per proposed token during draft sampling.
3. Acceptance kernel: per position, accept with min(1, p/q); first rejection resamples from normalize(max(0, p − q)) and stops; full acceptance samples the bonus token from the target distribution at K+1. Keep accept_greedy as the is_greedy() fast path — argmax equality is the degenerate case, so greedy stays bit-identical by construction.
4. Gate relaxation last: executor asserts and should_speculative_decode per-request (batch all-or-nothing relaxation stays a follow-up, as noted).

Correctness bar as stated: a statistical acceptance test vs plain decode at matched seeds/params (per-position exact-match under fixed seed for the accepted prefix + chi-squared over the resample distribution), plus the existing greedy A/B to pin no regression.

Heads-up that I'll verify on H100; will open the tracking PR once slice 1 measures clean.

[n-WN]

Progress:

• Slice 1 (kernel primitives) is up as feat(kernels): verify-side target probs + chain rejection sampling primitives (#512) #534: gpu_verify_probs_into (post-filter target distribution, filtered tokens exact zeros, distribution-equivalent to the fused sampler) + gpu_spec_accept_into (vendored ChainSpeculativeSampling; one-hot proposal derived on device — DFlash proposes greedily, so q = δ(x−draft) and acceptance reduces to min(1, p_target(draft)), no draft-side probability plumbing needed). 3/3 corner-point-deterministic GPU tests on H100. One trap pinned in doc+test: the kernel's accepted counter keeps counting past the first rejection (acceptance-rate telemetry) — emitted is the commit signal.
• Slice 2 (executor wiring) is implemented and building on a branch stacked on feat(kernels): verify-side target probs + chain rejection sampling primitives (#512) #534 (feat/spec-rejection-verify): sampled requests take chain rejection sampling over the verify forward's all-position logits (greedy rows keep the batched argmax untouched), scratch lives in VerifyGraphBuffers (no per-step allocation), and the gates open to greedy || (min_p == 0 && seed.is_none()) — min_p stays out (sampling-time mask, target distribution not representable as a renorm) and per-request seeds stay out (the chain kernel draws its own philox stream, so the replay contract can't be honored speculatively). Full openinfer-qwen3 lib test target builds on CUDA 12.9; unit seams pass; the remote_fetch suite is regression-clean.
• Remaining before the wiring PR: GPU behavior runs + the statistical equivalence gate vs plain decode (per-position fixed-seed exact-match on the accepted prefix + distribution test over resamples) — needs Qwen3 weights + the DFlash draft model on my box; will open the wiring PR once feat(kernels): verify-side target probs + chain rejection sampling primitives (#512) #534 settles and those measure clean.