diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index 6a7e82706..2766818a7 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 000000000..b3c9ea8fd --- /dev/null +++ b/afd_plugin/compat/patches/dp_coordinator_timeout.py @@ -0,0 +1,81 @@ +# 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 +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 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)), + ) + 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 000000000..2b23126ea --- /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 21e4f86b1..235d177d0 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 cfd154307..187070993 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,32 @@ 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.""" + # 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, + ) + 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 +680,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 +691,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 f633dc595..756ebb3e9 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 b559be68e..2dd1f32f1 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 4582d229f..08f86ab0e 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 4800eab98..66784a4bd 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/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index 5d7dd1d84..d8d89f044 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/execution_platforms.md b/docs/design/module/execution_platforms.md index 605c6a9db..38a069b1c 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/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index 189de1beb..778599175 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 26b3c77ec..dfe19ea0e 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/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh b/recipe/gpu/GpuAsyncAFDConnector/deepseek_v4_flash/2a2f_async.sh new file mode 100755 index 000000000..a426b2d1f --- /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 000000000..ade2e03f5 --- /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:-/path/model_weights/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/compat/patches/test_dp_coordinator_timeout.py b/tests/unit/compat/patches/test_dp_coordinator_timeout.py new file mode 100644 index 000000000..07931130a --- /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 000000000..495f36f0e --- /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 000000000..42afca5ec --- /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/model_executor/models/test_deepseek_v4_weight_policy.py b/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py index fa39bd24e..092d1cd22 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 136aee083..6e9ad29f4 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 5f9401284..7a8cbe03a 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_connector_driven_loop.py b/tests/unit/v1/worker/test_ffn_connector_driven_loop.py new file mode 100644 index 000000000..1b53a449c --- /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() diff --git a/tests/unit/v1/worker/test_ffn_model_runner.py b/tests/unit/v1/worker/test_ffn_model_runner.py index c0d198d04..b7c1eff72 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):