Implementation:Sgl project Sglang Sampling Kernels

Knowledge Sources	Sgl_project_Sglang
Domains	GPU Kernels, Sampling, LLM Inference
Last Updated	2026-02-10 00:00 GMT

Overview

Python interface for GPU-accelerated sampling kernels supporting top-k, top-p (nucleus), min-p, and combined sampling strategies adapted from the FlashInfer project.

Description

sampling.py provides fused GPU sampling kernels that run entirely on the GPU, avoiding the latency of CPU-based sampling and enabling efficient batched generation with diverse per-request sampling parameters. The implementations use GPU-based rejection sampling without explicit sorting.

Probability Renormalization:

• top_k_renorm_probs -- Renormalizes probabilities by keeping only the top-k values, zeroing the rest, and renormalizing. The threshold can be a scalar (same for all requests) or a per-request tensor of shape (batch_size,). Returns renormalized probabilities of shape (batch_size, num_classes).
• top_p_renorm_probs -- Renormalizes probabilities by top-p (nucleus) thresholding: masks out probabilities below the threshold where cumulative sum of probs[probs >= threshold] equals top_p. Supports scalar or per-request tensor thresholds.

Direct Sampling from Probabilities:

• top_p_sampling_from_probs -- Fused top-p (nucleus) sampling from probabilities using GPU rejection sampling in a single CUDA kernel. Supports optional indices tensor for reusing the same probability distribution for multiple outputs, deterministic mode, custom random generators, and NaN checking.
• top_k_top_p_sampling_from_probs -- Combined top-k and top-p sampling with configurable application order: "top_k_first" applies top-k renormalization then top-p sampling, while "joint" applies both filters simultaneously in each rejection round. Returns int32 sampled categories.
• min_p_sampling_from_probs -- Min-p sampling that keeps tokens with probability above min_p * max_prob, providing an adaptive threshold that scales with confidence.

Logit-Level Operations:

• top_k_mask_logits -- Masks logits by keeping only the top-k values and setting the rest to negative infinity. This is the logit-space equivalent of top_k_renorm_probs (applying softmax after masking yields the same result as renormalization).
• top_k_top_p_sampling_from_logits -- End-to-end sampling from pre-softmax logits: applies top-k masking, softmax, then top-p sampling. Supports the same filter_apply_order options as the probability-based variant.

Parameter Flexibility: All threshold parameters (top_k, top_p, min_p) accept either:

• A scalar value applied uniformly to all requests in the batch
• A tensor of shape (batch_size,) for per-request thresholds

This is handled internally by _to_tensor_scalar_tuple which splits the parameter into an optional tensor array and a scalar default value.

Implementation Details:

• All functions expect float32 probability inputs; tensors are cast internally
• Output samples are int32 category indices
• The internal _*_internal functions handle the actual kernel dispatch
• Legacy aliases top_k_renorm_prob and top_p_renorm_prob (singular) are provided

Usage

Use these functions during the token generation phase of LLM inference. They are called after logits have been computed, to select the next token according to the configured sampling strategy. The batched interface allows different sampling parameters per request in a single kernel launch.

Code Reference

Source Location

• Repository: Sgl_project_Sglang
• File: sgl-kernel/python/sgl_kernel/sampling.py
• Lines: 1-543

Signature

def top_k_renorm_probs(
    probs: torch.Tensor,
    top_k: Union[torch.Tensor, int],
) -> torch.Tensor:

def top_p_renorm_probs(
    probs: torch.Tensor,
    top_p: Union[torch.Tensor, float],
) -> torch.Tensor:

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:

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:

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:

def top_k_mask_logits(
    logits: torch.Tensor,
    top_k: Union[torch.Tensor, int],
) -> torch.Tensor:

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:

Import

from sgl_kernel import top_k_renorm_probs, top_p_sampling_from_probs
from sgl_kernel import top_k_top_p_sampling_from_probs, min_p_sampling_from_probs
from sgl_kernel import top_k_mask_logits, top_k_top_p_sampling_from_logits

I/O Contract

Inputs

Name	Type	Required	Description
probs	torch.Tensor (float32)	Yes	Probability distribution: (batch_size, num_classes)
logits	torch.Tensor (float32)	Yes (logit funcs)	Pre-softmax logits: (batch_size, num_classes)
top_k	Union[torch.Tensor, int]	Yes (top-k funcs)	Top-k threshold, scalar or (batch_size,) tensor
top_p	Union[torch.Tensor, float]	Yes (top-p funcs)	Top-p threshold in (0, 1), scalar or (batch_size,) tensor
min_p	Union[torch.Tensor, float]	Yes (min-p func)	Min-p threshold, scalar or (batch_size,) tensor
indices	Optional[torch.Tensor]	No	Maps output index to probability row: (batch_size,)
filter_apply_order	str	No	"top_k_first" (default) or "joint" for combined filtering
deterministic	bool	No	Use deterministic kernel, default True
generator	Optional[torch.Generator]	No	Custom random number generator
check_nan	bool	No	Check for NaN in input probabilities, default False

Outputs

Name	Type	Description
renorm_probs	torch.Tensor (float32)	Renormalized probabilities: (batch_size, num_classes)
samples	torch.Tensor (int32)	Sampled token indices: (batch_size,)
masked_logits	torch.Tensor (float32)	Masked logits with -inf for filtered tokens: (batch_size, num_classes)

Usage Examples

import torch
from sgl_kernel import (
    top_k_renorm_probs,
    top_p_sampling_from_probs,
    top_k_top_p_sampling_from_probs,
    min_p_sampling_from_probs,
    top_k_top_p_sampling_from_logits,
)

batch_size, vocab_size = 32, 32000

# Top-k + top-p sampling from logits (most common usage)
logits = torch.randn(batch_size, vocab_size, device="cuda", dtype=torch.float32)
samples = top_k_top_p_sampling_from_logits(
    logits, top_k=50, top_p=0.9, deterministic=True
)  # shape: (32,), int32

# Per-request top-p sampling from probabilities
probs = torch.softmax(logits, dim=-1)
per_request_top_p = torch.tensor([0.9, 0.95, 0.8, ...], device="cuda")
samples = top_p_sampling_from_probs(probs, top_p=per_request_top_p)

# Min-p sampling (adaptive threshold)
samples = min_p_sampling_from_probs(probs, min_p=0.05)

# Renormalize probabilities by top-k (useful for two-step sampling)
renormed = top_k_renorm_probs(probs, top_k=50)
# Then sample from renormed probs with any method

# Shared probability distribution for multiple outputs
indices = torch.tensor([0, 0, 1, 1, 2, 2], device="cuda", dtype=torch.int32)
samples = top_p_sampling_from_probs(
    probs[:3],  # only 3 unique distributions
    top_p=0.9,
    indices=indices,  # 6 outputs mapped to 3 distributions
)

Related Pages

• Environment:Sgl_project_Sglang_CUDA_GPU_Runtime