s04Inference core16 min read

Choose the Next Token

Turn logits into a token choice

Trace model output through sampling into request state that can be committed.

Sampling results and request-state updates are separate stages linked by metadata.

01Model output is not yet request output

A model forward returns scores and associated state. Sampling combines request parameters, constraints, and draft information into choices or acceptance results; the request layer then determines what can be committed. Trace sampler_inputs from sequence metadata in ModelRunner.execute_sample before entering the model’s sample path.

s04_sampling.py
# Framework pseudocode for the sampling boundary.
sampler_inputs = prepare_sampler_inputs(
    sequence_metadata, input_metadata,
)
sampler_output = model.sample(
    model_output=main_output,
    sampler_inputs=sampler_inputs,
)
step_state = StepState(sequence_metadata, sampler_output)

# Sampling can choose tokens or verify draft candidates.
# Post-processing commits accepted positions and stop states.
ready_outputs = post_process_when_ready(step_state)
# Under ZOS, readiness can cross a round boundary.

02Selection policies must retain request identity

Greedy choice, temperature, top-p, and constrained output affect different selection stages. Batching must preserve each request’s parameters and random state; draft verification also links candidate positions to accepted lengths. This excerpt describes the framework boundary rather than treating argmax or a random draw as the full upstream sampler.

03Committing state waits for its dependencies

Post-processing commits sampling results to sequences, including accepted positions, stop flags, and public outputs. With ZOS, it may run after the next round’s forward is launched. Distinguish produced samples, committed state, and published outputs to explain why a round can return results from earlier sampling.

04Structured output carries grammar state

Official production design: The official deployment article places structured output inside generation: the confirmed prefix determines the grammar state, which determines the legal next-token set. Sampling must respect that set. This constrains syntax and structure; factual correctness of field values still depends on the task. Public-code check: The pinned sampler.py advances each request’s matcher with accepted tokens, temporarily advances through draft positions to build masks, then rolls those temporary advances back. Masks apply before sampling and draft verification. Grammar state follows the same rule as other request state: only the confirmed prefix persists.

TRY IT YOURSELF

Make it yours

For a request with a stop condition, draw dependencies between model_output, sampler_output, sequence state, and RequestOutput.

Add this mechanism to your toolkit.