Skip to main content
afdEP-comm

Expert-parallel communication for Attention-FFN Disaggregation. DeepEP V2 kernels specialized for the two roles: Attention GPUs only send in dispatch, FFN GPUs only send in combine.

APIDesign
0entry barriers in one NVLink domain
50%registered buffers vs the union, at 2A2F
-68%small-batch cross-node round trip, best case

Quick start

pip install .
export AFDEP_NCCL_ROOT=/path/to/nccl
export LD_LIBRARY_PATH="$AFDEP_NCCL_ROOT/lib:$LD_LIBRARY_PATH"
python -m afdep doctor --build
comm = Buffer(group=group, config=Config(
    capacity_tokens_per_rank=2048, hidden_size=4096, num_experts=128,
    top_k=8, ffn_ranks=2, attention_ranks=2))
if comm.role == "attention":
    comm.source_dispatch(x, expert_ids, weights, bucket=2048)
    output = comm.source_collect()
else:
    received = comm.worker_receive(bucket=2048)
    comm.worker_reply(run_experts(received), received.receipt)
comm.close()