For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python class

GenerateIterator

GenerateIterator

class max.pipelines.lib.generate.GenerateIterator(pipeline, input_ids, attention_mask=None, *, max_new_tokens, min_new_tokens=None, temperature=None, top_k=None, top_p=None, min_p=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, ignore_eos=False, stop_token_ids=None, seed=0, output_logits=False, output_logprobs=0, caller='generate()')

source

Bases: object

Steps a fixed batch of prompts through a pipeline, one execute at a time.

GenerateIterator runs a pipeline’s forward passes over exactly the batch you pass, with nothing added on top of the pipeline API, which also makes it useful for reproducing and localizing kernel and model accuracy bugs. Serving features (scheduling and in-flight batching, chunked prefill, prefix caching, LoRA, and structured output) are tested through max serve. Requesting one that GenerateIterator cannot run raises an error that gives the matching max serve command.

Constructing the batch claims KV cache for every row. Use it as a context manager so an exception releases the rows. To run contexts built elsewhere, such as by the pipeline’s tokenizer, use from_contexts().

In a batch built from input_ids, row i samples with seed + i * 2**32. Serve deliberately makes requests with the same pinned seed draw identically. GenerateIterator instead gives each row its own stream, as Hugging Face users expect, so identical prompts draw different samples, a rerun reproduces them, and row 0 matches serve for the same seed. A prompt’s draw still depends on its row index. from_contexts() keeps each context’s own seed, so it draws as serve does.

On a synchronous pipeline, each step() returns its own forward’s tokens and the device is idle between steps. A pipeline that is not synchronous (an overlap pipeline with overlap active) returns the previous execute’s tokens instead: the first step returns {}, a row finishes one step after its last token is sampled, ignore_eos takes max_new_tokens + 1 steps, and a forward can still be running between steps and after done. Such a pipeline allows one open batch at a time, and if its execute raises, reload it: closing the batch frees the rows but cannot repair the pipeline’s pending state.

ignore_eos=True is the batch mode closest to Hugging Face generate, which keeps finished rows in the batch as padding. Prefill shapes still differ, because MAX batches prompts ragged where Hugging Face pads them on the left, so compare logits with a tolerance.

With prefix caching on, a later batch can reuse KV an earlier one committed and prefill fewer tokens. Turn prefix caching off in the KV cache config when prefill shapes must repeat across runs.

On a pipeline that ran warmup_graph_capture(), the first (prefill) step runs eagerly and later steps replay captured graphs, as long as each replica has at most min(max_batch_size, 128) rows. Warm up before creating a GenerateIterator. With device graph capture on but no warmup, GenerateIterator logs a warning and runs every step eagerly. A logits processor receives the captured output buffer, which the next replay overwrites, so it must copy what it keeps. output_logits copies a float32 [vocab] logits row to host for every generated token, so host memory grows by 4 * vocab bytes per token.

with GenerateIterator(pipeline, input_ids, max_new_tokens=20) as batch:
    for step_tokens in batch:
        print(step_tokens)
sequences = batch.sequences

Parameters:

• pipeline (BatchPipeline) – The pipeline to run.
• input_ids (Iterable[Sequence[int]]) – One token-id prompt per row, as lists or a 2D array.
• attention_mask (Iterable[Sequence[int]] | None) – Zeros mark padding in input_ids, which GenerateIterator drops, since it runs prompts unpadded. A tokenizer returns it with the ids.
• max_new_tokens (int) – The most tokens to generate for each row.
• min_new_tokens (int | None) – The fewest tokens a row generates before EOS can end it. Defaults to the model’s generation config.
• temperature (float | None) – The sampling temperature. Defaults to the pipeline’s runtime default, then the model’s generation config.
• top_k (int | None) – The top-k limit, with the same fallback as temperature.
• top_p (float | None) – The nucleus sampling threshold. Defaults to the model’s generation config.
• min_p (float | None) – The minimum token probability relative to the most likely token. Defaults to 0 (off), since the generation config does not set it.
• repetition_penalty (float | None) – The divisor for logits of repeated tokens. Defaults to the model’s generation config.
• frequency_penalty (float | None) – The penalty proportional to a token’s count. Defaults to 0 (off), since the generation config does not set it.
• presence_penalty (float | None) – The flat penalty for a token already present. Defaults to 0 (off), since the generation config does not set it.
• ignore_eos (bool) – Keeps generating past EOS, so every row runs max_new_tokens tokens and the batch size never changes.
• stop_token_ids (Sequence[int] | None) – Token ids that end a row, in addition to EOS.
• seed (int) – The sampling seed for row 0. Row i uses seed + i * 2**32.
• output_logits (bool) – Copies each row’s next-token logits to host on every step, for logits. Needs a synchronous pipeline.
• output_logprobs (int) – Records the log probability of each generated token and of this many top candidates, for log_probs. A model that does not compute log probabilities fails on the first step.
• caller (str) – What the user called, which errors name, such as "AutoPipeline".

Raises:

• TypeError – If input_ids or attention_mask is not a list of integer lists.
• ValueError – If the request needs a feature GenerateIterator leaves to max serve: LoRA, device graph capture with data_parallel_degree > 1 on a warmed-up pipeline, or a batch that only fits by scheduling or chunked prefill (more rows than the planned max batch size, or a replica’s prompts over its per-replica prefill budget or the batch total-token cap). Also raised for a prompt with no room to generate within the planned max length, an invalid sampling value, penalties on a pipeline built without enable_penalties, max_new_tokens below 1 or below min_new_tokens, attention_mask with a different shape than input_ids, ignore_eos with a row that does not fit or with speculative decoding, image inputs, or a model that does not generate text. Also raised for output_logprobs outside [0, 7], output_logits on a pipeline that is not synchronous, output_logits or output_logprobs with speculative decoding, or output_logprobs on an OverlapTextGenerationPipeline.
• RuntimeError – If the pipeline is not synchronous and another batch on it is still open.

close()

close()

source

Releases every live row, then drops the pipeline reference.

Safe to call more than once. Every row is attempted even if one fails, and the first error is re-raised afterwards with the pipeline still held.

Return type:

None

done

property done: bool

source

Whether every row has finished or been removed.

from_contexts()

classmethod from_contexts(pipeline, contexts, *, caller='GenerateIterator.from_contexts()')

source

Claims a batch over contexts already built, one row per context.

Each context keeps its own sampling params, seed, logits processors, stop criteria and max length, so this runs the contexts a pipeline’s tokenizer builds, image contexts included. On a speculative pipeline one step can return several tokens per row. It runs base-model contexts and contexts that name an active LoRA adapter. The batch records no logits or log_probs.

Parameters:

• pipeline (BatchPipeline) – The pipeline to run.
• contexts (Iterable[TextContext]) – One built context per row, for example from tokenizer.new_context, claimed in order.
• caller (str) – What the user called, which errors name, such as "max generate".

Returns:

A batch with every row claimed.

Raises:

• ValueError – If a context names a LoRA adapter that is loaded but not active, or for device graph capture with data_parallel_degree > 1 on a warmed-up pipeline, or if the batch does not fit the pipeline’s plan, as described for GenerateIterator.
• RuntimeError – If the pipeline is not synchronous and another batch on it is still open.

Return type:

GenerateIterator

log_probs

property log_probs: list[LogProbabilities]

source

The log probabilities behind each row’s new tokens, in input order.

For row i, token_log_probabilities[t] is the log probability of the row’s t-th new token, and top_log_probabilities[t] maps the output_logprobs most likely token ids, plus that token if it is not among them, to their log probabilities. sampled_token_ids holds the row’s new tokens. They match log_softmax(logits[i][t]) when the sampler applies no penalties and no min_new_tokens mask. With the default sampler on the model device, both edit the logits before log probabilities are taken. With enable_echo or sample_on_host, log probabilities come from the unpenalized logits.

Raises:

RuntimeError – If the batch was built without output_logprobs, as every from_contexts() batch is.

logits

property logits: list[ndarray[tuple[Any, ...], dtype[float32]]]

source

The next-token logits behind each row’s new tokens, in input order.

Row i has one [vocab] row per new token, and logits[i][t] holds the raw model output that produced the row’s t-th new token, before temperature, top-k, penalties or token masks. A row with no new tokens has shape [0, 0].

Raises:

RuntimeError – If the batch was built without output_logits, as every from_contexts() batch is.

remove()

remove(row)

source

Releases row so later steps run without it.

Removing a row that already finished is a no-op.

Parameters:

row (int) – The input index of the row to remove.

Raises:

IndexError – If row is not an input index.

Return type:

None

sequences

property sequences: list[list[int]]

source

Each row’s prompt followed by its new tokens so far, in input order.

step()

step()

source

Runs exactly one pipeline.execute over the live rows.

Rows the execute reports done are released before this returns. On a pipeline that is not synchronous the tokens lag one step, as GenerateIterator describes.

Returns:

The tokens the execute reported for each live row, keyed by row index.

Raises:

RuntimeError – If no rows are live, or if the pipeline returns log probabilities or logits rows that do not line up one per generated token, as when the model does not compute log probabilities.

Return type:

dict[int, list[int]]

synchronous

property synchronous: bool

source

Whether each step returns its own forward’s tokens, with no lag.