36 · Content revision 5

Shape the choices, then draw once

Learn how positive temperature, stable top-k filtering, and a restored random-generator state turn decoder logits into controlled, replayable uncached LLM generation.

Start with four logits and make every choice visible

The decoder supplies final-position logits in token-ID order: $[0,1,1,2]$ . Their stable descending-logit order, with ascending token ID as the tie-breaker, is $[3,1,2,0]$ . Token $3$ is unambiguously first. Tokens $1$ and $2$ have equal logits, so this implementation places the smaller ID first. That local rule makes the tied candidate order deterministic.

First keep all four tokens and vary only the positive temperature. At $\tau=0.5$ , token $3$ receives probability $0.775803492574$ ; at $\tau=1$ , it receives $0.534446645389$ ; at $\tau=2$ , it receives $0.387455619000$ . Lower temperature sharpens existing logit differences. Higher temperature flattens them. Neither changes the stable rank.

Now set $\tau=1$ and $k=2$ . The survivor order is $[3,1]$ . Token $1$ survives the tied boundary and receives $q_1=0.268941421370$ ; token $2$ is removed and receives exactly $q_2=0$ . Token $3$ receives $q_3=0.731058578630$ . The retained probabilities sum to $1$ .

With $k=3$ , SplitMix64 seed $36$ produces the reproducible sequence $[3,2,2,2,3,3,3,3]$ . For example, the second unit draw is $u=0.338833394523$ . It lies in token $2$ ‘s half-open interval $[0.211941557617,0.423883115234)$ , so that draw selects token $2$ .

Filter the candidate set, then renormalize

For finite logits, finite positive temperature, and $1\le k\le V$ , the final token probability is

q_i^{(\tau,k)}=\frac{\mathbf{1}[i\in K_k]\exp(\ell_i/\tau)}{\sum_j\mathbf{1}[j\in K_k]\exp(\ell_j/\tau)}

The implementation subtracts the largest retained scaled logit before applying the exponential. This does not alter probability ratios, but prevents a large common offset from overflowing. Removed IDs keep exact probability $0$ rather than a tiny approximation.

Temperature must be strictly positive. The mathematical limit $\tau\to0^+$ helps explain concentration toward a unique maximum, but literal $\tau=0$ would divide by zero. Greedy decoding is therefore a separate policy: it validates the same logits, chooses the first stable rank, and consumes no random draw. Stochastic $k=1$ chooses the same ID but deliberately consumes one draw, because it still follows the sampling policy.

Keep logits, candidates, and probabilities distinct

$q_i^{(\tau,k)}$ is token $i$ ‘s final probability after all three operations.
$\tau$ is a finite positive temperature. Lower values sharpen; higher values flatten.
$k$ is the exact number of retained token IDs.
$V$ is the vocabulary size, so the valid candidate count satisfies $1\le k\le V$ .
$K_k$ is the stable top- $k$ set. Equal logits use ascending token ID here.
$\mathbf{1}[i\in K_k]$ is $1$ for a retained ID and $0$ for a removed ID.
$\ell_i$ is the decoder logit for token $i$ at the final prefix position.
$i$ names the candidate whose final probability is being described.
$j$ ranges over the vocabulary in the denominator.
$\exp$ turns each shifted scaled logit into a positive softmax weight.

These names protect an important boundary. A logit is an unrestricted decoder score, stable rank determines membership, and a probability is normalized mass. Temperature changes score gaps before normalization. Top-k changes which terms exist in the denominator. The categorical draw acts only after both decisions.

The sampler traverses retained mass in ascending token-ID order with half-open intervals $[a_i,b_i)$ . A unit draw $u$ selects token $i$ exactly when $a_i\le u<b_i$ . Here $a_i$ is token $i$ ‘s cumulative lower endpoint and $b_i$ is its cumulative upper endpoint. The final interval accepts the tiny endpoint discrepancy that can arise from floating-point summation; it does not silently revive a removed token.

From constrained search to open-ended LLM sampling

Likelihood-maximizing beam decoding is useful when a source tightly constrains the target, but open-ended continuation admits many plausible futures; beam output can become generic or repetitive, while unrestricted sampling can admit an unreliable low-probability tail.

Fan, Lewis, and Dauphin make that open-ended boundary concrete. Fan, Lewis, and Dauphin sample at each step from the ten most likely words, tune a generation-time softmax temperature, and report that this task-bounded strategy works better for their open-ended stories than beam search, while unrestricted random sampling can introduce damaging unlikely words.

The GPT-2 report then shows the policy in a large Transformer language model. The GPT-2 report uses top-k random sampling with $k=2$ for one summarization setup and $k=40$ for open WebText continuations, showing truncated stochastic decoding in large Transformer language-model practice without claiming one $k$ is universal.

Holtzman and colleagues examine the trade-off directly. Holtzman and colleagues compare maximization and stochastic decoders on GPT-2, define top-k as sampling from the $k$ highest-probability tokens after renormalization, give the temperature-scaled softmax, and show why flat and peaked contexts make a fixed $k$ an imperfect compromise.

Open-ended story systems combined softmax temperature with top-k sampling, GPT-2 used top-k for summaries and long continuations, and later GPT-2 analysis made the truncation and renormalization trade-off explicit while showing why one fixed $k$ cannot fit every context.

Controlled stochastic decoding turns an autoregressive LLM distribution into an adjustable diversity-versus-concentration distribution; a fixed random-generator state, deterministic tie-breaking, and deterministic interval traversal can replay its choices, while later methods can replace the fixed candidate count without changing the decoder logits.

Top-k is therefore a clear first controlled stochastic decoder, not a universal quality guarantee, a hallucination defense, or the endpoint of decoding research. Stable token-ID ties, seed replay, invalid-setting behavior, EOS, and context stops are this implementation’s reproducibility rules rather than claims about the cited systems.

The executable contrast measures the mechanism on the chapter’s four logits. Greedy chooses token $3$ without advancing the random generator. At $\tau=1$ , fixed $k=3$ keeps IDs $[3,1,2]$ : those candidates held $0.927670511871$ of the full-softmax mass, while truncation removes $0.072329488129$ before renormalization. That measurement does not determine which policy produces better text; it makes the fixed-cardinality boundary visible.

Measure greedy selection and probability mass removed by fixed top-k truncation rust/demos/ch36-temperature-top-k/src/lib.rs#historical-decoding-contrast

/// Measures how greedy choice and fixed top-k truncation differ on one logit row.
pub fn historical_decoding_contrast() -> Result<HistoricalDecodingContrast, FixtureError> {
    let full = sampling_distribution(
        &LOGITS,
        SamplingMode::TemperatureTopK {
            temperature: 1.0,
            top_k: LOGITS.len(),
        },
    )?;
    let truncated = sampling_distribution(
        &LOGITS,
        SamplingMode::TemperatureTopK {
            temperature: 1.0,
            top_k: SAMPLE_TOP_K,
        },
    )?;
    let retained_token_ids = truncated.survivors().to_vec();
    let retained_full_probability_mass = full
        .candidates()
        .iter()
        .filter(|candidate| retained_token_ids.contains(&candidate.token_id()))
        .map(|candidate| candidate.probability())
        .sum::<f64>();
    let removed_full_probability_mass = full
        .candidates()
        .iter()
        .filter(|candidate| !retained_token_ids.contains(&candidate.token_id()))
        .map(|candidate| candidate.probability())
        .sum::<f64>();
    let mut greedy_rng = SplitMix64::from_seed(SAMPLE_SEED);
    let initial_rng_state = greedy_rng.state();
    let greedy = sample_next_token(&LOGITS, SamplingMode::Greedy, &mut greedy_rng)?;

    require(
        retained_token_ids.len() == SAMPLE_TOP_K,
        "historical top-k candidate count changed",
    )?;
    require(
        (retained_full_probability_mass + removed_full_probability_mass - 1.0).abs() < 1e-12,
        "historical full-softmax mass no longer sums to one",
    )?;

    Ok(HistoricalDecodingContrast {
        greedy_token: greedy.token_id(),
        greedy_rng_advanced: greedy_rng.state() != initial_rng_state,
        top_k: SAMPLE_TOP_K,
        retained_token_ids,
        retained_full_probability_mass,
        removed_full_probability_mass,
    })
}

Validate first, rank stably, and advance the random stream once

SamplingMode::Greedy and SamplingMode::TemperatureTopK are explicit choices. Both reject empty or nonfinite logits. The stochastic policy additionally rejects nonfinite or nonpositive temperature and any $k$ outside $1\le k\le V$ . All validation happens before the SplitMix64 state changes.

All three public operations start with the same calculation: validate the inputs, rank finite logits by descending value and ascending token ID, filter to the requested candidate count, and turn the retained logits into normalized probabilities. sample_next_token returns only the selected token ID, optional unit draw, and selected half-open interval. It does not build the token-by-token candidate records or the rank-ordered survivor list used for inspection. sample_next_token_with_trace uses the same calculated probabilities and the same interval-selection code, but builds that full inspectable distribution before any draw. sampling_distribution builds the distribution without drawing. The ordinary call still needs temporary arrays of ranked token IDs and probabilities to perform the algorithm; it avoids only the additional inspection records.

In an inspectable distribution, candidates appear in token-ID order and the separate survivor list appears in stable rank order. Removed IDs have exact zero probability after max-shifted normalization. Categorical selection visits positive intervals in ascending token-ID order. Greedy consumes no draw; every valid stochastic call consumes exactly one unit draw, including top-k with $k=1$ .

Prepare sampling once, then return compact selection evidence or an explicitly requested full trace rust/crates/llm-from-scratch/src/generation/sampling.rs#sampling-policy

/// The two intentionally distinct next-token policies taught in Chapter 36.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum SamplingMode {
    /// Select the highest logit, resolving equal logits by lower token ID.
    Greedy,
    /// Sample after positive-temperature scaling and stable top-k truncation.
    TemperatureTopK { temperature: f64, top_k: usize },
}

/// One vocabulary entry after ranking and probability construction.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct SamplingCandidate {
    token_id: u32,
    logit: f64,
    rank: usize,
    retained: bool,
    probability: f64,
}

impl SamplingCandidate {
    pub const fn token_id(self) -> u32 {
        self.token_id
    }

    pub const fn logit(self) -> f64 {
        self.logit
    }

    /// One-based rank under descending logit and ascending token-ID ties.
    pub const fn rank(self) -> usize {
        self.rank
    }

    pub const fn retained(self) -> bool {
        self.retained
    }

    pub const fn probability(self) -> f64 {
        self.probability
    }
}

/// The complete token-ID-ordered distribution plus rank-ordered survivors.
#[derive(Clone, Debug, PartialEq)]
pub struct SamplingDistribution {
    mode: SamplingMode,
    candidates: Vec<SamplingCandidate>,
    survivors: Vec<u32>,
}

impl SamplingDistribution {
    pub const fn mode(&self) -> SamplingMode {
        self.mode
    }

    /// Candidates are always returned in ascending token-ID order.
    pub fn candidates(&self) -> &[SamplingCandidate] {
        &self.candidates
    }

    /// Survivors are returned in stable descending-logit rank order.
    pub fn survivors(&self) -> &[u32] {
        &self.survivors
    }

    pub fn probability_sum(&self) -> f64 {
        compensated_sum(
            self.candidates
                .iter()
                .map(|candidate| candidate.probability),
        )
    }

    pub fn candidate(&self, token_id: u32) -> Option<SamplingCandidate> {
        self.candidates
            .get(usize::try_from(token_id).ok()?)
            .copied()
    }
}

/// One selected token and the half-open categorical interval that selected it.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct SampledToken {
    token_id: u32,
    unit_draw: Option<f64>,
    interval_start: f64,
    interval_end: f64,
}

impl SampledToken {
    pub const fn token_id(self) -> u32 {
        self.token_id
    }

    pub const fn unit_draw(self) -> Option<f64> {
        self.unit_draw
    }

    pub const fn interval_start(self) -> f64 {
        self.interval_start
    }

    pub const fn interval_end(self) -> f64 {
        self.interval_end
    }
}

/// A compact selection paired with the complete distribution used to produce it.
#[derive(Clone, Debug, PartialEq)]
pub struct SamplingDecision {
    sampled: SampledToken,
    distribution: SamplingDistribution,
}

impl SamplingDecision {
    pub const fn token_id(&self) -> u32 {
        self.sampled.token_id()
    }

    pub const fn unit_draw(&self) -> Option<f64> {
        self.sampled.unit_draw()
    }

    pub const fn interval_start(&self) -> f64 {
        self.sampled.interval_start()
    }

    pub const fn interval_end(&self) -> f64 {
        self.sampled.interval_end()
    }

    pub const fn distribution(&self) -> &SamplingDistribution {
        &self.distribution
    }
}

/// A setting or numerical input that cannot define a sampling distribution.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum SamplingError {
    EmptyLogits,
    VocabularyTooLarge {
        classes: usize,
    },
    NonFiniteLogit {
        token_id: usize,
        value: f64,
    },
    InvalidTemperature {
        value: f64,
    },
    InvalidTopK {
        top_k: usize,
        vocabulary_size: usize,
    },
    AllocationFailed {
        values: usize,
    },
    InvalidProbabilitySum {
        value: f64,
    },
}

impl fmt::Display for SamplingError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::EmptyLogits => formatter.write_str("sampling needs at least one finite logit"),
            Self::VocabularyTooLarge { classes } => write!(
                formatter,
                "sampling vocabulary of {classes} classes does not fit u32 token IDs"
            ),
            Self::NonFiniteLogit { token_id, value } => {
                write!(
                    formatter,
                    "logit for token {token_id} is not finite: {value}"
                )
            }
            Self::InvalidTemperature { value } => write!(
                formatter,
                "sampling temperature must be finite and positive, received {value}"
            ),
            Self::InvalidTopK {
                top_k,
                vocabulary_size,
            } => write!(
                formatter,
                "top-k must be in 1..={vocabulary_size}, received {top_k}"
            ),
            Self::AllocationFailed { values } => {
                write!(
                    formatter,
                    "cannot allocate sampling evidence for {values} values"
                )
            }
            Self::InvalidProbabilitySum { value } => write!(
                formatter,
                "sampling probabilities did not normalize to one: {value}"
            ),
        }
    }
}

impl Error for SamplingError {}

fn retained_count(mode: SamplingMode, vocabulary_size: usize) -> Result<usize, SamplingError> {
    match mode {
        SamplingMode::Greedy => Ok(1),
        SamplingMode::TemperatureTopK { temperature, top_k } => {
            if !temperature.is_finite() || temperature <= 0.0 {
                return Err(SamplingError::InvalidTemperature { value: temperature });
            }
            if top_k == 0 || top_k > vocabulary_size {
                return Err(SamplingError::InvalidTopK {
                    top_k,
                    vocabulary_size,
                });
            }
            Ok(top_k)
        }
    }
}

fn compensated_sum(values: impl IntoIterator<Item = f64>) -> f64 {
    let mut sum = 0.0;
    let mut compensation = 0.0;
    for value in values {
        let corrected = value - compensation;
        let next = sum + corrected;
        compensation = (next - sum) - corrected;
        sum = next;
    }
    sum
}

fn stable_ranks(logits: &[f64]) -> Result<Vec<usize>, SamplingError> {
    if logits.is_empty() {
        return Err(SamplingError::EmptyLogits);
    }
    if u32::try_from(logits.len() - 1).is_err() {
        return Err(SamplingError::VocabularyTooLarge {
            classes: logits.len(),
        });
    }
    for (token_id, &value) in logits.iter().enumerate() {
        if !value.is_finite() {
            return Err(SamplingError::NonFiniteLogit { token_id, value });
        }
    }

    let mut ranked = Vec::new();
    ranked
        .try_reserve_exact(logits.len())
        .map_err(|_| SamplingError::AllocationFailed {
            values: logits.len(),
        })?;
    ranked.extend(0..logits.len());
    ranked.sort_unstable_by(|&left, &right| {
        logits[right]
            .partial_cmp(&logits[left])
            .unwrap_or(Ordering::Equal)
            .then_with(|| left.cmp(&right))
    });
    Ok(ranked)
}

fn scaled_gap(logit: f64, maximum: f64, temperature: f64) -> f64 {
    if temperature < 1.0 {
        (logit - maximum) / temperature
    } else {
        logit / temperature - maximum / temperature
    }
}

struct PreparedSampling {
    mode: SamplingMode,
    ranked: Vec<usize>,
    probabilities: Vec<f64>,
    keep: usize,
}

fn prepare_sampling(logits: &[f64], mode: SamplingMode) -> Result<PreparedSampling, SamplingError> {
    let ranked = stable_ranks(logits)?;
    let keep = retained_count(mode, logits.len())?;
    let maximum = logits[ranked[0]];

    let mut probabilities = Vec::new();
    probabilities
        .try_reserve_exact(logits.len())
        .map_err(|_| SamplingError::AllocationFailed {
            values: logits.len(),
        })?;
    probabilities.resize(logits.len(), 0.0);

    match mode {
        SamplingMode::Greedy => probabilities[ranked[0]] = 1.0,
        SamplingMode::TemperatureTopK { temperature, .. } => {
            for (position, &token_id) in ranked[..keep].iter().enumerate() {
                let weight = if position == 0 {
                    1.0
                } else {
                    scaled_gap(logits[token_id], maximum, temperature).exp()
                };
                probabilities[token_id] = weight;
            }
            let weight_sum = compensated_sum(
                ranked[..keep]
                    .iter()
                    .map(|&token_id| probabilities[token_id]),
            );
            if !weight_sum.is_finite() || weight_sum <= 0.0 {
                return Err(SamplingError::InvalidProbabilitySum { value: weight_sum });
            }
            for &token_id in &ranked[..keep] {
                probabilities[token_id] /= weight_sum;
            }
        }
    }

    let probability_sum = compensated_sum(probabilities.iter().copied());
    if !probability_sum.is_finite() || (probability_sum - 1.0).abs() > PROBABILITY_TOLERANCE {
        return Err(SamplingError::InvalidProbabilitySum {
            value: probability_sum,
        });
    }

    Ok(PreparedSampling {
        mode,
        ranked,
        probabilities,
        keep,
    })
}

fn materialize_distribution(
    logits: &[f64],
    prepared: &PreparedSampling,
) -> Result<SamplingDistribution, SamplingError> {
    let PreparedSampling {
        mode,
        ranked,
        probabilities,
        keep,
    } = prepared;

    let mut rank_by_token = Vec::new();
    rank_by_token
        .try_reserve_exact(logits.len())
        .map_err(|_| SamplingError::AllocationFailed {
            values: logits.len(),
        })?;
    rank_by_token.resize(logits.len(), 0);
    for (position, &token_id) in ranked.iter().enumerate() {
        rank_by_token[token_id] = position + 1;
    }

    let mut candidates = Vec::new();
    candidates
        .try_reserve_exact(logits.len())
        .map_err(|_| SamplingError::AllocationFailed {
            values: logits.len(),
        })?;
    for (token_id, ((&logit, &probability), &rank)) in logits
        .iter()
        .zip(probabilities)
        .zip(&rank_by_token)
        .enumerate()
    {
        candidates.push(SamplingCandidate {
            token_id: u32::try_from(token_id).expect("validated token ID must fit u32"),
            logit,
            rank,
            retained: rank <= *keep,
            probability,
        });
    }

    let mut survivors = Vec::new();
    survivors
        .try_reserve_exact(*keep)
        .map_err(|_| SamplingError::AllocationFailed { values: *keep })?;
    for &token_id in &ranked[..*keep] {
        survivors.push(u32::try_from(token_id).expect("validated token ID must fit u32"));
    }
    Ok(SamplingDistribution {
        mode: *mode,
        candidates,
        survivors,
    })
}

fn select_prepared(prepared: &PreparedSampling, rng: &mut SplitMix64) -> SampledToken {
    if prepared.mode == SamplingMode::Greedy {
        return SampledToken {
            token_id: u32::try_from(prepared.ranked[0]).expect("validated token ID must fit u32"),
            unit_draw: None,
            interval_start: 0.0,
            interval_end: 1.0,
        };
    }

    let draw = rng.next_unit_f64();
    let final_id = prepared
        .probabilities
        .iter()
        .enumerate()
        .rev()
        .find(|(_, probability)| **probability > 0.0)
        .map(|(token_id, _)| token_id)
        .expect("a normalized distribution must have a positive survivor");
    let mut start = 0.0;
    for (token_id, &probability) in prepared
        .probabilities
        .iter()
        .enumerate()
        .filter(|(_, probability)| **probability > 0.0)
    {
        let end = if token_id == final_id {
            1.0
        } else {
            (start + probability).min(1.0)
        };
        if draw < end || token_id == final_id {
            return SampledToken {
                token_id: u32::try_from(token_id).expect("validated token ID must fit u32"),
                unit_draw: Some(draw),
                interval_start: start,
                interval_end: end,
            };
        }
        start = end;
    }
    unreachable!("the final positive survivor covers every unit draw")
}

fn sample_with_observer<T>(
    logits: &[f64],
    mode: SamplingMode,
    rng: &mut SplitMix64,
    observe: impl FnOnce(&PreparedSampling) -> Result<T, SamplingError>,
) -> Result<(SampledToken, T), SamplingError> {
    let prepared = prepare_sampling(logits, mode)?;
    let observation = observe(&prepared)?;
    let sampled = select_prepared(&prepared, rng);
    Ok((sampled, observation))
}

/// Builds a complete, inspectable distribution without consuming randomness.
pub fn sampling_distribution(
    logits: &[f64],
    mode: SamplingMode,
) -> Result<SamplingDistribution, SamplingError> {
    let prepared = prepare_sampling(logits, mode)?;
    materialize_distribution(logits, &prepared)
}

/// Selects one token without retaining the complete inspectable distribution.
pub fn sample_next_token(
    logits: &[f64],
    mode: SamplingMode,
    rng: &mut SplitMix64,
) -> Result<SampledToken, SamplingError> {
    sample_with_observer(logits, mode, rng, |_| Ok(())).map(|(sampled, ())| sampled)
}

/// Selects one token and records the complete distribution used for inspection.
pub fn sample_next_token_with_trace(
    logits: &[f64],
    mode: SamplingMode,
    rng: &mut SplitMix64,
) -> Result<SamplingDecision, SamplingError> {
    let (sampled, distribution) = sample_with_observer(logits, mode, rng, |prepared| {
        materialize_distribution(logits, prepared)
    })?;
    Ok(SamplingDecision {
        sampled,
        distribution,
    })
}

generate_uncached retains the complete prefix. Before each decoder call it checks that the prefix fits the configured context. It evaluates without an autodiff graph, takes the final vocabulary row, samples one token, appends that token, and records only compact step evidence: prefix length, selected ID, draw, and interval. Because it calls the result-only sampler, it neither constructs nor retains the complete inspectable candidate distribution for each generated token. EOS is included before the loop stops. A valid capacity- $2$ prefix can therefore emit a next token even though the resulting length- $3$ sequence cannot be submitted for another call.

Recompute each complete prefix and stop at EOS, token budget, or context capacity rust/crates/llm-from-scratch/src/generation/sampling.rs#uncached-generation

/// Settings for one bounded, uncached autoregressive call.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct GenerationConfig {
    mode: SamplingMode,
    eos_token: Option<u32>,
    max_new_tokens: usize,
}

impl GenerationConfig {
    pub const fn new(mode: SamplingMode, eos_token: Option<u32>, max_new_tokens: usize) -> Self {
        Self {
            mode,
            eos_token,
            max_new_tokens,
        }
    }

    pub const fn mode(self) -> SamplingMode {
        self.mode
    }

    pub const fn eos_token(self) -> Option<u32> {
        self.eos_token
    }

    pub const fn max_new_tokens(self) -> usize {
        self.max_new_tokens
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GenerationStop {
    Eos,
    TokenLimit,
    ContextLimit,
}

#[derive(Clone, Debug, PartialEq)]
pub struct GenerationStep {
    prefix_length: usize,
    token_id: u32,
    unit_draw: Option<f64>,
    interval_start: f64,
    interval_end: f64,
}

impl GenerationStep {
    pub const fn prefix_length(&self) -> usize {
        self.prefix_length
    }

    pub const fn token_id(&self) -> u32 {
        self.token_id
    }

    pub const fn unit_draw(&self) -> Option<f64> {
        self.unit_draw
    }

    pub const fn interval_start(&self) -> f64 {
        self.interval_start
    }

    pub const fn interval_end(&self) -> f64 {
        self.interval_end
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct GenerationResult {
    prompt: Vec<u32>,
    generated: Vec<u32>,
    steps: Vec<GenerationStep>,
    stop: GenerationStop,
    full_prefix_calls: usize,
}

impl GenerationResult {
    pub fn prompt(&self) -> &[u32] {
        &self.prompt
    }

    pub fn generated(&self) -> &[u32] {
        &self.generated
    }

    pub fn steps(&self) -> &[GenerationStep] {
        &self.steps
    }

    pub const fn stop(&self) -> GenerationStop {
        self.stop
    }

    pub const fn full_prefix_calls(&self) -> usize {
        self.full_prefix_calls
    }
}

#[derive(Debug)]
pub enum GenerationError {
    Sampling(SamplingError),
    Model(DecoderModelError),
    EmptyPrompt,
    PromptTooLong {
        tokens: usize,
        max_positions: usize,
    },
    PromptTokenOutOfBounds {
        position: usize,
        token_id: u32,
        vocabulary_size: usize,
    },
    EosTokenOutOfBounds {
        token_id: u32,
        vocabulary_size: usize,
    },
    LogitCountMismatch {
        expected: usize,
        actual: usize,
    },
    AllocationFailed {
        values: usize,
    },
}

impl fmt::Display for GenerationError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::Sampling(error) => error.fmt(formatter),
            Self::Model(error) => error.fmt(formatter),
            Self::EmptyPrompt => formatter.write_str("generation needs a nonempty prompt"),
            Self::PromptTooLong {
                tokens,
                max_positions,
            } => write!(
                formatter,
                "prompt has {tokens} tokens, exceeding context capacity {max_positions}"
            ),
            Self::PromptTokenOutOfBounds {
                position,
                token_id,
                vocabulary_size,
            } => write!(
                formatter,
                "prompt token {token_id} at position {position} is out of bounds for vocabulary {vocabulary_size}"
            ),
            Self::EosTokenOutOfBounds {
                token_id,
                vocabulary_size,
            } => write!(
                formatter,
                "EOS token {token_id} is out of bounds for vocabulary {vocabulary_size}"
            ),
            Self::LogitCountMismatch { expected, actual } => write!(
                formatter,
                "last-position logits need {expected} values, received {actual}"
            ),
            Self::AllocationFailed { values } => {
                write!(
                    formatter,
                    "cannot allocate generation evidence for {values} values"
                )
            }
        }
    }
}

impl Error for GenerationError {
    fn source(&self) -> Option<&(dyn Error + 'static)> {
        match self {
            Self::Sampling(error) => Some(error),
            Self::Model(error) => Some(error),
            _ => None,
        }
    }
}

impl From<SamplingError> for GenerationError {
    fn from(error: SamplingError) -> Self {
        Self::Sampling(error)
    }
}

impl From<DecoderModelError> for GenerationError {
    fn from(error: DecoderModelError) -> Self {
        Self::Model(error)
    }
}

fn generate_with<F>(
    vocabulary_size: usize,
    max_positions: usize,
    prompt: &[u32],
    config: GenerationConfig,
    rng: &mut SplitMix64,
    mut last_logits: F,
) -> Result<GenerationResult, GenerationError>
where
    F: FnMut(&[u32]) -> Result<Vec<f64>, GenerationError>,
{
    if prompt.is_empty() {
        return Err(GenerationError::EmptyPrompt);
    }
    if prompt.len() > max_positions {
        return Err(GenerationError::PromptTooLong {
            tokens: prompt.len(),
            max_positions,
        });
    }
    for (position, &token_id) in prompt.iter().enumerate() {
        if usize::try_from(token_id)
            .ok()
            .is_none_or(|token| token >= vocabulary_size)
        {
            return Err(GenerationError::PromptTokenOutOfBounds {
                position,
                token_id,
                vocabulary_size,
            });
        }
    }
    if let Some(token_id) = config.eos_token
        && usize::try_from(token_id)
            .ok()
            .is_none_or(|token| token >= vocabulary_size)
    {
        return Err(GenerationError::EosTokenOutOfBounds {
            token_id,
            vocabulary_size,
        });
    }
    retained_count(config.mode, vocabulary_size)?;

    let planned_steps = config.max_new_tokens.min(
        max_positions
            .checked_sub(prompt.len())
            .and_then(|remaining| remaining.checked_add(1))
            .ok_or(GenerationError::AllocationFailed { values: usize::MAX })?,
    );
    let prefix_capacity = prompt
        .len()
        .checked_add(planned_steps)
        .ok_or(GenerationError::AllocationFailed { values: usize::MAX })?;

    let mut prompt_copy = Vec::new();
    prompt_copy
        .try_reserve_exact(prompt.len())
        .map_err(|_| GenerationError::AllocationFailed {
            values: prompt.len(),
        })?;
    prompt_copy.extend_from_slice(prompt);

    let mut prefix = Vec::new();
    prefix
        .try_reserve_exact(prefix_capacity)
        .map_err(|_| GenerationError::AllocationFailed {
            values: prefix_capacity,
        })?;
    prefix.extend_from_slice(prompt);
    let mut generated = Vec::new();
    generated
        .try_reserve_exact(planned_steps)
        .map_err(|_| GenerationError::AllocationFailed {
            values: planned_steps,
        })?;
    let mut steps = Vec::new();
    steps
        .try_reserve_exact(planned_steps)
        .map_err(|_| GenerationError::AllocationFailed {
            values: planned_steps,
        })?;
    let mut full_prefix_calls = 0usize;

    if config.max_new_tokens == 0 {
        return Ok(GenerationResult {
            prompt: prompt_copy,
            generated,
            steps,
            stop: GenerationStop::TokenLimit,
            full_prefix_calls,
        });
    }

    let stop = loop {
        let prefix_length = prefix.len();
        let logits = last_logits(&prefix)?;
        full_prefix_calls += 1;
        if logits.len() != vocabulary_size {
            return Err(GenerationError::LogitCountMismatch {
                expected: vocabulary_size,
                actual: logits.len(),
            });
        }
        let decision = sample_next_token(&logits, config.mode, rng)?;
        let token_id = decision.token_id();
        let unit_draw = decision.unit_draw();
        let interval_start = decision.interval_start();
        let interval_end = decision.interval_end();
        prefix.push(token_id);
        generated.push(token_id);
        steps.push(GenerationStep {
            prefix_length,
            token_id,
            unit_draw,
            interval_start,
            interval_end,
        });

        if config.eos_token == Some(token_id) {
            break GenerationStop::Eos;
        }
        if generated.len() == config.max_new_tokens {
            break GenerationStop::TokenLimit;
        }
        if prefix.len() > max_positions {
            break GenerationStop::ContextLimit;
        }
    };

    Ok(GenerationResult {
        prompt: prompt_copy,
        generated,
        steps,
        stop,
        full_prefix_calls,
    })
}

/// Recomputes the complete decoder prefix for every selected token.
pub fn generate_uncached(
    model: &DecoderModel,
    prompt: &[u32],
    config: GenerationConfig,
    rng: &mut SplitMix64,
) -> Result<GenerationResult, GenerationError> {
    let model_config = model.config();
    let vocabulary_size = model_config.vocabulary_size();
    generate_with(
        vocabulary_size,
        model_config.max_positions(),
        prompt,
        config,
        rng,
        |prefix| {
            let forward = no_grad(|| model.forward(prefix, &[1, prefix.len()]))?;
            let logits = forward.logits().value();
            let expected = prefix
                .len()
                .checked_mul(vocabulary_size)
                .ok_or(GenerationError::AllocationFailed { values: usize::MAX })?;
            if logits.len() != expected {
                return Err(GenerationError::LogitCountMismatch {
                    expected,
                    actual: logits.len(),
                });
            }
            let start = expected - vocabulary_size;
            let mut final_logits = Vec::new();
            final_logits
                .try_reserve_exact(vocabulary_size)
                .map_err(|_| GenerationError::AllocationFailed {
                    values: vocabulary_size,
                })?;
            final_logits.extend_from_slice(&logits.as_slice()[start..]);
            Ok(final_logits)
        },
    )
}

The fixture loads the exact Chapter 35 checkpoint bytes. It first records the saved random state, then consumes the checkpoint and moves its already-owned model buffers into the decoder; no later operation needs the checkpoint object. This ownership transfer does not change any parameter value or sampling rule. Prompt $[0]$ emits $[4,4]$ from prefix lengths $[1,2]$ , performs two complete-prefix calls, and reports a context stop. Repeating from that saved starting state reproduces the same generated sequence and final random state. With token $4$ configured as EOS, the result includes $[4]$ and stops after one call.

Assemble exact temperature, tie, replay, checkpoint, EOS, context, and error evidence rust/demos/ch36-temperature-top-k/src/lib.rs#learner-evidence

/// Loads the Chapter 35 checkpoint and records sampling plus full-prefix stops.
pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
    let temperatures = temperature_evidence()?;
    let boundary = sampling_distribution(
        &LOGITS,
        SamplingMode::TemperatureTopK {
            temperature: 1.0,
            top_k: BOUNDARY_TOP_K,
        },
    )?;
    let mut greedy_rng = SplitMix64::from_seed(SAMPLE_SEED);
    let greedy_state = greedy_rng.state();
    let greedy = sample_next_token_with_trace(&LOGITS, SamplingMode::Greedy, &mut greedy_rng)?;
    require(
        greedy_rng.state() == greedy_state,
        "greedy sampling unexpectedly consumed RNG state",
    )?;
    let seeded_decisions = seeded_decisions()?;

    let prior = checkpoint_evidence()?;
    let loaded_checkpoint_bytes = prior.encoded.bytes().len();
    let checkpoint = Checkpoint::from_bytes(prior.encoded.bytes())?;
    let loaded_rng_state = checkpoint.rng_state();
    let model = checkpoint.into_model()?;
    let model_config = model.config();
    let loaded_vocabulary_size = model_config.vocabulary_size();
    let loaded_context = model_config.max_positions();
    let generation_config = GenerationConfig::new(
        SamplingMode::TemperatureTopK {
            temperature: 1.0,
            top_k: 3,
        },
        None,
        4,
    );
    let mut loaded_rng = SplitMix64::from_state(loaded_rng_state);
    let loaded = generate_uncached(&model, &LOADED_PROMPT, generation_config, &mut loaded_rng)?;
    let mut replay_rng = SplitMix64::from_state(loaded_rng_state);
    let replay = generate_uncached(&model, &LOADED_PROMPT, generation_config, &mut replay_rng)?;
    let loaded_replay_identical = loaded == replay && loaded_rng.state() == replay_rng.state();
    let first_token = *loaded.generated().first().ok_or(FixtureError::Invariant(
        "loaded generation emitted no token",
    ))?;
    let mut eos_rng = SplitMix64::from_state(loaded_rng_state);
    let eos = generate_uncached(
        &model,
        &LOADED_PROMPT,
        GenerationConfig::new(generation_config.mode(), Some(first_token), 4),
        &mut eos_rng,
    )?;
    let errors = error_evidence();
    let history = historical_decoding_contrast()?;

    require(
        boundary.survivors() == [3, 1],
        "stable tied top-k boundary changed",
    )?;
    require(
        seeded_decisions
            .iter()
            .map(SamplingDecision::token_id)
            .eq([3, 2, 2, 2, 3, 3, 3, 3]),
        "seeded sampling sequence changed",
    )?;
    require(
        loaded
            .steps()
            .iter()
            .map(|step| step.prefix_length())
            .eq([1, 2])
            && loaded.full_prefix_calls() == 2
            && loaded.stop() == GenerationStop::ContextLimit,
        "loaded uncached context evidence changed",
    )?;
    require(
        loaded_replay_identical,
        "loaded checkpoint and RNG no longer replay generation",
    )?;
    require(
        eos.generated() == [first_token]
            && eos.stop() == GenerationStop::Eos
            && eos.full_prefix_calls() == 1,
        "EOS stopping evidence changed",
    )?;
    require(
        errors.zero_temperature_rejected
            && errors.zero_top_k_rejected
            && errors.nonfinite_logit_rejected
            && errors.rng_unchanged,
        "invalid settings changed RNG state or escaped validation",
    )?;

    Ok(LearnerEvidence {
        temperatures,
        boundary,
        greedy,
        seeded_decisions,
        loaded_checkpoint_bytes,
        loaded_rng_state,
        loaded_vocabulary_size,
        loaded_context,
        generation_max_new_tokens: generation_config.max_new_tokens(),
        loaded,
        loaded_replay_identical,
        eos,
        errors,
        history,
    })
}

Run cargo run --quiet --locked -p ch36-temperature-top-k. The key evidence is printed directly by the Rust program:

top_k=k:2 survivors:[3,1] tied_boundary:keep:1 remove:2 sum:1.000000000000
sample=seed:36 top_k:3 sequence:[3,2,2,2,3,3,3,3] draws:8 greedy_token:3 greedy_draw:none
checkpoint=loaded_bytes:6330 rng_state:0x9e3779b97f4a7c38 vocabulary:5 context:2 eos:none max_new_tokens:4 prompt:[0] generated:[4,4] prefixes:[1,2] stop:context-limit full_prefix_calls:2 replay_identical:true
eos=vocabulary:5 context:2 eos_token:4 max_new_tokens:4 generated:[4] stop:eos full_prefix_calls:1
errors=temperature_zero:true top_k_zero:true nonfinite_logit:true rng_unchanged:true
history=greedy_token:3 greedy_rng_advanced:false top_k:3 survivors:[3,1,2] retained_full_mass:0.927670511871 removed_full_mass:0.072329488129

Print the exact Chapter 36 learner report rust/demos/ch36-temperature-top-k/src/main.rs

fn main() -> Result<(), Box<dyn std::error::Error>> {
    print!("{}", ch36_temperature_top_k::learner_report()?);
    Ok(())
}

Read the distribution from left to right

The figure begins with three aligned views of the same four-token logits. It then isolates the $k=2$ tie and explicitly switches to the $\tau=1$ , $k=3$ , seed- $36$ draw policy before following each recorded interval. A labeled fixture boundary then changes from synthetic vocabulary size $V=4$ to the loaded decoder’s $V=5$ , with no EOS configured for the context-stop run and token $4$ configured as EOS for the EOS run. Double and dashed borders repeat retained and removed states in text. The table preserves each token’s exact logit, stable rank, probability, and retained state so the border bars can be checked against numeric evidence.

Temperature reshapes; top-k removes; one draw selects

The exact Rust trace compares three temperatures, exposes a stable tied boundary, follows eight seeded half-open intervals, and records checkpoint generation stops.

retained — double border
removed — dashed border
selected by the draw

Compare temperature on the same logits

Synthetic four-token fixture $V=4$

All four tokens remain; only the probability ratios change. Each bar carries its exact value as text.

Sharper

$\tau=0.500000$

$i=0$ $q_i=0.014209336619$
$i=1$ $q_i=0.104993585404$
$i=2$ $q_i=0.104993585404$
$i=3$ $q_i=0.775803492574$

Original scale

$\tau=1.000000$

$i=0$ $q_i=0.072329488129$
$i=1$ $q_i=0.196611933241$
$i=2$ $q_i=0.196611933241$
$i=3$ $q_i=0.534446645389$

Flatter

$\tau=2.000000$

$i=0$ $q_i=0.142536956597$
$i=1$ $q_i=0.235003712202$
$i=2$ $q_i=0.235003712202$
$i=3$ $q_i=0.387455619000$

Keep exactly two stable ranks

At equal logits, ascending token ID breaks the boundary tie: token 1 remains and token 2 becomes exactly zero.

Stable top-k candidate table
Token	Logit	Stable rank	Top-k status	Final probability
$i=0$	$\ell_i=0.000000$	$r_i=4$	removed — dashed border	$q_i=0.000000000000$
$i=1$	$\ell_i=1.000000$	$r_i=2$	retained — double border	$q_i=0.268941421370$
$i=2$	$\ell_i=1.000000$	$r_i=3$	removed — dashed border	$q_i=0.000000000000$
$i=3$	$\ell_i=2.000000$	$r_i=1$	retained — double border	$q_i=0.731058578630$

Replay the seeded categorical draws

This section changes the illustrated candidate count from two retained candidates to the three-candidate policy with temperature 1 and seed 36. Sampling visits retained intervals in ascending token-ID order.

Synthetic four-token fixture

\tau=1.000000,\ k=3,\ V=4

seed=36 survivors=[3,1,2] sum=1.000000000000

Draw 1 $u=0.912888894097$ Half-open interval $[0.423883115234,1.000000000000)$ selects token $i=3$
Draw 2 $u=0.338833394523$ Half-open interval $[0.211941557617,0.423883115234)$ selects token $i=2$
Draw 3 $u=0.295371378932$ Half-open interval $[0.211941557617,0.423883115234)$ selects token $i=2$
Draw 4 $u=0.350092047261$ Half-open interval $[0.211941557617,0.423883115234)$ selects token $i=2$
Draw 5 $u=0.578054529784$ Half-open interval $[0.423883115234,1.000000000000)$ selects token $i=3$
Draw 6 $u=0.660097051275$ Half-open interval $[0.423883115234,1.000000000000)$ selects token $i=3$
Draw 7 $u=0.836130632904$ Half-open interval $[0.423883115234,1.000000000000)$ selects token $i=3$
Draw 8 $u=0.657589579642$ Half-open interval $[0.423883115234,1.000000000000)$ selects token $i=3$

Carry the policy into uncached generation

This fixture boundary changes from four synthetic token IDs to the loaded five-token decoder. The exact EOS and context policies remain visible beside each run.

Loaded five-token decoder fixture $V=5$ context=2

Explicit greedy mode

$i=3$ — highest logit, lowest ID on a tie

draw=none

Greedy leaves the random stream untouched.

Loaded checkpoint

$[0]\to[4,4]$

eos=none max_new_tokens=4 prefixes=[1,2] calls=2

stops at context capacity; the restored stream replays exactly

EOS boundary

$i_{\mathrm{EOS}}=4$

eos=4 max_new_tokens=4 generated=[4] calls=1

EOS remains in the emitted sequence

Transactional errors

temperature_zero=true top_k_zero=true nonfinite_logit=true rng_unchanged=true

Invalid settings return before the random stream advances.

Predict before revealing the trace

Which ID wins greedy for logits $[0,1,1,2]$ ?
Which equal-logit ID survives the $k=2$ boundary?
Does raising $\tau$ change stable rank?
Why is literal $\tau=0$ rejected even though $\tau\to0^+$ is informative?
Does stochastic $k=1$ consume a random draw?
Which token owns draw $u=0.338833394523$ in the $k=3$ fixture?
Is EOS included in the emitted sequence?
Why can a capacity- $2$ prefix emit one token before a context stop?

Check the eight predictions

Token $3$ wins because its logit $2$ is the unique maximum.
Token $1$ survives because equal logits use ascending token ID, leaving token $2$ at exact zero.
No. Positive temperature changes probability ratios but preserves logit order.
The limit describes concentration, while literal zero would require division by zero; explicit greedy owns that policy.
Yes. It selects the greedy ID but remains a stochastic policy with exactly one draw.
Token $2$ owns $[0.211941557617,0.423883115234)$ , which contains the draw.
Yes. Token $4$ is appended before the EOS stop is reported.
The valid length- $2$ prefix predicts a next token; only the following decoder call would exceed capacity.

Preserve this uncached sequence when generation becomes incremental

The cumulative decoder can now load its selected checkpoint, turn each final-position logit row into a controlled next-token distribution, replay choices from a restored random-generator state, stop at EOS or context capacity, and expose the uncached reference sequence that Chapter 37 will preserve incrementally.

Chapter 37 will cache one attention layer’s earlier key and value vectors. Its newest-position output must match the complete-prefix computation here before Chapter 38 extends caching across the whole decoder.