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

EPBatchManager

EPBatchManager

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

source

Bases: object

Batch manager for Expert Parallelism (EP).

This module manages two groups of SHMEM buffers in the graph. It switches between the two groups to avoid racing.

Initialize the EP batch manager.

Parameters:

config (EPConfig) – EP configuration.

atomic_counters

property atomic_counters: list[list[BufferValue]]

source

bind_comm_buffers()

bind_comm_buffers(comm)

source

Rebinds the buffer fields from a threaded EPCommBuffers.

Called at the top of the EP MoE forward so the kernel-dispatch helpers, which read self._send_buf_ptrs and so on, see the (subgraph-rebound) buffer values passed in as a forward argument.

Parameters:

comm (EPCommBuffers)

Return type:

None

comm_buffers()

comm_buffers(input_vals)

source

Fetches the EP buffers and wraps them as an EPCommBuffers.

Parameters:

input_vals (Iterable[Value[Any] | Tensor]) – The inputs declared by input_types(), in order, as graph values or single-device tensors.

Returns:

Dataclass containing the EP comm tensors.

Return type:

EPCommBuffers

config

config: EPConfig

source

Configuration for the Expert Parallelism (EP).

ep_combine()

ep_combine(input_tokens, router_weight, device_id, topk_ids=None)

source

Execute fused Expert Parallelism token combine (async + wait).

This method launches the fused EP combine kernel that combines both combine_async and combine_wait functionality in a single kernel launch. It sends expert outputs back to original devices, waits for all transfers to complete, and computes the weighted sum of routed expert outputs.

Note: For fused_shared_expert mode with the fused combine kernel, the shared expert outputs in input_tokens are automatically added to the reduced routed expert outputs.

Parameters:

• input_tokens (TensorLike) – Expert output tensors from the current device. A TensorValue with shape (max_recv_tokens, hidden_size). For fused_shared_expert mode, the shared expert outputs are stored at the start.
• router_weight (TensorLike) – Router weights for the current device. A TensorValue with shape (num_local_tokens, top_k).
• device_id (int) – Device ID for the current device.
• topk_ids (TensorLike | None) – Top-k expert IDs for the current device. Need to be provided for allreduce mode.

Returns:

Final output tensor with shape (num_local_tokens, hidden_size).

Return type:

TensorLike

ep_combine_all()

ep_combine_all(input_tokens, router_weights, device_ids)

source

Multi-device fused EP combine across all devices.

Launches a single mo.distributed.ep.combine graph op that combines expert outputs back to their original devices on all GPUs simultaneously.

Parameters:

• input_tokens (Sequence[TensorLike]) – Per-device expert output tokens.
• router_weights (Sequence[TensorLike]) – Per-device router weight tensors.
• device_ids (list[int]) – Device IDs corresponding to each input.

Returns:

Per-device combined output tensors.

Return type:

list[TensorLike]

ep_combine_async()

ep_combine_async(input_tokens, device_id)

source

Initiate Expert Parallelism combine phase (async).

This method launches the async combine phase of Expert Parallelism, sending expert outputs back to their original devices based on source routing information stored during the dispatch phase.

Parameters:

• input_tokens (TensorLike) – Expert output tensors from the current device. A TensorValue with shape (max_recv_tokens, hidden_size).
• device_id (int) – Device ID for the current device.

Return type:

None

ep_combine_wait()

ep_combine_wait(router_weight, device_id)

source

Wait for Expert Parallelism combine phase completion.

This method waits for all expert output transfers to complete, then organizes the received tokens back into their original format and positions for the current device.

Parameters:

• expert_weights – Router weights for the current device. A TensorValue with shape (num_local_tokens, top_k).
• device_id (int) – Device ID for the current device.
• router_weight (TensorLike)

Returns:

Final output tensor with shape (num_local_tokens, hidden_size).

Return type:

TensorLike

ep_dispatch()

ep_dispatch(input_tokens, topk_ids, device_id, input_scales=None)

source

Execute fused Expert Parallelism token dispatch (async + wait).

This method launches the fused EP dispatch kernel that combines both dispatch_async and dispatch_wait functionality in a single kernel launch. It distributes input tokens to expert devices, waits for all tokens to arrive, and organizes received tokens for grouped matmul.

For FP8 dispatch, input tokens are quantized to FP8 format during dispatch and the output includes both FP8 tokens and their scales.

Parameters:

• input_tokens (TensorLike) – Input tokens for the current device. A TensorValue with shape (num_local_tokens, hidden_size).
• topk_ids (TensorLike) – Top-k expert IDs for the current device. A TensorValue with shape (num_local_tokens, top_k).
• device_id (int) – Device ID for the current device.
• input_scales (TensorLike | None) – Optional input scales tensor. Needed for NVFP4 dispatch.

Returns:

• output_tokens: Aggregated tokens ready for grouped matmul. Shape: (max_recv_tokens, hidden_size).
• For FP8: output_scales: Scales for the FP8 tokens. Shape: (hidden_size // block_size, max_recv_tokens).
• With nvfp4_dyn_global_scales: output_rowwise_scales, each received row’s inverse global scale. Shape: (max_recv_tokens,).
• expert_start_indices: Row offsets for grouped matmul. Shape: (n_local_experts + 1,).
• expert_ids: Local expert IDs for the grouped operation. Shape: (n_local_experts,).
• expert_usage_stats: Statistics for the grouped matmul. Shape: (2,).

Return type:

A tuple containing

ep_dispatch_all()

ep_dispatch_all(input_tokens, topk_ids, device_ids, input_scales=None)

source

Multi-device fused EP dispatch across all devices.

Launches a single multi-device dispatch graph op (BF16, FP8, or NVFP4 depending on config) that dispatches tokens on all devices simultaneously.

Parameters:

• input_tokens (Sequence[TensorLike]) – Per-device input token tensors.
• topk_ids (Sequence[TensorLike]) – Per-device top-k expert ID tensors.
• device_ids (list[int]) – Device IDs corresponding to each input.
• input_scales (Sequence[TensorLike] | None) – Per-device input scales (required for NVFP4).

Returns:

Per-device output tuples. The last element of each tuple is always src_info; the remaining elements are the dispatch outputs followed by the grouped matmul metadata.

Return type:

list[tuple[TensorLike, …]]

ep_dispatch_async()

ep_dispatch_async(input_tokens, topk_ids, device_id, input_scales=None)

source

Initiate Expert Parallelism token dispatch phase (async).

This function launches the EP async dispatch kernel that distributes input tokens to expert devices based on top-k routing decisions.

Parameters:

• input_tokens (TensorLike) – Input tokens for the current device. A TensorValue with shape (num_local_tokens, hidden_size).
• topk_ids (TensorLike) – Top-k expert IDs for the current device. A TensorValue with shape (num_local_tokens, top_k).
• device_id (int) – Device ID for the current device.
• input_scales (TensorLike | None) – Optional input scales tensor. Required for NVFP4 dispatch.

Return type:

None

ep_dispatch_wait()

ep_dispatch_wait(device_id)

source

Wait for Expert Parallelism token dispatch phase completion.

This function launches the EP dispatch wait kernel that waits for all transfers to complete for the current GPU, then organizes the received tokens into a format suitable for grouped matmul computation.

Parameters:

device_id (int) – Device ID for the current device.

Returns:

• output_tokens: Aggregated tokens ready for grouped matmul computation. Shape: (max_recv_tokens, hidden_size).
• output_scales: Aggregated scales ready for grouped matmul computation. Only returned for quantized dispatch. Shape depends on format.
• expert_start_indices: Row offsets for grouped matmul computation. Shape: (n_local_experts + 1,).
• expert_ids: Local expert IDs for the grouped computation. Shape: (n_local_experts,).
• expert_usage_stats: Statistics for the grouped matmul computation. Shape: (2,).

Return type:

A tuple containing

fetch_buffers()

fetch_buffers(_input_vals)

source

Extract and organize communication buffers from graph input values.

Parameters:

• input_vals – List of input values containing all buffer references.
• _input_vals (Iterable[Value[Any]])

Return type:

None

input_tree()

input_tree()

source

Names the EP graph inputs, in the order the graph declares them.

atomic_counters is group-major ([NUM_GROUPS][n_gpus_per_node]); the three pointer fields hold one host tensor per group. The EPLB fields are present only when EPLB is enabled, one buffer per GPU.

Return type:

OrderedDict[str, tuple[TensorType | BufferType, …]]

input_types()

input_types()

source

Get the input types for the MoE graph.

Return type:

list[TensorType | BufferType]

mega_ffn_ep_combine_send()

mega_ffn_ep_combine_send(device_id, *, n_local_experts, max_recv_tokens, hidden_states, gate_up_weight, gate_up_a_scales, gate_up_b_scales, down_weight, down_b_scales, expert_start_indices, expert_ids, a_scale_offsets, gate_up_expert_scales, down_expert_scales, c_input_scales, estimated_total_m, num_active_experts, activation=None)

source

Runs the MoE FFN with the combine send fused into its epilogue.

Replaces the FFN’s two grouped matmuls and ep_combine_async() with one launch, so the FFN’s output is never staged in a local tensor. Follow it with ep_combine_wait(), which reads only the symmetric receive buffers.

Parameters:

• device_id (int) – Device ID for the current device.
• n_local_experts (int) – Experts owned by this device; the leading dim of the weights.
• max_recv_tokens (int) – Row capacity of the dispatched-token tensor.
• hidden_states (TensorValue) – Dispatched tokens.
• gate_up_weight (TensorValue) – Pre-permuted gate/up weights.
• gate_up_a_scales (TensorValue) – Token A-scale tile for the gate/up matmul.
• gate_up_b_scales (TensorValue) – Pre-permuted gate/up weight scale tile.
• down_weight (TensorValue) – Down-projection weights.
• down_b_scales (TensorValue) – Down-projection weight scale tile.
• expert_start_indices (TensorValue) – Per-expert prefix-sum token offsets.
• expert_ids (TensorValue) – Active local expert IDs.
• a_scale_offsets (TensorValue) – Per-expert scale-block offsets.
• gate_up_expert_scales (TensorValue) – Per-expert scaling for the gate/up leg.
• down_expert_scales (TensorValue) – Per-expert scaling for the down leg.
• c_input_scales (TensorValue) – Per-expert input scale for the re-quant.
• estimated_total_m (TensorValue) – Estimated total non-padded tokens.
• num_active_experts (TensorValue) – Active expert slots.
• activation (GatedActivation | None) – The gated activation the kernel fuses; see fused_silu().

Return type:

None

mega_ffn_ep_fused()

mega_ffn_ep_fused(device_id, *, input_tokens, topk_ids, router_weights, input_scales, gate_up_weight, gate_up_b_scales, down_weight, down_b_scales, gate_up_expert_scales, down_expert_scales, c_input_scales, activation=None)

source

One-launch fused EP MoE: dispatch -> MegaFFN -> weighted combine.

Replaces this layer’s ep_dispatch_async() / ep_dispatch_wait(), local FFN and ep_combine_async() / ep_combine_wait() on device_id. The fused workspace EP init set up is shared by every fused layer on the device: the kernel’s generation protocol makes consecutive launches safe, and the in-place group-0 counters chain them in the graph.

Parameters:

• device_id (int) – Device ID for the current device.
• input_tokens (TensorValue) – Pre-dispatch BF16 tokens (T, hidden).
• topk_ids (TensorValue) – Routed global expert ids (T, top_k) int32.
• router_weights (TensorValue) – Router weights (T, top_k) f32.
• input_scales (TensorValue) – Inverted NVFP4 dispatch input scale (element 0 read; 1.0 at MXFP8).
• gate_up_weight (TensorValue) – Sigma-permuted NVFP4 or MXFP8 gate/up weights.
• gate_up_b_scales (TensorValue) – Gate/up weight scales (6D).
• down_weight (TensorValue) – NVFP4 or MXFP8 down weights.
• down_b_scales (TensorValue) – Down weight scales (6D).
• gate_up_expert_scales (TensorValue) – Per-local-expert L1 alpha.
• down_expert_scales (TensorValue) – Per-local-expert L2 alpha.
• c_input_scales (TensorValue) – Per-local-expert L1->L2 requant scale.
• activation (GatedActivation | None) – The gated activation the kernel fuses; see fused_silu().

Returns:

The routed-expert MoE output (T, hidden) BF16.

Return type:

TensorValue

recv_buf_ptrs

property recv_buf_ptrs: list[TensorValue]

source

recv_count_ptrs

property recv_count_ptrs: list[TensorValue]

source

send_buf_ptrs

property send_buf_ptrs: list[TensorValue]

source