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.legacy.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.

Compiles init_module() as the ep_init graph and runs it with run_init().

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.

init_module()

init_module()

source

Prepares this process for EP initialization.

Enables peer access between the local GPUs and sets the NVSHMEM environment, so call it before compiling the module it returns. Then run the compiled graph with run_init().

Returns:

The module whose graph sets up SHMEM and allocates the EP buffers.

Return type:

EPInit

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.

run_init()

run_init(init_model)

source

Runs the compiled EP initialization graph.

Keeps the device pointers to the SHMEM buffers it allocates, and sets the config’s node ID.

Parameters:

init_model (Model) – The graph of init_module(), compiled and initialized.

Return type:

None

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