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

EPCommInitializer

EPCommInitializer

class max.nn.comm.ep.EPCommInitializer(config)

source

Bases: object

Helper class for initializing buffers for Expert Parallelism (EP).

This class handles the initialization of the SHMEM communication infrastructure required for Expert Parallelism. It creates and manages atomic counters, initializes the SHMEM library, and allocates symmetric memory buffers.

Initialize the EP communication initializer.

Parameters:

config (EPConfig) – EP configuration.

atomic_counters

atomic_counters: list[Buffer]

source

List of atomic counters used for synchronization.

config

config: EPConfig

source

EP configuration.

ep_init()

ep_init(session)

source

Initialize Expert Parallelism communication infrastructure.

Parameters:

session (InferenceSession) – Inference session used to compile and execute the graph.

Return type:

None

fused_moe_workspace

fused_moe_workspace: tuple[int, int] | None = None

source

Bytes per device and token block of the one-launch fused MoE workspace EP init set up, or None when it set none up.

init_model

init_model: Model

source

Compiled model that sets up the SHMEM library context for local GPUs and allocates the SHMEM memory for the send, receive, and receive count buffers.

model_inputs()

model_inputs()

source

Get the model inputs for the MoE model.

Returns:

List of all tensors needed as model inputs, in the order of EPBatchManager.input_types().

Return type:

list[Buffer]

Raises:

ValueError – If EPLB is enabled but no placement was supplied.

recv_buf_ptrs

recv_buf_ptrs: list[Buffer]

source

List of device pointers for the receive buffer.

recv_count_ptrs

recv_count_ptrs: list[Buffer]

source

List of device pointers for the receive count buffer.

send_buf_ptrs

send_buf_ptrs: list[Buffer]

source

List of device pointers for the send buffer.

set_eplb_placement()

set_eplb_placement(log2phy, logcnt)

source

Supplies the per-device EPLB tables model_inputs() returns.

Parameters:

• log2phy (Sequence[Buffer]) – One [num_moe_layers, num_logical, max_replicas] buffer per local GPU.
• logcnt (Sequence[Buffer]) – One [num_moe_layers, num_logical] buffer per local GPU.

Raises:

ValueError – If EPLB is disabled in the config or a list does not have one buffer per local GPU.

Return type:

None