diff --git a/afd_plugin/connectors/base.py b/afd_plugin/connectors/base.py index 4bf86a7e2..096f42753 100644 --- a/afd_plugin/connectors/base.py +++ b/afd_plugin/connectors/base.py @@ -15,6 +15,7 @@ if TYPE_CHECKING: import torch from vllm.config import VllmConfig + from vllm.distributed.eplb.eplb_state import EplbState from afd_plugin.connectors.metadata import ( AFDA2FTransferPayload, @@ -328,6 +329,7 @@ def init_afd_connector( *, moe_layer_indices: Sequence[int] | None = None, local_experts: Mapping[int, FFNLocalExperts] | None = None, + eplb_state: EplbState | None = None, ) -> None: """Collective: check what every rank loaded, then create resources. @@ -336,6 +338,9 @@ def init_afd_connector( ``None`` skips the post-load checks for this rank. local_experts: FFN only: this rank's experts per MoE layer, as vLLM loaded them. + eplb_state: FFN only, EPLB deployments: the runner's registered + vLLM ``EplbState``, whose expert maps and rearrangement + schedule the connector shares with the Attention ranks. """ raise NotImplementedError @@ -358,6 +363,28 @@ def synchronize_lanes(self) -> None: """ raise NotImplementedError + def apply_eplb_routing( + self, + topk_ids: torch.Tensor, + layer_idx: int, + ) -> torch.Tensor: + """Rewrite Attention-side logical top-k ids for expert placement. + + Connectors that place experts dynamically (EPLB) override this to map + logical ids to physical ids and record expert load; the default is the + identity, for connectors whose FFN routes or places experts linearly. + """ + return topk_ids + + def step_ffn_eplb(self) -> None: + """Advance FFN-side EPLB state after one connector step. + + The FFN runner calls this once per step, real or dummy. Connectors + that balance experts dynamically implement it; the default (static + expert placement) does nothing. + """ + return None + class AFDControlPlane(ABC): """DP metadata control plane.""" diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index c0d2fa917..5906325fb 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -83,6 +83,11 @@ AFDRoutedExpertsConnector, ConnectorExtraInfo, ) +from afd_plugin.connectors.gpu.deepep_m2n.eplb import ( + DeepEPM2NAttentionEplb, + DeepEPM2NFFNEplb, + recv_initial_eplb_state, +) from afd_plugin.connectors.gpu.deepep_m2n.experts import ( EXPERT_BACKENDS, DeepEPM2NExperts, @@ -433,6 +438,10 @@ class HandshakeResult: attn_tp_size: int capacity_tokens_per_rank: int bucket_ladder: tuple[int, ...] + # Reported by the FFN ranks; Attention adopts it (EPLB adds redundant + # physical expert slots on top of the model's logical experts). + num_physical_experts: int + enable_eplb: bool def expert_weight_format(quant_config: QuantizationConfig | None) -> str: @@ -450,6 +459,7 @@ def expert_weight_format(quant_config: QuantizationConfig | None) -> str: def ffn_deployment_info(vllm_config: VllmConfig) -> dict[str, Any]: """FFN-side vLLM settings that decide which experts each rank loads.""" parallel_config = vllm_config.parallel_config + num_physical_experts = _ffn_num_physical_experts(vllm_config) return { "expert_parallel": parallel_config.enable_expert_parallel, "ep_size": parallel_config.data_parallel_size @@ -457,11 +467,24 @@ def ffn_deployment_info(vllm_config: VllmConfig) -> dict[str, Any]: "tensor_parallel_size": parallel_config.tensor_parallel_size, "expert_placement_strategy": parallel_config.expert_placement_strategy, "eplb": parallel_config.enable_eplb, + "num_physical_experts": num_physical_experts, "all2all_backend": parallel_config.all2all_backend, "weight_format": expert_weight_format(vllm_config.quant_config), } +def _ffn_num_physical_experts(vllm_config: VllmConfig) -> int: + """Physical expert slots the FFN loads: logical experts plus EPLB + redundant copies. DeepEP routes rows to slots, so every rank must agree + on this count.""" + num_experts = vllm_config.model_config.get_num_experts() + if not vllm_config.parallel_config.enable_eplb: + return num_experts + return num_experts + int( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + + def local_handshake_info( *, role: str, @@ -527,11 +550,15 @@ def _check_ffn_deployment( deployment: dict[str, Any] | None, layout: DeepEPM2NLayout, payload_format: str, + num_logical_experts: int, ) -> None: """vLLM must place experts on FFN ranks exactly as DeepEP routes them. Otherwise rows reach a rank that holds other experts and the results are silently wrong. Every rank runs this, so Attention fails with the FFN. + With EPLB the FFN loads ``num_physical_experts`` slots (logical experts + plus redundant copies), and DeepEP routes by slot id, so Attention adopts + that count for its own layout after the handshake. """ if deployment is None: raise RuntimeError(f"FFN world rank {world_rank} sent no deployment info") @@ -540,7 +567,6 @@ def _check_ffn_deployment( "ep_size": layout.ffn_size, "tensor_parallel_size": 1, "expert_placement_strategy": "linear", - "eplb": False, "all2all_backend": IDLE_ALL2ALL_BACKEND, "weight_format": PAYLOAD_WEIGHT_FORMATS[payload_format], } @@ -555,6 +581,19 @@ def _check_ffn_deployment( f"FFN world rank {world_rank} is deployed with {mismatches}; " f"DeepEP M2N with payload_format={payload_format!r} needs {required}" ) + num_physical = int(deployment["num_physical_experts"]) + if deployment["eplb"]: + if num_physical < num_logical_experts: + raise RuntimeError( + f"FFN world rank {world_rank} enables EPLB with " + f"{num_physical} physical experts for {num_logical_experts} " + "logical experts" + ) + elif num_physical != num_logical_experts: + raise RuntimeError( + f"FFN world rank {world_rank} reports {num_physical} physical " + f"experts for {num_logical_experts} logical experts without EPLB" + ) def validate_handshake( @@ -588,11 +627,36 @@ def validate_handshake( ) for world_rank, info in enumerate(infos[: layout.ffn_size]): _check_ffn_deployment( - world_rank, info["ffn_deployment"], layout, info["payload_format"] + world_rank, + info["ffn_deployment"], + layout, + info["payload_format"], + int(info["num_experts"]), ) alignments = {info["expert_alignment"] for info in infos[: layout.ffn_size]} if len(alignments) != 1 or None in alignments: raise RuntimeError(f"FFN ranks disagree on expert_alignment: {alignments}") + # EPLB adds redundant physical expert slots; every FFN rank must load the + # same count, and it must split evenly over the FFN ranks so DeepEP's + # slot-to-rank routing stays contiguous. + deployments = [info["ffn_deployment"] for info in infos[: layout.ffn_size]] + enable_eplb = {bool(deployment["eplb"]) for deployment in deployments} + if len(enable_eplb) != 1: + raise RuntimeError(f"FFN ranks disagree on EPLB: {sorted(enable_eplb)}") + num_physical_experts = { + int(deployment["num_physical_experts"]) for deployment in deployments + } + if len(num_physical_experts) != 1: + raise RuntimeError( + f"FFN ranks disagree on num_physical_experts: {num_physical_experts}" + ) + num_physical_experts_value = num_physical_experts.pop() + if num_physical_experts_value % layout.ffn_size != 0: + raise RuntimeError( + f"EPLB gives {num_physical_experts_value} physical experts, which " + f"does not split evenly over {layout.ffn_size} FFN ranks; adjust " + "num_redundant_experts" + ) attn_infos = infos[layout.ffn_size :] tp_sizes = {info["attn_tp_size"] for info in attn_infos} if len(tp_sizes) != 1 or None in tp_sizes: @@ -636,6 +700,8 @@ def validate_handshake( attn_tp_size=attn_tp_size, capacity_tokens_per_rank=int(capacities.pop()), bucket_ladder=ladders.pop() or (), + num_physical_experts=num_physical_experts_value, + enable_eplb=enable_eplb.pop(), ) @@ -915,11 +981,27 @@ def __init__( self.layout = DeepEPM2NLayout( ffn_size=afd_config.num_ffn_ranks, attn_size=afd_config.num_attention_ranks, - num_experts=vllm_config.model_config.get_num_experts(), + # FFN knows its own EPLB slot count up front; Attention adopts + # the FFN's count from the handshake below. + num_experts=( + _ffn_num_physical_experts(vllm_config) + if afd_config.role == "ffn" + else vllm_config.model_config.get_num_experts() + ), ) + # EPLB coordination (see deepep_m2n/eplb.py): an Attention-side load + # recorder or an FFN-side rearrangement driver, built during + # ``init_afd_connector`` once the handshake reports the FFN's EPLB + # deployment. ``None`` when the FFN runs without EPLB. + self.eplb: DeepEPM2NAttentionEplb | DeepEPM2NFFNEplb | None = None + self._handshake_enable_eplb = False self.attn_size = self.layout.attn_size self.ffn_size = self.layout.ffn_size self.is_ffn = afd_config.role == "ffn" + # The model's logical expert count, unlike the layout's physical slot + # count when EPLB adds redundant experts: the handshake compares this + # across roles, and both roles know it from the model config. + self.num_logical_experts = int(text_config.num_experts) self.is_fp8_payload = self.config.payload_format == FP8_PAYLOAD self.payload_dtype = FP8_DTYPE if self.is_fp8_payload else self.dtype if self.is_fp8_payload and self.hidden_size % FP8_GROUP_SIZE != 0: @@ -1015,7 +1097,7 @@ def _handshake(self) -> None: config=self.config, hidden_size=self.hidden_size, dtype=self.dtype, - num_experts=self.layout.num_experts, + num_experts=self.num_logical_experts, top_k=self.top_k, num_lanes=self.num_lanes, capacity_tokens_per_rank=self.capacity_tokens_per_rank, @@ -1041,6 +1123,16 @@ def _handshake(self) -> None: self.attn_tp_size = result.attn_tp_size self.capacity_tokens_per_rank = result.capacity_tokens_per_rank self.bucket_ladder = result.bucket_ladder + self._handshake_enable_eplb = result.enable_eplb + # With EPLB the FFN loads physical expert slots (logical experts plus + # redundant copies) and DeepEP routes by slot id, so Attention adopts + # the FFN's slot count for its own routing layout. + if not self.is_ffn and self.layout.num_experts != result.num_physical_experts: + self.layout = DeepEPM2NLayout( + ffn_size=self.layout.ffn_size, + attn_size=self.layout.attn_size, + num_experts=result.num_physical_experts, + ) if self.world_rank == 0: logger.info("DeepEP M2N network plan: %s", self.network_plan.summary()) # DeepEP reads both when each buffer is created: the GIN switch, and @@ -1054,6 +1146,7 @@ def init_afd_connector( *, moe_layer_indices: Sequence[int] | None = None, local_experts: Mapping[int, FFNLocalExperts] | None = None, + eplb_state: Any | None = None, ) -> None: """Collective: check what every rank loaded, then create buffers. @@ -1062,6 +1155,9 @@ def init_afd_connector( ``None`` skips the post-load checks for this rank. local_experts: FFN only: this rank's experts per MoE layer, as vLLM loaded them. + eplb_state: FFN only, EPLB deployments: the runner's registered + ``EplbState``. Its expert maps and rearrangement schedule are + shared with the Attention ranks (see ``deepep_m2n/eplb.py``). """ if self._initialized: return @@ -1084,6 +1180,29 @@ def init_afd_connector( hidden_size=self.hidden_size, payload_format=self.config.payload_format, ) + if self._handshake_enable_eplb: + if self.is_ffn: + if eplb_state is None: + raise RuntimeError( + "DeepEP M2N EPLB needs the FFN runner's EplbState; " + "pass it to init_afd_connector(eplb_state=...)" + ) + self.eplb = DeepEPM2NFFNEplb(self, eplb_state) + self.eplb.broadcast_initial_state() + else: + if moe_layer_indices is None: + raise RuntimeError( + "DeepEP M2N EPLB needs the Attention model's MoE layer " + "indices at init_afd_connector" + ) + schedule, initial_map, initial_count = recv_initial_eplb_state(self) + self.eplb = DeepEPM2NAttentionEplb( + self, + schedule, + tuple(moe_layer_indices), + initial_map, + initial_count, + ) # Never the lanes' bootstrap group: the control plane's close() aborts # this group to cancel a pending receive. self.group = init_afd_process_group( @@ -1155,6 +1274,14 @@ def close(self) -> None: def is_initialized(self) -> bool: return self._initialized + @property + def attention_tp_rank(self) -> int: + """This rank's position inside its Attention TP group.""" + if self.is_ffn: + raise RuntimeError("FFN ranks have no Attention TP rank") + assert self.attn_tp_size is not None + return (self.world_rank - self.ffn_size) % self.attn_tp_size + @property def domain(self) -> tuple[int, int]: """DeepEP ``(scale-out ranks, scale-up ranks)``: nodes and GPUs per node @@ -1376,6 +1503,28 @@ def send_ffn_output( done.record(lane_stream) self._lane_done[state.lane] = done + def apply_eplb_routing( + self, + topk_ids: torch.Tensor, + layer_idx: int, + ) -> torch.Tensor: + """Map logical top-k ids to physical expert ids and record the load. + + Identity without EPLB; with EPLB the Attention-side recorder maps + through the placement the FFN committed and counts rows per physical + expert (see ``deepep_m2n/eplb.py``). + """ + self._check_initialized() + if not isinstance(self.eplb, DeepEPM2NAttentionEplb): + return topk_ids + return self.eplb.apply_routing(topk_ids, layer_idx) + + def step_ffn_eplb(self) -> None: + """Advance the FFN-side rearrangement counter; see ``eplb.py``.""" + self._check_initialized() + if isinstance(self.eplb, DeepEPM2NFFNEplb): + self.eplb.step() + def synchronize_lanes(self) -> None: """Make the caller's stream wait for every lane's outstanding combine. @@ -1585,6 +1734,13 @@ def update_state_from_dp_metadata(self, payload: AFDControlPayload) -> None: f"{len(payload.dp_metadata_list)} stages exceed " f"num_lanes={connector.num_lanes}" ) + # EPLB recording runs where routing runs, on Attention: step the + # recorder for the transaction that just finished (this rank's + # previous payload) and prepare the recording state for the one about + # to run. The FFN steps its side after its forward, so both roles + # reach each rearrangement boundary together. + if isinstance(connector.eplb, DeepEPM2NAttentionEplb): + connector.eplb.on_control_step(payload) graph_mode = payload.is_graph_capturing or payload.is_graph_replaying if ( connector.attn_tp_size is None diff --git a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py new file mode 100644 index 000000000..c37240420 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py @@ -0,0 +1,551 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""EPLB support for the DeepEP M2N connector. + +With DeepEP M2N, routing (gate and top-k) runs on Attention, so the two EPLB +duties split across the AFD roles: + +- **Attention** maps logical top-k ids to physical expert ids before dispatch + (``eplb_map_to_physical_and_record``) and records the expert load of its own + rows. :class:`DeepEPM2NAttentionEplb` owns that recorder: per-layer load + buffers, a sliding load window, and the current logical-to-physical maps. +- **FFN** holds the experts, so it owns vLLM's ``EplbState`` (registered by + ``GPUFFNModelRunner.load_model``) and runs the rearrangement. + :class:`DeepEPM2NFFNEplb` drives it: it never calls ``EplbState.step`` + (nothing on the FFN records load), it steps its own rearrangement counter + and calls ``EplbState.rearrange`` directly at each boundary. + +Both sides count every connector step (real, warmup and graph-capture alike, +matching the FFN runner's ``eplb_step`` cadence) so the boundary is detected +in lockstep. All exchanges run on the connector's CPU bootstrap group (gloo) +with CPU-staged tensors: rearrangements are rare, so a host sync there is +cheap and keeps DeepEP's NCCL communicators out of it. + +Boundary exchange (every rank enters, in this order): + +1. ``reduce``: every Attention rank contributes its recorded load window; FFN + ranks contribute zeros. FFN world rank 0 receives the global window and + copies it into its ``EplbState``; the other FFN ranks' windows stay zero, + so the all-reduce inside ``EplbState.rearrange`` (over the FFN-only EP + group) distributes exactly one copy of the global window to every FFN rank. +2. FFN only: ``EplbState.rearrange`` computes the new placement, moves the + expert weights and commits the new maps in place. +3. ``broadcast``: FFN world rank 0 sends the committed + ``logical_to_physical_map`` and ``logical_replica_count``; Attention copies + them into its stable map tensors, whose addresses CUDA graphs captured. + +Attention blocks on the exchange before it routes the next step's tokens, and +the FFN runs it after its step's last combine, so Attention can never dispatch +rows placed by an older map than the FFN's weights. + +Loss logging (``log_balancedness``, one line every +``log_balancedness_interval`` steps) reuses that exchange shape without +changing anything: Attention ranks reduce their load window aggregated into +logical experts, and the FFN's world rank 0 scores the deployed map on it with +EPLB's own load model (an expert's window load split evenly over its replicas). +The result is the quantity ``DefaultEplbPolicy`` optimises, so the series moves +only when a rearrangement changes the placement. At each boundary the same +window scores the replaced map and the committed one, which is what tells +whether the rearrangement helped. It is opt-in because it adds a host sync on +every participating rank; the map exchange above is unaffected and still runs +only at a rearrangement boundary. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch +import torch.distributed as dist +from vllm.distributed.eplb.eplb_state import EplbModelState, EplbState +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) + +from afd_plugin.connectors.metadata import AFDControlPayload +from afd_plugin.distributed.eplb import ( + describe_eplb_placement, + describe_eplb_rank_load, + eplb_map_loss, + logical_load_from_window, +) + +if TYPE_CHECKING: + from afd_plugin.connectors.gpu.deepep_m2n.connector import DeepEPM2NAFDConnector + +logger = init_logger(__name__) + +# Keys of the schedule FFN world rank 0 broadcasts to every rank at connector +# init. Everything the Attention recorder needs to mirror the FFN's +# rearrangement cadence and allocate identically shaped tensors. +_EPLB_SCHEDULE_FIELDS = ( + "num_moe_layers", + "num_logical_experts", + "num_physical_experts", + "map_slots", + "window_size", + "step_interval", + "initial_rearrangement_step", + "log_interval", +) + + +def recv_initial_eplb_state( + connector: DeepEPM2NAFDConnector, +) -> tuple[dict[str, int], torch.Tensor, torch.Tensor]: + """Collective: Attention side of ``broadcast_initial_state``. + + Returns the FFN's EPLB schedule and the initial + ``(logical_to_physical_map, logical_replica_count)`` CPU tensors. + """ + holder: list[dict[str, int] | None] = [None] + dist.broadcast_object_list(holder, src=0, group=connector.bootstrap_group) + schedule = holder[0] + assert schedule is not None + missing = [key for key in _EPLB_SCHEDULE_FIELDS if key not in schedule] + if missing: + raise RuntimeError(f"FFN EPLB schedule is missing {missing}: {schedule}") + map_cpu = torch.empty( + ( + int(schedule["num_moe_layers"]), + int(schedule["num_logical_experts"]), + int(schedule["map_slots"]), + ), + dtype=torch.int64, + ) + dist.broadcast(map_cpu, src=0, group=connector.bootstrap_group) + count_cpu = torch.empty( + (int(schedule["num_moe_layers"]), int(schedule["num_logical_experts"])), + dtype=torch.int64, + ) + dist.broadcast(count_cpu, src=0, group=connector.bootstrap_group) + return schedule, map_cpu, count_cpu + + +def real_rows_in_tp_slice( + num_real_tokens: int, + *, + tp_size: int, + tp_rank: int, +) -> int: + """Real rows one Attention TP rank's dispatch slice holds. + + The DP rank's ``num_real_tokens`` rows are padded to a multiple of + ``tp_size`` and chunked; a chunk's real rows are the ones before + ``num_real_tokens``. Eager steps are exact. A graph-replayed batch is + padded to its capture size, so with TP > 1 the chunk boundary is computed + from the real count: the total recorded load stays exact (per-expert + attribution may shift by the padding rows), and TP 1 is exact in every + mode. Padding rows route to no expert and are never recorded. + """ + chunk = -(-num_real_tokens // tp_size) + return min(max(num_real_tokens - tp_rank * chunk, 0), chunk) + + +def _single_model_state(eplb_state: EplbState) -> EplbModelState: + """The one ``EplbModelState`` an AFD FFN registers (speculative decoding + and draft models are rejected by AFD validation, and a state without any + model was never registered).""" + if not eplb_state.model_states: + raise RuntimeError( + "DeepEP M2N EPLB needs the FFN model registered with EPLB; " + "the EplbState holds no model" + ) + if len(eplb_state.model_states) > 1: + raise RuntimeError( + "DeepEP M2N EPLB supports exactly one registered model, got " + f"{len(eplb_state.model_states)}" + ) + return next(iter(eplb_state.model_states.values())) + + +class DeepEPM2NAttentionEplb: + """Attention-side EPLB recorder for one DeepEP M2N deployment. + + Owns stable-address GPU tensors that the routing path (inside CUDA graphs) + reads and writes: the per-layer expert load buffer, the sliding load + window, the logical-to-physical maps adopted from the FFN, and the + per-step ``record_enabled`` / ``num_unpadded_tokens`` scalars. The + connector's control plane calls :meth:`on_control_step` once per + transaction, before the transaction's forward. + """ + + def __init__( + self, + connector: DeepEPM2NAFDConnector, + schedule: dict[str, int], + moe_layer_indices: tuple[int, ...], + initial_logical_map: torch.Tensor, + initial_replica_count: torch.Tensor, + ) -> None: + self._connector = connector + self._device = connector.device + self._step_interval = int(schedule["step_interval"]) + self._log_interval = int(schedule["log_interval"]) + window_size = int(schedule["window_size"]) + num_layers = int(schedule["num_moe_layers"]) + num_logical = int(schedule["num_logical_experts"]) + num_physical = int(schedule["num_physical_experts"]) + self._map_slots = int(schedule["map_slots"]) + if num_layers != len(moe_layer_indices): + raise RuntimeError( + f"FFN reports {num_layers} MoE layers, this Attention model " + f"runs {len(moe_layer_indices)}" + ) + self._layer_rows = { + layer_idx: row for row, layer_idx in enumerate(moe_layer_indices) + } + # Read and written by the routing kernel; stable addresses, updated + # in place so captured CUDA graphs follow the current placement. + self.expert_load_pass = torch.zeros( + (num_layers, num_physical), dtype=torch.int32, device=self._device + ) + self.expert_load_window = torch.zeros( + (window_size, num_layers, num_physical), + dtype=torch.int32, + device=self._device, + ) + self.logical_to_physical_map = torch.zeros( + (num_layers, num_logical, self._map_slots), + dtype=torch.int64, + device=self._device, + ) + self.logical_replica_count = torch.zeros( + (num_layers, num_logical), dtype=torch.int64, device=self._device + ) + self.record_enabled = torch.ones((), dtype=torch.bool, device=self._device) + # DBO is rejected with EPLB, so one ubatch is enough. + self.num_unpadded_tokens = [ + torch.zeros((), dtype=torch.int32, device=self._device) + ] + self.logical_to_physical_map.copy_(initial_logical_map.to(self._device)) + self.logical_replica_count.copy_(initial_replica_count.to(self._device)) + self._window_step = 0 + self._rearrange_step = int(schedule["initial_rearrangement_step"]) + self._previous_step_is_dummy: bool | None = None + + # ============================== + # Control-plane stepping + # ============================== + + def on_control_step(self, payload: AFDControlPayload) -> None: + """Step the recorder for the transaction that just finished, then + prepare the recording state for the transaction about to run. + + The FFN steps EPLB after its forward, while the Attention control + plane publishes before its forward, so the finished transaction's + flags are the previous payload's. The very first transaction has no + predecessor and only prepares its own recording state. + """ + if self._previous_step_is_dummy is not None: + self._step_finished_transaction(self._previous_step_is_dummy) + self._prepare_current_transaction(payload) + self._previous_step_is_dummy = payload.is_warmup or payload.is_graph_capturing + + def _step_finished_transaction(self, is_dummy: bool) -> None: + """Mirror the FFN's per-step EPLB transitions. + + Dummy steps (warmup, graph capture) never enter the window; every + step advances the rearrangement counter, exactly like vLLM's + ``EplbState.step``, so both roles detect the boundary together. + """ + if not is_dummy: + # Record every real step: the window then always holds the last + # ``window_size`` real steps, the same set vLLM's + # should_record schedule would have recorded before a rearrange. + self.expert_load_window[self._window_step].copy_(self.expert_load_pass) + window_size = self.expert_load_window.shape[0] + self._window_step = (self._window_step + 1) % window_size + self._rearrange_step += 1 + if self._log_interval and self._rearrange_step % self._log_interval == 0: + self._exchange_window_load() + # The pass holds this step's load until here: dummy steps recorded + # nothing, so dropping it is a no-op for them. + self.expert_load_pass.zero_() + if self._rearrange_step >= self._step_interval: + self._rearrange_step = 0 + self._exchange_at_boundary() + + def _exchange_window_load(self) -> None: + """Join the loss-logging exchange. + + Contributes this rank's load window aggregated into logical experts, + which is all EPLB's policy looks at, so the FFN's world rank 0 can + score the deployed map on the same traffic that drives the next + rearrangement. Shipped every ``log_interval`` steps instead of at + every boundary, which keeps the figures a window average rather than + one step's noise. The tensor is CPU-staged because the connector's + bootstrap group is gloo. + """ + load_cpu = logical_load_from_window( + self.expert_load_window, + self.logical_to_physical_map, + self.logical_replica_count, + ).cpu() + dist.reduce(load_cpu, dst=0, group=self._connector.bootstrap_group) + + def _prepare_current_transaction(self, payload: AFDControlPayload) -> None: + """Gate recording and fill the real token count for the upcoming + forward from the payload this rank is about to execute.""" + is_dummy = payload.is_warmup or payload.is_graph_capturing + self.record_enabled.fill_(not is_dummy) + self.num_unpadded_tokens[0].fill_( + real_rows_in_tp_slice( + payload.num_real_tokens, + tp_size=self._connector.attn_tp_size or 1, + tp_rank=self._connector.attention_tp_rank, + ) + ) + + def _exchange_at_boundary(self) -> None: + """Join the boundary exchange: contribute the load window, adopt the + FFN's new maps. Blocks until the FFN finished its rearrangement.""" + group = self._connector.bootstrap_group + # Copied, not just moved: the window is cleared below once the new map + # is in place, and on a CPU device ``.cpu()`` would alias it. + window_cpu = torch.empty( + self.expert_load_window.shape, dtype=self.expert_load_window.dtype + ) + window_cpu.copy_(self.expert_load_window) + dist.reduce(window_cpu, dst=0, group=group) + map_cpu = torch.empty(self.logical_to_physical_map.shape, dtype=torch.int64) + dist.broadcast(map_cpu, src=0, group=group) + count_cpu = torch.empty(self.logical_replica_count.shape, dtype=torch.int64) + dist.broadcast(count_cpu, src=0, group=group) + self.logical_to_physical_map.copy_(map_cpu.to(self._device)) + self.logical_replica_count.copy_(count_cpu.to(self._device)) + # Slot identities changed with the placement: load recorded under the + # old map cannot be read through the new one, and EPLB does not need + # it either (a window is fully refreshed before the next boundary + # whenever window_size <= step_interval). Start the next window clean + # so every logged figure describes traffic routed by one map. + self.expert_load_window.zero_() + self._window_step = 0 + + # ============================== + # Routing path + # ============================== + + def apply_routing(self, topk_ids: torch.Tensor, layer_idx: int) -> torch.Tensor: + """Map logical top-k ids to physical ids and record the expert load.""" + row = self._layer_rows[layer_idx] + return eplb_map_to_physical_and_record( + topk_ids=topk_ids, + expert_load_view=self.expert_load_pass[row], + logical_to_physical_map=self.logical_to_physical_map[row], + logical_replica_count=self.logical_replica_count[row], + record_enabled=self.record_enabled, + num_unpadded_tokens=self.num_unpadded_tokens[0], + ) + + +class DeepEPM2NFFNEplb: + """FFN-side EPLB rearrangement driver for one DeepEP M2N deployment. + + Steps a rearrangement counter (seeded from the ``EplbState`` so the first + rearrangement keeps vLLM's initial three-quarter progress) and, at each + boundary, runs the exchange described in the module docstring and calls + ``EplbState.rearrange`` directly: on the routed-experts path nothing on + the FFN records load, so ``EplbState.step``'s window bookkeeping has no + input and its rearrangement trigger is replaced by this counter. With + ``log_balancedness`` the same counter also paces the loss log, which + changes no state. + """ + + def __init__( + self, + connector: DeepEPM2NAFDConnector, + eplb_state: EplbState, + ) -> None: + self._connector = connector + self._state = eplb_state + self._model_state = _single_model_state(eplb_state) + self._step_interval = eplb_state.expert_rearrangement_step_interval + self._rearrange_step = eplb_state.expert_rearrangement_step + # vLLM's log_balancedness flag gates placement logging: on this path + # no EplbState.step runs, so the flag cannot log from there. It also + # enables the loss log: every ``log_balancedness_interval`` steps the + # deployed map is scored on the Attention ranks' load window, and at + # each rearrangement the replaced and the committed map are scored on + # that same window. + eplb_config = connector.vllm_config.parallel_config.eplb_config + self._log_placement = bool(eplb_config.log_balancedness) + self._log_interval = ( + int(eplb_config.log_balancedness_interval) if self._log_placement else 0 + ) + + def broadcast_initial_state(self) -> None: + """Collective: tell every rank the EPLB schedule and initial maps. + + Called from every rank's ``init_afd_connector``; FFN world rank 0 is + the source. The Attention recorder is built from the result. + """ + schedule: dict[str, int] = { + "num_moe_layers": int(self._model_state.model.num_moe_layers), + "num_logical_experts": int(self._model_state.model.num_logical_experts), + "num_physical_experts": int(self._model_state.model.num_physical_experts), + "map_slots": int(self._model_state.logical_to_physical_map.shape[2]), + "window_size": int(self._model_state.expert_load_window.shape[0]), + "step_interval": int(self._step_interval), + "initial_rearrangement_step": int(self._rearrange_step), + "log_interval": int(self._log_interval), + } + holder = [schedule] + dist.broadcast_object_list(holder, src=0, group=self._connector.bootstrap_group) + map_cpu = self._model_state.logical_to_physical_map.cpu() + dist.broadcast(map_cpu, src=0, group=self._connector.bootstrap_group) + count_cpu = self._model_state.logical_replica_count.cpu() + dist.broadcast(count_cpu, src=0, group=self._connector.bootstrap_group) + + def step(self) -> None: + """One connector step; logs the map's loss or rearranges when due.""" + self._rearrange_step += 1 + if self._log_interval and self._rearrange_step % self._log_interval == 0: + self._log_window_loss() + if self._rearrange_step < self._step_interval: + return + self._rearrange_step = 0 + self._rearrange() + + def _log_window_loss(self) -> None: + """Join the loss-logging exchange and log it on world rank 0. + + Attention ranks contribute their load window aggregated into logical + experts; FFN ranks hold no load, so they only contribute the zero + tensor the reduce writes into on rank 0. + """ + num_logical = self._model_state.logical_to_physical_map.shape[1] + window_dtype = self._model_state.expert_load_window.dtype + staging = torch.zeros( + (int(self._model_state.model.num_moe_layers), num_logical), + dtype=window_dtype, + ) + dist.reduce(staging, dst=0, group=self._connector.bootstrap_group) + if self._connector.world_rank == 0: + self._log_deployed_loss(staging) + + def _log_deployed_loss(self, logical_load: torch.Tensor) -> None: + """Log the deployed map's EPLB loss on the recorded window. + + ``logical_load`` is the reduced per-logical-expert load, + shape ``(num_moe_layers, num_logical)``. The loss is the figure + ``DefaultEplbPolicy`` optimises, so the series moves when a + rearrangement changes the placement and stays flat in between. + """ + summary = eplb_map_loss( + logical_load, + self._model_state.physical_to_logical_map.cpu(), + ffn_size=self._connector.ffn_size, + ) + if summary.max_tokens <= 0: + # Dummy steps (warmup, graph capture) record no load. + return + window_size = self._model_state.expert_load_window.shape[0] + recorded_steps = min(self._rearrange_step, window_size) + logger.info( + "AFD FFN EPLB step %d: loss=%.4f over %d steps " + "(avg_tokens=%.2f, max_tokens=%d), " + "steps until the next rearrangement: %d", + self._rearrange_step, + summary.balancedness, + recorded_steps, + summary.avg_tokens, + int(summary.max_tokens), + self._step_interval - self._rearrange_step, + ) + + def _rearrange(self) -> None: + window = self._model_state.expert_load_window + # FFN ranks contribute nothing: rank 0's tensor receives the Attention + # window sum, the others stay zero for the EP-group all-reduce inside + # EplbState.rearrange to redistribute exactly one copy. + staging = torch.zeros(window.shape, dtype=window.dtype) + dist.reduce(staging, dst=0, group=self._connector.bootstrap_group) + if self._connector.world_rank == 0: + window.copy_(staging.to(window.device)) + logs_loss = self._log_placement and self._connector.world_rank == 0 + deployed_maps = ( + ( + self._model_state.physical_to_logical_map.cpu(), + self._model_state.logical_to_physical_map.cpu(), + self._model_state.logical_replica_count.cpu(), + ) + if logs_loss + else None + ) + self._state.rearrange() + if logs_loss: + # Rearrangement commits the new map in place, so the loss log + # compares it with the copies saved above. + self._log_chosen_experts(staging) + self._log_map_loss(staging, deployed_maps) + map_cpu = self._model_state.logical_to_physical_map.cpu() + dist.broadcast(map_cpu, src=0, group=self._connector.bootstrap_group) + count_cpu = self._model_state.logical_replica_count.cpu() + dist.broadcast(count_cpu, src=0, group=self._connector.bootstrap_group) + + def _log_chosen_experts(self, load_window: torch.Tensor) -> None: + """Log the placement EPLB just committed and the load that chose it. + + ``load_window`` is the reduced Attention window (physical expert + load per recorded step) this rearrangement decided from, shape + ``(window_size, num_moe_layers, num_physical_experts)``. + """ + ffn_size = self._connector.ffn_size + lines = describe_eplb_placement( + self._model_state.physical_to_logical_map.cpu(), + ffn_size=ffn_size, + ) + logger.info( + "DeepEP M2N EPLB chose experts (recorded load over %d steps: %s):\n%s", + load_window.shape[0], + describe_eplb_rank_load(load_window, ffn_size=ffn_size), + "\n".join(lines), + ) + + def _log_map_loss( + self, + load_window: torch.Tensor, + deployed_maps: tuple[torch.Tensor, torch.Tensor, torch.Tensor], + ) -> None: + """Log what the rearrangement did to the quantity EPLB optimises. + + Both figures use EPLB's own load model (a logical expert's window load + divided evenly over its replicas) on the same window, so the gain is + the policy's improvement and not workload drift. ``deployed_maps`` + holds the physical-to-logical map, the logical-to-physical map and the + replica counts in force while the window was recorded; the state has + already committed the replacements. + """ + ffn_size = self._connector.ffn_size + deployed_map, recording_map, recording_counts = deployed_maps + logical_load = logical_load_from_window( + load_window, recording_map, recording_counts + ) + committed_map = self._model_state.physical_to_logical_map.cpu() + before = eplb_map_loss(logical_load, deployed_map, ffn_size=ffn_size) + after = eplb_map_loss(logical_load, committed_map, ffn_size=ffn_size) + gain = ( + 0.0 + if before.balancedness == 0 + else 100.0 * (after.balancedness / before.balancedness - 1.0) + ) + logger.info( + "AFD FFN EPLB loss: before=%.4f, after=%.4f, gain=%+.2f%% " + "(same %d-step window; per-layer mean/max of the load EPLB " + "estimates by splitting each expert over its replicas)", + before.balancedness, + after.balancedness, + gain, + load_window.shape[0], + ) + + +__all__ = [ + "DeepEPM2NAttentionEplb", + "DeepEPM2NFFNEplb", + "describe_eplb_placement", + "real_rows_in_tp_slice", + "recv_initial_eplb_state", +] diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 02351e7c8..74f0ad671 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -51,7 +51,7 @@ def sp_local_sizes( sequence_parallel_size: int, pcp_size: int = 1, use_ep: bool = False, - ) -> Generator[list[int]]: + ) -> Generator[list[int], None, None]: num_tokens = self.num_tokens_across_dp_cpu if pcp_size > 1: num_tokens = ( @@ -87,7 +87,7 @@ def chunked_sizes( sequence_parallel_size: int, max_chunk_size_per_rank: int, chunk_idx: int, - ) -> Generator[list[int]]: + ) -> Generator[list[int], None, None]: sp_tokens = ( ( (self.num_tokens_across_dp_cpu + sequence_parallel_size - 1) @@ -141,6 +141,10 @@ class AFDControlPayload: is_profile: Whether the payload belongs to an initial memory-profile forward. NPU FFN workers use this to preserve Ascend's balanced dummy MoE routing in the split-process AFD execution model. + num_real_tokens: The publishing Attention rank's real (unpadded) token + count for the first stage. Graph capture and replay payloads carry + padded DP counts, so connectors that need real counts (DeepEP M2N + EPLB load recording) read this instead. ``0`` when unknown. """ dp_metadata_list: dict[int, AFDDPMetadata] @@ -148,6 +152,7 @@ class AFDControlPayload: is_warmup: bool is_graph_replaying: bool = False is_profile: bool = False + num_real_tokens: int = 0 def __post_init__(self) -> None: self.dp_metadata_list = { @@ -402,6 +407,7 @@ def encode_control_payload(payload: AFDControlPayload) -> bytes: "is_warmup": bool(payload.is_warmup), "is_graph_replaying": bool(payload.is_graph_replaying), "is_profile": bool(payload.is_profile), + "num_real_tokens": int(payload.num_real_tokens), } return json.dumps(wire_payload, separators=(",", ":"), sort_keys=True).encode( "utf-8", @@ -432,6 +438,7 @@ def decode_control_payload(payload_bytes: bytes) -> AFDControlPayload: is_warmup=bool(payload.get("is_warmup", False)), is_graph_replaying=bool(payload.get("is_graph_replaying", False)), is_profile=bool(payload.get("is_profile", False)), + num_real_tokens=int(payload.get("num_real_tokens", 0)), ) diff --git a/afd_plugin/distributed/eplb.py b/afd_plugin/distributed/eplb.py new file mode 100644 index 000000000..503321ea0 --- /dev/null +++ b/afd_plugin/distributed/eplb.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Shared helpers for expert-parallel load balancing (EPLB) observability. + +Used by both EPLB paths: the P2P connector's FFN runner (routing and load +recording run on the FFN, so vLLM's ``EplbState`` is stepped natively) and +the DeepEP M2N connector's coordinators (routing runs on Attention, so the +plugin drives the rearrangement itself). +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import torch + + +@dataclass(frozen=True) +class EplbLoadSummary: + """Expert-parallel load figures, summed over MoE layers. + + ``avg_tokens`` and ``max_tokens`` are per-layer means and maxima across + ranks, summed over layers; ``balancedness`` is their ratio, so 1.0 is + perfectly balanced. The window the figures describe is whatever the caller + scored: one step, or an EPLB load window. + """ + + avg_tokens: float + max_tokens: float + balancedness: float + + +def describe_eplb_placement( + physical_to_logical_map: torch.Tensor, + *, + ffn_size: int, +) -> list[str]: + """One line per MoE layer describing the placement EPLB committed. + + Each line names, per FFN rank, the logical experts its physical slots + hold, and which logical experts are replicated (hold more than one slot). + ``physical_to_logical_map`` is the committed map, shape + ``(num_moe_layers, num_physical_experts)``; FFN rank ``g`` owns the + contiguous slot block ``[g * P / F, (g + 1) * P / F)``. + """ + num_physical = physical_to_logical_map.shape[1] + slots_per_rank = num_physical // ffn_size + lines: list[str] = [] + for layer_idx, layer_map in enumerate(physical_to_logical_map.tolist()): + parts = [] + for rank in range(ffn_size): + slots = layer_map[rank * slots_per_rank : (rank + 1) * slots_per_rank] + experts = ", ".join(map(str, sorted(slots))) + parts.append(f"rank {rank}=[{experts}]") + replica_counts: dict[int, int] = {} + for logical in layer_map: + if logical >= 0: + replica_counts[logical] = replica_counts.get(logical, 0) + 1 + replicated = { + logical: count + for logical, count in sorted(replica_counts.items()) + if count > 1 + } + replicas = ( + "none" + if not replicated + else ", ".join( + f"{logical}x{count}" for logical, count in replicated.items() + ) + ) + lines.append(f"layer {layer_idx}: {' | '.join(parts)} | replicas: {replicas}") + return lines + + +def summarize_rank_load(num_tokens_per_rank: torch.Tensor) -> EplbLoadSummary: + """Balancedness of a per-rank load table, as vLLM's EPLB computes it. + + ``num_tokens_per_rank`` has shape ``(num_moe_layers, ffn_size)``. Per layer + the mean and the max rank load are taken across ranks and both are summed + over layers; balancedness is their ratio, so 1.0 is perfectly balanced. + Recorded load is integer, so the mean is taken in float, as vLLM does. + """ + num_tokens_per_rank = num_tokens_per_rank.float() + avg_tokens = float(num_tokens_per_rank.mean(dim=1).sum()) + max_tokens = float(num_tokens_per_rank.max(dim=1).values.sum()) + balancedness = avg_tokens / max_tokens if max_tokens > 0 else 0.0 + return EplbLoadSummary( + avg_tokens=avg_tokens, + max_tokens=max_tokens, + balancedness=balancedness, + ) + + +def logical_load_from_window( + expert_load_window: torch.Tensor, + logical_to_physical_map: torch.Tensor, + logical_replica_count: torch.Tensor, +) -> torch.Tensor: + """Aggregate a recorded load window into logical experts. + + ``expert_load_window`` is the recorded load, shape + ``(window_size, num_moe_layers, num_physical_experts)``, indexed by the + slots of ``logical_to_physical_map`` (``(num_moe_layers, num_logical, + max_replicas)``, valid entries designated by ``logical_replica_count``). + The window is summed over its steps and gathered through the map, giving + ``(num_moe_layers, num_logical)`` tokens per logical expert, which is the + only load EPLB's policy looks at. + """ + num_layers, num_logical, max_replicas = logical_to_physical_map.shape + slot_is_used = torch.arange( + max_replicas, device=logical_to_physical_map.device + ).unsqueeze(0).unsqueeze(1) < logical_replica_count.unsqueeze(-1) + used_slots = logical_to_physical_map.masked_fill(~slot_is_used, 0).long() + window_load = expert_load_window.sum(dim=0) + per_slot = window_load.unsqueeze(1).expand(num_layers, num_logical, -1) + per_replica = per_slot.gather(2, used_slots).masked_fill(~slot_is_used, 0) + return per_replica.sum(dim=-1).to(expert_load_window.dtype) + + +def eplb_map_loss( + logical_load: torch.Tensor, + candidate_map: torch.Tensor, + *, + ffn_size: int, +) -> EplbLoadSummary: + """Score one placement with the quantity ``DefaultEplbPolicy`` minimises. + + The policy never sees realised per-slot counts: it aggregates the window + into logical experts, divides each expert's load evenly over the replicas + the *candidate* map gives it, and packs those per-replica loads into ranks. + This function evaluates that model in the per-layer mean/max form EPLB logs + balancedness in, so a deployed map and the policy's proposal are directly + comparable on one window. + + ``logical_load`` is ``(num_moe_layers, num_logical)`` tokens per logical + expert, as ``logical_load_from_window`` produces. ``candidate_map`` is the + placement to score, ``(num_moe_layers, num_physical_experts)``; slots + mapped to ``-1`` carry no load and no replica. + """ + logical_load = logical_load.cpu().float() + num_layers, num_logical = logical_load.shape + num_slots = candidate_map.shape[1] + # Unused slots are gathered into a trailing guard column so they neither + # add load nor inflate an expert's replica count. + slot_is_used = candidate_map.cpu() >= 0 + candidate_index = torch.where( + slot_is_used, + candidate_map.cpu().long(), + torch.full_like(candidate_map.cpu().long(), num_logical), + ) + logical_load = torch.cat( + [logical_load, torch.zeros((num_layers, 1), dtype=logical_load.dtype)], dim=1 + ) + replica_count = torch.zeros_like(logical_load) + replica_count.scatter_add_(1, candidate_index, slot_is_used.float()) + load_per_replica = logical_load.div(replica_count.clamp_min(1)).gather( + 1, candidate_index + ) + load_per_rank = ( + load_per_replica.mul(slot_is_used.float()) + .reshape(num_layers, ffn_size, num_slots // ffn_size) + .sum(dim=-1) + ) + return summarize_rank_load(load_per_rank) + + +def describe_eplb_rank_load( + expert_load_window: torch.Tensor, + *, + ffn_size: int, +) -> str: + """Per-FFN-rank token load recorded in one EPLB load window. + + ``expert_load_window`` is the recorded load, shape + ``(window_size, num_moe_layers, num_physical_experts)``; a rank's load is + the sum over its contiguous slot block. + """ + slots_per_rank = expert_load_window.shape[2] // ffn_size + rank_loads = expert_load_window.sum(dim=(0, 1)).tolist() + parts = [] + for rank in range(ffn_size): + slots = rank_loads[rank * slots_per_rank : (rank + 1) * slots_per_rank] + parts.append(f"rank {rank}={int(sum(slots))} tokens") + return ", ".join(parts) + + +__all__ = [ + "EplbLoadSummary", + "describe_eplb_placement", + "describe_eplb_rank_load", + "eplb_map_loss", + "logical_load_from_window", + "summarize_rank_load", +] diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 54f115f34..b0313f77f 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -34,6 +34,7 @@ from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDForwardContextMetadata, + AFDRoutedExpertsConnector, AFDTransferContext, AFDTransferMetadata, ) @@ -153,9 +154,11 @@ class RemoteFFNProxy(nn.Module): on Attention (an M2N routed-experts connector such as DeepEP M2N) sets ``gate``, ``top_k``, ``renormalize`` and ``tp_size`` after construction, from its model and parallel config. The proxy then computes the gate and - vLLM's default fused top-k (``FusedTopKRouter``; EPLB is not supported) - and sends ``topk_ids`` and ``topk_weights``; combine returns the top-k - weighted sum, which replaces the native MoE block's output. + vLLM's default fused top-k (``FusedTopKRouter``) and sends ``topk_ids`` + and ``topk_weights``; combine returns the top-k weighted sum, which + replaces the native MoE block's output. When the FFN runs EPLB, the + connector maps the logical top-k ids to physical expert ids here + (``AFDRoutedExpertsConnector.apply_eplb_routing``). With Attention TP > 1 every TP rank holds the same tokens after the attention all-reduce; dispatching them from every rank would make the FFN @@ -218,6 +221,15 @@ def _send( ) -> int: """Send this layer's Attention output; return the stage to receive.""" afd_metadata = self._afd_metadata() + # With EPLB the connector places experts dynamically: map the logical + # top-k ids to physical slots (and record load) before dispatch. + if "topk_ids" in send_kwargs: + connector = afd_metadata.connector + if isinstance(connector, AFDRoutedExpertsConnector): + send_kwargs["topk_ids"] = connector.apply_eplb_routing( + send_kwargs["topk_ids"], + self.layer_idx, + ) forward_context = get_forward_context() stage_idx = int( getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), diff --git a/afd_plugin/model_executor/models/qwen3_5.py b/afd_plugin/model_executor/models/qwen3_5.py index fbfb9ae0a..463d0d7fa 100644 --- a/afd_plugin/model_executor/models/qwen3_5.py +++ b/afd_plugin/model_executor/models/qwen3_5.py @@ -111,7 +111,10 @@ def __init__( if parallel_config.use_sequence_parallel_moe: raise RuntimeError("AFD Qwen3.5/3.6 does not support SP MoE") if parallel_config.enable_eplb: - raise RuntimeError("AFD Qwen3.5/3.6 does not support EPLB") + raise RuntimeError( + "AFD Qwen3.5/3.6 remote experts do not support EPLB on the " + "Attention role", + ) self.tp_size = next_native.get_tensor_model_parallel_world_size() self.ep_group = next_native.get_ep_group().device_group @@ -294,8 +297,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): raise RuntimeError("AFD Qwen model requires AFD activation") if vllm_config.parallel_config.use_sequence_parallel_moe: raise RuntimeError("AFD Qwen3.5/3.6 does not support SP MoE") - if vllm_config.parallel_config.enable_eplb: - raise RuntimeError("AFD Qwen3.5/3.6 does not support EPLB") + if vllm_config.parallel_config.enable_eplb and afd_config.role != "ffn": + raise RuntimeError("AFD Qwen3.5/3.6 supports EPLB only on the FFN role") self.afd_config = afd_config # ### PATCH END @@ -304,8 +307,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_text_config self.config = config self.quant_config = vllm_config.quant_config - # ### PATCH START: remote experts do not participate in EPLB. - self.num_redundant_experts = 0 + # ### PATCH START: mirror the native loader-facing EPLB metadata; the + # Attention process runs without EPLB flags so its value stays zero. + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) # ### PATCH END self.vocab_size = config.vocab_size @@ -479,8 +485,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"): parallel_config = vllm_config.parallel_config if parallel_config.use_sequence_parallel_moe: raise ValueError("AFD Qwen3.5/3.6 CUDA does not support SP MoE") - if parallel_config.enable_eplb: - raise ValueError("AFD Qwen3.5/3.6 CUDA does not support EPLB") + if parallel_config.enable_eplb and afd_config.role != "ffn": + raise ValueError("AFD Qwen3.5/3.6 supports EPLB only on the FFN role") if parallel_config.pipeline_parallel_size != 1: raise ValueError( "AFD Qwen3.5/3.6 CUDA supports pipeline_parallel_size=1 only", diff --git a/afd_plugin/model_executor/models/qwen3_moe.py b/afd_plugin/model_executor/models/qwen3_moe.py index 54dc0f406..db18488e0 100644 --- a/afd_plugin/model_executor/models/qwen3_moe.py +++ b/afd_plugin/model_executor/models/qwen3_moe.py @@ -144,8 +144,8 @@ def _validate_supported_config( raise RuntimeError( "AFD Qwen3 MoE does not support sequence-parallel MoE", ) - if parallel_config.enable_eplb: - raise RuntimeError("AFD Qwen3 MoE does not support EPLB") + if parallel_config.enable_eplb and afd_config.role != "ffn": + raise RuntimeError("AFD Qwen3 MoE supports EPLB only on the FFN role") if parallel_config.pipeline_parallel_size != 1: raise RuntimeError("AFD Qwen3 MoE does not support pipeline parallelism") if vllm_config.speculative_config is not None: diff --git a/afd_plugin/v1/worker/attention_metadata.py b/afd_plugin/v1/worker/attention_metadata.py index 6f87ffb14..692a01b2d 100644 --- a/afd_plugin/v1/worker/attention_metadata.py +++ b/afd_plugin/v1/worker/attention_metadata.py @@ -105,12 +105,21 @@ def send_dp_metadata( # callers compatible with runners created before this mixin existed. is_graph_capturing = getattr(self, "_afd_is_graph_capturing", False) is_graph_replaying = getattr(self, "_afd_is_graph_replaying", False) + # Graph capture and replay payloads carry padded DP counts; the real + # count of this rank's transaction lives on the pending sidecar. + # Connectors that record per-expert load (DeepEP M2N EPLB) read it. + num_real_tokens = 0 + if self._afd_pending_metadata is not None: + unpadded = self._afd_pending_metadata.tokens_unpadded_lens + if unpadded: + num_real_tokens = int(unpadded[0]) payload = AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, is_graph_replaying=is_graph_replaying, is_profile=self._afd_is_profile, + num_real_tokens=num_real_tokens, ) self.connector.control_plane.update_state_from_dp_metadata(payload) self.connector.control_plane.send_dp_metadata_list(payload) diff --git a/afd_plugin/v1/worker/ffn_metadata.py b/afd_plugin/v1/worker/ffn_metadata.py index 50a37f638..759ee076e 100644 --- a/afd_plugin/v1/worker/ffn_metadata.py +++ b/afd_plugin/v1/worker/ffn_metadata.py @@ -12,6 +12,31 @@ from vllm.forward_context import ForwardContext +def _expand_counts_across_attention_ranks( + attention_counts: tuple[int, ...], + *, + attention_size: int, +) -> tuple[int, ...]: + """Expand per-DP-rank counts to one count per Attention rank. + + vLLM DP metadata carries one count per DP rank while the AFD Attention + role counts TP-derived worker ranks, so each count covers the TP peers + of its DP rank. Counts already covering every Attention rank are kept. + """ + + if ( + attention_counts + and len(attention_counts) < attention_size + and attention_size % len(attention_counts) == 0 + ): + attention_ranks_per_count = attention_size // len(attention_counts) + return tuple( + attention_counts[rank // attention_ranks_per_count] + for rank in range(attention_size) + ) + return attention_counts + + def aggregate_ffn_token_counts( attention_counts: tuple[int, ...], *, @@ -34,17 +59,10 @@ def aggregate_ffn_token_counts( if ffn_size <= 0 or attention_size < ffn_size: return fallback_counts - expanded_counts = attention_counts - if ( - attention_counts - and len(attention_counts) < attention_size - and attention_size % len(attention_counts) == 0 - ): - attention_ranks_per_count = attention_size // len(attention_counts) - expanded_counts = tuple( - attention_counts[rank // attention_ranks_per_count] - for rank in range(attention_size) - ) + expanded_counts = _expand_counts_across_attention_ranks( + attention_counts, + attention_size=attention_size, + ) # Each FFN rank sums the block that ``build_rank_mapping`` assigns to it. return tuple( @@ -60,6 +78,28 @@ def aggregate_ffn_token_counts( ) +def real_attention_token_total( + attention_counts: tuple[int, ...], + *, + attention_size: int, +) -> int: + """Total real Attention tokens for one stage, excluding placeholder rows. + + The connector transfers at least one row per Attention peer, so + zero-token peers contribute placeholder rows to the FFN batch. This + helper counts only real tokens: FFN-side EPLB uses the total as the + router's unpadded token count, because every FFN rank's router view + covers these tokens either through TP-replicated delivery or through + the MoE-level all-gather. + """ + + expanded_counts = _expand_counts_across_attention_ranks( + attention_counts, + attention_size=attention_size, + ) + return sum(max(0, int(count)) for count in expanded_counts) + + def project_ffn_token_counts_to_dp( ffn_counts: tuple[int, ...], *, @@ -78,7 +118,11 @@ def project_ffn_token_counts_to_dp( return tuple(ffn_counts[rank * ffn_ranks_per_dp] for rank in range(dp_size)) -__all__ = ["aggregate_ffn_token_counts", "project_ffn_token_counts_to_dp"] +__all__ = [ + "aggregate_ffn_token_counts", + "project_ffn_token_counts_to_dp", + "real_attention_token_total", +] def _set_moe_layer_index(forward_context: ForwardContext, layer_idx: int) -> None: diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index bab6ac42a..b26e61811 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -11,11 +11,17 @@ from vllm.compilation.monitor import set_cudagraph_capturing_enabled from vllm.config import CUDAGraphMode, VllmConfig from vllm.config import update_config as update_vllm_config -from vllm.distributed.parallel_state import get_world_group, graph_capture +from vllm.distributed.parallel_state import ( + get_ep_group, + get_world_group, + graph_capture, +) from vllm.forward_context import DPMetadata, get_forward_context, set_forward_context +from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import _ROPE_DICT from vllm.model_executor.model_loader import get_model_loader from vllm.utils.mem_utils import DeviceMemoryProfiler +from vllm.v1.worker.gpu.eplb_utils import EPLBController from vllm.v1.worker.workspace import reset_workspace_manager from afd_plugin.compat.profiler import ( @@ -30,6 +36,10 @@ AFDControlPlane, AFDDPMetadata, ) +from afd_plugin.distributed.eplb import ( + describe_eplb_placement, + describe_eplb_rank_load, +) from afd_plugin.v1.worker.cuda_graph import ( AFDGraphRunMode, graph_run_mode, @@ -40,6 +50,7 @@ _set_moe_layer_index, aggregate_ffn_token_counts, project_ffn_token_counts_to_dp, + real_attention_token_total, ) from afd_plugin.validation import validate_gpu_model_runner_v2_config @@ -48,6 +59,8 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec +logger = init_logger(__name__) + class GPUFFNModelRunner: """FFN model runner for AFD GPU execution. @@ -91,6 +104,10 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: self.model: Any | None = None self.model_memory_usage = 0 self.num_layers = int(self.model_config.hf_text_config.num_hidden_layers) + # Expert parallelism load balancer. EPLB runs on the FFN role only: + # the Attention role holds no experts and no MoE routing. The + # controller no-ops unless ``parallel_config.enable_eplb`` is set. + self.eplb = EPLBController(vllm_config.parallel_config, device) self.use_cuda_graph = bool( self.afd_cudagraph_policy.enable_ffn_graph_cache, ) @@ -109,9 +126,10 @@ def initialize_afd_connector(self) -> None: self.connector.init_afd_connector() def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None: - """Load the vLLM model.""" + """Load the vLLM model and register it with FFN-side EPLB.""" model_loader = get_model_loader(self.load_config) + self.eplb.prepare_load() with DeviceMemoryProfiler() as profiler: if self.model is None: self.model = model_loader.load_model( @@ -124,6 +142,11 @@ def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None model_config=self.model_config, ) self.model_memory_usage = profiler.consumed_memory + eplb_models_added = self.eplb.maybe_register_model( + self.model, + self.model_config, + ) + self.eplb.maybe_start_async_loop(eplb_models_added, load_dummy_weights) def profile_run(self) -> None: pass @@ -163,6 +186,7 @@ def execute_model( step_afd_gpu_profiler(self.prof) if dp_metadata_list is None: raise RuntimeError("GPUFFNModelRunner requires dp_metadata_list") + self._eplb_prepare_forward(dp_metadata_list) graph_key = make_ffn_graph_key(dp_metadata_list) cuda_graph_info = self._cuda_graphs.get(graph_key) run_mode = graph_run_mode( @@ -174,6 +198,7 @@ def execute_model( ) if run_mode is AFDGraphRunMode.REPLAY: self._cuda_graphs[graph_key]["graph"].replay() + self.eplb_step(is_dummy=is_warmup or is_graph_capturing) return None # An eager fallback for a step Attention replays keeps the replay flag, @@ -185,8 +210,117 @@ def execute_model( is_warmup=is_warmup, is_graph_replaying=is_graph_replaying, ) + self.eplb_step(is_dummy=is_warmup or is_graph_capturing) return None + # Patch reason: vLLM steps EPLB from GPUModelRunner methods + # (sample_tokens/pool/dummy_run via step_eplb_after); the FFN runner is + # connector-driven and has no native runner methods to decorate, so the + # connector step paths call this seam directly after each forward. + # Patch functionality: step the FFN-side EPLB state once per connector + # step. Warmup and graph-capture steps are dummy steps: they keep the + # rearrangement counters of all FFN ranks in lockstep without recording + # their (garbage) expert load. + # Signature: mirrors GPUModelRunner.eplb_step; no added parameters. + # Upstream: vllm/v1/worker/gpu/eplb_utils.py, + # step_eplb_after. + def eplb_step(self, is_dummy: bool = False) -> None: + """Step the FFN-side EPLB state after one connector-driven step.""" + + placement_before = self._snapshot_eplb_placement_if_due() + self.eplb.step(is_dummy=is_dummy) + if placement_before is not None: + self._log_eplb_placement_if_changed(placement_before) + + # Placement logging (log_balancedness): a rearrangement commits new + # expert maps inside EplbState.step, at a counter boundary in sync mode + # or from the async worker's results while a transfer is pending, so the + # map is snapshotted only on steps where either can happen and the + # placement is logged when the map actually changed. Every FFN EP rank + # holds the same global map, so only EP rank 0 logs. + def _snapshot_eplb_placement_if_due(self) -> dict[str, torch.Tensor] | None: + """CPU snapshots of the committed expert maps when this step may + commit a rearrangement; ``None`` keeps the common step sync-free.""" + state = self.eplb.state + if state is None: + return None + # vLLM logs balancedness stats from EplbState.step with this flag; + # it also gates this runner's committed-placement log. + if not self.vllm_config.parallel_config.eplb_config.log_balancedness: + return None + # Every FFN EP rank holds the same global map; only rank 0 logs. + # ``rank_in_group`` is an attribute in vLLM's GroupCoordinator, + # the rank inside the EP group. + if get_ep_group().rank_in_group != 0: + return None + boundary_due = ( + state.expert_rearrangement_step + 1 + >= state.expert_rearrangement_step_interval + ) + async_pending = any( + model_state.rebalanced for model_state in state.model_states.values() + ) + if not boundary_due and not async_pending: + return None + return { + key: model_state.physical_to_logical_map.cpu().clone() + for key, model_state in state.model_states.items() + } + + def _log_eplb_placement_if_changed( + self, + placement_before: dict[str, torch.Tensor], + ) -> None: + state = self.eplb.state + assert state is not None + ep_size = self.vllm_config.parallel_config.data_parallel_size * ( + self.vllm_config.parallel_config.tensor_parallel_size + ) + for key, model_state in state.model_states.items(): + current = model_state.physical_to_logical_map.cpu() + if torch.equal(placement_before[key], current): + continue + logger.info( + "AFD FFN EPLB chose experts for model %s (recorded load over " + "%d steps: %s):\n%s", + model_state.model_name, + model_state.expert_load_window.shape[0], + describe_eplb_rank_load( + model_state.expert_load_window, ffn_size=ep_size + ), + "\n".join(describe_eplb_placement(current, ffn_size=ep_size)), + ) + + def _eplb_prepare_forward( + self, + dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], + ) -> None: + """Fill the EPLB unpadded token count for the router's record kernel. + + The count is the total real Attention tokens of the step: each FFN + rank's router view covers those tokens either through TP-replicated + delivery or through the MoE-level all-gather, so the total matches + the rows the map-and-record kernel processes. Placeholder rows that + idle Attention peers transfer are excluded. + """ + + if not self.vllm_config.parallel_config.enable_eplb: + return + if len(dp_metadata_list) != 1: + raise RuntimeError( + "AFD FFN EPLB supports exactly one stage per step; DBO and " + f"ubatching are not supported, got stages {sorted(dp_metadata_list)}", + ) + stage_metadata = next(iter(dp_metadata_list.values())) + attention_counts = tuple( + int(count) for count in stage_metadata.num_tokens_across_dp_cpu + ) + num_unpadded_tokens = real_attention_token_total( + attention_counts, + attention_size=int(self.connector.attn_size), + ) + self.eplb.prepare_forward(self.model_config, num_unpadded_tokens) + def _ffn_forward( self, *, @@ -386,6 +520,7 @@ def capture_model( return 0 if dp_metadata_list is None: raise RuntimeError("GPUFFNModelRunner.capture_model requires metadata") + self._eplb_prepare_forward(dp_metadata_list) start_free_gpu_memory = torch.cuda.mem_get_info()[0] if self._graph_memory_pool is None: @@ -421,6 +556,7 @@ def capture_model( end_free_gpu_memory = torch.cuda.mem_get_info()[0] cuda_graph_size = start_free_gpu_memory - end_free_gpu_memory + self.eplb_step(is_dummy=True) return int(cuda_graph_size) def sample_tokens(self, grammar_output: Any = None) -> Any: diff --git a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py index 55c26f41d..1f80aee77 100644 --- a/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py +++ b/afd_plugin/v1/worker/ffn_routed_experts_model_runner.py @@ -141,9 +141,16 @@ def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None } def initialize_afd_connector(self) -> None: + """Collective: share the loaded experts and, with EPLB, the expert + maps and rearrangement schedule the whole world routes by.""" self.routed_connector.init_afd_connector( moe_layer_indices=self.moe_layer_indices, local_experts=self.local_experts, + eplb_state=( + self.eplb.state + if self.vllm_config.parallel_config.enable_eplb + else None + ), ) def _ffn_forward( @@ -188,5 +195,17 @@ def _ffn_forward( connector.synchronize_lanes() return ffn_output + # The FFN role of an EPLB deployment never routes tokens, so nothing on + # this side records expert load and EplbState.step has no input. Instead + # the connector's FFN EPLB driver steps the rearrangement counter (every + # connector step, real or dummy, exactly like the Attention recorder) and + # at each boundary exchanges the Attention-recorded load window, runs + # EplbState.rearrange and broadcasts the committed expert maps. Without + # EPLB there is nothing to step. `is_dummy` is part of the base seam's + # signature; the counter advances on dummy steps too, so it is unused. + def eplb_step(self, is_dummy: bool = False) -> None: + """Step FFN-side EPLB for the routed-experts path.""" + self.routed_connector.step_ffn_eplb() + __all__ = ["GPUFFNRoutedExpertsModelRunner", "local_experts_from_vllm"] diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index bfef19f77..b826d1fca 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -162,10 +162,35 @@ def validate_gpu_model_runner_v2_config( if ( not parallel.enable_expert_parallel or parallel.enable_elastic_ep - or parallel.enable_eplb or parallel.use_sequence_parallel_moe ): raise RuntimeError("AFD ModelRunnerV2 requires static expert parallelism") + if parallel.enable_eplb: + # EPLB runs on the FFN role only: Attention holds no experts and no + # MoE routing, and its CUDA remote-expert shells reject EPLB. + if expected_role != "ffn": + raise RuntimeError( + "AFD ModelRunnerV2 supports EPLB only on the FFN role", + ) + # The FFN-side record kernel selects its per-ubatch unpadded token + # tensor through vLLM's DBO thread context, which the FFN daemon loop + # never installs, so DBO steps would read the wrong token count. + if parallel.enable_dbo or parallel.use_ubatching: + raise RuntimeError( + "AFD ModelRunnerV2 FFN EPLB does not support DBO or ubatching", + ) + # vLLM 0.32 async EPLB commits new expert maps from a background + # worker outside the connector step boundary: AFD steps EPLB from + # the connector loop, and the M2N exchange adopts the committed maps + # only at that boundary. Migration batching is async-only, so this + # gate covers it too. + if parallel.eplb_config.use_async: + raise RuntimeError( + "AFD ModelRunnerV2 FFN EPLB does not support async EPLB: " + "rearrangement commits expert maps from a background worker " + "outside the connector step boundary. Pass use_async=false " + "in --eplb-config.", + ) if ( parallel.enable_dbo or parallel.use_ubatching ) and parallel.num_ubatches != GPU_MRV2_NUM_UBATCHES: diff --git a/docker/Dockerfile.m2n b/docker/Dockerfile.m2n new file mode 100644 index 000000000..c865e20f0 --- /dev/null +++ b/docker/Dockerfile.m2n @@ -0,0 +1,121 @@ +# Target vLLM version: keep VLLM_BASE_TAG aligned with +# [project.optional-dependencies].vllm in pyproject.toml. +# +# Serving image for DeepEPM2NAFDConnector deployments (see +# docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md): the AFD plugin on top of the +# vLLM OpenAI image, plus the connector dependencies: +# - afdEP-comm, at the revision pinned by the plugin's `afdep` extra +# (uv.lock), and +# - DeepGEMM, built from source because it is not a plugin extra. +# +# Build from the repository root: +# docker build -f docker/Dockerfile.m2n -t afd-m2n . +# Optionally pin DeepGEMM for a reproducible image: +# docker build -f docker/Dockerfile.m2n --build-arg DEEPGEMM_REF= \ +# -t afd-m2n . +# +# The afdEP-comm transport extension and the DeepGEMM kernels compile just in +# time at first use on the GPUs (`python -m afdep doctor --build` needs a GPU, +# so it cannot run here). This image ships their build inputs instead: nvcc +# and the compilers from the base image, the NCCL wheel afdEP-comm builds +# against, and the elfutils headers (libdw-dev) that DeepGEMM's JIT needs. The +# first server start on a node therefore compiles kernels and is slower than +# later starts. +ARG VLLM_BASE_IMAGE=vllm/vllm-openai +ARG VLLM_BASE_TAG=v0.32.0 + +# deps stage: everything except the real project sources, so most PRs can +# cache-hit on the heavy dependency layer. +FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} AS deps + +ARG APP_DIR=/opt/afd-plugin +# DeepGEMM git ref to build (https://github.com/deepseek-ai/DeepGEMM). Empty +# builds the default branch, as the user guide documents. +ARG DEEPGEMM_REF= +# GPU architectures the DeepGEMM extension is compiled for. The connector +# requires SM90 or newer on every rank. +ARG DEEPGEMM_TORCH_CUDA_ARCH_LIST=9.0a;10.0a +WORKDIR ${APP_DIR} + +RUN ln -sf /usr/bin/python3 /usr/bin/python + +# Placeholder package tree so editable install can register APP_DIR. +# The final stage COPY --link fills this directory with the real sources. +RUN mkdir -p afd_plugin && touch afd_plugin/__init__.py + +# libdw-dev: elfutils headers (libdwfl.h) for DeepGEMM's JIT, needed both to +# build and at runtime. git: cloning DeepGEMM. +RUN apt-get update && \ + apt-get install -y --no-install-recommends git libdw-dev && \ + apt-get clean && \ + rm -rf /var/lib/apt/lists/* + +COPY pyproject.toml setup.py uv.lock README.md LICENSE ./ + +# Docker build context has no .git (see .dockerignore), so setuptools-scm +# cannot infer a version; pretend one for image installs. +ENV PYTHONDONTWRITEBYTECODE=1 \ + PYTHONUNBUFFERED=1 \ + AFD_BUILD_ASCEND_OPS=0 \ + PIP_DISABLE_PIP_VERSION_CHECK=1 \ + UV_LINK_MODE=copy \ + SETUPTOOLS_SCM_PRETEND_VERSION_FOR_VLLM_AFD_PLUGIN=0.0.0+m2n \ + SETUPTOOLS_SCM_PRETEND_VERSION=0.0.0+m2n + +# afdEP-comm, at the revision pinned by the `afdep` extra in uv.lock. +# --no-deps: its dependencies (torch, ninja, packaging, setuptools) are +# already provided by the base image. +RUN uv export --locked --extra afdep --no-group dev --no-group e2e-tests \ + --no-emit-project --no-header -o /tmp/requirements-afdep.txt && \ + uv pip install --system --no-deps \ + "$(grep -m1 '^afdep-comm @' /tmp/requirements-afdep.txt)" && \ + uv pip install --system -e . wheel + +RUN python -c "import torch; from importlib.metadata import version; \ + print('vLLM:', version('vllm'), 'PyTorch:', torch.__version__, 'CUDA:', torch.version.cuda); \ + assert version('vllm') == '0.32.0', 'M2N image requires vLLM 0.32.0'" + +# afdEP-comm builds against the NCCL wheel: point AFDEP_NCCL_ROOT at it and +# keep its lib directory on LD_LIBRARY_PATH for the served processes, as the +# user guide requires on both roles. The symlink keeps the path stable across +# base image updates. +RUN NCCL=$(python -c "import nvidia.nccl; print(nvidia.nccl.__path__[0])") && \ + ln -s "$NCCL" /opt/nccl +ENV AFDEP_NCCL_ROOT=/opt/nccl \ + LD_LIBRARY_PATH=/opt/nccl/lib:${LD_LIBRARY_PATH} + +# torch pip headers include cusparse.h and other CUDA headers shipped by the +# nvidia-* wheels; collect them once for the DeepGEMM build below and for +# DeepGEMM's runtime JIT. +RUN NVIDIA_DIR=$(python -c "import nvidia; print(nvidia.__path__[0])") && \ + mkdir -p /opt/nvidia-include && \ + cp -R "$NVIDIA_DIR"/*/include/. /opt/nvidia-include/ +ENV CPLUS_INCLUDE_PATH=/opt/nvidia-include \ + C_INCLUDE_PATH=/opt/nvidia-include + +# DeepGEMM (the connector's default `deepgemm` expert backend). +# --no-deps on install: DeepGEMM's own install path would reinstall torch. +RUN git clone https://github.com/deepseek-ai/DeepGEMM.git /tmp/DeepGEMM && \ + cd /tmp/DeepGEMM && \ + if [ -n "$DEEPGEMM_REF" ]; then git checkout "$DEEPGEMM_REF"; fi && \ + git submodule update --init --recursive && \ + TORCH_CUDA_ARCH_LIST="${DEEPGEMM_TORCH_CUDA_ARCH_LIST}" \ + python setup.py bdist_wheel && \ + uv pip install --system --no-deps dist/*.whl && \ + rm -rf /tmp/DeepGEMM + +# GPU-less build check: the packages import, afdEP-comm finds nvcc and the +# NCCL wheel, and the elfutils headers DeepGEMM's JIT needs are present. +RUN python -c "import afd_plugin, afdep, deep_gemm" && \ + python -c "from afdep._build import nccl_paths, resolve_nvcc; \ + nccl_paths(); resolve_nvcc()" && \ + test -f /usr/include/elfutils/libdwfl.h + +# Final stage: overlay workspace sources onto the cached deps image. +FROM deps +ARG APP_DIR=/opt/afd-plugin +WORKDIR ${APP_DIR} + +COPY --link . . + +ENTRYPOINT [] diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index dcaf453bc..356679129 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -66,7 +66,7 @@ depend on role worker implementations. | Base surface and lazy factory | [`base.py`](../../../afd_plugin/connectors/base.py), [`factory.py`](../../../afd_plugin/connectors/factory.py) | [P2P factory and construction tests](../../../tests/unit/connectors/test_p2p_connector.py), [metadata contracts](../../../tests/unit/connectors/test_metadata.py) | | Payloads and control-plane codec | [`metadata.py`](../../../afd_plugin/connectors/metadata.py) | Connector unit tests and [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | | CUDA P2P | [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py), [`topology.py`](../../../afd_plugin/distributed/topology.py) | [`test_p2p_connector.py`](../../../tests/unit/connectors/test_p2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | -| CUDA DeepEP M2N | [`gpu/deepep_m2n/`](../../../afd_plugin/connectors/gpu/deepep_m2n/), `AFDRoutedExpertsConnector` in [`base.py`](../../../afd_plugin/connectors/base.py) | [`test_deepep_m2n_connector.py`](../../../tests/unit/connectors/test_deepep_m2n_connector.py), [GPU device tests](../../../tests/gpu/README.md), [Qwen3 MoE E2E](../../../tests/e2e/models/qwen3_moe/test_qwen3_moe.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) (`test_deepseek_v2_lite_deepep_m2n`) | +| CUDA DeepEP M2N | [`gpu/deepep_m2n/`](../../../afd_plugin/connectors/gpu/deepep_m2n/), `AFDRoutedExpertsConnector` in [`base.py`](../../../afd_plugin/connectors/base.py) | [`test_deepep_m2n_connector.py`](../../../tests/unit/connectors/test_deepep_m2n_connector.py), [`test_deepep_m2n_eplb.py`](../../../tests/unit/connectors/test_deepep_m2n_eplb.py), [GPU device tests](../../../tests/gpu/README.md), [Qwen3 MoE E2E](../../../tests/e2e/models/qwen3_moe/test_qwen3_moe.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) (`test_deepseek_v2_lite_deepep_m2n`) | | Ascend CAMP2P | [`npu/camp2p.py`](../../../afd_plugin/connectors/npu/camp2p.py) | [`test_camp2p_connector.py`](../../../tests/unit/connectors/test_camp2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | | Ascend CAM async | [`npu/async_cam.py`](../../../afd_plugin/connectors/npu/async_cam.py), [`Push/Pull sessions`](../../../afd_plugin/connectors/npu/async_cam_runtime/async_cam_session.py) | [Connector unit tests](../../../tests/unit/connectors/test_async_cam_connector.py), [CAM host contracts](../../../tests/unit/connectors/npu/test_async_cam.py), [fused routing device test](../../../tests/npu/test_async_cam_connector.py), [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py) | | Process-group construction | [`afd_process_group.py`](../../../afd_plugin/distributed/afd_process_group.py) | Connector initialization tests plus platform E2E paths | diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index 2f1af04f4..5b8424cb0 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -176,7 +176,10 @@ does not adopt native V2 request state or scheduler-driven execution: Attention role before constructing communication resources. The V2 pair requires PP/PCP/DCP size 1, configured role ranks equal to DP x TP, -static EP, and a registered AFD model. CUDA uses `P2pNcclAFDConnector` with +static EP on Attention, and a registered AFD model. The FFN role may enable +vLLM's expert-parallel load balancing (`--enable-eplb`): EPLB is FFN-only +because Attention holds no experts and no MoE routing, and FFN EPLB rejects +DBO/ubatching (see below). CUDA uses `P2pNcclAFDConnector` with model-specific gate placement and supports two native microbatches with Attention DP1 or DP > 1. See the [runtime matrix](execution_platforms.md#tested-runtime-matrix) for tested @@ -319,6 +322,107 @@ of the step it calls `connector.synchronize_lanes()`. vLLM's MoE forward never runs on this runner. Dense (non-MoE) layers and shared experts have no owner on this path, so the model adapters reject them. +With EPLB the routed runner replaces the base EPLB step: nothing on this +FFN routes tokens, so no load is recorded on it. Its `eplb_step` delegates +to the connector's FFN EPLB driver +(`afd_plugin/connectors/gpu/deepep_m2n/eplb.py`), which counts every +connector step, and at each rearrangement boundary collects the +Attention-recorded load window over the connector's bootstrap group, runs +`EplbState.rearrange` directly and broadcasts the committed maps; Attention +adopts them before it routes the next step. The Attention side of that +module owns the recorder: it maps logical top-k ids to physical slots in the +model proxy's routing path and records load with stable-address tensors +CUDA graphs capture. With `log_balancedness` the connector logs EPLB's loss +for the deployed placement every `log_balancedness_interval` steps (see +below). + +## FFN-side EPLB + +The FFN role can enable vLLM's expert-parallel load balancing +(`--enable-eplb` plus `--eplb-config`). `GPUFFNModelRunner` owns the native +`EPLBController` and follows the upstream `GPUModelRunner` wiring: + +- `load_model` calls `prepare_load`, registers the loaded model through + `maybe_register_model` (the AFD FFN model inherits the native + `MixtureOfExperts` protocol, so the native `DeepseekV2MoE` layers receive + `EplbLayerState` maps and buffers), and starts the async transfer loop. +- Every connector step prepares the router's unpadded token count and then + steps the EPLB state: eager steps and graph replays are real steps that + record expert load; warmup and graph-capture steps are dummy steps that + keep every FFN rank's rearrangement counter in lockstep. +- Rearrangement runs entirely inside the native `EplbState`: load windows + are all-reduced across the FFN EP group, new maps are committed in place + (CUDA-graph compatible), and expert weights move over the EPLB + communicator (auto-selected: nixl, falling back to torch_gloo). + +With `log_balancedness` in `--eplb-config`, vLLM logs its balancedness +stats from `EplbState.step`, and the runner additionally logs the committed +placement — which logical experts each FFN EP rank holds, per MoE layer, and +the replicated ones — whenever a step commits a rearrangement (`eplb_step` +snapshots the map only on steps where the counter boundary or a pending +async transfer can commit, so ordinary steps stay sync-free). The same +placement log runs on the DeepEP M2N path (see the routed-experts note +above); the formatters live in `afd_plugin/distributed/eplb.py`. + +The unpadded token count is the step's total real Attention tokens +(`real_attention_token_total` in `ffn_metadata.py`), excluding the +placeholder rows idle Attention peers transfer. Every FFN rank's router +view covers those tokens, either through TP-replicated delivery or through +the MoE-level all-gather, so the count matches the rows the router's +map-and-record kernel processes. + +Limits: EPLB is rejected on the Attention role (its CUDA remote-expert +shells hold no experts), and FFN EPLB rejects DBO/ubatching because the +FFN daemon loop never installs vLLM's DBO thread context, so the record +kernel could not select the per-ubatch token count. EPLB recipes live under +`recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/`. + +### Loss logging on DeepEP M2N + +On the M2N path no `EplbState.step` runs, so vLLM's balancedness line never +appears. `DeepEPM2NFFNEplb` therefore logs EPLB's own objective itself, and it +logs it at two cadences (both opt-in through `log_balancedness`): + +- **every `log_balancedness_interval` steps**, the deployed map's loss on the + Attention ranks' load window; +- **at each rearrangement**, the loss of the map being replaced and of the map + just committed, on the same window. + +Both use `eplb_map_loss`, i.e. the quantity `DefaultEplbPolicy` optimises: a +logical expert's window load divided evenly over the replicas the candidate map +gives it, packed into ranks and reduced to the per-layer mean/max ratio vLLM +uses. A per-step figure was deliberately not kept: it measures *realised* load, +which the policy never sees and which moves with the batch rather than with the +placement. + +``` +AFD FFN EPLB step 1400: loss=0.9312 over 1000 steps (avg_tokens=..., max_tokens=...), steps until the next rearrangement: 1600 +AFD FFN EPLB loss: before=0.9421, after=0.9709, gain=+3.06% +``` + +Mechanics: + +- The trigger is the same rearrangement counter that paces the boundaries, so + both roles log at the same step: the FFN after its forward, Attention when it + steps the finished transaction, exactly like the boundary exchange. +- Attention ranks aggregate their own load window into logical experts + (`logical_load_from_window`, the only load the policy ever reads) and the FFN + ranks contribute zeros; the connector's bootstrap group reduces the sum to FFN + world rank 0, which scores the deployed map and logs it. A window that + recorded nothing (warmup, graph capture) logs nothing. +- `over N steps` is the number of recorded steps in the window, which is + `min(steps since the last rearrangement, window_size)`: the recorder clears + its window when Attention adopts a new map, because slot identities change + with the placement and load recorded under the old map cannot be read through + the new one. A window is fully refreshed before the next boundary whenever + `window_size <= step_interval`, so rearrangements are unaffected. +- Each logged step costs a host sync on every rank, which is why the interval + exists; a rearrangement boundary is unaffected and still runs the map exchange + only. +- Because both figures of the rearrangement line come from one window, `gain` is + the policy's own improvement rather than workload drift. On CUDA the proposal + commits unconditionally, so a negative `gain` is visible rather than gated. + ## Role-aware model loading The FFN process uses vLLM's model loader, while plugin-owned model integration diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 6bbb97883..1d08c809f 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -201,8 +201,11 @@ is required. With DeepEP M2N, Attention owns gate/top-k and FFN runs routed experts only; this requires `compute_gate_on_attention=true`, every layer to be MoE, and no shared expert. Native weight loading follows role-based filtering. -Both paths are CUDA-only and reject sequence-parallel MoE, EPLB, pipeline -parallelism, speculative decoding, and LoRA during construction. +Both paths are CUDA-only and reject sequence-parallel MoE, EPLB on the +Attention role, pipeline parallelism, speculative decoding, and LoRA during +construction; the FFN role supports FFN-side EPLB (see +[FFN runtime](ffn_runtime.md#ffn-side-eplb)) through the native +`Qwen3MoeSparseMoeBlock` construction. ### Qwen3.5/3.6 MoE CUDA boundary @@ -213,8 +216,11 @@ gate, routed experts, shared expert, and shared-expert gate. This path is CUDA-only, text-only, and requires `--language-model-only`, `compute_gate_on_attention=false`, and `pipeline_parallel_size=1`. It rejects -multimodal execution, sequence-parallel MoE, EPLB, speculative decoding, and -LoRA before language-model construction; NPU model-config resolution fails +multimodal execution, sequence-parallel MoE, EPLB on the Attention role, +speculative decoding, and LoRA before language-model construction; the FFN +role supports FFN-side EPLB (see +[FFN runtime](ffn_runtime.md#ffn-side-eplb)) through the native +`Qwen3NextSparseMoeBlock` construction. NPU model-config resolution fails before model loading. GDN DBO is unsupported under native hybrid/multimodal DBO validation; ordinary eager and graph execution remain supported. The 2A1F suite keeps static EP and PP1. Attention-side router transport, async diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index f5cbd55fb..0245ab2c1 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -75,9 +75,11 @@ Before allocating transport buffers, all ranks exchange a handshake and fail together if they disagree. The handshake checks the payload format, model shapes, lane count (DBO), `network`, bucket settings, GPU architecture (SM90 or newer), and the FFN deployment (expert parallel, EP -size, FFN TP=1, linear expert placement, no EPLB, all2all backend and expert -weight format). The FFN adopts the Attention TP size, dispatch capacity and -bucket ladder reported by the Attention ranks. +size, FFN TP=1, linear expert placement, EPLB off or on with a matching +physical expert count, all2all backend and expert weight format). The FFN +adopts the Attention TP size, dispatch capacity and bucket ladder reported by +the Attention ranks; with EPLB, Attention adopts the FFN's physical expert +count and expert maps. ## Installation @@ -213,7 +215,7 @@ that sets `capacity_tokens_per_rank` explicitly must agree with Attention. | Expert placement | FFN | `linear` (vLLM's default `--expert-placement-strategy`). | | `--all2all-backend` | FFN | `allgather_reducescatter` (vLLM's default). The connector moves the rows itself, so vLLM's MoE all-to-all never runs; this backend allocates no buffers and keeps vLLM's own DeepEP out of the process. | | `--moe-backend` | FFN | `triton` or `auto`. `auto` is pinned to `triton`; any other backend is rejected because it reorders the expert weights that the connector reads. | -| `--enable-eplb` | Both | Not supported. | +| `--enable-eplb` | FFN | Optional (see [EPLB](#eplb)). Async EPLB (`use_async`) is not supported. Attention runs no EPLB flags. | | `--enable-dbo` | Both | Optional, but must match on both roles. DBO uses exactly two micro-batches. | | CUDA graphs | Both | Eager or `FULL_DECODE_ONLY`. | | Pipeline / context parallelism, sequence-parallel MoE | Both | Not supported. | @@ -286,12 +288,88 @@ and FP8 examples, see the - SM90 or newer on every rank; the handshake rejects older GPUs. - Qwen3 MoE with every layer MoE and no shared expert, or DeepSeek-V2. -- `compute_gate_on_attention=true` on both roles; EPLB is not supported. +- `compute_gate_on_attention=true` on both roles; EPLB is supported on the + FFN role (see [EPLB](#eplb)). - FFN TP must be 1; FFN EP size equals `num_ffn_ranks`. -- `num_experts` must be divisible by `num_ffn_ranks`. +- `num_experts` must be divisible by `num_ffn_ranks` (with EPLB, + `num_experts + num_redundant_experts` must be). - CUDA graphs are limited to `FULL_DECODE_ONLY`; DBO uses exactly two micro-batches and must be enabled on both roles or neither. - Every rank must take part in every MoE layer and stage. Missing ranks, mismatched settings, or a crashed peer fail at the handshake or after the transport timeouts (`cpu_timeout_s`, `gpu_timeout_s`). - There is no fallback to another transport. + +## EPLB + +The FFN server may run vLLM's expert-parallel load balancing +(`--enable-eplb` plus `--eplb-config`). Because routing runs on Attention, +the two EPLB duties split across the roles: + +- Attention maps its logical top-k ids through the FFN's expert placement + before dispatch and records the per-expert load of its rows. It needs no + EPLB flags; the placement and schedule arrive over the connector's + handshake and control plane. +- The FFN owns the experts and vLLM's `EplbState`: at each rearrangement it + collects the Attention-recorded load, moves the expert weights, and + publishes the new placement, which Attention adopts before it routes the + next step. Captured CUDA graphs follow the placement because the maps are + updated in place. + +Constraints: async EPLB (`use_async`) and DBO are rejected with EPLB. EPLB +exchanges run on the connector's CPU bootstrap group at each rearrangement +boundary, so a rearrangement adds a host sync on every rank. + +To see what EPLB is doing, set `"log_balancedness": true` in the FFN's +`--eplb-config`. It turns on two logs: + +- at every rearrangement, per MoE layer, which logical experts each FFN rank + holds, which are replicated, and the recorded token load per rank that drove + the choice; +- the loss, i.e. the quantity `DefaultEplbPolicy` actually minimises, logged + every `"log_balancedness_interval"` steps for the deployed placement and + twice at each rearrangement: + +``` +AFD FFN EPLB step 1400: loss=0.9312 over 1000 steps (avg_tokens=812.34, max_tokens=1180), steps until the next rearrangement: 1600 +AFD FFN EPLB loss: before=0.9421, after=0.9709, gain=+3.06% +``` + +`loss` is the ratio of the average to the maximum per-rank token load, computed +per MoE layer and summed over layers, so 1.0 is perfectly even. Unlike realised +load, it uses the policy's own estimate: a logical expert's window load is +divided evenly over the replicas the placement gives it. Because it is a window +average and not one step's sample, it stays flat while the placement is +unchanged and moves when a rearrangement changes it — that series is how you +watch EPLB work. + +At a rearrangement, `before` scores the map that served the window and `after` +the map just committed, both on that same window, so `gain` is the improvement +the rearrangement achieved rather than workload drift. Track `gain` across +rearrangements to see whether EPLB keeps converging: near-zero gain means the +placement has converged, and since CUDA commits the proposal unconditionally a +negative gain is visible rather than gated. + +Two operational notes. Each logged step adds a host sync on every rank, so keep +`"log_balancedness_interval"` well above 1 for throughput-sensitive runs, or set +it equal to `"step_interval"` to log only one sample per rearrangement. And the +Attention recorder clears its window when it adopts a new map — slot identities +change with the placement — so `over N steps` restarts at the rearrangement; it +reads `window_size` once `window_size <= step_interval`, as with the defaults +1000 and 3000. + +The placement also lives in tensors at runtime: the FFN's +`EplbState.physical_to_logical_map` (via the runner's `eplb.state`), and on +Attention the connector's `logical_to_physical_map` (the same map, adopted +from the FFN). + +Set `"step_interval"` to control how many steps run between rearrangements +(default 3000; the first one happens after three quarters of that, matching +vLLM's initial progress). + +Run the recipe from the repository root: + +```bash +export MODEL_PATH=/path/model_weights/Qwen3-30B-A3B +bash recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_eplb_dp2tp1.sh +``` diff --git a/recipe/gpu/DeepEPM2NAFDConnector/README.md b/recipe/gpu/DeepEPM2NAFDConnector/README.md index 5151e899a..71fde56b3 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -16,7 +16,7 @@ supported layouts. │ └── prefill_decode_colocation/ # 2A2F: eager, graph, graph + DBO ├── qwen3_30b_a3b/ │ ├── README.md -│ └── prefill_decode_colocation/ # 2A2F: eager, graph, graph + DBO, Attention TP2 +│ └── prefill_decode_colocation/ # 2A2F: eager, graph, graph + EPLB, graph + DBO, Attention TP2 └── qwen3_235b_a22b_fp8/ ├── README.md └── prefill_decode_colocation/ # 4A2F FP8 payload: graph, graph + DBO diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1_eplb.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1_eplb.sh new file mode 100755 index 000000000..2d2b58865 --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1_eplb.sh @@ -0,0 +1,313 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# +# Qwen3-235B-A22B-FP8 4A2F prefill/decode colocation on one node with 6 GPUs: +# Attention DP4/TP1 (or DP1/TP4 with ATTN_TP=4) and FFN DP2/TP1 (expert +# parallel), FULL_DECODE_ONLY CUDA graphs. Every knob is an environment +# variable (default in parentheses): +# +# MODEL_PATH weights dir or HF id (/path/model_weights/Qwen3-235B-A22B-FP8) +# AFD_CONNECTOR DeepEPM2NAFDConnector | P2pNcclAFDConnector +# (DeepEPM2NAFDConnector) +# M2N_PAYLOAD_FORMAT DeepEPM2NAFDConnector dispatch format (fp8_block128_fp32) +# M2N_EXTRA_CONFIG extra connector_extra_config JSON members for +# DeepEPM2NAFDConnector, e.g. "expert_num_sms": 130 (empty) +# M2N_CPU_TIMEOUT_S DeepEP host-wait timeout. The FFN enters its first +# dispatch wait right after the roles join, while +# Attention is still loading weights and compiling +# this model's backbone; DeepEP's 120 s default is +# below that skew on a cold compile cache (900) +# M2N_GPU_TIMEOUT_S DeepEP kernel-spin timeout. Attention's first +# dispatch arrives only after its cold compile, so +# the FFN transport kernels must spin that long as +# well, or they give up and leave the elastic +# workspace in a broken state (900) +# VLLM_ENGINE_READY_TIMEOUT_S how long the API server waits for its engine +# cores to finish loading and compiling (1800; +# vLLM's default of 600 s is exceeded by this model) +# KV_CONNECTOR none | decode_bench (none); see the README +# KV_TRANSFER_CONFIG raw --kv-transfer-config JSON for Attention; +# overrides KV_CONNECTOR (unset) +# ENABLE_PREFIX_CACHING 1 | 0 (1) +# ENABLE_DBO 1 | 0 (0) +# DBO_DECODE_TOKEN_THRESHOLD / DBO_PREFILL_TOKEN_THRESHOLD (2 / 2) +# ENABLE_EPLB enable expert-parallel load balancing on the FFN +# role (1) +# EPLB_NUM_REDUNDANT_EXPERTS redundant experts; 16 on the 128 routed experts +# of Qwen3-235B-A22B gives 144 physical experts, 72 +# per FFN rank at EP size 2 (16) +# EPLB_STEP_INTERVAL EPLB rearrangement interval in steps; vLLM's +# default rearranges every 3000 steps (3000) +# EPLB_LOG_INTERVAL EPLB loss logging interval in steps, decoupled +# from the rearrangement interval (100) +# AFD_LOG_LEVEL level for the vLLM and afd_plugin loggers; the +# AFD connector and EPLB logs are INFO lines (INFO) +# MAX_NUM_SEQS per Attention DP rank (256) +# MAX_NUM_BATCHED_TOKENS prefill chunk per step (3072) +# CUDAGRAPH_CAPTURE_SIZES comma-separated decode graph sizes +# (8,16,32,48,64,96,112,128,192,256) +# MAX_MODEL_LEN (4096) +# GPU_MEMORY_UTILIZATION (unset: vLLM default) +# ATTN_DEVICES / FFN_DEVICES (0,1,2,3 / 4,5) +# ATTN_TP Attention tensor-parallel size; Attention DP is the +# number of ATTN_DEVICES divided by it (1) +# AFD_HOST / AFD_PORT AFD connector rendezvous (127.0.0.1 / 6239) +# AFD_PROCESS_GROUP_TIMEOUT_S AFD join timeout; the roles can finish +# loading minutes apart (3600) +# ATTN_PORT / FFN_PORT HTTP ports (18000 / 18001) +# EXTRA_ARGS extra `vllm serve` args for both roles (empty) +# ATTN_EXTRA_ARGS / FFN_EXTRA_ARGS extra `vllm serve` args for one role (empty) +# +# KV_CONNECTOR=decode_bench only (tools/benchmarks/decode_bench.py): +# DECODE_BENCH_NUM_PROMPTS pool prompts per Attention rank; 0 = dummy KV (500) +# DECODE_BENCH_PROMPT_LEN pool prompt length (1025) +# DECODE_BENCH_MAX_OUTPUT_LEN longest measured output; sizes the KV check (256) +# DECODE_BENCH_SEED pool token seed (0) +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../../../.." && pwd)" + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-235B-A22B-FP8} +AFD_CONNECTOR=${AFD_CONNECTOR:-DeepEPM2NAFDConnector} +M2N_PAYLOAD_FORMAT=${M2N_PAYLOAD_FORMAT:-fp8_block128_fp32} +M2N_EXTRA_CONFIG=${M2N_EXTRA_CONFIG:-} +M2N_CPU_TIMEOUT_S=${M2N_CPU_TIMEOUT_S:-900} +M2N_GPU_TIMEOUT_S=${M2N_GPU_TIMEOUT_S:-900} +KV_CONNECTOR=${KV_CONNECTOR:-none} +KV_TRANSFER_CONFIG=${KV_TRANSFER_CONFIG:-} +ENABLE_PREFIX_CACHING=${ENABLE_PREFIX_CACHING:-1} +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_DECODE_TOKEN_THRESHOLD=${DBO_DECODE_TOKEN_THRESHOLD:-2} +DBO_PREFILL_TOKEN_THRESHOLD=${DBO_PREFILL_TOKEN_THRESHOLD:-2} +ENABLE_EPLB=${ENABLE_EPLB:-1} +EPLB_NUM_REDUNDANT_EXPERTS=${EPLB_NUM_REDUNDANT_EXPERTS:-16} +EPLB_STEP_INTERVAL=${EPLB_STEP_INTERVAL:-3000} +EPLB_LOG_INTERVAL=${EPLB_LOG_INTERVAL:-100} +AFD_LOG_LEVEL=${AFD_LOG_LEVEL:-INFO} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-3072} +CUDAGRAPH_CAPTURE_SIZES=${CUDAGRAPH_CAPTURE_SIZES:-8,16,32,48,64,96,112,128,192,256} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-4096} +GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-} +ATTN_DEVICES=${ATTN_DEVICES:-0,1,2,3} +FFN_DEVICES=${FFN_DEVICES:-4,5} +ATTN_TP=${ATTN_TP:-1} +AFD_HOST=${AFD_HOST:-127.0.0.1} +AFD_PORT=${AFD_PORT:-6239} +AFD_PROCESS_GROUP_TIMEOUT_S=${AFD_PROCESS_GROUP_TIMEOUT_S:-3600} +ATTN_PORT=${ATTN_PORT:-18000} +FFN_PORT=${FFN_PORT:-18001} +EXTRA_ARGS=${EXTRA_ARGS:-} +ATTN_EXTRA_ARGS=${ATTN_EXTRA_ARGS:-} +FFN_EXTRA_ARGS=${FFN_EXTRA_ARGS:-} +DECODE_BENCH_NUM_PROMPTS=${DECODE_BENCH_NUM_PROMPTS:-500} +DECODE_BENCH_PROMPT_LEN=${DECODE_BENCH_PROMPT_LEN:-1025} +DECODE_BENCH_MAX_OUTPUT_LEN=${DECODE_BENCH_MAX_OUTPUT_LEN:-256} +DECODE_BENCH_SEED=${DECODE_BENCH_SEED:-0} + +export VLLM_USE_V2_MODEL_RUNNER=1 +# This model's cold torch.compile pass (plus weight load) outlives vLLM's +# 600 s engine-core startup budget; the pod keeps the pod-level startup +# budget (55 min) as the outer bound. +export VLLM_ENGINE_READY_TIMEOUT_S=${VLLM_ENGINE_READY_TIMEOUT_S:-1800} +# vLLM's default logging config attaches its handler to the `vllm` logger only, +# so the plugin's own `afd_plugin.*` loggers fall through to the unconfigured +# root logger and drop every INFO line: the connector plan and the EPLB +# balancedness/placement statistics never reach the log. Extend the same +# handler to the plugin hierarchy, and keep the uvicorn loggers vLLM hands to +# uvicorn (it passes this same file as uvicorn's log config, so leaving them +# out also hides `Application startup complete`). A custom dictConfig is +# authoritative over --log-level, hence the level is templated here. +AFD_LOG_CONFIG=${AFD_LOG_CONFIG:-/tmp/afd-plugin-logging.json} +cat >"$AFD_LOG_CONFIG" <&2 + exit 1 +fi +ATTN_DP=$((NUM_ATTN_RANKS / ATTN_TP)) + +case "$AFD_CONNECTOR" in + # DeepEP M2N routes on Attention and sends FP8 tokens to the FFN experts. + DeepEPM2NAFDConnector) + COMPUTE_GATE_ON_ATTENTION=true + CONNECTOR_EXTRA_CONFIG=", \"connector_extra_config\": {\"payload_format\": \"${M2N_PAYLOAD_FORMAT}\", \"cpu_timeout_s\": ${M2N_CPU_TIMEOUT_S}, \"gpu_timeout_s\": ${M2N_GPU_TIMEOUT_S}${M2N_EXTRA_CONFIG:+, ${M2N_EXTRA_CONFIG}}}" + ;; + # P2P routes on the FFN and sends BF16 hidden states. Do not run its FFN + # with PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True: the first step hangs. + P2pNcclAFDConnector) + COMPUTE_GATE_ON_ATTENTION=false + CONNECTOR_EXTRA_CONFIG="" + ;; + *) + echo "unsupported AFD_CONNECTOR=${AFD_CONNECTOR}" >&2 + exit 1 + ;; +esac + +if [[ -z "$KV_TRANSFER_CONFIG" ]]; then + case "$KV_CONNECTOR" in + none) ;; + decode_bench) + KV_TRANSFER_CONFIG=$(cat <&2 + exit 1 + ;; + esac +fi +if [[ "$KV_TRANSFER_CONFIG" == *tools.benchmarks.* ]]; then + # tools/benchmarks is not in the wheel; the vLLM scheduler and worker + # subprocesses import it from the repository root. + export PYTHONPATH="${REPO_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" +fi + +afd_config() { + cat <&2 + exit 1 +fi + +MAX_CUDAGRAPH_CAPTURE_SIZE=$(tr ',' '\n' <<<"$CUDAGRAPH_CAPTURE_SIZES" | sort -n | tail -n 1) +COMMON_ARGS=( + --enable-expert-parallel + --enable-ep-weight-filter + --max-num-seqs "$MAX_NUM_SEQS" + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" + --max-model-len "$MAX_MODEL_LEN" + --max-cudagraph-capture-size "$MAX_CUDAGRAPH_CAPTURE_SIZE" + --compilation-config "{\"cudagraph_mode\": \"FULL_DECODE_ONLY\", \"cudagraph_capture_sizes\": [${CUDAGRAPH_CAPTURE_SIZES}]}" + # Loading the weights can outlast the default timeouts, especially from + # cold storage. + --cpu-distributed-timeout-seconds 3600 + --trust-remote-code +) +if [[ "$ENABLE_DBO" == 1 ]]; then + COMMON_ARGS+=( + --enable-dbo + --dbo-decode-token-threshold "$DBO_DECODE_TOKEN_THRESHOLD" + --dbo-prefill-token-threshold "$DBO_PREFILL_TOKEN_THRESHOLD" + ) +fi +if [[ -n "$GPU_MEMORY_UTILIZATION" ]]; then + COMMON_ARGS+=(--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION") +fi +# shellcheck disable=SC2206 +COMMON_ARGS+=($EXTRA_ARGS) + +ATTN_ARGS=(--data-parallel-size "$ATTN_DP" --tensor-parallel-size "$ATTN_TP") +if [[ "$ENABLE_PREFIX_CACHING" == 1 ]]; then + ATTN_ARGS+=(--enable-prefix-caching) +else + ATTN_ARGS+=(--no-enable-prefix-caching) +fi +if [[ -n "$KV_TRANSFER_CONFIG" ]]; then + ATTN_ARGS+=(--kv-transfer-config "$KV_TRANSFER_CONFIG") +fi +# shellcheck disable=SC2206 +ATTN_ARGS+=($ATTN_EXTRA_ARGS) + +FFN_ARGS=(--data-parallel-size "$NUM_FFN_RANKS" --tensor-parallel-size 1) +# EPLB runs on the FFN role only: Attention maps its logical top-k ids +# through the FFN's placement, so it needs no EPLB flags. Async EPLB is not +# supported with this connector (use_async=false). Experts rearrange every +# EPLB_STEP_INTERVAL steps; log_balancedness plus log_balancedness_interval +# log the deployed map's EPLB loss every EPLB_LOG_INTERVAL steps independently +# of the rearrangement cadence, and each rearrangement also logs the loss of +# the map it replaced and the one it committed (logging forces a host sync, +# hence it is opt-in). +if [[ "$ENABLE_EPLB" == 1 ]]; then + FFN_ARGS+=( + --enable-eplb + --eplb-config "{\"num_redundant_experts\": ${EPLB_NUM_REDUNDANT_EXPERTS}, \"step_interval\": ${EPLB_STEP_INTERVAL}, \"use_async\": false, \"log_balancedness\": true, \"log_balancedness_interval\": ${EPLB_LOG_INTERVAL}}" + ) +fi +# shellcheck disable=SC2206 +FFN_ARGS+=($FFN_EXTRA_ARGS) + +CUDA_VISIBLE_DEVICES=$FFN_DEVICES vllm serve "$MODEL_PATH" \ + --additional-config "$(afd_config ffn)" \ + "${COMMON_ARGS[@]}" \ + "${FFN_ARGS[@]}" \ + --host 127.0.0.1 \ + --port "$FFN_PORT" > ffn.log 2>&1 & + +CUDA_VISIBLE_DEVICES=$ATTN_DEVICES vllm serve "$MODEL_PATH" \ + --additional-config "$(afd_config attention)" \ + "${COMMON_ARGS[@]}" \ + "${ATTN_ARGS[@]}" \ + --host 127.0.0.1 \ + --port "$ATTN_PORT" > attn.log 2>&1 & + +wait diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md index 184806b5e..27db12de3 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/README.md @@ -21,10 +21,11 @@ Launch scripts for Qwen3-30B-A3B (BF16) on `DeepEPM2NAFDConnector` with vLLM ```text prefill_decode_colocation/ -├── 2a2f_eager_dp2tp1.sh # Attention DP2 x TP1, eager -├── 2a2f_graph_dp2tp1.sh # Attention DP2 x TP1, FULL_DECODE_ONLY graph -├── 2a2f_graph_dbo_dp2tp1.sh # Attention DP2 x TP1, graph + DBO -└── 2a2f_graph_dp1tp2.sh # Attention DP1 x TP2, FULL_DECODE_ONLY graph +├── 2a2f_eager_dp2tp1.sh # Attention DP2 x TP1, eager +├── 2a2f_graph_dp2tp1.sh # Attention DP2 x TP1, FULL_DECODE_ONLY graph +├── 2a2f_graph_eplb_dp2tp1.sh # Attention DP2 x TP1, graph + FFN EPLB +├── 2a2f_graph_dbo_dp2tp1.sh # Attention DP2 x TP1, graph + DBO +└── 2a2f_graph_dp1tp2.sh # Attention DP1 x TP2, FULL_DECODE_ONLY graph ``` ## Configuration notes @@ -41,6 +42,13 @@ prefill_decode_colocation/ - The dispatch capacity is derived from `--max-num-batched-tokens` and the largest capture size, so lowering `--max-num-batched-tokens` shrinks the transport buffers. +- The EPLB script enables `--enable-eplb` on the FFN server only, with 32 + redundant experts (160 physical experts, 80 per FFN rank). Attention needs + no EPLB flags: it records expert load while routing and adopts the FFN's + expert placement. Async EPLB (`use_async`) and DBO are not supported with + EPLB on this connector; the first rearrangement happens after + `step_interval - step_interval // 4` steps (vLLM's `--eplb-config` + defaults: 2250 steps). ## Checking a deployment diff --git a/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_eplb_dp2tp1.sh b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_eplb_dp2tp1.sh new file mode 100644 index 000000000..529c013ee --- /dev/null +++ b/recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_eplb_dp2tp1.sh @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +# 2A2F FULL_DECODE_ONLY graph with EPLB: Attention DP2 x TP1 on GPUs 0,1; +# FFN DP2 on GPUs 2,3. Qwen3-30B-A3B has 128 routed experts, so 32 redundant +# experts yield 160 physical experts (80 per FFN rank at EP size 2); the FFN +# rearranges them every step_interval steps from the load Attention records +# while routing. EPLB runs on the FFN role only: Attention maps the logical +# top-k ids through the FFN's placement, so it needs no EPLB flags. Async EPLB +# is not supported with this connector (the recipe pins use_async=false). + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3-30B-A3B} +export VLLM_USE_V2_MODEL_RUNNER=1 + +CUDA_VISIBLE_DEVICES=2,3 vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --enable-eplb \ + --eplb-config '{"num_redundant_experts": 32, "use_async": false}' \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "DeepEPM2NAFDConnector", + "host": "127.0.0.1", + "port": 6239, + "num_attention_ranks": 2, + "num_ffn_ranks": 2, + "compute_gate_on_attention": true + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 4096 \ + --max-model-len 4096 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ + --host 127.0.0.1 \ + --port 18001 > ffn.log 2>&1 & + +CUDA_VISIBLE_DEVICES=0,1 vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "DeepEPM2NAFDConnector", + "host": "127.0.0.1", + "port": 6239, + "num_attention_ranks": 2, + "num_ffn_ranks": 2, + "compute_gate_on_attention": true + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 4096 \ + --max-model-len 4096 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{"cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes": [8, 16, 32, 64]}' \ + --host 127.0.0.1 \ + --port 18000 > attn.log 2>&1 & + +wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_eplb_dp1tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_eplb_dp1tp2.sh new file mode 100644 index 000000000..baf8ebdba --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_eplb_dp1tp2.sh @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +export VLLM_USE_V2_MODEL_RUNNER=1 + +# EPLB is enabled on the FFN role only: the Attention role ships hidden +# states without computing MoE routing (compute_gate_on_attention=false) +# and its CUDA remote-expert shells reject enable_eplb. DeepSeek-V2-Lite +# has 64 routed experts, so 32 redundant experts yield 96 physical experts +# (48 per FFN rank at EP size 2). + +CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 2 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "P2pNcclAFDConnector", + "host": "127.0.0.1", + "port": 6269, + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 64 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{ + "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] + }' \ + --host 127.0.0.1 \ + --port 18305 \ + --trust-remote-code > attn.log 2>&1 & + +CUDA_VISIBLE_DEVICES=2,3 uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 2 \ + --enable-expert-parallel \ + --enable-eplb \ + --eplb-config '{"num_redundant_experts": 32}' \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "P2pNcclAFDConnector", + "host": "127.0.0.1", + "port": 6269, + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 64 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{ + "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] + }' \ + --host 127.0.0.1 \ + --port 18305 \ + --trust-remote-code > ffn.log 2>&1 & + +wait diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_eplb_dp1tp2.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_eplb_dp1tp2.sh new file mode 100644 index 000000000..8fa0ce48a --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_eplb_dp1tp2.sh @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +# Use .agents/skills/deploy-afd-k8s/SKILL.md to deploy this recipe +# on 2 pods using the setup below. +# For pod 1: POD=ATTENTION_0 +# For pod 2: POD=FFN_0 + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/DeepSeek-V2-Lite} +export VLLM_USE_V2_MODEL_RUNNER=1 + +# EPLB is enabled on the FFN role only: the Attention role ships hidden +# states without computing MoE routing (compute_gate_on_attention=false) +# and its CUDA remote-expert shells reject enable_eplb. DeepSeek-V2-Lite +# has 64 routed experts, so 32 redundant experts yield 96 physical experts +# (48 per FFN rank at EP size 2). + +AFD_CONNECTOR_HOST=${AFD_CONNECTOR_HOST:-vllm-ffn-p2p-service} +AFD_CONNECTOR_PORT=${AFD_CONNECTOR_PORT:-6269} + +if [ "$POD" == "ATTENTION_0" ]; then + CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 2 \ + --enable-expert-parallel \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "P2pNcclAFDConnector", + "host": "'"${AFD_CONNECTOR_HOST}"'", + "port": '"${AFD_CONNECTOR_PORT}"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 64 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{ + "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] + }' \ + --host 127.0.0.1 \ + --port 18305 \ + --trust-remote-code > attn.log 2>&1 & +fi + +if [ "$POD" == "FFN_0" ]; then + CUDA_VISIBLE_DEVICES=0,1 uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 1 \ + --tensor-parallel-size 2 \ + --enable-expert-parallel \ + --enable-eplb \ + --eplb-config '{"num_redundant_experts": 32}' \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "P2pNcclAFDConnector", + "host": "'"${AFD_CONNECTOR_HOST}"'", + "port": '"${AFD_CONNECTOR_PORT}"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 64 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{ + "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] + }' \ + --host 127.0.0.1 \ + --port 18305 \ + --trust-remote-code > ffn.log 2>&1 & +fi + +wait diff --git a/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/eplb_4a_2f_graph.sh b/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/eplb_4a_2f_graph.sh new file mode 100644 index 000000000..429af1278 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/eplb_4a_2f_graph.sh @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/Qwen3.5-122B-A10B-FP8} +export VLLM_USE_V2_MODEL_RUNNER=1 + +# EPLB is enabled on the FFN role only: the Attention role ships hidden +# states without computing MoE routing (compute_gate_on_attention=false) +# and its remote-expert shells reject enable_eplb. Qwen3.5-122B-A10B has +# 256 routed experts, so 32 redundant experts yield 288 physical experts +# (144 per FFN rank at EP size 2). + +CUDA_VISIBLE_DEVICES=0,1,2,3 uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 4 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --dtype bfloat16 \ + --language-model-only \ + --max-model-len 114688 \ + --mamba-cache-mode align \ + --all2all-backend allgather_reducescatter \ + --seed 0 \ + --additional-config '{ + "afd": { + "role": "attention", + "connector": "P2pNcclAFDConnector", + "host": "127.0.0.1", + "port": 6269, + "num_attention_ranks": 4, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 8192 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{ + "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] + }' \ + --host 127.0.0.1 \ + --port 18305 \ + --trust-remote-code > attn.log 2>&1 & + +CUDA_VISIBLE_DEVICES=4,5 uv run vllm serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --enable-eplb \ + --eplb-config '{"num_redundant_experts": 32}' \ + --dtype bfloat16 \ + --language-model-only \ + --max-model-len 114688 \ + --mamba-cache-mode align \ + --all2all-backend allgather_reducescatter \ + --seed 0 \ + --additional-config '{ + "afd": { + "role": "ffn", + "connector": "P2pNcclAFDConnector", + "host": "127.0.0.1", + "port": 6269, + "num_attention_ranks": 4, + "num_ffn_ranks": 2 + } + }' \ + --max-num-seqs 64 \ + --max-num-batched-tokens 8192 \ + --max-cudagraph-capture-size 64 \ + --compilation-config '{ + "cudagraph_mode": "FULL_DECODE_ONLY", "cudagraph_capture_sizes":[64] + }' \ + --host 127.0.0.1 \ + --port 18305 \ + --trust-remote-code > ffn.log 2>&1 & + +wait diff --git a/tests/e2e/mrv2_evidence.py b/tests/e2e/mrv2_evidence.py index b416e23c4..0642a9d31 100644 --- a/tests/e2e/mrv2_evidence.py +++ b/tests/e2e/mrv2_evidence.py @@ -8,7 +8,7 @@ of successful live transactions, pairing native A input IDs with the actual first local Hash-layer F input IDs. Capture retains references without D2H. CUDA target: 5df63a7203cdb9328d824a27b682e77c7b9928be (0.32.0). -Ascend retains the existing 0.30 call without CUDA's randomize_inputs argument. +Ascend retains the existing native call without CUDA's randomize_inputs argument. Remove these wrappers when upstream exposes equivalent live Attention/FFN replay and microbatch counters. """ @@ -396,7 +396,7 @@ def replay( # Patch functionality: count completed live replay and nonempty microbatches. # Signature/source: GPUModelRunner.execute_model, CUDA vLLM 5df63a7, # vllm/v1/worker/gpu/model_runner.py. Delegation is observation-only. - # Ascend retains its 0.30 native call; randomize_inputs is CUDA-only. + # Ascend retains its native call; randomize_inputs is CUDA-only. def execute( self, scheduler_output: SchedulerOutput, diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index 6d8a53e0c..f7b504233 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -191,6 +191,24 @@ def test_routed_transport_preserves_runner_and_gate_policy( ) +def test_gpu_validator_rejects_async_eplb(mrv2_config): + mrv2_config.parallel_config.enable_eplb = True + mrv2_config.parallel_config.eplb_config = SimpleNamespace(use_async=False) + mrv2_config.parallel_config.enable_dbo = False + mrv2_config.parallel_config.use_ubatching = False + mrv2_config.additional_config["afd"]["role"] = "ffn" + + validate_gpu_model_runner_v2_config( + mrv2_config, expected_role="ffn", device_type="cuda" + ) + + mrv2_config.parallel_config.eplb_config = SimpleNamespace(use_async=True) + with pytest.raises(RuntimeError, match="does not support async EPLB"): + validate_gpu_model_runner_v2_config( + mrv2_config, expected_role="ffn", device_type="cuda" + ) + + @pytest.mark.parametrize("role", ["attention", "ffn"]) @pytest.mark.parametrize( "architecture, accepted", @@ -512,7 +530,48 @@ def _expect(error): {"par.enable_elastic_ep": True}, _err("static expert"), id="reject-elastic" ), pytest.param( - {"par.enable_eplb": True}, _err("static expert"), id="reject-eplb" + {"par.enable_eplb": True}, + _err("EPLB only on the FFN role"), + id="reject-attention-eplb", + ), + pytest.param( + { + "afd.role": "ffn", + "hf.architectures": ["AFDDeepseekV2ForCausalLM"], + "par.enable_eplb": True, + }, + None, + id="accept-ffn-eplb", + ), + pytest.param( + { + "afd.role": "ffn", + "hf.architectures": ["AFDDeepseekV2ForCausalLM"], + "par.enable_eplb": True, + "par.eplb_config.use_async": True, + }, + _err("async EPLB"), + id="reject-ffn-eplb-async", + ), + pytest.param( + { + "afd.role": "ffn", + "hf.architectures": ["AFDDeepseekV2ForCausalLM"], + "par.enable_eplb": True, + "par.enable_dbo": True, + }, + _err("DBO or ubatching"), + id="reject-ffn-eplb-dbo", + ), + pytest.param( + { + "afd.role": "ffn", + "hf.architectures": ["AFDDeepseekV2ForCausalLM"], + "par.enable_eplb": True, + "par.use_ubatching": True, + }, + _err("DBO or ubatching"), + id="reject-ffn-eplb-ubatching", ), pytest.param( {"par.use_sequence_parallel_moe": True}, diff --git a/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 7e8c3cc4a..d66e762b3 100644 --- a/tests/unit/connectors/test_deepep_m2n_connector.py +++ b/tests/unit/connectors/test_deepep_m2n_connector.py @@ -89,14 +89,23 @@ def _connector( return AFDConnectorFactory.create_connector(0, local_rank, vllm_config, afd_config) -def _ffn_deployment(layout: DeepEPM2NLayout, weight_format: str = "bf16") -> dict: +def _ffn_deployment( + layout: DeepEPM2NLayout, + weight_format: str = "bf16", + *, + eplb: bool = False, + num_physical_experts: int | None = None, +) -> dict: """A valid FFN deployment: linear EP over the FFN ranks, idle all2all.""" return { "expert_parallel": True, "ep_size": layout.ffn_size, "tensor_parallel_size": 1, "expert_placement_strategy": "linear", - "eplb": False, + "eplb": eplb, + "num_physical_experts": ( + layout.num_experts if num_physical_experts is None else num_physical_experts + ), "all2all_backend": "allgather_reducescatter", "weight_format": weight_format, } @@ -982,7 +991,6 @@ def test_expert_weight_format_names_what_vllm_loads(): ("ep_size", 4), ("tensor_parallel_size", 2), ("expert_placement_strategy", "round_robin"), - ("eplb", True), ("all2all_backend", "deepep_low_latency"), ("weight_format", "fp8_block128"), ], @@ -1004,6 +1012,50 @@ def test_handshake_rejects_ffn_deployments_that_misplace_experts(key, value): validate_handshake(infos, layout, NetworkConfig()) +def test_handshake_accepts_eplb_ffn_deployment(): + """EPLB on the FFN is valid when every FFN rank reports the same, + evenly splitting physical expert count. + + Why: DeepEP routes by physical slot id, so the deployment is consistent + as long as the slot count matches what the FFN loads. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + infos = _handshake(layout) + for rank in range(layout.ffn_size): + infos[rank]["ffn_deployment"] = _ffn_deployment( + layout, eplb=True, num_physical_experts=160 + ) + + result = validate_handshake(infos, layout, NetworkConfig()) + + assert result.enable_eplb is True + assert result.num_physical_experts == 160 + + +@pytest.mark.parametrize( + ("num_physical", "match"), + [ + (127, "physical experts for 128 logical experts"), + (129, "does not split evenly"), + ], +) +def test_handshake_rejects_broken_eplb_slot_counts(num_physical, match): + """EPLB with fewer slots than logical experts, or a slot count that does + not split over the FFN ranks, fails the handshake. + + Why: DeepEP's slot-to-rank routing requires contiguous, even blocks. + """ + layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) + infos = _handshake(layout) + for rank in range(layout.ffn_size): + infos[rank]["ffn_deployment"] = _ffn_deployment( + layout, eplb=True, num_physical_experts=num_physical + ) + + with pytest.raises(RuntimeError, match=match): + validate_handshake(infos, layout, NetworkConfig()) + + def test_handshake_accepts_fp8_block_weights_for_fp8_payload(): layout = DeepEPM2NLayout(ffn_size=2, attn_size=2, num_experts=128) infos = _handshake( diff --git a/tests/unit/connectors/test_deepep_m2n_eplb.py b/tests/unit/connectors/test_deepep_m2n_eplb.py new file mode 100644 index 000000000..5a875dd25 --- /dev/null +++ b/tests/unit/connectors/test_deepep_m2n_eplb.py @@ -0,0 +1,534 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU-only tests for the DeepEP M2N connector's EPLB coordination. + +The distributed boundary exchange is exercised with a recording stand-in for +``torch.distributed``: what is asserted is the protocol (which tensors are +reduced, which are adopted), not the transport. +""" + +from __future__ import annotations + +import logging +from types import SimpleNamespace +from typing import Any + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.connectors import AFDControlPayload, AFDDPMetadata # noqa: E402 +from afd_plugin.connectors.gpu.deepep_m2n import eplb as eplb_module # noqa: E402 +from afd_plugin.connectors.gpu.deepep_m2n.eplb import ( # noqa: E402 + DeepEPM2NAttentionEplb, + DeepEPM2NFFNEplb, + real_rows_in_tp_slice, +) +from tests.unit.fakes.connector import FakeRoutedConnector # noqa: E402 + +MOE_LAYERS = (1, 3) +NUM_LOGICAL = 12 +NUM_PHYSICAL = 12 +WINDOW_SIZE = 4 +STEP_INTERVAL = 100 +INITIAL_REARRANGE_STEP = 8 + + +def _schedule(**overrides) -> dict[str, int]: + schedule = { + "num_moe_layers": len(MOE_LAYERS), + "num_logical_experts": NUM_LOGICAL, + "num_physical_experts": NUM_PHYSICAL, + "map_slots": 4, + "window_size": WINDOW_SIZE, + "step_interval": STEP_INTERVAL, + "initial_rearrangement_step": INITIAL_REARRANGE_STEP, + "log_interval": 0, + } + schedule.update(overrides) + return schedule + + +def _attention_connector(**overrides): + attrs = { + "device": torch.device("cpu"), + "attn_tp_size": 1, + "attention_tp_rank": 0, + "world_rank": 2, + "bootstrap_group": object(), + } + attrs.update(overrides) + return SimpleNamespace(**attrs) + + +def _recorder(*, log_interval: int = 0, **connector_overrides): + initial_map = torch.zeros((len(MOE_LAYERS), NUM_LOGICAL, 4), dtype=torch.int64) + initial_count = torch.ones((len(MOE_LAYERS), NUM_LOGICAL), dtype=torch.int64) + return DeepEPM2NAttentionEplb( + _attention_connector(**connector_overrides), + _schedule(log_interval=log_interval), + MOE_LAYERS, + initial_map, + initial_count, + ) + + +def _payload(*, is_warmup=False, is_graph_capturing=False, num_real_tokens=5): + return AFDControlPayload( + dp_metadata_list={ + 0: AFDDPMetadata( + num_tokens_across_dp_cpu=torch.tensor([num_real_tokens]), + ), + }, + is_graph_capturing=is_graph_capturing, + is_warmup=is_warmup, + num_real_tokens=num_real_tokens, + ) + + +class _FakeDist: + """Records the EPLB exchange collectives; the map broadcast fills the + receiving tensor so adoption is observable.""" + + def __init__(self, *, map_value: int = 7) -> None: + self.reduces: list[tuple[Any, int]] = [] + self.broadcasts: list[Any] = [] + self._map_value = map_value + + def reduce(self, tensor, dst, group): + self.reduces.append((tensor, dst)) + + def broadcast(self, tensor, src, group): + tensor.fill_(self._map_value) + self.broadcasts.append(tensor) + + +def test_real_rows_in_tp_slice_splits_the_dp_rank_tokens(): + """Each TP rank's slice holds its share of the DP rank's real tokens; + padding rows (beyond the real count) are excluded.""" + assert real_rows_in_tp_slice(10, tp_size=1, tp_rank=0) == 10 + assert real_rows_in_tp_slice(10, tp_size=2, tp_rank=0) == 5 + assert real_rows_in_tp_slice(10, tp_size=2, tp_rank=1) == 5 + assert real_rows_in_tp_slice(9, tp_size=2, tp_rank=0) == 5 + assert real_rows_in_tp_slice(9, tp_size=2, tp_rank=1) == 4 + assert real_rows_in_tp_slice(1, tp_size=4, tp_rank=2) == 0 + assert real_rows_in_tp_slice(0, tp_size=2, tp_rank=0) == 0 + + +def test_recorder_rotates_the_window_and_drops_dummy_load(): + """Real steps move the recorded load into the window and zero the pass; + dummy steps drop their load; the upcoming transaction's flags gate + recording and set its real token count.""" + recorder = _recorder() + pass_sum = lambda: int(recorder.expert_load_pass.sum()) # noqa: E731 + recorder.expert_load_pass.fill_(3) + + recorder.on_control_step(_payload()) + # The first transaction has no predecessor: only its own state is set. + assert recorder.record_enabled.item() is True + assert recorder.num_unpadded_tokens[0].item() == 5 + assert pass_sum() == 3 * NUM_PHYSICAL * len(MOE_LAYERS) + + recorder.on_control_step(_payload(num_real_tokens=7)) + # The real transaction rotated its load into the window and reset the pass. + assert int(recorder.expert_load_window[0].sum()) == 3 * NUM_PHYSICAL * len( + MOE_LAYERS + ) + assert pass_sum() == 0 + assert recorder.num_unpadded_tokens[0].item() == 7 + + # A warmup (dummy) transaction's load is dropped, not recorded. + recorder.on_control_step(_payload(is_warmup=True, num_real_tokens=2)) + assert int(recorder.expert_load_window[1].sum()) == 0 + recorder.expert_load_pass.fill_(9) + recorder.on_control_step(_payload()) + assert pass_sum() == 0 + assert int(recorder.expert_load_window[2].sum()) == 0 + # Dummy transactions do not record. + recorder.on_control_step(_payload(is_graph_capturing=True)) + assert recorder.record_enabled.item() is False + + +def test_recorder_steps_tp_rank_real_counts(): + """With Attention TP the recorder counts only this rank's real rows.""" + recorder = _recorder(attn_tp_size=2, attention_tp_rank=1) + + recorder.on_control_step(_payload(num_real_tokens=9)) + + assert recorder.num_unpadded_tokens[0].item() == 4 + + +def test_recorder_exchanges_window_and_adopts_maps_at_boundary( + monkeypatch: pytest.MonkeyPatch, +): + """At a rearrangement boundary the recorder contributes its window to the + reduce aimed at FFN world rank 0 and adopts the broadcast maps.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + initial_map = torch.zeros((len(MOE_LAYERS), NUM_LOGICAL, 4), dtype=torch.int64) + initial_count = torch.ones((len(MOE_LAYERS), NUM_LOGICAL), dtype=torch.int64) + recorder = DeepEPM2NAttentionEplb( + _attention_connector(), + _schedule(step_interval=10), + MOE_LAYERS, + initial_map, + initial_count, + ) + recorder.expert_load_pass.fill_(2) + + # initial_rearrangement_step=8, interval=10: the third step reaches it. + recorder.on_control_step(_payload()) + recorder.on_control_step(_payload()) + recorder.on_control_step(_payload()) + + assert len(fake_dist.reduces) == 1 + window, dst = fake_dist.reduces[0] + assert dst == 0 + # The window carries the load recorded before the boundary: the first + # real transaction's pass (filled with 2), the second recorded nothing. + assert int(window.sum()) == 2 * NUM_PHYSICAL * len(MOE_LAYERS) + # Both maps were adopted from the broadcast (filled with 7). + assert int(recorder.logical_to_physical_map.sum()) == 7 * len(MOE_LAYERS) * ( + NUM_LOGICAL * 4 + ) + assert int(recorder.logical_replica_count.sum()) == 7 * NUM_LOGICAL * len( + MOE_LAYERS + ) + # Slot identities changed with the map, so the next window starts empty. + assert int(recorder.expert_load_window.sum()) == 0 + assert int(recorder.expert_load_window[0].sum()) == 0 + + +def test_recorder_reduces_the_logical_window_at_each_log_interval( + monkeypatch: pytest.MonkeyPatch, +): + """Every log interval the recorder joins a reduce carrying its load window + aggregated into logical experts -- the only load EPLB's policy sees -- and + it changes no EPLB state.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + recorder = _recorder(log_interval=2) + per_logical_expert = NUM_LOGICAL * len(MOE_LAYERS) + + recorder.on_control_step(_payload()) + for step_load in (1, 2, 3, 4): + recorder.expert_load_pass.fill_(step_load) + recorder.on_control_step(_payload()) + + # Counters 10 and 12 are log steps. All the recorder's slots hold logical + # 0 (its map is all zeros), so each expert's load is slot 0's window + # total: 1+2 tokens at the first log step, then 1+2+3+4. + assert [dst for _load, dst in fake_dist.reduces] == [0, 0] + assert [load.shape for load, _dst in fake_dist.reduces] == [ + (len(MOE_LAYERS), NUM_LOGICAL), + ] * 2 + assert [int(load.sum()) for load, _dst in fake_dist.reduces] == [ + 3 * per_logical_expert, + 10 * per_logical_expert, + ] + assert fake_dist.broadcasts == [] + + +def test_recorder_keeps_rearranging_without_the_loss_exchange( + monkeypatch: pytest.MonkeyPatch, +): + """Without log_balancedness the recorder reduces nothing until the + rearrangement boundary, which still runs the map exchange.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + recorder = _recorder() + + for _ in range(STEP_INTERVAL - INITIAL_REARRANGE_STEP + 1): + recorder.on_control_step(_payload()) + + assert len(fake_dist.reduces) == 1 + assert len(fake_dist.broadcasts) == 2 + + +def _logical_view(physical_to_logical_map: torch.Tensor): + """The logical-to-physical map and replica counts a placement implies.""" + num_layers = physical_to_logical_map.shape[0] + logical_to_physical = torch.zeros((num_layers, NUM_LOGICAL, 4), dtype=torch.int64) + replica_count = torch.zeros((num_layers, NUM_LOGICAL), dtype=torch.int64) + for layer in range(num_layers): + for slot, logical in enumerate(physical_to_logical_map[layer].tolist()): + if logical < 0: + continue + used = int(replica_count[layer, logical]) + logical_to_physical[layer, logical, used] = slot + replica_count[layer, logical] = used + 1 + return logical_to_physical, replica_count + + +def _ffn_state(): + model = SimpleNamespace( + num_moe_layers=len(MOE_LAYERS), + num_logical_experts=NUM_LOGICAL, + num_physical_experts=NUM_PHYSICAL, + ) + # Two layers of a 12-slot world over 2 FFN ranks: layer 0 keeps the + # identity placement with logical 8 replicated in a redundant slot, + # layer 1 swaps logical 5 and 6 across the ranks. + layer0 = [0, 1, 2, 3, 4, 8, 6, 7, 8, 9, 10, 11] + layer1 = [0, 1, 2, 3, 4, 6, 5, 7, 8, 9, 10, 11] + physical_to_logical = torch.tensor([layer0, layer1]) + logical_to_physical, replica_count = _logical_view(physical_to_logical) + model_state = SimpleNamespace( + model=model, + expert_load_window=torch.zeros((WINDOW_SIZE, len(MOE_LAYERS), NUM_PHYSICAL)), + physical_to_logical_map=physical_to_logical, + logical_to_physical_map=logical_to_physical, + logical_replica_count=replica_count, + ) + rearrangements: list[None] = [] + + def rearrange(): + rearrangements.append(None) + + state = SimpleNamespace( + model_states={"model-hash": model_state}, + expert_rearrangement_step_interval=STEP_INTERVAL, + expert_rearrangement_step=INITIAL_REARRANGE_STEP, + rearrange=rearrange, + ) + return state, model_state, rearrangements + + +def _ffn_connector( + world_rank: int, *, log_placement: bool = False, log_interval: int = 0 +): + return SimpleNamespace( + bootstrap_group=object(), + world_rank=world_rank, + ffn_size=2, + vllm_config=SimpleNamespace( + parallel_config=SimpleNamespace( + eplb_config=SimpleNamespace( + log_balancedness=log_placement, + log_balancedness_interval=log_interval, + ), + ), + ), + ) + + +def test_ffn_driver_rearranges_only_at_the_boundary(monkeypatch: pytest.MonkeyPatch): + """The FFN driver counts every step and runs the exchange (reduce, + EplbState.rearrange, map broadcasts) only when the counter reaches the + interval; FFN ranks reduce zeros and only rank 0 adopts the window.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + state, model_state, rearrangements = _ffn_state() + state.expert_rearrangement_step_interval = 10 + connector = _ffn_connector(world_rank=1) + driver = DeepEPM2NFFNEplb(connector, state) + + for _ in range(10 - INITIAL_REARRANGE_STEP - 1): + driver.step() + + assert rearrangements == [] + assert fake_dist.reduces == [] + + driver.step() + + assert len(rearrangements) == 1 + # FFN ranks contribute an empty window; rank 0 is the reduce destination. + window, dst = fake_dist.reduces[0] + assert dst == 0 + assert int(window.sum()) == 0 + # Rank != 0 keeps its (empty) state window; maps went to every rank. + assert int(model_state.expert_load_window.sum()) == 0 + assert len(fake_dist.broadcasts) == 2 + + +def test_ffn_driver_rank_zero_adopts_the_reduced_window( + monkeypatch: pytest.MonkeyPatch, +): + """FFN world rank 0 copies the reduced Attention window into its + EplbState before the rearrangement reads it.""" + + def reduce(tensor, dst, group): + tensor.fill_(5) + + monkeypatch.setattr(eplb_module.dist, "reduce", reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", _FakeDist().broadcast) + state, model_state, rearrangements = _ffn_state() + state.expert_rearrangement_step_interval = 10 + connector = _ffn_connector(world_rank=0) + driver = DeepEPM2NFFNEplb(connector, state) + + for _ in range(10 - INITIAL_REARRANGE_STEP): + driver.step() + + assert len(rearrangements) == 1 + assert ( + int(model_state.expert_load_window.sum()) + == 5 * WINDOW_SIZE * len(MOE_LAYERS) * NUM_PHYSICAL + ) + + +def test_routed_connector_base_seams_default_to_no_op(): + """A routed connector without dynamic expert placement maps ids + unchanged and steps nothing.""" + connector = FakeRoutedConnector() + ids = torch.tensor([[1, 2]]) + + assert torch.equal(connector.apply_eplb_routing(ids, 0), ids) + connector.step_ffn_eplb() + assert connector.events == [] + + +def test_ffn_driver_logs_chosen_experts_when_balancedness_logging_is_on( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """With log_balancedness the FFN world rank 0 logs the committed + placement and the recorded load that chose it at each boundary.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + state, model_state, rearrangements = _ffn_state() + state.expert_rearrangement_step_interval = 10 + driver = DeepEPM2NFFNEplb(_ffn_connector(0, log_placement=True), state) + + with caplog.at_level(logging.INFO, logger=eplb_module.logger.name): + for _ in range(10 - INITIAL_REARRANGE_STEP): + driver.step() + + assert len(rearrangements) == 1 + message = caplog.messages[0] + assert "EPLB chose experts" in message + assert "rank 0=0 tokens, rank 1=0 tokens" in message + assert "layer 0: rank 0=[0, 1, 2, 3, 4, 8]" in message + + +def test_ffn_driver_stays_silent_without_balancedness_logging( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """Without log_balancedness no placement is logged, on any rank.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + for world_rank in (0, 1): + state, _model_state, _rearrangements = _ffn_state() + state.expert_rearrangement_step_interval = 10 + driver = DeepEPM2NFFNEplb(_ffn_connector(world_rank), state) + + with caplog.at_level(logging.INFO, logger=eplb_module.logger.name): + for _ in range(10 - INITIAL_REARRANGE_STEP): + driver.step() + + assert caplog.messages == [] + + +def _rank_heavy_reduce(tensor, dst, group): + """Attention ranks contribute a window whose load sits on the logical + experts FFN rank 0's slots hold (the fake map's first four).""" + tensor.zero_() + tensor[:, :4] = 10 + + +def test_ffn_driver_logs_the_deployed_map_loss_at_each_log_interval( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """Every log interval FFN world rank 0 scores the deployed map on the + Attention-reduced window with EPLB's own load model.""" + monkeypatch.setattr(eplb_module.dist, "reduce", _rank_heavy_reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", _FakeDist().broadcast) + state, _model_state, rearrangements = _ffn_state() + state.expert_rearrangement_step_interval = STEP_INTERVAL + state.expert_rearrangement_step = 0 + driver = DeepEPM2NFFNEplb( + _ffn_connector(0, log_placement=True, log_interval=2), state + ) + + with caplog.at_level(logging.INFO, logger=eplb_module.logger.name): + for _ in range(5): + driver.step() + + assert rearrangements == [] + assert len(caplog.messages) == 2 + message = caplog.messages[0] + assert "AFD FFN EPLB step 2:" in message + # Rank 0's slots hold 40 of the 80 tokens in every layer: half balanced. + assert "loss=0.5000 over 2 steps (avg_tokens=40.00, max_tokens=80)" in message + assert "steps until the next rearrangement: 98" in message + assert "AFD FFN EPLB step 4:" in caplog.messages[1] + + +def test_ffn_driver_skips_loss_logging_without_recorded_load( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """A log step that recorded no load (warmup, graph capture) logs nothing, + though every rank still joined the reduce.""" + fake_dist = _FakeDist() + monkeypatch.setattr(eplb_module.dist, "reduce", fake_dist.reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) + state, _model_state, _rearrangements = _ffn_state() + state.expert_rearrangement_step = 0 + driver = DeepEPM2NFFNEplb( + _ffn_connector(0, log_placement=True, log_interval=2), state + ) + + with caplog.at_level(logging.INFO, logger=eplb_module.logger.name): + for _ in range(2): + driver.step() + + assert len(fake_dist.reduces) == 1 + assert caplog.messages == [] + + +def test_ffn_driver_logs_the_loss_of_the_map_it_replaced( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """At a boundary the FFN world rank 0 logs EPLB's own objective for the + deployed map and for the proposal on the same window, so the gain is the + improvement the rearrangement made rather than workload drift.""" + + # Both layers recorded rank 0 holding 80 of the 100 tokens: logical 0 + # and 1 are its hot experts, logical 6 and 7 are rank 1's cold ones. + def reduce(tensor, dst, group): + tensor.zero_() + tensor[:, :, 0] = 50 + tensor[:, :, 1] = 30 + tensor[:, :, 6] = 10 + tensor[:, :, 7] = 10 + + monkeypatch.setattr(eplb_module.dist, "reduce", reduce) + monkeypatch.setattr(eplb_module.dist, "broadcast", _FakeDist().broadcast) + state, model_state, _rearrangements = _ffn_state() + state.expert_rearrangement_step_interval = 10 + # Both layers serve the same placement while the window is recorded. + deployed = [0, 1, 2, 3, 4, 8, 6, 7, 8, 9, 10, 11] + model_state.physical_to_logical_map = torch.tensor([deployed, deployed]) + ( + model_state.logical_to_physical_map, + model_state.logical_replica_count, + ) = _logical_view(model_state.physical_to_logical_map) + # The proposal moves the cold expert 7 onto rank 0 and the hot expert 1 + # onto rank 1: 60/40 instead of 80/20. + balanced = [0, 7, 2, 3, 4, 8, 1, 6, 8, 9, 10, 11] + + def rearrange(): + model_state.physical_to_logical_map = torch.tensor([balanced, balanced]) + + state.rearrange = rearrange + driver = DeepEPM2NFFNEplb(_ffn_connector(0, log_placement=True), state) + + with caplog.at_level(logging.INFO, logger=eplb_module.logger.name): + for _ in range(10 - INITIAL_REARRANGE_STEP): + driver.step() + + assert ( + "AFD FFN EPLB loss: before=0.6250, after=0.8333, gain=+33.33%" + in (caplog.messages[-1]) + ) diff --git a/tests/unit/connectors/test_metadata.py b/tests/unit/connectors/test_metadata.py index 79c5901f5..82f85b7f0 100644 --- a/tests/unit/connectors/test_metadata.py +++ b/tests/unit/connectors/test_metadata.py @@ -212,12 +212,14 @@ def test_p2p_dp_metadata_serialization_uses_json_payload(): is_warmup=False, is_graph_replaying=True, is_profile=True, + num_real_tokens=42, ), ) decoded_payload = decode_control_payload(payload) decoded = decoded_payload.dp_metadata_list assert payload.startswith(b"{") + assert decoded_payload.num_real_tokens == 42 assert isinstance(decoded[7], AFDDPMetadata) assert _tolist(decoded[7].num_tokens_across_dp_cpu) == [3, 5] assert int(decoded[7].max_tokens_across_dp_cpu) == 5 diff --git a/tests/unit/distributed/test_eplb.py b/tests/unit/distributed/test_eplb.py new file mode 100644 index 000000000..2a24f691c --- /dev/null +++ b/tests/unit/distributed/test_eplb.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU-only tests for the shared EPLB observability helpers.""" + +from __future__ import annotations + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.distributed.eplb import ( # noqa: E402 + describe_eplb_placement, + describe_eplb_rank_load, + eplb_map_loss, + logical_load_from_window, + summarize_rank_load, +) + +# Two layers of a 12-slot world over 2 FFN ranks: layer 0 keeps the identity +# placement with logical 8 replicated in a redundant slot, layer 1 swaps +# logical 5 and 6 across the ranks and drops the replica. +LAYER_IDENTITY = [0, 1, 2, 3, 4, 8, 6, 7, 8, 9, 10, 11] +LAYER_SWAPPED = [0, 1, 2, 3, 4, 6, 5, 7, 8, 9, 10, 11] + + +def test_describe_eplb_placement_lists_experts_per_rank(): + """The placement description names each rank's logical experts and the + replicated ones, per MoE layer.""" + placement = torch.tensor([LAYER_IDENTITY, LAYER_SWAPPED]) + + lines = describe_eplb_placement(placement, ffn_size=2) + + assert len(lines) == 2 + assert lines[0] == ( + "layer 0: rank 0=[0, 1, 2, 3, 4, 8] | rank 1=[6, 7, 8, 9, 10, 11] " + "| replicas: 8x2" + ) + assert lines[1] == ( + "layer 1: rank 0=[0, 1, 2, 3, 4, 6] | rank 1=[5, 7, 8, 9, 10, 11] " + "| replicas: none" + ) + + +def test_describe_eplb_rank_load_sums_each_ranks_slots(): + """A rank's recorded load is the sum over its contiguous slot block, + across the window steps and layers.""" + # Two window steps, one layer: rank 0's slots (0-5) carry 1 each, rank + # 1's slots (6-11) carry 2 each. + layer_load = [1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2] + window = torch.tensor([layer_load, layer_load]).unsqueeze(1) + + load = describe_eplb_rank_load(window, ffn_size=2) + + assert load == "rank 0=12 tokens, rank 1=24 tokens" + + +def test_logical_load_from_window_sums_the_window_through_the_map(): + """The window is summed over its steps and gathered through the placement + that recorded it, giving tokens per logical expert; replicas beyond an + expert's count are ignored rather than double counted.""" + step_load = [4, 2, 1, 1] + window = torch.tensor([step_load, step_load, step_load]).unsqueeze(1) + # Logical 0 is replicated on slots 0 and 1, logical 1 on 1 and 0, and the + # remaining slots are all logicals 2 and 3 hold. + logical_to_physical = torch.tensor([[[0, 1], [1, 0], [2, -1], [3, -1]]]) + replica_count = torch.tensor([[2, 2, 1, 1]]) + + logical_load = logical_load_from_window(window, logical_to_physical, replica_count) + + assert logical_load.tolist() == [[3 * (4 + 2), 3 * (2 + 4), 3 * 1, 3 * 1]] + + +def test_summarize_rank_load_matches_vllms_balancedness(): + """Balancedness is the ratio of the per-layer mean and max rank load, + each summed over layers.""" + per_rank = torch.tensor([[18, 6], [12, 24]]) + + summary = summarize_rank_load(per_rank) + + assert summary.avg_tokens == pytest.approx(30.0) + assert summary.max_tokens == pytest.approx(42.0) + assert summary.balancedness == pytest.approx(30.0 / 42.0) + + +def test_summarize_rank_load_reports_zero_without_tokens(): + """A step that recorded no load has no ratio to report.""" + summary = summarize_rank_load(torch.zeros((2, 2))) + + assert (summary.avg_tokens, summary.max_tokens, summary.balancedness) == (0, 0, 0) + + +def test_summarize_rank_load_reaches_one_when_perfectly_balanced(): + """Equal rank load is a balancedness of 1.""" + per_rank = torch.tensor([[5, 5], [7, 7]]) + + assert summarize_rank_load(per_rank).balancedness == pytest.approx(1.0) + + +def test_eplb_map_loss_splits_a_replicated_expert_over_its_replicas(): + """EPLB estimates a logical expert's load as divided evenly over its + replicas, so an idle replica still counts as half of the load.""" + # Layer 0 of LAYER_IDENTITY: logical 8 sits in slots 5 (rank 0) and 8 + # (rank 1), and the whole of its load was recorded on one replica. + logical_load = torch.zeros((1, 12)) + logical_load[0, 8] = 20 + + placement = torch.tensor([LAYER_IDENTITY]) + loss = eplb_map_loss(logical_load, placement, ffn_size=2) + + # 20 tokens over 2 replicas -> 10 per rank, which is perfectly balanced, + # even though the recorded per-slot load sits on rank 0 alone. + assert loss.avg_tokens == pytest.approx(10.0) + assert loss.max_tokens == pytest.approx(10.0) + assert loss.balancedness == pytest.approx(1.0) + + +def test_eplb_map_loss_rises_when_a_hot_expert_moves_to_the_light_rank(): + """The metric the policy minimises improves when the hot expert stops + sharing a rank with the other hot expert.""" + identity = [0, 1, 2, 3] + # Logical 0 (50 tokens) and logical 1 (30) start on rank 0, which also + # gives logical 2 (10) and logical 3 (10) to rank 1. + logical_load = torch.tensor([[50, 30, 10, 10]], dtype=torch.float32) + balanced = [0, 3, 2, 1] + + before = eplb_map_loss(logical_load, torch.tensor([identity]), ffn_size=2) + after = eplb_map_loss(logical_load, torch.tensor([balanced]), ffn_size=2) + + assert before.balancedness == pytest.approx(50.0 / 80.0) + assert after.balancedness == pytest.approx(50.0 / 60.0) + + +def test_eplb_map_loss_ignores_slots_no_logical_expert_uses(): + """Slots mapped to -1 carry no load and no replica, so they neither add + to a rank's total nor dilute an expert's load.""" + logical_load = torch.tensor([[12.0, 6.0, 3.0]]) + placement = torch.tensor([[0, 1, 2, -1]]) + + loss = eplb_map_loss(logical_load, placement, ffn_size=2) + + # Rank 0 holds the first two slots (18 tokens), rank 1 the third (3). + # Read without the guard the -1 slot would wrap to the last logical + # expert and hand rank 1 its 3 tokens twice, a 0.6667 loss. + assert loss.avg_tokens == pytest.approx(10.5) + assert loss.max_tokens == pytest.approx(18.0) + assert loss.balancedness == pytest.approx(10.5 / 18.0) diff --git a/tests/unit/fakes/connector.py b/tests/unit/fakes/connector.py index 7222e2270..7a2469359 100644 --- a/tests/unit/fakes/connector.py +++ b/tests/unit/fakes/connector.py @@ -74,10 +74,17 @@ def is_initialized(self) -> bool: def close(self) -> None: self.events.append(("close",)) - def init_afd_connector(self, *, moe_layer_indices=None, local_experts=None): + def init_afd_connector( + self, + *, + moe_layer_indices=None, + local_experts=None, + eplb_state=None, + ): self.init_kwargs = { "moe_layer_indices": moe_layer_indices, "local_experts": local_experts, + "eplb_state": eplb_state, } self.initialized = True @@ -86,7 +93,7 @@ def update_state_from_dp_metadata(self, payload) -> None: self.control_updates.append(payload) def send_attn_output(self, hidden_states, context, **kwargs) -> None: - self.events.append(("send_attn",)) + self.events.append(("send_attn", dict(kwargs))) def recv_ffn_output(self, ref_tensor, ubatch_idx=0, **kwargs): return ref_tensor diff --git a/tests/unit/fakes/eplb.py b/tests/unit/fakes/eplb.py new file mode 100644 index 000000000..3c8532218 --- /dev/null +++ b/tests/unit/fakes/eplb.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A recording stand-in for vLLM's ``EPLBController`` in FFN runner tests.""" + +from __future__ import annotations + + +class FakeEPLBController: + """Records the EPLB calls the FFN runner makes per connector step.""" + + def __init__(self): + self.state = None + self.prepare_loads = 0 + self.registrations = [] + self.async_loop_starts = [] + self.steps = [] + self.prepared_forwards = [] + + def prepare_load(self): + self.prepare_loads += 1 + + def maybe_register_model(self, model, model_config): + self.registrations.append((model, model_config)) + return True + + def maybe_start_async_loop(self, eplb_models_added, load_dummy_weights=False): + self.async_loop_starts.append((eplb_models_added, load_dummy_weights)) + + def step(self, is_dummy=False, is_profile=False): + self.steps.append((is_dummy, is_profile)) + + def prepare_forward(self, model_config, num_unpadded_tokens, ubatch_slices=None): + self.prepared_forwards.append((model_config, num_unpadded_tokens)) diff --git a/tests/unit/fakes/ubatching.py b/tests/unit/fakes/ubatching.py index c5a4957ac..cb16399aa 100644 --- a/tests/unit/fakes/ubatching.py +++ b/tests/unit/fakes/ubatching.py @@ -137,7 +137,7 @@ def finish(): class LegacyNPUAttentionMetadataState: - """Expose the vLLM 0.30 attention-preparation signature on CPU.""" + """Expose the legacy vLLM attention-preparation signature on CPU.""" def __init__(self): self.prepare_error = None diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index e67a400da..e88aef2b0 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -335,6 +335,7 @@ def mrv2_vllm_config( enable_expert_parallel=True, enable_elastic_ep=False, enable_eplb=False, + eplb_config=SimpleNamespace(use_async=False), use_sequence_parallel_moe=False, is_moe_model=True, enable_dbo=False, diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index 2bac7facd..596597097 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -13,8 +13,11 @@ from afd_plugin.config import AFD_ASYNC_CONNECTOR, AFDConfig # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 -from tests.unit.fakes.connector import ProxyConnector as _FakeConnector # noqa: E402 -from tests.unit.fakes.connector import install_proxy_context # noqa: E402 +from tests.unit.fakes.connector import ( # noqa: E402 + FakeRoutedConnector, + ProxyConnector, + install_proxy_context, +) from tests.unit.fakes.model_layers import GateLinear # noqa: E402 ROUTING_HIDDEN_SIZE = 16 @@ -183,7 +186,7 @@ def test_remote_experts_proxy_without_external_routing_ships_hidden_states( ROUTE_WEIGHT = 0.5 -class _InPlaceCombineConnector(_FakeConnector): +class _InPlaceCombineConnector(ProxyConnector): """Combine writes the routed output into the Attention buffer, as DeepEP M2N does.""" @@ -279,6 +282,59 @@ def test_remote_proxy_requires_forward_metadata(monkeypatch): adapter.RemoteFFNProxy(layer_idx=0)(torch.ones(1, 4)) +class _EplbRoutedConnector(FakeRoutedConnector): + """A routed connector whose FFN runs EPLB: dispatch carries physical ids.""" + + def apply_eplb_routing(self, topk_ids, layer_idx): + self.events.append(("eplb_route", layer_idx)) + return topk_ids + 100 + + +def test_gate_proxy_maps_topk_ids_through_connector_eplb_routing(monkeypatch): + """A gate-bearing proxy on a routed-experts connector maps its logical + top-k ids through the connector's EPLB placement before dispatch. + + Why: with EPLB the FFN's experts sit in physical slots, so dispatch by + logical id would reach the wrong experts with no error. + """ + connector = _EplbRoutedConnector() + afd_metadata = SimpleNamespace(connector=connector, stage_idx=0) + monkeypatch.setattr( + adapter, + "get_afd_metadata_from_forward_context", + lambda: afd_metadata, + ) + monkeypatch.setattr( + adapter, + "get_forward_context", + lambda: SimpleNamespace(ubatch_idx=0), + ) + monkeypatch.setattr( + adapter, + "maybe_apply_dbo_yield", + lambda hidden_states, *, role: hidden_states, + ) + topk_weights = torch.tensor([[0.6, 0.4]]) + topk_ids = torch.tensor([[1, 3]]) + + def fake_fused_topk(**kwargs): + return topk_weights, topk_ids, None + + monkeypatch.setattr(adapter, "fused_topk", fake_fused_topk) + proxy = adapter.RemoteFFNProxy(layer_idx=5) + proxy.gate = nn.Linear(4, 8, bias=False) + proxy.top_k = 2 + proxy.renormalize = False + proxy.tp_size = 1 + + proxy(torch.ones(2, 4)) + + assert ("eplb_route", 5) in connector.events + send = next(event for event in connector.events if event[0] == "send_attn") + assert torch.equal(send[1]["topk_ids"], topk_ids + 100) + assert send[1]["topk_weights"] is topk_weights + + def test_remote_proxy_exchanges_cam_during_profile(monkeypatch): events: list[tuple] = [] install_proxy_context(adapter, monkeypatch, events) @@ -481,7 +537,7 @@ def test_routing_proxy_owns_gate_at_native_checkpoint_path(): assert set(dict(proxy.named_parameters())) == {"gate.weight"} -# Fake combine output scale of ``_FakeConnector.recv_ffn_output``. +# Fake combine output scale of ``ProxyConnector.recv_ffn_output``. COMBINE_SCALE = 0.25 # Added to the peer rank's rows by the fake all-gather, to tell them apart. PEER_OFFSET = 100.0 diff --git a/tests/unit/model_executor/models/test_qwen3_5_contract.py b/tests/unit/model_executor/models/test_qwen3_5_contract.py index ca48798cd..4f938351e 100644 --- a/tests/unit/model_executor/models/test_qwen3_5_contract.py +++ b/tests/unit/model_executor/models/test_qwen3_5_contract.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project from __future__ import annotations +from contextlib import nullcontext from types import SimpleNamespace import pytest @@ -231,3 +232,81 @@ def test_qwen_rejects_unsupported_configuration_before_construction( adapter.AFDQwen3_5MoeForConditionalGeneration( vllm_config=conditional_config.config ) + + +def test_qwen_conditional_model_accepts_ffn_role_eplb(monkeypatch): + """FFN-role EPLB passes the feature gates that Attention EPLB fails.""" + + class FakeCausalLM(nn.Module): + def __init__(self): + super().__init__() + self.make_empty_intermediate_tensors = object() + + model_config = SimpleNamespace( + hf_config=SimpleNamespace( + vision_config=SimpleNamespace(out_hidden_size=8), + ), + multimodal_config=SimpleNamespace( + language_model_only=True, + mm_encoder_tp_mode="weights", + is_multimodal_pruning_enabled=lambda: False, + video_pruning_rate=0.2, + ), + ) + model_config.get_multimodal_config = lambda: model_config.multimodal_config + vllm_config = SimpleNamespace( + lora_config=None, + model_config=model_config, + quant_config=None, + parallel_config=SimpleNamespace( + enable_eplb=True, + eplb_config=SimpleNamespace(num_redundant_experts=32), + pipeline_parallel_size=1, + use_sequence_parallel_moe=False, + ), + speculative_config=None, + ) + monkeypatch.setattr( + adapter, + "parse_optional_afd_config", + lambda *_args, **_kwargs: SimpleNamespace( + role="ffn", + compute_gate_on_attention=False, + ), + ) + monkeypatch.setattr( + adapter.AFDQwen3_5MoeForConditionalGeneration, + "_mark_tower_model", + lambda *_args, **_kwargs: nullcontext(), + ) + monkeypatch.setattr( + adapter.AFDQwen3_5MoeForConditionalGeneration, + "_mark_language_model", + lambda *_args, **_kwargs: nullcontext(), + ) + monkeypatch.setattr( + adapter.native, + "Qwen3_VisionTransformer", + lambda *_args, **_kwargs: nn.Identity(), + ) + monkeypatch.setattr( + adapter.native, + "cached_tokenizer_from_config", + lambda *_args, **_kwargs: object(), + ) + monkeypatch.setattr( + adapter, + "AFDQwen3_5MoeForCausalLM", + lambda **_kwargs: FakeCausalLM(), + ) + monkeypatch.setattr( + adapter.AFDQwen3_5MoeForConditionalGeneration, + "set_moe_parameters", + lambda _self: None, + ) + + model = adapter.AFDQwen3_5MoeForConditionalGeneration( + vllm_config=vllm_config, + ) + + assert model.afd_role == "ffn" diff --git a/tests/unit/model_executor/models/test_qwen3_moe_construction.py b/tests/unit/model_executor/models/test_qwen3_moe_construction.py index bdae43e44..c0d0b439e 100644 --- a/tests/unit/model_executor/models/test_qwen3_moe_construction.py +++ b/tests/unit/model_executor/models/test_qwen3_moe_construction.py @@ -326,6 +326,22 @@ def test_unsupported_modes_fail_closed( adapter._validate_supported_config(vllm_config, afd_config) +def test_ffn_role_eplb_passes_supported_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """FFN-role EPLB passes the gate that Attention-role EPLB fails.""" + + vllm_config = _model_config() + vllm_config.parallel_config.enable_eplb = True + monkeypatch.setattr( + adapter, + "current_platform", + _platform(is_cuda=True), + ) + + adapter._validate_supported_config(vllm_config, AFDConfig(role="ffn")) + + def test_gate_on_attention_rejects_dense_layers( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/unit/v1/worker/test_ffn_metadata.py b/tests/unit/v1/worker/test_ffn_metadata.py index fb1422e40..283574352 100644 --- a/tests/unit/v1/worker/test_ffn_metadata.py +++ b/tests/unit/v1/worker/test_ffn_metadata.py @@ -10,6 +10,7 @@ from afd_plugin.v1.worker.ffn_metadata import ( aggregate_ffn_token_counts, project_ffn_token_counts_to_dp, + real_attention_token_total, ) @@ -41,3 +42,28 @@ def test_aggregate_follows_the_subgroup_partition( def test_project_collapses_tp_ranks_to_dp_ranks(): assert project_ffn_token_counts_to_dp((8, 8, 13, 13), dp_size=2) == (8, 13) assert project_ffn_token_counts_to_dp((3, 2), dp_size=2) == (3, 2) + + +@pytest.mark.parametrize( + ("attention_counts", "attention", "expected"), + [ + # Zero-token peers transfer placeholder rows but contribute no + # real tokens (docstring example counterpart of the aggregation). + ((0, 4, 5, 6), 4, 15), + ((3, 5), 2, 8), + ((), 4, 0), + ((), 1, 0), + # TP-expanded counts: one count per DP rank, two TP ranks each. + ((4, 6), 4, 20), + ((7,), 2, 14), + # Counts already covering every Attention rank are kept as-is. + ((1, 2, 3), 3, 6), + ], +) +def test_real_token_total_excludes_placeholder_rows( + attention_counts, attention, expected +): + assert ( + real_attention_token_total(attention_counts, attention_size=attention) + == expected + ) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index 93a299f8a..595f54bb4 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -38,6 +38,7 @@ DPMetadataValues as _FakeDPMetadata, ) from tests.unit.fakes.dp_metadata import dp_tokens as _tokens # noqa: E402 +from tests.unit.fakes.eplb import FakeEPLBController # noqa: E402 from tests.unit.fakes.graph import Graph as _FakeGraph # noqa: E402 from tests.unit.fakes.profiler import StepProfiler as _StepProfiler # noqa: E402 @@ -274,7 +275,9 @@ def _runner_with_connector_and_model(model, *, num_layers=1): use_v2_model_runner=True, parallel_config=SimpleNamespace( data_parallel_size=1, + enable_eplb=False, is_moe_model=True, + tensor_parallel_size=1, use_sequence_parallel_moe=False, ), compilation_config=SimpleNamespace( @@ -284,12 +287,17 @@ def _runner_with_connector_and_model(model, *, num_layers=1): ) runner.connector = _FakeConnector() runner.model = model + runner.model_config = SimpleNamespace() runner.afd_config = SimpleNamespace(compute_gate_on_attention=False) runner.afd_cudagraph_policy = SimpleNamespace(enabled=False) runner.num_layers = num_layers runner.use_cuda_graph = False runner._cuda_graphs = {} + runner._graph_memory_pool = None + runner.load_config = SimpleNamespace() + runner.device = torch.device("cpu") runner.prof = None + runner.eplb = FakeEPLBController() return runner @@ -670,6 +678,273 @@ def test_ffn_runner_cuda_graph_miss_falls_back_to_eager(): ] +def test_ffn_runner_loads_model_and_registers_eplb(monkeypatch): + runner = _runner_with_connector_and_model(_FakeModel()) + model = runner.model + runner.model = None + loader = SimpleNamespace( + load_model=lambda **_kwargs: model, + load_weights=lambda *_args, **_kwargs: None, + ) + monkeypatch.setattr( + ffn_model_runner_module, "get_model_loader", lambda _config: loader + ) + + runner.load_model() + + assert runner.model is model + assert runner.eplb.prepare_loads == 1 + assert runner.eplb.registrations == [(model, runner.model_config)] + assert runner.eplb.async_loop_starts == [(True, False)] + + +def test_fake_eplb_controller_matches_vllm_eplb_controller_api(): + """The fake must mirror vLLM's EPLBController call surface. + + A drifted fake hides upstream signature changes: vLLM 0.32 removed + ``load_dummy_weights`` from ``maybe_register_model`` and moved it to + ``maybe_start_async_loop``, and the stale fake kept the runner's broken + three-argument call green. + """ + import inspect + + from vllm.v1.worker.gpu.eplb_utils import EPLBController + + for name in ( + "prepare_load", + "maybe_register_model", + "maybe_start_async_loop", + "step", + "prepare_forward", + ): + real = inspect.signature(getattr(EPLBController, name)) + fake = inspect.signature(getattr(FakeEPLBController, name)) + real_params = [ + (param.name, param.default is not inspect.Parameter.empty) + for param in real.parameters.values() + ] + fake_params = [ + (param.name, param.default is not inspect.Parameter.empty) + for param in fake.parameters.values() + ] + assert fake_params == real_params, name + + +def test_ffn_runner_steps_eplb_once_per_real_step(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.connector.attn_outputs.append(_payload("hidden", _metadata())) + + runner.execute_model(dp_metadata_list={0: _FakeDPMetadata([1])}) + + assert runner.eplb.steps == [(False, False)] + + +def test_ffn_runner_marks_warmup_and_capturing_steps_dummy(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.connector.attn_outputs.extend( + _payload("hidden", _metadata()) for _ in range(2) + ) + + runner.execute_model(dp_metadata_list={0: _FakeDPMetadata([1])}, is_warmup=True) + runner.execute_model( + dp_metadata_list={0: _FakeDPMetadata([1])}, is_graph_capturing=True + ) + + assert runner.eplb.steps == [(True, False), (True, False)] + + +def test_ffn_runner_replay_steps_eplb_as_real_step(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.use_cuda_graph = True + graph = _FakeGraph() + dp_metadata = {0: _FakeDPMetadata([1])} + runner._cuda_graphs = { + make_ffn_graph_key(dp_metadata): {"graph": graph}, + } + + runner.execute_model(dp_metadata_list=dp_metadata, is_graph_replaying=True) + + assert graph.replay_count == 1 + assert runner.eplb.steps == [(False, False)] + + +def test_ffn_runner_prepares_eplb_unpadded_tokens_from_dp_metadata(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.vllm_config.parallel_config.enable_eplb = True + runner.connector.attn_size = 4 + runner.connector.attn_outputs.append(_payload("hidden", _metadata())) + + runner.execute_model(dp_metadata_list={0: _FakeDPMetadata([4, 6])}) + + # Two DP counts of 4 and 6 tokens, each replicated across two Attention + # TP peers, give 20 real tokens for the router's record kernel. + assert runner.eplb.prepared_forwards == [(runner.model_config, 20)] + + +def test_ffn_runner_skips_eplb_prepare_when_eplb_is_disabled(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.connector.attn_outputs.append(_payload("hidden", _metadata())) + + runner.execute_model(dp_metadata_list={0: _FakeDPMetadata([4, 6])}) + + assert runner.eplb.prepared_forwards == [] + + +def test_ffn_runner_rejects_multistage_eplb_payload(): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.vllm_config.parallel_config.enable_eplb = True + two_stage = {0: _FakeDPMetadata([1]), 1: _FakeDPMetadata([1])} + + with pytest.raises(RuntimeError, match="exactly one stage"): + runner.execute_model(dp_metadata_list=two_stage) + + +def test_ffn_runner_capture_prepared_and_steps_eplb_as_dummy(monkeypatch): + runner = _runner_with_connector_and_model(_FakeModel()) + runner.use_cuda_graph = True + runner.vllm_config.parallel_config.enable_eplb = True + monkeypatch.setattr(runner, "_dummy_run", lambda **_kwargs: None) + free_gpu_memory = [100, 90] + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (free_gpu_memory.pop(0), 0)) + monkeypatch.setattr(torch.cuda, "graph_pool_handle", lambda: object()) + monkeypatch.setattr( + ffn_model_runner_module, "graph_capture", lambda **_kwargs: nullcontext() + ) + monkeypatch.setattr( + ffn_model_runner_module, "set_cudagraph_capturing_enabled", lambda _v: None + ) + + cuda_graph_size = runner.capture_model(dp_metadata_list={0: _FakeDPMetadata([2])}) + + assert cuda_graph_size == 10 + assert runner.eplb.steps == [(True, False)] + assert runner.eplb.prepared_forwards == [(runner.model_config, 2)] + + +class _FakeEplbModelState: + """The EplbModelState fields the placement log reads.""" + + def __init__(self, placement: list[list[int]]): + self.physical_to_logical_map = torch.tensor(placement) + self.expert_load_window = torch.zeros( + (2, len(placement), len(placement[0])), dtype=torch.int32 + ) + self.model_name = "test-model" + self.rebalanced = False + + +class _FakeEplbState: + """Steps the counter and commits a new placement at the boundary.""" + + def __init__(self, placements: list[list[list[int]]]): + self.model_states = {"model-hash": _FakeEplbModelState(placements[0])} + self.expert_rearrangement_step_interval = 2 + self.expert_rearrangement_step = 0 + self._placements = placements + self._rearrangements = 0 + + def step(self, is_dummy=False, is_profile=False): + self.expert_rearrangement_step += 1 + if self.expert_rearrangement_step < self.expert_rearrangement_step_interval: + return + self.expert_rearrangement_step = 0 + self._rearrangements += 1 + self.model_states["model-hash"].physical_to_logical_map = torch.tensor( + self._placements[min(self._rearrangements, len(self._placements) - 1)] + ) + + +def _runner_for_placement_logging(state, *, log_placement: bool): + runner = _runner_with_connector_and_model(_FakeModel()) + # The runner's EPLB seam needs only the state and its step; two FFN EP + # ranks split the four-slot placement. + runner.eplb = SimpleNamespace(state=state, step=state.step) + runner.vllm_config.parallel_config.data_parallel_size = 2 + runner.vllm_config.parallel_config.eplb_config = SimpleNamespace( + log_balancedness=log_placement, + ) + return runner + + +def test_ffn_runner_logs_placement_when_a_rearrangement_commits( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """With log_balancedness the FFN logs the committed expert placement + when EplbState.step rearranges, and only then. + + Why: vLLM's own log lists load stats, not which expert went where. + """ + monkeypatch.setattr( + ffn_model_runner_module, + "get_ep_group", + lambda: SimpleNamespace(rank_in_group=0), + ) + state = _FakeEplbState( + [ + [[0, 1, 2, 3], [0, 1, 2, 3]], + [[0, 2, 1, 3], [0, 2, 1, 3]], + ] + ) + runner = _runner_for_placement_logging(state, log_placement=True) + + with caplog.at_level(logging.INFO, logger=ffn_model_runner_module.logger.name): + runner.eplb_step() + # Not due yet: this step only advanced the counter. + assert caplog.messages == [] + runner.eplb_step() + + message = caplog.messages[0] + assert "AFD FFN EPLB chose experts for model test-model" in message + assert "rank 0=[0, 2] | rank 1=[1, 3]" in message + # A step without a further rearrangement stays silent. + caplog.clear() + with caplog.at_level(logging.INFO, logger=ffn_model_runner_module.logger.name): + runner.eplb_step() + assert caplog.messages == [] + + +def test_ffn_runner_placement_logging_needs_balancedness_flag( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """Without log_balancedness no placement is logged, even at a + rearrangement boundary.""" + monkeypatch.setattr( + ffn_model_runner_module, + "get_ep_group", + lambda: SimpleNamespace(rank_in_group=0), + ) + state = _FakeEplbState([[[0, 1, 2, 3], [0, 1, 2, 3]]]) + runner = _runner_for_placement_logging(state, log_placement=False) + + with caplog.at_level(logging.INFO, logger=ffn_model_runner_module.logger.name): + runner.eplb_step() + runner.eplb_step() + + assert caplog.messages == [] + + +def test_ffn_runner_placement_logging_only_on_ep_rank_zero( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +): + """Every FFN EP rank holds the same map; only EP rank 0 logs it.""" + monkeypatch.setattr( + ffn_model_runner_module, + "get_ep_group", + lambda: SimpleNamespace(rank_in_group=1), + ) + state = _FakeEplbState([[[0, 1, 2, 3], [0, 1, 2, 3]]]) + runner = _runner_for_placement_logging(state, log_placement=True) + + with caplog.at_level(logging.INFO, logger=ffn_model_runner_module.logger.name): + runner.eplb_step() + runner.eplb_step() + + assert caplog.messages == [] + + def test_ffn_runner_releases_owned_runtime_state_on_shutdown(monkeypatch): runner = _runner_with_connector_and_model(_FakeModel()) runner.prof = _StepProfiler() diff --git a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py index a28a99820..c1d340939 100644 --- a/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_routed_experts_model_runner.py @@ -20,6 +20,7 @@ from tests.unit.fakes.connector import ( # noqa: E402 FakeRoutedConnector, ) +from tests.unit.fakes.eplb import FakeEPLBController # noqa: E402 MOE_LAYERS = (1, 3) NUM_LOCAL_EXPERTS = 2 @@ -30,6 +31,10 @@ def _runner(connector=None): runner = object.__new__(GPUFFNRoutedExpertsModelRunner) + runner.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(enable_eplb=False), + ) + runner.eplb = FakeEPLBController() runner.connector = connector or FakeRoutedConnector() runner.moe_layer_indices = MOE_LAYERS runner.local_experts = { @@ -133,6 +138,42 @@ def test_eager_fallback_keeps_attention_replay_flag(): assert payload.is_graph_capturing is False +def test_eplb_step_delegates_to_the_connector(): + """The routed runner's EPLB step is the connector's FFN EPLB driver. + + Why: the FFN never routes on this path, so the connector, not vLLM's + EPLBController, owns the rearrangement cadence. + """ + + class _SteppingConnector(FakeRoutedConnector): + def step_ffn_eplb(self) -> None: + self.events.append(("eplb_step",)) + + connector = _SteppingConnector() + runner = _runner(connector) + + runner.eplb_step(is_dummy=True) + runner.eplb_step() + + assert connector.events == [("eplb_step",), ("eplb_step",)] + + +@pytest.mark.parametrize("enable_eplb", [False, True]) +def test_connector_init_receives_eplb_state_only_when_eplb_is_enabled(enable_eplb): + """The connector builds its FFN EPLB driver from the runner's state. + + Why: without the state the connector cannot rearrange, and with EPLB off + it must not try. + """ + runner = _runner() + runner.vllm_config.parallel_config.enable_eplb = enable_eplb + + runner.initialize_afd_connector() + + expected = runner.eplb.state if enable_eplb else None + assert runner.connector.init_kwargs["eplb_state"] is expected + + def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): def base_init(self, vllm_config, device): self.connector = SimpleNamespace(close=lambda: None)