Sampling¶

Table of Contents¶

• What Is Sampling?
• The Vocabulary Distribution
• Greedy Decoding
• Temperature Scaling
• Top-K Sampling
• Top-P (Nucleus) Sampling
• Min-P Sampling
• XTC (Exclude Top Choices)
• P-less Sampling
• How Filters Combine
• Penalty-Based Methods
• Repetition Penalty
• Frequency and Presence Penalties
• DRY Penalty
• Seed Control for Deterministic Output
• Implementation: The Sampling Pipeline
• Speculative Decoding
• Configuration Reference

What Is Sampling?¶

At each decode step, the model produces a logit score for every token in its vocabulary: a raw, unnormalized real number that represents how probable the model considers that token to be next. A softmax function converts these logits into a probability distribution.

Sampling is the process of selecting one token from this distribution to continue the generation.

The simplest approach is to always pick the highest-probability token. This is deterministic and tends to produce fluent, predictable output. However, always selecting the peak probability token can lead to output that is repetitive, overly cautious, and lacking in variety. For creative tasks, instruction following, or diverse responses, some randomness is desirable.

Sampling strategies exist on a spectrum from fully deterministic (greedy) to fully random (sample from the raw softmax distribution). The strategies available in mlxcel let you tune where on this spectrum each generation sits.

The Vocabulary Distribution¶

After a forward pass, the model returns logits of shape [batch, seq_len, vocab_size]. The sampling function first extracts the logits for the last sequence position (the position that represents the prediction for the next token), producing a [batch, vocab_size] tensor.

All probability transformations are applied to this last-position logit vector before sampling. The order of operations matters: penalty adjustments run first, then temperature, then vocabulary filters, then the final categorical sample.

Greedy Decoding¶

Greedy decoding selects the token with the highest logit at every step, equivalent to applying argmax over the vocabulary. It is deterministic for the same checkpoint, prompt, build, device, execution path, and external state; changes in kernels, batching, precision, or model state can change ties or logits.

In mlxcel, greedy decoding is activated when temperature <= 0.0. The build_sampling_config() function in src/execution/sampling.rs detects this case and sets temperature = 0.0 and top_k = 1.

Greedy decoding is appropriate for tasks where correctness matters more than diversity: code completion, structured output, factual questions. The default temperature in the mlxcel generate CLI is 0.0, meaning greedy decoding is used unless you set --temp explicitly.

Temperature Scaling¶

Temperature controls how sharply the probability distribution is peaked. It does this by dividing the logits by the temperature value before applying softmax.

P(token_i) = softmax(logits / T)[i]

• T close to 0: The distribution becomes very sharp; the highest-logit token has nearly all the probability. At T=0, this becomes greedy decoding.
• T = 1.0: The distribution is unchanged from the model's raw output.
• T > 1.0: The distribution becomes flatter; lower-probability tokens get relatively more weight.

xychart-beta
    title "Effect of Temperature on Token Probabilities"
    x-axis ["token A", "token B", "token C", "token D", "token E"]
    y-axis "Probability" 0 --> 0.8
    line [0.72, 0.18, 0.06, 0.02, 0.01]
    line [0.40, 0.25, 0.18, 0.10, 0.07]
    line [0.25, 0.22, 0.20, 0.17, 0.15]

The three curves represent approximately T=0.3 (top, sharp), T=1.0 (middle), and T=2.0 (bottom, flat) applied to the same underlying logits.

A temperature of 0.7 to 0.9 is a common starting point for chat and instruction models. Higher temperatures like 1.2 to 1.5 are used for creative writing.

Top-K Sampling¶

Top-K sampling restricts the sampling pool to the K tokens with the highest probability. Tokens outside the top K have their logits set to negative infinity before the softmax, effectively giving them zero probability.

For example, with K=50, only the 50 most likely tokens can be sampled. The relative probabilities within those 50 tokens are preserved.

Top-K alone can be problematic because K is a fixed count regardless of the shape of the distribution. Sometimes the top 50 tokens cover nearly all the probability mass; other times they cover only 30%. Top-P was designed to address this.

In mlxcel, top_k = 0 disables top-K filtering (all tokens remain eligible), and a top_k at or above the model's vocabulary size behaves the same way, as in llama-server. The default in the server is top_k = 40.

Top-P (Nucleus) Sampling¶

Top-P (nucleus) sampling selects the smallest set of tokens whose cumulative probability exceeds the threshold P. Tokens outside this nucleus have their logits set to negative infinity.

For example, with P=0.9, tokens are sorted by probability from highest to lowest. The algorithm includes tokens until their cumulative probability reaches 0.9, then zeros out the rest.

The key advantage over Top-K is that the size of the nucleus adapts to the distribution's shape: - When the model is very confident (peaked distribution), the nucleus may contain only 2-3 tokens. - When the model is uncertain (flat distribution), the nucleus may contain hundreds of tokens.

This means sampling is more constrained when the model is sure and more exploratory when the model is uncertain, which generally produces better-calibrated output.

A common setting is top_p = 0.9 or top_p = 0.95.

Min-P Sampling¶

Min-P sampling sets a dynamic minimum probability threshold based on the highest-probability token in the distribution. A token is included in the sampling pool only if its probability is at least min_p * max_probability.

For example, with min_p = 0.05, if the most probable token has probability 0.4, then only tokens with probability >= 0.02 (5% of 0.4) are eligible.

This has properties similar to Top-P: it adapts to the distribution's peak. When the model is very confident (high max probability), the threshold rises and fewer tokens are eligible. When the model is uncertain (low max probability), more tokens remain eligible.

Min-P tends to be more aggressive than Top-P at pruning implausible tokens while preserving diversity in uncertain situations. min_p = 0.05 to min_p = 0.1 are common values.

In mlxcel, Min-P is implemented in src/lib/mlxcel-core/src/sampling.rs via min_p_filter(). The fused_sample() C++ function handles temperature, Top-K, Top-P, and Min-P in a single operation to reduce FFI round-trips.

Two kernels back the draw itself, both on by default:

• Categorical sampling uses a softmax-free Gumbel-max kernel, drawing from the logits directly instead of materializing a softmax. MLXCEL_SAMPLING_GUMBEL=0 falls back to the softmax path.
• Top-K, Top-P and Min-P are resolved by dual-pivot rejection rather than a full-vocabulary sort. MLXCEL_SAMPLING_REJECTION=0 falls back to the sorting path.

Both fallbacks exist for A/B measurement and rollback; the defaults are the faster paths.

XTC (Exclude Top Choices)¶

XTC inverts the usual goal of a truncation filter. Where Top-K, Top-P and Min-P remove the unlikely tail, XTC probabilistically removes the most likely tokens, leaving a plausible but less predictable one. It is aimed at creative writing, where the highest-probability continuation is often the blandest.

Two parameters control it:

• xtc_probability (0.0 to 1.0, default 0.0) is the chance that the filter fires on any given token. At 0.0 XTC is disabled entirely.
• xtc_threshold (0.0 to 0.5, default 0.1) is the probability a token must exceed to be a removal candidate.

When the filter fires, every token above the threshold except the least likely of them is removed, so the model keeps a coherent option rather than falling through to noise. Special tokens, including end-of-sequence, are never removed, so XTC cannot prevent a model from stopping.

XTC is per-request and opt-in: it has no server-level CLI default, and an absent request field always resolves to disabled. It is available on /v1/chat/completions, /v1/completions, /v1/responses, and the native completion endpoint.

P-less Sampling¶

P-less sampling is a truncation filter with no hyperparameter. It computes L = sum(p^2) over the temperature-scaled distribution and keeps only the tokens whose probability is at least L. A peaked distribution has a large L, so few tokens survive; a flat distribution has a small L, so many do. The most likely token always survives, so the filter never empties the distribution.

P-less is opt-in per request through the boolean p_less field (default false) on /v1/chat/completions, /v1/completions, /v1/responses, and the native /completion endpoint. It has no server-wide default and is not echoed in /props. On the CLI, mlxcel generate --p-less and mlxcel run --p-less turn it on. Greedy requests (temperature = 0 or top_k = 1) and mirostat requests skip it.

How Filters Combine¶

mlxcel applies the sampler stages in a fixed order that matches the llama-server chain:

1. Token bias: logit_bias, language bias, and token suppression.
2. History penalties: repetition, DRY, then frequency and presence.
3. Row filters: top-n-sigma, then p-less, then typical-p, on the penalized, untempered logits (p-less reads the temperature to compute its threshold).
4. Top-K, Top-P, Min-P: evaluated on the untempered distribution.
5. XTC: applied to the renormalized filtered row.
6. Temperature and categorical sample: one temperature scaling, applied only to the final draw.

When both Top-K and Top-P are active, a token must survive both filters to be eligible. Min-P applies independently. Because the truncation filters read the untempered distribution, temperature changes how the surviving tokens are weighted but not which tokens survive.

Setting any filter to its disabled value (topk=0, topp=1.0, min_p=0.0) effectively removes that filter from the pipeline.

Penalty-Based Methods¶

Penalty methods adjust logit scores based on the token history, suppressing tokens that have recently appeared. They run before the distribution filters.

Repetition Penalty¶

Repetition penalty is a multiplicative adjustment that discourages the model from reusing tokens that have appeared in the context. For each token in the history:

• If the logit is positive: divide by the penalty factor (making it smaller).
• If the logit is negative: multiply by the penalty factor (making it more negative).

This asymmetric treatment means both positive and negative logit tokens are pushed away from their original values in a direction that reduces their probability.

A penalty of 1.0 is neutral (no effect). Values like 1.1 to 1.3 provide mild discouragement of repetition without harshly prohibiting it. Values above 1.5 can distort the output for models trained with repetition.

Repetition penalty considers all distinct tokens in the history (deduplication is applied before the penalty calculation). It does not account for how many times a token appeared.

Frequency and Presence Penalties¶

These are OpenAI-style penalties imported for API compatibility.

Frequency penalty subtracts frequency_penalty * count(token) from the logit for each token, where count(token) is the number of times that token appeared in the history. This linearly scales the penalty with frequency, more strongly suppressing tokens that have been overused.

Presence penalty subtracts presence_penalty from the logit for any token that has appeared at all, regardless of count. This provides a flat discouragement of any token that has been seen before.

Both penalties can be applied simultaneously. Their combined effect is:

logit[i] -= frequency_penalty * count(i) + presence_penalty * (1 if i in history else 0)

Typical values are in the range 0.0 to 2.0. Values above 2.0 can produce incoherent output.

DRY Penalty¶

DRY (Don't Repeat Yourself) is a more sophisticated penalty that targets repeated phrases rather than individual tokens. It penalizes tokens that would continue a sequence of tokens that has already appeared earlier in the context.

The algorithm works as follows:

1. Scan the token history for occurrences of the most recently generated token (the "last token").
2. For each prior occurrence of that token, extend the match backwards: how many preceding tokens match the current context's preceding tokens?
3. If the match length exceeds dry_allowed_length, apply a penalty to the token that followed the previous occurrence (the token that would be selected to continue the repeated phrase).
4. The penalty magnitude is: dry_multiplier * dry_base ^ (match_length - dry_allowed_length).

The exponential scaling with dry_base means that longer repeated phrases receive exponentially larger penalties. A match of 2 tokens above the allowed length is penalized much less than a match of 5 tokens, so short repetitions receive mild penalties and long repetitions receive strong penalties.

DRY sequence breakers are token IDs (typically newlines, punctuation, or formatting tokens) that terminate a match. If a breaker appears in the backward scan, the match stops. This prevents the penalty from extending across natural sentence boundaries where repetition is appropriate.

DRY is disabled when dry_multiplier = 0.0 (the default). Typical values when enabled are dry_multiplier = 0.8 to 1.0 with dry_base = 1.75 and dry_allowed_length = 2.

The dry_penalty_last_n parameter limits the lookback window, following llama-server's sentinels: -1 considers the full history, 0 disables DRY, and a value like 1024 limits lookback to the most recent 1024 tokens. mlxcel generate defaults to -1; the server and mlxcel run default to 64.

Seed Control for Deterministic Output¶

For reproducible sampling (when temperature > 0), a seed can be provided via SamplingConfig::seed. When a seed is set, mlxcel calls ffi::seed_rng() before generation begins to initialize MLX's random number generator.

A fixed seed makes sampling reproducible when the checkpoint, prompt, parameters, build, hardware, and execution path are also held fixed. Batch shape or kernel changes can consume random values differently, so a seed is not a cross-version or cross-device byte-identity guarantee.

When seed = None, the RNG state is not modified and generation uses whatever state is current, producing different results on each run.

Implementation: The Sampling Pipeline¶

Every decode path samples a row through one step, RowSampler in src/lib/mlxcel-core/src/sampling_row_step.rs. The decode engine draws through it for the server's decode rows and prefill completions and for the CLI's generate and run, so the CLI and the server run the same chain. Token bias is composed once per request by compose_token_bias in sampling_token_bias.rs, and the chain itself lives in src/lib/mlxcel-core/src/sampling.rs:

flowchart LR
    LOGITS["Logits\n[batch, seq, vocab]"] --> SLICE["last-position slice\n[1, vocab]"]
    SLICE --> MASK["structured-output mask\n(server, when set)"]
    MASK --> BIAS["token bias\nlogit_bias, language bias, suppression"]
    BIAS --> REP["repetition penalty"]
    REP --> DRY["DRY penalty"]
    DRY --> FP["frequency/presence penalty"]
    FP --> ROW["row filters\ntop-n-sigma, p-less, typical-p"]
    ROW --> FUSED["fused_sample (C++)\ntop_k + top_p + min_p (untempered)\n+ XTC + temperature + draw"]
    FUSED --> TOKEN["sampled token ID"]

Each penalty and filter stage is skipped without adding graph nodes when its parameter is at the disabled value. The penalty stages read the token history; before a step that reads history is drawn, the history ends with the token emitted by the previous step, so a pipelined caller reads that token on the host first. A sampler that carries feedback state but reads no history (mirostat, adaptive-p) keeps building the next draw straight from the lazy graph.

The draw returns the sampled token without forcing GPU evaluation, so the decode loop can submit the next step's forward before waiting for the GPU. When every row of a server decode batch is eligible for fusion and shares the same parameters, the scheduler draws the whole [B, vocab] batch in one fused call instead of one draw per row.

Speculative Decoding¶

Speculative decoding asks a faster drafter for several candidates and verifies them with the target model. The exactness contract depends on the drafter mode.

Prompt lookup (mlxcel generate --prompt-lookup, server --draft-kind prompt-lookup) and the full-model drafter (--draft-model with an ordinary checkpoint) accept a proposed token only when it equals a fresh draw from the target's sampler, which is lossless for their one-hot proposals; at temperature 0 the draw is the argmax. At temperature 0, block verification is not guaranteed to reproduce the plain-decode stream (#2778). MTP and DFlash use separate round loops and have mode-specific greedy exactness, penalty, qualification, and engagement conditions; they must not be covered by a blanket distribution-preserving claim. See Speculative Decoding.

An ordinary smaller checkpoint drafts through FullModelDrafter (src/lib/mlxcel-core/src/speculative/full_model_drafter.rs), which proposes its argmax from its own KV caches and is verified by the same Engine::verify_round that prompt lookup uses. There is no separate draft-model decode loop. --draft-model on mlxcel generate and mlxcel run, and --draft-model or --model-draft on mlxcel serve and mlxcel-server, auto-detect an ordinary model directory (no DFlash-family marker, no MTP model_type) as full-model; --draft-kind <dflash|mtp|full-model> selects explicitly, and prompt-lookup is also accepted on the server entry points. --draft and --draft-max are llama-server-compatible aliases where documented. --draft-block-size is the single proposal-width control for every drafter kind. A --draft-model whose config.json is missing or malformed fails at startup with the config error.

A useful full-model drafter is substantially faster than the target and uses a compatible tokenizer and vocabulary. Always measure accepted tokens per round and end-to-end latency: verification overhead can make speculation slower.

Configuration Reference¶

All sampling parameters are assembled into a SamplingConfig in src/execution/sampling.rs.

When temperature <= 0.0, build_sampling_config() activates the greedy path regardless of other filter settings. Penalty parameters are still applied in greedy mode.