[Feat]: support multistream for CAMP2p

afd_plugin/connectors/npu/camp2p.py

@@ -51,6 +51,9 @@
     send_control_payload,
 )
 from afd_plugin.distributed import init_afd_process_group, topology_from_config
+from afd_plugin.v1.worker.npu.multistream import (
+    npu_stream_switch_within_graph,
+)
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -63,6 +66,8 @@
         "core_num",
         "attn_core_num",
         "ffn_core_num",
+        "is_attn_multistream",
+        "is_ffn_multistream",
         "compute_gate_on_attention",
         "quant_mode",
     },
@@ -77,13 +82,19 @@ class CAMP2PExtraInfo(ConnectorExtraInfo):
         core_num: Default number of AIV cores used by each AFD role.
         attn_core_num: Optional Attention-role override for ``core_num``.
         ffn_core_num: Optional FFN-role override for ``core_num``.
+        is_attn_multistream: Run Attention-side A2F communication on a
+            dedicated NPU stream.
+        is_ffn_multistream: Run FFN-side F2A communication on a dedicated NPU
+            stream.
         compute_gate_on_attention: Whether Attention computes MoE gate outputs.
         quant_mode: CAM quantization mode; the current runtime supports only 0.
     """
 
     core_num: int = 8
     attn_core_num: int | None = None
     ffn_core_num: int | None = None
+    is_attn_multistream: bool = False
+    is_ffn_multistream: bool = False
     compute_gate_on_attention: bool = False
     quant_mode: int = 0
 
@@ -117,6 +128,14 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> CAMP2PExtraInfo:
                 raw.get("ffn_core_num"),
                 field_name="ffn_core_num",
             ),
+            is_attn_multistream=coerce_extra_bool(
+                raw.get("is_attn_multistream", False),
+                field_name="is_attn_multistream",
+            ),
+            is_ffn_multistream=coerce_extra_bool(
+                raw.get("is_ffn_multistream", False),
+                field_name="is_ffn_multistream",
+            ),
             compute_gate_on_attention=coerce_extra_bool(
                 raw.get("compute_gate_on_attention", False),
                 field_name="compute_gate_on_attention",
@@ -145,6 +164,8 @@ def validate_supported(self) -> None:
     def to_mapping(self) -> dict[str, Any]:
         result: dict[str, Any] = {
             "core_num": self.core_num,
+            "is_attn_multistream": self.is_attn_multistream,
+            "is_ffn_multistream": self.is_ffn_multistream,
             "compute_gate_on_attention": self.compute_gate_on_attention,
             "quant_mode": self.quant_mode,
         }
@@ -595,24 +616,35 @@ def send_ffn_output(
         if states.atten_batch_size is None:
             raise RuntimeError("CAMP2P FFN side is missing A2E atten_batch_size")
         ubatch_idx = int(kwargs.get("ubatch_idx", context.metadata.stage_idx))
+        multistream_enable = bool(kwargs.get("multistream_enable", False))
+        comm_stream = kwargs.get("comm_stream")
+        comm_event = kwargs.get("comm_event")
         group_ep = _get_group_ep(
             ubatch_idx,
             self.hccl_comm_name,
             self.hccl_comm_name2,
             self.hccl_comm_name3,
         )
-        torch.ops.afd_ascend.e2a(
-            ffn_output,
-            states.atten_batch_size,
-            states.batch_size,
-            states.h,
-            states.k,
-            self.ffn_size,
-            self.attn_size,
-            self.world_rank,
-            group_ep,
-            states.aiv_num,
-        )
+        current_stream = torch.npu.current_stream()
+        with npu_stream_switch_within_graph(
+            current_stream,
+            comm_stream,
+            multistream_enable,
+        ):
+            torch.ops.afd_ascend.e2a(
+                ffn_output,
+                states.atten_batch_size,
+                states.batch_size,
+                states.h,
+                states.k,
+                self.ffn_size,
+                self.attn_size,
+                self.world_rank,
+                group_ep,
+                states.aiv_num,
+            )
+            if multistream_enable and comm_event is not None:
+                comm_event.record(comm_stream)
         return None
 
 
@@ -845,20 +877,34 @@ def send_attn_output_impl(
             hccl_comm_name3,
         )
 
-        outputs = torch.ops.afd_ascend.a2e(
-            hidden_states,
-            None,
-            None,
-            transfer_state.batch_size,
-            transfer_state.h,
-            transfer_state.k,
-            ffn_size,
-            attn_size,
-            world_rank,
-            group_ep,
-            transfer_state.aiv_num,
-            compute_gate,
+        forward_context = get_forward_context()
+        multistream_enable = bool(
+            getattr(forward_context, "afd_multistream_enabled", False)
         )
+        comm_stream = getattr(forward_context, "afd_comm_stream", None)
+        comm_event = getattr(forward_context, "afd_comm_event", None)
+        current_stream = torch.npu.current_stream()
+        with npu_stream_switch_within_graph(
+            current_stream,
+            comm_stream,
+            multistream_enable,
+        ):
+            outputs = torch.ops.afd_ascend.a2e(
+                hidden_states,
+                None,
+                None,
+                transfer_state.batch_size,
+                transfer_state.h,
+                transfer_state.k,
+                ffn_size,
+                attn_size,
+                world_rank,
+                group_ep,
+                transfer_state.aiv_num,
+                compute_gate,
+            )
+            if multistream_enable and comm_event is not None:
+                comm_event.record(comm_stream)
         transfer_state.atten_batch_size = outputs[3]
         forward_context = get_forward_context()
         forward_context.cam_afdtransfer_state = transfer_state
@@ -907,6 +953,12 @@ def recv_ffn_output_impl(
             hccl_comm_name2,
             hccl_comm_name3,
         )
+        forward_context = get_forward_context()
+        if bool(getattr(forward_context, "afd_multistream_enabled", False)):
+            comm_event = getattr(forward_context, "afd_comm_event", None)
+            if comm_event is None:
+                raise RuntimeError("CAMP2P Attention multistream requires an event")
+            comm_event.wait(torch.npu.current_stream())
         output = torch.ops.afd_ascend.e2a(
             ref_tensor,
             transfer_state.atten_batch_size,

afd_plugin/v1/worker/npu/attention_model_runner.py

@@ -161,6 +161,12 @@ def _model_forward(self, *args: Any, **kwargs: Any) -> Any:
         self._install_afd_metadata_on_forward_context(forward_context)
         self._install_async_moe_ubatch_metadata_on_forward_context(forward_context)
 
+        # Keep all attention DP ranks aligned before entering model forward.
+        # In graph mode this is immediately before graph replay, preventing a
+        # late attention rank from surfacing as A2E/Dispatch waits on FFN.
+        if self.afd_config.connector == "CAMP2pAFDConnector":
+            dist.barrier(group=get_dp_group().cpu_group)
+
         (
             num_tokens_padded,
             input_ids,
@@ -180,20 +186,29 @@ def _model_forward(self, *args: Any, **kwargs: Any) -> Any:
         }
         run_model = partial(self.model, **model_inputs)
 
+        # The Ascend ubatch wrapper captures and replays all stage forwards in
+        # one NPUGraph. Its stage-local contexts intentionally use
+        # CUDAGraphMode.NONE, so they do not register entries in Ascend's
+        # standard full-graph parameter tables. Updating the outer context would
+        # therefore treat the per-stage metadata list as a single-batch dict.
+        update_standard_full_graph = forward_context.ubatch_slices is None
+
         if self.enable_enpu:
-            self._update_full_graph_params_if_needed(
-                forward_context,
-                num_tokens_padded,
-                positions,
-            )
+            if update_standard_full_graph:
+                self._update_full_graph_params_if_needed(
+                    forward_context,
+                    num_tokens_padded,
+                    positions,
+                )
             hidden_states = run_model()
         else:
             hidden_states = run_model()
-            self._update_full_graph_params_if_needed(
-                forward_context,
-                num_tokens_padded,
-                positions,
-            )
+            if update_standard_full_graph:
+                self._update_full_graph_params_if_needed(
+                    forward_context,
+                    num_tokens_padded,
+                    positions,
+                )
 
         if (
             forward_context.flash_comm_v1_enabled

afd_plugin/v1/worker/npu/ffn_model_runner.py

@@ -32,6 +32,7 @@
     AFDForwardContextMetadata,
     AFDTransferContext,
 )
+from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo
 from afd_plugin.v1.worker.attention_model_runner import (
     _resolve_world_ranks,
     _with_dp_derived_afd_rank,
@@ -76,6 +77,20 @@ def __init__(self, vllm_config: VllmConfig, device: object) -> None:
             self.afd_config,
         )
         self.num_layers = int(self.model_config.hf_config.num_hidden_layers)
+        connector_extra_info = self.connector.extra_info
+        self.ffn_multistream_enabled = (
+            isinstance(connector_extra_info, CAMP2PExtraInfo)
+            and connector_extra_info.is_ffn_multistream
+        )
+        configured_ubatches = self.parallel_config.num_ubatches or 1
+        self.ffn_comm_stream = (
+            torch.npu.Stream(device=device) if self.ffn_multistream_enabled else None
+        )
+        self.ffn_comm_events = (
+            [torch.npu.Event() for _ in range(configured_ubatches)]
+            if self.ffn_multistream_enabled
+            else []
+        )
         self.use_aclgraph = _use_npu_aclgraph(vllm_config, self)
         self._acl_graphs: dict[tuple, dict[str, Any]] = {}
         self.graph_pool = (
@@ -223,6 +238,11 @@ def _ffn_forward(
         )
         num_tokens = _ffn_token_count_for_rank(self.connector, num_tokens_across_dp)
         rank_ffn_output = None
+        multistream_enabled = self.ffn_multistream_enabled and num_stages > 1
+        required_event_count = max(stage_ids) + 1
+        while multistream_enabled and len(self.ffn_comm_events) < required_event_count:
+            self.ffn_comm_events.append(torch.npu.Event())
+        event_recorded = [False] * required_event_count
 
         # Build DP-level token counts for vLLM's forward context.
         # num_tokens_across_dp has ffn_size entries (AFD-level, one per
@@ -242,7 +262,12 @@ def _ffn_forward(
             aclgraph_runtime_mode=aclgraph_runtime_mode,
         ) as forward_context:
             for layer_idx in _ffn_layer_indices(self):
+                layer_multistream = multistream_enabled and layer_idx > 0
                 for stage_idx in stage_ids:
+                    if multistream_enabled and event_recorded[stage_idx]:
+                        self.ffn_comm_events[stage_idx].wait(
+                            torch.npu.current_stream(),
+                        )
                     payload = self.connector.recv_attn_output(
                         ubatch_idx=stage_idx,
                         layer_idx=layer_idx,
@@ -268,7 +293,20 @@ def _ffn_forward(
                         rank_ffn_output,
                         context,
                         stage_idx=stage_idx,
+                        multistream_enable=layer_multistream,
+                        comm_stream=self.ffn_comm_stream,
+                        comm_event=self.ffn_comm_events[stage_idx]
+                        if layer_multistream
+                        else None,
                     )
+                    if layer_multistream:
+                        event_recorded[stage_idx] = True
+
+            if multistream_enabled:
+                current_stream = torch.npu.current_stream()
+                for stage_idx in stage_ids:
+                    if event_recorded[stage_idx]:
+                        self.ffn_comm_events[stage_idx].wait(current_stream)
         return rank_ffn_output
 
     def _ffn_forward_connector_driven(self) -> Any:
@@ -435,20 +473,27 @@ def _send_ffn_output(
     context: AFDTransferContext,
     *,
     stage_idx: int,
+    multistream_enable: bool = False,
+    comm_stream=None,
+    comm_event=None,
 ) -> None:
-    if not isinstance(ffn_output, AFDF2ATransferPayload):
-        connector.send_ffn_output(
-            ffn_output,
-            context,
-            ubatch_idx=stage_idx,
+    kwargs: dict[str, object] = {"ubatch_idx": stage_idx}
+    if multistream_enable:
+        kwargs.update(
+            multistream_enable=True,
+            comm_stream=comm_stream,
+            comm_event=comm_event,
         )
-        return
 
-    kwargs: dict[str, object] = {"ubatch_idx": stage_idx}
-    if ffn_output.shared_output is not None:
-        kwargs["expand_x_shared"] = ffn_output.shared_output
+    if isinstance(ffn_output, AFDF2ATransferPayload):
+        output_to_send = ffn_output.routed_output
+        if ffn_output.shared_output is not None:
+            kwargs["expand_x_shared"] = ffn_output.shared_output
+    else:
+        output_to_send = ffn_output
+
     connector.send_ffn_output(
-        ffn_output.routed_output,
+        output_to_send,
         context,
         **kwargs,
     )

afd_plugin/v1/worker/npu/forward_context.py

@@ -27,6 +27,9 @@ def create_ascend_forward_context(
     cudagraph_runtime_mode: CUDAGraphMode | None = None,
     batch_descriptor: BatchDescriptor | None = None,
     skip_compiled: bool = False,
+    afd_comm_stream=None,
+    afd_comm_event=None,
+    afd_multistream_enabled: bool = False,
 ) -> ForwardContext:
     if cudagraph_runtime_mode is None:
         cudagraph_runtime_mode = CUDAGraphMode.NONE
@@ -130,6 +133,9 @@ def create_ascend_forward_context(
         new_forward_context.mc2_mask = mc2_mask
 
     new_forward_context.dbo_enabled = True
+    new_forward_context.afd_comm_stream = afd_comm_stream
+    new_forward_context.afd_comm_event = afd_comm_event
+    new_forward_context.afd_multistream_enabled = afd_multistream_enabled
     return new_forward_context
 
 

afd_plugin/v1/worker/npu/multistream.py

@@ -0,0 +1,32 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""NPU stream helpers used by CAMP2P communication."""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, nullcontext
+
+import torch
+
+
+def npu_stream_switch_within_graph(
+    current_stream: torch.npu.Stream | None,
+    target_stream: torch.npu.Stream | None,
+    enabled: bool,
+) -> AbstractContextManager[None]:
+    """Switch to ``target_stream`` after ``current_stream`` when enabled.
+
+    Return a no-op context when multi-stream execution is disabled. Both
+    streams are required when it is enabled.
+    """
+    if not enabled:
+        return nullcontext()
+    if current_stream is None or target_stream is None:
+        raise RuntimeError(
+            "CAMP2P multistream requires compute and communication streams",
+        )
+    target_stream.wait_stream(current_stream)
+    return torch.npu.stream(target_stream)
+
+
+__all__ = ["npu_stream_switch_within_graph"]