Distributed: parallelism, KV transfer, EPLB

Active contributors: Wentao Ye, Tyler Michael Smith, Robert Shaw, Wensheng Tang, Nick Hill.

Purpose

vllm/distributed/ is the home of every cross-process or cross-host concern: tensor / pipeline / data / expert / context parallelism, the device communicators (NCCL, custom all-reduce, custom quickreduce, IPC, NIXL, etc.), KV / EC / weight transfer, and the elastic / EPLB controllers.

Directory layout

vllm/distributed/
├── __init__.py
├── communication_op.py     # broadcast/all-reduce primitives
├── parallel_state.py       # The big one: TP/PP/DP/EP groups (~1,800 lines)
├── stateless_coordinator.py
├── utils.py                # ProcessGroup helpers
├── nixl_utils.py           # NIXL transport helpers
├── kv_events.py            # KV cache event publisher (for connectors)
├── device_communicators/   # Per-device communicators (CUDA, ROCm, XPU, CPU, Gloo, custom_all_reduce, IPC)
├── kv_transfer/            # KV connectors (V1)
│   ├── kv_connector/v1/    # Connector v1 base + impls (NIXL, LMCache, Mooncake, Redis, ...)
│   ├── kv_connector/utils.py # KVOutputAggregator
│   └── kv_transfer_state.py
├── ec_transfer/            # Encoder-cache connectors (mirror of KV transfer)
│   └── ec_connector/...
├── weight_transfer/        # Weight reload across nodes (RL workflows)
├── eplb/                   # Expert-Parallel Load Balancer
└── elastic_ep/             # Elastic EP scaling middleware

Key abstractions

Abstraction	File	Role
parallel_state module	vllm/distributed/parallel_state.py	Manages all process groups (TP, PP, DP, EP, KV, etc.)
GroupCoordinator	vllm/distributed/parallel_state.py	Wraps a torch ProcessGroup with vLLM extras
DeviceCommunicator family	vllm/distributed/device_communicators/	Per-device fast collectives (custom all-reduce, IPC)
KVConnectorBase_V1	vllm/distributed/kv_transfer/kv_connector/v1/base.py	KV transport interface
KVConnectorFactory	vllm/distributed/kv_transfer/kv_connector/factory.py	Discovery + instantiation
ECConnectorFactory	vllm/distributed/ec_transfer/ec_connector/factory.py	Encoder-cache equivalent
EventPublisher	vllm/distributed/kv_events.py	Publishes per-block KV cache events for connectors
KVOutputAggregator	vllm/distributed/kv_transfer/kv_connector/utils.py	Aggregates per-rank connector outputs
EPLBConfig + reroute logic	vllm/distributed/eplb/	Token-level expert load balancing
ScalingMiddleware	vllm/entrypoints/serve/elastic_ep/middleware.py	API-server-side knob for elastic-EP scaling

Process groups

parallel_state is the central registry of process groups. Concepts:

• TP (tensor-parallel) — each transformer block is sharded across tp_size ranks. Linear layers split row/column-wise; attention heads split across ranks.
• PP (pipeline-parallel) — layers are partitioned into pp_size stages; activations flow rank-to-rank.
• DP (data-parallel) — independent replicas of the engine; load-balanced by the coordinator.
• EP (expert-parallel) — MoE experts sharded across ranks; each rank owns a subset.
• CP (context-parallel) — long-sequence partitioning (used by some attention backends; see vllm/v1/worker/cp_utils.py).

graph TB
    subgraph "World"
      direction LR
      subgraph "DP replica 0"
        direction LR
        subgraph "TP × PP grid"
          R0[rank 0]
          R1[rank 1]
          R2[rank 2]
          R3[rank 3]
        end
      end
      subgraph "DP replica 1"
        direction LR
        subgraph "TP × PP grid 2"
          R4[rank 4]
          R5[rank 5]
          R6[rank 6]
          R7[rank 7]
        end
      end
    end
    Coord[DPCoordinator]
    Coord --> R0
    Coord --> R4

Group construction lives in parallel_state.initialize_model_parallel. The order of group creation matters because torch.distributed allocates communicator handles in declaration order.

Communicators

Per-device communicators implement collective ops with hardware-aware fast paths:

• cuda_communicator.py — NCCL + custom all-reduce kernels (csrc/custom_all_reduce.cu) for small messages
• custom_quickreduce.py — Faster all-reduce via NVIDIA NVLink P2P
• pynccl.py — NCCL through pynccl for explicit stream control
• cpu_communicator.py — Gloo / shared memory
• xpu_communicator.py — Intel XPU (oneCCL)
• tpu_communicator.py — TPU XLA
• rocm_communicator.py
• ipc_communicator.py — Cross-process IPC (used for DP coordination)

The right communicator is picked by the active platform; users rarely select it directly.

KV connectors

vllm/distributed/kv_transfer/kv_connector/v1/ contains:

• base.py — KVConnectorBase_V1, plus SupportsHMA (heterogeneous memory addressing) capability flag
• Concrete implementations (NIXL, Mooncake, LMCache, Redis, etc.)
• metrics.py — KVConnectorStats

Connector roles:

• SCHEDULER — runs on the EngineCore process; is asked which requests have remote KV available, what to load/save.
• WORKER — runs on each worker; performs the actual block reads/writes.

The factory (factory.py) instantiates the right pair based on KVTransferConfig.kv_connector (null, nixl, mooncake, custom). Disaggregated prefill and KV offloading both use this interface — the choice of connector decides the destination.

For the original V0 disaggregated prefill workflow image, see vllm/distributed/kv_transfer/disagg_prefill_workflow.jpg.

EPLB

vllm/distributed/eplb/ implements Expert-Parallel Load Balancing: at runtime, when a few experts get hot, the EPLB controller redistributes them across ranks to keep each rank's compute roughly balanced. Configuration is in EPLBConfig (vllm/config/parallel.py).

Elastic EP

vllm/distributed/elastic_ep/ and vllm/entrypoints/serve/elastic_ep/middleware.py together let the EP topology grow or shrink while the engine is serving. The ReconfigureDistributedRequest message (vllm/v1/engine/__init__.py) is what plumbs the change end-to-end.

Weight transfer

vllm/distributed/weight_transfer/ provides WeightTransferInitRequest / WeightTransferUpdateRequest and a base transport. RL frameworks (e.g., for online policy updates) use this to push fresh weights into a running engine without restarting it.

Key source files

File	Purpose
vllm/distributed/parallel_state.py	Process group state + helpers
vllm/distributed/utils.py	Stateless group setup, env propagation
vllm/distributed/communication_op.py	High-level broadcast/all-reduce primitives
vllm/distributed/kv_events.py	KV-block event publisher
vllm/distributed/device_communicators/cuda_communicator.py	CUDA collectives
vllm/distributed/eplb/__init__.py	EPLB entry
vllm/distributed/elastic_ep/__init__.py	Elastic-EP entry
csrc/custom_all_reduce.cu / csrc/custom_quickreduce.cu	Custom collective kernels

Entry points for modification

• New transport: implement KVConnectorBase_V1 (or ECConnectorBase) and register through the factory.
• New collective: add a method on the right DeviceCommunicator and surface it through communication_op.py.
• New parallelism axis: extend parallel_state.py with the new group, plumb sizes through ParallelConfig, and update the executor.
• EPLB policy: subclass the rebalance algorithm in vllm/distributed/eplb/.

For how parallelism is exposed at the user level, see Configuration. For how parallel state is consumed by layers, see Model executor.