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
EPConfig
EPConfig
class max.nn.comm.ep.EPConfig(dispatch_dtype, combine_dtype, hidden_size, top_k, n_experts, max_tokens_per_rank, n_gpus_per_node, n_nodes, node_id=-1, dispatch_quant_config=None, nvfp4_dyn_global_scales=False, fused_shared_expert=False, fuse_ffn_combine_send=False, moe_dim=0, fused_moe_ready=False, use_allreduce=False, max_batch_size=0, eplb_enabled=False, num_moe_layers=0, max_replicas=1, eplb_phy2log_plan=None, num_logical_experts=0)
Bases: object
Configuration for Expert Parallelism (EP) communication.
The EP kernel requires n_gpus_per_node * n_nodes >= 2.
Single-GPU MoE configs pass Python construction but fail at
Mojo kernel compile (constraint failed) after 2-3 minutes.
For single-GPU MoE, instantiate the MoE block without an EP
manager and use ShardingStrategy.tensor_parallel(1).
-
Parameters:
-
- dispatch_dtype (DType)
- combine_dtype (DType)
- hidden_size (int)
- top_k (int)
- n_experts (int)
- max_tokens_per_rank (int)
- n_gpus_per_node (int)
- n_nodes (int)
- node_id (int)
- dispatch_quant_config (QuantConfig | None)
- nvfp4_dyn_global_scales (bool)
- fused_shared_expert (bool)
- fuse_ffn_combine_send (bool)
- moe_dim (int)
- fused_moe_ready (bool)
- use_allreduce (bool)
- max_batch_size (int)
- eplb_enabled (bool)
- num_moe_layers (int)
- max_replicas (int)
- eplb_phy2log_plan (Any)
- num_logical_experts (int)
combine_dtype
combine_dtype: DType
Data type used for combining expert outputs.
dispatch_dtype
dispatch_dtype: DType
Data type used for dispatching tokens to experts.
dispatch_quant_config
dispatch_quant_config: QuantConfig | None = None
Quantization configuration used for dispatch token quantization.
eplb_enabled
eplb_enabled: bool = False
When true, EPBatchManager exposes log2phy / logcnt input buffer.
eplb_phy2log_plan
eplb_phy2log_plan: Any = None
Optional pre-computed EPLB plan (phy2log numpy array, shape [num_moe_layers, num_phy]). When non-None, EPBatchManager picks it up at construction time so that MoE.shard sees the plan during the model’s first __init__ pass (no re-shard needed).
estimate_memory_usage()
estimate_memory_usage()
Estimate the EP communication memory usage per device per buffer group.
-
Returns:
-
Estimated memory usage in bytes per device per buffer group.
-
Return type:
fuse_ffn_combine_send
fuse_ffn_combine_send: bool = False
Fuse the combine send into the MoE FFN’s own epilogue.
The FFN’s down-projection epilogue scatters each finished row straight
into the owning rank’s receive buffer and releases the arrival signal
itself, so the FFN’s output is never staged in a local tensor: the
(max_recv_tokens, hidden_size) buffer the unfused chain has to
materialize – the largest single activation in the MoE region – is not
allocated. The caller drops its combine_async and goes straight to
the wait, which reads only the symmetric receive buffers.
Opt-in per model rather than always-on: the send pays for itself at
prefill-scale tokens per expert and is neutral below, so it is a recipe
choice rather than a universal default. Incompatible with
use_allreduce (no peer buffers to write) and with
fused_shared_expert (which keeps the shared expert’s rows in the very
tensor being elided); both raise rather than degrade silently.
fused_moe_ready
fused_moe_ready: bool = False
Whether EP init set up the backend’s one-launch fused MoE for this
configuration. Set by EPCommInitializer when
MODULAR_EP_FUSED_MOE=1 requests it and the backend serves the
configuration; layers whose experts the fused op serves then run
dispatch, FFN and combine as one launch.
fused_shared_expert
fused_shared_expert: bool = False
Whether to fuse the shared expert computation with the routed experts.
get_max_recv_tokens()
get_max_recv_tokens()
Get the maximum number of tokens that can be received by a single device.
-
Return type:
hidden_size
hidden_size: int
Size of the hidden dimension in the model.
max_batch_size
max_batch_size: int = 0
Decode concurrency cap sizing the AMD MXFP4 preb direct grid.y on the decode bands. 0 disables (full-stride fallback).
max_replicas
max_replicas: int = 1
Largest replica count across all (layer, logical) cells. Also the trailing dim of log2phy.
max_tokens_per_rank
max_tokens_per_rank: int
Maximum number of tokens processed per GPU rank.
moe_dim
moe_dim: int = 0
Intermediate width of the routed experts’ FFN, when the model provides
it. A config that gives it gets the fused MoE counters’ reserve in its EP
sync-counter buffers. It does not request the fused MoE, which needs
MODULAR_EP_FUSED_MOE=1 and a configuration the backend serves.
n_experts
n_experts: int
Total number of experts in the model.
n_gpus_per_node
n_gpus_per_node: int
Number of GPUs available per node.
n_nodes
n_nodes: int
Total number of nodes in the distributed setup.
node_id
node_id: int = -1
ID of the current node. Will be set by the EPCommInitializer.
num_logical_experts
num_logical_experts: int = 0
Number of logical experts per layer. When EPLB is disabled this equals
n_experts. When EPLB is enabled, n_experts is inflated to the
physical slot count and this preserves the original logical count, which
is the leading dim of log2phy / logcnt.
num_moe_layers
num_moe_layers: int = 0
Number of MoE layers; sizes the per layer EPLB buffer.
nvfp4_dyn_global_scales
nvfp4_dyn_global_scales: bool = False
Quantize each NVFP4 dispatch token against its own global scale.
Instead of the checkpoint’s static input scale, the dispatch kernel first
reduces each token’s absolute max, then quantizes the token against
2688 / rowmax. The scale’s BF16 inverse travels with the token, and
the dispatch outputs gain output_rowwise_scales – one per received
row, directly after the block scales – so a received row dequantizes as
fp4 * block_scale * rowwise_scale. Requires NVFP4 dispatch.
top_k
top_k: int
Number of top experts to route each token to.
use_allreduce
use_allreduce: bool = False
Whether to use allreduce for the cross-device communication.