feat(kernels): verify-side target probs + chain rejection sampling primitives (#512)

[n-WN]

Slice 1 of #512 — the two GPU primitives non-greedy speculative decoding needs, as a self-contained kernels-layer PR. Executor wiring (verify-path integration, gate relaxation, the statistical equivalence gate) follows as the next slice on top.

gpu_verify_probs_into — the target distribution itself

The batched sampling pipeline's gather + softmax + top-k/top-p renorm, stopped before the terminal sampling kernel: the post-filter probabilities land in the caller's n_rows × vocab buffer, filtered tokens as exact zeros (that matters for the rejection residual). Distribution-equivalent to the fused sampling fast path — it filters at draw time, this filters then draws; the law over tokens is identical, which is all rejection sampling's correctness needs.

min_p rows are rejected fail-loud: min_p is applied inside the sampling kernel, not as a renorm, so that target distribution is not representable here — such requests stay off the speculative path (gate keeps them out in the next slice).

gpu_spec_accept_into — chain rejection sampling

Wraps the vendored FlashInfer ChainSpeculativeSampling: accept draft i with min(1, p_target/q_draft); the first rejection resamples from relu(target − draft) renormalized and stops; full acceptance emits the bonus token from the target at position K; the tail is -1-filled — exactly the "longest accepted prefix + one model token" contract accept_greedy already has, so the executor seam stays shape-compatible.

onehot_draft: DFlash proposes greedily, i.e. the proposal is the degenerate q(x) = δ(x − draft). Under that proposal, acceptance reduces to min(1, p_target(draft)) and the residual is the target with the draft token's mass removed — still distribution-exact. Deriving the one-hot rows on device means no draft-side proposal-probability plumbing is needed at all for the current proposer (real q support is already in the signature for a future sampled drafter).

A semantic trap, pinned in doc + test

The kernel returns two counters with different meanings: emitted is the accepted-prefix length — the commit signal; accepted keeps counting hypothetical acceptances past the first rejection (FlashInfer's acceptance-rate telemetry). Committing on accepted would emit wrong tokens; the wrapper doc says so and the test asserts the distinction.

Verification (H100, CUDA 12.9) — 3/3, corner-point deterministic (not statistical)

• verify_probs_renorm_matches_the_sampling_law: closed-form softmax of {2,1,0} matches to 5e-3; top_k=2 renorm leaves the filtered token an exact 0.0 and renormalizes survivors to {0.7311, 0.2689}; rows sum to 1.
• verify_probs_rejects_min_p_rows: fail-loud path.
• spec_accept_full_acceptance_and_certain_rejection: a target with mass 1.0 on each draft forces full acceptance + bonus ([100, 200, 777], emitted=2); a target with zero mass on the draft forces first-step rejection where the residual is a single token, making the resample exact ([555, -1, -1], emitted=0).

🤖 Generated with Claude Code

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: f162433592

ℹ️ 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".

[chatgpt-codex-connector]

Re-export the speculative sampling entry points

These new primitives are declared pub inside the private sampling module, but openinfer-kernels/src/ops.rs only re-exports the older sampling APIs (gpu_sample_batch_into, etc.). In that setup model crates such as the upcoming Qwen verify-path integration cannot call either gpu_verify_probs_into or gpu_spec_accept_into; they are only reachable from this module’s tests. Please add them to the public ops re-export list (and core re-exports if needed) so the kernels-layer API added here is usable outside openinfer-kernels.

Useful? React with 👍 / 👎.

[n-WN]

Right — they were pub only inside the private module, unreachable from model crates. Re-exported through ops in d0de3b1 (the upcoming verify-path slice consumes them from there).

[xiaguan]

Thanks for this — the approach is sound: reusing FlashInfer's renorm pipeline + ChainSpeculativeSampling instead of reinventing rejection sampling, and the onehot_draft idea elegantly avoids q-probability plumbing for a greedy proposer. The emitted vs accepted semantic trap is well documented and tested. A few things that would make this land cleaner:

Suggested changes

1. sampling.rs exceeds 1k lines (would benefit from a split)

openinfer-kernels/src/ops/sampling.rs goes from 996 → 1234 lines with this PR, past our 1k soft limit. The two new functions + their tests (~336 lines) are self-contained spec-decoding concerns and would read better in a dedicated module like ops/spec_sampling.rs. Happy to help with the split if useful — it's mostly cut/paste + re-export.

2. C + Rust both duplicate the existing gpu_sample_batch renorm pipeline

The gather + softmax + top-k/top-p renorm sequence appears twice at each layer:

• C: gpu_verify_probs_flashinfer_cuda (flashinfer_sampling.cu:120-162) and gpu_sample_batch_flashinfer_cuda (flashinfer_sampling.cu:164-242) share the same gather-grid construction, gather_cast_logits_f32_kernel launch, OnlineSoftmax, RadixTopKRenormProbMultiCTA, and TopPRenormProb — roughly 30 lines duplicated.
• Rust: gpu_verify_probs_into (sampling.rs:134-245) and sample_uniform_batch_into (sampling.rs:334-493) share the same capacity checks, row/temperature/top_p validation, top_k clamping (if r.top_k > 0 && r.top_k < vocab), and the four memcpy_htod calls — roughly 35 lines duplicated.

A shared C helper renorm_probs_flashinfer(...) and a Rust helper prepare_sampling_params(...) returning (row_indices, temperature, top_k, top_p, has_top_k_filter, has_top_p_filter) would let both entry points share the pipeline. This also makes future renorm-stage changes (e.g. min_p support) a single edit.

3. onehot_draft hardcoded to 1

sampling.rs:316 passes a literal 1 for onehot_draft, even though the FFI signature already accepts it (ffi/shared.rs:140). The PR description notes real-q support is intended for a future sampled drafter — exposing onehot_draft: bool in the Rust API now would avoid a signature change in slice 2, and true as i32 is at least self-documenting.

Minor notes (optional)

• gpu_spec_accept_into doesn't check vocab > 0 (sampling.rs:284-294): with vocab == 0 the buffer-size checks all pass (batch * K * 0 = 0), but ChainSpeculativeSampling's behavior on vocab_size=0 is unverified. A one-line ensure!(vocab > 0, ...) would match gpu_verify_probs_into's implicit guarantee via logits.hidden_dim == scratch.vocab.
• Test coverage is corner-point only: spec_accept_full_acceptance_and_certain_rejection proves the kernel doesn't crash and handles the deterministic extremes, but rejection sampling's core guarantee is distributional correctness. A partial-acceptance test (draft with ~50% target probability, asserting acceptance frequency converges over many seeds) would strengthen confidence — though understandable if deferred given the difficulty of statistical tests on GPU.

The seam design (shape-compatible with accept_greedy's "longest accepted prefix + one model token" contract) is clean. Once the file split and the duplication are addressed, this looks good to merge.

[n-WN]

Thanks for the careful review — all four points addressed, plus the statistical test you suggested. Rebased onto main (dd7019e).

1. File split. Moved gpu_verify_probs_into + gpu_spec_accept_into (and their tests) into a new ops/spec_sampling.rs. sampling.rs drops from 1336 → 1060; the remainder is essentially the pre-existing sampler + upstream's new argmax_bf16_* — the speculative code no longer contributes to the overage.

2. Duplication.

• C: extracted gather_softmax_renorm_flashinfer(...) (gather → softmax → optional top-k/top-p renorm). gpu_verify_probs_flashinfer_cuda now is that call; gpu_sample_batch_flashinfer_cuda's min_p path calls it with the filters on, the fast path with them off — behaviour-identical, ~30 duplicated lines gone.
• Rust: extracted prepare_sampling_params(rows, seq_len, vocab) -> SamplingParams (validated row_indices/temperature/top_k/top_p + filter flags), shared by the batched sampler and the verify path. min_p stays caller-side since the two paths disagree on it (sampler packs it, verify rejects it).

3. onehot_draft. Now a bool parameter on gpu_spec_accept_into instead of the hardcoded 1, so the slice-2 verify path can pass a real drafter's draft_probs (onehot_draft=false) without touching the kernel.

4. vocab > 0. Added to the gpu_spec_accept_into guard.

5. Statistical test. Added spec_accept_partial_acceptance_matches_the_rate_law: a onehot proposer whose draft token holds 0.7 of the target mass, run over 512 i.i.d. rows (per-row philox decorrelates them), asserting the empirical acceptance rate lands in [0.63, 0.77] — 3σ around the exact min(1, 0.7) law, and checking accepted rows emit [A, bonus] while rejected rows emit the residual [B, -1]. This is the p∈(0,1) case the corner-point tests didn't cover.

[xiaguan]

Thanks for the contribution — the math checks out, the slice boundary is clean, and the corner-point tests are well designed. I also ran the suite locally on an RTX 5070 Ti (SM120): 11/11 pass, so the consumer-GPU side is covered too.

A few things I'd like addressed before merge:

1. Missing FFI guard on the two new extern "C" entry points. gpu_chain_speculative_sampling_cuda and gpu_verify_probs_flashinfer_cuda lack OPENINFER_FFI_GUARD_BEGIN/END, while gpu_sample_batch_flashinfer_cuda in the same file has it. The guard exists precisely because FlashInfer can throw from host-side dispatch, and a C++ exception crossing the FFI boundary aborts the process with no message (see 8028dfa). Note the shared gather_softmax_renorm_flashinfer currently runs guarded on the sampler path but unguarded on the verify path.

2. Return a named struct instead of (Vec<i32>, Vec<i32>). The PR body itself calls out that committing on accepted would emit wrong tokens — but the API returns two same-typed Vecs positionally, so one swapped destructuring compiles fine and silently corrupts output. Something like SpecAcceptCounts { emitted, acceptance_telemetry } encodes the invariant in the type and the trap disappears.

3. Seeded-request replay. ChainSpeculativeSampling folds the row index into the philox subsequence, so a fixed-seed request's draw stream depends on batch composition — the exact trap the sampler comment in this file documents and works around with n_rows=1 calls. The doc says "same discipline as the batched sampler", but no such per-row path exists here. Please state explicitly that per-request seeded replay is not representable and seeded requests must stay off the speculative path (and make sure the slice-2 gate enforces it, like the min_p ensure! already does).

4. Per-call alloc_zeros for the counters. This function will run every verify step in slice 2; the codebase convention is allocate-once scratch (BatchSamplingScratch / SampleScratch: "never reallocate per step"). Moving the two counter buffers into a scratch now avoids a signature change later.

Two smaller notes, fine to defer:

• The accepted/emitted distinction test asserts accepted[1] >= 0, which is trivially true; assert_eq!(accepted[1], 1) holds with the same confidence as the other corner-point assertions. The onehot_draft=false path also has no coverage yet.
• For slice 2: please route the executor through openinfer-sample (the policy crate every model routes through, see dd58508) rather than importing these primitives from openinfer_kernels::ops directly — a one-line re-export in openinfer-sample keeps the single-entry invariant. The dense one-hot materialization (memset of batch·K·vocab·4B per step) is also worth an A/B measurement in slice 2 before we rely on it.

Looking forward to the next revision!

[n-WN]

All four addressed (plus both minor notes), rebased onto current main, re-verified on H100 — 12/12 (batch_sampling 7, spec_accept 3, verify_probs 2):

1. FFI guards: both new extern "C" entries now wrapped in OPENINFER_FFI_GUARD_BEGIN/END, and the Rust callers append ffi_exception_message — so the shared gather_softmax_renorm_flashinfer runs guarded on both paths.
2. Named struct: gpu_spec_accept_into returns SpecAcceptCounts { emitted, acceptance_telemetry }; the doc on acceptance_telemetry says "never commit on this" and the swapped-destructuring trap is gone by construction.
3. Seeded replay: documented as not representable — the chain kernel folds the row index into the philox subsequence, so a fixed-seed request's stream depends on batch composition, and no per-row path exists here. The doc now states seeded requests must stay off the speculative path, enforced by the slice-2 gate exactly like the existing min_p ensure! (spec_verify_supported already checks seed.is_none() on the slice-2 branch; I'll carry that through when slice 2 goes up).
4. Scratch counters: SpecAcceptScratch::new(ctx, max_batch) (allocate-once, memset per call) — signature changed now so slice 2 doesn't have to.

Minor notes: the telemetry corner assertion is now assert_eq!(accepted[1], 1) (with a comment deriving why 1 exactly), and there's a new onehot_draft=false corner-point test (explicit proposal: certain accept at q=0.5/p=1.0, certain reject → residual resample). For slice 2 I'll route the executor through openinfer-sample re-exports per dd58508 and A/B the dense one-hot materialization before relying on it.

[xiaguan]

Thanks for the fast turnaround on the review — the revision reads well.

One change of plan on scoping, though: let's not land the kernels slice standalone. Please carry this straight through slice 2 (executor wiring + gate relaxation), either by extending this PR or stacking the slice-2 PR on top of it immediately — whichever you prefer, but they merge together or back-to-back. Corner-point + rate-law unit tests are good evidence for the primitives, but #512 is a 0.2.0 milestone item and what we accept it on is end-to-end accuracy and performance, which only exist once the executor is wired.

Concretely, the acceptance bar for the combined work:

Accuracy (e2e, real model — Qwen3-4B + the DFlash drafter):

• Greedy speculative path byte-identical to today (existing gates stay green — this change must be invisible to greedy).
• Sampled spec-on vs spec-off statistical equivalence on a fixed prompt set across the representative sampling configs (temperature / top_k / top_p): distribution-level evidence, not eyeballing — e.g. many-seed per-position token-frequency comparison, plus an lm-eval sampled-score parity run as the coarse gate.
• Gate tests proving min_p and seeded requests demonstrably stay off the speculative path (the two not-representable cases), and hf_golden_gate unaffected.

Performance (e2e):

• Non-greedy single-stream + concurrent A/B (c1/c8/c16, same protocol as the DFlash greedy numbers in docs/models/qwen3/dflash-speculative-decoding.md): spec-on must beat plain decode for sampled requests — that speedup is the entire point of feat(qwen3): rejection sampling for speculative decoding (non-greedy requests) #512, so it's a merge gate, not a nice-to-have.
• A mixed batch (greedy + sampled) must no longer knock the whole batch off speculation, and greedy-only TPOT must show no regression.
• The dense one-hot materialization A/B you already planned (batch·K·vocab·4B memset per step) — if it costs real TPOT, we need the alternative before merge, not after.

Plumbing stays as agreed: executor goes through openinfer-sample re-exports, and the slice-2 gate enforces seed.is_none() the way the min_p ensure! does.

If any of this is under-specified, raise it here before writing code — happy to pin down the exact statistical test / bench protocol.

[n-WN]

Agreed — merging the primitives without the executor evidence would leave #512's actual claim untested. Plan: stacked slice-2 PR on top of this one, merged back-to-back (slice 2 already exists as a branch: executor wiring through the openinfer-sample re-exports, scratch hoisted to executor init, greedy path untouched; lib tests passed against the pre-rebase head, and I'll restack it on the rebased one). I'll keep this PR as the reviewed-primitives base and open slice 2 once the e2e evidence below is in hand.

Two things done meanwhile:

• Rebased this branch onto current main (post-perf(glm52): graph the dspark draft round — markov chain + piecewise forward (−6.8% bs=1) #591): only union conflicts in the ops.rs re-export list; the spec-sampling sources are byte-identical before/after. Re-running the H100 suite on the rebased head before I push.
• The remaining slice-2 work items per your bar: per-request spec routing (mixed batch), the seed.is_none() gate, and the e2e accuracy/perf evidence.

Before writing the evidence harness, pinning down the two under-specified pieces — please confirm or adjust:

Statistical equivalence (sampled spec-on vs spec-off):

• Prompt set: 8 fixed prompts (from the sharegpt slice used for the DFlash greedy bench), 32 generated positions each.
• Configs: {t=0.8, top_p=0.95}, {t=1.0, top_k=50}, {t=0.7, top_k=20, top_p=0.9} — one pure-top_p, one pure-top_k, one combined.
• Per (prompt, config): 256 seeds per arm, per-position token histograms; two-sample chi-squared (merged support, expected-count ≥5 binning), Benjamini–Hochberg FDR at 5% across all positions×prompts×configs. Gate: zero rejections after FDR, and report the max per-position total-variation distance as the effect-size readout. (~25k short generations per arm — a few H100-hours.)
• Coarse gate: lm-eval GSM8K sampled (t=0.8), 3 seeds per arm, gate: overlapping 95% CIs.
• Sanity teeth: the same harness run spec-on-vs-spec-on (different seeds) must pass, and a deliberately biased acceptor (ε-inflated draft prob) must fail — so the test demonstrably detects real distribution shift.

Perf:

• Same protocol as docs/models/qwen3/dflash-speculative-decoding.md: c1/c8/c16, sampled at {t=0.8, top_p=0.95}, spec-on vs spec-off TPOT; merge gate = spec-on wins at every concurrency.
• Mixed batch: c8 with 4 greedy + 4 sampled — gates: greedy TPOT no regression vs all-greedy c8, sampled streams keep a speedup vs spec-off.
• One-hot A/B: measured TPOT delta of the per-step batch·K·vocab·4B memset; if it costs >1% TPOT at c8, I switch to writing only the K accepted one-hot rows via a small scatter kernel (indices are already on device) before merge.

If the seed/prompt counts or the lm-eval task should be different (or you want a specific prompt list checked in for reproducibility), say the word — otherwise I'll build exactly the above and check the harness in next to the smoke.