The sampler

Names: R is the number of selected rows, V is vocab_size, N is DEEPINFER_SAMPLER_FORCE_TOPK. With force top-k off, read N as V.

Worked example (prefill). Two prompts of 41 and 100 tokens are in one prefill step. The model output has 141 rows. selected_token_indices = [40, 140], so R = 2 (csrc/engine/prepare_sample.cpp:197-199). sample_positions = [41, 100], the seqlen_k of each prompt (csrc/engine/prepare_sample.cpp:198). The logits use 2 × V × 4 bytes. With the inferred V = 129280, that is 1,034,240 bytes, about 1 MB.

Worked example (decode). Three requests have seqlen_q = 6, 6, 4. Then R = total_q = 16 and mtp_first_token_positions = [0, 6, 12] (csrc/engine/prepare_sample.cpp:152-153). prepare_verification_metadata makes three tensors of shape [16]:

row_to_req = [0,0,0,0,0,0, 1,1,1,1,1,1, 2,2,2,2]
local_step = [0,1,2,3,4,5, 0,1,2,3,4,5, 0,1,2,3]
is_verify  = [T,T,T,T,T,F, T,T,T,T,T,F, T,T,T,F]   # the last row of each request is a fresh sample

The reference implementation shows the rule (deepinfer/model_executor/layers/tilelang_ops/sampler.py:13-25). expand_sampling_params then copies the per-sequence parameters to the 16 rows (deepinfer/model_executor/layers/tilelang_ops/sampler.py:56-66). The logits use 16 × 129280 × 4 bytes, about 8.3 MB.

Every random number of the sampler comes from a hash of two values: the 64-bit seq_seed of the request and a 64-bit sequence number. The hash is SplitMix64, and the result is a float64 in [0, 1) (deepinfer/model_executor/layers/tilelang_ops/random_sampling.py:58-67). The sequence number is 3 × position + case:

Which position.

• Prefill: sample_position = seqlen_k, the prompt length after this chunk (csrc/engine/prepare_sample.cpp:198).
• Decode: sample_position = seqlen_k - seqlen_q + 1 per request (csrc/engine/prepare_sample.cpp:206). Row j of the request uses sample_position + j × 10000019, so every row has its own stream (deepinfer/model_executor/layers/tilelang_ops/sampler.py:65, deepinfer/model_executor/layers/tilelang_ops/common.py:52).
• Draft model: the seed gets an offset of 1000000007 (tag "mtp"), so drafts and target use different streams (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:15, deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:93). Draft i (1 to 5) uses position sample_position + last_valid_offset + i (deepinfer/model_executor/models/mtp/dflash.py:290-294).

Worked example. Request "1" of the tour Tour: Life of a request has seq_seed = 1 and a 41-token prompt. Its prefill token uses position 41, case 0. In a decode step with input positions 49 to 54, seqlen_k = 55 and seqlen_q = 6, so sample_position = 50. Row 2 then uses position 50 + 2 × 10000019 = 20000088.

No seed given. SamplingParams takes 64 random bits from Python's random module (deepinfer/sampling_params.py:85-91). So two runs differ, unless you pass seq_seed. The README gives the full recipe for equal outputs: DEEPINFER_BATCH_INVARIANT=1 DEEPINFER_ENABLE_MTP_VARLEN=0 and a fixed seq_seed, on NVIDIA only (README.md:28). See Determinism and batch invariance.

Parameter checks. SamplingParams raises ValueError for a negative temperature, for top_p outside (0, 1], and for top_k that is 0, below -1 or above the int32 maximum (deepinfer/sampling_params.py:93-99). A temperature below SAMPLING_EPS = 1e-5 becomes 1e-5, and top-p and top-k are then turned off (deepinfer/sampling_params.py:76-79). So "temperature 0" is very close to greedy, but it is still a seeded draw.

Order. The reference applies top-p first and top-k second. Both use a stable sort, so equal values keep their token order (deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:121-150):

1. Sort the probabilities in descending order.
2. Remove a token if the sum of the probabilities before it is larger than top_p (line 139). So the kept set is the shortest prefix whose sum is larger than top_p, or the full row.
3. On the remaining logits, remove every token with rank ≥ top_k (lines 145-150). top_k = -1 removes nothing.
4. log_softmax normalizes the kept tokens again.

Worked example. After temperature, five tokens have probabilities [0.50, 0.25, 0.15, 0.06, 0.04]. With top_p = 0.8, the sums before each token are [0, 0.50, 0.75, 0.90, 0.96]. Tokens 4 and 5 have sums above 0.8, so they go. With top_k = 2, token 3 also goes. The result is [0.667, 0.333, 0, 0, 0].

The logprobs are of the filtered distribution. All later values use these processed logprobs: the token logprob, the top logprobs and the prompt logprobs. A token that top-p removed has logprob minus infinity.

A full row has about 129280 floats (vocabulary size inferred). Sorting it for top-p is expensive. With force top-k, the sampler first keeps only the N largest logits of each row, with their ids, and does all later work on N columns:

    vocab_size = logits.shape[1]
    if force_topk:
        logits, token_index = topk(logits, force_topk, False, indices_type=torch.int32)  # [total_q, force_topk]
    else:
        token_index = None

    # Align reject sampling inputs when force_topk>0: draft top-k space -> fresh top-k space.
    if force_topk and last_nextn_token_logprobs is not None:
        last_logprob, last_next_n_token_ids, mtp_info_read_rows = remap_draft_topk(
            last_nextn_token_logprobs_index,
            last_nextn_token_logprobs,
            draft_token_ids_pool,
            token_index,
            mtp_info_read_rows,
            vocab_size,
        )
    else:
        last_logprob = last_nextn_token_logprobs
        last_next_n_token_ids = draft_token_ids_pool

1. Line 167: topk(..., sorted=False) returns float32 values [R, N] and int32 token_index [R, N]. The function comes from the vendor library DeepSelect. On Ascend, an import hook replaces it (Hardware platforms: NVIDIA and Ascend).
2. Lines 172-180: the drafts of the previous step live in the top-N space of the draft model. remap_draft_topk moves them into the top-N space of the target row (deepinfer/model_executor/layers/tilelang_ops/sampler.py:165-203). A draft token that is not in the target top-N gets id -1, and -1 forces a reject (deepinfer/model_executor/layers/tilelang_ops/sampler.py:221-222, deepinfer/model_executor/layers/tilelang_ops/fused_sample.py:46-48).
3. After the kernel, map_lookup turns column indices back into vocabulary ids (deepinfer/model_executor/layers/sampler.py:460-467).

What it changes. A token outside the top N can never be sampled, and its logprob is minus infinity. A script that needs full-vocabulary logprobs sets sampler_force_topk=0 for this reason (examples/yoco/yoco_equiv.py:559).

Memory. The draft logprobs live in the MTP info pool, with (max_num_seqs + 1) × nextn rows of N floats, plus N int32 ids with force top-k (deepinfer/model_executor/layers/mtp_info_pool.py:20-23). Example with max_num_seqs = 64 and nextn = 5: 325 rows. With N = V = 129280, the pool needs 325 × 129280 × 4 bytes, about 168 MB. With N = 4096, it needs 325 × 4096 × 8 bytes, about 10.6 MB.

Limits and settings. remap_draft_topk asserts N ≤ 16384 on NVIDIA and N ≤ 4096 on Ascend (deepinfer/model_executor/layers/tilelang_ops/sampler.py:189). The code default is 0, off (deepinfer/environ.py:150). examples/offline_inference.py uses 4096 (examples/offline_inference.py:17), and examples/benchmark.py uses 1024 (examples/benchmark.py:54).

Token logprob. For each sampled token, the sampler looks up its processed logprob with map_lookup (deepinfer/model_executor/layers/sampler.py:494-499).

Top logprobs. _prepare_top_info takes the k largest logprobs of each row, sorted (deepinfer/model_executor/layers/sampler.py:636-648). k is the largest top_logprobs of all requests in the batch (csrc/engine/prepare_sample.cpp:157). So the tensor has the same width for all rows.

sum_smaller_probs. With DEEPINFER_COMPUTE_SUM_SMALLER_PROBS=1, the sampler also returns one more value per row. It is the probability sum of all tokens that are less likely than the sampled token, plus the sampled token itself. Equal probabilities are ordered by column index:

def sum_smaller_probs_ref(
    logprobs: torch.Tensor,
    sample_results: torch.Tensor,
) -> torch.Tensor:
    invalid = sample_results < 0
    sample_results = sample_results.clone()
    sample_results.masked_fill_(invalid, 0)
    pivot = logprobs.gather(1, sample_results.unsqueeze(1))
    mask_less = logprobs < pivot
    col_idx = torch.arange(logprobs.shape[1], device=logprobs.device).unsqueeze(0)
    mask_equal_before = (logprobs == pivot) & (col_idx <= sample_results.unsqueeze(1))
    out = (logprobs.exp() * (mask_less | mask_equal_before)).sum(dim=1)
    out.masked_fill_(invalid, 0.0)
    return out

Prompt logprobs. With SamplingParams(prompt_logprobs=True), a prefill step also returns the logprob of each prompt token, given the tokens before it:

1. C++ selects every row of the chunk except the last row of the last chunk, and stores the next prompt token as the target (csrc/engine/prepare_sample.cpp:12-18, csrc/engine/prepare_sample.cpp:181-189).
2. compute_prompt_logprobs works in chunks of 8192 rows, so the float32 logits never exceed 8192 × V (deepinfer/model_executor/layers/sampler_utils.py:17, deepinfer/model_executor/layers/sampler_utils.py:93-110).
3. Each chunk runs fused_sample without a draw, with the temperature, top-p and top-k of the request (deepinfer/model_executor/layers/sampler_utils.py:69-78).
4. A post function splits the flat result back into one list per request (deepinfer/model_executor/layers/sampler.py:205-210).

Example: a 41-token prompt in one chunk gives 40 prompt logprobs. With 8192 × 129280 × 4 bytes, one chunk of logits needs about 4.2 GB. YOCO truncation turns prompt logprobs off (deepinfer/engine/llm_engine.py:529-537).

The docstring of the class names this feature "golden teacher forcing" (deepinfer/model_executor/layers/sampler.py:92). A request can pass golden_token_ids, a list of output tokens (deepinfer/sampling_params.py:60). The sampler then outputs these tokens instead of its own samples, while it still computes all logprobs.

1. The C++ sequence stores the list (csrc/engine/pybind_sequence.cpp:36). At each step, next_golden_tokens returns 2 × nextn + 1 values from the current output length, padded with -1 (csrc/engine/sequence.hpp:155-165).
2. prepare_sampler_inputs writes up to min(seqlen_q, nextn + 1) of them into golden_output, one per row; other rows get -1 (csrc/engine/prepare_sample.cpp:229-248).
3. Prefill: _update_golden_tokens replaces the sample where golden_output != -1, and looks up the logprob of the golden token (deepinfer/model_executor/layers/sampler.py:665-677).
4. Decode: a row with a golden token accepts its draft only if the draft equals the golden token, and outputs the golden token (deepinfer/model_executor/layers/tilelang_ops/sampler.py:239-242).

Result: the request follows the golden text exactly. The logprobs show how likely the model finds that text, and the accept counts show how well the drafts match it.

_get_so_mask runs on the CPU while the device computes the logits (deepinfer/model_executor/layers/sampler.py:135-138). It returns nothing when no request in the batch has a grammar (deepinfer/model_executor/layers/sampler.py:318-319).

• Layout. One pinned int32 bit mask [batch, ceil(V / 32)] per step: max_MTP_value + 1 masks, so 6 by default (deepinfer/model_executor/layers/sampler.py:323-327). Mask 0 is for row 0, whose input is the last real token. Mask j (1 to 5) is for row j, whose input is draft j.
• Drafts. The matcher accepts the drafts one by one and fills one mask per accepted draft. It stops at the first draft that the grammar forbids. Then rollback undoes the trial (deepinfer/model_executor/layers/sampler.py:366-376).
• Copy. A side stream copies the masks to the device and records an event. The main stream waits for that event before it applies the masks (deepinfer/model_executor/layers/sampler.py:391-404, deepinfer/model_executor/layers/sampler.py:163-164).
• Errors. A matcher error kills only its own request (so_status = -1). A comment explains why: an exception here would stop all requests and time out the NCCL barrier (deepinfer/model_executor/layers/sampler.py:377-386).

The page Structured outputs explains the grammar matcher.

With context parallelism, only the driver rank samples. On a follower rank, sampler_inputs.driver is None (deepinfer/model_executor/layers/sampler.py:136-137). The follower still does two things, because they contain collective operations (deepinfer/model_executor/layers/sampler.py:145-161):

• It runs compute_prompt_logprobs. Translated comment: "The result is valid only on the driver; the follower only takes part in the collective communication of the embedding gather."
• In a decode step, it runs DFlash.decode_forward without a driver input. It gets the sampled tokens of the driver by a broadcast (deepinfer/model_executor/models/mtp/dflash.py:358-365).

Then it returns an empty SamplerOutput.