diff --git a/openinfer-kernels/csrc/shared/flashinfer_sampling.cu b/openinfer-kernels/csrc/shared/flashinfer_sampling.cu index 92d35952..e1b8cde6 100644 --- a/openinfer-kernels/csrc/shared/flashinfer_sampling.cu +++ b/openinfer-kernels/csrc/shared/flashinfer_sampling.cu @@ -66,67 +66,162 @@ extern "C" size_t gpu_sample_topk_renorm_row_states_bytes_cuda() { // their own n_rows=1 calls by the Rust layer, with the request seed and step // mixed into `seed` — blockIdx is then always 0 and replay is independent of // batch composition. -extern "C" int gpu_sample_batch_flashinfer_cuda( - const __nv_bfloat16* logits, const int* row_indices, float* probs_scratch, - const float* temperature_arr, const int* top_k_arr, const float* top_p_arr, - const float* min_p_arr, uint8_t* topk_row_states_scratch, - uint8_t* valid_scratch, int* output, void* softmax_workspace, - size_t softmax_workspace_bytes, int n_rows, int vocab_size, int has_top_k_filter, - int has_top_p_filter, uint64_t seed, uint64_t offset, cudaStream_t stream) { +// One-hot proposal distributions for a greedy proposer: draft_probs is +// zeroed and draft_probs[i * vocab + draft_token_ids[i]] = 1.0 for each of +// the batch*K proposed tokens. A deterministic (argmax) draft is the +// degenerate proposal q(x) = delta(x - draft), under which chain rejection +// sampling accepts with min(1, p_target(draft)) and resamples from the +// target with the draft token's mass removed — still distribution-exact. +__global__ void onehot_rows_kernel(float* probs, const int* token_ids, int vocab) { + const int row = blockIdx.x; + const int id = token_ids[row]; + if (id >= 0 && id < vocab && threadIdx.x == 0) { + probs[static_cast(row) * vocab + id] = 1.0f; + } +} + +// Chain speculative (rejection) sampling over one verify span per row: +// accepts each draft token with min(1, p_target/q_draft), resamples the first +// rejection from relu(target - draft) renormalized, appends the bonus token +// on full acceptance, and -1-fills the tail. accepted/emitted counters are +// ACCUMULATED by the kernel, so the caller must zero them first. +extern "C" int gpu_chain_speculative_sampling_cuda( + float* draft_probs, const int* draft_token_ids, float* target_probs, + int* output_token_ids, int* output_accepted_num, int* output_emitted_num, + int batch_size, int num_speculative_tokens, int vocab_size, int onehot_draft, + uint64_t seed, uint64_t offset, cudaStream_t stream) { OPENINFER_FFI_GUARD_BEGIN + if (onehot_draft) { + const size_t total = static_cast(batch_size) * num_speculative_tokens; + cudaError_t err = cudaMemsetAsync( + draft_probs, 0, total * vocab_size * sizeof(float), stream); + if (err != cudaSuccess) { + return static_cast(err); + } + onehot_rows_kernel<<(total), 32, 0, stream>>>( + draft_probs, draft_token_ids, vocab_size); + err = cudaGetLastError(); + if (err != cudaSuccess) { + return static_cast(err); + } + } + cudaError_t err = flashinfer::sampling::ChainSpeculativeSampling( + draft_probs, const_cast(draft_token_ids), target_probs, output_token_ids, + output_accepted_num, output_emitted_num, batch_size, num_speculative_tokens, + vocab_size, /*deterministic=*/true, /*seed_arr=*/nullptr, seed, + /*offset_arr=*/nullptr, offset, stream); + return static_cast(err); + OPENINFER_FFI_GUARD_END(-1) +} + +// Shared front half of the batched sampling pipeline: gather the requested +// bf16 rows into f32 `probs`, softmax with per-row temperature, then optionally +// top-k and top-p renormalize in place — leaving `probs` a proper distribution +// with filtered tokens as exact zeros. The verify path stops here; the min_p +// sampling path draws afterward. Returns the cudaError as int (0 == success). +static int gather_softmax_renorm_flashinfer( + const __nv_bfloat16* logits, const int* row_indices, float* probs, + const float* temperature_arr, const int* top_k_arr, const float* top_p_arr, + uint8_t* topk_row_states_scratch, void* softmax_workspace, + size_t softmax_workspace_bytes, int n_rows, int vocab_size, bool do_top_k, + bool do_top_p, cudaStream_t stream) { dim3 gather_grid( (vocab_size + GATHER_CAST_BLOCK * GATHER_CAST_ELEMS_PER_THREAD - 1) / (GATHER_CAST_BLOCK * GATHER_CAST_ELEMS_PER_THREAD), n_rows); gather_cast_logits_f32_kernel<<>>( - logits, row_indices, probs_scratch, vocab_size); + logits, row_indices, probs, vocab_size); cudaError_t err = cudaGetLastError(); if (err != cudaSuccess) { return static_cast(err); } - // In-place: phase 2 of both OnlineSoftmax strategies is an elementwise // read-then-write of the same index. err = flashinfer::sampling::OnlineSoftmax( - probs_scratch, probs_scratch, n_rows, vocab_size, const_cast(temperature_arr), + probs, probs, n_rows, vocab_size, const_cast(temperature_arr), /*temperature_val=*/1.0f, softmax_workspace, softmax_workspace_bytes, /*enable_pdl=*/false, stream); if (err != cudaSuccess) { return static_cast(err); } + if (do_top_k) { + auto* row_states = + reinterpret_cast(topk_row_states_scratch); + if (row_states == nullptr) { + return static_cast(cudaErrorInvalidValue); + } + // In-place: the renorm kernels reduce a threshold first, then rewrite + // each element from its own index. + err = flashinfer::sampling::RadixTopKRenormProbMultiCTA( + probs, probs, const_cast(top_k_arr), n_rows, + /*top_k_val=*/0, vocab_size, row_states, stream); + if (err != cudaSuccess) { + return static_cast(err); + } + } + if (do_top_p) { + err = flashinfer::sampling::TopPRenormProb( + probs, probs, const_cast(top_p_arr), n_rows, + /*top_p_val=*/0.0f, vocab_size, stream); + } + return static_cast(err); +} + +// The verify-side half of the batched sampling pipeline: gather + softmax + +// top-k/top-p renorm, WITHOUT the terminal sampling kernel — the renormalized +// probabilities stay in probs_out for a downstream consumer (speculative +// rejection sampling needs the target distribution itself, filtered tokens as +// exact zeros). Distribution-equivalent to the fused sampling fast path: the +// rejection sampler filters at draw time, the renorm filters then draws — the +// resulting law over tokens is identical. +extern "C" int gpu_verify_probs_flashinfer_cuda( + const __nv_bfloat16* logits, const int* row_indices, float* probs_out, + const float* temperature_arr, const int* top_k_arr, const float* top_p_arr, + uint8_t* topk_row_states_scratch, void* softmax_workspace, + size_t softmax_workspace_bytes, int n_rows, int vocab_size, int has_top_k_filter, + int has_top_p_filter, cudaStream_t stream) { + OPENINFER_FFI_GUARD_BEGIN + return gather_softmax_renorm_flashinfer( + logits, row_indices, probs_out, temperature_arr, top_k_arr, top_p_arr, + topk_row_states_scratch, softmax_workspace, softmax_workspace_bytes, n_rows, + vocab_size, has_top_k_filter != 0, has_top_p_filter != 0, stream); + OPENINFER_FFI_GUARD_END(-1) +} +extern "C" int gpu_sample_batch_flashinfer_cuda( + const __nv_bfloat16* logits, const int* row_indices, float* probs_scratch, + const float* temperature_arr, const int* top_k_arr, const float* top_p_arr, + const float* min_p_arr, uint8_t* topk_row_states_scratch, + uint8_t* valid_scratch, int* output, void* softmax_workspace, + size_t softmax_workspace_bytes, int n_rows, int vocab_size, int has_top_k_filter, + int has_top_p_filter, uint64_t seed, uint64_t offset, cudaStream_t stream) { + OPENINFER_FFI_GUARD_BEGIN bool* valid = reinterpret_cast(valid_scratch); if (min_p_arr != nullptr) { - if (has_top_k_filter) { - auto* row_states = - reinterpret_cast(topk_row_states_scratch); - if (row_states == nullptr) { - return static_cast(cudaErrorInvalidValue); - } - // In-place: the renorm kernels reduce a threshold first, then rewrite - // each element from its own index. - err = flashinfer::sampling::RadixTopKRenormProbMultiCTA( - probs_scratch, probs_scratch, const_cast(top_k_arr), n_rows, - /*top_k_val=*/0, vocab_size, row_states, stream); - if (err != cudaSuccess) { - return static_cast(err); - } + // min_p reads a renormalized distribution, so filter before drawing. + int rc = gather_softmax_renorm_flashinfer( + logits, row_indices, probs_scratch, temperature_arr, top_k_arr, top_p_arr, + topk_row_states_scratch, softmax_workspace, softmax_workspace_bytes, n_rows, + vocab_size, has_top_k_filter != 0, has_top_p_filter != 0, stream); + if (rc != 0) { + return rc; } - if (has_top_p_filter) { - err = flashinfer::sampling::TopPRenormProb( - probs_scratch, probs_scratch, const_cast(top_p_arr), n_rows, - /*top_p_val=*/0.0f, vocab_size, stream); - if (err != cudaSuccess) { - return static_cast(err); - } - } - err = flashinfer::sampling::MinPSamplingFromProb( + cudaError_t err = flashinfer::sampling::MinPSamplingFromProb( probs_scratch, const_cast(min_p_arr), output, valid, /*indices=*/nullptr, n_rows, /*min_p_val=*/0.0f, vocab_size, /*deterministic=*/true, /*seed_arr=*/nullptr, seed, /*offset_arr=*/nullptr, offset, stream); return static_cast(err); } + // Fast path: gather + softmax only — the fused sampler filters at draw time. + int rc = gather_softmax_renorm_flashinfer( + logits, row_indices, probs_scratch, temperature_arr, top_k_arr, top_p_arr, + topk_row_states_scratch, softmax_workspace, softmax_workspace_bytes, n_rows, + vocab_size, /*do_top_k=*/false, /*do_top_p=*/false, stream); + if (rc != 0) { + return rc; + } + cudaError_t err; if (has_top_k_filter) { err = flashinfer::sampling::TopKTopPSamplingFromProb( probs_scratch, const_cast(top_k_arr), const_cast(top_p_arr), output, valid, diff --git a/openinfer-kernels/src/ffi/shared.rs b/openinfer-kernels/src/ffi/shared.rs index f6765732..ed16926f 100644 --- a/openinfer-kernels/src/ffi/shared.rs +++ b/openinfer-kernels/src/ffi/shared.rs @@ -128,6 +128,39 @@ unsafe extern "C" { pub fn gpu_sample_topk_renorm_row_states_bytes_cuda() -> usize; + pub fn gpu_chain_speculative_sampling_cuda( + draft_probs: *mut f32, + draft_token_ids: *const i32, + target_probs: *mut f32, + output_token_ids: *mut i32, + output_accepted_num: *mut i32, + output_emitted_num: *mut i32, + batch_size: i32, + num_speculative_tokens: i32, + vocab_size: i32, + onehot_draft: i32, + seed: u64, + offset: u64, + stream: CUstream, + ) -> i32; + + pub fn gpu_verify_probs_flashinfer_cuda( + logits: *const Half, + row_indices: *const i32, + probs_out: *mut f32, + temperature_arr: *const f32, + top_k_arr: *const i32, + top_p_arr: *const f32, + topk_row_states_scratch: *mut u8, + softmax_workspace: *mut u8, + softmax_workspace_bytes: usize, + n_rows: i32, + vocab_size: i32, + has_top_k_filter: i32, + has_top_p_filter: i32, + stream: CUstream, + ) -> i32; + pub fn gpu_sample_batch_flashinfer_cuda( logits: *const Half, row_indices: *const i32, diff --git a/openinfer-kernels/src/ops.rs b/openinfer-kernels/src/ops.rs index 9dc626ac..e6feb392 100644 --- a/openinfer-kernels/src/ops.rs +++ b/openinfer-kernels/src/ops.rs @@ -15,6 +15,7 @@ mod linear; mod lora; mod norm; mod sampling; +mod spec_sampling; pub use attention::{ PrefillPagedPlan, SUPPORTED_GQA_GROUP_SIZES, dflash_qk_norm_rope_into, @@ -72,6 +73,9 @@ pub use sampling::{ argmax_bf16_split_into, flashinfer_top1_batch_into, flashinfer_top1_row_states_bytes, gpu_sample_batch_into, markov_step_argmax_into, markov_step_argmax_partials_len, }; +pub use spec_sampling::{ + SpecAcceptCounts, SpecAcceptScratch, gpu_spec_accept_into, gpu_verify_probs_into, +}; /// Calling thread's last FFI exception message, ready to append to an error; /// empty unless `result` is the -1 sentinel set by the C++ guard. Public for diff --git a/openinfer-kernels/src/ops/sampling.rs b/openinfer-kernels/src/ops/sampling.rs index e955e0aa..03f249fa 100644 --- a/openinfer-kernels/src/ops/sampling.rs +++ b/openinfer-kernels/src/ops/sampling.rs @@ -24,15 +24,17 @@ pub struct BatchSamplingRow { /// Device buffers for `gpu_sample_batch_into`, sized for `max_rows` x `vocab`. pub struct BatchSamplingScratch { probs: CudaSlice, - row_indices: CudaSlice, - temperature: CudaSlice, - top_k: CudaSlice, - top_p: CudaSlice, + // The parameter and workspace buffers are read by the sibling + // `spec_sampling` module's verify path, so they are visible to `ops`. + pub(super) row_indices: CudaSlice, + pub(super) temperature: CudaSlice, + pub(super) top_k: CudaSlice, + pub(super) top_p: CudaSlice, min_p: CudaSlice, - topk_row_states: CudaSlice, + pub(super) topk_row_states: CudaSlice, valid: CudaSlice, out: CudaSlice, - softmax_workspace: CudaSlice, + pub(super) softmax_workspace: CudaSlice, max_rows: usize, vocab: usize, } @@ -71,6 +73,84 @@ impl BatchSamplingScratch { pub fn max_rows(&self) -> usize { self.max_rows } + + pub(super) fn vocab(&self) -> usize { + self.vocab + } +} + +/// Validated per-row sampling parameters, packed into the layout the +/// FlashInfer FFI expects. `min_p` is deliberately excluded: the verify path +/// forbids it and the sampling path packs it separately, so the shared +/// preparation stops at the temperature/top-k/top-p triple. +pub(super) struct SamplingParams { + pub row_indices: Vec, + pub temperature: Vec, + pub top_k: Vec, + pub top_p: Vec, + pub has_top_k_filter: bool, + pub has_top_p_filter: bool, +} + +/// Validate `rows` against `seq_len`/`vocab` and pack the shared +/// temperature/top-k/top-p parameters. `top_k` is clamped to `vocab` when +/// disabled (FlashInfer reads "k >= vocab" as off). Rejects greedy rows +/// (`temperature <= 0`) and out-of-range `top_p` — both the batched sampler +/// and the verify-probs path share these invariants. `min_p` is left to the +/// caller: the sampler allows it, the verify path rejects it. +pub(super) fn prepare_sampling_params( + rows: &[BatchSamplingRow], + seq_len: usize, + vocab: usize, +) -> Result { + let n = rows.len(); + let vocab_i32 = i32::try_from(vocab)?; + let mut row_indices = Vec::with_capacity(n); + let mut temperature = Vec::with_capacity(n); + let mut top_k = Vec::with_capacity(n); + let mut top_p = Vec::with_capacity(n); + let mut has_top_k_filter = false; + let mut has_top_p_filter = false; + for r in rows { + ensure!( + r.row < seq_len, + "sampling row {} out of arena range {}", + r.row, + seq_len + ); + ensure!( + r.temperature > 0.0 && r.temperature.is_finite(), + "sampling temperature {} must be finite and > 0 (greedy rows take the argmax path)", + r.temperature + ); + ensure!( + r.top_p > 0.0 && r.top_p <= 1.0, + "sampling top_p {} must be in (0, 1]", + r.top_p + ); + row_indices.push(i32::try_from(r.row)?); + temperature.push(r.temperature); + // FlashInfer reads top_k as u32; "disabled" is any k >= vocab. + let clamped_top_k = if r.top_k > 0 && r.top_k < vocab_i32 { + has_top_k_filter = true; + r.top_k + } else { + vocab_i32 + }; + top_k.push(clamped_top_k); + if r.top_p < 1.0 { + has_top_p_filter = true; + } + top_p.push(r.top_p); + } + Ok(SamplingParams { + row_indices, + temperature, + top_k, + top_p, + has_top_k_filter, + has_top_p_filter, + }) } /// Batched temperature/top-k/top-p sampling: gathers the requested bf16 arena @@ -142,62 +222,30 @@ fn sample_uniform_batch_into( scratch.vocab ); - let mut row_indices = Vec::with_capacity(n); - let mut temperature = Vec::with_capacity(n); - let mut top_k = Vec::with_capacity(n); - let mut top_p = Vec::with_capacity(n); + let params = prepare_sampling_params(rows, logits.seq_len, scratch.vocab)?; + let has_top_k_filter = params.has_top_k_filter; + let has_top_p_filter = params.has_top_p_filter; + // min_p is sampler-only (the verify path forbids it): validate and pack it + // here rather than in the shared parameter helper. let mut min_p = Vec::with_capacity(n); - let mut has_top_k_filter = false; - let mut has_top_p_filter = false; let mut has_min_p_filter = false; for r in rows { - ensure!( - r.row < logits.seq_len, - "batch sampling row {} out of arena range {}", - r.row, - logits.seq_len - ); - ensure!( - r.temperature > 0.0 && r.temperature.is_finite(), - "batch sampling temperature {} must be finite and > 0 (greedy rows take the argmax path)", - r.temperature - ); - ensure!( - r.top_p > 0.0 && r.top_p <= 1.0, - "batch sampling top_p {} must be in (0, 1]", - r.top_p - ); ensure!( (0.0..1.0).contains(&r.min_p) && r.min_p.is_finite(), "batch sampling min_p {} must be in [0, 1)", r.min_p ); - row_indices.push(i32::try_from(r.row)?); - temperature.push(r.temperature); - // FlashInfer reads top_k as u32; "disabled" is any k >= vocab. - let vocab = i32::try_from(scratch.vocab)?; - let clamped_top_k = if r.top_k > 0 && r.top_k < vocab { - has_top_k_filter = true; - r.top_k - } else { - vocab - }; - top_k.push(clamped_top_k); - if r.top_p < 1.0 { - has_top_p_filter = true; - } - top_p.push(r.top_p); if r.min_p > 0.0 { has_min_p_filter = true; } min_p.push(r.min_p); } ctx.stream - .memcpy_htod(&row_indices, &mut scratch.row_indices)?; + .memcpy_htod(¶ms.row_indices, &mut scratch.row_indices)?; ctx.stream - .memcpy_htod(&temperature, &mut scratch.temperature)?; - ctx.stream.memcpy_htod(&top_k, &mut scratch.top_k)?; - ctx.stream.memcpy_htod(&top_p, &mut scratch.top_p)?; + .memcpy_htod(¶ms.temperature, &mut scratch.temperature)?; + ctx.stream.memcpy_htod(¶ms.top_k, &mut scratch.top_k)?; + ctx.stream.memcpy_htod(¶ms.top_p, &mut scratch.top_p)?; if has_min_p_filter { ctx.stream.memcpy_htod(&min_p, &mut scratch.min_p)?; } @@ -689,20 +737,23 @@ pub fn flashinfer_top1_batch_into( Ok(()) } +/// Test fixtures shared with the sibling `spec_sampling` module: a poisoned +/// logits arena and a flat-row builder. `pub(crate)` so `spec_sampling`'s +/// tests gather the same arena rows without duplicating the helper. #[cfg(test)] -mod tests { +pub(crate) mod test_support { use half::bf16; - use super::*; + use crate::tensor::{DeviceContext, HiddenStates}; - const VOCAB: usize = 32768; // >= 24576 so OnlineSoftmax takes the vocab-splitting path - const ARENA_ROWS: usize = 8; + pub(crate) const VOCAB: usize = 32768; // >= 24576 so OnlineSoftmax takes the vocab-splitting path + pub(crate) const ARENA_ROWS: usize = 8; /// Arena where every row not under test is poisoned with a dominant logit /// at `POISON_TOKEN` — a broken row gather makes every assertion fail. const POISON_TOKEN: usize = 7777; - fn arena_with_rows(ctx: &DeviceContext, rows: &[(usize, Vec)]) -> HiddenStates { + pub(crate) fn arena_with_rows(ctx: &DeviceContext, rows: &[(usize, Vec)]) -> HiddenStates { let mut host = vec![bf16::from_f32(0.0); ARENA_ROWS * VOCAB]; for r in 0..ARENA_ROWS { host[r * VOCAB + POISON_TOKEN] = bf16::from_f32(20.0); @@ -721,9 +772,15 @@ mod tests { } } - fn flat_row(fill: f32) -> Vec { + pub(crate) fn flat_row(fill: f32) -> Vec { vec![fill; VOCAB] } +} + +#[cfg(test)] +mod tests { + use super::test_support::{ARENA_ROWS, VOCAB, arena_with_rows, flat_row}; + use super::*; #[test] fn batch_sampling_honors_top_k_top_p_and_gathers_rows() { diff --git a/openinfer-kernels/src/ops/spec_sampling.rs b/openinfer-kernels/src/ops/spec_sampling.rs new file mode 100644 index 00000000..f478a57e --- /dev/null +++ b/openinfer-kernels/src/ops/spec_sampling.rs @@ -0,0 +1,510 @@ +//! Speculative-decoding sampling primitives: the verify-side target +//! distribution (`gpu_verify_probs_into`) and chain rejection sampling +//! (`gpu_spec_accept_into`). Split out of `sampling.rs` to keep that file +//! under the module size budget; the shared parameter packing and scratch +//! buffers live in the sibling `sampling` module. + +use anyhow::{Result, ensure}; +use cudarc::driver::{CudaSlice, DevicePtr, DevicePtrMut}; + +use super::sampling::{BatchSamplingRow, BatchSamplingScratch, prepare_sampling_params}; +use crate::ffi; +use crate::tensor::{DeviceContext, HiddenStatesRef}; + +/// The target distribution at each requested arena row: gather + softmax + +/// top-k/top-p renorm, written to `probs_out` as `n_rows x vocab` f32 — +/// filtered tokens are exact zeros. This is the verify-side input to +/// speculative rejection sampling; it is distribution-equivalent to the +/// sampling fast path (the fused sampler filters at draw time, this filters +/// then draws). +/// +/// `min_p` rows are rejected: the min_p mask is applied inside the sampling +/// kernel, not as a renorm, so a min_p target distribution is not +/// representable here yet — callers must keep such requests off the +/// speculative path. +pub fn gpu_verify_probs_into( + ctx: &DeviceContext, + logits: HiddenStatesRef<'_>, + rows: &[BatchSamplingRow], + probs_out: &mut CudaSlice, + scratch: &mut BatchSamplingScratch, +) -> Result<()> { + let n = rows.len(); + ensure!(n > 0, "verify probs requires at least one row"); + ensure!( + n <= scratch.max_rows(), + "verify probs scratch too small: {n} rows > capacity {}", + scratch.max_rows() + ); + ensure!( + logits.hidden_dim == scratch.vocab(), + "verify probs vocab mismatch: logits {} vs scratch {}", + logits.hidden_dim, + scratch.vocab() + ); + ensure!( + probs_out.len() >= n * scratch.vocab(), + "verify probs output {} < {n} x {}", + probs_out.len(), + scratch.vocab() + ); + // A min_p target distribution is not representable as a renorm, so reject + // such rows before packing the shared parameters. + for r in rows { + ensure!( + r.min_p == 0.0, + "verify probs cannot represent a min_p target distribution (min_p {})", + r.min_p + ); + } + + let params = prepare_sampling_params(rows, logits.seq_len, scratch.vocab())?; + let has_top_k_filter = params.has_top_k_filter; + let has_top_p_filter = params.has_top_p_filter; + ctx.stream + .memcpy_htod(¶ms.row_indices, &mut scratch.row_indices)?; + ctx.stream + .memcpy_htod(¶ms.temperature, &mut scratch.temperature)?; + ctx.stream.memcpy_htod(¶ms.top_k, &mut scratch.top_k)?; + ctx.stream.memcpy_htod(¶ms.top_p, &mut scratch.top_p)?; + + let vocab = scratch.vocab(); + let softmax_workspace_bytes = scratch.softmax_workspace.len(); + let (logits_ptr, _gl) = logits.data.device_ptr(&ctx.stream); + let (indices_ptr, _gi) = scratch.row_indices.device_ptr(&ctx.stream); + let (probs_ptr, _gp) = probs_out.device_ptr_mut(&ctx.stream); + let (temp_ptr, _gt) = scratch.temperature.device_ptr(&ctx.stream); + let (top_k_ptr, _gk) = scratch.top_k.device_ptr(&ctx.stream); + let (top_p_ptr, _gtp) = scratch.top_p.device_ptr(&ctx.stream); + let row_states = if has_top_k_filter { + Some(scratch.topk_row_states.device_ptr_mut(&ctx.stream)) + } else { + None + }; + let row_states_ptr = row_states.as_ref().map_or(0, |(ptr, _guard)| *ptr); + let (ws_ptr, _gw) = scratch.softmax_workspace.device_ptr_mut(&ctx.stream); + let err = unsafe { + ffi::gpu_verify_probs_flashinfer_cuda( + logits_ptr as *const ffi::Half, + indices_ptr as *const i32, + probs_ptr as *mut f32, + temp_ptr as *const f32, + top_k_ptr as *const i32, + top_p_ptr as *const f32, + row_states_ptr as *mut u8, + ws_ptr as *mut u8, + softmax_workspace_bytes, + n as i32, + vocab as i32, + i32::from(has_top_k_filter), + i32::from(has_top_p_filter), + crate::tensor::active_cu_stream(ctx), + ) + }; + ensure!(err == 0, "verify probs kernel failed: cudaError {err}"); + Ok(()) +} + +/// Chain speculative (rejection) sampling over per-row verify spans. +/// +/// For each of the `batch` rows: accept draft token `i` with probability +/// `min(1, p_target(x_i) / q_draft(x_i))`; the first rejection resamples from +/// `relu(target − draft)` renormalized and stops; full acceptance emits the +/// bonus token sampled from the target at position K. Output row layout is +/// `K+1` token ids, `-1`-filled after the stop — exactly the "longest +/// accepted prefix + one model token" contract `accept_greedy` has. +/// +/// `onehot_draft` selects the proposal semantics. With `onehot_draft == true` +/// (a greedy/argmax proposer) `draft_probs` is derived on-device from +/// `draft_token_ids` — the degenerate proposal `q(x) = δ(x − draft)`, under +/// which acceptance is `min(1, p_target(draft))` and the residual is the +/// target with the draft token's mass removed. With `onehot_draft == false` +/// the caller must supply the true `draft_probs` the proposer sampled from; +/// rejection sampling stays distribution-exact either way. +/// +/// `target_probs` must be post-filter probabilities (`gpu_verify_probs_into`), +/// `batch x (K+1) x vocab`; `draft_probs` is `batch x K x vocab` scratch. +/// +/// **Seeded requests are not representable here.** The chain kernel folds the +/// row index into the philox subsequence, so a fixed-seed request's draw +/// stream depends on batch composition — the exact replay trap the batched +/// sampler documents and dodges with per-row calls, and a per-row path does +/// not exist for the chain. Callers must keep seeded requests off the +/// speculative path entirely (a gate the verify scheduler enforces the same +/// way it already keeps min_p rows off it). Pass a fresh `seed`/`offset` per +/// verify step for the unseeded rows. +#[allow(clippy::too_many_arguments)] +pub fn gpu_spec_accept_into( + ctx: &DeviceContext, + draft_probs: &mut CudaSlice, + draft_token_ids: &CudaSlice, + target_probs: &mut CudaSlice, + output_token_ids: &mut CudaSlice, + batch: usize, + num_spec_tokens: usize, + vocab: usize, + onehot_draft: bool, + seed: u64, + offset: u64, + scratch: &mut SpecAcceptScratch, +) -> Result { + ensure!( + batch > 0 && num_spec_tokens > 0 && vocab > 0, + "spec accept requires batch > 0, K > 0, and vocab > 0" + ); + ensure!( + batch <= scratch.max_batch, + "spec accept scratch too small: batch {batch} > capacity {}", + scratch.max_batch + ); + ensure!( + draft_probs.len() >= batch * num_spec_tokens * vocab + && target_probs.len() >= batch * (num_spec_tokens + 1) * vocab + && draft_token_ids.len() >= batch * num_spec_tokens + && output_token_ids.len() >= batch * (num_spec_tokens + 1), + "spec accept buffer too small for batch {batch} x K {num_spec_tokens} x vocab {vocab}" + ); + // The kernel accumulates into the counters, so they must start at zero on + // every call — memset instead of a fresh allocation (this runs every + // verify step; scratch buffers are allocate-once by convention). + ctx.stream.memset_zeros(&mut scratch.accepted)?; + ctx.stream.memset_zeros(&mut scratch.emitted)?; + { + let (dp, _g1) = draft_probs.device_ptr_mut(&ctx.stream); + let (ids, _g2) = draft_token_ids.device_ptr(&ctx.stream); + let (tp, _g3) = target_probs.device_ptr_mut(&ctx.stream); + let (out, _g4) = output_token_ids.device_ptr_mut(&ctx.stream); + let (acc, _g5) = scratch.accepted.device_ptr_mut(&ctx.stream); + let (emi, _g6) = scratch.emitted.device_ptr_mut(&ctx.stream); + let err = unsafe { + ffi::gpu_chain_speculative_sampling_cuda( + dp as *mut f32, + ids as *const i32, + tp as *mut f32, + out as *mut i32, + acc as *mut i32, + emi as *mut i32, + batch as i32, + num_spec_tokens as i32, + vocab as i32, + i32::from(onehot_draft), + seed, + offset, + crate::tensor::active_cu_stream(ctx), + ) + }; + ensure!( + err == 0, + "chain speculative sampling failed with error {err}{}", + crate::ops::ffi_exception_message(err) + ); + } + let accepted_view = scratch.accepted.slice(..batch); + let emitted_view = scratch.emitted.slice(..batch); + let acceptance_telemetry = ctx.stream.clone_dtoh(&accepted_view)?; + let emitted = ctx.stream.clone_dtoh(&emitted_view)?; + ctx.sync()?; + Ok(SpecAcceptCounts { + emitted, + acceptance_telemetry, + }) +} + +/// Per-row results of one chain-rejection call, named so the commit signal +/// and the telemetry cannot be swapped silently (they are both `Vec`). +pub struct SpecAcceptCounts { + /// The accepted-prefix length per row — **the commit signal**: the row's + /// output holds `emitted` accepted drafts plus one resampled/bonus token, + /// then `-1` filler. + pub emitted: Vec, + /// FlashInfer's acceptance-rate telemetry: keeps counting hypothetical + /// acceptances *past* the first rejection. Never commit on this. + pub acceptance_telemetry: Vec, +} + +/// Allocate-once device counters for [`gpu_spec_accept_into`] (it runs every +/// verify step; per-call `alloc_zeros` is against the scratch convention). +pub struct SpecAcceptScratch { + accepted: CudaSlice, + emitted: CudaSlice, + max_batch: usize, +} + +impl SpecAcceptScratch { + pub fn new(ctx: &DeviceContext, max_batch: usize) -> Result { + ensure!(max_batch > 0, "spec accept scratch requires max_batch > 0"); + Ok(Self { + accepted: ctx.stream.alloc_zeros(max_batch)?, + emitted: ctx.stream.alloc_zeros(max_batch)?, + max_batch, + }) + } +} + +#[cfg(test)] +mod tests { + use cudarc::driver::CudaSlice; + + use super::super::sampling::test_support::{ARENA_ROWS, VOCAB, arena_with_rows, flat_row}; + use super::super::sampling::{BatchSamplingRow, BatchSamplingScratch}; + use super::{SpecAcceptScratch, gpu_spec_accept_into, gpu_verify_probs_into}; + use crate::tensor::DeviceContext; + + #[test] + fn verify_probs_renorm_matches_the_sampling_law() { + let ctx = DeviceContext::new().expect("create CUDA context"); + // Row with a known 3-token support: P = {0.6652, 0.2447, 0.0900} + // (softmax of {2, 1, 0} with the rest at -120), plus top_k=2 on a + // second view of the same row so the renorm must zero token 300 and + // renormalize the survivors to {0.7311, 0.2689}. + let mut row = flat_row(-120.0); + row[100] = 2.0; + row[200] = 1.0; + row[300] = 0.0; + let logits = arena_with_rows(&ctx, &[(1, row.clone()), (5, row)]); + let mut scratch = BatchSamplingScratch::new(&ctx, ARENA_ROWS, VOCAB).expect("scratch"); + let rows = [ + BatchSamplingRow { + row: 1, + temperature: 1.0, + top_k: -1, + top_p: 1.0, + min_p: 0.0, + }, + BatchSamplingRow { + row: 5, + temperature: 1.0, + top_k: 2, + top_p: 1.0, + min_p: 0.0, + }, + ]; + let mut probs: CudaSlice = ctx.stream.alloc_zeros(2 * VOCAB).expect("probs"); + gpu_verify_probs_into(&ctx, logits.as_ref(), &rows, &mut probs, &mut scratch) + .expect("verify probs"); + let host = ctx.stream.clone_dtoh(&probs).expect("D2H"); + ctx.sync().expect("sync"); + + // Unfiltered row: sums to 1, matches the closed-form softmax. + let r0 = &host[..VOCAB]; + let sum0: f32 = r0.iter().sum(); + assert!((sum0 - 1.0).abs() < 1e-3, "row0 sum {sum0}"); + assert!((r0[100] - 0.6652).abs() < 5e-3, "P(100) {}", r0[100]); + assert!((r0[200] - 0.2447).abs() < 5e-3, "P(200) {}", r0[200]); + assert!((r0[300] - 0.0900).abs() < 5e-3, "P(300) {}", r0[300]); + + // top_k=2 row: token 300 is an exact zero (renorm, not epsilon), the + // survivors renormalize, and the sum is 1 again. + let r1 = &host[VOCAB..2 * VOCAB]; + assert_eq!(r1[300], 0.0, "top-k filtered token must be exactly zero"); + assert!((r1[100] - 0.7311).abs() < 5e-3, "renorm P(100) {}", r1[100]); + assert!((r1[200] - 0.2689).abs() < 5e-3, "renorm P(200) {}", r1[200]); + let sum1: f32 = r1.iter().sum(); + assert!((sum1 - 1.0).abs() < 1e-3, "row1 sum {sum1}"); + } + + #[test] + fn verify_probs_rejects_min_p_rows() { + let ctx = DeviceContext::new().expect("create CUDA context"); + let logits = arena_with_rows(&ctx, &[(0, flat_row(0.0))]); + let mut scratch = BatchSamplingScratch::new(&ctx, ARENA_ROWS, VOCAB).expect("scratch"); + let rows = [BatchSamplingRow { + row: 0, + temperature: 1.0, + top_k: -1, + top_p: 1.0, + min_p: 0.1, + }]; + let mut probs: CudaSlice = ctx.stream.alloc_zeros(VOCAB).expect("probs"); + let err = gpu_verify_probs_into(&ctx, logits.as_ref(), &rows, &mut probs, &mut scratch) + .expect_err("min_p rows must be rejected"); + assert!(err.to_string().contains("min_p"), "unexpected error: {err}"); + } + + #[test] + fn spec_accept_full_acceptance_and_certain_rejection() { + let ctx = DeviceContext::new().expect("create CUDA context"); + let k = 2usize; + // Hand-built target probs, batch=2, K+1=3 positions each. + // Row 0: target puts prob 1.0 on the draft token at every position → + // certain acceptance; bonus position puts 1.0 on token 777. + // Row 1: target puts 0.0 on the draft at position 0 (prob 1.0 on 555) + // → certain first-step rejection; the residual is exactly {555: 1.0}. + let mut target = vec![0.0f32; 2 * (k + 1) * VOCAB]; + target[100] = 1.0; // row0 pos0 -> 100 + target[VOCAB + 200] = 1.0; // row0 pos1 -> 200 + target[2 * VOCAB + 777] = 1.0; // row0 bonus -> 777 + let r1 = 3 * VOCAB; + target[r1 + 555] = 1.0; // row1 pos0: all mass on 555, draft is 100 + target[r1 + VOCAB + 200] = 1.0; // never reached + target[r1 + 2 * VOCAB + 300] = 1.0; // never reached + let mut target_d: CudaSlice = + ctx.stream.clone_htod(&target).expect("target probs H2D"); + let drafts: Vec = vec![100, 200, /* row1 */ 100, 200]; + let drafts_d = ctx.stream.clone_htod(&drafts).expect("draft ids H2D"); + let mut draft_probs: CudaSlice = + ctx.stream.alloc_zeros(2 * k * VOCAB).expect("draft probs"); + let mut out: CudaSlice = ctx.stream.alloc_zeros(2 * (k + 1)).expect("out"); + + let mut scratch = SpecAcceptScratch::new(&ctx, 2).expect("scratch"); + let counts = gpu_spec_accept_into( + &ctx, + &mut draft_probs, + &drafts_d, + &mut target_d, + &mut out, + 2, + k, + VOCAB, + /*onehot_draft=*/ true, + 0x5eed, + 0, + &mut scratch, + ) + .expect("spec accept"); + let (accepted, emitted) = (counts.acceptance_telemetry, counts.emitted); + let out_h = ctx.stream.clone_dtoh(&out).expect("D2H"); + ctx.sync().expect("sync"); + + // Row 0: both drafts accepted (p_target = 1), bonus token 777 emitted. + assert_eq!(&out_h[..3], &[100, 200, 777], "row0 {:?}", &out_h[..3]); + assert_eq!(emitted[0], 2, "row0 emitted prefix"); + assert_eq!(accepted[0], 2, "row0 acceptance telemetry"); + // Row 1: draft rejected at pos 0 (p_target = 0); the residual + // relu(target - onehot(100)) is exactly {555: 1.0}, so the resample is + // deterministic; the tail is -1-filled. `emitted` is the commit + // signal; `accepted` may exceed it (it keeps counting hypothetical + // acceptances past the rejection - telemetry, asserted only >= 0). + assert_eq!(&out_h[3..], &[555, -1, -1], "row1 {:?}", &out_h[3..]); + assert_eq!(emitted[1], 0, "row1 emitted prefix"); + // Telemetry keeps counting past the rejection: pos 1's draft (200) is + // hypothetically accepted (target pos1 mass is 1.0 on 200), so the + // acceptance counter reads exactly 1 while the commit signal reads 0. + assert_eq!(accepted[1], 1, "row1 acceptance telemetry"); + } + + #[test] + fn spec_accept_partial_acceptance_matches_the_rate_law() { + // A onehot (greedy) proposer whose draft token holds 0.7 of the target + // mass must be accepted with probability exactly min(1, 0.7) = 0.7. Run + // one big batch of i.i.d. rows (per-row philox decorrelates them) and + // check the empirical acceptance rate converges — the acceptance + // probability is the whole correctness claim of rejection sampling, and + // the corner-point tests above only exercise p ∈ {0, 1}. + let ctx = DeviceContext::new().expect("create CUDA context"); + let vocab = 128usize; + let batch = 512usize; + let k = 1usize; + let (tok_a, tok_b, tok_c) = (10usize, 20usize, 30usize); + + // Every row identical: pos0 target = {A: 0.7, B: 0.3}; the draft is A, + // so on rejection the residual relu(target - onehot(A)) is {B: 1.0}. + // pos1 (the bonus) is deterministic on C, so an accepted row emits + // [A, C] and a rejected row emits [B, -1]. + let mut target = vec![0.0f32; batch * (k + 1) * vocab]; + for row in 0..batch { + let base = row * (k + 1) * vocab; + target[base + tok_a] = 0.7; + target[base + tok_b] = 0.3; + target[base + vocab + tok_c] = 1.0; + } + let mut target_d: CudaSlice = + ctx.stream.clone_htod(&target).expect("target probs H2D"); + let drafts: Vec = vec![tok_a as i32; batch * k]; + let drafts_d = ctx.stream.clone_htod(&drafts).expect("draft ids H2D"); + let mut draft_probs: CudaSlice = ctx + .stream + .alloc_zeros(batch * k * vocab) + .expect("draft probs"); + let mut out: CudaSlice = ctx.stream.alloc_zeros(batch * (k + 1)).expect("out"); + + let mut scratch = SpecAcceptScratch::new(&ctx, batch).expect("scratch"); + let counts = gpu_spec_accept_into( + &ctx, + &mut draft_probs, + &drafts_d, + &mut target_d, + &mut out, + batch, + k, + vocab, + /*onehot_draft=*/ true, + 0xC0FFEE, + 0, + &mut scratch, + ) + .expect("spec accept"); + let emitted = counts.emitted; + let out_h = ctx.stream.clone_dtoh(&out).expect("D2H"); + ctx.sync().expect("sync"); + + let mut accept_count = 0usize; + for row in 0..batch { + let o = &out_h[row * (k + 1)..row * (k + 1) + 2]; + if emitted[row] == 1 { + assert_eq!(o, &[tok_a as i32, tok_c as i32], "accepted row {row}"); + accept_count += 1; + } else { + assert_eq!(emitted[row], 0, "row {row} emitted"); + assert_eq!(o, &[tok_b as i32, -1], "rejected row {row}"); + } + } + let rate = accept_count as f64 / batch as f64; + // Bernoulli(0.7) over 512 draws: std ≈ 0.020, so 0.7 ± 0.07 is > 3σ. + assert!( + (0.63..=0.77).contains(&rate), + "acceptance rate {rate} outside [0.63, 0.77] (expected 0.70)" + ); + } + + #[test] + fn spec_accept_explicit_draft_probs_corner_points() { + // onehot_draft = false: the caller supplies the true proposal + // distribution. Row 0: q puts 0.5 on the draft while the target puts + // 1.0 there -> accept probability min(1, 1.0/0.5) = 1, certain accept; + // the bonus position is deterministic on 777. Row 1: q puts 1.0 on the + // draft while the target puts 0 there -> certain rejection; the + // residual relu(target - draft) is exactly {555: 1.0}. + let ctx = DeviceContext::new().expect("create CUDA context"); + let k = 1usize; + let mut target = vec![0.0f32; 2 * (k + 1) * VOCAB]; + target[100] = 1.0; // row0 pos0 -> all mass on the draft + target[VOCAB + 777] = 1.0; // row0 bonus -> 777 + let r1 = 2 * VOCAB; + target[r1 + 555] = 1.0; // row1 pos0: no mass on the draft + target[r1 + VOCAB + 300] = 1.0; // never reached + let mut target_d: CudaSlice = + ctx.stream.clone_htod(&target).expect("target probs H2D"); + let mut draft = vec![0.0f32; 2 * k * VOCAB]; + draft[100] = 0.5; // row0: q(draft) = 0.5 + draft[200] = 0.5; + draft[VOCAB + 100] = 1.0; // row1: q(draft) = 1.0 + let mut draft_d: CudaSlice = ctx.stream.clone_htod(&draft).expect("draft probs H2D"); + let drafts_d = ctx.stream.clone_htod(&[100i32, 100]).expect("draft ids"); + let mut out: CudaSlice = ctx.stream.alloc_zeros(2 * (k + 1)).expect("out"); + let mut scratch = SpecAcceptScratch::new(&ctx, 2).expect("scratch"); + + let counts = gpu_spec_accept_into( + &ctx, + &mut draft_d, + &drafts_d, + &mut target_d, + &mut out, + 2, + k, + VOCAB, + /*onehot_draft=*/ false, + 0xFACE, + 0, + &mut scratch, + ) + .expect("spec accept"); + let out_h = ctx.stream.clone_dtoh(&out).expect("D2H"); + ctx.sync().expect("sync"); + + assert_eq!(&out_h[..2], &[100, 777], "row0 {:?}", &out_h[..2]); + assert_eq!(counts.emitted[0], 1, "row0 emitted prefix"); + assert_eq!(&out_h[2..], &[555, -1], "row1 {:?}", &out_h[2..]); + assert_eq!(counts.emitted[1], 0, "row1 emitted prefix"); + } +}