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

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<size_t>(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<size_t>(batch_size) * num_speculative_tokens;
+    cudaError_t err = cudaMemsetAsync(
+        draft_probs, 0, total * vocab_size * sizeof(float), stream);
+    if (err != cudaSuccess) {
+      return static_cast<int>(err);
+    }
+    onehot_rows_kernel<<<static_cast<unsigned>(total), 32, 0, stream>>>(
+        draft_probs, draft_token_ids, vocab_size);
+    err = cudaGetLastError();
+    if (err != cudaSuccess) {
+      return static_cast<int>(err);
+    }
+  }
+  cudaError_t err = flashinfer::sampling::ChainSpeculativeSampling<float, int>(
+      draft_probs, const_cast<int*>(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<int>(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<<<gather_grid, GATHER_CAST_BLOCK, 0, stream>>>(
-      logits, row_indices, probs_scratch, vocab_size);
+      logits, row_indices, probs, vocab_size);
   cudaError_t err = cudaGetLastError();
   if (err != cudaSuccess) {
     return static_cast<int>(err);
   }
-
   // In-place: phase 2 of both OnlineSoftmax strategies is an elementwise
   // read-then-write of the same index.
   err = flashinfer::sampling::OnlineSoftmax<float>(
-      probs_scratch, probs_scratch, n_rows, vocab_size, const_cast<float*>(temperature_arr),
+      probs, probs, n_rows, vocab_size, const_cast<float*>(temperature_arr),
       /*temperature_val=*/1.0f, softmax_workspace, softmax_workspace_bytes,
       /*enable_pdl=*/false, stream);
   if (err != cudaSuccess) {
     return static_cast<int>(err);
   }
+  if (do_top_k) {
+    auto* row_states =
+        reinterpret_cast<flashinfer::sampling::RadixRowState*>(topk_row_states_scratch);
+    if (row_states == nullptr) {
+      return static_cast<int>(cudaErrorInvalidValue);
+    }
+    // In-place: the renorm kernels reduce a threshold first, then rewrite
+    // each element from its own index.
+    err = flashinfer::sampling::RadixTopKRenormProbMultiCTA<float, int>(
+        probs, probs, const_cast<int*>(top_k_arr), n_rows,
+        /*top_k_val=*/0, vocab_size, row_states, stream);
+    if (err != cudaSuccess) {
+      return static_cast<int>(err);
+    }
+  }
+  if (do_top_p) {
+    err = flashinfer::sampling::TopPRenormProb<float>(
+        probs, probs, const_cast<float*>(top_p_arr), n_rows,
+        /*top_p_val=*/0.0f, vocab_size, stream);
+  }
+  return static_cast<int>(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<bool*>(valid_scratch);
   if (min_p_arr != nullptr) {
-    if (has_top_k_filter) {
-      auto* row_states =
-          reinterpret_cast<flashinfer::sampling::RadixRowState*>(topk_row_states_scratch);
-      if (row_states == nullptr) {
-        return static_cast<int>(cudaErrorInvalidValue);
-      }
-      // In-place: the renorm kernels reduce a threshold first, then rewrite
-      // each element from its own index.
-      err = flashinfer::sampling::RadixTopKRenormProbMultiCTA<float, int>(
-          probs_scratch, probs_scratch, const_cast<int*>(top_k_arr), n_rows,
-          /*top_k_val=*/0, vocab_size, row_states, stream);
-      if (err != cudaSuccess) {
-        return static_cast<int>(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<float>(
-          probs_scratch, probs_scratch, const_cast<float*>(top_p_arr), n_rows,
-          /*top_p_val=*/0.0f, vocab_size, stream);
-      if (err != cudaSuccess) {
-        return static_cast<int>(err);
-      }
-    }
-    err = flashinfer::sampling::MinPSamplingFromProb<float, int>(
+    cudaError_t err = flashinfer::sampling::MinPSamplingFromProb<float, int>(
         probs_scratch, const_cast<float*>(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<int>(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<float, int>(
         probs_scratch, const_cast<int*>(top_k_arr), const_cast<float*>(top_p_arr), output, valid,

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,

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