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

MulticastPool

MulticastPool

class max.nn.comm.MulticastPool

source

Bases: object

Symmetric multicast pool backing NVSwitch (NVLS) allreduce.

Unlike Signals, the pool is a single driver-owned allocation registered once per device set; the kernel reaches it through the driver registry, never through a graph operand. Capacity is a cap, not a fit: a payload that does not fit falls back to peer-to-peer allreduce.

ENABLED

ENABLED = False

source

MAX_BYTES

MAX_BYTES = 536870912

source

MIN_DEVICES

MIN_DEVICES = 4

source

allocate()

static allocate(devices, byte_size)

source

Allocates the pool for devices, or returns None.

Every precondition is a fallback rather than an error: zero capacity, virtual devices, no multicast support, or a failed allocation all yield None and callers run peer-to-peer allreduce.

Parameters:

• devices (Sequence[Device]) – Driver devices the pool should span.
• byte_size (int) – Per-device bytes, as charged by memory planning.

Returns:

The registered pool, or None when NVLS is unavailable.

Return type:

SymmetricPool | None

capacity_bytes()

static capacity_bytes(num_devices, num_tokens, hidden_size)

source

Returns the usable per-device pool bytes, or 0 for none.

Sizes for the dominant tensor-parallel allreduce, whose operand is num_tokens x hidden_size in bf16. Larger reductions (vocab-parallel logits, MoE router) exceed it and take peer-to-peer instead. The driver rounds the reservation up to its granularity (2 MiB), which memory planning does not charge.

Parameters:

• num_devices (int) – Ranks the pool would span.
• num_tokens (int) – Per-batch token budget (max_batch_input_tokens), not max_batch_size, which memory planning derives from this estimate.
• hidden_size (int) – Model hidden width, or 0 when unknown.

Returns:

Usable per-device bytes, or 0 when no pool should be allocated.

Return type:

int