[3/5] feat: run DeepSeek-V4 on the async GPU connector

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",

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"]

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"]

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 (