IMPORTANT: To view this page as Markdown, append `.md` to the URL (e.g. /get-started.md). For the complete documentation index, see llms.txt.
Skip to main content
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

OverlapTextGenerationPipeline

OverlapTextGenerationPipeline​

final class max.pipelines.lib.OverlapTextGenerationPipeline(pipeline_config, pipeline_model, weight_adapters, tokenizer, memory_plan, disable_overlap=False)

source

Bases: TextGenerationPipelineInterface[TextGenerationContextType], Generic[TextGenerationContextType]

Overlap text generation pipeline.

Initialize a text generation pipeline instance.

This sets up devices, the inference session, tokenizer, KV-cache manager, sampling kernel, and loads model weights and adapters.

Parameters:

  • pipeline_config (PipelineConfig) – Configuration for the pipeline and runtime behavior.
  • pipeline_model (type[PipelineModel[Any]]) – Concrete model implementation to use for execution.
  • weight_adapters (dict[WeightsFormat, WeightsAdapter]) – Mapping from weights format to adapter implementation.
  • tokenizer (PipelineTokenizer[TextGenerationContextType, npt.NDArray[np.integer[Any]], TextGenerationRequest]) – Tokenizer implementation used to build contexts and decode.
  • memory_plan (_MemoryPlan) – Memory plan from the registry containing max_batch_size and other resolved memory parameters.
  • disable_overlap (bool) – When this flag is set, the overlap scheduler will immediately synchronize after model execution. This removes any potential cpu / gpu overlap.

Raises:

ValueError – If quantization_encoding is not configured in pipeline_config.model or if structured output is requested without a valid tokenizer delegate.

batch_spec_decode_metrics()​

batch_spec_decode_metrics()

source

Returns the per-batch draft token acceptance metrics for the most recent batch.

Return type:

_SpeculativeDecodingMetrics | None

batch_video_metrics()​

batch_video_metrics()

source

Returns video encoder metrics for the most recent batch.

Returns None for models with no video support and for batches that did no video encoding. Video encoding has no pipeline-owned cache equivalent to VisionEncoderCache, so this only ever comes from a model implementing SupportsPooledVisionMetrics.

Return type:

VideoEncoderMetrics | None

batch_vision_metrics()​

batch_vision_metrics()

source

Returns vision encoder metrics for the most recent batch.

Returns None for text-only models and for batches that did no vision encoding (e.g. decode steps). The metrics come from the pipeline-owned VisionEncoderCache, if this pipeline has one; otherwise, for a model that owns its encoder cache internally, from SupportsPooledVisionMetrics.

Return type:

VisionEncoderMetrics | None

execute()​

execute(inputs)

source

Executes a batch of requests asynchronously on the GPU.

This method returns before the outputs for the current batch are ready, so the outputs it returns belong to the previous batch. To drain the outputs for the final batch, call execute() again with an empty batch.

The batch of requests is a TextGenerationInputs, which wraps one or more TextContext objects. Build a batch on CPU like this:

import numpy as np
from max.pipelines.context import TextContext, TokenBuffer
from max.pipelines.modeling.types import (
    RequestID,
    TextGenerationInputs,
)

contexts = [
    TextContext(
        request_id=RequestID(),
        max_length=32,
        tokens=TokenBuffer(np.arange(8, dtype=np.int64)),
    )
    for _ in range(4)
]
inputs = TextGenerationInputs(batches=[contexts])
empty_inputs = TextGenerationInputs(batches=[[]])
assert len(inputs.flat_batch) == 4
assert len(empty_inputs.flat_batch) == 0

Given a loaded pipeline, the first execute returns no outputs (they belong to a not-yet-submitted previous batch); a second, empty call drains the first batch’s outputs:

output_a = pipeline.execute(inputs)
assert len(output_a) == 0

output_b = pipeline.execute(empty_inputs)
assert len(output_b) == len(inputs.flat_batch)

Parameters:

inputs (TextGenerationInputs[TextGenerationContextType]) – The inputs for the batch.

Returns:

A dictionary of request IDs to outputs. The outputs do not correspond to the requests in the input batch. Instead they are from the previous batch.

Return type:

dict[RequestID, TextGenerationOutput]

has_pending_outputs()​

has_pending_outputs()

source

Returns True if there are pending outputs for the previous batch.

If this is True, the caller should call execute() even with empty inputs to retrieve the outputs for the previous batch.

Return type:

bool

initialize_bitmask()​

initialize_bitmask(batch)

source

Allocates a per-request token bitmask for structured decoding.

Parameters:

batch (list[TextGenerationContextType]) – The generation contexts for the batch.

Returns:

A bitmask array of shape [batch_size, vocab_size] if structured output is enabled; otherwise None.

Return type:

ndarray[tuple[Any, …], dtype[int32]] | None

kv_manager​

property kv_manager: PagedKVCacheManager

source

Returns the KV cache manager for this pipeline.

max_batch_size​

property max_batch_size: int

source

Maximum number of requests that can be processed in a single batch.

overlap_active​

property overlap_active: bool

source

Whether CPU/GPU overlap is actually in effect.

When overlap is active, execute() defers synchronization of the current batch until the next call, so wall-clock time measured around execute() reflects the previous batch’s execution, not the current one.

pipeline_config​

property pipeline_config: PipelineConfig

source

Returns the pipeline configuration.

release()​

release(request_id)

source

Mark the context as complete, releasing the cache slot from the KV manager.

Note: Primary KV cache lifecycle is managed by the scheduler. This method handles extra KV caches managed by the pipeline model (e.g., indexer cache for DeepSeekV3.2).

Parameters:

request_id (RequestID)

Return type:

None

take_completed_batch_stats()​

take_completed_batch_stats()

source

Returns and clears stats for the most recently completed batch.

When overlap is active, a batch’s outputs are synchronized one execute() call after it was enqueued, so the scheduler cannot attribute execution time to the batch it just submitted. After each execute() call that synchronized a batch, this returns that batch’s composition and timing; the scheduler should publish execution-time and throughput telemetry from this record instead of from its own wall-clock measurement. Returns None when no batch completed since the last call.

Return type:

CompletedBatchStats | None

tokenizer​

property tokenizer: PipelineTokenizer[TextGenerationContextType, ndarray[tuple[Any, ...], dtype[integer[Any]]], TextGenerationRequest]

source

Returns the tokenizer used for building contexts and decoding.

update_for_structured_output()​

update_for_structured_output(context, bitmask, index)

source

Update context and logits bitmask for structured output.

If a json_schema is present and no matcher is set, this compiles a grammar matcher and installs it on the context, then fills the per-request token bitmask used to constrain the next-token distribution.

Parameters:

  • context (TextGenerationContextType) – Request context to update.
  • bitmask (ndarray[tuple[Any, ...], dtype[int32]]) – Optional preallocated bitmask buffer; updated in-place.
  • index (int) – Global position into the bitmask for this request.

Raises:

InputError – If a JSON schema is provided but structured output is not enabled via sampling configuration.

Return type:

None

warmup_graph_capture()​

warmup_graph_capture()

source

Initializes and runs overlap device graph capture warmup.

Return type:

None