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

EPCommBuffers

EPCommBuffers

class max.nn.comm.ep.EPCommBuffers(atomic_counters, send_buf_ptrs, recv_buf_ptrs, recv_count_ptrs, eplb_log2phy=None, eplb_logcnt=None)

source

Bases: object

SHMEM communication buffers for one Expert Parallelism MoE forward pass.

Bundles the per-device send, receive, and synchronization tensors produced by EPBatchManager.comm_buffers() into a single value so they can be passed through the subgraph pytree machinery as one forward argument. The __tree_flatten__ / __tree_unflatten__ hooks expose the Tensor leaves to that machinery.

Parameters:

• atomic_counters (list[list[Tensor]])
• send_buf_ptrs (list[Tensor])
• recv_buf_ptrs (list[Tensor])
• recv_count_ptrs (list[Tensor])
• eplb_log2phy (dict[int, Tensor] | None)
• eplb_logcnt (dict[int, Tensor] | None)

atomic_counters

atomic_counters: list[list[Tensor]]

source

Atomic synchronization counters, indexed as [group][device], used to coordinate thread blocks during the dispatch and combine phases.

eplb_log2phy

eplb_log2phy: dict[int, Tensor] | None = None

source

Per-device Expert Parallelism load-balancing (EPLB) logical-to-physical expert map, keyed by device ID. Populated only when EPLB is enabled; otherwise None.

eplb_logcnt

eplb_logcnt: dict[int, Tensor] | None = None

source

Per-device count of physical replicas per logical expert (EPLB), keyed by device ID. Populated only when EPLB is enabled; otherwise None.

recv_buf_ptrs

recv_buf_ptrs: list[Tensor]

source

Per-group SHMEM receive-buffer device pointers, one entry per buffer group. Each tensor holds the per-device addresses of the buffers for incoming tokens from remote devices.

recv_count_ptrs

recv_count_ptrs: list[Tensor]

source

Per-group SHMEM receive-count device pointers, one entry per buffer group. Each tensor holds the per-device addresses of the buffers that signal transfer completion.

send_buf_ptrs

send_buf_ptrs: list[Tensor]

source

Per-group SHMEM send-buffer device pointers, one entry per buffer group. Each tensor holds the per-device addresses of the staging buffers for outgoing tokens.