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

EplbStatsAccumulator

EplbStatsAccumulator

class max.pipelines.lib.eplb_stats.EplbStatsAccumulator(metadata, devices)

source

Bases: object

Owns per-GPU persistent int64 histogram buffers.

Lives in the model worker. Each forward atomically increments cells in these buffers via on-GPU scatter_nd_add ops emitted by _ep_forward. snapshot() pulls them to host and sums across GPUs.

Initializes a zero-valued accumulator.

Parameters:

• metadata (EplbStatsMetadata) – Static shape / semantics descriptor.
• devices (list[Device]) – List of devices to allocate buffers on.

device_buffers

property device_buffers: list[Buffer]

source

Layer-major X device-major flat list, matches input_types order.

metadata

property metadata: EplbStatsMetadata

source

Static shape/semantics descriptor for the accumulator.

record_batch_total_tokens()

record_batch_total_tokens(num_tokens)

source

Increments the cumulative token count by num_tokens.

Parameters:

num_tokens (int)

Return type:

None

reset()

reset()

source

Zeros every per-(layer, device) counter buffer and the token count.

Zeros in place instead of reallocating: device-graph capture pins the pre-captured decode graph to the exact Buffer objects this accumulator owned at warmup. Reallocating orphans those buffers, so decode-step accumulations (which run via captured-graph replay) land in the orphaned buffers while snapshot() reads the new ones. Preserving buffer identity keeps the captured decode graph writing the live buffers. Runs on the model-worker thread between forward steps, so the zero cannot race a forward pass.

Return type:

None

snapshot()

snapshot(hostname=None)

source

Returns a snapshot of the current per-layer routing histograms.

Pulls each per-(layer, device) counter buffer to host, sums across devices, and packages the result with the cumulative token count and optional worker hostname.

Parameters:

hostname (str | None) – Optional worker hostname to tag the snapshot with.

Returns:

An immutable EplbStatsSnapshot.

Return type:

EplbStatsSnapshot