feat: Added tvm binding for sampling kernel (#958)

Author: annanyapr

flashinfer/jit/__init__.py

@@ -44,6 +44,7 @@
     gen_customize_single_prefill_module as gen_customize_single_prefill_module,
 )
 from .attention import gen_pod_module as gen_pod_module
+from .attention import gen_sampling_tvm_binding as gen_sampling_tvm_binding
 from .attention import gen_single_decode_module as gen_single_decode_module
 from .attention import gen_single_prefill_module as gen_single_prefill_module
 from .attention import get_batch_decode_mla_uri as get_batch_decode_mla_uri

flashinfer/jit/attention/__init__.py

@@ -48,3 +48,4 @@
 from .tvm import (
     gen_customize_batch_prefill_tvm_binding as gen_customize_batch_prefill_tvm_binding,
 )
+from .tvm import gen_sampling_tvm_binding as gen_sampling_tvm_binding

flashinfer/jit/attention/tvm.py

@@ -35,6 +35,22 @@
 from .utils import generate_additional_params
 
 
+def gen_sampling_tvm_binding(uri: str):
+    gen_directory = FLASHINFER_GEN_SRC_DIR / uri
+    os.makedirs(gen_directory, exist_ok=True)
+
+    source_paths = []
+    for filename in ["sampling.cu", "sampling_jit_tvm_binding.cu"]:
+        src_path = FLASHINFER_TVM_BINDING_DIR / filename
+        dest_path = gen_directory / filename
+        source_paths.append(dest_path)
+        with open(src_path, "r") as f:
+            source = f.read()
+        write_if_different(dest_path, source)
+
+    return uri, source_paths
+
+
 def gen_customize_batch_prefill_tvm_binding(
     backend: str,
     uri: str,

tvm_binding/sampling.cu

@@ -0,0 +1,32 @@
+#include <flashinfer/attention/hopper/attention_updater.cuh>
+#include <flashinfer/attention/hopper/variant_helper.cuh>
+#include <flashinfer/cutlass_utils.cuh>
+#include <flashinfer/layout.cuh>
+#include <flashinfer/math.cuh>
+#include <flashinfer/sampling.cuh>
+
+#include "tvm_binding_utils.h"
+
+using namespace flashinfer;
+
+// TODO: change the philox seeds and offsets to DLTensor once the underlying API for sampling
+// changes to support multiple seeds
+void SamplingFromProbs(DLTensor* probs, DLTensor* output, DLTensor* maybe_indices,
+                       bool deterministic, uint64_t philox_seed, uint64_t philox_offset,
+                       int64_t cuda_stream) {
+  CHECK(probs->ndim == 2) << "Probs should have 2 dimensions";
+  unsigned int batch_size = output->shape[0];
+  unsigned int vocab_size = probs->shape[1];
+
+  cudaStream_t stream = reinterpret_cast<cudaStream_t>(cuda_stream);
+  float* probs_cast = static_cast<float*>(probs->data) + probs->byte_offset;
+  int* output_cast = static_cast<int*>(output->data) + output->byte_offset;
+  int* maybe_indices_cast =
+      maybe_indices ? static_cast<int*>(maybe_indices->data) + maybe_indices->byte_offset : nullptr;
+
+  cudaError_t status =
+      sampling::SamplingFromProb(probs_cast, output_cast, maybe_indices_cast, batch_size,
+                                 vocab_size, deterministic, philox_seed, philox_offset, stream);
+  CHECK(status == cudaSuccess) << "SamplingFromProbs failed with error "
+                               << cudaGetErrorString(status);
+}

tvm_binding/sampling_jit_tvm_binding.cu

@@ -0,0 +1,7 @@
+#include "tvm_binding_utils.h"
+
+void SamplingFromProbs(DLTensor* probs, DLTensor* output, DLTensor* maybe_indices,
+                       bool deterministic, uint64_t philox_seed, uint64_t philox_offset,
+                       int64_t cuda_stream);
+
+TVM_DLL_EXPORT_TYPED_FUNC(sampling_from_probs, SamplingFromProbs);