From 9195be2ffa0ed5e636e26efea03d09a77b817d9e Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Mon, 26 Oct 2026 16:00:18 +0200 Subject: [PATCH 01/12] initial eplb support Signed-off-by: Shlomit Koyfman --- afd_plugin/v1/worker/ffn_metadata.py | 68 ++++++-- afd_plugin/v1/worker/ffn_model_runner.py | 66 +++++++- afd_plugin/validation.py | 15 +- docs/design/module/ffn_runtime.md | 37 ++++- .../2a2f_graph_eplb_dp1tp2.sh | 63 ++++++++ .../multipod_2a_2f_graph_eplb_dp1tp2.sh | 75 +++++++++ tests/unit/v1/worker/test_ffn_metadata.py | 26 ++++ tests/unit/v1/worker/test_ffn_model_runner.py | 145 ++++++++++++++++++ tests/unit/v1/worker/test_model_runner_v2.py | 33 +++- 9 files changed, 512 insertions(+), 16 deletions(-) create mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/2a2f_graph_eplb_dp1tp2.sh create mode 100644 recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/prefill_decode_colocation/multipod_2a_2f_graph_eplb_dp1tp2.sh 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 e31d6c239..d0f61bcce 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -16,6 +16,7 @@ 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 ( @@ -40,6 +41,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 @@ -91,6 +93,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 +115,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 +131,12 @@ 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, + load_dummy_weights, + ) + self.eplb.maybe_start_async_loop(eplb_models_added) def profile_run(self) -> None: pass @@ -163,6 +176,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 +188,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 self._ffn_forward( @@ -181,8 +196,55 @@ def execute_model( is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, ) + self.eplb_step(is_dummy=is_warmup or is_graph_capturing) return None + # Patch reason: vLLM 0.30.0 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 v0.30.0, vllm/v1/worker/gpu/eplb_utils.py, + # step_eplb_after; commit ced6857afa0ea7b2e3f0846a62e1394e90f15607. + def eplb_step(self, is_dummy: bool = False) -> None: + """Step the FFN-side EPLB state after one connector-driven step.""" + + self.eplb.step(is_dummy=is_dummy) + + 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, *, @@ -365,6 +427,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: @@ -400,6 +463,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/validation.py b/afd_plugin/validation.py index ad4ce5c27..97842c31f 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -118,11 +118,24 @@ 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 or vllm_config.compilation_config.pass_config.enable_sp ): 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", + ) if ( parallel.enable_dbo or parallel.use_ubatching ) and parallel.num_ubatches != GPU_MRV2_NUM_UBATCHES: diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index bb28bd4b4..c87e83d41 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -153,7 +153,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. H20 validation covers DeepSeek-V2-Lite DP1/TP2 Graph + DBO with 128 GSM8K samples and live two-stage/full-replay assertions. @@ -300,6 +303,38 @@ expected by the upstream forward context. When TP expands Attention rank counts beyond DP metadata width, token counts are replicated per TP rank and then aggregated for each FFN rank. +## FFN-side EPLB + +The FFN role can enable vLLM 0.30'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). + +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/`. + ## Role-aware model loading The FFN process uses vLLM's model loader, while plugin-owned model integration 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/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 eea8a4aa2..0931401d2 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -118,6 +118,33 @@ def __init__(self): self.control_plane = None +class _FakeEplbController: + """Records the EPLB calls the FFN runner makes per connector step.""" + + def __init__(self): + 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, load_dummy_weights): + self.registrations.append((model, model_config, load_dummy_weights)) + return True + + def maybe_start_async_loop(self, eplb_models_added): + self.async_loop_starts.append(eplb_models_added) + + 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)) + + class _FakeModel: def get_experts_layer_indices(self): return () @@ -156,6 +183,7 @@ 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, use_sequence_parallel_moe=False, ), @@ -166,12 +194,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 @@ -182,6 +215,7 @@ def test_ffn_runner_owns_connector_execution_without_native_runner_inheritance() def test_ffn_constructor_defers_connector_init_and_owns_model_and_graphs(monkeypatch): config = SimpleNamespace( use_v2_model_runner=True, + parallel_config=SimpleNamespace(enable_eplb=False), model_config=SimpleNamespace( dtype=torch.bfloat16, hf_text_config=SimpleNamespace(num_hidden_layers=2), @@ -727,6 +761,117 @@ def test_ffn_runner_steps_gpu_profiler(): assert runner.prof.steps == 1 +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, False)] + assert runner.eplb.async_loop_starts == [True] + + +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)] + + 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_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index a6c68d0e1..2b882d74a 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -485,7 +485,38 @@ 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.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}, From 2aeb62f6eb39bc5ecde36127591dc8a4d6e7b5d0 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Wed, 28 Oct 2026 16:51:49 +0200 Subject: [PATCH 02/12] support EPLB in Qwen models Signed-off-by: Shlomit Koyfman --- afd_plugin/model_executor/models/qwen3_5.py | 20 +++-- afd_plugin/model_executor/models/qwen3_moe.py | 4 +- docs/design/module/model_integration.md | 14 +++- .../eplb_4a_2f_graph.sh | 75 ++++++++++++++++++ .../models/test_qwen3_5_contract.py | 77 +++++++++++++++++++ .../models/test_qwen3_moe_construction.py | 16 ++++ 6 files changed, 193 insertions(+), 13 deletions(-) create mode 100644 recipe/gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/eplb_4a_2f_graph.sh diff --git a/afd_plugin/model_executor/models/qwen3_5.py b/afd_plugin/model_executor/models/qwen3_5.py index 356a9326c..cce13bee5 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 @@ -293,8 +296,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 @@ -303,8 +306,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 @@ -471,8 +477,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 a7969b6ce..22445cc99 100644 --- a/afd_plugin/model_executor/models/qwen3_moe.py +++ b/afd_plugin/model_executor/models/qwen3_moe.py @@ -89,8 +89,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/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 14050d0d7..e7cb0f154 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -178,8 +178,11 @@ by vLLM. Checkpoint weights are filtered once by layer-stage path before the native loader consumes them. This path is CUDA-only and requires `compute_gate_on_attention=false`. It fails -during construction for sequence-parallel MoE, EPLB, pipeline parallelism, -speculative decoding, LoRA, or NPU. +during construction for sequence-parallel MoE, EPLB on the Attention role, +pipeline parallelism, speculative decoding, LoRA, or NPU; 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 @@ -190,8 +193,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. Earlier H20 Graph + DBO and recurrent-state/CUDA Graph parity checks describe the 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/unit/model_executor/models/test_qwen3_5_contract.py b/tests/unit/model_executor/models/test_qwen3_5_contract.py index cb97711ce..41fa3aa6f 100644 --- a/tests/unit/model_executor/models/test_qwen3_5_contract.py +++ b/tests/unit/model_executor/models/test_qwen3_5_contract.py @@ -260,6 +260,83 @@ def test_qwen_rejects_unsupported_parallel_feature_before_construction( adapter.AFDQwen3_5MoeForConditionalGeneration(vllm_config=vllm_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, + ), + ) + 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" + + def test_qwen_conditional_model_rejects_pipeline_parallelism_before_construction( monkeypatch, ): 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 dee09f2e9..06c4aee57 100644 --- a/tests/unit/model_executor/models/test_qwen3_moe_construction.py +++ b/tests/unit/model_executor/models/test_qwen3_moe_construction.py @@ -314,6 +314,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_rocm_platform_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( adapter, From c5d3bb5b8f42dfebab99bab22b3f884646915c2a Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Thu, 29 Oct 2026 12:45:51 +0200 Subject: [PATCH 03/12] support EPLB for M2N connector Signed-off-by: Shlomit Koyfman --- afd_plugin/connectors/base.py | 27 ++ .../connectors/gpu/deepep_m2n/connector.py | 176 ++++++++- afd_plugin/connectors/gpu/deepep_m2n/eplb.py | 361 ++++++++++++++++++ afd_plugin/connectors/metadata.py | 7 + .../model_executor/models/deepseek_v2.py | 18 +- afd_plugin/v1/worker/attention_metadata.py | 9 + .../worker/ffn_routed_experts_model_runner.py | 19 + docs/design/module/connector_contracts.md | 2 +- docs/design/module/ffn_runtime.md | 12 + docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 44 ++- recipe/gpu/DeepEPM2NAFDConnector/README.md | 2 +- .../qwen3_30b_a3b/README.md | 16 +- .../2a2f_graph_eplb_dp2tp1.sh | 63 +++ .../connectors/test_deepep_m2n_connector.py | 58 ++- tests/unit/connectors/test_deepep_m2n_eplb.py | 288 ++++++++++++++ tests/unit/connectors/test_p2p_connector.py | 2 + tests/unit/fakes/connector.py | 11 +- .../models/test_deepseek_v2_proxy.py | 54 +++ .../test_ffn_routed_experts_model_runner.py | 21 + tests/unit/v1/worker/test_model_runner_v2.py | 1 + 20 files changed, 1167 insertions(+), 24 deletions(-) create mode 100644 afd_plugin/connectors/gpu/deepep_m2n/eplb.py create mode 100644 recipe/gpu/DeepEPM2NAFDConnector/qwen3_30b_a3b/prefill_decode_colocation/2a2f_graph_eplb_dp2tp1.sh create mode 100644 tests/unit/connectors/test_deepep_m2n_eplb.py diff --git a/afd_plugin/connectors/base.py b/afd_plugin/connectors/base.py index d0bb01770..06aad44d3 100644 --- a/afd_plugin/connectors/base.py +++ b/afd_plugin/connectors/base.py @@ -22,6 +22,7 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.distributed.eplb.eplb_state import EplbState @dataclass(frozen=True) @@ -281,6 +282,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. @@ -289,6 +291,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 @@ -311,6 +316,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 adfc5c620..7d571dba5 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -81,6 +81,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, @@ -439,6 +444,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: @@ -456,6 +465,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 @@ -463,11 +473,25 @@ 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.""" + text_config = vllm_config.model_config.hf_text_config + num_experts = int(text_config.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, @@ -536,11 +560,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") @@ -549,7 +577,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], } @@ -564,6 +591,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( @@ -597,11 +637,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: @@ -645,6 +710,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(), ) @@ -924,11 +991,38 @@ def __init__( self.layout = DeepEPM2NLayout( ffn_size=afd_config.num_ffn_ranks, attn_size=afd_config.num_attention_ranks, - num_experts=text_config.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 int(text_config.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 + if ( + afd_config.role == "ffn" + and vllm_config.parallel_config.enable_eplb + and vllm_config.parallel_config.eplb_config.use_async + ): + raise RuntimeError( + "DeepEP M2N EPLB does not support async EPLB: the async " + "worker commits new expert maps outside the connector step " + "boundary the Attention ranks adopt them at. Pass " + "use_async=false in --eplb-config." + ) 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: @@ -1030,7 +1124,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, @@ -1056,6 +1150,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 @@ -1069,6 +1173,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. @@ -1077,6 +1182,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 @@ -1099,6 +1207,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, + ) self.group = ( self.bootstrap_group if self.config.control_backend == "gloo" @@ -1170,6 +1301,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 @@ -1391,6 +1530,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 +1746,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..d11fef779 --- /dev/null +++ b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py @@ -0,0 +1,361 @@ +# 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. +""" + +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.model_executor.layers.fused_moe.router.base_router import ( + eplb_map_to_physical_and_record, +) + +from afd_plugin.connectors.metadata import AFDControlPayload + +if TYPE_CHECKING: + from afd_plugin.connectors.gpu.deepep_m2n.connector import DeepEPM2NAFDConnector + +# 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", +) + + +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"]) + 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 is_dummy: + self.expert_load_pass.zero_() + else: + # 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) + self.expert_load_pass.zero_() + window_size = self.expert_load_window.shape[0] + self._window_step = (self._window_step + 1) % window_size + self._rearrange_step += 1 + if self._rearrange_step >= self._step_interval: + self._rearrange_step = 0 + self._exchange_at_boundary() + + 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 + window_cpu = self.expert_load_window.cpu() + 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)) + + # ============================== + # 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. + """ + + 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 + + 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), + } + 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; runs the boundary exchange when due.""" + self._rearrange_step += 1 + if self._rearrange_step < self._step_interval: + return + self._rearrange_step = 0 + self._rearrange() + + 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)) + self._state.rearrange() + 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) + + +__all__ = [ + "DeepEPM2NAttentionEplb", + "DeepEPM2NFFNEplb", + "real_rows_in_tp_slice", + "recv_initial_eplb_state", +] diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 30c63db0f..c78fab9b9 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -132,6 +132,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] @@ -139,6 +143,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 = { @@ -393,6 +398,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", @@ -423,6 +429,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/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 6433f6f92..89694e8d4 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -31,6 +31,7 @@ from afd_plugin.connectors import ( AFDExpertRoutingSpec, AFDF2ATransferPayload, + AFDRoutedExpertsConnector, AFDTransferContext, AFDTransferMetadata, ) @@ -151,9 +152,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 @@ -212,6 +215,15 @@ def _send_and_receive( afd_metadata = get_afd_metadata_from_forward_context() if afd_metadata is None: raise RuntimeError("RemoteFFNProxy requires AFD forward 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/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_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/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index b74b76fec..59b7cb3d2 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -63,7 +63,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) | [`test_base_factory.py`](../../../tests/unit/connectors/test_base_factory.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) | +| 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) | | 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) | [`test_async_cam_connector.py`](../../../tests/unit/connectors/test_async_cam_connector.py), [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v2_lite/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 a78fa665d..a0afa44be 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -327,6 +327,18 @@ 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. + ## FFN-side EPLB The FFN role can enable vLLM 0.30's expert-parallel load balancing diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 0d076c624..777e02a0c 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -68,9 +68,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), `control_backend`, `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 @@ -207,7 +209,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. | @@ -279,12 +281,42 @@ and FP8 examples, see the - SM90 or newer on every rank; the handshake rejects older GPUs. - Qwen3 MoE only, with every layer MoE and no shared expert. -- `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, and +`log_balancedness` has no effect (no FFN-side step runs). EPLB exchanges run +on the connector's CPU bootstrap group at each rearrangement boundary, so a +rearrangement adds a host sync on every rank. + +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 da6ffbcbb..8b14924ec 100644 --- a/recipe/gpu/DeepEPM2NAFDConnector/README.md +++ b/recipe/gpu/DeepEPM2NAFDConnector/README.md @@ -13,7 +13,7 @@ supported layouts. . ├── 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_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/tests/unit/connectors/test_deepep_m2n_connector.py b/tests/unit/connectors/test_deepep_m2n_connector.py index 9a47fbdf2..b2933cfd5 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, } @@ -1056,7 +1065,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"), ], @@ -1078,6 +1086,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..d107c6a75 --- /dev/null +++ b/tests/unit/connectors/test_deepep_m2n_eplb.py @@ -0,0 +1,288 @@ +# 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 + +from types import SimpleNamespace + +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 = 8 +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, + } + 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(**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(), + 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[torch.Tensor, int]] = [] + self.broadcasts: list[torch.Tensor] = [] + 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 + ) + + +def _ffn_state(): + model = SimpleNamespace( + num_moe_layers=len(MOE_LAYERS), + num_logical_experts=NUM_LOGICAL, + num_physical_experts=NUM_PHYSICAL, + ) + model_state = SimpleNamespace( + model=model, + expert_load_window=torch.zeros((WINDOW_SIZE, len(MOE_LAYERS), NUM_PHYSICAL)), + logical_to_physical_map=torch.zeros((len(MOE_LAYERS), NUM_LOGICAL, 4)), + logical_replica_count=torch.zeros((len(MOE_LAYERS), NUM_LOGICAL)), + ) + 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 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 = SimpleNamespace(bootstrap_group=object(), 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 = SimpleNamespace(bootstrap_group=object(), 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 == [] diff --git a/tests/unit/connectors/test_p2p_connector.py b/tests/unit/connectors/test_p2p_connector.py index 04f25048c..9a80794fc 100644 --- a/tests/unit/connectors/test_p2p_connector.py +++ b/tests/unit/connectors/test_p2p_connector.py @@ -467,12 +467,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 = module.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/fakes/connector.py b/tests/unit/fakes/connector.py index 5795633fd..f7bd339c2 100644 --- a/tests/unit/fakes/connector.py +++ b/tests/unit/fakes/connector.py @@ -37,10 +37,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 @@ -49,7 +56,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/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index 6ede9f77b..c2378deaf 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -16,6 +16,7 @@ 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 FakeRoutedConnector # noqa: E402 ROUTING_HIDDEN_SIZE = 16 ROUTING_NUM_EXPERTS = 8 @@ -228,6 +229,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_fake_forward_context(monkeypatch, events) 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 75f4fb81d..4bdd9a913 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 @@ -188,9 +188,30 @@ def test_initialize_passes_moe_layers_and_experts_to_the_connector(): assert runner.connector.init_kwargs == { "moe_layer_indices": MOE_LAYERS, "local_experts": runner.local_experts, + "eplb_state": None, } +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",)] + + def test_constructor_rejects_connectors_without_routed_rows(monkeypatch): def base_init(self, vllm_config, device): self.connector = SimpleNamespace(close=lambda: None) diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 450e07de3..35e3c806a 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -2095,6 +2095,7 @@ def test_v2_load_model_reports_moe_layers_to_routed_connector(monkeypatch): assert connector.init_kwargs == { "moe_layer_indices": (1, 3), "local_experts": None, + "eplb_state": None, } From 0bb1e8a3121f0cd0306b534dc9a5103cf66b4299 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Thu, 29 Oct 2026 12:58:38 +0200 Subject: [PATCH 04/12] fix pre-commit failures in the M2N EPLB change The pre-commit mypy run (pinned 1.11.1, follow-imports=skip for tests) flagged two issues the local dev env run missed: - The new M2N EPLB test annotated fake-dist buffers with torch.Tensor; under follow-imports=skip torch is not resolvable. Drop the annotations. - metadata.py, touched by the EPLB change, still carried pre-existing one-argument Generator[...] annotations that mypy 1.11 rejects; complete them to Generator[..., None, None]. --- afd_plugin/connectors/metadata.py | 6 ++++-- tests/unit/connectors/test_deepep_m2n_eplb.py | 5 +++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index c78fab9b9..8c60895a1 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -46,7 +46,9 @@ def __post_init__(self) -> None: ) @contextmanager - def sp_local_sizes(self, sequence_parallel_size: int) -> Generator[list[int]]: + def sp_local_sizes( + self, sequence_parallel_size: int + ) -> Generator[list[int], None, None]: local_sizes: list[int] = ( ( (self.num_tokens_across_dp_cpu + sequence_parallel_size - 1) @@ -78,7 +80,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) diff --git a/tests/unit/connectors/test_deepep_m2n_eplb.py b/tests/unit/connectors/test_deepep_m2n_eplb.py index d107c6a75..19a062574 100644 --- a/tests/unit/connectors/test_deepep_m2n_eplb.py +++ b/tests/unit/connectors/test_deepep_m2n_eplb.py @@ -10,6 +10,7 @@ from __future__ import annotations from types import SimpleNamespace +from typing import Any import pytest @@ -89,8 +90,8 @@ class _FakeDist: receiving tensor so adoption is observable.""" def __init__(self, *, map_value: int = 7) -> None: - self.reduces: list[tuple[torch.Tensor, int]] = [] - self.broadcasts: list[torch.Tensor] = [] + self.reduces: list[tuple[Any, int]] = [] + self.broadcasts: list[Any] = [] self._map_value = map_value def reduce(self, tensor, dst, group): From 157aeaa22c6d679475d96b43b51a6b7759484780 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Thu, 29 Oct 2026 15:09:00 +0200 Subject: [PATCH 05/12] log the experts EPLB chooses on the M2N path vLLM logs EPLB balancedness from EplbState.step, which the M2N driver never runs (routing, and load recording, are on Attention). Wire the existing log_balancedness eplb-config flag to the M2N boundary instead: when set, FFN world rank 0 logs, per MoE layer, the logical experts each FFN rank holds after a rearrangement, the replicated ones, and the recorded per-rank token load that drove the choice. --- afd_plugin/connectors/gpu/deepep_m2n/eplb.py | 79 ++++++++++++++++ docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 15 ++- tests/unit/connectors/test_deepep_m2n_eplb.py | 91 ++++++++++++++++++- 3 files changed, 179 insertions(+), 6 deletions(-) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py index d11fef779..35597e35e 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py @@ -46,6 +46,7 @@ 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, ) @@ -55,6 +56,8 @@ 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. @@ -121,6 +124,48 @@ def real_rows_in_tp_slice( return min(max(num_real_tokens - tp_rank * chunk, 0), chunk) +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 _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 @@ -306,6 +351,11 @@ def __init__( 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. + self._log_placement = bool( + connector.vllm_config.parallel_config.eplb_config.log_balancedness + ) def broadcast_initial_state(self) -> None: """Collective: tell every rank the EPLB schedule and initial maps. @@ -347,15 +397,44 @@ def _rearrange(self) -> None: if self._connector.world_rank == 0: window.copy_(staging.to(window.device)) self._state.rearrange() + if self._log_placement and self._connector.world_rank == 0: + self._log_chosen_experts(staging) 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 + slots_per_rank = load_window.shape[2] // ffn_size + rank_loads = load_window.sum(dim=(0, 1)).tolist() + per_rank_parts = [] + for rank in range(ffn_size): + slots = rank_loads[rank * slots_per_rank : (rank + 1) * slots_per_rank] + per_rank_parts.append(f"rank {rank}={int(sum(slots))} tokens") + per_rank = ", ".join(per_rank_parts) + 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], + per_rank, + "\n".join(lines), + ) + __all__ = [ "DeepEPM2NAttentionEplb", "DeepEPM2NFFNEplb", + "describe_eplb_placement", "real_rows_in_tp_slice", "recv_initial_eplb_state", ] diff --git a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 777e02a0c..32108851b 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -309,10 +309,17 @@ the two EPLB duties split across the roles: 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, and -`log_balancedness` has no effect (no FFN-side step runs). EPLB exchanges run -on the connector's CPU bootstrap group at each rearrangement boundary, so a -rearrangement adds a host sync on every rank. +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 which experts EPLB chose, set `"log_balancedness": true` in the FFN's +`--eplb-config`: at every rearrangement the FFN then logs, 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 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). Run the recipe from the repository root: diff --git a/tests/unit/connectors/test_deepep_m2n_eplb.py b/tests/unit/connectors/test_deepep_m2n_eplb.py index 19a062574..18cac598b 100644 --- a/tests/unit/connectors/test_deepep_m2n_eplb.py +++ b/tests/unit/connectors/test_deepep_m2n_eplb.py @@ -9,6 +9,7 @@ from __future__ import annotations +import logging from types import SimpleNamespace from typing import Any @@ -22,6 +23,7 @@ from afd_plugin.connectors.gpu.deepep_m2n.eplb import ( # noqa: E402 DeepEPM2NAttentionEplb, DeepEPM2NFFNEplb, + describe_eplb_placement, real_rows_in_tp_slice, ) from tests.unit.fakes.connector import FakeRoutedConnector # noqa: E402 @@ -202,9 +204,15 @@ def _ffn_state(): 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] model_state = SimpleNamespace( model=model, expert_load_window=torch.zeros((WINDOW_SIZE, len(MOE_LAYERS), NUM_PHYSICAL)), + physical_to_logical_map=torch.tensor([layer0, layer1]), logical_to_physical_map=torch.zeros((len(MOE_LAYERS), NUM_LOGICAL, 4)), logical_replica_count=torch.zeros((len(MOE_LAYERS), NUM_LOGICAL)), ) @@ -222,6 +230,19 @@ def rearrange(): return state, model_state, rearrangements +def _ffn_connector(world_rank: int, *, log_placement: bool = False): + 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), + ), + ), + ) + + 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 @@ -231,7 +252,7 @@ def test_ffn_driver_rearranges_only_at_the_boundary(monkeypatch: pytest.MonkeyPa monkeypatch.setattr(eplb_module.dist, "broadcast", fake_dist.broadcast) state, model_state, rearrangements = _ffn_state() state.expert_rearrangement_step_interval = 10 - connector = SimpleNamespace(bootstrap_group=object(), world_rank=1) + connector = _ffn_connector(world_rank=1) driver = DeepEPM2NFFNEplb(connector, state) for _ in range(10 - INITIAL_REARRANGE_STEP - 1): @@ -265,7 +286,7 @@ def reduce(tensor, dst, group): monkeypatch.setattr(eplb_module.dist, "broadcast", _FakeDist().broadcast) state, model_state, rearrangements = _ffn_state() state.expert_rearrangement_step_interval = 10 - connector = SimpleNamespace(bootstrap_group=object(), world_rank=0) + connector = _ffn_connector(world_rank=0) driver = DeepEPM2NFFNEplb(connector, state) for _ in range(10 - INITIAL_REARRANGE_STEP): @@ -287,3 +308,69 @@ def test_routed_connector_base_seams_default_to_no_op(): assert torch.equal(connector.apply_eplb_routing(ids, 0), ids) connector.step_ffn_eplb() assert connector.events == [] + + +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.""" + _state, model_state, _rearrangements = _ffn_state() + + lines = describe_eplb_placement( + model_state.physical_to_logical_map, + 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" + ) + # Layer 1 swaps logical 5 and 6 between the ranks and drops the replica. + 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_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 == [] From 95f8d149cb0a90a60628ef04b6d31dc28936d993 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Thu, 29 Oct 2026 15:29:21 +0200 Subject: [PATCH 06/12] log the experts EPLB chooses on the P2P path too The placement log added for DeepEP M2N now also runs on the P2P path, where EplbState.step commits rearrangements itself: GPUFFNModelRunner's eplb_step snapshots the committed map only on steps where the counter boundary or a pending async transfer can commit, and logs the placement (plus vLLM's balancedness stats, which log_balancedness already drove) when the map actually changed, on EP rank 0 only. The placement formatters move to afd_plugin/distributed/eplb.py so the shared runner and the DeepEP M2N connector module use one implementation without the runner importing an M2N-specific module. --- afd_plugin/connectors/gpu/deepep_m2n/eplb.py | 55 +------- afd_plugin/distributed/eplb.py | 81 ++++++++++++ afd_plugin/v1/worker/ffn_model_runner.py | 73 ++++++++++- docs/design/module/ffn_runtime.md | 9 ++ tests/unit/connectors/test_deepep_m2n_eplb.py | 23 ---- tests/unit/distributed/test_eplb.py | 52 ++++++++ tests/unit/fakes/eplb.py | 1 + tests/unit/v1/worker/test_ffn_model_runner.py | 119 ++++++++++++++++++ 8 files changed, 339 insertions(+), 74 deletions(-) create mode 100644 afd_plugin/distributed/eplb.py create mode 100644 tests/unit/distributed/test_eplb.py diff --git a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py index 35597e35e..6b20b0dd6 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py @@ -52,6 +52,10 @@ ) from afd_plugin.connectors.metadata import AFDControlPayload +from afd_plugin.distributed.eplb import ( + describe_eplb_placement, + describe_eplb_rank_load, +) if TYPE_CHECKING: from afd_plugin.connectors.gpu.deepep_m2n.connector import DeepEPM2NAFDConnector @@ -124,48 +128,6 @@ def real_rows_in_tp_slice( return min(max(num_real_tokens - tp_rank * chunk, 0), chunk) -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 _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 @@ -412,13 +374,6 @@ def _log_chosen_experts(self, load_window: torch.Tensor) -> None: ``(window_size, num_moe_layers, num_physical_experts)``. """ ffn_size = self._connector.ffn_size - slots_per_rank = load_window.shape[2] // ffn_size - rank_loads = load_window.sum(dim=(0, 1)).tolist() - per_rank_parts = [] - for rank in range(ffn_size): - slots = rank_loads[rank * slots_per_rank : (rank + 1) * slots_per_rank] - per_rank_parts.append(f"rank {rank}={int(sum(slots))} tokens") - per_rank = ", ".join(per_rank_parts) lines = describe_eplb_placement( self._model_state.physical_to_logical_map.cpu(), ffn_size=ffn_size, @@ -426,7 +381,7 @@ def _log_chosen_experts(self, load_window: torch.Tensor) -> None: logger.info( "DeepEP M2N EPLB chose experts (recorded load over %d steps: %s):\n%s", load_window.shape[0], - per_rank, + describe_eplb_rank_load(load_window, ffn_size=ffn_size), "\n".join(lines), ) diff --git a/afd_plugin/distributed/eplb.py b/afd_plugin/distributed/eplb.py new file mode 100644 index 000000000..13e07e5a7 --- /dev/null +++ b/afd_plugin/distributed/eplb.py @@ -0,0 +1,81 @@ +# 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 + +import torch + + +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 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__ = [ + "describe_eplb_placement", + "describe_eplb_rank_load", +] diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index 5aa546a07..df43fe648 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -11,8 +11,13 @@ 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 @@ -31,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, @@ -50,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. @@ -217,7 +228,67 @@ def execute_model( 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. + if get_ep_group().rank() != 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, diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index a0afa44be..bce773211 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -358,6 +358,15 @@ The FFN role can enable vLLM 0.30's expert-parallel load balancing (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 +below); 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 diff --git a/tests/unit/connectors/test_deepep_m2n_eplb.py b/tests/unit/connectors/test_deepep_m2n_eplb.py index 18cac598b..45ac4dd36 100644 --- a/tests/unit/connectors/test_deepep_m2n_eplb.py +++ b/tests/unit/connectors/test_deepep_m2n_eplb.py @@ -23,7 +23,6 @@ from afd_plugin.connectors.gpu.deepep_m2n.eplb import ( # noqa: E402 DeepEPM2NAttentionEplb, DeepEPM2NFFNEplb, - describe_eplb_placement, real_rows_in_tp_slice, ) from tests.unit.fakes.connector import FakeRoutedConnector # noqa: E402 @@ -310,28 +309,6 @@ def test_routed_connector_base_seams_default_to_no_op(): assert connector.events == [] -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.""" - _state, model_state, _rearrangements = _ffn_state() - - lines = describe_eplb_placement( - model_state.physical_to_logical_map, - 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" - ) - # Layer 1 swaps logical 5 and 6 between the ranks and drops the replica. - 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_ffn_driver_logs_chosen_experts_when_balancedness_logging_is_on( monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture, diff --git a/tests/unit/distributed/test_eplb.py b/tests/unit/distributed/test_eplb.py new file mode 100644 index 000000000..9a7a53e50 --- /dev/null +++ b/tests/unit/distributed/test_eplb.py @@ -0,0 +1,52 @@ +# 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, +) + +# 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" diff --git a/tests/unit/fakes/eplb.py b/tests/unit/fakes/eplb.py index d9e0cbeea..777c1b0ee 100644 --- a/tests/unit/fakes/eplb.py +++ b/tests/unit/fakes/eplb.py @@ -9,6 +9,7 @@ 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 = [] diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index cbb6f2632..6e56d3294 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -159,6 +159,7 @@ def _runner_with_connector_and_model(model, *, num_layers=1): data_parallel_size=1, enable_eplb=False, is_moe_model=True, + tensor_parallel_size=1, use_sequence_parallel_moe=False, ), compilation_config=SimpleNamespace( @@ -846,6 +847,124 @@ def test_ffn_runner_capture_prepared_and_steps_eplb_as_dummy(monkeypatch): 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=lambda: 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=lambda: 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=lambda: 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() From 433d840027394ed2aead337f98a5f46c5d9b7ac1 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Thu, 29 Oct 2026 17:40:50 +0200 Subject: [PATCH 07/12] bug fix Signed-off-by: Shlomit Koyfman --- afd_plugin/v1/worker/ffn_model_runner.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index df43fe648..a8362d2c0 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -250,7 +250,9 @@ def _snapshot_eplb_placement_if_due(self) -> dict[str, torch.Tensor] | None: 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. - if get_ep_group().rank() != 0: + # ``rank_in_group`` is an attribute in vLLM 0.30'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 From f7b65cdbc73b7e1e1a0c8457890f722620421c17 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Sun, 1 Nov 2026 15:56:41 +0200 Subject: [PATCH 08/12] add dockerfile for m2n connector addition Signed-off-by: Shlomit Koyfman --- docker/Dockerfile.m2n | 117 ++++++++++++++++++ tests/unit/v1/worker/test_ffn_model_runner.py | 12 +- 2 files changed, 126 insertions(+), 3 deletions(-) create mode 100644 docker/Dockerfile.m2n diff --git a/docker/Dockerfile.m2n b/docker/Dockerfile.m2n new file mode 100644 index 000000000..05c02fc80 --- /dev/null +++ b/docker/Dockerfile.m2n @@ -0,0 +1,117 @@ +# 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.30.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 + +# 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/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index 6e56d3294..229ec9bac 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -902,7 +902,9 @@ def test_ffn_runner_logs_placement_when_a_rearrangement_commits( 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=lambda: 0) + ffn_model_runner_module, + "get_ep_group", + lambda: SimpleNamespace(rank_in_group=0), ) state = _FakeEplbState( [ @@ -935,7 +937,9 @@ def test_ffn_runner_placement_logging_needs_balancedness_flag( """Without log_balancedness no placement is logged, even at a rearrangement boundary.""" monkeypatch.setattr( - ffn_model_runner_module, "get_ep_group", lambda: SimpleNamespace(rank=lambda: 0) + 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) @@ -953,7 +957,9 @@ def test_ffn_runner_placement_logging_only_on_ep_rank_zero( ): """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=lambda: 1) + 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) From 47e3fcfe112c038b7182cfc224d5ccbb1e0038b2 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Mon, 2 Nov 2026 11:35:08 +0200 Subject: [PATCH 09/12] update dockerfile for m2n connectors due to vllm version upgrade Signed-off-by: Shlomit Koyfman --- docker/Dockerfile.m2n | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/docker/Dockerfile.m2n b/docker/Dockerfile.m2n index 05c02fc80..c865e20f0 100644 --- a/docker/Dockerfile.m2n +++ b/docker/Dockerfile.m2n @@ -22,7 +22,7 @@ # 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.30.0 +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. @@ -71,6 +71,10 @@ RUN uv export --locked --extra afdep --no-group dev --no-group e2e-tests \ "$(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 From 91fba9838cc9c1952eb35621249c380d0dbd4c48 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Mon, 2 Nov 2026 15:59:09 +0200 Subject: [PATCH 10/12] changes caused by vllm version upgrade Signed-off-by: Shlomit Koyfman --- .../connectors/gpu/deepep_m2n/connector.py | 11 ------ afd_plugin/v1/worker/ffn_model_runner.py | 3 +- afd_plugin/validation.py | 12 +++++++ tests/unit/config/test_validation.py | 18 ++++++++++ tests/unit/fakes/eplb.py | 8 ++--- tests/unit/v1/worker/test_ffn_model_runner.py | 36 +++++++++++++++++-- tests/unit/v1/worker/test_model_runner_v2.py | 11 ++++++ 7 files changed, 80 insertions(+), 19 deletions(-) diff --git a/afd_plugin/connectors/gpu/deepep_m2n/connector.py b/afd_plugin/connectors/gpu/deepep_m2n/connector.py index da554b01e..5906325fb 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/connector.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/connector.py @@ -995,17 +995,6 @@ def __init__( # deployment. ``None`` when the FFN runs without EPLB. self.eplb: DeepEPM2NAttentionEplb | DeepEPM2NFFNEplb | None = None self._handshake_enable_eplb = False - if ( - afd_config.role == "ffn" - and vllm_config.parallel_config.enable_eplb - and vllm_config.parallel_config.eplb_config.use_async - ): - raise RuntimeError( - "DeepEP M2N EPLB does not support async EPLB: the async " - "worker commits new expert maps outside the connector step " - "boundary the Attention ranks adopt them at. Pass " - "use_async=false in --eplb-config." - ) self.attn_size = self.layout.attn_size self.ffn_size = self.layout.ffn_size self.is_ffn = afd_config.role == "ffn" diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index f654c723b..b26e61811 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -145,9 +145,8 @@ def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None eplb_models_added = self.eplb.maybe_register_model( self.model, self.model_config, - load_dummy_weights, ) - self.eplb.maybe_start_async_loop(eplb_models_added) + self.eplb.maybe_start_async_loop(eplb_models_added, load_dummy_weights) def profile_run(self) -> None: pass diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index 9c18b7451..be079b5c6 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -177,6 +177,18 @@ def validate_gpu_model_runner_v2_config( 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/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index b1db5afa7..d2b028bb2 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(mrv2_config, role): ) +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", diff --git a/tests/unit/fakes/eplb.py b/tests/unit/fakes/eplb.py index 777c1b0ee..3c8532218 100644 --- a/tests/unit/fakes/eplb.py +++ b/tests/unit/fakes/eplb.py @@ -19,12 +19,12 @@ def __init__(self): def prepare_load(self): self.prepare_loads += 1 - def maybe_register_model(self, model, model_config, load_dummy_weights): - self.registrations.append((model, model_config, load_dummy_weights)) + 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): - self.async_loop_starts.append(eplb_models_added) + 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)) diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index fee94ab15..c578a6618 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -871,8 +871,40 @@ def test_ffn_runner_loads_model_and_registers_eplb(monkeypatch): assert runner.model is model assert runner.eplb.prepare_loads == 1 - assert runner.eplb.registrations == [(model, runner.model_config, False)] - assert runner.eplb.async_loop_starts == [True] + 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(): diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index dc72a6636..68995a5ef 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -156,6 +156,7 @@ def _v2_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, @@ -508,6 +509,16 @@ def _expect(error): 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", From b37fe3d636eb07acc988c762dd3eba83f0d76c92 Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Mon, 2 Nov 2026 15:59:42 +0200 Subject: [PATCH 11/12] changes caused by vllm version upgrade Signed-off-by: Shlomit Koyfman --- tests/unit/model_executor/models/test_qwen3_5_contract.py | 1 + 1 file changed, 1 insertion(+) 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 21a9d89f3..fbe3fccc1 100644 --- a/tests/unit/model_executor/models/test_qwen3_5_contract.py +++ b/tests/unit/model_executor/models/test_qwen3_5_contract.py @@ -295,6 +295,7 @@ def __init__(self): 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, From 9ca1dd2edd78ef4a00595a8191702b1ffd6ac32a Mon Sep 17 00:00:00 2001 From: Shlomit Koyfman Date: Wed, 11 Nov 2026 11:58:35 +0200 Subject: [PATCH 12/12] log eplb gain every log_balancedness_interval steps Signed-off-by: Shlomit Koyfman --- afd_plugin/connectors/gpu/deepep_m2n/eplb.py | 178 +++++++++- afd_plugin/distributed/eplb.py | 113 +++++++ docs/design/module/ffn_runtime.md | 52 ++- docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md | 53 ++- .../4a2f_graph_dp4tp1_eplb.sh | 313 ++++++++++++++++++ tests/unit/connectors/test_deepep_m2n_eplb.py | 197 ++++++++++- tests/unit/distributed/test_eplb.py | 95 ++++++ 7 files changed, 973 insertions(+), 28 deletions(-) create mode 100755 recipe/gpu/DeepEPM2NAFDConnector/qwen3_235b_a22b_fp8/prefill_decode_colocation/4a2f_graph_dp4tp1_eplb.sh diff --git a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py index 6b20b0dd6..c37240420 100644 --- a/afd_plugin/connectors/gpu/deepep_m2n/eplb.py +++ b/afd_plugin/connectors/gpu/deepep_m2n/eplb.py @@ -37,6 +37,18 @@ 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 @@ -55,6 +67,8 @@ from afd_plugin.distributed.eplb import ( describe_eplb_placement, describe_eplb_rank_load, + eplb_map_loss, + logical_load_from_window, ) if TYPE_CHECKING: @@ -73,6 +87,7 @@ "window_size", "step_interval", "initial_rearrangement_step", + "log_interval", ) @@ -167,6 +182,7 @@ def __init__( 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"]) @@ -234,21 +250,41 @@ def _step_finished_transaction(self, is_dummy: bool) -> None: step advances the rearrangement counter, exactly like vLLM's ``EplbState.step``, so both roles detect the boundary together. """ - if is_dummy: - self.expert_load_pass.zero_() - else: + 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) - self.expert_load_pass.zero_() 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.""" @@ -266,7 +302,12 @@ 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 - window_cpu = self.expert_load_window.cpu() + # 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) @@ -274,6 +315,13 @@ def _exchange_at_boundary(self) -> None: 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 @@ -300,7 +348,9 @@ class DeepEPM2NFFNEplb: 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. + 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__( @@ -314,9 +364,15 @@ def __init__( 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. - self._log_placement = bool( - connector.vllm_config.parallel_config.eplb_config.log_balancedness + # 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: @@ -333,6 +389,7 @@ def broadcast_initial_state(self) -> None: "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) @@ -342,13 +399,62 @@ def broadcast_initial_state(self) -> None: dist.broadcast(count_cpu, src=0, group=self._connector.bootstrap_group) def step(self) -> None: - """One connector step; runs the boundary exchange when due.""" + """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 @@ -358,9 +464,22 @@ def _rearrange(self) -> None: 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 self._log_placement and self._connector.world_rank == 0: + 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() @@ -385,6 +504,43 @@ def _log_chosen_experts(self, load_window: torch.Tensor) -> None: "\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", diff --git a/afd_plugin/distributed/eplb.py b/afd_plugin/distributed/eplb.py index 13e07e5a7..503321ea0 100644 --- a/afd_plugin/distributed/eplb.py +++ b/afd_plugin/distributed/eplb.py @@ -10,9 +10,26 @@ 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, *, @@ -55,6 +72,98 @@ def describe_eplb_placement( 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, *, @@ -76,6 +185,10 @@ def describe_eplb_rank_load( __all__ = [ + "EplbLoadSummary", "describe_eplb_placement", "describe_eplb_rank_load", + "eplb_map_loss", + "logical_load_from_window", + "summarize_rank_load", ] diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index 24438e1f4..5b8424cb0 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -332,7 +332,9 @@ Attention-recorded load window over the connector's bootstrap group, runs 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. +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 @@ -360,7 +362,7 @@ 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 -below); the formatters live in `afd_plugin/distributed/eplb.py`. +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 @@ -375,6 +377,52 @@ 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/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md index 3fb29d70e..0245ab2c1 100644 --- a/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/DEEPEP_M2N_CONNECTOR_USER_GUIDE.md @@ -320,13 +320,52 @@ 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 which experts EPLB chose, set `"log_balancedness": true` in the FFN's -`--eplb-config`: at every rearrangement the FFN then logs, 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 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). +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: 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/tests/unit/connectors/test_deepep_m2n_eplb.py b/tests/unit/connectors/test_deepep_m2n_eplb.py index 45ac4dd36..5a875dd25 100644 --- a/tests/unit/connectors/test_deepep_m2n_eplb.py +++ b/tests/unit/connectors/test_deepep_m2n_eplb.py @@ -28,7 +28,7 @@ from tests.unit.fakes.connector import FakeRoutedConnector # noqa: E402 MOE_LAYERS = (1, 3) -NUM_LOGICAL = 8 +NUM_LOGICAL = 12 NUM_PHYSICAL = 12 WINDOW_SIZE = 4 STEP_INTERVAL = 100 @@ -44,6 +44,7 @@ def _schedule(**overrides) -> dict[str, int]: "window_size": WINDOW_SIZE, "step_interval": STEP_INTERVAL, "initial_rearrangement_step": INITIAL_REARRANGE_STEP, + "log_interval": 0, } schedule.update(overrides) return schedule @@ -61,12 +62,12 @@ def _attention_connector(**overrides): return SimpleNamespace(**attrs) -def _recorder(**connector_overrides): +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(), + _schedule(log_interval=log_interval), MOE_LAYERS, initial_map, initial_count, @@ -195,6 +196,72 @@ def test_recorder_exchanges_window_and_adopts_maps_at_boundary( 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(): @@ -208,12 +275,14 @@ def _ffn_state(): # 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=torch.tensor([layer0, layer1]), - logical_to_physical_map=torch.zeros((len(MOE_LAYERS), NUM_LOGICAL, 4)), - logical_replica_count=torch.zeros((len(MOE_LAYERS), NUM_LOGICAL)), + physical_to_logical_map=physical_to_logical, + logical_to_physical_map=logical_to_physical, + logical_replica_count=replica_count, ) rearrangements: list[None] = [] @@ -229,14 +298,19 @@ def rearrange(): return state, model_state, rearrangements -def _ffn_connector(world_rank: int, *, log_placement: bool = False): +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), + eplb_config=SimpleNamespace( + log_balancedness=log_placement, + log_balancedness_interval=log_interval, + ), ), ), ) @@ -351,3 +425,110 @@ def test_ffn_driver_stays_silent_without_balancedness_logging( 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/distributed/test_eplb.py b/tests/unit/distributed/test_eplb.py index 9a7a53e50..2a24f691c 100644 --- a/tests/unit/distributed/test_eplb.py +++ b/tests/unit/distributed/test_eplb.py @@ -12,6 +12,9 @@ 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 @@ -50,3 +53,95 @@ def test_describe_eplb_rank_load_sums_each_ranks_slots(): 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)