Implementation:InternLM Lmdeploy SamplingKernels

Knowledge Sources	InternLM_Lmdeploy
Domains	GPU_Kernels, Sampling
Last Updated	2026-02-07 15:00 GMT

Overview

Main sampling kernel interface that draws token samples from filtered logit distributions using cuRAND random states.

Description

This header defines the SamplingParams struct and the invokeSampling() function template. SamplingParams encapsulates all state needed for the final token sampling step: the logit buffer (after top-k/top-p filtering), stride between batch elements, sorted indices, kept-token counts, cuRAND random states, batch size, and output buffers for sampled token IDs, sequence lengths, sampled log probabilities, sampled indexes, and sampled counts. invokeSampling() performs the actual sampling by drawing from the filtered distribution using the per-sequence cuRAND state.

Usage

Use this kernel as the final sampling step after top-k and/or top-p filtering has narrowed the candidate token set. It draws one token per sequence from the remaining probability distribution.

Code Reference

Source Location

• Repository: InternLM_Lmdeploy
• File: src/turbomind/kernels/sampling_kernels.h

Signature

struct SamplingParams {
    void*          logits;
    int            stride;
    int*           indices;
    int*           kept;
    curandState_t* curandstate;
    size_t         batch_size;
    int*           output_ids;
    int*           sequence_length;
    void*          sampled_logprobs;
    int*           sampled_indexes;
    int*           sampled_nums;
};

template<typename T>
void invokeSampling(SamplingParams& params, cudaStream_t stream);

Import

#include "src/turbomind/kernels/sampling_kernels.h"

I/O Contract

Inputs

Name	Type	Required	Description
logits	void*	Yes	Filtered logit/probability buffer
stride	int	Yes	Stride between consecutive batch rows in the logit buffer
indices	int*	Yes	Sorted token indices from filtering stage
kept	int*	Yes	Number of kept candidates per sequence
curandstate	curandState_t*	Yes	Per-sequence cuRAND random state
batch_size	size_t	Yes	Number of sequences in the batch

Outputs

Name	Type	Description
output_ids	int*	Sampled token ID per sequence
sequence_length	int*	Updated sequence lengths
sampled_logprobs	void*	Log probability of the sampled token
sampled_indexes	int*	Index of the sampled token in sorted order
sampled_nums	int*	Number of tokens sampled per sequence

Usage Examples

using namespace turbomind;

SamplingParams params;
params.logits = filtered_logits;
params.stride = vocab_size_padded;
params.indices = sorted_indices;
params.kept = kept_counts;
params.curandstate = rand_states;
params.batch_size = batch_size;
params.output_ids = output_ids;
params.sequence_length = seq_lengths;

invokeSampling<half>(params, stream);

Related Pages

• Environment:InternLM_Lmdeploy_CUDA_GPU_Runtime