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

BatchContextLengths

BatchContextLengths

class max.pipelines.lib.registration_runtime.BatchContextLengths(data_parallel_degree, page_size, *, devices_per_replica=1, dtype=int32)

source

Bases: object

Each replica’s page-aligned KV context length, refreshed per batch.

With devices_per_replica, each replica’s length fills that many buffers in a row, one for each of its devices. The buffers persist across batches, and each batch rewrites them in place, so a captured device graph keeps reading the live values.

Parameters:

• data_parallel_degree (int)
• page_size (int)
• devices_per_replica (int)
• dtype (DType)

update()

update(replica_batches)

source

Writes each replica’s context length and returns the buffers.

Parameters:

replica_batches (Sequence[Sequence[TextContext]])

Return type:

list[Buffer]