[WIP] feat(npu): DeepSeek V4 over the synchronous CAMP2P boundary

afd_plugin/compat/npu/feature_validation.py

@@ -8,6 +8,7 @@
 
 from afd_plugin.config import (
     AFD_ASYNC_CONNECTOR,
+    CAMP2P_CONNECTOR,
     AFDConfig,
     is_afd_async_dp,
     parse_afd_config,
@@ -130,8 +131,18 @@ def _fail_if_unsupported_dsv4_async_features(
 
 
 def _fail_if_unsupported_dsv4_connector(afd_config: AFDConfig) -> None:
-    if afd_config.connector != AFD_ASYNC_CONNECTOR:
-        raise RuntimeError("DSV4 NPU AFD supports only CAMAsyncAFDConnector")
+    """Reject DSV4 on connectors that cannot carry token ids.
+
+    A DSV4 Hash layer routes by token identity, so whichever role owns the gate
+    needs the tokens' ids. CAM async reads them from its dispatch metadata;
+    CAMP2P moves them over the A2E ids channel with the gate left on FFN.
+    """
+
+    if afd_config.connector not in (AFD_ASYNC_CONNECTOR, CAMP2P_CONNECTOR):
+        raise RuntimeError(
+            "DSV4 NPU AFD supports only CAMAsyncAFDConnector and "
+            f"CAMP2pAFDConnector; got {afd_config.connector!r}",
+        )
 
 
 def _fail_if_unsupported_npu_afd_async_features(

afd_plugin/compat/npu/forward_context.py

@@ -27,8 +27,15 @@ def ascend_forward_context(
     in_profile_run: bool = False,
     aclgraph_runtime_mode: CUDAGraphMode | None = None,
     skip_mc2_mask: bool = False,
+    input_ids: torch.Tensor | None = None,
 ) -> Iterator[ForwardContext]:
-    """Create the minimal forward context needed by connector-driven FFN steps."""
+    """Create the minimal forward context needed by connector-driven FFN steps.
+
+    ``input_ids`` carries the token ids of exactly the tokens this FFN rank
+    computes on. Models whose router is keyed by token identity read it from the
+    forward context, so a connector that transports ids must install them here
+    before running the FFN compute.
+    """
 
     from vllm.config import CUDAGraphMode
     from vllm.forward_context import get_forward_context
@@ -60,6 +67,8 @@ def ascend_forward_context(
             forward_context = get_forward_context()
             forward_context.additional_kwargs["afd_metadata"] = afd_metadata
             forward_context.additional_kwargs["model_instance"] = model_instance
+            if input_ids is not None:
+                forward_context.input_ids = input_ids
             yield forward_context
         return
 
@@ -81,6 +90,8 @@ def ascend_forward_context(
         if forward_context.additional_kwargs is None:
             forward_context.additional_kwargs = {}
         forward_context.additional_kwargs["afd_metadata"] = afd_metadata
+        if input_ids is not None:
+            forward_context.input_ids = input_ids
         yield forward_context
 
 

afd_plugin/config.py

@@ -17,12 +17,13 @@
 
 AFD_ADDITIONAL_CONFIG_KEY: Final[str] = "afd"
 AFD_ASYNC_CONNECTOR: Final[str] = "CAMAsyncAFDConnector"
+CAMP2P_CONNECTOR: Final[str] = "CAMP2pAFDConnector"
 AFDRole = Literal["attention", "ffn"]
 
 SUPPORTED_AFD_ROLES: Final[tuple[str, ...]] = ("attention", "ffn")
 SUPPORTED_AFD_CONNECTORS: Final[tuple[str, ...]] = (
     "P2pNcclAFDConnector",
-    "CAMP2pAFDConnector",
+    CAMP2P_CONNECTOR,
     AFD_ASYNC_CONNECTOR,
 )
 
@@ -335,6 +336,7 @@ def validate_afd_config(
     "afd_config_from_mapping",
     "AFD_ADDITIONAL_CONFIG_KEY",
     "AFDRole",
+    "CAMP2P_CONNECTOR",
     "SUPPORTED_AFD_CONNECTORS",
     "SUPPORTED_AFD_ROLES",
     "connector_extra_config_from_mapping",

afd_plugin/connectors/npu/camp2p.py

@@ -178,6 +178,13 @@ class CAMP2PTransferState(AFDTransferState):
     A2E-returned Attention token count that the FFN-to-Attention send requires.
     ``x_active_mask`` and ``cam_p2p_ep_name`` are the A2E-returned active-token
     mask and HCCL endpoint name captured on the receive path.
+
+    ``input_ids`` holds the token-aligned ids that Attention sent alongside the
+    hidden states, as received by the FFN rank. It is populated only when the
+    receiving rank declared ``recv_input_ids``, which ``compute_gate_mode``
+    records. ``compute_gate_mode`` is the operator's ids mode for this transfer
+    and has to equal the mode the sending rank selected, because the operator
+    only writes the ids slot in that mode.
     """
 
     aiv_num: int = 8
@@ -187,6 +194,8 @@ class CAMP2PTransferState(AFDTransferState):
     atten_batch_size: torch.Tensor | None = None
     x_active_mask: torch.Tensor | None = None
     cam_p2p_ep_name: str | None = None
+    input_ids: torch.Tensor | None = None
+    compute_gate_mode: int = 0
 
 
 @dataclass(frozen=True, slots=True)
@@ -224,6 +233,90 @@ def is_attn_top_min_size_rank(self) -> bool:
         return self.ffn_size <= self.world_rank < self.ffn_size + self.min_size
 
 
+def prepare_token_id_transfer(
+    input_ids: torch.Tensor,
+    *,
+    topk: int,
+    expected_tokens: int,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Pack token ids and inert scales for the A2E ids channel.
+
+    The operator transports an ``int32`` ids tensor and a ``float32`` scales
+    tensor of shape ``(batch, topk)``. The id repeats across the columns, which
+    is what the receiving side collapses back. The scales carry no routing
+    weight: this channel moves token identity, not router output.
+
+    Args:
+        input_ids: Token-aligned ids for the local Attention tokens.
+        topk: Number of routed experts per token; the operator's column count.
+        expected_tokens: Token count the transfer metadata declares.
+
+    Returns:
+        The ``(ids, scales)`` pair to hand to the operator.
+
+    Raises:
+        ValueError: If the ids do not describe exactly ``expected_tokens``
+            tokens, which is the alignment invariant for this channel.
+    """
+
+    num_tokens = int(input_ids.numel())
+    if num_tokens != expected_tokens:
+        raise ValueError(
+            f"input_ids token count {num_tokens} does not match the AFD "
+            f"transfer token count {expected_tokens}",
+        )
+
+    ids = (
+        input_ids.reshape(-1)
+        .to(dtype=torch.int32)
+        .unsqueeze(1)
+        .expand(-1, topk)
+        .contiguous()
+    )
+    scales = torch.zeros(
+        (num_tokens, topk),
+        dtype=torch.float32,
+        device=input_ids.device,
+    )
+    return ids, scales
+
+
+def received_token_ids(
+    sim_expert_ids: torch.Tensor,
+    *,
+    expected_tokens: int,
+) -> torch.Tensor:
+    """Collapse the A2E ids channel back to a token-aligned id vector.
+
+    The operator returns ``(N, topk)`` with every column of a row repeating that
+    token's id, and sizes ``N`` to the capacity it was given, so this trims to
+    the tokens the FFN rank actually computes on.
+
+    Args:
+        sim_expert_ids: The operator's ids output.
+        expected_tokens: Token count derived from the FFN rank's DP metadata,
+            which is what the FFN compute will actually run on.
+
+    Returns:
+        A one-dimensional ``int32`` tensor of length ``expected_tokens``.
+
+    Raises:
+        ValueError: If fewer ids arrived than the FFN rank computes on. That is
+            the misalignment this channel has to catch: a token-keyed router
+            would otherwise select experts for the wrong tokens.
+    """
+
+    received_tokens = int(sim_expert_ids.shape[0])
+    if received_tokens < expected_tokens:
+        raise ValueError(
+            f"received {received_tokens} token ids but the FFN rank computes on "
+            f"{expected_tokens} tokens; the ids channel is not aligned with the "
+            "FFN token layout",
+        )
+    # The columns are replicas of the same id, so the first one carries it.
+    return sim_expert_ids[:expected_tokens, 0].contiguous()
+
+
 class CAMP2pAFDConnector(AFDConnectorBase):
     """Move model data between Attention and FFN workers on Ascend NPU.
 
@@ -429,12 +522,16 @@ def send_attn_output(
             hidden_states: Model data with shape ``(tokens, hidden_size)``.
             context: Transfer context whose ``metadata`` supplies the layer
                 number, ubatch number, and token count for this transfer.
-            **kwargs: Extra arguments accepted for interface compatibility.
+            **kwargs: An optional token-aligned ``input_ids`` tensor. When it is
+                supplied, the transfer runs with ``compute_gate=1`` so the ids
+                reach the FFN rank through the operator's ids channel, and the
+                local FFN-side receive exposes them on ``CAMP2PTransferState``.
 
         Raises:
             RuntimeError: If the communication groups are not ready.
             ValueError: If the number of tokens in ``hidden_states`` does not
-                match ``context.metadata`` outside a ``torch.compile`` trace.
+                match ``context.metadata`` outside a ``torch.compile`` trace, or
+                if a supplied ``input_ids`` tensor is malformed.
         """
         if not self._initialized:
             raise RuntimeError("CAMP2P connector is not initialized")
@@ -446,6 +543,17 @@ def send_attn_output(
                 f"hidden_states shape {hidden_states.shape!r} does not match "
                 f"CAMP2P metadata token count {metadata.total_tokens}",
             )
+        input_ids = cast(torch.Tensor | None, kwargs.get("input_ids"))
+        expert_ids: torch.Tensor | None = None
+        expert_scales: torch.Tensor | None = None
+        compute_gate = 0
+        if input_ids is not None:
+            expert_ids, expert_scales = prepare_token_id_transfer(
+                input_ids,
+                topk=self.num_experts_per_tok,
+                expected_tokens=metadata.total_tokens,
+            )
+            compute_gate = 1
         transfer_state = CAMP2PTransferState(
             aiv_num=self.aiv_num,
             batch_size=metadata.total_tokens,
@@ -469,7 +577,9 @@ def send_attn_output(
             self.attn_size,
             self.world_rank,
             transfer_state.aiv_num,
-            0,
+            compute_gate,
+            expert_ids,
+            expert_scales,
         )
         return None
 
@@ -526,7 +636,12 @@ def recv_attn_output(
         Args:
             ubatch_idx: Ubatch number, starting from ``0``.
             **kwargs: May provide existing transfer information or the layer
-                number needed to create it.
+                number needed to create it. ``recv_input_ids`` states that this
+                FFN rank expects token ids on the transfer, which selects the
+                operator's ids mode and makes the connector validate and expose
+                the operator's ids slot. The sending rank has to select the same
+                mode, so only request ids for a run whose Attention role
+                transports them.
 
         Returns:
             The received hidden states and the information FFN needs to process
@@ -535,11 +650,14 @@ def recv_attn_output(
         Raises:
             RuntimeError: If communication is not ready, transfer information
                 is missing, or the requested ubatch group does not exist.
+            ValueError: If ids were requested but do not align with the FFN
+                rank's token layout.
         """
         if not self._initialized:
             raise RuntimeError("CAMP2P connector is not initialized")
         layer_idx: int = kwargs.get("layer_idx", 0)
         max_num_tokens: int = kwargs.get("max_num_tokens", 0)
+        recv_input_ids: bool = bool(kwargs.get("recv_input_ids", False))
         batch_size = _num_tokens_for_ffn_rank(
             self.dp_metadata_list,
             ubatch_idx,
@@ -558,6 +676,7 @@ def recv_attn_output(
             batch_size=batch_size,
             h=self.hidden_size,
             k=self.num_experts_per_tok,
+            compute_gate_mode=1 if recv_input_ids else 0,
         )
         context = AFDTransferContext(
             metadata=metadata,
@@ -582,11 +701,23 @@ def recv_attn_output(
             self.world_rank,
             group_ep,
             custom_states.aiv_num,
-            0,
+            custom_states.compute_gate_mode,
         )
         custom_states.atten_batch_size = outputs[3]
         custom_states.x_active_mask = outputs[4]
         custom_states.cam_p2p_ep_name = self.hccl_comm_name1
+        # The ids slot is only written in the operator's ids mode, so the mode has
+        # to match the sending rank's ``compute_gate``. It is declared here by the
+        # receiving FFN rank through ``recv_input_ids`` rather than inferred from
+        # the returned tensor: a genuine single-token layer would otherwise be
+        # indistinguishable from the operator's placeholder. Reading the slot in
+        # the other mode would hand the model uninitialised device memory as token
+        # ids, which a token-keyed router turns into an out-of-range table read.
+        if custom_states.compute_gate_mode == 1:
+            custom_states.input_ids = received_token_ids(
+                outputs[1],
+                expected_tokens=batch_size,
+            )
         return AFDA2FTransferPayload(
             hidden_states=outputs[0],
             context=context,
@@ -850,6 +981,8 @@ def send_attn_output_impl(
         world_rank: int,
         aiv_num: int,
         compute_gate: int,
+        expert_ids: torch.Tensor | None,
+        expert_scales: torch.Tensor | None,
     ) -> torch.Tensor:
         transfer_state = getattr(get_forward_context(), "cam_afdtransfer_state", None)
         if transfer_state is None:
@@ -867,8 +1000,8 @@ def send_attn_output_impl(
 
         outputs = torch.ops.afd_ascend.a2e(
             hidden_states,
-            None,
-            None,
+            expert_ids,
+            expert_scales,
             transfer_state.batch_size,
             transfer_state.h,
             transfer_state.k,
@@ -897,6 +1030,8 @@ def send_attn_output_fake_impl(
         world_rank: int,
         aiv_num: int,
         compute_gate: int,
+        expert_ids: torch.Tensor | None,
+        expert_scales: torch.Tensor | None,
     ) -> torch.Tensor:
         """Return the input unchanged while PyTorch inspects the send operation."""
         return hidden_states
@@ -970,6 +1105,8 @@ def recv_ffn_output_fake_impl(
         "world_rank": int,
         "aiv_num": int,
         "compute_gate": int,
+        "expert_ids": torch.Tensor | None,
+        "expert_scales": torch.Tensor | None,
         "return": torch.Tensor,
     }
     recv_annotations = {