Sync fork main validation changes

afd_plugin/connectors/__init__.py

@@ -5,6 +5,7 @@
 from afd_plugin.connectors.base import AFDConnectorBase
 from afd_plugin.connectors.factory import AFDConnectorFactory
 from afd_plugin.connectors.metadata import (
+    AFDConnectorData,
     AFDConnectorMetadata,
     AFDDPMetadata,
     AFDDPMetadataPayload,
@@ -15,6 +16,7 @@
 
 __all__ = [
     "AFDConnectorBase",
+    "AFDConnectorData",
     "AFDConnectorFactory",
     "AFDConnectorMetadata",
     "AFDDPMetadata",

afd_plugin/connectors/base.py

@@ -48,12 +48,6 @@ def init_afd_connector(self) -> None:
     def is_initialized(self) -> bool:
         raise NotImplementedError
 
-    def get_connector_rank(self) -> int:
-        return self.rank
-
-    def get_connector_local_rank(self) -> int:
-        return self.local_rank
-
     @abstractmethod
     def send_attn_output(
         self,
@@ -84,12 +78,10 @@ def send_ffn_output(
     ) -> None:
         raise NotImplementedError
 
+    @abstractmethod
     def update_state_from_dp_metadata(
         self,
-        dp_metadata_list: dict[int, Any],
-        *,
-        is_graph_capturing: bool = False,
-        is_warmup: bool = False,
+        payload: AFDDPMetadataPayload,
     ) -> None:
         raise NotImplementedError
 

afd_plugin/connectors/gpu/p2p.py

@@ -160,19 +160,16 @@ def is_initialized(self) -> bool:
 
     def update_state_from_dp_metadata(
         self,
-        dp_metadata_list: dict[int, Any],
-        *,
-        is_graph_capturing: bool = False,
-        is_warmup: bool = False,
+        payload: AFDDPMetadataPayload,
     ) -> None:
-        self.dp_metadata_list = dp_metadata_list
-        self.is_graph_capturing = is_graph_capturing
-        self.is_warmup = is_warmup
+        self.dp_metadata_list = payload.dp_metadata_list
+        self.is_graph_capturing = payload.is_graph_capturing
+        self.is_warmup = payload.is_warmup
         self._tensor_metadata_list = {}
         device = torch.device(f"cuda:{self.local_rank}")
         dtype = self.vllm_config.model_config.dtype
         dp_rank = int(self.vllm_config.parallel_config.data_parallel_rank)
-        for stage_idx, dp_metadata in dp_metadata_list.items():
+        for stage_idx, dp_metadata in payload.dp_metadata_list.items():
             token_count = dp_metadata.num_tokens_across_dp_cpu[dp_rank]
             item = getattr(token_count, "item", None)
             num_tokens = int(item() if callable(item) else token_count)
@@ -203,15 +200,14 @@ def update_state_from_dp_metadata(
                         device=tensor_metadata.device,
                     )
 
-    def is_attn_top_min_size_rank(self, world_rank: int) -> bool:
-        return self.ffn_size <= world_rank < self.ffn_size + self.min_size
-
     def send_dp_metadata_list(
         self,
         payload: AFDDPMetadataPayload,
     ) -> None:
         if self.p2p_pg is None:
             return
+        if not (self.ffn_size <= self.world_rank < self.ffn_size + self.min_size):
+            return
         device = torch.device(f"cuda:{self.local_rank}")
         object_bytes = _encode_dp_metadata_payload(payload)
         object_tensor_cpu = torch.frombuffer(bytearray(object_bytes), dtype=torch.uint8)

afd_plugin/connectors/metadata.py

@@ -173,15 +173,19 @@ def _compute_sp_num_tokens(
     return local_sizes
 
 
+@dataclass(slots=True)
+class AFDConnectorData:
+    """Base class for backend-specific connector metadata payloads."""
+
+
 @dataclass(slots=True)
 class AFDConnectorMetadata:
     """Communication metadata for one AFD Attention/FFN exchange."""
 
     layer_idx: int
     stage_idx: int
     seq_lens: list[int]
-    recv_handle_list: list[Any] | None = None
-    connector_data: Any = None
+    connector_data: AFDConnectorData | None = None
 
     def __post_init__(self) -> None:
         if not self.seq_lens:
@@ -268,6 +272,7 @@ def clone(self) -> AFDMetadata:
 
 
 __all__ = [
+    "AFDConnectorData",
     "AFDConnectorMetadata",
     "AFDDPMetadata",
     "AFDDPMetadataPayload",

afd_plugin/connectors/npu/camp2p.py

@@ -20,6 +20,7 @@
 from afd_plugin.config import AFDConfig
 from afd_plugin.connectors.base import AFDConnectorBase
 from afd_plugin.connectors.metadata import (
+    AFDConnectorData,
     AFDConnectorMetadata,
     AFDDPMetadataPayload,
     AFDRecvOutput,
@@ -34,7 +35,7 @@
 
 
 @dataclass(slots=True)
-class CAMP2PAFDConnectorMetadata:
+class CAMP2PAFDConnectorMetadata(AFDConnectorData):
     """CAMP2P payload metadata carried between recv and send phases.
 
     This mirrors ``vllm_ascend.distributed.metadata.CAMP2PAFDConnectorMetadata``
@@ -220,24 +221,20 @@ def close(self) -> None:
 
     def update_state_from_dp_metadata(
         self,
-        dp_metadata_list: dict[int, Any],
-        *,
-        is_graph_capturing: bool = False,
-        is_warmup: bool = False,
+        payload: AFDDPMetadataPayload,
     ) -> None:
-        self.dp_metadata_list = dict(dp_metadata_list)
-        self.is_graph_capturing = bool(is_graph_capturing)
-        self.is_warmup = bool(is_warmup)
-
-    def is_attn_top_min_size_rank(self, world_rank: int) -> bool:
-        return self.ffn_size <= int(world_rank) < self.ffn_size + self.min_size
+        self.dp_metadata_list = dict(payload.dp_metadata_list)
+        self.is_graph_capturing = payload.is_graph_capturing
+        self.is_warmup = payload.is_warmup
 
     def send_dp_metadata_list(
         self,
         payload: AFDDPMetadataPayload,
     ) -> None:
         if self.p2p_pg is None:
             return
+        if not self.topology.is_attn_top_min_size_rank:
+            return
         for dst in self.dst_list:
             _send_object(payload, dst=dst, group=self.p2p_pg)
 

afd_plugin/v1/worker/ascend/attention_model_runner.py

@@ -1036,31 +1036,21 @@ def _send_dp_metadata(self, dp_metadata: Any, ubatch_slices: Any) -> None:
             dp_metadata_list = {0: dp_metadata}
         is_warmup = bool(self._is_warmup)
         is_graph_capturing = bool(self._afd_is_graph_capturing)
-        self.afd_connector.update_state_from_dp_metadata(
-            dp_metadata_list,
+        payload = AFDDPMetadataPayload(
+            dp_metadata_list=dp_metadata_list,
             is_graph_capturing=is_graph_capturing,
             is_warmup=is_warmup,
         )
-        should_send = self.afd_connector.is_attn_top_min_size_rank(
-            self.afd_connector.world_rank,
-        )
+        self.afd_connector.update_state_from_dp_metadata(payload)
         logger.warning(
             "AFD NPU Attention send_dp_metadata decision; world_rank=%d "
-            "key=%s should_send=%s is_graph_capturing=%s is_warmup=%s",
+            "key=%s is_graph_capturing=%s is_warmup=%s",
             self.afd_connector.world_rank,
             _dp_metadata_debug_key(dp_metadata_list),
-            should_send,
             is_graph_capturing,
             is_warmup,
         )
-        if should_send:
-            self.afd_connector.send_dp_metadata_list(
-                AFDDPMetadataPayload(
-                    dp_metadata_list=dp_metadata_list,
-                    is_graph_capturing=is_graph_capturing,
-                    is_warmup=is_warmup,
-                ),
-            )
+        self.afd_connector.send_dp_metadata_list(payload)
 
     def _ensure_dp_metadata(self, dp_metadata: Any) -> Any:
         if dp_metadata is not None:

afd_plugin/v1/worker/ascend/ffn_model_runner.py

@@ -28,6 +28,7 @@
 from afd_plugin.connectors import (
     AFDConnectorFactory,
     AFDConnectorMetadata,
+    AFDDPMetadataPayload,
     AFDMetadata,
     AFDRecvOutput,
 )
@@ -182,8 +183,10 @@ def _ffn_forward(
     ) -> Any:
         if update_connector_state:
             self.connector.update_state_from_dp_metadata(
-                dp_metadata_list,
-                is_graph_capturing=is_graph_capturing,
+                _make_dp_metadata_payload(
+                    dp_metadata_list,
+                    is_graph_capturing=is_graph_capturing,
+                ),
             )
         num_stages = max(len(dp_metadata_list), 1)
         afd_metadata = AFDMetadata(
@@ -240,11 +243,6 @@ def _ffn_forward(
                         )
                         _set_moe_layer_index(forward_context, layer_idx)
 
-                    if metadata.recv_handle_list is not None:
-                        for work in metadata.recv_handle_list:
-                            work.wait()
-                        metadata.recv_handle_list = None
-
                     rank_ffn_output = self._run_ffn_computation(
                         hidden_states=hidden_states,
                         layer_idx=layer_idx,
@@ -328,8 +326,10 @@ def _capture_graphs(
         graph = torch.npu.NPUGraph()
         logger.debug("AFD NPU FFN created NPUGraph for key=%s", graph_key)
         self.connector.update_state_from_dp_metadata(
-            dp_metadata_list,
-            is_graph_capturing=is_attn_graph_capturing,
+            _make_dp_metadata_payload(
+                dp_metadata_list,
+                is_graph_capturing=is_attn_graph_capturing,
+            ),
         )
         logger.debug("AFD NPU FFN updated connector state for key=%s", graph_key)
         with torch.npu.graph(graph, pool=self.graph_pool):
@@ -439,6 +439,19 @@ def _normalize_recv_output(
     return hidden_states, metadata, recv_output
 
 
+def _make_dp_metadata_payload(
+    dp_metadata_list: dict[int, Any],
+    *,
+    is_graph_capturing: bool = False,
+    is_warmup: bool = False,
+) -> AFDDPMetadataPayload:
+    return AFDDPMetadataPayload(
+        dp_metadata_list=dp_metadata_list,
+        is_graph_capturing=is_graph_capturing,
+        is_warmup=is_warmup,
+    )
+
+
 def _ffn_token_counts_across_ranks(
     connector: Any,
     dp_metadata_list: dict[int, Any],

afd_plugin/v1/worker/attention_model_runner.py

@@ -120,24 +120,13 @@ def _send_dp_metadata(self, dp_metadata: Any, ubatch_slices: Any) -> None:
             dp_metadata_list = {0: dp_metadata}
         is_warmup = bool(self._is_warmup)
         is_graph_capturing = bool(getattr(self, "_afd_is_graph_capturing", False))
-        self.afd_connector.update_state_from_dp_metadata(
-            dp_metadata_list,
+        payload = AFDDPMetadataPayload(
+            dp_metadata_list=dp_metadata_list,
             is_graph_capturing=is_graph_capturing,
             is_warmup=is_warmup,
         )
-
-        should_send = True
-        rank = self.afd_connector.world_rank
-        should_send = bool(self.afd_connector.is_attn_top_min_size_rank(rank))
-
-        if should_send:
-            self.afd_connector.send_dp_metadata_list(
-                AFDDPMetadataPayload(
-                    dp_metadata_list=dp_metadata_list,
-                    is_graph_capturing=is_graph_capturing,
-                    is_warmup=is_warmup,
-                ),
-            )
+        self.afd_connector.update_state_from_dp_metadata(payload)
+        self.afd_connector.send_dp_metadata_list(payload)
 
     def load_model(self, *args: Any, **kwargs: Any) -> Any:
         use_ubatching = bool(self.vllm_config.parallel_config.use_ubatching)

afd_plugin/v1/worker/ffn_model_runner.py

@@ -26,6 +26,7 @@
 from afd_plugin.connectors import (
     AFDConnectorFactory,
     AFDConnectorMetadata,
+    AFDDPMetadataPayload,
     AFDRecvOutput,
 )
 from afd_plugin.v1.worker.attention_model_runner import (
@@ -153,8 +154,10 @@ def _ffn_forward(
     ) -> Any:
         if update_connector_state:
             self.connector.update_state_from_dp_metadata(
-                dp_metadata_list,
-                is_graph_capturing=is_graph_capturing,
+                _make_dp_metadata_payload(
+                    dp_metadata_list,
+                    is_graph_capturing=is_graph_capturing,
+                ),
             )
 
         rank_ffn_output = None
@@ -177,11 +180,6 @@ def _ffn_forward(
                         )
                         forward_context.additional_kwargs["afd_metadata"] = metadata
                         _set_moe_layer_index(forward_context, layer_idx)
-                    recv_handle_list = metadata.recv_handle_list
-                    if recv_handle_list is not None:
-                        for work in recv_handle_list:
-                            work.wait()
-                        metadata.recv_handle_list = None
                     rank_ffn_output = self._execute_eager_mode(hidden_states, layer_idx)
                     self.connector.send_ffn_output(rank_ffn_output, metadata)
         return rank_ffn_output
@@ -228,8 +226,10 @@ def _dummy_run(
             # DP metadata receive/update is a control-plane side effect and must
             # complete before CUDA graph capture starts.
             self.connector.update_state_from_dp_metadata(
-                dp_metadata_list,
-                is_graph_capturing=is_attn_graph_capturing,
+                _make_dp_metadata_payload(
+                    dp_metadata_list,
+                    is_graph_capturing=is_attn_graph_capturing,
+                ),
             )
             with torch.cuda.graph(cudagraph, pool=self._graph_memory_pool):
                 output = self._ffn_forward(
@@ -268,8 +268,11 @@ def capture_model(
             with graph_capture(device=self.device):
                 if is_warmup:
                     self.connector.update_state_from_dp_metadata(
-                        dp_metadata_list,
-                        is_graph_capturing=False,
+                        _make_dp_metadata_payload(
+                            dp_metadata_list,
+                            is_graph_capturing=False,
+                            is_warmup=True,
+                        ),
                     )
                     self._ffn_forward(
                         dp_metadata_list=dp_metadata_list,
@@ -343,6 +346,19 @@ def _set_moe_layer_index(forward_context: object, layer_idx: int) -> None:
             return
 
 
+def _make_dp_metadata_payload(
+    dp_metadata_list: dict[int, Any],
+    *,
+    is_graph_capturing: bool = False,
+    is_warmup: bool = False,
+) -> AFDDPMetadataPayload:
+    return AFDDPMetadataPayload(
+        dp_metadata_list=dp_metadata_list,
+        is_graph_capturing=is_graph_capturing,
+        is_warmup=is_warmup,
+    )
+
+
 def _normalize_recv_output(
     recv_output: Any,
     *,

tests/unit/connectors/test_base_factory.py

@@ -111,6 +111,9 @@ def recv_attn_output(self, ubatch_idx=None):
     def send_ffn_output(self, ffn_output, metadata):
         return None
 
+    def update_state_from_dp_metadata(self, payload):
+        return None
+
     def send_dp_metadata_list(self, payload):
         return None