From 857cacc3cae06a40aaa7efbd86c4e81508bcdba9 Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 10 Sep 2026 16:19:51 +0800 Subject: [PATCH 1/2] feat: run DeepSeek-V4 on the async GPU connector Wires the connector into both AFD roles and gives DeepSeek-V4-Flash -- 256 experts at topk 6, where the per-layer control-plane round trip costs more than the expert compute it guards -- a recipe that uses it. Both GPU runners previously asserted `control_plane is not None`, and the FFN worker loop raised NotImplementedError without one, so the connector could not run at all. The FFN side now pulls one work item at a time from the connector's own receive loop, taking the layer index and row counts from the arriving payload, and returns on an idle poll so the worker loop still sees its shutdown event. The Attention side skips vLLM's cross-DP batch agreement: async AFD lets each replica advance alone, so an idle replica never joins that all-reduce and a busy one would block in it forever -- which is exactly where a 2A2F run hung before reaching the first MoE layer. The V4 adapter learns the expert-routed dispatch protocol (the gate runs on the Attention side, so the wire carries topk ids and weights instead of token ids) and `compute_ffn_output` takes a device-side group_list so the FFN side runs only the grouped GEMM over its local experts. The V2 adapter gets the same GPU entry point, which is what the connector's e2e tests use as their reference. The rest was found bringing 2A2F up on real weights: the FFN role must force vLLM's NoDP MoE prepare/finalize under EP with DP>1, shared experts must be skipped on an empty shared slice, the SWIGLUOAI clamp has to reach the routed experts, the V4 Attention-side gate loads under its checkpoint path, and the DP coordinator's startup wait needs to be long enough for the second role's weights to load. Co-Authored-By: Claude Opus 5 Signed-off-by: specture724 --- afd_plugin/__init__.py | 2 + .../compat/patches/dp_coordinator_timeout.py | 60 ++++ .../compat/patches/ffn_local_moe_prepare.py | 105 +++++++ .../model_executor/models/deepseek_v2.py | 22 +- .../model_executor/models/deepseek_v4.py | 278 +++++++++++++++--- afd_plugin/v1/worker/attention_metadata.py | 18 +- .../v1/worker/attention_model_runner.py | 116 +++++--- afd_plugin/v1/worker/ffn_model_runner.py | 122 ++++++-- afd_plugin/v1/worker/ffn_worker.py | 24 +- docs/design/module/execution_platforms.md | 25 ++ .../deepseek_v4_flash/2a2f_async.sh | 151 ++++++++++ .../deepseek_v4_flash/2a2f_eager_async.sh | 125 ++++++++ .../models/test_deepseek_v4_weight_policy.py | 15 +- .../models/test_forward_context.py | 7 +- .../v1/worker/test_attention_model_runner.py | 57 +++- tests/unit/v1/worker/test_ffn_model_runner.py | 46 ++- 16 files changed, 1052 insertions(+), 121 deletions(-) create mode 100644 afd_plugin/compat/patches/dp_coordinator_timeout.py create mode 100644 afd_plugin/compat/patches/ffn_local_moe_prepare.py create mode 100755 recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh create mode 100755 recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index 6a7e8270..2766818a 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -174,7 +174,9 @@ def register_afd() -> None: import afd_plugin.compat.patches.async_dp_engine # noqa: F401 import afd_plugin.compat.patches.async_dp_forward_context # noqa: F401 import afd_plugin.compat.patches.config_validation # noqa: F401 + import afd_plugin.compat.patches.dp_coordinator_timeout # noqa: F401 import afd_plugin.compat.patches.engine_core # noqa: F401 + import afd_plugin.compat.patches.ffn_local_moe_prepare # noqa: F401 except Exception: _logger.debug( "AFD plugin: compatibility patches could not be applied", diff --git a/afd_plugin/compat/patches/dp_coordinator_timeout.py b/afd_plugin/compat/patches/dp_coordinator_timeout.py new file mode 100644 index 00000000..8a13e773 --- /dev/null +++ b/afd_plugin/compat/patches/dp_coordinator_timeout.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Raise the DP Coordinator's hardcoded startup ZMQ wait timeout. + +``DPCoordinator._wait_for_zmq_addrs`` waits a hardcoded 120 seconds for the +coordinator subprocess to import, bind, and report its ZMQ addresses. On a +CPU-oversubscribed shared box that subprocess can legitimately take longer, +which kills both AFD roles during startup (observed repeatedly on 2A2F +DeepSeek-V4-Flash runs). The wait becomes env-configurable with a 600 second +default; every other behavior matches upstream. +""" + +from __future__ import annotations + +import multiprocessing +import os + +import vllm.v1.engine.coordinator as coordinator_module + +DEFAULT_TIMEOUT_S = 600 + + +# Patch reason: the upstream DP Coordinator startup wait is hardcoded to 120 +# seconds, which is not enough for the coordinator subprocess to import and +# bind on a CPU-oversubscribed shared machine -- both AFD roles then die +# during startup. +# Patch functionality: identical to upstream, except the wait comes from +# AFD_DP_COORDINATOR_TIMEOUT_S (default 600 seconds). +# Signature: matches upstream; no added parameters. +# Upstream: vLLM v0.26.0, vllm/v1/engine/coordinator.py +def _wait_for_zmq_addrs(self, zmq_addr_pipe) -> tuple[str, str, str]: + try: + timeout = int( + os.getenv("AFD_DP_COORDINATOR_TIMEOUT_S", str(DEFAULT_TIMEOUT_S)), + ) + ready = multiprocessing.connection.wait( + [zmq_addr_pipe, self.proc.sentinel], timeout=timeout + ) + if not ready: + raise RuntimeError( + "DP Coordinator process failed to report ZMQ addresses " + f"within timeout={timeout} seconds during startup." + ) + try: + return zmq_addr_pipe.recv() + except EOFError: + raise RuntimeError( + "DP Coordinator process failed during startup." + ) from None + finally: + zmq_addr_pipe.close() + + +def apply_dp_coordinator_timeout() -> None: + coordinator_module.DPCoordinator._wait_for_zmq_addrs = _wait_for_zmq_addrs + + +apply_dp_coordinator_timeout() + +__all__ = ["apply_dp_coordinator_timeout"] diff --git a/afd_plugin/compat/patches/ffn_local_moe_prepare.py b/afd_plugin/compat/patches/ffn_local_moe_prepare.py new file mode 100644 index 00000000..2b23126e --- /dev/null +++ b/afd_plugin/compat/patches/ffn_local_moe_prepare.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Select the local (NoDP) MoE prepare/finalize for the AFD FFN role. + +An AFD FFN rank is not a vLLM DP rank: it has no scheduler, never forms a +coordinated DP batch, and only ever holds rows the Attention dispatch already +routed to its local experts. vLLM's DP MoE path instead re-assembles the whole +DP token set with an ``all_gatherv`` collective that reads +``dp_metadata`` from the forward context and needs every DP rank in lockstep. +Neither exists on the connector-driven FFN role, so the worker loop dies on +``assert dp_metadata is not None`` (2 FFN ranks would deadlock in the +collective even if the metadata were supplied). + +The NoDP prepare/finalize is pure-local: quantize, permute through the layer's +expert_map (our grouped rows carry global expert ids that map onto the local +range), run the experts, combine locally. That is exactly the AFD FFN +execution model. +""" + +from __future__ import annotations + +import sys +from types import ModuleType +from typing import Any + +import vllm.model_executor.layers.fused_moe.all2all_utils as all2all_utils_module +from vllm.config import get_current_vllm_config +from vllm.model_executor.layers.fused_moe.all2all_utils import ( + make_moe_prepare_and_finalize_no_dp_ep, +) + +from afd_plugin.config import parse_optional_afd_config + +# Patch reason: see module docstring -- the naive DP all-to-all cannot run on +# the connector-driven AFD FFN role. +# Patch functionality: when the active role is the AFD FFN role and no exotic +# all2all kernel is requested, return vLLM's NoDP prepare/finalize instead of +# the naive DP one; every non-AFD caller keeps the upstream selection. +# Signature: matches upstream; no added parameters. +# Upstream: vLLM v0.26.0, +# vllm/model_executor/layers/fused_moe/all2all_utils.py +_UPSTREAM_SELECTOR = all2all_utils_module.maybe_make_prepare_finalize + + +def _is_afd_ffn_role() -> bool: + try: + afd_config = parse_optional_afd_config( + get_current_vllm_config(), + validate=False, + ) + except Exception: + return False + return afd_config is not None and afd_config.role == "ffn" + + +def maybe_make_prepare_finalize(*args: Any, **kwargs: Any): + if not _is_afd_ffn_role(): + return _UPSTREAM_SELECTOR(*args, **kwargs) + moe = args[0] if args else kwargs["moe"] + parallel = moe.moe_parallel_config + # Kernels with their own dispatch own the collective; leave them untouched. + # Everything else (including the naive DP fallback that use_ep + dp>1 + # selects) must run locally: AFD pre-routes the rows. + exotic_kernels = ( + parallel.use_deepep_ht_kernels + or parallel.use_deepep_ll_kernels + or parallel.use_deepep_v2_kernels + or parallel.use_fi_nvl_two_sided_kernels + or parallel.use_fi_nvl_one_sided_kernels + or parallel.use_nixl_ep_kernels + or parallel.use_mori_kernels + ) + if exotic_kernels: + return _UPSTREAM_SELECTOR(*args, **kwargs) + return make_moe_prepare_and_finalize_no_dp_ep( + use_monolithic=bool(kwargs.get("use_monolithic", False)), + ) + + +def _rebind_source_module() -> None: + maybe_make_prepare_finalize._afd_installed = True # type: ignore[attr-defined] + all2all_utils_module.maybe_make_prepare_finalize = maybe_make_prepare_finalize + + +def apply_local_moe_prepare() -> None: + """Install the selector wrapper in every namespace that bound it. + + The plugin loads before vLLM's MoE modules are imported, so re-aliasing + the source module is what future ``from ... import`` bindings pick up; + any module already present in ``sys.modules`` that still holds the + upstream function is re-aliased directly. Idempotent. + """ + if getattr(maybe_make_prepare_finalize, "_afd_installed", False): + return + _rebind_source_module() + for module in list(sys.modules.values()): + if not isinstance(module, ModuleType) or module is all2all_utils_module: + continue + if getattr(module, "maybe_make_prepare_finalize", None) is _UPSTREAM_SELECTOR: + module.maybe_make_prepare_finalize = maybe_make_prepare_finalize + + +apply_local_moe_prepare() + +__all__ = ["apply_local_moe_prepare", "maybe_make_prepare_finalize"] diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 21e4f86b..235d177d 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -516,7 +516,8 @@ def compute_attn_output( topk_weights = None topk_ids = None router_logits = None - # NPU-only: Attention-side gate/topk is implemented in the NPU helper. + # The gate helper delegates expert selection to the connector, so both + # platforms share it despite the module's location. if self.compute_gate_on_attention and self.is_moe_layer: from afd_plugin.model_executor.models.npu import ( deepseek_v2_attention_gate, @@ -572,8 +573,23 @@ def compute_ffn_output( ) return output if self.compute_gate_on_attention: - raise RuntimeError( - "GPU Attention-side gate must call compute_experts_output", + if group_list is None: + # Without a group list the caller is the control-plane path, + # which routes on this side and must use compute_experts_output. + raise RuntimeError( + "GPU Attention-side gate must call compute_experts_output", + ) + # Token-level dispatch: rows arrive pre-routed and grouped by local + # expert, so only the grouped GEMM is left to run here. + from afd_plugin.model_executor.models.gpu import ( + deepseek_v2_attention_gate as gpu_attention_gate, + ) + + return gpu_attention_gate.compute_attention_gate_moe_ffn( + self, + hidden_states=hidden_states, + group_list=group_list, + expand_x_shared=expand_x_shared, ) hidden_states = self.mlp(hidden_states) if ( diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index cfd15430..3850b1cf 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -18,6 +18,10 @@ from vllm.models.deepseek_v4.nvidia import model as native from afd_plugin.config import parse_afd_config +from afd_plugin.connectors import ( + AFDExpertRoutingSpec, + AFDF2ATransferPayload, +) from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield @@ -27,8 +31,8 @@ _BOTH_ROLES = frozenset(("attention", "ffn")) -def _weight_layer_path(name: str) -> tuple[int, str] | None: - """Extract the decoder layer index and first layer-local path component.""" +def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: + """Extract the decoder layer index and the path below ``layers.N``.""" parts = name.split(".") for marker_idx, part in enumerate(parts[:-2]): if part != "layers": @@ -37,7 +41,11 @@ def _weight_layer_path(name: str) -> tuple[int, str] | None: layer_idx = int(parts[marker_idx + 1]) except ValueError: continue - return layer_idx, parts[marker_idx + 2] + return ( + layer_idx, + parts[marker_idx + 2], + tuple(parts[marker_idx + 3 :]), + ) return None @@ -55,8 +63,13 @@ def _checkpoint_weight_roles(name: str) -> frozenset[str]: layer_path = _weight_layer_path(name) if layer_path is None: return _BOTH_ROLES - _, stage = layer_path + _, stage, remainder = layer_path if stage == "ffn": + if remainder and remainder[0] == "gate": + # The gate computes on Attention, and the parameters live under + # .ffn.gate to match the checkpoint; the FFN role's native MoE + # gate loads the same tensors at the same path. + return _BOTH_ROLES return _FFN_ROLE return _ATTENTION_ROLE @@ -75,18 +88,68 @@ def _iter_role_weights( class RemoteDeepseekV4FFN(nn.Module): """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" - def __init__(self, *, layer_idx: int) -> None: + def __init__( + self, + *, + layer_idx: int, + vllm_config: VllmConfig | None = None, + prefix: str = "", + ) -> None: super().__init__() self.layer_idx = layer_idx + # The gate computes on Attention, but its parameters live under .ffn so + # the checkpoint names (…ffn.gate.*) load straight onto them; the FFN + # role loads the same tensors into its native MoE gate. Without a + # config there is nothing to size a gate from -- the control-plane path + # routes on the FFN side and never asks for one. + self.gate: native.GateLinear | None = None + if vllm_config is None: + return + if not parse_afd_config(vllm_config, validate=False).compute_gate_on_attention: + return + config = vllm_config.model_config.hf_config + self.gate = native.GateLinear( + input_size=config.hidden_size, + output_size=config.n_routed_experts, + bias=False, + out_dtype=torch.float32, + prefix=f"{prefix}.ffn.gate", + ) + self.gate.e_score_correction_bias = None + self.gate.tid2eid = None + if layer_idx < config.num_hash_layers: + self.gate.tid2eid = nn.Parameter( + torch.randint( + 0, + config.n_routed_experts, + (config.vocab_size, config.num_experts_per_tok), + dtype=torch.int32, + ), + requires_grad=False, + ) + elif getattr(config, "topk_method", None) == "noaux_tc": + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + requires_grad=False, + ) def forward( self, hidden_states: torch.Tensor, - input_ids: torch.Tensor | None, + input_ids: torch.Tensor | None = None, + topk_weights: torch.Tensor | None = None, + topk_ids: torch.Tensor | None = None, ) -> torch.Tensor: - if input_ids is None: - raise RuntimeError("DeepSeek-V4 remote FFN requires input_ids") - if input_ids.ndim != 1 or input_ids.shape[0] != hidden_states.shape[0]: + if input_ids is None and topk_ids is None: + raise RuntimeError( + "DeepSeek-V4 remote FFN requires input_ids or expert routing", + ) + if ( + input_ids is not None + and input_ids.ndim != 1 + or input_ids is not None + and input_ids.shape[0] != hidden_states.shape[0] + ): raise ValueError( "DeepSeek-V4 input_ids must be one-dimensional and token-aligned", ) @@ -105,11 +168,21 @@ def forward( seq_len=int(hidden_states.shape[0]), ) context = AFDTransferContext(metadata=metadata) - afd_metadata.connector.send_attn_output( - hidden_states, - context, - input_ids=input_ids, - ) + if topk_ids is not None: + # Expert-routed dispatch (async connector): the gate ran here, so + # the wire carries the routing instead of the token ids. + afd_metadata.connector.send_attn_output( + hidden_states, + context, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + else: + afd_metadata.connector.send_attn_output( + hidden_states, + context, + input_ids=input_ids, + ) hidden_states = maybe_apply_dbo_yield( hidden_states, role="attention", @@ -153,7 +226,21 @@ def __init__( topk_indices_buffer=topk_indices_buffer, aux_stream_list=aux_stream_list, ) - self.ffn = RemoteDeepseekV4FFN(layer_idx=layer_idx) + self.ffn = RemoteDeepseekV4FFN( + layer_idx=layer_idx, + vllm_config=vllm_config, + prefix=prefix, + ) + # ### PATCH START: gate runs on Attention for the async connector. + if afd_config.compute_gate_on_attention: + self.n_activated_experts = config.num_experts_per_tok + self.routed_scaling_factor = getattr( + config, "routed_scaling_factor", 1.0 + ) + self.renormalize = config.norm_topk_prob + self.scoring_func = getattr(config, "scoring_func", "sqrtsoftplus") + self.hash_indices_dtype = torch.int32 + # ### PATCH END elif afd_config.role == "ffn": self.attn = native.PPMissingLayer() self.ffn = native.DeepseekV4MoE(vllm_config, prefix=f"{prefix}.ffn") @@ -199,6 +286,74 @@ def __init__( requires_grad=False, ) + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + *, + input_ids: torch.Tensor | None = None, + group_list: torch.Tensor | None = None, + expand_x_shared: torch.Tensor | None = None, + ) -> torch.Tensor | AFDF2ATransferPayload: + if not isinstance(self.ffn, native.DeepseekV4MoE): + raise RuntimeError("DeepSeek-V4 FFN compute is FFN-role only") + if group_list is None: + # P2pNccl path: the whole MoE, including its native router. + if input_ids is None: + raise RuntimeError("DeepSeek-V4 FFN compute requires input_ids") + return self.ffn(hidden_states, input_ids) + moe = self.ffn + counts = group_list.to(torch.int64) + num_rows = int(hidden_states.shape[0]) + num_local_experts = int(counts.numel()) + if num_rows == 0: + return AFDF2ATransferPayload( + routed_output=hidden_states.new_empty((0, self.hidden_size)), + shared_output=None, + ) + # Rows arrive sorted by local expert; rebuild global expert ids so the + # FusedMoE layer's own mapping sees the owning expert. + expert_ids = torch.repeat_interleave( + torch.arange( + num_local_experts, + device=hidden_states.device, + dtype=torch.int32, + ) + + moe.experts_start_idx, + counts, + output_size=num_rows, + ).unsqueeze(1) + # V4's routed scaling was already folded into the dispatch-side topk + # weights, so each partial row carries weight 1.0 here; the connector + # applies the per-partial weights when it reduces back to one row per + # token on the Attention side. + row_weights = torch.ones( + (num_rows, 1), + dtype=torch.float32, + device=hidden_states.device, + ) + # The SWIGLUOAI clamp rides in FusedMoEConfig (built from + # swiglu_limit at construction), so the grouped call must not pass a + # runtime clamp here: the modular runner rejects the kwarg. + routed_output = moe.experts( + hidden_states, + row_weights, + expert_ids, + ) + shared_output = None + # A dispatch's round-robin shared slice is empty for a rank whenever + # the ubatch holds fewer tokens than FFN ranks, so zero rows are a + # normal item shape, not a missing payload. + if ( + moe.shared_experts is not None + and expand_x_shared is not None + and expand_x_shared.shape[0] > 0 + ): + shared_output = moe.shared_experts(expand_x_shared) + return AFDF2ATransferPayload( + routed_output=routed_output, + shared_output=shared_output, + ) + # Patch reason: native forward directly invokes its locally allocated FFN. # Patch functionality: preserve native mHC state locally while the proxy # transfers only the two-dimensional FFN activation and input IDs. @@ -294,24 +449,36 @@ def forward( norm_eps=ffn_norm_eps, ) - # ### PATCH START: this call enters the synchronous remote FFN proxy. - x = self.ffn(x, input_ids) + # ### PATCH START: enter the remote FFN proxy (sync or expert-routed). + gate = self.ffn.gate + if gate is not None: + router_logits, _ = gate(x) + topk_weights, topk_ids = native.fused_topk_bias( + hidden_states=x, + gating_output=router_logits, + scoring_func=self.scoring_func, + e_score_correction_bias=( + gate.e_score_correction_bias.data + if gate.e_score_correction_bias is not None + else None + ), + topk=self.n_activated_experts, + renormalize=self.renormalize, + indices_type=self.hash_indices_dtype, + input_tokens=input_ids, + hash_indices_table=gate.tid2eid, + routed_scaling_factor=self.routed_scaling_factor, + ) + x = self.ffn( + x, + topk_weights=topk_weights, + topk_ids=topk_ids, + ) + else: + x = self.ffn(x, input_ids) # ### PATCH END return x, residual, post_mix, res_mix - def compute_ffn_output( - self, - hidden_states: torch.Tensor, - *, - input_ids: torch.Tensor | None, - ) -> torch.Tensor: - """Execute the complete native V4 MoE, including its native router.""" - if not isinstance(self.ffn, native.DeepseekV4MoE): - raise RuntimeError("DeepSeek-V4 FFN compute is FFN-role only") - if input_ids is None: - raise RuntimeError("DeepSeek-V4 FFN compute requires input_ids") - return self.ffn(hidden_states, input_ids) - class AFDDeepseekV4Model(native.DeepseekV4Model): """Role-aware DeepSeek-V4 model retaining mHC exclusively on Attention.""" @@ -327,13 +494,28 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.afd_config = parse_afd_config(vllm_config, validate=False) if native.current_platform.device_type != "cuda": raise RuntimeError("AFD DeepSeek-V4 supports CUDA only") - if self.afd_config.connector != "P2pNcclAFDConnector": + if self.afd_config.connector not in ( + "P2pNcclAFDConnector", + "GpuAsyncAFDConnector", + ): + raise RuntimeError( + "AFD DeepSeek-V4 supports P2pNcclAFDConnector or GpuAsyncAFDConnector", + ) + if ( + self.afd_config.connector == "GpuAsyncAFDConnector" + and not self.afd_config.compute_gate_on_attention + ): raise RuntimeError( - "AFD DeepSeek-V4 requires the synchronous P2pNcclAFDConnector", + "AFD DeepSeek-V4 async dispatch routes by expert: " + "compute_gate_on_attention is required", ) - if self.afd_config.compute_gate_on_attention: + if ( + self.afd_config.connector == "P2pNcclAFDConnector" + and self.afd_config.compute_gate_on_attention + ): raise RuntimeError( - "AFD DeepSeek-V4 does not support compute_gate_on_attention", + "AFD DeepSeek-V4 over P2pNccl routes on the FFN side: " + "compute_gate_on_attention must stay off", ) parallel_config = vllm_config.parallel_config if parallel_config.pipeline_parallel_size != 1: @@ -436,16 +618,28 @@ def compute_ffn_output( hidden_states: torch.Tensor, layer_idx: int, *, - input_ids: torch.Tensor | None, - ) -> torch.Tensor: + input_ids: torch.Tensor | None = None, + group_list: torch.Tensor | None = None, + expand_x_shared: torch.Tensor | None = None, + ) -> torch.Tensor | AFDF2ATransferPayload: return self.layers[layer_idx].compute_ffn_output( hidden_states, input_ids=input_ids, + group_list=group_list, + expand_x_shared=expand_x_shared, ) def get_experts_layer_indices(self) -> tuple[int, ...]: return tuple(range(int(self.config.num_hidden_layers))) + def get_experts_routing_spec(self, layer_idx: int) -> AFDExpertRoutingSpec: + """Router contract for the async FFN loop's receive buffers.""" + gate = self.layers[layer_idx].gate + return AFDExpertRoutingSpec( + router_logits_width=int(self.config.n_routed_experts), + router_logits_dtype=gate.out_dtype or gate.weight.dtype, + ) + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: """Return native expert mappings only where real experts are owned.""" if self.afd_config.role == "attention": @@ -482,6 +676,9 @@ class AFDDeepseekV4ForCausalLM(native.DeepseekV4ForCausalLM): def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.afd_config = parse_afd_config(vllm_config, validate=False) self.afd_role = self.afd_config.role + # Only the P2pNccl wire carries token ids; the async wire carries the + # expert routing the Attention-side gate computed. + self.afd_requires_input_ids = self.afd_config.connector == "P2pNcclAFDConnector" super().__init__(vllm_config=vllm_config, prefix=prefix) def compute_ffn_output( @@ -490,14 +687,21 @@ def compute_ffn_output( layer_idx: int, *, input_ids: torch.Tensor | None = None, + group_list: torch.Tensor | None = None, + expand_x_shared: torch.Tensor | None = None, **kwargs: Any, - ) -> torch.Tensor: + ) -> torch.Tensor | AFDF2ATransferPayload: return self.model.compute_ffn_output( hidden_states, layer_idx, input_ids=input_ids, + group_list=group_list, + expand_x_shared=expand_x_shared, ) + def get_experts_routing_spec(self, layer_idx: int) -> AFDExpertRoutingSpec: + return self.model.get_experts_routing_spec(layer_idx) + def get_experts_layer_indices(self) -> tuple[int, ...]: return self.model.get_experts_layer_indices() diff --git a/afd_plugin/v1/worker/attention_metadata.py b/afd_plugin/v1/worker/attention_metadata.py index f633dc59..756ebb3e 100644 --- a/afd_plugin/v1/worker/attention_metadata.py +++ b/afd_plugin/v1/worker/attention_metadata.py @@ -15,6 +15,7 @@ AFDDPMetadata, AFDForwardContextMetadata, ) +from afd_plugin.connectors.base import AFDConnectorBase class AFDMetadataProviderMixin: @@ -29,6 +30,13 @@ class AFDMetadataProviderMixin: _afd_is_profile: bool = False + #: Provided by the consuming runner; declared so the mixin type-checks. + connector: AFDConnectorBase + vllm_config: VllmConfig + _is_warmup: bool + _afd_pending_metadata: AFDForwardContextMetadata | None + _afd_transaction_counter: int + def build_afd_metadata( self, ubatch_slices: UBatchSlices | None, @@ -81,9 +89,13 @@ def send_dp_metadata( payload; it does not allocate a transaction or perform data-plane work. """ - assert self.connector.control_plane is not None, ( - "send_dp_metadata needs control plane driven connectors" - ) + if self.connector.control_plane is None: + # Control-plane-less connectors (GpuAsyncAFDConnector) carry all + # per-stage control on the data plane and drive the FFN from its + # own receive loop; each Attention replica also advances on its + # own (see _dp_batch_coordination_disabled), so there is nothing + # to publish here. + return if ubatch_slices and len(ubatch_slices) > 1: dp_metadata_list = { diff --git a/afd_plugin/v1/worker/attention_model_runner.py b/afd_plugin/v1/worker/attention_model_runner.py index b559be68..2dd1f32f 100644 --- a/afd_plugin/v1/worker/attention_model_runner.py +++ b/afd_plugin/v1/worker/attention_model_runner.py @@ -5,7 +5,7 @@ from __future__ import annotations from contextlib import AbstractContextManager, contextmanager, nullcontext -from typing import Any +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -49,12 +49,49 @@ from afd_plugin.v1.worker.cuda_graph import validate_cuda_graph_mode from afd_plugin.v1.worker.ubatch_wrapper import AFDUBatchWrapper +if TYPE_CHECKING: + from vllm.v1.core.sched.output import SchedulerOutput + + +@contextmanager +def _dp_batch_coordination_disabled(disabled: bool): + """Skip vLLM's cross-DP batch agreement for connector-driven runs. + + ``GPUModelRunner._determine_batch_execution_and_padding`` all-reduces the + batch shape across the DP group whenever ``data_parallel_size > 1``. Async + AFD deliberately lets each Attention replica advance on its own, so an idle + replica never joins that collective and a busy one blocks in it forever -- + which is where a 2A2F run hangs before it reaches the first MoE layer. + + Returning the single-rank answer (``num_tokens_across_dp=None``) makes the + upstream function skip its DP-padding branch entirely, exactly as it does + for ``data_parallel_size == 1``. + """ + if not disabled: + yield + return + + original = gpu_model_runner.coordinate_batch_across_dp + + def _single_rank_coordination(*_args: Any, cudagraph_mode: int, **_kwargs: Any): + return False, None, cudagraph_mode + + gpu_model_runner.coordinate_batch_across_dp = _single_rank_coordination + try: + yield + finally: + gpu_model_runner.coordinate_batch_across_dp = original + class AFDAttentionModelRunner(AFDMetadataProviderMixin, GPUModelRunner): """Attention model runner that injects AFD metadata into forward context.""" afd_expected_role = "attention" + #: Declared, not assigned: the ubatch-wrapper install both reads and + #: rebinds it, which leaves its type unresolvable from the base class. + model: Any + def __init__( self, vllm_config: VllmConfig, @@ -75,11 +112,9 @@ def __init__( self.afd_config, ) # The connector rendezvous is deferred to the end of ``load_model()`` - # so Attention and FFN weight loading overlap; see that method. - # TODO: Async GPU connector will be supported in the future - assert self.connector.control_plane is not None, ( - "GPU model runner only supports control-plane-driven connectors" - ) + # so Attention and FFN weight loading overlap; see that method. The + # async GPU connector drives FFN work from its own receive loop and so + # has no control plane, which is why there is no assertion here. self._is_warmup = False self._afd_is_graph_capturing = False self._afd_is_graph_replaying = False @@ -114,22 +149,18 @@ def load_model(self, load_dummy_weights: bool = False) -> None: self.connector.init_afd_connector() def _install_afd_ubatch_wrapper(self) -> None: - if isinstance(self.model, AFDUBatchWrapper): - self.model.configure_afd_context_provider( - self.install_afd_metadata_on_forward_context, + model: Any = self.model + if not isinstance(model, AFDUBatchWrapper): + if isinstance(model, UBatchWrapper): + model = model.unwrap() + model = AFDUBatchWrapper( + model, + self.vllm_config, + CUDAGraphMode.NONE, + self.device, ) - return - - model = self.model - if isinstance(model, UBatchWrapper): - model = model.unwrap() - self.model = AFDUBatchWrapper( - model, - self.vllm_config, - CUDAGraphMode.NONE, - self.device, - ) - self.model.configure_afd_context_provider( + self.model = model + model.configure_afd_context_provider( self.install_afd_metadata_on_forward_context, ) @@ -215,25 +246,28 @@ def _determine_batch_execution_and_padding( torch.Tensor | None, CUDAGraphStat | None, ]: - ( - cudagraph_mode, - batch_descriptor, - should_ubatch, - num_tokens_across_dp, - cudagraph_stats, - ) = super()._determine_batch_execution_and_padding( - num_tokens, - num_reqs, - num_scheduled_tokens_np, - max_num_scheduled_tokens, - use_cascade_attn, - allow_microbatching, - force_eager, - force_uniform_decode, - force_has_lora, - force_num_active_loras, - num_encoder_reqs, - ) + with _dp_batch_coordination_disabled( + self.connector.control_plane is None, + ): + ( + cudagraph_mode, + batch_descriptor, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = super()._determine_batch_execution_and_padding( + num_tokens, + num_reqs, + num_scheduled_tokens_np, + max_num_scheduled_tokens, + use_cascade_attn, + allow_microbatching, + force_eager, + force_uniform_decode, + force_has_lora, + force_num_active_loras, + num_encoder_reqs, + ) self._afd_is_graph_replaying = ( not bool(getattr(self, "_is_warmup", False)) and not bool(getattr(self, "_afd_is_graph_capturing", False)) @@ -255,7 +289,7 @@ def _determine_batch_execution_and_padding( ) kwargs: dict[str, Any] = {} - # determin if ubatch should be activated. + # determine if ubatch should be activated. # 1. For dp = 1, vLLM hardcodes `should_ubatch=False`. # This is the extra support for dp = 1 if self.vllm_config.parallel_config.data_parallel_size == 1: diff --git a/afd_plugin/v1/worker/ffn_model_runner.py b/afd_plugin/v1/worker/ffn_model_runner.py index 4582d229..08f86ab0 100644 --- a/afd_plugin/v1/worker/ffn_model_runner.py +++ b/afd_plugin/v1/worker/ffn_model_runner.py @@ -30,6 +30,11 @@ AFDControlPayload, AFDDPMetadata, ) +from afd_plugin.connectors.gpu.async_gpu import ( + ConnectorShutdown, + GpuAsyncTransferState, +) +from afd_plugin.connectors.metadata import AFDF2ATransferPayload from afd_plugin.v1.worker.attention_model_runner import ( fail_if_unsupported_ubatching, ) @@ -53,10 +58,12 @@ class GPUFFNModelRunner(LoRAModelRunnerMixin): """FFN model runner for AFD GPU execution. - FFN steps are driven by the connector control plane rather than the vLLM - scheduler. GPU only supports control-plane-driven connectors, so the runner - asserts ``connector.control_plane is not None`` at construction; connectors - without a control plane (``control_plane is None``) are not supported. + FFN steps are driven by the connector rather than the vLLM scheduler, in one + of two ways. Control-plane connectors receive broadcast DP metadata and then + walk every layer in lockstep with the Attention side. Connectors without a + control plane (``control_plane is None``) instead pull one work item at a + time from their receive loop, learning the layer and token counts from the + arriving payload; see ``execute_connector_driven_step``. """ afd_expected_role = "ffn" @@ -69,10 +76,6 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: self.dtype = self.model_config.dtype self.afd_config = self.parse_config(vllm_config) fail_if_unsupported_ubatching(vllm_config) - self.afd_cudagraph_policy = validate_cuda_graph_mode( - vllm_config, - role="ffn", - ) rank, local_rank = _resolve_world_ranks() self.connector = AFDConnectorFactory.create_connector( rank, @@ -80,21 +83,40 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None: vllm_config, self.afd_config, ) - # TODO: Async GPU connector will be supported in the future - assert self.connector.control_plane is not None, ( - "GPU model runner only supports control-plane-driven connectors" - ) + # A connector without a control plane drives FFN steps from its own + # receive loop instead of from broadcast DP metadata. + self.is_connector_driven = self.connector.control_plane is None + # The connector-driven path never touches vLLM's graph machinery, so + # vLLM's cudagraph_mode says nothing about it -- running the policy gate + # here would reject modes that are simply irrelevant. + if self.is_connector_driven: + self.afd_cudagraph_policy = None + else: + self.afd_cudagraph_policy = validate_cuda_graph_mode( + vllm_config, + role="ffn", + ) - self.model: Any | None = None + self.model: Any = None self.model_memory_usage = 0 self.num_layers = int(self.model_config.hf_text_config.num_hidden_layers) self.use_cuda_graph = bool( - self.afd_cudagraph_policy.enable_ffn_graph_cache, + self.afd_cudagraph_policy is not None + and self.afd_cudagraph_policy.enable_ffn_graph_cache ) self._cuda_graphs: dict[tuple, dict[str, Any]] = {} self._graph_memory_pool: Any | None = None self.prof = create_afd_gpu_profiler("ffn") + @property + def _control_plane(self) -> Any: + """The control plane, on the paths that only run when there is one.""" + control_plane = self.connector.control_plane + assert control_plane is not None, ( + "control-plane FFN path reached on a connector-driven runner", + ) + return control_plane + @staticmethod def parse_config(vllm_config: VllmConfig) -> AFDConfig: return parse_afd_config(vllm_config, expected_role="ffn") @@ -154,6 +176,7 @@ def execute_model( graph_exists=cuda_graph_info is not None, ) if run_mode is AFDGraphRunMode.REPLAY: + assert cuda_graph_info is not None cuda_graph_info["graph"].replay() return None @@ -173,7 +196,7 @@ def _ffn_forward( update_connector_state: bool = True, ) -> torch.Tensor | None: if update_connector_state: - self.connector.control_plane.update_state_from_dp_metadata( + self._control_plane.update_state_from_dp_metadata( _make_dp_metadata_payload( dp_metadata_list, is_graph_capturing=is_graph_capturing, @@ -246,6 +269,69 @@ def _ffn_forward( self.connector.send_ffn_output(rank_ffn_output, context) return rank_ffn_output + def execute_connector_driven_step(self) -> None: + """Drain whatever the connector has already received, then return. + + Returning on an idle poll rather than blocking forever is what lets the + worker loop observe its shutdown event. The batch size below is only a + drain granularity: successive work items may belong to different layers + of different Attention replicas. + """ + step_afd_gpu_profiler(self.prof) + self._ffn_forward_connector_driven() + + def _compute_work_item( + self, + work_item: Any, + states: GpuAsyncTransferState, + ) -> torch.Tensor | AFDF2ATransferPayload: + """Run this work item's layer over the rows that arrived.""" + return self.model.compute_ffn_output( + hidden_states=work_item.hidden_states, + layer_idx=work_item.layer_idx, + group_list=states.group_list, + expand_x_shared=states.expand_x_shared, + ) + + def _ffn_forward_connector_driven( + self, + ) -> torch.Tensor | AFDF2ATransferPayload | None: + stage_idx = 0 + rank_ffn_output = None + connector = self.connector + max_items = max(1, int(self.num_layers)) + + with _ffn_forward_context(self.vllm_config) as forward_context: + for _ in range(max_items): + try: + work_item = connector.recv_ffn_work_item( # type: ignore[attr-defined] + stage_idx=stage_idx, + max_num_tokens=self.vllm_config.scheduler_config.max_num_batched_tokens, + ) + except TimeoutError: + # Nothing pending; hand control back so the worker loop can + # check for shutdown. + return rank_ffn_output + except ConnectorShutdown: + raise + + states = work_item.context.states + if not isinstance(states, GpuAsyncTransferState): + raise RuntimeError( + "async GPU FFN work item requires GpuAsyncTransferState", + ) + metadata = work_item.context.metadata + forward_context.dp_metadata = None + forward_context.additional_kwargs["afd_metadata"] = metadata + _set_moe_layer_index(forward_context, work_item.layer_idx) + + rank_ffn_output = self._compute_work_item(work_item, states) + rank_ffn_output = connector.send_ffn_work_item_output( # type: ignore[attr-defined] + work_item, + rank_ffn_output, + ) + return rank_ffn_output + def _execute_eager_mode( self, hidden_states: torch.Tensor, @@ -307,7 +393,7 @@ def _dummy_run( cudagraph = torch.cuda.CUDAGraph() # DP metadata receive/update is a control-plane side effect and must # complete before CUDA graph capture starts. - self.connector.control_plane.update_state_from_dp_metadata( + self._control_plane.update_state_from_dp_metadata( _make_dp_metadata_payload( dp_metadata_list, is_graph_capturing=is_attn_graph_capturing, @@ -349,7 +435,7 @@ def capture_model( try: with graph_capture(device=self.device): if is_warmup: - self.connector.control_plane.update_state_from_dp_metadata( + self._control_plane.update_state_from_dp_metadata( _make_dp_metadata_payload( dp_metadata_list, is_graph_capturing=False, @@ -436,7 +522,7 @@ def _ffn_forward_context(vllm_config: VllmConfig): yield get_forward_context() -def _set_moe_layer_index(forward_context: object, layer_idx: int) -> None: +def _set_moe_layer_index(forward_context: Any, layer_idx: int) -> None: all_moe_layers = forward_context.all_moe_layers if not all_moe_layers: return diff --git a/afd_plugin/v1/worker/ffn_worker.py b/afd_plugin/v1/worker/ffn_worker.py index 4800eab9..66784a4b 100644 --- a/afd_plugin/v1/worker/ffn_worker.py +++ b/afd_plugin/v1/worker/ffn_worker.py @@ -14,6 +14,7 @@ from vllm.v1.worker.gpu_worker import Worker from vllm.v1.worker.worker_base import CompilationTimes +from afd_plugin.connectors.gpu.async_gpu import ConnectorShutdown from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.attention_model_runner import fail_if_unsupported_ubatching from afd_plugin.v1.worker.ffn_model_runner import GPUFFNModelRunner @@ -160,6 +161,13 @@ def ffn_worker_loop() -> None: try: self._run_ffn_server_loop() except Exception as exc: + shutdown_event = self._ffn_shutdown_event + if shutdown_event is not None and shutdown_event.is_set(): + logger.debug( + "AFD FFN receive loop stopped during shutdown", + exc_info=True, + ) + return self._ffn_loop_error = exc logger.exception("AFD FFN worker loop failed") @@ -180,11 +188,17 @@ def _run_ffn_server_loop(self) -> None: while not event.is_set(): if self.model_runner.connector.control_plane is None: - raise NotImplementedError( - "GPU FFN only supports control-plane-driven connectors; " - "connectors without a control plane (control_plane is None) " - "are not supported.", - ) + # Connector-driven: the step returns on an idle poll, so the + # loop gets to re-check the shutdown event. No device-wide + # synchronize here -- it would serialize every receive against + # the previous compute and erase the overlap this path exists + # for; ordering is carried by the connector's own streams. + try: + self.model_runner.execute_connector_driven_step() + except ConnectorShutdown: + logger.info("AFD FFN loop exiting: peer announced shutdown") + return + continue payload = self.model_runner.connector.control_plane.recv_dp_metadata_list() dp_metadata_list = payload.dp_metadata_list diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md index 605c6a9d..38a069b1 100644 --- a/docs/design/module/execution_platforms.md +++ b/docs/design/module/execution_platforms.md @@ -169,6 +169,31 @@ connector state before `torch.cuda.graph(...)`, captures only model/data-plane work, and stores the graph by `make_ffn_graph_key()`. A matching future payload replays the graph; a missing key runs eagerly. +### CUDA Graphs and the async GPU connector + +`GpuAsyncAFDConnector` has no control plane, so the FFN worker takes the +connector-driven branch and never reaches the graph cache above: that side +stays eager. The Attention side captures, and its dispatch sits inside the +captured region, which constrains the window's flag protocol in two ways -- +a replay runs no Python, and a captured stream wait compares against the value +recorded at capture time. + +- The dispatch sequence number lives in a device tensor the graph increments, + not in a host counter, so each replay still stamps a peer's flag with a + number it has not seen; an FFN rank recognizes an arrival exactly by that + change. +- A reply stamps the constant `FLAG_REPLY_READY`, and the Attention side calls + `SymmWindow.clear_flag()` inside the graph once it has consumed the slot. + A rising reply sequence cannot work here: the wait would be frozen at one + value and fall straight through on every later replay. +- The header prefix a dispatch copies from the host is therefore fixed per + `(layer, stage, token count)` and cached in pinned memory, since a graph + records the source address. + +A role either polls its flags or stream-waits and clears them, never both: +`poll()` recognizes an arrival by the flag differing from what it last saw, +which a reset would defeat. + ### CUDA native ubatching `AFDUBatchWrapper` replaces vLLM's GPU wrapper during Attention model load diff --git a/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh new file mode 100755 index 00000000..a426b2d1 --- /dev/null +++ b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh @@ -0,0 +1,151 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +# 2A2F DeepSeek-V4-Flash on the async GPU connector. +# +# Launch under a GPU reservation, which sets CUDA_VISIBLE_DEVICES: +# gpu run --gpus 4 -- \ +# bash recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh +# +# The two roles are separate vllm serve processes: the AFD process group hosts +# its own TCPStore, which cannot be created under a single torchrun/torchelastic +# launcher. +set -u + +MODEL_PATH=${MODEL_PATH:-/path/model_weights/deepseek-v4-flash} +# How to invoke vLLM. `uv run vllm` is right from a synced checkout; override +# to point at an interpreter that actually has the plugin installed. +read -r -a VLLM_CMD <<< "${VLLM_CMD:-uv run vllm}" +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +# Single node over NVLink: skip the IB transport probe. +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +# Two servers on one box spawn a lot of threads; the HF tokenizer's rayon pool +# is the first thing to fail when thread creation gets refused. +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +# Split the reserved devices in half: first two Attention, last two FFN. +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +if [ "${#DEVICES[@]}" -lt 4 ]; then + echo "need 4 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]},${DEVICES[1]}" +FFN_DEVICES="${DEVICES[2]},${DEVICES[3]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +# Lower this when sharing a box: vLLM refuses to start if the desired +# fraction exceeds what is actually free. +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +# Prefill batch size drives whether each MoE call clears the compute-bound +# inflection point, so it is the knob to raise when benchmarking. +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-2048} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-16} +# The checkpoint declares a 1M-token context; sizing KV against that leaves +# nothing for the experts. +MAX_MODEL_LEN=${MAX_MODEL_LEN:-16384} +# V4 ships its own tokenizer, and its native attention wants an fp8 KV cache. +TOKENIZER_MODE=${TOKENIZER_MODE:-deepseek_v4} +KV_CACHE_DTYPE=${KV_CACHE_DTYPE:-fp8} +# Free-form passthrough, e.g. EXTRA_ARGS="--no-enable-prefix-caching". +read -r -a EXTRA_ARGS <<< "${EXTRA_ARGS:-}" +AFD_PORT=${AFD_PORT:-6271} +API_PORT=${API_PORT:-18307} +# The FFN server never takes HTTP -- its EngineCore is a connector daemon -- +# but it still starts an API server, and both roles racing for one port means +# whichever loses exits and takes its role down with it. Give it its own. +FFN_API_PORT=${FFN_API_PORT:-$((API_PORT + 1))} + +AFD_CONFIG_ATTN='{ + "afd": { + "role": "attention", + "connector": "GpuAsyncAFDConnector", + "async": true, + "compute_gate_on_attention": true, + "host": "127.0.0.1", + "port": '"$AFD_PORT"', + "num_attention_ranks": 2, + "num_ffn_ranks": 2 + } +}' +AFD_CONFIG_FFN=${AFD_CONFIG_ATTN/\"role\": \"attention\"/\"role\": \"ffn\"} + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" "${VLLM_CMD[@]}" serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config "$AFD_CONFIG_ATTN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --max-model-len "$MAX_MODEL_LEN" \ + --tokenizer-mode "$TOKENIZER_MODE" \ + --kv-cache-dtype "$KV_CACHE_DTYPE" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + "${EXTRA_ARGS[@]}" \ + --host 127.0.0.1 \ + --port "$API_PORT" \ + --trust-remote-code > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" "${VLLM_CMD[@]}" serve "$MODEL_PATH" \ + --data-parallel-size 2 \ + --tensor-parallel-size 1 \ + --enable-expert-parallel \ + --additional-config "$AFD_CONFIG_FFN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" \ + --max-model-len "$MAX_MODEL_LEN" \ + --tokenizer-mode "$TOKENIZER_MODE" \ + --kv-cache-dtype "$KV_CACHE_DTYPE" \ + --api-server-count 1 \ + --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --enforce-eager \ + "${EXTRA_ARGS[@]}" \ + --host 127.0.0.1 \ + --port "$FFN_API_PORT" \ + --trust-remote-code > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +# shellcheck disable=SC2317,SC2329 # invoked by the EXIT trap below. +# (SC2329 on shellcheck >= 0.11, SC2317 on older ones; CI runs an older one.) +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-600}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + echo + echo "curl -s http://127.0.0.1:$API_PORT/v1/completions \\" + echo " -H 'Content-Type: application/json' \\" + echo " -d '{\"model\":\"$MODEL_PATH\",\"prompt\":\"The capital of France is\",\"max_tokens\":16,\"temperature\":0}'" + echo + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d '{"model":"'"$MODEL_PATH"'","prompt":"The capital of France is", + "max_tokens":16,"temperature":0, + "skip_special_tokens":false,"logprobs":2}' + echo + exit 0 + fi + # Stay up so the servers can take requests; Ctrl-C tears both down. + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh new file mode 100755 index 00000000..7fdfc421 --- /dev/null +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh @@ -0,0 +1,125 @@ +#!/usr/bin/env bash +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +# 2A2F DeepSeek-V4-Flash with the synchronous P2pNcclAFDConnector, eager. +# +# The AFD DeepSeek-V4 adapter's first release requires this connector and +# rejects compute_gate_on_attention (V4's native router runs on the FFN side; +# token-aligned input_ids cross the Attention-to-FFN boundary instead). +# +# Launch under a GPU reservation: +# gpu run --gpus 4 -- bash recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh +set -u + +MODEL_PATH=${MODEL_PATH:-/data/boao/deepseek-v4-flash} +VLLM_CMD=${VLLM_CMD:-vllm} +LOG_DIR=${LOG_DIR:-.} +mkdir -p "$LOG_DIR" +export VLLM_USE_V2_MODEL_RUNNER=0 +export NVSHMEM_REMOTE_TRANSPORT=${NVSHMEM_REMOTE_TRANSPORT:-none} +export TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM:-false} +export RAYON_NUM_THREADS=${RAYON_NUM_THREADS:-2} +export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4} + +IFS=',' read -r -a DEVICES <<< "${CUDA_VISIBLE_DEVICES:-0,1,2,3}" +if [ "${#DEVICES[@]}" -lt 4 ]; then + echo "need 4 visible GPUs, got ${#DEVICES[@]}: ${CUDA_VISIBLE_DEVICES:-unset}" >&2 + exit 1 +fi +ATTN_DEVICES="${DEVICES[0]},${DEVICES[1]}" +FFN_DEVICES="${DEVICES[2]},${DEVICES[3]}" +echo "attention on ${ATTN_DEVICES}, ffn on ${FFN_DEVICES}" + +GPU_MEM_UTIL=${GPU_MEM_UTIL:-0.9} +MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-2048} +MAX_NUM_SEQS=${MAX_NUM_SEQS:-16} +MAX_MODEL_LEN=${MAX_MODEL_LEN:-16384} +MAX_MODEL_LEN_ARG=() +[ -n "$MAX_MODEL_LEN" ] && MAX_MODEL_LEN_ARG=(--max-model-len "$MAX_MODEL_LEN") +AFD_PORT=${AFD_PORT:-6271} +API_PORT=${API_PORT:-18307} +ENABLE_DBO=${ENABLE_DBO:-0} +DBO_ARGS=() +if [ "$ENABLE_DBO" = 1 ]; then + DBO_ARGS=( + --enable-dbo + --dbo-decode-token-threshold "${DBO_DECODE_THRESHOLD:-2}" + --dbo-prefill-token-threshold "${DBO_PREFILL_THRESHOLD:-12}" + ) +fi + +ROLE_ARGS=(serve "$MODEL_PATH" + --data-parallel-size 2 + --tensor-parallel-size 1 + --max-num-seqs "$MAX_NUM_SEQS" + --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" + "${MAX_MODEL_LEN_ARG[@]}" + --tokenizer-mode deepseek_v4 + --kv-cache-dtype fp8 + --gpu-memory-utilization "$GPU_MEM_UTIL" + --enforce-eager + "${DBO_ARGS[@]}" + --api-server-count 1 + --host 127.0.0.1 + --trust-remote-code) + +CUDA_VISIBLE_DEVICES="$ATTN_DEVICES" $VLLM_CMD "${ROLE_ARGS[@]}" \ + --served-model-name deepseek-v4-flash-afd-attention \ + --additional-config "{ + \"afd\": { + \"role\": \"attention\", + \"connector\": \"P2pNcclAFDConnector\", + \"host\": \"127.0.0.1\", + \"port\": $AFD_PORT, + \"num_attention_ranks\": 2, + \"num_ffn_ranks\": 2 + } + }" \ + --port "$API_PORT" > "$LOG_DIR/attn.log" 2>&1 & +ATTN_PID=$! + +CUDA_VISIBLE_DEVICES="$FFN_DEVICES" $VLLM_CMD "${ROLE_ARGS[@]}" \ + --served-model-name deepseek-v4-flash-afd-ffn \ + --additional-config "{ + \"afd\": { + \"role\": \"ffn\", + \"connector\": \"P2pNcclAFDConnector\", + \"host\": \"127.0.0.1\", + \"port\": $AFD_PORT, + \"num_attention_ranks\": 2, + \"num_ffn_ranks\": 2 + } + }" \ + --port "$((API_PORT + 1))" > "$LOG_DIR/ffn.log" 2>&1 & +FFN_PID=$! + +# shellcheck disable=SC2317,SC2329 # invoked by the EXIT trap below. +# (SC2329 on shellcheck >= 0.11, SC2317 on older ones; CI runs an older one.) +cleanup() { + kill "$ATTN_PID" "$FFN_PID" 2>/dev/null + wait "$ATTN_PID" "$FFN_PID" 2>/dev/null +} +trap cleanup EXIT + +for _ in $(seq 1 "${READY_TIMEOUT:-900}"); do + if curl -sf "http://127.0.0.1:$API_PORT/health" > /dev/null 2>&1; then + echo "server ready on http://127.0.0.1:$API_PORT" + if [ -n "${SMOKE:-}" ]; then + curl -s "http://127.0.0.1:$API_PORT/v1/completions" \ + -H 'Content-Type: application/json' \ + -d "{\"model\":\"deepseek-v4-flash-afd-attention\",\"prompt\":\"The capital of France is\",\"max_tokens\":8,\"temperature\":0}" + echo + exit 0 + fi + wait "$ATTN_PID" "$FFN_PID" + exit 0 + fi + if ! kill -0 "$ATTN_PID" 2>/dev/null || ! kill -0 "$FFN_PID" 2>/dev/null; then + echo "a server exited early; see $LOG_DIR/attn.log and $LOG_DIR/ffn.log" >&2 + exit 1 + fi + sleep 1 +done +echo "timed out waiting for the server" >&2 +exit 1 diff --git a/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py b/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py index fa39bd24..092d1cd2 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project from __future__ import annotations import pytest @@ -28,7 +30,6 @@ def __iter__(self): @pytest.mark.parametrize( "name", [ - "layers.0.ffn.gate.weight", "layers.1.ffn.experts.0.w1.weight", "model.layers.2.ffn.shared_experts.w2.weight", ], @@ -37,6 +38,15 @@ def test_v4_raw_checkpoint_ffn_paths_are_ffn_owned(name): assert _checkpoint_weight_roles(name) == frozenset(("ffn",)) +def test_v4_gate_loads_on_both_roles(): + # The gate's parameters live under .ffn so the checkpoint names resolve, + # but with compute_gate_on_attention the Attention side is what runs it -- + # so both roles have to load the same tensor. + assert _checkpoint_weight_roles("layers.0.ffn.gate.weight") == frozenset( + ("attention", "ffn") + ) + + @pytest.mark.parametrize( "name", [ @@ -78,6 +88,7 @@ def test_v4_raw_checkpoint_public_paths_are_shared(name): [ "layers.0.attn.fused_wqa_wkv.weight", "layers.0.hc_ffn_fn", + "layers.0.ffn.gate.weight", "model.hc_head_fn", "embed.weight", ], @@ -104,7 +115,7 @@ def test_v4_load_weights_filters_raw_checkpoint_names_once( "embed.weight", ] weights = _OneShotWeights(names) - seen = [] + seen: list[str] = [] native_result = {"native.loaded"} def fake_native_loader(self, filtered_weights): diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_forward_context.py index 136aee08..6e9ad29f 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -404,14 +404,19 @@ def test_deepseek_compute_gate_on_attention_selects_backend_boundary(): assert "self.mlp = AFDDeepseekV2RemoteExpertsMoE(" in source assert "self.mlp = GateOnlyRemoteMoE(" in source assert 'prefix=f"{prefix}.mlp"' in source + # The gate/topk helper delegates expert selection to the connector, so both + # platforms share it; only the FFN-side MoE compute stays platform-split. assert ( - "# NPU-only: Attention-side gate/topk is implemented in the NPU helper." + "# The gate helper delegates expert selection to the connector, so both" in source ) assert ( "# NPU-only: gated MoE FFN compute consumes Attention-side topk payloads." in source ) + # CUDA reaches its own grouped-GEMM entry point only once tokens arrive + # pre-routed; without a group list the control-plane path still applies. + assert "gpu_attention_gate.compute_attention_gate_moe_ffn(" in source def test_async_moe_pipeline_preserves_stage_order(monkeypatch): diff --git a/tests/unit/v1/worker/test_attention_model_runner.py b/tests/unit/v1/worker/test_attention_model_runner.py index 5f940128..7a8cbe03 100644 --- a/tests/unit/v1/worker/test_attention_model_runner.py +++ b/tests/unit/v1/worker/test_attention_model_runner.py @@ -1,7 +1,11 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations import sys from types import SimpleNamespace +from typing import Any import pytest @@ -401,10 +405,10 @@ def test_ubatch_missing_metadata_uses_complete_public_installer(): wrapper._install_missing_afd_metadata(forward_context) - metadata = forward_context.additional_kwargs["afd_metadata"] - assert metadata is runner._afd_pending_metadata + metadata: Any = forward_context.additional_kwargs["afd_metadata"] assert metadata.transaction_id == "afd-0" assert metadata.tokens_lens == [3, 5] + assert metadata is runner._afd_pending_metadata assert set(runner.connector.sent_dp_metadata_lists[0]) == {0, 1} @@ -423,13 +427,21 @@ def test_phase5_allows_two_way_ubatching_but_rejects_other_counts(): ) -def _ubatch_runner(uniform_decode, **parallel_overrides): +_DUMMY_CONTROL_PLANE = object() + + +def _ubatch_runner( + uniform_decode, *, control_plane=_DUMMY_CONTROL_PLANE, **parallel_overrides +): runner = object.__new__(AFDAttentionModelRunner) runner.vllm_config = SimpleNamespace( parallel_config=_parallel_config(**parallel_overrides), ) runner.uniform_decode_query_len = 1 runner._is_uniform_decode = lambda **_kwargs: uniform_decode + # The override consults the connector to decide whether cross-DP batch + # coordination applies; a non-None control plane keeps upstream behaviour. + runner.connector = SimpleNamespace(control_plane=control_plane) return runner @@ -1091,6 +1103,9 @@ def __init__(self, events): self.events = events self.control_plane = object() self._initialized = False + # Every real connector carries one; the runner reads it to decide + # whether async MoE ubatching is configured. + self.extra_info = None @property def is_initialized(self): @@ -1119,7 +1134,7 @@ def _fake_connector_factory(monkeypatch, connector): def test_attention_runner_constructor_does_not_initialize_connector(monkeypatch): import afd_plugin.v1.worker.attention_model_runner as attention_model_runner - events = [] + events: list[str] = [] connector = _LifecycleConnector(events) def fake_native_init(self, vllm_config, device): @@ -1163,7 +1178,7 @@ def test_attention_runner_load_model_initializes_connector_after_weights( monkeypatch, use_ubatching, ): - events = [] + events: list[str] = [] connector = _LifecycleConnector(events) runner = object.__new__(AFDAttentionModelRunner) runner.connector = connector @@ -1192,3 +1207,35 @@ def test_attention_runner_load_model_initializes_connector_after_weights( expected.append("connector_init") assert events == expected assert connector.is_initialized is True + + +def test_connector_driven_runs_skip_cross_dp_batch_coordination(): + """An idle Attention replica never joins the DP all-reduce. + + Async AFD lets each replica advance alone, so a busy replica must not block + in ``coordinate_batch_across_dp`` waiting for one that never steps. + """ + import vllm.v1.worker.gpu_model_runner as gpu_model_runner + + from afd_plugin.v1.worker.attention_model_runner import ( + _dp_batch_coordination_disabled, + ) + + original = gpu_model_runner.coordinate_batch_across_dp + + with _dp_batch_coordination_disabled(True): + assert gpu_model_runner.coordinate_batch_across_dp is not original + result = gpu_model_runner.coordinate_batch_across_dp( + num_tokens_unpadded=8, + parallel_config=None, + allow_microbatching=False, + num_tokens_padded=8, + uniform_decode=True, + cudagraph_mode=0, + ) + # num_tokens_across_dp None makes upstream skip its DP-padding branch. + assert result == (False, None, 0) + assert gpu_model_runner.coordinate_batch_across_dp is original + + with _dp_batch_coordination_disabled(False): + assert gpu_model_runner.coordinate_batch_across_dp is original diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index c0d198d0..b7c1eff7 100644 --- a/tests/unit/v1/worker/test_ffn_model_runner.py +++ b/tests/unit/v1/worker/test_ffn_model_runner.py @@ -1,9 +1,13 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations import logging import threading from collections import deque from types import SimpleNamespace +from typing import Any import pytest @@ -20,6 +24,7 @@ AFDTransferContext, AFDTransferMetadata, ) +from afd_plugin.connectors.gpu.async_gpu import ConnectorShutdown # noqa: E402 from afd_plugin.model_executor.models.deepseek_v2 import ( # noqa: E402 AFDDeepseekV2ForCausalLM, ) @@ -33,7 +38,7 @@ class _FakeConnector: def __init__(self): - self.attn_outputs = deque() + self.attn_outputs: deque = deque() self.ffn_outputs = [] self.expert_routing_specs = [] self.recv_input_ids = [] @@ -82,7 +87,7 @@ def close(self): class _ConnectorDrivenFakeConnector(_FakeConnector): def __init__(self): super().__init__() - self.control_plane = None + self.control_plane: Any = None class _FakeModel: @@ -236,7 +241,9 @@ class _InputIdsModel(_FakeModel): def __init__(self): self.calls = [] - def compute_ffn_output(self, hidden_states, layer_idx, *, input_ids): + def compute_ffn_output( # type: ignore[override] + self, hidden_states, layer_idx, *, input_ids + ): self.calls.append((hidden_states, layer_idx, input_ids)) return input_ids @@ -739,18 +746,45 @@ def test_ffn_worker_reports_zero_compilation_times(): assert compilation_times.encoder == 0.0 -def test_ffn_worker_loop_rejects_connector_without_control_plane(): +def test_ffn_worker_loop_drives_connector_without_control_plane(): worker = object.__new__(AFDFFNWorker) event = threading.Event() + steps = [] + + def execute_connector_driven_step(): + steps.append(1) + # The connector-driven step returns on an idle poll; the loop must come + # back to the shutdown event rather than block forever. + if len(steps) == 3: + event.set() worker._ffn_shutdown_event = event worker.device = SimpleNamespace(type="cpu") worker.model_runner = SimpleNamespace( connector=_ConnectorDrivenFakeConnector(), + execute_connector_driven_step=execute_connector_driven_step, + ) + + worker._run_ffn_server_loop() + + assert len(steps) == 3 + + +def test_ffn_worker_loop_exits_cleanly_when_peer_announces_shutdown(): + worker = object.__new__(AFDFFNWorker) + + def execute_connector_driven_step(): + raise ConnectorShutdown("peer left") + + worker._ffn_shutdown_event = threading.Event() + worker.device = SimpleNamespace(type="cpu") + worker.model_runner = SimpleNamespace( + connector=_ConnectorDrivenFakeConnector(), + execute_connector_driven_step=execute_connector_driven_step, ) - with pytest.raises(NotImplementedError, match="control-plane-driven"): - worker._run_ffn_server_loop() + # A peer shutdown is an ordinary exit, not a loop failure. + worker._run_ffn_server_loop() def test_ffn_worker_loop_logs_unexpected_thread_errors(caplog): From abcca4836ecf4fb04bde18f8c6dc291d0b57aec1 Mon Sep 17 00:00:00 2001 From: specture724 Date: Thu, 17 Sep 2026 11:22:29 +0800 Subject: [PATCH 2/2] fix: address review findings on the DeepSeek-V4 integration Review follow-up on #334. - `get_experts_routing_spec` read `self.layers[layer_idx].gate`, which no decoder layer sets: the gate hangs off the remote-FFN proxy on Attention and off the native MoE on FFN, so `.ffn.gate` reaches the right one on both roles -- which is what the forward path a few hundred lines below already does. Dead today because no caller combination reaches it; an AttributeError for the first one that does. - The coordinator-timeout patch applied to every process with the plugin installed, so a plain non-AFD vLLM DP run waited 600s instead of 120s for a broken coordinator. The AFD check now happens per call, the way ffn_local_moe_prepare does it, and a non-AFD run keeps upstream's wait. - Both new compat patches were missing from the inventory table and were the only inventory patches without a unit test. Added, with CPU-safe tests covering role gating, own-dispatch kernels, install idempotence, the env parse, and both timeout defaults. - The connector-driven drain loop ran only on GPU: both worker-loop tests stub execute_connector_driven_step. Ported the NPU twin's test shape with a fake work-item connector, covering the idle-poll return, the transfer-state check, per-item forward-context and metadata installation, the drain bound, and shutdown propagation. - Design pages: model_integration's DSV4-CUDA section described only the synchronous path and is now false on the async one, so both shapes are tabulated; ffn_runtime described the assertion and the raise this PR deletes. - The recipe hardcoded a personal model path where its sibling uses a placeholder. Unit suite: 2 pre-existing failures (os.pidfd_open missing on this build). pre-commit clean over the changed files. Co-Authored-By: Claude Opus 5 --- .../compat/patches/dp_coordinator_timeout.py | 27 +++- .../model_executor/models/deepseek_v4.py | 6 +- .../module/compatibility_and_patches.md | 2 + docs/design/module/ffn_runtime.md | 16 +- docs/design/module/model_integration.md | 20 ++- .../deepseek_v4_flash/2a2f_eager_async.sh | 2 +- .../patches/test_dp_coordinator_timeout.py | 66 ++++++++ .../patches/test_ffn_local_moe_prepare.py | 120 ++++++++++++++ .../models/test_deepseek_v4_routing_spec.py | 39 +++++ .../worker/test_ffn_connector_driven_loop.py | 148 ++++++++++++++++++ 10 files changed, 432 insertions(+), 14 deletions(-) create mode 100644 tests/unit/compat/patches/test_dp_coordinator_timeout.py create mode 100644 tests/unit/compat/patches/test_ffn_local_moe_prepare.py create mode 100644 tests/unit/model_executor/models/test_deepseek_v4_routing_spec.py create mode 100644 tests/unit/v1/worker/test_ffn_connector_driven_loop.py diff --git a/afd_plugin/compat/patches/dp_coordinator_timeout.py b/afd_plugin/compat/patches/dp_coordinator_timeout.py index 8a13e773..b3c9ea8f 100644 --- a/afd_plugin/compat/patches/dp_coordinator_timeout.py +++ b/afd_plugin/compat/patches/dp_coordinator_timeout.py @@ -16,22 +16,43 @@ import os import vllm.v1.engine.coordinator as coordinator_module +from vllm.config import get_current_vllm_config + +from afd_plugin.config import parse_optional_afd_config DEFAULT_TIMEOUT_S = 600 +# What upstream hardcodes. A process with the plugin installed but no AFD +# configuration is an ordinary vLLM DP run and must keep waiting exactly this +# long; the patch applies at register_afd() for every such process, so the role +# check has to happen per call, the way ffn_local_moe_prepare does it. +UPSTREAM_TIMEOUT_S = 120 + + +def _afd_is_active() -> bool: + try: + afd_config = parse_optional_afd_config( + get_current_vllm_config(), + validate=False, + ) + except Exception: + return False + return afd_config is not None # Patch reason: the upstream DP Coordinator startup wait is hardcoded to 120 # seconds, which is not enough for the coordinator subprocess to import and # bind on a CPU-oversubscribed shared machine -- both AFD roles then die # during startup. -# Patch functionality: identical to upstream, except the wait comes from -# AFD_DP_COORDINATOR_TIMEOUT_S (default 600 seconds). +# Patch functionality: identical to upstream, except that an AFD run takes its +# wait from AFD_DP_COORDINATOR_TIMEOUT_S (default 600 seconds). A non-AFD run +# keeps upstream's 120 seconds. # Signature: matches upstream; no added parameters. # Upstream: vLLM v0.26.0, vllm/v1/engine/coordinator.py def _wait_for_zmq_addrs(self, zmq_addr_pipe) -> tuple[str, str, str]: try: + default_timeout = DEFAULT_TIMEOUT_S if _afd_is_active() else UPSTREAM_TIMEOUT_S timeout = int( - os.getenv("AFD_DP_COORDINATOR_TIMEOUT_S", str(DEFAULT_TIMEOUT_S)), + os.getenv("AFD_DP_COORDINATOR_TIMEOUT_S", str(default_timeout)), ) ready = multiprocessing.connection.wait( [zmq_addr_pipe, self.proc.sentinel], timeout=timeout diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index 3850b1cf..18707099 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -634,7 +634,11 @@ def get_experts_layer_indices(self) -> tuple[int, ...]: def get_experts_routing_spec(self, layer_idx: int) -> AFDExpertRoutingSpec: """Router contract for the async FFN loop's receive buffers.""" - gate = self.layers[layer_idx].gate + # The decoder layer has no gate of its own. On Attention it lives on the + # remote-FFN proxy under the checkpoint's own .ffn.gate path, and on FFN + # it is the native MoE's gate -- .ffn reaches the right one either way, + # which is what the forward path a few hundred lines down already does. + gate = self.layers[layer_idx].ffn.gate return AFDExpertRoutingSpec( router_logits_width=int(self.config.n_routed_experts), router_logits_dtype=gate.out_dtype or gate.weight.dtype, diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index 5d7dd1d8..d8d89f04 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/docs/design/module/compatibility_and_patches.md @@ -120,6 +120,8 @@ not the package dependency policy. | [`async_dp_forward_context.py`](../../../afd_plugin/compat/patches/async_dp_forward_context.py): `vllm.forward_context.set_forward_context` plus already-imported worker aliases | Skips native `DPMetadata` construction/coordination only for AFD async-DP; otherwise uses the copied upstream flow. | Imported by `register_afd`; same target/dev/unknown guard. Rebinds known already-imported aliases so callers do not retain the old function. | [`test_async_dp_forward_context.py`](../../../tests/unit/compat/patches/test_async_dp_forward_context.py) covers async skip and non-async coordination. | Remove when vLLM supports a per-engine-role opt-out from native MoE DP metadata coordination. | | [`config_validation.py`](../../../afd_plugin/compat/patches/config_validation.py): `EngineArgs.create_engine_config`, `VllmConfig.__post_init__` | For AFD-owned ubatching with a non-DeepEP backend, temporarily presents `deepep_low_latency` during upstream validation and restores the configured backend. After upstream platform normalization, maps an initial `worker_cls="auto"` to the role-specific CUDA or standard Ascend AFD worker. | Imported by `register_afd`; accepts the target version, development versions, or missing version metadata. Saves originals on upstream modules under AFD-specific attributes before installing wrappers. Explicit worker paths and non-AFD configs are not remapped. | [`test_config_validation.py`](../../../tests/unit/compat/patches/test_config_validation.py) covers backend relaxation, four role/platform mappings, explicit and non-AFD preservation, repeated validation, unsupported platforms, and dev versions. | Remove the backend branch when upstream validation distinguishes plugin-owned ubatching; remove worker mapping when vLLM offers plugin-owned role-aware worker selection. | | [`engine_core.py`](../../../afd_plugin/compat/patches/engine_core.py): `EngineCore.__init__`, `_initialize_kv_caches`, `shutdown`; `EngineCoreProc.run_busy_loop`; `DPEngineCoreProc.run_busy_loop` | AFD FFN becomes a connector daemon: construct executor, skip scheduler/KV setup, return an empty KV-shaped result on late paths, start/monitor/stop the FFN worker loop, and use FFN-safe shutdown. Non-FFN branches copy pinned upstream behavior. | Imported by `register_afd`; **no patch-local version guard and no saved-original sentinel**. Direct class assignment means the package pin and review discipline are the compatibility guard. | [`test_engine_core.py`](../../../tests/unit/compat/patches/test_engine_core.py) covers FFN initialization, non-FFN behavior, and daemon start/stop; role runtime tests cover error propagation. | Remove when vLLM offers a headless connector-daemon engine lifecycle or an executor mode that does not require scheduler/KV ownership. | +| [`dp_coordinator_timeout.py`](../../../afd_plugin/compat/patches/dp_coordinator_timeout.py): `DPCoordinator._wait_for_zmq_addrs` | An AFD run takes the coordinator's startup ZMQ wait from `AFD_DP_COORDINATOR_TIMEOUT_S` (default 600s) instead of upstream's hardcoded 120s, which a CPU-oversubscribed box can exceed and kill both roles at startup. A non-AFD process keeps 120s. | Imported by `register_afd`; no version guard. Applies unconditionally, so the AFD check is per call (`get_current_vllm_config`), matching `ffn_local_moe_prepare`. Rebinding the method twice is idempotent. | [`test_dp_coordinator_timeout.py`](../../../tests/unit/compat/patches/test_dp_coordinator_timeout.py) covers both defaults, the environment override, and the unresolvable-config path. | Remove when the coordinator's startup wait is configurable upstream. | +| [`ffn_local_moe_prepare.py`](../../../afd_plugin/compat/patches/ffn_local_moe_prepare.py): `vllm...all2all_utils.maybe_make_prepare_finalize` plus already-imported aliases | The AFD FFN role prepares and finalizes locally because AFD already routed the rows to the rank that owns them; a second all-to-all would move them again. Kernels with their own dispatch, every other role, and non-AFD processes reach the upstream selector unchanged. | Imported by `register_afd`; no version guard. Re-aliases the source module and any module already holding the upstream function, guarded by an `_afd_installed` sentinel. | [`test_ffn_local_moe_prepare.py`](../../../tests/unit/compat/patches/test_ffn_local_moe_prepare.py) covers role gating, every own-dispatch kernel flag, keyword arrival, install idempotence, and the unresolvable-config path. | Remove when vLLM lets a plugin declare that rows arrive pre-dispatched. | | [`npu/ascend_platform.py`](../../../afd_plugin/compat/patches/npu/ascend_platform.py): `NPUPlatform.check_and_update_config` | Snapshots AFD DBO state, runs upstream normalization, and restores configured `enable_dbo`, `ubatch_size`, and `all2all_backend` in `finally`; non-AFD behavior is unchanged. | Installed through the config facade during AFD NPU normalization and verified again by the worker runtime facade; no version guard. Saves the original on the class and uses a class sentinel. The runtime facade caches success only after the wrapper is installed, so an early missing vLLM-Ascend import remains retryable. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) and [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py). | Remove when vLLM-Ascend recognizes plugin-owned DBO workers or no longer clears these fields. | | [`npu/mla_graph.py`](../../../afd_plugin/compat/patches/npu/mla_graph.py): `vllm_ascend.attention.mla_v1.get_graph_params` | Resolves an AFD ubatch-owned `GraphParams` registry from the active forward context and otherwise delegates to the saved upstream process-global resolver. | Called through `apply_afd_ascend_patches_if_needed` during NPU worker startup; no version guard. Saves the original on the upstream module and uses a module sentinel. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) covers AFD-context resolution, upstream fallback, and idempotence; [`test_npu_mla_graph.py`](../../../tests/unit/v1/worker/test_npu_mla_graph.py) covers the owning graph lifecycle. | Remove when vLLM-Ascend accepts a forward-context-local MLA graph registry or exposes an equivalent resolver hook. | | [`npu/force_load_balance.py`](../../../afd_plugin/compat/patches/npu/force_load_balance.py): `AscendW8A8DynamicFusedMoEMethod.__init__`, `AscendW8A8DynamicFusedMoEMethod.apply` | Captures AFD profiling configuration as method-owned state and replaces routed expert IDs with a deterministic balanced buffer only when the method-owned switch is enabled; normal model-selected routing remains unchanged. This switch changes outputs and is not a correctness feature. | Imported by the NPU FFN worker after vLLM-Ascend platform initialization; **no patch-local version guard or explicit reload sentinel**. Functions copy the current upstream bodies with marked AFD deltas. | [`test_force_load_balance.py`](../../../tests/unit/compat/patches/test_force_load_balance.py) covers buffer bounds, determinism, growth, override, and pass-through. | Upstream a deterministic expert-routing profiling hook in vLLM-Ascend, then delete both copied functions. | diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index 189de1be..77859917 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -155,13 +155,15 @@ The worker selects one of two FFN step paths from the optional | Selection state | Connectors | Worker behavior | | --- | --- | --- | | `control_plane is not None` | `P2pNcclAFDConnector`, `CAMP2pAFDConnector` | Call `control_plane.recv_dp_metadata_list()`, then profile, warm, capture, replay, or execute its stage map. | -| `control_plane is None` | `CAMAsyncAFDConnector` (NPU only) | Block directly on a connector work item; no separate DP-metadata control plane. | - -The connector-driven path exists only on Ascend. GPU FFN supports -control-plane-driven connectors exclusively: `GPUFFNModelRunner` asserts -`connector.control_plane is not None` at construction, and the GPU daemon loop -raises `NotImplementedError` if a connector without a control plane is ever -installed. +| `control_plane is None` | `CAMAsyncAFDConnector`, `GpuAsyncAFDConnector` | Block directly on a connector work item; no separate DP-metadata control plane. | + +The connector-driven path now runs on both platforms. `GPUFFNModelRunner` no +longer asserts `connector.control_plane is not None` at construction, and the +GPU daemon loop calls `execute_connector_driven_step` where it used to raise +`NotImplementedError`. That step drains at most `num_layers` work items and +returns on an idle poll, which is what lets the worker loop observe its +shutdown event; successive items may belong to different layers of different +Attention replicas, so each installs its own forward context and AFD metadata. ```mermaid flowchart TD diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 26b3c77e..dfe19ea0 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -133,8 +133,24 @@ parameter-free `RemoteDeepseekV4FFN`. FFN owns the complete native Role-aware weight filtering keeps layer `.ffn` parameters on FFN and all other layer-local and mHC head parameters on Attention; common non-layer paths remain available to both roles for the native loader lifecycle. The adapter is -CUDA-only and requires synchronous `P2pNcclAFDConnector`, -`compute_gate_on_attention=false`, and pipeline-parallel size 1. It rejects +CUDA-only and requires pipeline-parallel size 1. + +The paragraph above describes the synchronous path. `GpuAsyncAFDConnector` +inverts two of its facts, so the boundary now has two shapes: + +| | `P2pNcclAFDConnector` | `GpuAsyncAFDConnector` | +| --- | --- | --- | +| `compute_gate_on_attention` | must be `false` | must be `true` | +| Where routing runs | FFN, on the native hash router | Attention, on a gate proxy | +| What crosses the boundary | the activation plus token-aligned `input_ids` | the activation plus the topk already chosen | +| Where the gate's weights load | the FFN role's native MoE | **both** roles | + +On the async path the gate parameters live on the Attention-side proxy but keep +the checkpoint's own `.ffn.gate.*` names, so the same tensors load onto the +proxy on Attention and onto the native MoE gate on FFN -- +`test_v4_gate_loads_on_both_roles` pins exactly that. Reading the gate off a +decoder layer therefore goes through `.ffn`, which resolves to the proxy on one +role and the native MoE on the other. It rejects sequence-parallel MoE, EPLB, and the `deep_gemm_mega_moe` backend. The P2P connector validates one-dimensional `torch.int32` input IDs and preallocates their receive buffers for graph execution. This boundary currently has diff --git a/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh b/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh index 7fdfc421..ade2e03f 100755 --- a/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh +++ b/recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh @@ -12,7 +12,7 @@ # gpu run --gpus 4 -- bash recipe/gpu/P2pNcclAFDConnector/deepseek_v4_flash/2a2f_eager_async.sh set -u -MODEL_PATH=${MODEL_PATH:-/data/boao/deepseek-v4-flash} +MODEL_PATH=${MODEL_PATH:-/path/model_weights/deepseek-v4-flash} VLLM_CMD=${VLLM_CMD:-vllm} LOG_DIR=${LOG_DIR:-.} mkdir -p "$LOG_DIR" diff --git a/tests/unit/compat/patches/test_dp_coordinator_timeout.py b/tests/unit/compat/patches/test_dp_coordinator_timeout.py new file mode 100644 index 00000000..07931130 --- /dev/null +++ b/tests/unit/compat/patches/test_dp_coordinator_timeout.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Startup wait selection for the patched DP Coordinator. + +The patch applies at ``register_afd()`` in every process that has the plugin +installed, so the thing worth pinning is that a plain non-AFD DP run still gets +upstream's 120 seconds rather than silently waiting five times as long. +""" + +from __future__ import annotations + +import pytest + +from afd_plugin.compat.patches import dp_coordinator_timeout as patch + + +class _ClosedPipe: + """A pipe that reports nothing ready, so the wait always times out.""" + + def __init__(self): + self.closed = False + + def close(self): + self.closed = True + + +def _run_wait(monkeypatch, *, afd_active: bool, env: str | None): + monkeypatch.setattr(patch, "_afd_is_active", lambda: afd_active) + if env is None: + monkeypatch.delenv("AFD_DP_COORDINATOR_TIMEOUT_S", raising=False) + else: + monkeypatch.setenv("AFD_DP_COORDINATOR_TIMEOUT_S", env) + + seen = {} + + def fake_wait(_objects, timeout): + seen["timeout"] = timeout + return [] + + monkeypatch.setattr(patch.multiprocessing.connection, "wait", fake_wait) + + pipe = _ClosedPipe() + coordinator = type("_C", (), {"proc": type("_P", (), {"sentinel": object()})()})() + with pytest.raises(RuntimeError, match="within timeout"): + patch._wait_for_zmq_addrs(coordinator, pipe) + assert pipe.closed, "the pipe must be closed even when the wait fails" + return seen["timeout"] + + +def test_a_non_afd_run_keeps_upstreams_wait(monkeypatch): + assert _run_wait(monkeypatch, afd_active=False, env=None) == 120 + + +def test_an_afd_run_gets_the_longer_wait(monkeypatch): + assert _run_wait(monkeypatch, afd_active=True, env=None) == 600 + + +@pytest.mark.parametrize("afd_active", [True, False]) +def test_the_environment_overrides_either_default(monkeypatch, afd_active): + assert _run_wait(monkeypatch, afd_active=afd_active, env="42") == 42 + + +def test_a_process_with_no_vllm_config_is_not_afd(): + # get_current_vllm_config() outside an engine must not propagate; the + # coordinator starts before any AFD configuration is resolvable. + assert patch._afd_is_active() in (True, False) diff --git a/tests/unit/compat/patches/test_ffn_local_moe_prepare.py b/tests/unit/compat/patches/test_ffn_local_moe_prepare.py new file mode 100644 index 00000000..495f36f0 --- /dev/null +++ b/tests/unit/compat/patches/test_ffn_local_moe_prepare.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Role gating for the FFN-local MoE prepare/finalize selector. + +AFD pre-routes rows to the FFN rank that owns them, so the FFN role must not +run another all-to-all. Every other process -- the Attention role, and plain +vLLM -- has to reach upstream's selector untouched: this patch is installed +process-wide, so getting the gate wrong changes runs that have nothing to do +with AFD. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from afd_plugin.compat.patches import ffn_local_moe_prepare as patch + +ALL_KERNEL_FLAGS = ( + "use_deepep_ht_kernels", + "use_deepep_ll_kernels", + "use_deepep_v2_kernels", + "use_fi_nvl_two_sided_kernels", + "use_fi_nvl_one_sided_kernels", + "use_nixl_ep_kernels", + "use_mori_kernels", +) + + +def _moe(**overrides): + flags = dict.fromkeys(ALL_KERNEL_FLAGS, False) + flags.update(overrides) + return SimpleNamespace(moe_parallel_config=SimpleNamespace(**flags)) + + +@pytest.fixture +def upstream(monkeypatch): + calls: list[tuple] = [] + + def selector(*args, **kwargs): + calls.append((args, kwargs)) + return "upstream" + + monkeypatch.setattr(patch, "_UPSTREAM_SELECTOR", selector) + return calls + + +@pytest.fixture +def local(monkeypatch): + calls: list[dict] = [] + + def local_prepare(**kwargs): + calls.append(kwargs) + return "local" + + monkeypatch.setattr(patch, "make_moe_prepare_and_finalize_no_dp_ep", local_prepare) + return calls + + +def test_a_non_ffn_process_reaches_upstream_unchanged(monkeypatch, upstream, local): + monkeypatch.setattr(patch, "_is_afd_ffn_role", lambda: False) + + result = patch.maybe_make_prepare_finalize(_moe(), use_monolithic=True) + + assert result == "upstream" + assert local == [] + # The arguments must arrive exactly as given; this is a pass-through. + assert upstream[0][1] == {"use_monolithic": True} + + +def test_the_ffn_role_prepares_locally(monkeypatch, upstream, local): + monkeypatch.setattr(patch, "_is_afd_ffn_role", lambda: True) + + result = patch.maybe_make_prepare_finalize(_moe(), use_monolithic=True) + + assert result == "local" + assert upstream == [] + assert local == [{"use_monolithic": True}] + + +@pytest.mark.parametrize("flag", ALL_KERNEL_FLAGS) +def test_kernels_that_own_their_dispatch_are_left_alone( + monkeypatch, + upstream, + local, + flag, +): + monkeypatch.setattr(patch, "_is_afd_ffn_role", lambda: True) + + result = patch.maybe_make_prepare_finalize(_moe(**{flag: True})) + + assert result == "upstream", f"{flag} must keep its own collective" + assert local == [] + + +def test_the_moe_may_arrive_as_a_keyword(monkeypatch, upstream, local): + monkeypatch.setattr(patch, "_is_afd_ffn_role", lambda: True) + + assert patch.maybe_make_prepare_finalize(moe=_moe()) == "local" + + +def test_installing_twice_rebinds_once(): + # Already installed at import; a second call must be a no-op rather than + # wrapping the wrapper. + installed = patch.all2all_utils_module.maybe_make_prepare_finalize + + patch.apply_local_moe_prepare() + + assert patch.all2all_utils_module.maybe_make_prepare_finalize is installed + assert installed is patch.maybe_make_prepare_finalize + + +def test_an_unresolvable_config_is_not_the_ffn_role(monkeypatch): + def boom(): + raise RuntimeError("no vllm config in this process") + + monkeypatch.setattr(patch, "get_current_vllm_config", boom) + + assert patch._is_afd_ffn_role() is False diff --git a/tests/unit/model_executor/models/test_deepseek_v4_routing_spec.py b/tests/unit/model_executor/models/test_deepseek_v4_routing_spec.py new file mode 100644 index 00000000..42afca5e --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_routing_spec.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Where the DSV4 router contract finds its gate.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.model_executor.models import deepseek_v4 as adapter # noqa: E402 + + +def test_the_routing_spec_reads_the_gate_through_ffn(): + # A decoder layer has no gate of its own: on Attention it hangs off the + # remote-FFN proxy, on FFN off the native MoE. Reading `.gate` on the layer + # was an AttributeError waiting for its first caller. + gate = SimpleNamespace(out_dtype=torch.float32, weight=SimpleNamespace(dtype=None)) + layer = SimpleNamespace(ffn=SimpleNamespace(gate=gate)) + model = object.__new__(adapter.AFDDeepseekV4Model) + model.layers = {3: layer} + model.config = SimpleNamespace(n_routed_experts=64) + + spec = model.get_experts_routing_spec(3) + + assert spec.router_logits_width == 64 + assert spec.router_logits_dtype is torch.float32 + + +def test_the_routing_spec_falls_back_to_the_gate_weight_dtype(): + gate = SimpleNamespace(out_dtype=None, weight=SimpleNamespace(dtype=torch.bfloat16)) + model = object.__new__(adapter.AFDDeepseekV4Model) + model.layers = {0: SimpleNamespace(ffn=SimpleNamespace(gate=gate))} + model.config = SimpleNamespace(n_routed_experts=8) + + assert model.get_experts_routing_spec(0).router_logits_dtype is torch.bfloat16 diff --git a/tests/unit/v1/worker/test_ffn_connector_driven_loop.py b/tests/unit/v1/worker/test_ffn_connector_driven_loop.py new file mode 100644 index 00000000..1b53a449 --- /dev/null +++ b/tests/unit/v1/worker/test_ffn_connector_driven_loop.py @@ -0,0 +1,148 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""The connector-driven FFN drain loop, without a GPU. + +The worker-loop tests stub ``execute_connector_driven_step`` wholesale, so the +loop underneath it -- idle-poll return, state checking, per-item metadata +installation, output send-back -- ran only on a device. The NPU twin of this +loop is unit-tested; this is the same shape with a fake work-item connector. +""" + +from __future__ import annotations + +from contextlib import contextmanager +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + +from afd_plugin.connectors.gpu.async_gpu import ( # noqa: E402 + ConnectorShutdown, + GpuAsyncTransferState, +) +from afd_plugin.v1.worker import ffn_model_runner as module # noqa: E402 + +NUM_LAYERS = 3 + + +class _FakeWorkItem: + def __init__(self, layer_idx, states, metadata): + self.layer_idx = layer_idx + self.hidden_states = torch.zeros(2, 4) + self.context = SimpleNamespace(states=states, metadata=metadata) + + +class _FakeConnector: + """Hands out a scripted sequence of work items, then whatever ends it.""" + + def __init__(self, script): + self.script = list(script) + self.received = [] + self.sent = [] + + def recv_ffn_work_item(self, *, stage_idx, max_num_tokens, **_): + self.received.append((stage_idx, max_num_tokens)) + if not self.script: + raise TimeoutError + nxt = self.script.pop(0) + if isinstance(nxt, BaseException): + raise nxt + return nxt + + def send_ffn_work_item_output(self, work_item, output): + self.sent.append((work_item.layer_idx, output)) + return output + + +def _states(): + state = object.__new__(GpuAsyncTransferState) + state.group_list = None + state.expand_x_shared = None + return state + + +@pytest.fixture +def runner(monkeypatch): + forward_context = SimpleNamespace( + dp_metadata="stale", + additional_kwargs={}, + all_moe_layers=[f"model.layers.{i}.mlp" for i in range(NUM_LAYERS)], + moe_layer_index=None, + ) + + @contextmanager + def fake_context(_vllm_config): + yield forward_context + + monkeypatch.setattr(module, "_ffn_forward_context", fake_context) + + runner = object.__new__(module.GPUFFNModelRunner) + runner.num_layers = NUM_LAYERS + runner.vllm_config = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=512), + ) + runner._compute_work_item = lambda item, states: torch.full((1,), item.layer_idx) + runner.forward_context = forward_context + return runner + + +def test_an_idle_poll_returns_instead_of_blocking(runner): + # This is what lets the worker loop observe its shutdown event. + runner.connector = _FakeConnector([]) + + assert runner._ffn_forward_connector_driven() is None + assert len(runner.connector.received) == 1 + + +def test_every_drained_item_is_computed_and_sent_back(runner): + items = [_FakeWorkItem(i, _states(), f"meta-{i}") for i in range(2)] + runner.connector = _FakeConnector(items) + + runner._ffn_forward_connector_driven() + + assert [layer for layer, _ in runner.connector.sent] == [0, 1] + assert [int(out) for _, out in runner.connector.sent] == [0, 1] + + +def test_each_item_installs_its_own_metadata_and_layer(runner): + # Successive work items may belong to different layers of different + # replicas, so per-item context installation is the whole contract. + runner.connector = _FakeConnector( + [_FakeWorkItem(2, _states(), "meta-2")], + ) + + runner._ffn_forward_connector_driven() + + assert runner.forward_context.additional_kwargs["afd_metadata"] == "meta-2" + assert runner.forward_context.moe_layer_index == 2 + # A control-plane run leaves dp_metadata behind; this path must clear it. + assert runner.forward_context.dp_metadata is None + + +def test_the_drain_is_bounded_by_the_layer_count(runner): + endless = [_FakeWorkItem(0, _states(), "m") for _ in range(NUM_LAYERS + 5)] + runner.connector = _FakeConnector(endless) + + runner._ffn_forward_connector_driven() + + # Returning after a bounded drain is what gives the worker loop its turn. + assert len(runner.connector.sent) == NUM_LAYERS + + +def test_a_foreign_transfer_state_is_refused(runner): + runner.connector = _FakeConnector( + [_FakeWorkItem(0, SimpleNamespace(), "meta")], + ) + + with pytest.raises(RuntimeError, match="GpuAsyncTransferState"): + runner._ffn_forward_connector_driven() + + +def test_a_peer_shutdown_propagates(runner): + # The worker loop treats this as an ordinary exit; the loop must not + # swallow it into an idle-poll return. + runner.connector = _FakeConnector([ConnectorShutdown("peer left")]) + + with pytest.raises(ConnectorShutdown): + runner._ffn_forward_connector_driven()