refactor: decouple shared FFN components from model adapters

afd_plugin/compat/patches/gpu_deepseek_v4_dbo.py

@@ -42,7 +42,7 @@
 from vllm.v1.worker.utils import AttentionGroup
 from vllm.v1.worker.workspace import WorkspaceManager, current_workspace_manager
 
-from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context
+from afd_plugin.forward_context import get_afd_metadata_from_forward_context
 
 if TYPE_CHECKING:
     from afd_plugin.v1.worker.attention_model_runner_v2 import AFDAttentionModelRunnerV2

afd_plugin/envs.py

@@ -7,14 +7,21 @@
 import os
 
 AFD_FORCE_BALANCED_TOPK_IDS = "AFD_FORCE_BALANCED_TOPK_IDS"
+AFD_ASYNC_MOE_LAYOUT_LOG = "AFD_ASYNC_MOE_LAYOUT_LOG"
 ENV_TRUE_VALUES = frozenset({"1", "true", "yes", "on"})
 
 
 def force_balanced_topk_ids_enabled() -> bool:
     return os.environ.get(AFD_FORCE_BALANCED_TOPK_IDS, "").lower() in ENV_TRUE_VALUES
 
 
+def async_moe_layout_log_enabled() -> bool:
+    return os.environ.get(AFD_ASYNC_MOE_LAYOUT_LOG, "").lower() in ENV_TRUE_VALUES
+
+
 __all__ = [
+    "AFD_ASYNC_MOE_LAYOUT_LOG",
     "AFD_FORCE_BALANCED_TOPK_IDS",
+    "async_moe_layout_log_enabled",
     "force_balanced_topk_ids_enabled",
 ]

afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py

@@ -24,9 +24,9 @@
     V4CheckpointDescription,
 )
 from afd_plugin.ffn_faas.platforms.npu.settings import validate_static_experts
-from afd_plugin.model_executor.npu.async_cam_layered import (
-    AsyncCAMLayeredExecutor,
-    LayeredFFNWeights,
+from afd_plugin.model_executor.ffn.npu.expert_weights import (
+    QuantizedExpertWeights,
+    validate_expert_weights,
 )
 
 
@@ -158,7 +158,7 @@ def routed_experts(self, layer_index: int) -> RoutedExperts:
         """Keep the native V4 module path inside the model adapter."""
         return self.model.layers[str(layer_index)].mlp.experts.routed_experts
 
-    def weight_views(self) -> tuple[LayeredFFNWeights, ...]:
+    def weight_views(self) -> tuple[QuantizedExpertWeights, ...]:
         from vllm_ascend.quantization.quant_type import QuantType
 
         weights = []
@@ -174,7 +174,7 @@ def weight_views(self) -> tuple[LayeredFFNWeights, ...]:
             scheme = routed.quant_method.quant_method
             if scheme.quant_type != QuantType.W4A8:
                 raise ValueError("V4 FFN requires native W4A8 routed parameters")
-            view = LayeredFFNWeights(
+            view = QuantizedExpertWeights(
                 layer_idx=layer.layer_idx,
                 w13=routed.w13_weight,
                 w2=routed.w2_weight,
@@ -187,7 +187,7 @@ def weight_views(self) -> tuple[LayeredFFNWeights, ...]:
                 routed_scaling_factor=1.0,
                 quantization="W4A8",
             )
-            AsyncCAMLayeredExecutor._validate_layer(view)
+            validate_expert_weights([view])
             weights.append(view)
         if not weights or any(w.signature() != weights[0].signature() for w in weights):
             raise ValueError("V4 FFN requires homogeneous native W4A8 layers")

afd_plugin/ffn_faas/platforms/npu/model.py

@@ -22,7 +22,7 @@
     from vllm.config import VllmConfig
     from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts
 
-    from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights
+    from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights
 
 
 class RoutedWeightOwner(Protocol):
@@ -32,7 +32,7 @@ def create_layer(self, layer_index: int) -> nn.Module: ...
     def routed_experts(self, layer_index: int) -> RoutedExperts: ...
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ...
     def finish_layer(self, layer_index: int) -> None: ...
-    def weight_views(self) -> tuple[LayeredFFNWeights, ...]: ...
+    def weight_views(self) -> tuple[QuantizedExpertWeights, ...]: ...
     def parameters(self) -> Iterator[nn.Parameter]: ...
 
 
@@ -65,7 +65,7 @@ def __init__(
         self.adapter, self.ffn_ranks, self.ffn_rank = adapter, ffn_ranks, ffn_rank
         self.local_experts = adapter.num_experts // ffn_ranks
         self._owner: RoutedWeightOwner | None = None
-        self._views: tuple[LayeredFFNWeights, ...] | None = None
+        self._views: tuple[QuantizedExpertWeights, ...] | None = None
         self._loading = False
         self._loaded = False
         self._parameter_bytes = adapter.estimated_parameter_bytes(self.local_experts)
@@ -79,7 +79,7 @@ def parameter_bytes(self) -> int:
         return self._parameter_bytes
 
     @property
-    def native_view(self) -> tuple[LayeredFFNWeights, ...]:
+    def native_view(self) -> tuple[QuantizedExpertWeights, ...]:
         if not self._loaded or self._views is None:
             raise RuntimeError("FFN weights have not been loaded")
         return self._views
@@ -177,7 +177,7 @@ def _load_native(
     def get_experts_layer_indices(self) -> tuple[int, ...]:
         return self.adapter.layer_indices
 
-    def get_async_cam_layered_weights(self) -> tuple[list[LayeredFFNWeights], str]:
+    def get_quantized_expert_weights(self) -> tuple[list[QuantizedExpertWeights], str]:
         return list(self.native_view), ""
 
     def compute_ffn_output(

afd_plugin/model_executor/models/forward_context.py → afd_plugin/forward_context.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
-"""Helpers for plugin-owned model wrappers to read AFD forward metadata."""
+"""AFD forward metadata installation and access shared by runners and models."""
 
 from __future__ import annotations
 
@@ -20,8 +20,8 @@ def get_afd_metadata_from_forward_context(
 ) -> AFDForwardContextMetadata | None:
     """Return AFD metadata from vLLM ``ForwardContext.additional_kwargs``.
 
-    Model wrappers use this helper so AFD metadata stays outside vLLM's
-    ``ForwardContext`` schema.
+    Models and remote FFN proxies read the metadata installed by runners
+    without extending vLLM's ``ForwardContext`` schema.
     """
 
     if forward_context is None:

afd_plugin/model_executor/ffn/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Shared remote FFN calls and backend-specific FFN implementations."""

afd_plugin/model_executor/ffn/npu/__init__.py

@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""NPU Async CAM layouts, quantized weights and native FFN compute."""

afd_plugin/model_executor/models/npu/async_cam_layout.py → afd_plugin/model_executor/ffn/npu/async_cam.py

@@ -1,20 +1,25 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
-"""Attention TP token-layout conversion at the Async CAM boundary."""
+"""Model-side Async CAM tensor layout and native DBO receive handoff."""
 
 from __future__ import annotations
 
-import os
+from collections.abc import Callable
 from dataclasses import dataclass
-from typing import Final
 
 import torch
 from vllm.distributed import tensor_model_parallel_all_gather
 from vllm.distributed.parallel_state import get_tp_group
 from vllm.logger import init_logger
+from vllm.v1.worker.ubatching import (
+    dbo_enabled,
+    dbo_maybe_run_recv_hook,
+    dbo_register_recv_hook,
+)
 
-ASYNC_MOE_LAYOUT_LOG_ENV: Final[str] = "AFD_ASYNC_MOE_LAYOUT_LOG"
-_TRUE_ENV_VALUES: Final[frozenset[str]] = frozenset({"1", "true", "yes", "on"})
+from afd_plugin.connectors import AFDConnectorBase, AFDTransferContext
+from afd_plugin.envs import AFD_ASYNC_MOE_LAYOUT_LOG as ASYNC_MOE_LAYOUT_LOG_ENV
+from afd_plugin.envs import async_moe_layout_log_enabled
 
 logger = init_logger(__name__)
 
@@ -184,7 +189,7 @@ def _log_cam_layout(
     phase: str,
     layout: CAMDispatchLayout,
 ) -> None:
-    if os.environ.get(ASYNC_MOE_LAYOUT_LOG_ENV, "").lower() not in _TRUE_ENV_VALUES:
+    if not async_moe_layout_log_enabled():
         return
     logger.warning(
         "AFD Async CAM layout; phase=%s sequence_parallel=%s "
@@ -223,10 +228,46 @@ def _pad_first_dim(
     )
 
 
+def send_async_cam_remote_call(
+    connector: AFDConnectorBase,
+    hidden_states: torch.Tensor,
+    context: AFDTransferContext,
+    **send_kwargs: torch.Tensor | None,
+) -> Callable[[], torch.Tensor]:
+    """Send once and let the peer submit raw receive before its next send.
+
+    The returned closure belongs to this call. Layout restoration, TP
+    collectives and shared/HC/scaling remain with the resumed owner.
+    """
+    overlapping = dbo_enabled()
+    if overlapping:
+        dbo_maybe_run_recv_hook()
+    connector.send_attn_output(hidden_states, context, **send_kwargs)
+    owner_stage = context.metadata.stage_idx
+    receive_started = False
+    output: torch.Tensor | None = None
+
+    def receive() -> torch.Tensor:
+        nonlocal receive_started, output
+        if not receive_started:
+            receive_started = True
+            output = connector.recv_ffn_output(
+                ref_tensor=hidden_states, ubatch_idx=owner_stage, context=context
+            )
+        if output is None:
+            raise RuntimeError("Async CAM receive failed; restart the A/F worker group")
+        return output
+
+    if overlapping:
+        dbo_register_recv_hook(receive)
+    return receive
+
+
 __all__ = [
     "ASYNC_MOE_LAYOUT_LOG_ENV",
     "CAMDispatchLayout",
     "CAMDispatchPayload",
     "prepare_cam_dispatch_payload",
     "restore_cam_dispatch_output",
+    "send_async_cam_remote_call",
 ]