The generation loop of a decoder with a KV cache. Transmuter.Text.generation/2, Transmuter.Text.text_to_text/2, Transmuter.Vision.image_to_text/2, Transmuter.Vision.image_text_to_text/2 and Transmuter.Audio.speech_to_text/2 use it.
Generation runs the model in steps. A step takes new tokens and the KV cache, and gives the logits and the next cache. The prompt runs in steps of power-of-two lengths (45 tokens run as 32, 8, 4 and 1), and each new token is a step of length 1, so a small number of compiled shapes covers all prompts.
The cache works in one of three modes, from the structure of the graph:
:in_place- the graph has GroupQueryAttention, which counts the valid rows of the cache from the values of the attention mask. The cache is a buffer as long as the attention mask, and the operator writes the new keys and values after the valid rows.:append- the graph declaresposition_ids. The cache is a buffer whose free rows the attention mask hides. The present cache is the buffer and the new rows after it, and the new rows go into the buffer.:grow- the graph computes the positions from the length of the cache. It has noposition_ids, or its GroupQueryAttention counts the valid rows from the length of the attention mask, as in Gemma 3. The cache has its exact length, so each length is a new shape that compiles on its first step. The first step takes the branch without cache of a merged decoder (use_cache_branch), because Nx has no tensors of length 0.
The decoder of an encoder-decoder model also takes the outputs of the encoder, such as encoder_hidden_states, in each step. Its first step computes the cross-attention cache (past_key_values.N.encoder.key and .value) from these outputs. Later steps give an empty present for this cache, so the cache keeps its value.
The token selection applies the logits processors of Python transformers that generation_config.json sets, in the order of _get_logits_processor: repetition_penalty, no_repeat_ngram_size, the bad words of one token in bad_words_ids, min_length or min_new_tokens, forced_bos_token_id, forced_eos_token_id, suppress_tokens, begin_suppress_tokens, the processors of the task, renormalize_logits, temperature, top_k and top_p.
Lengths
The processors count the tokens of each row as Python transformers counts the columns of input_ids: the prompt and the new tokens. The prompt of a decoder-only model is the text, and the prompt of the decoder of an encoder-decoder model starts with decoder_start_token_id. Each row has its own length, and the rows of a batch are not padded. Transmuter limits the new tokens with max_new_tokens, so the max_length of a row in Python is the length of its prompt plus max_new_tokens. Thus:
no_repeat_ngram_size- the n-grams of the prompt and the new tokens can occur only once.min_length- the eos tokens are removed while the prompt and the new tokens are fewer thanmin_length.min_new_tokens- the eos tokens are removed while the new tokens are fewer thanmin_new_tokens. As in Python,min_new_tokensreplacesmin_lengthwhen both are set.forced_bos_token_id- the token when the row has one token, which is the first new token of an encoder-decoder model. A decoder-only prompt of more than one token gets no forced token, as in Python.forced_eos_token_id- the token when the row hasmax_length - 1tokens, which is the last new token.begin_suppress_tokens- removed at the first new token, or at the second new token after a forced bos token of a prompt of one token.
Embeddings and image features
The decoder of a vision language model, such as SmolVLM or Qwen2-VL, declares inputs_embeds in place of input_ids. Its context has the model of the embeddings in :embed (embed_tokens.onnx) and the id of the image token in :image_token. Each step turns its tokens into embeddings with that model, so the steps are the same as for input_ids.
The task gives the features of the images to generate/5 in :image_features. They are the features of the image tokens of all rows, in order, as masked_scatter in Python transformers uses them. A row starts at the features after those of the rows before it. In the steps of the prompt, the features replace the embeddings of the image tokens, and a cursor of each row counts the image tokens that the row used. Parts of the prompt without image tokens and the steps of new tokens only embed their tokens.
Positions
Each token of a row has the position of its index in the row. The task can give other positions for the prompt of each row in :position_ids of generate/5, such as the 3-D positions of M-RoPE in Qwen2-VL. There, the image tokens of a prompt have the positions of their rows and columns in the image. A decoder whose position_ids has 3 axes takes 3 positions for each token. The tokens after the prompt continue from the largest position of the prompt: the token at index i has the position i + delta on all axes. delta is the largest position of the prompt, plus 1, minus the length of the prompt, as the rope delta of Python transformers.
Processors
The :processors key of the context is the list of processors, which apply in order after the repetition penalty and before the temperature. context/4 puts the processors of generation_config.json in it, and a task can add its own processors at the end. A processor is {state, process}. Before each step, state takes a map of one row and gives a tensor of the same shape at each step. The map has the token ids of the prompt in :prompt, the new token ids in :generated (the last token first), and the size of the vocabulary in :vocab. process takes the logits and that tensor of one row in the compiled selection, and gives the new logits.
The atoms :bad_words and :suppress_tokens in the list are the places of two masks that do not change between steps: the bad words of one token and suppress_tokens. Generation makes these masks once for each group of rows.
Log-probabilities
With token_logprobs: true in the options of context/4, each step also gives the log-probability of the token of each row: the log-softmax of the scores that the step selects from, at the token. This is the rule of compute_transition_scores of Python transformers with normalize_logits=True, for greedy selection, sampling and beam search (see generate/5). The compiled selection subtracts the log-sum-exp of the scores of the row from the score of the token. Without the option, the selection does not compute it.
The compiled selection gives one s32 output, because EMLX evaluates each output of a compiled function separately. The tokens and the indices of the candidates of beam search are integers in it, and the scores and log-probabilities are the bits of their f32 values (Nx.bitcast/2, see from_bits/1).
Beam search
With num_beams above 1, generation does the beam search of _beam_search in Python transformers 4.57. The beams of a prompt are rows of the batch: prompt i has the rows i * num_beams to i * num_beams + num_beams - 1, and its inputs, image features and positions repeat for each row. compile/3 compiles the decoder for batch_size * num_beams rows. At the start, all beams of a prompt are equal, and only the first beam counts.
Each step takes the log-softmax of the logits, applies the processors to it, and adds the score of the beam. Of the num_beams * vocabulary candidates of a prompt, the compiled selection keeps the best max(2, 1 + eos tokens) * num_beams. On the host:
- A candidate ends when its token is an eos token or when it is the last new token. The candidates that end among the first
num_beamsbecome finished hypotheses. The score of a hypothesis is its sum divided by its length to the power oflength_penalty, and its length counts the eos token. A prompt keeps its bestnum_beamshypotheses. - The best
num_beamscandidates that do not end are the next beams. Each row takes the cache and the tokens of the row of its beam. The cache entries that the step did not write, such as the cross-attention cache, are the same for all beams of a prompt, so they do not move. - A prompt stops when it cannot improve: it has
num_beamshypotheses, and its best beam divided by the best length to the power oflength_penaltyscores less than the worst hypothesis. The best length is the length now, or the longest length forearly_stopping: :neverwith a positivelength_penalty. Withearly_stopping: true, a prompt stops when it hasnum_beamshypotheses. A stopped prompt has padding rows until all prompts stop.
The result of a prompt is its best hypothesis. The scores are f32 as in transformers. Removed tokens have the smallest finite value in place of an infinite one, and a score never goes below it.
Batches
generate/4 takes a batch of prompts. Each row has its own length, sampling key and stop. A step gives the same number of tokens to each row. A row with fewer tokens gets padding tokens after its tokens, and a finished row gets only padding tokens. The next token of a row comes from the logits of its last real token. An empty prompt is a padding row, which gives no tokens.
In :in_place and :append, all rows share one buffer. Each row writes its keys and values after its own valid rows, so the padding tokens of a step go into rows that the next step overwrites. In :in_place, the attention mask of a row has ones for its valid rows and the new tokens, and GroupQueryAttention gets the write position of each row from the sum of the mask. In :append, the attention mask hides the free rows of each row, and the new rows go into the buffer at the length of each row. Both modes give the positions of each row in position_ids.
In :grow, the positions come from the length of the cache, which is the same for all rows. So only prompts of the same length share a cache. The prompts of each length run as one group, and the groups run one after the other. A group runs with all rows of the batch, and the rows of other lengths are padding rows of the group. So each length of the cache has one compiled shape, and a batch with prompts of three lengths runs the steps of three groups. Text-to-text, image-to-text and speech-to-text have prompts of one length, so all their rows run as one group.
Device memory
A step makes large device tensors that die at the next step: the logits of all positions of a part of the prompt, the present cache, and in :append the old buffers. The garbage collector of the BEAM does not see the size of a device buffer, so it frees them late, and a request of 24 tokens from a prompt with an image kept 13 GB of dead buffers on the EXLA CPU. So each step frees the tensors that it owns with Nx.backend_deallocate/1: the logits of the last step after their token is selected, the present after it is in the cache, and the old buffers of :append and of the reorder of beam search. A present that shares a buffer with a kept tensor stays, as when a merged decoder gives a past back. Generation never frees the tensors of the caller, such as the outputs of an encoder or the image features.
Summary
Functions
Compiles the model and the token selection of ctx for batches of batch_size prompts. With beam search, the decoder runs batch_size * num_beams rows. Call this function in the serving process, which owns the cache of compiled shapes.
Returns the context of generation for model: the cache mode and the settings from generation_config.json and opts. given are the names of other inputs that the caller gives to generate/5, such as the outputs of an encoder.
Returns the f32 value of bits, the signed 32-bit integer of its bits, as Nx.to_list/1 gives an f32: a float, :nan, :infinity or :neg_infinity. A list gives the value of each element. The compiled selection gives its scores and log-probabilities in this form (see the module documentation).
Returns the :max_new_tokens of a decoder whose max_length counts its prompt, such as the decoder of an encoder-decoder model. The rule is that of the pipelines in Python transformers 4
Returns the outputs of a decoder that generation reads: logits and the present cache. names are the names of the graph outputs. Give this function to the :outputs option of Transmuter.load!/2, so the decoder computes no other outputs, such as the attention probabilities of the Whisper exports of Xenova.
Returns the inputs of a decoder, other than input_ids, for one step over batch_size rows of length tokens with an empty cache, and the static inputs of that step. A task that runs a decoder once over padded text, such as an embedding, gives these inputs with the tokens. A model without a cache has none, so the result is two empty maps.
Returns the signed 32-bit integer of the bits of value in f32, the inverse of from_bits/1. NaN is the quiet NaN with the sign bit 0.
Functions
Compiles the model and the token selection of ctx for batches of batch_size prompts. With beam search, the decoder runs batch_size * num_beams rows. Call this function in the serving process, which owns the cache of compiled shapes.
Returns the context of generation for model: the cache mode and the settings from generation_config.json and opts. given are the names of other inputs that the caller gives to generate/5, such as the outputs of an encoder.
A decoder that declares inputs_embeds needs two more options. :embed is the model that turns token ids into embeddings, such as embed_tokens.onnx. :image_token is the id of the image token, whose embeddings the image features replace. See the module documentation.
Returns the f32 value of bits, the signed 32-bit integer of its bits, as Nx.to_list/1 gives an f32: a float, :nan, :infinity or :neg_infinity. A list gives the value of each element. The compiled selection gives its scores and log-probabilities in this form (see the module documentation).
Generates tokens after each prompt of prompts with a context from compile/3. A prompt is a list of token ids, and an empty prompt is a padding row. seeds has the seed of sampling of each row. inputs are the inputs of each step that context/4 names in given, with one row for each prompt.
inputs can also have the key :image_features, a tensor of shape {count, hidden} with the features of the image tokens of all prompts: the features of the first prompt, in the order of its image tokens, then those of the next prompt. It can have more rows than the prompts have image tokens. It is not an input of the decoder. See the module documentation.
inputs can also have the key :position_ids, a list with the positions of each prompt. The positions of a prompt are a list with the positions of each axis of the decoder's position_ids, one for each token. An empty list gives the default positions. See the module documentation.
on_token is nil or a function that receives %{token: tensor, finished?: tensor} after each step of a group, with one row for each prompt. :token is the new token of the row, or -1 when the row has no new token. :finished? is 1 when the row has no more tokens. With token_logprobs: true in the options of context/4, the map also has :logprob, the log-probability of the new token, or NaN when the row has no new token.
Returns the new tokens of each row, without the eos token. A token is the list of its values: [id], or [id, logprob] with token_logprobs: true in the options of context/4. The log-probability of a token is the log-softmax of the scores that the step selects from, at the token. These scores are the logits after the processors, also after the temperature, top-k and top-p of sampling. For beam search, they are the scores of the step before the sum with the score of the beam. This is compute_transition_scores of Python transformers with normalize_logits=True.
Returns the :max_new_tokens of a decoder whose max_length counts its prompt, such as the decoder of an encoder-decoder model. The rule is that of the pipelines in Python transformers 4:
:max_new_tokensofopts, elsemax_new_tokensofgeneration_config.json.- When
max_lengthis not in the configuration or is20, the default ofGenerationConfig, the value is 256. - Else
max_lengthminusprompt_length, such as 448 minus the prompt for Whisper and 511 for Marian.
Returns the outputs of a decoder that generation reads: logits and the present cache. names are the names of the graph outputs. Give this function to the :outputs option of Transmuter.load!/2, so the decoder computes no other outputs, such as the attention probabilities of the Whisper exports of Xenova.
Returns the inputs of a decoder, other than input_ids, for one step over batch_size rows of length tokens with an empty cache, and the static inputs of that step. A task that runs a decoder once over padded text, such as an embedding, gives these inputs with the tokens. A model without a cache has none, so the result is two empty maps.
The cache has length 0 and use_cache_branch is 0, as in the first step of transformers.js. Each row has the positions from 0. The attention mask counts all tokens, also the padding tokens after the text, as in generation. The decoder is causal, so the tokens of the text do not see the padding tokens.
Returns the signed 32-bit integer of the bits of value in f32, the inverse of from_bits/1. NaN is the quiet NaN with the sign bit 0.