dsv4 A5 NPU rebased

README.md

@@ -45,6 +45,7 @@ Model support:
 | Model family | Registered architectures | Plugin model wrappers | Notes |
 | --- | --- | --- | --- |
 | DeepSeekV2 / DeepSeekV3 / DeepSeekV3.2 | `DeepseekForCausalLM`, `DeepseekV2ForCausalLM`, `DeepseekV3ForCausalLM`, `DeepseekV32ForCausalLM` | `AFDDeepseekForCausalLM`, `AFDDeepseekV2ForCausalLM`, `AFDDeepseekV3ForCausalLM` | DeepSeekV3.2 uses `AFDDeepseekV3ForCausalLM`. Each AFD role constructs and loads only its role-required model components, while shared embedding, normalization, and output components remain available where required by the model lifecycle. |
+| DeepSeekV4 | `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` (CUDA), `AFDNPUDeepseekV4ForCausalLM` (Ascend) | CUDA uses NCCL P2P. The initial Ascend A5 path uses eager HCCL P2P and carries token IDs required by V4 hash routing; hardware E2E validation is still required. |
 | Qwen3 MoE | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | CUDA with `compute_gate_on_attention=false`. |
 | Qwen3.5 / Qwen3.6 MoE | `Qwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | Qwen3.5/Qwen3.6 adapter family. Repository CUDA E2E evidence currently covers text-only Qwen3.6-35B-A3B with `--language-model-only`, synchronous `P2pNcclAFDConnector`, native DP4/TP1/EP4 baseline, and AFD 2A1F eager/graph/graph+DBO. |
 
@@ -55,7 +56,7 @@ See the [recipe index](recipe/README.md) for deployment and benchmark examples.
 | Connector | Platform | Recommend Stage | Sync or Async | Graph Support | Notes |
 | --- | --- | --- | --- | --- | --- |
 | `P2pNcclAFDConnector` | CUDA | Decode | Sync | `FULL_DECODE_ONLY` CUDA graph | FFN ranks are ordered before Attention ranks. `num_attention_ranks` must be greater than or equal to `num_ffn_ranks` and divisible by it. See the [DeepSeek V2 Lite recipe](recipe/gpu/P2pNcclAFDConnector/deepseek_v2_lite/README.md). |
-| `CAMP2pAFDConnector` | Ascend NPU | Decode | Sync | `FULL_DECODE_ONLY` ACL graph | Uses HCCL/CAMP2P custom ops. Ascend ops build by default on NPU platforms. See the [synchronous DeepSeek V3.2 recipe](recipe/npu/CAMP2pAFDConnector/deepseek_v3_2/README.md). |
+| `CAMP2pAFDConnector` | Ascend NPU | Decode | Sync | A2/A3 custom-op graph; A5 V4 eager only | Uses HCCL/CAMP2P custom ops on established platforms. The A5 V4 route uses blocking HCCL P2P pending A2E/E2A operator support. See the [DeepSeek V3.2](recipe/npu/CAMP2pAFDConnector/deepseek_v3_2/README.md) and [A5 DeepSeek V4](recipe/npu/CAMP2pAFDConnector/deepseek_v4/README.md) recipes. |
 | `CAMAsyncAFDConnector` | Ascend NPU | Prefill / decode | Async | Not supported | Experimental v0.26 DP+TP/SP path with AFD-managed two-stage MoE ubatching; native DBO and PCP are unsupported. Post-fix DeepSeek-V3.2 DP2TP8+EP16 token split reached `0.9522` strict match on the complete GSM8K evaluation. The [legacy PCP8 recipe](recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) requires `release/v0.19.1rc1`. |
 
 Connector implementations are grouped by backend package:
@@ -68,6 +69,10 @@ Known gaps:
 - vLLM/vLLM-Ascend model runner v2 is not supported.
 - GPU and NPU E2E tests are opt-in and require real hardware plus model weights.
 - GPU CUDA graph support is limited to `FULL_DECODE_ONLY`.
+- Ascend A5 DeepSeek-V4 currently requires eager execution, synchronous
+  `CAMP2pAFDConnector`, gate-on-FFN, and HCCL P2P. It has unit coverage but no
+  repository hardware E2E evidence yet; A2E/E2A transport is an explicit
+  follow-up seam.
 - Native DBO is limited to exactly two ubatches and is not supported by
   `CAMAsyncAFDConnector`.
 - Qwen3 MoE currently rejects Attention-side gate placement, sequence-parallel

afd_plugin/__init__.py

@@ -139,6 +139,21 @@ def get_spawn_context(method: str | None = None):
     }
 )
 
+_NPU_DEEPSEEK_V4_REGISTRATION = (
+    "afd_plugin.model_executor.models.npu.deepseek_v4:AFDNPUDeepseekV4ForCausalLM"
+)
+
+
+def _model_registration_for_device(
+    model_arch: str,
+    model_cls: str,
+    device_type: str,
+) -> str:
+    """Select a backend wrapper while preserving the public AFD alias."""
+    if model_arch == "DeepseekV4ForCausalLM" and device_type == "npu":
+        return _NPU_DEEPSEEK_V4_REGISTRATION
+    return model_cls
+
 
 def register_afd() -> None:
     """Entry point for ``vllm.general_plugins``.
@@ -201,9 +216,15 @@ def register_afd() -> None:
     # worker startup, after vLLM-Ascend completes its platform initialization.
 
     from vllm.model_executor.models import ModelRegistry
+    from vllm.platforms import current_platform
 
     for model_arch, model_cls in _MODEL_REGISTRATIONS.items():
-        ModelRegistry.register_model(f"AFD{model_arch}", model_cls)
+        registration = _model_registration_for_device(
+            model_arch,
+            model_cls,
+            current_platform.device_type,
+        )
+        ModelRegistry.register_model(f"AFD{model_arch}", registration)
 
     _registered = True
 
@@ -220,7 +241,9 @@ def register_afd() -> None:
     "__version__",
     "_DEEPSEEK_MODEL_REGISTRATIONS",
     "_MODEL_REGISTRATIONS",
+    "_NPU_DEEPSEEK_V4_REGISTRATION",
     "_QWEN_MODEL_REGISTRATIONS",
     "_QWEN3_5_MODEL_REGISTRATIONS",
+    "_model_registration_for_device",
     "register_afd",
 ]

afd_plugin/compat/npu/feature_validation.py

@@ -130,8 +130,13 @@ 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")
+    if afd_config.connector not in {
+        AFD_ASYNC_CONNECTOR,
+        "CAMP2pAFDConnector",
+    }:
+        raise RuntimeError(
+            "DSV4 NPU AFD supports only CAMAsyncAFDConnector or CAMP2pAFDConnector",
+        )
 
 
 def _fail_if_unsupported_npu_afd_async_features(

afd_plugin/connectors/npu/camp2p.py

@@ -50,6 +50,14 @@
     recv_control_payload,
     send_control_payload,
 )
+from afd_plugin.connectors.npu.camp2p_a5 import (
+    attention_peers_for_ffn,
+    attention_token_counts,
+    dst_ffn_for_attention,
+    is_a5,
+    p2p_recv,
+    p2p_send,
+)
 from afd_plugin.distributed import (
     create_hccl_process_group_options,
     init_afd_process_group,
@@ -305,6 +313,21 @@ def is_initialized(self) -> bool:
         """Return ``True`` after all CAMP2p connections have been created."""
         return self._initialized
 
+    def _get_afd_pg(self, ubatch_idx: int) -> ProcessGroup:
+        """Return the AFD HCCL process group for a ubatch (A5 p2p path)."""
+        if not self.afd_pg_list:
+            raise RuntimeError("CAMP2P connector has no AFD process groups")
+        if ubatch_idx < 0:
+            raise RuntimeError(
+                f"CAMP2P ubatch index must be non-negative: {ubatch_idx}",
+            )
+        if ubatch_idx >= len(self.afd_pg_list):
+            raise RuntimeError(
+                f"CAMP2P ubatch {ubatch_idx} requires "
+                f"{ubatch_idx + 1} AFD process groups",
+            )
+        return self.afd_pg_list[ubatch_idx]
+
     def init_afd_connector(self) -> None:
         """Connect this process to the other Attention and FFN processes.
 
@@ -324,9 +347,13 @@ def init_afd_connector(self) -> None:
             return
         import torch_npu  # noqa: F401
 
-        ensure_cam_p2p_ops_available()
-
-        _register_camp2p_custom_ops()
+        if is_a5():
+            # The a2e/e2a custom ops are 910C-only (not registered for
+            # ascend950); Route B2 on A5 uses plain HCCL p2p instead.
+            logger.info("CAMP2P on A5: using HCCL p2p route (B2), skipping custom ops")
+        else:
+            ensure_cam_p2p_ops_available()
+            _register_camp2p_custom_ops()
 
         num_ubatches = max(1, self.vllm_config.parallel_config.num_ubatches)
         self.afd_pg_list = []
@@ -429,7 +456,10 @@ 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: May contain token-aligned ``input_ids``. The A5 P2P path
+                transfers these IDs after the hidden-state tensor. The custom
+                A2E path deliberately rejects them until its operator contract
+                supports the additional payload.
 
         Raises:
             RuntimeError: If the communication groups are not ready.
@@ -446,6 +476,41 @@ 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 = kwargs.get("input_ids")
+        if input_ids is not None and not torch.compiler.is_compiling():
+            if not isinstance(input_ids, torch.Tensor):
+                raise TypeError("CAMP2P input_ids must be a torch.Tensor")
+            if input_ids.ndim != 1 or input_ids.shape[0] != metadata.total_tokens:
+                raise ValueError(
+                    "CAMP2P input_ids must be one-dimensional and token-aligned",
+                )
+            if input_ids.dtype != torch.int32:
+                raise ValueError("CAMP2P input_ids must use torch.int32")
+            if input_ids.device != hidden_states.device:
+                raise ValueError(
+                    "CAMP2P input_ids and hidden_states must use the same device",
+                )
+        if is_a5():
+            # Route B2: plain HCCL p2p send to the mapped FFN rank. The
+            # a2e/e2a custom ops (and native MC2 ops) are unusable on A5.
+            if torch.compiler.is_compiling():
+                return None
+            ubatch_idx = metadata.stage_idx
+            get_forward_context().ubatch_idx = ubatch_idx
+            dst_ffn = dst_ffn_for_attention(
+                self.world_rank - self.ffn_size,
+                self.attn_size,
+                self.ffn_size,
+            )
+            p2p_send(self._get_afd_pg(ubatch_idx), hidden_states, dst_ffn)
+            if input_ids is not None:
+                p2p_send(self._get_afd_pg(ubatch_idx), input_ids, dst_ffn)
+            return None
+        if input_ids is not None:
+            raise NotImplementedError(
+                "CAMP2P A2E input_ids transport is not implemented; use the "
+                "A5 HCCL P2P route until the A2E/E2A operator accepts this payload",
+            )
         transfer_state = CAMP2PTransferState(
             aiv_num=self.aiv_num,
             batch_size=metadata.total_tokens,
@@ -496,6 +561,23 @@ def recv_ffn_output(
         """
         if not self._initialized:
             raise RuntimeError("CAMP2P connector is not initialized")
+        if is_a5():
+            # Route B2: recv the FFN result over HCCL p2p from the mapped FFN
+            # rank, sized by the ref tensor (same tokens as this rank sent).
+            if torch.compiler.is_compiling():
+                return ref_tensor
+            src_ffn = dst_ffn_for_attention(
+                self.world_rank - self.ffn_size,
+                self.attn_size,
+                self.ffn_size,
+            )
+            return p2p_recv(
+                self._get_afd_pg(ubatch_idx),
+                tuple(ref_tensor.shape),
+                ref_tensor.dtype,
+                ref_tensor.device,
+                src_ffn,
+            )
         transfer_state = getattr(get_forward_context(), "cam_afdtransfer_state", None)
         if transfer_state is None:
             raise RuntimeError("CAMP2P Attention side is missing connector data")
@@ -525,8 +607,9 @@ 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.
+            **kwargs: May provide existing transfer information, the layer
+                number needed to create it, and ``recv_input_ids=True`` for a
+                model whose FFN requires token IDs.
 
         Returns:
             The received hidden states and the information FFN needs to process
@@ -540,6 +623,7 @@ def recv_attn_output(
             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(kwargs.get("recv_input_ids", False))
         batch_size = _num_tokens_for_ffn_rank(
             self.dp_metadata_list,
             ubatch_idx,
@@ -548,6 +632,73 @@ def recv_attn_output(
             ffn_size=self.ffn_size,
             fallback=max_num_tokens,
         )
+        if is_a5():
+            # Route B2: recv each mapped Attention peer's token block over
+            # HCCL p2p and concatenate. Per-peer counts come from the DP
+            # metadata control plane (no equal-ratio split like e2a).
+            peers = attention_peers_for_ffn(
+                self.role_rank,
+                self.attn_size,
+                self.ffn_size,
+            )
+            counts = attention_token_counts(
+                self.dp_metadata_list,
+                ubatch_idx,
+                self.attn_size,
+            )
+            if counts is not None:
+                seq_lens = [max(1, counts[peer]) for peer in peers]
+            else:
+                # Metadata missing (e.g. warmup): even split of the total.
+                seq_lens = [max(1, batch_size // len(peers))] * len(peers)
+            dtype = self.vllm_config.model_config.dtype
+            pg = self._get_afd_pg(ubatch_idx)
+            blocks = []
+            input_id_blocks = []
+            for i, peer in enumerate(peers):
+                src_rank = self.ffn_size + peer
+                blocks.append(
+                    p2p_recv(
+                        pg,
+                        (seq_lens[i], self.hidden_size),
+                        dtype,
+                        torch.device("npu"),
+                        src_rank,
+                    ),
+                )
+                if recv_input_ids:
+                    input_id_blocks.append(
+                        p2p_recv(
+                            pg,
+                            (seq_lens[i],),
+                            torch.int32,
+                            torch.device("npu"),
+                            src_rank,
+                        ),
+                    )
+            hidden_states = torch.cat(blocks, dim=0)
+            input_ids = torch.cat(input_id_blocks, dim=0) if recv_input_ids else None
+            a5_metadata = AFDTransferMetadata.create_ffn_metadata(
+                layer_idx=layer_idx,
+                stage_idx=ubatch_idx,
+                seq_lens=seq_lens,
+            )
+            a5_states = CAMP2PTransferState(
+                aiv_num=self.aiv_num,
+                batch_size=int(sum(seq_lens)),
+                h=self.hidden_size,
+                k=self.num_experts_per_tok,
+            )
+            return AFDA2FTransferPayload(
+                hidden_states=hidden_states,
+                context=AFDTransferContext(metadata=a5_metadata, states=a5_states),
+                input_ids=input_ids,
+            )
+        if recv_input_ids:
+            raise NotImplementedError(
+                "CAMP2P A2E input_ids transport is not implemented; use the "
+                "A5 HCCL P2P route until the A2E/E2A operator accepts this payload",
+            )
         metadata = AFDTransferMetadata.create_ffn_metadata(
             layer_idx=layer_idx,
             stage_idx=ubatch_idx,
@@ -614,9 +765,31 @@ def send_ffn_output(
         if not self._initialized:
             raise RuntimeError("CAMP2P connector is not initialized")
         states = cast(CAMP2PTransferState, context.states)
+        ubatch_idx = int(kwargs.get("ubatch_idx", context.metadata.stage_idx))
+        if is_a5():
+            # Route B2: send each Attention peer its slice of the result back
+            # over HCCL p2p, split by the per-peer counts captured at recv.
+            if torch.compiler.is_compiling():
+                return None
+            peers = attention_peers_for_ffn(
+                self.role_rank,
+                self.attn_size,
+                self.ffn_size,
+            )
+            split_sizes = list(context.metadata.seq_lens)
+            if sum(split_sizes) != ffn_output.shape[0]:
+                # Inconsistent per-peer counts: fall back to an even split.
+                split_sizes = [ffn_output.shape[0] // len(peers)] * len(peers)
+                for i in range(ffn_output.shape[0] % len(peers)):
+                    split_sizes[i] += 1
+            pg = self._get_afd_pg(ubatch_idx)
+            offset = 0
+            for peer, n in zip(peers, split_sizes, strict=True):
+                p2p_send(pg, ffn_output[offset : offset + n], self.ffn_size + peer)
+                offset += n
+            return None
         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))
         group_ep = _get_group_ep(
             ubatch_idx,
             self.hccl_comm_name,