How does the Tunix sampler actually generate text?

In two compiled functions over one fixed-size buffer. The public entry point tokenizes every prompt, left-pads them all to the same length, allocates a token buffer of shape (batch, total_sampling_steps) filled with the pad id, and allocates a KV cache of shape (batch, cache_size, num_kv_heads, head_dim). Then it calls a jitted prefill that runs the whole prompt through the transformer in one pass, writes the prompt keys and values into the cache, and samples the first new token. Then it calls a jitted decode function whose body is a jax.lax.while_loop: each iteration reads the single last token, builds an attention mask for exactly that position, calls the transformer once, samples one token, and writes it into the buffer at decoding_step + 1. The loop condition is decoding_step < total_sampling_steps - 1 and jnp.any(not done), so it stops when the step budget runs out or every sequence in the batch has emitted an end-of-sequence token, whichever comes first.

Why does the Tunix sampler pad prompts to a power of two?

To keep the number of XLA compilations small. Both the prefill and the decode function are jitted, and a jitted function retraces and recompiles whenever an input shape changes. The prompt length is an input shape, so a batch of 41 tokens and a batch of 42 tokens would be two separate compiles. The sampler defends against that by rounding: if the caller did not pass max_prompt_length, or passed one smaller than the longest prompt in the batch, it sets max_prompt_length to utils.next_power_of_2 of the longest prompt. Every prompt is then left-padded to that length. There are only about a dozen powers of two in any realistic prompt range, so a long training run settles into a handful of compiled shapes instead of one per distinct prompt length. The cost is padding: a 33-token prompt is padded to 64, and the model does attention work over 31 pad positions the mask then discards.

What is total_sampling_steps, and why must it fit in the KV cache?

total_sampling_steps is max_prompt_length + max_generation_steps, computed once in __call__ before anything is allocated. It is the width of the token buffer, so the buffer holds the padded prompt and every generated token in one array with no resizing. Immediately after computing it the sampler checks it against cache_config.cache_size and raises a ValueError if it is larger. The reason is that the KV cache has a fixed second dimension of cache_size slots, one per position, and it is allocated before generation starts. Prefill writes the prompt positions into it and every decode step writes one more, so a run that needed more than cache_size positions would write past the end of an array that JAX will not grow. The check turns that into a loud error at the top of the call rather than a wrong answer or a shape failure deep inside a jitted while_loop. The surprising consequence is that the power-of-two rounding happens first, so a 600-token prompt with a 1024-slot cache fails at 128 generation steps: 600 rounds up to 1024, and 1024 + 128 is 1152.

Why does the sampler split the model into a graphdef and a state?

Because a Flax NNX module closed over by a jitted function would become a static argument, and every weight update would change it and force a retrace. So the constructor calls nnx.graphdef on the transformer to get its structure, nnx.variables to get its parameters, and keeps them apart. The structure is captured; the flattened parameter list is passed as the first runtime argument to both compiled functions, which re-merge them into a live module inside the trace. The comment in the file says the reason directly: state is passed as an argument so it is not treated as a static arg, which greatly reduces the size of the HLO and reduces compile time. It is also what makes weight sync cheap during RL training. The rollout worker assigns a fresh parameter tree through the transformer_state setter, which checks shape, dtype and sharding against the old tree and swaps the leaves; no reconstruction and no recompile, because nothing the jit keys on has changed.

What is donate_argnums doing on the sampler jits?

It lets XLA reuse the KV cache buffer in place instead of copying it every decode step. Both compiled functions are created with donate_argnums=(1,), which marks the sampling state, the argument that carries the cache, as donated. JAX arrays are immutable, so without donation an update to the cache at each step would allocate a fresh buffer and copy the old contents into it. The comment in the file spells out why that matters here specifically: the cache footprint scales with batch size and with prompt plus decoding length, reaching gigabytes, so the repeated allocation and copy causes heavy memory overhead and out-of-memory failures. Donating tells the compiler the caller will not use the input buffer again, so it can write the new keys and values into the same memory. The consequence for a caller is the usual donation rule: the sampling state you passed in is invalid after the call.

How do greedy, top-p and beam search differ inside the sampler?

They are three branches of one internal method, and which branch runs is decided once, before generation, by which keyword arguments you passed. init_sample_state builds the mode from a small precedence rule: if beam_size is not None the mode is beam_search, if top_p is not None the mode is top_p, and if neither was given the mode is greedy. Passing both raises a conflict error that names the current mode and the rule. Greedy is an argmax over the last position, plus a log-softmax when logprobs were requested. Top-p divides the logits by the temperature, takes the top k, softmaxes, takes a cumulative sum, masks everything past the p threshold to negative infinity, and draws from the survivors with a key folded from the seed and the decoding step, which makes each step independently random but the whole run reproducible from one seed. A top_p of 1.0 with no top-k skips the sort entirely and samples from the raw distribution. Beam search is the odd one out: it does not just pick a token, it also permutes the token buffer, the cache and the done flags, because the surviving beams may have come from different rows.

The Tunix Sampler, Line by Line: Generation as Two Compiled Functions

Q: Why does the Tunix sampler pad prompts to a power of two?

To keep the number of XLA compilations small. Both the prefill and the decode function are jitted, and a jitted function retraces and recompiles whenever an input shape changes. The prompt length is an input shape, so a batch of 41 tokens and a batch of 42 tokens would be two separate compiles. The sampler defends against that by rounding: if the caller did not pass max_prompt_length, or passed one smaller than the longest prompt in the batch, it sets max_prompt_length to utils.next_power_of_2 of the longest prompt. Every prompt is then left-padded to that length. There are only about a dozen powers of two in any realistic prompt range, so a long training run settles into a handful of compiled shapes instead of one per distinct prompt length. The cost is padding: a 33-token prompt is padded to 64, and the model does attention work over 31 pad positions the mask then discards.

Q: What is total_sampling_steps, and why must it fit in the KV cache?

total_sampling_steps is max_prompt_length + max_generation_steps, computed once in __call__ before anything is allocated. It is the width of the token buffer, so the buffer holds the padded prompt and every generated token in one array with no resizing. Immediately after computing it the sampler checks it against cache_config.cache_size and raises a ValueError if it is larger. The reason is that the KV cache has a fixed second dimension of cache_size slots, one per position, and it is allocated before generation starts. Prefill writes the prompt positions into it and every decode step writes one more, so a run that needed more than cache_size positions would write past the end of an array that JAX will not grow. The check turns that into a loud error at the top of the call rather than a wrong answer or a shape failure deep inside a jitted while_loop. The surprising consequence is that the power-of-two rounding happens first, so a 600-token prompt with a 1024-slot cache fails at 128 generation steps: 600 rounds up to 1024, and 1024 + 128 is 1152.

Q: Why does the sampler split the model into a graphdef and a state?

Because a Flax NNX module closed over by a jitted function would become a static argument, and every weight update would change it and force a retrace. So the constructor calls nnx.graphdef on the transformer to get its structure, nnx.variables to get its parameters, and keeps them apart. The structure is captured; the flattened parameter list is passed as the first runtime argument to both compiled functions, which re-merge them into a live module inside the trace. The comment in the file says the reason directly: state is passed as an argument so it is not treated as a static arg, which greatly reduces the size of the HLO and reduces compile time. It is also what makes weight sync cheap during RL training. The rollout worker assigns a fresh parameter tree through the transformer_state setter, which checks shape, dtype and sharding against the old tree and swaps the leaves; no reconstruction and no recompile, because nothing the jit keys on has changed.

Q: What is donate_argnums doing on the sampler jits?

It lets XLA reuse the KV cache buffer in place instead of copying it every decode step. Both compiled functions are created with donate_argnums=(1,), which marks the sampling state, the argument that carries the cache, as donated. JAX arrays are immutable, so without donation an update to the cache at each step would allocate a fresh buffer and copy the old contents into it. The comment in the file spells out why that matters here specifically: the cache footprint scales with batch size and with prompt plus decoding length, reaching gigabytes, so the repeated allocation and copy causes heavy memory overhead and out-of-memory failures. Donating tells the compiler the caller will not use the input buffer again, so it can write the new keys and values into the same memory. The consequence for a caller is the usual donation rule: the sampling state you passed in is invalid after the call.

Q: How do greedy, top-p and beam search differ inside the sampler?

They are three branches of one internal method, and which branch runs is decided once, before generation, by which keyword arguments you passed. init_sample_state builds the mode from a small precedence rule: if beam_size is not None the mode is beam_search, if top_p is not None the mode is top_p, and if neither was given the mode is greedy. Passing both raises a conflict error that names the current mode and the rule. Greedy is an argmax over the last position, plus a log-softmax when logprobs were requested. Top-p divides the logits by the temperature, takes the top k, softmaxes, takes a cumulative sum, masks everything past the p threshold to negative infinity, and draws from the survivors with a key folded from the seed and the decoding step, which makes each step independently random but the whole run reproducible from one seed. A top_p of 1.0 with no top-k skips the sort entirely and samples from the raw distribution. Beam search is the odd one out: it does not just pick a token, it also permutes the token buffer, the cache and the done flags, because the surviving beams may have come from different rows.

Every RL post-training step needs completions, and in Tunix the in-tree sampler is what produces them. It is one file, and its whole design falls out of one constraint: JAX wants fixed shapes. So the prompt gets left-padded to a power of two, the output lives in a buffer sized before the first token is generated, the KV cache is allocated once and donated to the compiler, and the decode loop is a lax.while_loop that runs until the slowest sequence in the batch stops. We read tunix/generate/sampler.py top to bottom at one pinned commit, then compute the decode timeline and the cache filling for a batch you dial yourself.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of The Tunix Sampler, Line by Line: Generation as Two Compiled Functions — built to be understood, not skimmed.

Questions

How does the Tunix sampler actually generate text?: In two compiled functions over one fixed-size buffer. The public entry point tokenizes every prompt, left-pads them all to the same length, allocates a token buffer of shape (batch, total_sampling_steps) filled with the pad id, and allocates a KV cache of shape (batch, cache_size, num_kv_heads, head_dim). Then it calls a jitted prefill that runs the whole prompt through the transformer in one pass, writes the prompt keys and values into the cache, and samples the first new token. Then it calls a jitted decode function whose body is a jax.lax.while_loop: each iteration reads the single last token, builds an attention mask for exactly that position, calls the transformer once, samples one token, and writes it into the buffer at decoding_step + 1. The loop condition is decoding_step < total_sampling_steps - 1 and jnp.any(not done), so it stops when the step budget runs out or every sequence in the batch has emitted an end-of-sequence token, whichever comes first.
Why does the Tunix sampler pad prompts to a power of two?: To keep the number of XLA compilations small. Both the prefill and the decode function are jitted, and a jitted function retraces and recompiles whenever an input shape changes. The prompt length is an input shape, so a batch of 41 tokens and a batch of 42 tokens would be two separate compiles. The sampler defends against that by rounding: if the caller did not pass max_prompt_length, or passed one smaller than the longest prompt in the batch, it sets max_prompt_length to utils.next_power_of_2 of the longest prompt. Every prompt is then left-padded to that length. There are only about a dozen powers of two in any realistic prompt range, so a long training run settles into a handful of compiled shapes instead of one per distinct prompt length. The cost is padding: a 33-token prompt is padded to 64, and the model does attention work over 31 pad positions the mask then discards.
What is total_sampling_steps, and why must it fit in the KV cache?: total_sampling_steps is max_prompt_length + max_generation_steps, computed once in __call__ before anything is allocated. It is the width of the token buffer, so the buffer holds the padded prompt and every generated token in one array with no resizing. Immediately after computing it the sampler checks it against cache_config.cache_size and raises a ValueError if it is larger. The reason is that the KV cache has a fixed second dimension of cache_size slots, one per position, and it is allocated before generation starts. Prefill writes the prompt positions into it and every decode step writes one more, so a run that needed more than cache_size positions would write past the end of an array that JAX will not grow. The check turns that into a loud error at the top of the call rather than a wrong answer or a shape failure deep inside a jitted while_loop. The surprising consequence is that the power-of-two rounding happens first, so a 600-token prompt with a 1024-slot cache fails at 128 generation steps: 600 rounds up to 1024, and 1024 + 128 is 1152.
Why does the sampler split the model into a graphdef and a state?: Because a Flax NNX module closed over by a jitted function would become a static argument, and every weight update would change it and force a retrace. So the constructor calls nnx.graphdef on the transformer to get its structure, nnx.variables to get its parameters, and keeps them apart. The structure is captured; the flattened parameter list is passed as the first runtime argument to both compiled functions, which re-merge them into a live module inside the trace. The comment in the file says the reason directly: state is passed as an argument so it is not treated as a static arg, which greatly reduces the size of the HLO and reduces compile time. It is also what makes weight sync cheap during RL training. The rollout worker assigns a fresh parameter tree through the transformer_state setter, which checks shape, dtype and sharding against the old tree and swaps the leaves; no reconstruction and no recompile, because nothing the jit keys on has changed.
What is donate_argnums doing on the sampler jits?: It lets XLA reuse the KV cache buffer in place instead of copying it every decode step. Both compiled functions are created with donate_argnums=(1,), which marks the sampling state, the argument that carries the cache, as donated. JAX arrays are immutable, so without donation an update to the cache at each step would allocate a fresh buffer and copy the old contents into it. The comment in the file spells out why that matters here specifically: the cache footprint scales with batch size and with prompt plus decoding length, reaching gigabytes, so the repeated allocation and copy causes heavy memory overhead and out-of-memory failures. Donating tells the compiler the caller will not use the input buffer again, so it can write the new keys and values into the same memory. The consequence for a caller is the usual donation rule: the sampling state you passed in is invalid after the call.
How do greedy, top-p and beam search differ inside the sampler?: They are three branches of one internal method, and which branch runs is decided once, before generation, by which keyword arguments you passed. init_sample_state builds the mode from a small precedence rule: if beam_size is not None the mode is beam_search, if top_p is not None the mode is top_p, and if neither was given the mode is greedy. Passing both raises a conflict error that names the current mode and the rule. Greedy is an argmax over the last position, plus a log-softmax when logprobs were requested. Top-p divides the logits by the temperature, takes the top k, softmaxes, takes a cumulative sum, masks everything past the p threshold to negative infinity, and draws from the survivors with a key folded from the seed and the decoding step, which makes each step independently random but the whole run reproducible from one seed. A top_p of 1.0 with no top-k skips the sort entirely and samples from the raw distribution. Beam search is the odd one out: it does not just pick a token, it also permutes the token buffer, the cache and the done flags, because the surviving beams may have come from different rows.

The Tunix Sampler, Line by Line: Generation as Two Compiled Functions

Questions

Related explainers