use flashinfer.sampling (#18696)

python/sglang/srt/layers/sampler.py

@@ -19,13 +19,14 @@ from sglang.srt.server_args import get_global_server_args
 from sglang.srt.utils.common import crash_on_warnings, get_bool_env_var, is_cuda, is_npu
 
 if is_cuda():
-    from sgl_kernel import (
+    from flashinfer.sampling import (
         min_p_sampling_from_probs,
-        top_k_renorm_prob,
         top_k_top_p_sampling_from_probs,
+    )
+    from sgl_kernel import (
+        top_k_renorm_prob,
         top_p_renorm_prob,
     )
-
 if is_npu():
     import torch_npu
 

sgl-kernel/CMakeLists.txt

@@ -329,7 +329,6 @@ set(SOURCES
 
     "${repo-flashinfer_SOURCE_DIR}/csrc/norm.cu"
     "${repo-flashinfer_SOURCE_DIR}/csrc/renorm.cu"
-    "${repo-flashinfer_SOURCE_DIR}/csrc/sampling.cu"
 
     "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_causal_sm80.cu"
     "${repo-flash-attention_SOURCE_DIR}/csrc/flash_attn/src/flash_fwd_sparse_hdim128_bf16_sm80.cu"

sgl-kernel/benchmark/bench_top_k_top_p_sampling.py

@@ -1,6 +1,7 @@
 import itertools
 import os
 
+import flashinfer.sampling
 import sgl_kernel
 import torch
 import triton
@@ -69,7 +70,7 @@ def calculate_diff(batch_size, vocab_size, p):
     torch_samples = torch_top_k_top_p_joint_sampling_from_probs(
         normalized_prob, top_k_tensor, top_p_tensor
     )
-    sglang_samples = sgl_kernel.top_k_top_p_sampling_from_probs(
+    sglang_samples = flashinfer.sampling.top_k_top_p_sampling_from_probs(
         normalized_prob, top_k_tensor, top_p_tensor, filter_apply_order="joint"
     )
 
@@ -120,7 +121,7 @@ def benchmark_sampling(batch_size, vocab_size, p, provider):
             normalized_prob.clone(), top_k_tensor, top_p_tensor
         )
     elif provider == "sglang":
-        fn = lambda: sgl_kernel.top_k_top_p_sampling_from_probs(
+        fn = lambda: flashinfer.sampling.top_k_top_p_sampling_from_probs(
             normalized_prob.clone(),
             top_k_tensor,
             top_p_tensor,

sgl-kernel/csrc/common_extension.cc

@@ -417,27 +417,12 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
       {at::Tag::needs_fixed_stride_order});
   m.impl("bmm_fp8", torch::kCUDA, &bmm_fp8);
 
-  m.def(
-      "min_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_min_p_arr, float "
-      "min_p_val, bool deterministic, Generator? gen) -> ()");
-  m.impl("min_p_sampling_from_probs", torch::kCUDA, &min_p_sampling_from_probs);
-
   m.def("top_k_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
   m.impl("top_k_renorm_probs", torch::kCUDA, &top_k_renorm_probs);
 
   m.def("top_p_renorm_probs(Tensor probs, Tensor! renorm_probs, Tensor? maybe_top_p_arr, float top_p_val) -> ()");
   m.impl("top_p_renorm_probs", torch::kCUDA, &top_p_renorm_probs);
 
-  m.def(
-      "top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? "
-      "maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()");
-  m.impl("top_p_sampling_from_probs", torch::kCUDA, &top_p_sampling_from_probs);
-
-  m.def(
-      "top_k_top_p_sampling_from_probs(Tensor probs, Tensor output, Tensor? maybe_indices, Tensor? maybe_top_k_arr, "
-      "float top_k_val, Tensor? maybe_top_p_arr, float top_p_val, bool deterministic, Generator? gen) -> ()");
-  m.impl("top_k_top_p_sampling_from_probs", torch::kCUDA, &top_k_top_p_sampling_from_probs);
-
   m.def("top_k_mask_logits(Tensor logits, Tensor mask_logits, Tensor? maybe_top_k_arr, int top_k_val) -> ()");
   m.impl("top_k_mask_logits", torch::kCUDA, &top_k_mask_logits);
 

sgl-kernel/include/sgl_kernel_ops.h

@@ -683,41 +683,12 @@ void store_kv_cache(at::Tensor k_cache, at::Tensor v_cache, at::Tensor out_loc,
 /*
  * From FlashInfer
  */
-void min_p_sampling_from_probs(
-    at::Tensor probs,
-    at::Tensor output,
-    std::optional<at::Tensor> maybe_indices,
-    std::optional<at::Tensor> maybe_min_p_arr,
-    double min_p_val,
-    bool deterministic,
-    std::optional<at::Generator> gen);
-
 void top_k_renorm_probs(
     at::Tensor probs, at::Tensor renorm_probs, std::optional<at::Tensor> maybe_top_k_arr, int64_t top_k_val);
 
 void top_p_renorm_probs(
     at::Tensor probs, at::Tensor renorm_probs, std::optional<at::Tensor> maybe_top_p_arr, double top_p_val);
 
-void top_k_top_p_sampling_from_probs(
-    at::Tensor probs,
-    at::Tensor output,
-    std::optional<at::Tensor> maybe_indices,
-    std::optional<at::Tensor> maybe_top_k_arr,
-    double top_k_val,
-    std::optional<at::Tensor> maybe_top_p_arr,
-    double top_p_val,
-    bool deterministic,
-    std::optional<at::Generator> gen);
-
-void top_p_sampling_from_probs(
-    at::Tensor probs,
-    at::Tensor output,
-    std::optional<at::Tensor> maybe_indices,
-    std::optional<at::Tensor> maybe_top_p_arr,
-    double top_p_val,
-    bool deterministic,
-    std::optional<at::Generator> gen);
-
 void top_k_mask_logits(
     at::Tensor logits, at::Tensor mask_logits, std::optional<at::Tensor> maybe_top_k_arr, int64_t top_k_val);
 

sgl-kernel/python/sgl_kernel/__init__.py

@@ -101,13 +101,9 @@ from sgl_kernel.quantization import (
     ggml_mul_mat_vec_a8,
 )
 from sgl_kernel.sampling import (
-    min_p_sampling_from_probs,
     top_k_mask_logits,
     top_k_renorm_prob,
-    top_k_top_p_sampling_from_logits,
-    top_k_top_p_sampling_from_probs,
     top_p_renorm_prob,
-    top_p_sampling_from_probs,
 )
 from sgl_kernel.speculative import (
     build_tree_kernel_efficient,

sgl-kernel/python/sgl_kernel/sampling.py

@@ -102,289 +102,6 @@ def top_p_renorm_probs(
 top_p_renorm_prob = top_p_renorm_probs
 
 
-def _top_p_sampling_from_probs_internal(
-    probs: torch.Tensor,
-    indices: Optional[torch.Tensor],
-    maybe_top_p_arr: Optional[torch.Tensor],
-    top_p_val: float,
-    deterministic: bool,
-    generator: Optional[torch.Generator],
-) -> torch.Tensor:
-    with probs.device as device:
-        probs = probs.float()
-        maybe_top_p_arr = (
-            maybe_top_p_arr.float() if maybe_top_p_arr is not None else None
-        )
-        samples = torch.empty(probs.size(0), dtype=torch.int32, device=device)
-        torch.ops.sgl_kernel.top_p_sampling_from_probs.default(
-            probs,
-            samples,
-            indices,
-            maybe_top_p_arr,
-            top_p_val,
-            deterministic,
-            generator,
-        )
-        return samples
-
-
-def top_p_sampling_from_probs(
-    probs: torch.Tensor,
-    top_p: Union[torch.Tensor, float],
-    indices: Optional[torch.Tensor] = None,
-    deterministic: bool = True,
-    generator: Optional[torch.Generator] = None,
-    check_nan: bool = False,
-) -> torch.Tensor:
-    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
-    Fused GPU kernel for top-p sampling (nucleus sampling) from probabilities,
-    this operator implements GPU-based rejection sampling without explicit sorting.
-    Check the `blog post <https://flashinfer.ai/2025/03/10/sampling.html>`_ for more details.
-
-    The multiple rounds of rejection sampling are implemented in a single CUDA kernel,
-    which is more efficient than the naive implementation that launches a series of kernels.
-
-    Parameters
-    ----------
-    probs: torch.Tensor
-        Probabilities for sampling. When indices is not provided, shape should be ``(batch_size, num_classes)``
-        and the i-th output will be sampled from the i-th row of probabilities. When indices is provided,
-        shape should be ``(unique_batch_size, num_classes)`` where unique_batch_size is the number of unique
-        probability distributions.
-    top_p: Union[torch.Tensor, float]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the threshold for top-p sampling.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-    indices: Optional[torch.Tensor]
-        Optional indices tensor of shape ``(batch_size,)`` that maps each output to a row in probs.
-        For example, if indices[i] = j, then the i-th output will be sampled from probs[j].
-        This allows reusing the same probability distribution for multiple outputs.
-        If indices is not provided, the i-th output will be sampled from the i-th row of probs.
-    deterministic: bool
-        Whether to use deterministic kernel implementation, default is ``True``.
-    generator: Optional[torch.Generator]
-        A random number generator for the operation.
-    check_nan: bool
-        Whether to check nan in :attr:`probs`, default is ``False``.
-
-    Returns
-    -------
-    samples: torch.Tensor
-        Sampled categories, shape ``(batch_size,)``.
-
-    Note
-    ----
-    This function expects float32 inputs, and the output is int32.
-
-    """
-    if check_nan:
-        if torch.any(torch.isnan(probs)):
-            raise ValueError("Input probs contains NaN.")
-    return _top_p_sampling_from_probs_internal(
-        probs, indices, *_to_tensor_scalar_tuple(top_p), deterministic, generator
-    )
-
-
-def _top_k_top_p_sampling_from_probs_internal(
-    probs: torch.Tensor,
-    indices: Optional[torch.Tensor],
-    maybe_top_k_arr: Optional[torch.Tensor],
-    top_k_val: int,
-    maybe_top_p_arr: Optional[torch.Tensor],
-    top_p_val: float,
-    deterministic: bool,
-    generator: Optional[torch.Generator],
-) -> torch.Tensor:
-    with probs.device as device:
-        probs = probs.float()
-        maybe_top_k_arr = maybe_top_k_arr.int() if maybe_top_k_arr is not None else None
-        maybe_top_p_arr = (
-            maybe_top_p_arr.float() if maybe_top_p_arr is not None else None
-        )
-        samples = torch.empty(probs.size(0), dtype=torch.int32, device=device)
-        torch.ops.sgl_kernel.top_k_top_p_sampling_from_probs.default(
-            probs,
-            samples,
-            indices,
-            maybe_top_k_arr,
-            top_k_val,
-            maybe_top_p_arr,
-            top_p_val,
-            deterministic,
-            generator,
-        )
-        return samples
-
-
-def top_k_top_p_sampling_from_probs(
-    probs: torch.Tensor,
-    top_k: Union[torch.Tensor, int],
-    top_p: Union[torch.Tensor, float],
-    indices: Optional[torch.Tensor] = None,
-    filter_apply_order: str = "top_k_first",
-    deterministic: bool = True,
-    generator: Optional[torch.Generator] = None,
-    check_nan: bool = False,
-) -> torch.Tensor:
-    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
-    Fused GPU kernel for top-k and top-p sampling from probabilities,
-
-    this operator implements GPU-based rejection sampling without explicit sorting.
-    Check the `blog post <https://flashinfer.ai/2025/03/10/sampling.html>`_ for more details.
-
-    The multiple rounds of rejection sampling are implemented in a single CUDA kernel,
-    which is more efficient than the naive implementation that launches a series of kernels.
-
-    Parameters
-    ----------
-    probs: torch.Tensor
-        Probabilities for sampling. When indices is not provided, shape should be ``(batch_size, num_classes)``
-        and the i-th output will be sampled from the i-th row of probabilities. When indices is provided,
-        shape should be ``(unique_batch_size, num_classes)`` where unique_batch_size is the number of unique
-        probability distributions.
-    top_k: Union[torch.Tensor, int]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the threshold for top-k sampling.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-    top_p: Union[torch.Tensor, float]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the threshold for top-p sampling.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-    indices: Optional[torch.Tensor]
-        Optional indices tensor of shape ``(batch_size,)`` that maps each output to a row in probs.
-        For example, if indices[i] = j, then the i-th output will be sampled from probs[j].
-        This allows reusing the same probability distribution for multiple outputs.
-        If indices is not provided, the i-th output will be sampled from the i-th row of probs.
-    filter_apply_order: str
-        The order of applying top-k and top-p sampling, should be either ``"top_k_first"`` or ``"joint"``.
-        If ``"top_k_first"``, we first apply top-k filter, then apply top-p sampling on the top-k results.
-        If ``"joint"``, we apply top-k and top-p filter simultaneously in each round. Default is ``"top_k_first"``.
-    deterministic: bool
-        Whether to use deterministic kernel implementation, default is ``True``.
-    generator: Optional[torch.Generator]
-        A random number generator for the operation.
-    check_nan: bool
-        Whether to check nan in :attr:`probs`, default is ``False``.
-
-    Returns
-    -------
-    samples: torch.Tensor
-        Sampled categories, shape ``(batch_size,)``.
-
-    Note
-    ----
-    This function expects float32 inputs, and the output is int32.
-
-    """
-    if filter_apply_order == "top_k_first":
-        renorm_probs = top_k_renorm_probs(probs, top_k)
-        return top_p_sampling_from_probs(
-            renorm_probs,
-            top_p,
-            indices,
-            deterministic,
-            check_nan=check_nan,
-            generator=generator,
-        )
-    elif filter_apply_order == "joint":
-        if check_nan:
-            if torch.any(torch.isnan(probs)):
-                raise ValueError("Input probs contains NaN.")
-        return _top_k_top_p_sampling_from_probs_internal(
-            probs,
-            indices,
-            *_to_tensor_scalar_tuple(top_k),
-            *_to_tensor_scalar_tuple(top_p),
-            deterministic,
-            generator,
-        )
-    else:
-        raise ValueError(f"Invalid filter_apply_order: {filter_apply_order}")
-
-
-def _min_p_sampling_from_probs_internal(
-    probs: torch.Tensor,
-    indices: Optional[torch.Tensor],
-    maybe_min_p_arr: Optional[torch.Tensor],
-    min_p_val: float,
-    deterministic: bool,
-    generator: Optional[torch.Generator],
-) -> torch.Tensor:
-    with probs.device as device:
-        probs = probs.float()
-        maybe_min_p_arr = (
-            maybe_min_p_arr.float() if maybe_min_p_arr is not None else None
-        )
-        samples = torch.empty(probs.size(0), dtype=torch.int32, device=device)
-        torch.ops.sgl_kernel.min_p_sampling_from_probs.default(
-            probs,
-            samples,
-            indices,
-            maybe_min_p_arr,
-            min_p_val,
-            deterministic,
-            generator,
-        )
-        return samples
-
-
-def min_p_sampling_from_probs(
-    probs: torch.Tensor,
-    min_p: Union[torch.Tensor, float],
-    indices: Optional[torch.Tensor] = None,
-    deterministic: bool = True,
-    generator: Optional[torch.Generator] = None,
-    check_nan: bool = False,
-) -> torch.Tensor:
-    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
-    Fused GPU kernel for `min_p sampling <https://arxiv.org/abs/2407.01082>`_ from probabilities,
-
-    this operator implements GPU-based rejection sampling without explicit sorting.
-    Check the `blog post <https://flashinfer.ai/2025/03/10/sampling.html>`_ for more details.
-
-    The multiple rounds of rejection sampling are implemented in a single CUDA kernel,
-    which is more efficient than the naive implementation that launches a series of kernels.
-
-    Parameters
-    ----------
-    probs: torch.Tensor
-        Probabilities for sampling. When indices is not provided, shape should be ``(batch_size, num_classes)``
-        and the i-th output will be sampled from the i-th row of probabilities. When indices is provided,
-        shape should be ``(unique_batch_size, num_classes)`` where unique_batch_size is the number of unique
-        probability distributions.
-    min_p: Union[torch.Tensor, float]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the threshold for min-p sampling.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-    indices: Optional[torch.Tensor]
-        Optional indices tensor of shape ``(batch_size,)`` that maps each output to a row in probs.
-        For example, if indices[i] = j, then the i-th output will be sampled from probs[j].
-        This allows reusing the same probability distribution for multiple outputs.
-        If indices is not provided, the i-th output will be sampled from the i-th row of probs.
-    deterministic: bool
-        Whether to use deterministic kernel implementation, default is ``True``.
-    generator: Optional[torch.Generator]
-        A random number generator for the operation.
-    check_nan: bool
-        Whether to check nan in :attr:`probs`, default is ``False``.
-
-    Returns
-    -------
-    samples: torch.Tensor
-        Sampled categories, shape ``(batch_size,)``.
-
-    Note
-    ----
-    This function expects float32 inputs, and the output is int32.
-    """
-    if check_nan:
-        if torch.any(torch.isnan(probs)):
-            raise ValueError("Input probs contains NaN.")
-    return _min_p_sampling_from_probs_internal(
-        probs, indices, *_to_tensor_scalar_tuple(min_p), deterministic, generator
-    )
-
-
 def _top_k_mask_logits_internal(
     logits: torch.Tensor,
     maybe_top_k_arr: Optional[torch.Tensor],
@@ -453,91 +170,3 @@ def top_k_mask_logits(
     top_k_renorm_probs
     """
     return _top_k_mask_logits_internal(logits, *_to_tensor_scalar_tuple(top_k))
-
-
-def top_k_top_p_sampling_from_logits(
-    logits: torch.Tensor,
-    top_k: Union[torch.Tensor, int],
-    top_p: Union[torch.Tensor, float],
-    indices: Optional[torch.Tensor] = None,
-    filter_apply_order: str = "top_k_first",
-    deterministic: bool = True,
-    generator: Optional[torch.Generator] = None,
-    check_nan: bool = False,
-) -> torch.Tensor:
-    r"""Adapt from https://github.com/flashinfer-ai/flashinfer/flashinfer/sampling.py
-    Fused GPU kernel for top-k and top-p sampling from probabilities,
-
-    this operator implements GPU-based rejection sampling without explicit sorting.
-    Check the `blog post <https://flashinfer.ai/2025/03/10/sampling.html>`_ for more details.
-
-    The multiple rounds of rejection sampling are implemented in a single CUDA kernel,
-    which is more efficient than the naive implementation that launches a series of kernels.
-
-    Parameters
-    ----------
-    logits: torch.Tensor
-        Pre-softmax logits for sampling. When indices is not provided, shape should be ``(batch_size, num_classes)``
-        and the i-th output will be sampled from the i-th row of logits. When indices is provided,
-        shape should be ``(unique_batch_size, num_classes)`` where unique_batch_size is the number of unique
-        probability distributions.
-    top_k: Union[torch.Tensor, int]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the threshold for top-k sampling.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-    top_p: Union[torch.Tensor, float]
-        Either a scalar or a tensor of shape ``(batch_size,)``, representing the threshold for top-p sampling.
-        If a scalar, the same threshold is used for all requests.
-        If a tensor, each request has its own threshold.
-    indices: Optional[torch.Tensor]
-        Optional indices tensor of shape ``(batch_size,)`` that maps each output to a row in probs.
-        For example, if indices[i] = j, then the i-th output will be sampled from probs[j].
-        This allows reusing the same probability distribution for multiple outputs.
-        If indices is not provided, the i-th output will be sampled from the i-th row of probs.
-    filter_apply_order: str
-        The order of applying top-k and top-p sampling, should be either ``"top_k_first"`` or ``"joint"``.
-        If ``"top_k_first"``, we first apply top-k filter, then apply top-p sampling on the top-k results.
-        If ``"joint"``, we apply top-k and top-p filter simultaneously in each round. Default is ``"top_k_first"``.
-    deterministic: bool
-        Whether to use deterministic kernel implementation, default is ``True``.
-    generator: Optional[torch.Generator]
-        A random number generator for the operation.
-    check_nan: bool
-        Whether to check nan in :attr:`probs`, default is ``False``.
-
-    Returns
-    -------
-    samples: torch.Tensor
-        Sampled categories, shape ``(batch_size,)``.
-
-    Note
-    ----
-    This function expects float32 inputs, and the output is int32.
-
-    """
-    if filter_apply_order == "top_k_first":
-        masked_logits = top_k_mask_logits(logits, top_k)
-        probs = torch.softmax(masked_logits, dim=-1)
-        return top_p_sampling_from_probs(
-            probs,
-            top_p,
-            indices,
-            deterministic,
-            check_nan=check_nan,
-            generator=generator,
-        )
-    elif filter_apply_order == "joint":
-        probs = torch.softmax(logits, dim=-1)
-        if check_nan:
-            if torch.any(torch.isnan(probs)):
-                raise ValueError("Input probs contains NaN.")
-        return _top_k_top_p_sampling_from_probs_internal(
-            probs,
-            indices,
-            *_to_tensor_scalar_tuple(top_k),
-            *_to_tensor_scalar_tuple(top_p),
-            deterministic,
-            generator,
-        )
-    else:
-        raise ValueError(f"Invalid filter_apply_order: {filter_apply_order}")

sgl-kernel/tests/test_sampling.py

@@ -1,5 +1,6 @@
 # Adapted from https://github.com/flashinfer-ai/flashinfer/blob/93e1a2634e22355b0856246b032b285ad1d1da6b/tests/test_sampling.py
 
+import flashinfer.sampling
 import pytest
 import sgl_kernel
 import torch
@@ -16,10 +17,10 @@ def test_top_k_top_p_sampling_from_probs_logits_top_k_first_alignment(
     logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5
     generator_logits = torch.Generator("cuda:0")
     generator_probs = generator_logits.clone_state()
-    samples = sgl_kernel.sampling.top_k_top_p_sampling_from_logits(
+    samples = flashinfer.sampling.top_k_top_p_sampling_from_logits(
         logits, k, p, filter_apply_order="top_k_first", generator=generator_logits
     )
-    samples_ref = sgl_kernel.sampling.top_k_top_p_sampling_from_probs(
+    samples_ref = flashinfer.sampling.top_k_top_p_sampling_from_probs(
         torch.softmax(logits, dim=-1),
         k,
         p,
@@ -40,10 +41,10 @@ def test_top_k_top_p_sampling_from_probs_logits_joint_alignment(
     logits = torch.randn(batch_size, vocab_size, device="cuda:0") * 5
     generator_logits = torch.Generator("cuda:0")
     generator_probs = generator_logits.clone_state()
-    samples = sgl_kernel.sampling.top_k_top_p_sampling_from_logits(
+    samples = flashinfer.sampling.top_k_top_p_sampling_from_logits(
         logits, k, p, filter_apply_order="joint", generator=generator_logits
     )
-    samples_ref = sgl_kernel.sampling.top_k_top_p_sampling_from_probs(
+    samples_ref = flashinfer.sampling.top_k_top_p_sampling_from_probs(
         torch.softmax(logits, dim=-1),
         k,
         p,
@@ -83,7 +84,7 @@ def test_top_k_top_p_joint_sampling_from_probs(batch_size, vocab_size, p):
 
     num_trails = 1000
     for _ in range(num_trails):
-        samples = sgl_kernel.top_k_top_p_sampling_from_probs(
+        samples = flashinfer.sampling.top_k_top_p_sampling_from_probs(
             normalized_prob,
             top_k_tensor,
             top_p_tensor,
@@ -167,7 +168,7 @@ def test_min_p_sampling(batch_size, vocab_size, p):
 
     num_trails = 1000
     for _ in range(num_trails):
-        samples = sgl_kernel.min_p_sampling_from_probs(
+        samples = flashinfer.sampling.min_p_sampling_from_probs(
             normalized_prob,
             min_p_tensor,
         )