integrations.expert_parallel.buffer
integrations.expert_parallel.buffer
DeepEP Buffer singleton, lazily constructed on first call.
A single Buffer is reused across all MoE layers in a model, since DeepEP’s
intranode kernels are sized by num_nvl_bytes which we set conservatively at
plugin init. Per-layer Buffer construction would burn memory.
Functions
| Name | Description |
|---|---|
| barrier_ep | Barrier on the EP group. Placed before each DeepEP combine so the fast ranks wait (on |
| configure_buffer | Stash params for lazy Buffer construction. Call from post_model_build. |
| get_buffer | Return the (lazily constructed) DeepEP Buffer. |
| get_combine_config | Recommended DeepEP combine Config for the EP group (see get_dispatch_config). |
| get_dispatch_config | Recommended DeepEP dispatch Config for the EP group. DeepEP’s own tests ALWAYS pass an |
| reset_buffer | Drop the cached Buffer + configs. Used in tests. |
barrier_ep
integrations.expert_parallel.buffer.barrier_ep()Barrier on the EP group. Placed before each DeepEP combine so the fast ranks wait (on
NCCL, no short timeout) for any rank still AUTOTUNING the local expert kernel — otherwise the
combine collective hits DeepEP’s short internal timeout (value=0) and aborts. Negligible cost
once kernels are cached (all ranks arrive together). Disable with AXOLOTL_EP_NO_BARRIER=1.
configure_buffer
integrations.expert_parallel.buffer.configure_buffer(
ep_group,
num_nvl_bytes=256 << 20,
num_rdma_bytes=0,
)Stash params for lazy Buffer construction. Call from post_model_build.
get_buffer
integrations.expert_parallel.buffer.get_buffer()Return the (lazily constructed) DeepEP Buffer.
get_combine_config
integrations.expert_parallel.buffer.get_combine_config()Recommended DeepEP combine Config for the EP group (see get_dispatch_config).
get_dispatch_config
integrations.expert_parallel.buffer.get_dispatch_config()Recommended DeepEP dispatch Config for the EP group. DeepEP’s own tests ALWAYS pass an explicit config to dispatch/combine; the no-config default deadlocks / launch-fails the combine on the singleton buffer’s reuse across MoE layers (cudaErrorLaunchFailure).
reset_buffer
integrations.expert_parallel.buffer.reset_buffer()Drop the cached Buffer + configs. Used in tests.