Implementation:Mlc ai Mlc llm CPU Sampler
| Knowledge Sources | |
|---|---|
| Domains | LLM Serving, Sampling, Speculative Decoding |
| Last Updated | 2026-02-09 19:00 GMT |
Overview
A CPU-based sampler implementation for MLC LLM that performs top-p sampling, probability renormalization, and draft token verification on CPU-resident probability distributions.
Description
The cpu_sampler.cc file provides the CPU implementation of the token sampling interface for the MLC LLM serving engine. It includes three major components:
SampleTopPFromProb is a standalone function that samples a token from a probability distribution using top-p (nucleus) sampling. It handles three cases with distinct optimizations:
- When top_p is 0, it performs an argmax operation with an early exit optimization (stops scanning when the remaining probability mass cannot exceed the current maximum).
- When top_p is 1.0, it performs standard categorical sampling by accumulating probabilities until the uniform random sample is reached.
- For intermediate top_p values, it uses a two-phase approach: first filtering probabilities using a cutoff of top_p/1024 (which by pigeonhole principle yields at most 1024 elements), sorting the filtered results, then sampling from the sorted top-p subset. If the filtered set is insufficient, it falls back to processing the full distribution.
RenormalizeProbByTopP renormalizes a probability distribution in-place according to a top-p threshold. It uses a multi-round partitioning strategy with progressively smaller cutoff values (top_p/256, top_p/8192, 0.0) to efficiently separate the distribution into upper and lower partitions. Only the upper partition is sorted. Values below the boundary probability are zeroed out and the remaining values are renormalized to sum to 1.
CPUSampler is the main class implementing the SamplerObj interface. It provides:
- BatchRenormalizeProbsByTopP: Copies probability distributions from device to CPU, then renormalizes each distribution by its respective top-p value using parallel execution.
- BatchSampleTokensWithProbBeforeTopP and BatchSampleTokensWithProbAfterTopP: Batch sampling with and without prior top-p application.
- BatchVerifyDraftTokensWithProbAfterTopP: Implements speculative decoding verification. For each draft token, it compares the model probability p(x) against the draft probability q(x). If p >= q, the token is accepted. Otherwise, it is accepted with probability p/q. If rejected, a new token is sampled from a corrected distribution max(p - q, 0), normalized.
The ComputeTopProbs helper function computes the top-k highest probability tokens using a selection sort algorithm, with template specialization for k values 1 through 5. This is used for log-probability reporting.
All batch operations use TVM's parallel threading backend for concurrent processing across samples.
Usage
The CPU sampler is created via Sampler::CreateCPUSampler() and is used as the sampling backend when GPU-based sampling is not available or not preferred. It is integrated into the serving engine's decode loop, where it is called after model inference to select the next token for each request.
Code Reference
Source Location
- Repository: Mlc_ai_Mlc_llm
- File: cpp/serve/sampler/cpu_sampler.cc
Signature
// Standalone sampling function
TokenProbPair SampleTopPFromProb(Tensor prob, int unit_offset, int input_prob_offset,
double top_p, double uniform_sample);
// Standalone renormalization function
void RenormalizeProbByTopP(Tensor prob, int unit_offset, double top_p, double eps);
// CPUSampler class (implements SamplerObj)
class CPUSampler : public SamplerObj {
public:
explicit CPUSampler(Optional<EventTraceRecorder> trace_recorder);
Tensor BatchRenormalizeProbsByTopP(Tensor probs_on_device,
const std::vector<int>& sample_indices,
const Array<String>& request_ids,
const Array<GenerationConfig>& generation_cfg) final;
std::vector<SampleResult> BatchSampleTokensWithProbBeforeTopP(
Tensor probs_on_device, const std::vector<int>& sample_indices,
const Array<String>& request_ids, const Array<GenerationConfig>& generation_cfg,
const std::vector<RandomGenerator*>& rngs) final;
std::vector<SampleResult> BatchSampleTokensWithProbAfterTopP(
Tensor probs_on_host, const std::vector<int>& sample_indices,
const Array<String>& request_ids, const Array<GenerationConfig>& generation_cfg,
const std::vector<RandomGenerator*>& rngs) final;
std::pair<std::vector<std::vector<SampleResult>>, std::vector<int>>
BatchVerifyDraftTokensWithProbAfterTopP(
Tensor probs_on_host, const Array<String>& request_ids,
const std::vector<int>& cum_verify_lengths,
const Array<GenerationConfig>& generation_cfg,
const std::vector<RandomGenerator*>& rngs,
const std::vector<std::vector<SampleResult>>& draft_output_tokens,
const std::vector<int64_t>& token_tree_parent_ptr,
Tensor draft_probs_on_device) final;
};
// Factory function
Sampler Sampler::CreateCPUSampler(Optional<EventTraceRecorder> trace_recorder);
Import
#include "sampler.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| probs_on_device | Tensor (n, vocab_size) | Yes | Batch of probability distributions, shape (n, v), dtype float32. |
| sample_indices | vector<int> | Yes | Indices mapping each sample to its probability distribution row. |
| request_ids | Array<String> | Yes | Request IDs for event tracing. |
| generation_cfg | Array<GenerationConfig> | Yes | Per-request generation configs (top_p, temperature, top_logprobs). |
| rngs | vector<RandomGenerator*> | Yes | Per-request random number generators. |
| top_p | double | For sampling | Nucleus sampling threshold in [0, 1]. |
| uniform_sample | double | For sampling | Uniform random number in [0, 1). |
| draft_output_tokens | vector<vector<SampleResult>> | For verification | Draft tokens to verify against model probabilities. |
| draft_probs_on_device | Tensor | For verification | Draft model probability distributions. |
Outputs
| Name | Type | Description |
|---|---|---|
| SampleResult | struct | Contains sampled_token_id (token, probability pair) and top_prob_tokens (top-k tokens with probabilities). |
| BatchRenormalizeProbsByTopP result | Tensor | CPU tensor of renormalized probability distributions. |
| BatchVerifyDraftTokensWithProbAfterTopP result | (vector<vector<SampleResult>>, vector<int>) | Accepted/resampled tokens per sequence, and the index of the last accepted draft token per sequence. |
Usage Examples
// Create a CPU sampler with optional event tracing
Sampler sampler = Sampler::CreateCPUSampler(trace_recorder);
// Batch sample tokens from probability distributions
std::vector<SampleResult> results = sampler->BatchSampleTokensWithProbBeforeTopP(
probs_on_device, sample_indices, request_ids, generation_cfg, rngs);
// Renormalize probabilities by top-p, then sample
Tensor renormed = sampler->BatchRenormalizeProbsByTopP(
probs_on_device, sample_indices, request_ids, generation_cfg);
results = sampler->BatchSampleTokensWithProbAfterTopP(
renormed, sample_indices, request_ids, generation_cfg, rngs);
// Verify draft tokens for speculative decoding
auto [verified_results, last_accepted] = sampler->BatchVerifyDraftTokensWithProbAfterTopP(
probs_on_host, request_ids, cum_verify_lengths, generation_cfg, rngs,
draft_output_tokens, token_tree_parent_ptr, draft_probs_on_device);