From 2ad475b4edfed009701737d199b9410540af63d7 Mon Sep 17 00:00:00 2001 From: pansicheng Date: Thu, 26 Feb 2026 10:02:38 +0800 Subject: [PATCH] use flashinfer.sampling (#18696) --- python/sglang/srt/layers/sampler.py | 7 +- sgl-kernel/CMakeLists.txt | 1 - .../benchmark/bench_top_k_top_p_sampling.py | 5 +- sgl-kernel/csrc/common_extension.cc | 15 - sgl-kernel/include/sgl_kernel_ops.h | 29 -- sgl-kernel/python/sgl_kernel/__init__.py | 4 - sgl-kernel/python/sgl_kernel/sampling.py | 371 ------------------ sgl-kernel/tests/test_sampling.py | 13 +- 8 files changed, 14 insertions(+), 431 deletions(-) diff --git a/python/sglang/srt/layers/sampler.py b/python/sglang/srt/layers/sampler.py index f78d83d79..a4c7c7db0 100644 --- a/python/sglang/srt/layers/sampler.py +++ b/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 diff --git a/sgl-kernel/CMakeLists.txt b/sgl-kernel/CMakeLists.txt index 680f41438..395980904 100644 --- a/sgl-kernel/CMakeLists.txt +++ b/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" diff --git a/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py b/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py index 278356c38..d84a34884 100644 --- a/sgl-kernel/benchmark/bench_top_k_top_p_sampling.py +++ b/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, diff --git a/sgl-kernel/csrc/common_extension.cc b/sgl-kernel/csrc/common_extension.cc index 087383f1d..54b58cca5 100644 --- a/sgl-kernel/csrc/common_extension.cc +++ b/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); diff --git a/sgl-kernel/include/sgl_kernel_ops.h b/sgl-kernel/include/sgl_kernel_ops.h index 9072072cf..8de01e017 100644 --- a/sgl-kernel/include/sgl_kernel_ops.h +++ b/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 maybe_indices, - std::optional maybe_min_p_arr, - double min_p_val, - bool deterministic, - std::optional gen); - void top_k_renorm_probs( at::Tensor probs, at::Tensor renorm_probs, std::optional maybe_top_k_arr, int64_t top_k_val); void top_p_renorm_probs( at::Tensor probs, at::Tensor renorm_probs, std::optional maybe_top_p_arr, double top_p_val); -void top_k_top_p_sampling_from_probs( - at::Tensor probs, - at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_k_arr, - double top_k_val, - std::optional maybe_top_p_arr, - double top_p_val, - bool deterministic, - std::optional gen); - -void top_p_sampling_from_probs( - at::Tensor probs, - at::Tensor output, - std::optional maybe_indices, - std::optional maybe_top_p_arr, - double top_p_val, - bool deterministic, - std::optional gen); - void top_k_mask_logits( at::Tensor logits, at::Tensor mask_logits, std::optional maybe_top_k_arr, int64_t top_k_val); diff --git a/sgl-kernel/python/sgl_kernel/__init__.py b/sgl-kernel/python/sgl_kernel/__init__.py index a24713a24..6366de7d5 100644 --- a/sgl-kernel/python/sgl_kernel/__init__.py +++ b/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, diff --git a/sgl-kernel/python/sgl_kernel/sampling.py b/sgl-kernel/python/sgl_kernel/sampling.py index 4ee6f24d3..6f98d4109 100644 --- a/sgl-kernel/python/sgl_kernel/sampling.py +++ b/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 `_ 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 `_ 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 `_ from probabilities, - - this operator implements GPU-based rejection sampling without explicit sorting. - Check the `blog post `_ 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 `_ 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}") diff --git a/sgl-kernel/tests/test_sampling.py b/sgl-kernel/tests/test_sampling.py index dc5734cb7..51b05c7ef 100644 --- a/sgl-kernel/tests/test_sampling.py +++ b/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, )