From 37d47921ff971140c709314690a0a52abbbe5cb9 Mon Sep 17 00:00:00 2001 From: ksiyuan Date: Tue, 1 Sep 2026 20:08:18 +0800 Subject: [PATCH 1/5] feat(npu): A5 (Ascend950) support for CAMP2p via HCCL p2p (Route B2) a2e/e2a custom ops and native MC2 dispatch/combine both fail on the AFD mixed group on A5 (507035 MTE out-of-range; remote HCCL windows unreachable). Replace the four a2e/e2a call sites on A5 with plain dist.send/recv over the per-ubatch 'afd' HCCL groups. Validated by a 2-rank eager 1A+1F DeepSeek-V2-Lite completion smoke on Ascend 950PR. - camp2p.py: A5 branches in send_attn_output/recv_ffn_output/recv_attn_output/send_ffn_output; init_afd_connector skips the custom-op loader on A5 (a2e/e2a register only ascend910_93). - camp2p_a5.py (new): is_a5(), contiguous rank mapping (attention i -> ffn i//(attn//ffn), matching _num_tokens_for_ffn_rank), per-peer DP token counts, p2p primitives; unit tests added. - Gate stays on FFN (only hidden_states crosses the wire). 910C custom-op path unchanged. - Known gaps: ACL graph capture, DBO, multi-rank validation, 910C regression. --- afd_plugin/connectors/npu/camp2p.py | 135 +++++++++++++++++- afd_plugin/connectors/npu/camp2p_a5.py | 121 ++++++++++++++++ .../unit/connectors/test_camp2p_connector.py | 46 ++++++ 3 files changed, 298 insertions(+), 4 deletions(-) create mode 100644 afd_plugin/connectors/npu/camp2p_a5.py diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index 8dcec3bf..a0ea3a25 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/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 = [] @@ -446,6 +473,20 @@ def send_attn_output( f"hidden_states shape {hidden_states.shape!r} does not match " f"CAMP2P metadata token count {metadata.total_tokens}", ) + 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) + return None transfer_state = CAMP2PTransferState( aiv_num=self.aiv_num, batch_size=metadata.total_tokens, @@ -496,6 +537,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") @@ -548,6 +606,53 @@ 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 = [ + p2p_recv( + pg, + (seq_lens[i], self.hidden_size), + dtype, + torch.device("npu"), + self.ffn_size + peer, + ) + for i, peer in enumerate(peers) + ] + hidden_states = torch.cat(blocks, dim=0) + 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), + ) metadata = AFDTransferMetadata.create_ffn_metadata( layer_idx=layer_idx, stage_idx=ubatch_idx, @@ -614,9 +719,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, diff --git a/afd_plugin/connectors/npu/camp2p_a5.py b/afd_plugin/connectors/npu/camp2p_a5.py new file mode 100644 index 00000000..c6114e78 --- /dev/null +++ b/afd_plugin/connectors/npu/camp2p_a5.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""A5 (Ascend950) HCCL p2p data movement for the CAMP2p connector. + +On Atlas A5 both the a2e/e2a custom ops and the torch_npu native MoE +dispatch/combine ops fail on the AFD mixed group with ``507035`` (MTE +out-of-range). This module provides the Route-B2 replacement: plain +``torch.distributed.send`` / ``recv`` over the plugin-owned ``afd`` HCCL +process group, moving hidden states between Attention and FFN ranks. The +FFN side then runs its MoE internally through the standard vLLM-Ascend EP +path, where the native ops are proven. + +Attention<->FFN rank mapping matches ``camp2p._num_tokens_for_ffn_rank``: +with ``group_size = attention_size // ffn_size``, Attention local rank ``i`` +maps to FFN rank ``i // group_size`` and FFN rank ``j`` receives from the +consecutive Attention local ranks ``[j*group_size, (j+1)*group_size)``. + +Only ``hidden_states`` crosses the wire (CAMP2p enforces gate-on-FFN, so +``router_logits`` never enters the connector). The per-peer token counts +come from the DP metadata control plane, so the FFN side can size its +receive buffers and split results back exactly (no equal-ratio split). +""" + +from __future__ import annotations + +from collections.abc import Mapping +from functools import cache + +import torch +import torch.distributed as dist + + +@cache +def is_a5() -> bool: + """Return whether this process runs on Atlas A5 (Ascend950).""" + try: + from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + + return get_ascend_device_type() == AscendDeviceType.A5 + except Exception: + # Not on an Ascend platform or vllm_ascend unavailable. + return False + + +def attention_group_size(attention_size: int, ffn_size: int) -> int: + """Number of Attention ranks mapped to each FFN rank.""" + return attention_size // ffn_size + + +def dst_ffn_for_attention( + attn_local_rank: int, + attention_size: int, + ffn_size: int, +) -> int: + """FFN rank (0-based) that an Attention rank sends its tokens to.""" + return attn_local_rank // attention_group_size(attention_size, ffn_size) + + +def attention_peers_for_ffn( + ffn_rank: int, + attention_size: int, + ffn_size: int, +) -> list[int]: + """Attention local ranks that an FFN rank receives from (ascending).""" + group_size = attention_group_size(attention_size, ffn_size) + start = ffn_rank * group_size + return list(range(start, start + group_size)) + + +def attention_token_counts( + dp_metadata_list: Mapping[int, object], + stage_idx: int, + attention_size: int, +) -> list[int] | None: + """Per-Attention-rank token counts for a stage (DP expanded to AFD ranks). + + Mirrors ``camp2p._num_tokens_for_ffn_rank``'s DP -> AFD expansion: when TP + creates several Attention workers per DP rank, the DP token count is + replicated ``tp_size`` times. Returns ``None`` when the metadata is + missing or cannot be expanded (the caller falls back to an even split). + """ + dp_metadata = dp_metadata_list.get(stage_idx) + if dp_metadata is None: + return None + token_counts = dp_metadata.num_tokens_across_dp_cpu + counts = token_counts.flatten().tolist() + if len(counts) < attention_size and attention_size % len(counts) == 0: + tp_size = attention_size // len(counts) + counts = [counts[i // tp_size] for i in range(attention_size)] + if len(counts) < attention_size: + return None + return counts + + +def p2p_send(pg, tensor: torch.Tensor, dst_rank: int) -> None: + """Blocking HCCL send of ``tensor`` to ``dst_rank`` over the AFD group.""" + dist.send(tensor.contiguous(), dst=dst_rank, group=pg) + + +def p2p_recv( + pg, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + src_rank: int, +) -> torch.Tensor: + """Blocking HCCL recv of a tensor with the given shape from ``src_rank``.""" + buf = torch.empty(shape, dtype=dtype, device=device) + dist.recv(buf, src=src_rank, group=pg) + return buf + + +__all__ = [ + "attention_group_size", + "attention_peers_for_ffn", + "attention_token_counts", + "dst_ffn_for_attention", + "is_a5", + "p2p_recv", + "p2p_send", +] diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index df513da6..016cef0e 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -26,6 +26,12 @@ CAMP2PTransferState, build_camp2p_topology, ) +from afd_plugin.connectors.npu.camp2p_a5 import ( + attention_group_size, + attention_peers_for_ffn, + attention_token_counts, + dst_ffn_for_attention, +) class _FakeDPMetadata: @@ -149,6 +155,46 @@ def test_camp2p_recv_attn_output_uses_original_contiguous_af_grouping(monkeypatc assert context0.states.k == 2 +def test_camp2p_a5_attention_to_ffn_mapping_is_contiguous_groups(): + # attn_size=4, ffn_size=2 -> group_size=2. Attention local ranks 0,1 map + # to FFN 0 and 2,3 to FFN 1 - the same grouping as _num_tokens_for_ffn_rank + # and the CAMP2P recv test (seq_lens [5] for FFN0, [12] for FFN1). + assert attention_group_size(4, 2) == 2 + assert dst_ffn_for_attention(0, 4, 2) == 0 + assert dst_ffn_for_attention(1, 4, 2) == 0 + assert dst_ffn_for_attention(2, 4, 2) == 1 + assert dst_ffn_for_attention(3, 4, 2) == 1 + assert attention_peers_for_ffn(0, 4, 2) == [0, 1] + assert attention_peers_for_ffn(1, 4, 2) == [2, 3] + + +def test_camp2p_a5_mapping_round_trip_covers_all_attention_ranks(): + for ffn_size, attn_size in [(1, 1), (1, 2), (2, 4), (2, 6), (4, 4)]: + groups = [ + attention_peers_for_ffn(j, attn_size, ffn_size) for j in range(ffn_size) + ] + assert sorted(p for group in groups for p in group) == list( + range(attn_size), + ) + for i in range(attn_size): + assert i in attention_peers_for_ffn( + dst_ffn_for_attention(i, attn_size, ffn_size), + attn_size, + ffn_size, + ) + + +def test_camp2p_a5_attention_token_counts_expands_dp_to_tp(): + # dp_size=2, attention_size=4 -> each DP count is replicated 2x. + counts = attention_token_counts({0: _FakeDPMetadata([2, 3])}, 0, 4) + assert counts == [2, 2, 3, 3] + + +def test_camp2p_a5_attention_token_counts_missing_metadata_returns_none(): + assert attention_token_counts({}, 0, 4) is None + assert attention_token_counts({0: None}, 0, 4) is None + + def test_camp2p_extra_info_rejects_unknown_mix_placement(): with pytest.raises(ValueError, match="unknown CAMP2P connector_extra_config"): CAMP2PExtraInfo.from_mapping({"mix_placement": True}) From b73cc5c3a4f41e31c5146b3fae40853b1463f0c0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=8B=E5=85=83=E7=9B=9F?= Date: Mon, 7 Sep 2026 11:26:20 +0800 Subject: [PATCH 2/5] feat(npu): support DeepSeek-V4 AFD on Ascend A5_trial --- README.md | 7 +- afd_plugin/__init__.py | 25 +++- afd_plugin/connectors/npu/camp2p.py | 70 +++++++++-- afd_plugin/connectors/npu/camp2p_a5.py | 15 +-- .../model_executor/models/deepseek_v4.py | 52 +------- .../models/deepseek_v4_common.py | 63 ++++++++++ afd_plugin/v1/worker/npu/ffn_model_runner.py | 16 +++ docs/design/module/connector_contracts.md | 7 ++ docs/design/module/model_integration.md | 24 +++- .../CAMP2pAFDConnector/deepseek_v4/README.md | 71 +++++++++++ .../deepseek_v4/afd_attention.sh | 49 ++++++++ .../CAMP2pAFDConnector/deepseek_v4/afd_ffn.sh | 49 ++++++++ .../unit/connectors/test_camp2p_connector.py | 111 ++++++++++++++++++ .../models/test_deepseek_v4_common.py | 25 ++++ .../models/test_deepseek_v4_weight_policy.py | 3 + .../models/test_npu_deepseek_v4_contract.py | 105 +++++++++++++++++ tests/unit/package/test_package.py | 33 ++++++ tests/unit/v1/worker/test_npu_runtime.py | 45 ++++++- 18 files changed, 697 insertions(+), 73 deletions(-) create mode 100644 afd_plugin/model_executor/models/deepseek_v4_common.py create mode 100644 recipe/npu/CAMP2pAFDConnector/deepseek_v4/README.md create mode 100644 recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_attention.sh create mode 100644 recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_ffn.sh create mode 100644 tests/unit/model_executor/models/test_deepseek_v4_common.py create mode 100644 tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py diff --git a/README.md b/README.md index bafc8b74..3331f1ba 100644 --- a/README.md +++ b/README.md @@ -43,6 +43,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. | @@ -53,7 +54,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: @@ -66,6 +67,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 diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index 6a7e8270..7e384790 100644 --- a/afd_plugin/__init__.py +++ b/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", ] diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index a0ea3a25..acf235cd 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -456,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. @@ -473,6 +476,20 @@ 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. @@ -486,7 +503,14 @@ def send_attn_output( 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, @@ -583,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 @@ -598,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, @@ -627,17 +653,31 @@ def recv_attn_output( 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 = [ - p2p_recv( - pg, - (seq_lens[i], self.hidden_size), - dtype, - torch.device("npu"), - self.ffn_size + peer, + 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, + ), ) - for i, peer in enumerate(peers) - ] + 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, @@ -652,6 +692,12 @@ def recv_attn_output( 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, diff --git a/afd_plugin/connectors/npu/camp2p_a5.py b/afd_plugin/connectors/npu/camp2p_a5.py index c6114e78..d03e3806 100644 --- a/afd_plugin/connectors/npu/camp2p_a5.py +++ b/afd_plugin/connectors/npu/camp2p_a5.py @@ -6,19 +6,20 @@ dispatch/combine ops fail on the AFD mixed group with ``507035`` (MTE out-of-range). This module provides the Route-B2 replacement: plain ``torch.distributed.send`` / ``recv`` over the plugin-owned ``afd`` HCCL -process group, moving hidden states between Attention and FFN ranks. The -FFN side then runs its MoE internally through the standard vLLM-Ascend EP -path, where the native ops are proven. +process group, moving hidden states and optional token IDs between Attention +and FFN ranks. The FFN side then runs its MoE internally through the standard +vLLM-Ascend EP path, where the native ops are proven. Attention<->FFN rank mapping matches ``camp2p._num_tokens_for_ffn_rank``: with ``group_size = attention_size // ffn_size``, Attention local rank ``i`` maps to FFN rank ``i // group_size`` and FFN rank ``j`` receives from the consecutive Attention local ranks ``[j*group_size, (j+1)*group_size)``. -Only ``hidden_states`` crosses the wire (CAMP2p enforces gate-on-FFN, so -``router_logits`` never enters the connector). The per-peer token counts -come from the DP metadata control plane, so the FFN side can size its -receive buffers and split results back exactly (no equal-ratio split). +``hidden_states`` always crosses the wire. DeepSeek-V4 additionally sends its +token-aligned int32 ``input_ids`` for hash routing. CAMP2p enforces gate-on-FFN, +so ``router_logits`` never enters this connector. The per-peer token counts +come from the DP metadata control plane, so the FFN side can size its receive +buffers and split results back exactly (no equal-ratio split). """ from __future__ import annotations diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index cfd15430..7b375575 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -8,7 +8,7 @@ FFN activation is the sole FFN-to-Attention tensor. """ -from collections.abc import Iterable, Iterator +from collections.abc import Iterable from typing import Any import torch @@ -20,57 +20,9 @@ from afd_plugin.config import parse_afd_config from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.models.deepseek_v4_common import _iter_role_weights from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield -_ATTENTION_ROLE = frozenset(("attention",)) -_FFN_ROLE = frozenset(("ffn",)) -_BOTH_ROLES = frozenset(("attention", "ffn")) - - -def _weight_layer_path(name: str) -> tuple[int, str] | None: - """Extract the decoder layer index and first layer-local path component.""" - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2] - return None - - -def _checkpoint_weight_roles(name: str) -> frozenset[str]: - """Classify a native DeepSeek-V4 checkpoint path by execution owner.""" - if name in { - "hc_head_fn", - "hc_head_base", - "hc_head_scale", - "model.hc_head_fn", - "model.hc_head_base", - "model.hc_head_scale", - }: - return _ATTENTION_ROLE - layer_path = _weight_layer_path(name) - if layer_path is None: - return _BOTH_ROLES - _, stage = layer_path - if stage == "ffn": - return _FFN_ROLE - return _ATTENTION_ROLE - - -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str, -) -> Iterator[tuple[str, torch.Tensor]]: - """Consume a checkpoint iterator once and retain only role-owned paths.""" - for name, loaded_weight in weights: - if role in _checkpoint_weight_roles(name): - yield name, loaded_weight - class RemoteDeepseekV4FFN(nn.Module): """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" diff --git a/afd_plugin/model_executor/models/deepseek_v4_common.py b/afd_plugin/model_executor/models/deepseek_v4_common.py new file mode 100644 index 00000000..18670466 --- /dev/null +++ b/afd_plugin/model_executor/models/deepseek_v4_common.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Backend-neutral DeepSeek-V4 checkpoint ownership helpers.""" + +from __future__ import annotations + +from collections.abc import Iterable, Iterator +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + import torch + +_ATTENTION_ROLE = frozenset(("attention",)) +_FFN_ROLE = frozenset(("ffn",)) +_BOTH_ROLES = frozenset(("attention", "ffn")) + + +def _weight_layer_path(name: str) -> tuple[int, str] | None: + """Extract the decoder layer index and first layer-local path component.""" + parts = name.split(".") + for marker_idx, part in enumerate(parts[:-2]): + if part != "layers": + continue + try: + layer_idx = int(parts[marker_idx + 1]) + except ValueError: + continue + return layer_idx, parts[marker_idx + 2] + return None + + +def _checkpoint_weight_roles(name: str) -> frozenset[str]: + """Classify a native GPU or Ascend V4 checkpoint path by AFD owner.""" + if name in { + "hc_head_fn", + "hc_head_base", + "hc_head_scale", + "model.hc_head_fn", + "model.hc_head_base", + "model.hc_head_scale", + }: + return _ATTENTION_ROLE + layer_path = _weight_layer_path(name) + if layer_path is None: + return _BOTH_ROLES + _, stage = layer_path + if stage in {"ffn", "mlp"}: + return _FFN_ROLE + return _ATTENTION_ROLE + + +def _iter_role_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + role: str, +) -> Iterator[tuple[str, torch.Tensor]]: + """Consume a checkpoint iterator once and retain only role-owned paths.""" + for name, loaded_weight in weights: + if role in _checkpoint_weight_roles(name): + yield name, loaded_weight + + +__all__ = ["_checkpoint_weight_roles", "_iter_role_weights"] diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 60ad17c3..5e6f1beb 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -235,6 +235,9 @@ def _ffn_forward( ) stage_ids = sorted(int(stage_idx) for stage_idx in dp_metadata_list) or [0] rank_ffn_output = None + requires_input_ids = bool( + getattr(self.model, "afd_requires_input_ids", False), + ) for layer_idx in _ffn_layer_indices(self): for stage_idx in stage_ids: @@ -264,10 +267,14 @@ def _ffn_forward( in_profile_run=is_profile, aclgraph_runtime_mode=aclgraph_runtime_mode, ) as forward_context: + recv_kwargs: dict[str, Any] = {} + if requires_input_ids: + recv_kwargs["recv_input_ids"] = True payload = self.connector.recv_attn_output( ubatch_idx=stage_idx, layer_idx=layer_idx, max_num_tokens=self.max_num_tokens, + **recv_kwargs, ) context = payload.context metadata = context.metadata @@ -280,9 +287,18 @@ def _ffn_forward( assert states, "Context.states must not be None" _set_moe_layer_index(forward_context, layer_idx) + compute_kwargs: dict[str, Any] = {} + if requires_input_ids: + if payload.input_ids is None: + raise RuntimeError( + "AFD model requires input_ids but the connector " + "did not return them", + ) + compute_kwargs["input_ids"] = payload.input_ids rank_ffn_output = self.model.compute_ffn_output( hidden_states=hidden_states, layer_idx=layer_idx, + **compute_kwargs, ) _send_ffn_output( self.connector, diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 14a24b92..00ac5d08 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -109,6 +109,13 @@ synchronous NPU runtime requires both common and connector-local values to be | `CAMP2pAFDConnector` | Ascend | FFN ranks, then Attention ranks | `CAMP2pAFDControlPlane`; stage DP metadata over Gloo plus HCCL data groups | `connector.control_plane is not None` | | `CAMAsyncAFDConnector` | Ascend | Attention ranks, then FFN ranks | `None`; routing/token metadata travels with CAM dispatch payloads | `connector.control_plane is None` | +On A5, `CAMP2pAFDConnector` has a backend-local blocking HCCL P2P route. It +always transfers hidden states and can additionally transfer token-aligned +int32 `input_ids` for DeepSeek-V4 hash routing. The `AFDA2FTransferPayload` +field is backend-neutral: future A5 A2E/E2A operators must populate the same +field so the model and FFN runner remain transport-independent. Until then, +the custom A2E path rejects requests for `input_ids` rather than dropping them. + The CUDA P2P mapping requires `num_attention_ranks >= num_ffn_ranks` and an integral A/F ratio. Each FFN rank owns a subgroup containing itself and consecutive Attention peers. CAMP2P also diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 26b3c77e..bc8f90d1 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -77,7 +77,7 @@ registers lazy AFD wrapper paths under `AFD`-prefixed aliases. | `DeepseekV2ForCausalLM` | `AFDDeepseekV2ForCausalLM` | `AFDDeepseekV2ForCausalLM` | | `DeepseekV3ForCausalLM` | `AFDDeepseekV3ForCausalLM` | `AFDDeepseekV3ForCausalLM` | | `DeepseekV32ForCausalLM` | `AFDDeepseekV32ForCausalLM` | `AFDDeepseekV3ForCausalLM` | -| `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | +| `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | CUDA: `AFDDeepseekV4ForCausalLM`; Ascend: `AFDNPUDeepseekV4ForCausalLM` | | `GlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | | `Qwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | @@ -140,6 +140,28 @@ connector validates one-dimensional `torch.int32` input IDs and preallocates their receive buffers for graph execution. This boundary currently has focused unit coverage but no repository model or accuracy E2E case. +### DeepSeek V4 Ascend A5 boundary + +On Ascend, the same `AFDDeepseekV4ForCausalLM` registry alias resolves to the +backend-local `AFDNPUDeepseekV4ForCausalLM`. It subclasses vLLM-Ascend's +native DeepSeek-V4 implementation, retaining DSA Attention, NPU mHC operators, +and the native `DeepseekV4MoE`. Attention owns all residual-stream and +normalization state and uses `RemoteNPUDeepseekV4FFN`; FFN owns gate, hash +router, shared experts, and routed experts. + +The A5 P2P payload is hidden states followed by token-aligned int32 +`input_ids`; FFN returns the output hidden states. The NPU FFN runner requests +the optional IDs only for models declaring `afd_requires_input_ids`, keeping +other model contracts unchanged. The non-P2P CAMP2P branch currently rejects +that payload explicitly. Its existing optional payload interface is the seam +where the planned A5 A2E/E2A operators should be connected. + +The first implementation requires A5, `CAMP2pAFDConnector`, gate-on-FFN, +pipeline/context-parallel size 1, no SP MoE, EPLB/elastic EP, LoRA, or +speculative decoding, and `--enforce-eager`. It is based on vLLM-Ascend commit +`4fe7ddbf94bc28bcd2b9f3d2d93f0fe1f0499cf5` and still requires hardware +accuracy and stability validation. + ### Qwen3 MoE CUDA boundary `AFDQwen3MoeModel` uses the native `decoder_layer_type` injection hook. The diff --git a/recipe/npu/CAMP2pAFDConnector/deepseek_v4/README.md b/recipe/npu/CAMP2pAFDConnector/deepseek_v4/README.md new file mode 100644 index 00000000..f7a439bc --- /dev/null +++ b/recipe/npu/CAMP2pAFDConnector/deepseek_v4/README.md @@ -0,0 +1,71 @@ +# DeepSeek-V4 AFD on Ascend A5 with HCCL P2P + +This is the first functional NPU path for DeepSeek-V4 AFD. It reuses the +native vLLM-Ascend DeepSeek-V4 DSA, mHC, and MoE implementation and separates +execution immediately before each decoder FFN. + +The checked-in launchers default to: + +```text +/mnt/share/weight/DeepSeek-V4-Flash +``` + +Override it with `MODEL_PATH` when required. + +## Current support boundary + +- Hardware: Ascend A5 (Ascend 950). +- Communication: synchronous blocking HCCL P2P through + `CAMP2pAFDConnector`. +- Payload A to F: normalized hidden states followed by token-aligned + `torch.int32` input IDs. The IDs are required by V4 hash routing. +- Payload F to A: FFN output hidden states. +- Gate and the complete native MoE run on the FFN role. +- Pipeline/context parallelism, sequence-parallel MoE, EPLB/elastic EP, LoRA, + speculative decoding, and graph execution are intentionally rejected for + this first path. +- `--enforce-eager` is mandatory while the connector uses blocking P2P. +- The reference launchers explicitly select model runner V1 for the initial + smoke and accuracy pass. + +The connector API already exposes optional `input_ids` on its backend-neutral +receive payload. When the A5 A2E/E2A operators are ready, implement the same +payload contract in the non-P2P branch of `CAMP2pAFDConnector`; the model and +FFN runner do not need another interface change. + +## Reference topology + +The scripts use an 8A8F reference: one 8-NPU Attention node and one 8-NPU FFN +node, each with TP8. Set `TP_SIZE`, `ATTENTION_RANKS`, `FFN_RANKS`, and +`ASCEND_RT_VISIBLE_DEVICES` to match the real cluster. The total AFD rank +counts must describe all workers across both roles, and `ATTENTION_RANKS` +must be divisible by `FFN_RANKS` for the current contiguous P2P mapping. + +Install this branch in the same environment as its matching vLLM and +vLLM-Ascend checkouts: + +```bash +cd /path/to/afd-plugin +pip install -e . --no-build-isolation -v +``` + +Start the FFN node first: + +```bash +cd recipe/npu/CAMP2pAFDConnector/deepseek_v4 +AFD_HOST= NIC_NAME= bash afd_ffn.sh +``` + +Then start the Attention node with the same rendezvous address and port: + +```bash +cd recipe/npu/CAMP2pAFDConnector/deepseek_v4 +AFD_HOST= NIC_NAME= bash afd_attention.sh +``` + +Send requests only to the Attention API port (`8900` by default). Both roles +must use identical model weights, rank counts, `AFD_HOST`, and `AFD_PORT`. + +Start with a short prompt and one request. Before performance measurement, +compare greedy output against a non-AFD native vLLM-Ascend deployment using +the same checkpoint and serving parameters. diff --git a/recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_attention.sh b/recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_attention.sh new file mode 100644 index 00000000..b2aed3ec --- /dev/null +++ b/recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_attention.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL_PATH="${MODEL_PATH:-/mnt/share/weight/DeepSeek-V4-Flash}" +: "${AFD_HOST:?Set AFD_HOST to the FFN node IP}" +: "${NIC_NAME:?Set NIC_NAME to the HCCL/Gloo network interface}" + +AFD_PORT="${AFD_PORT:-29666}" +SERVER_PORT="${SERVER_PORT:-8900}" +TP_SIZE="${TP_SIZE:-8}" +ATTENTION_RANKS="${ATTENTION_RANKS:-8}" +FFN_RANKS="${FFN_RANKS:-8}" + +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export HCCL_BUFFSIZE="${HCCL_BUFFSIZE:-1024}" +export HCCL_OP_EXPANSION_MODE="${HCCL_OP_EXPANSION_MODE:-AIV}" +export OMP_PROC_BIND="${OMP_PROC_BIND:-false}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-10}" +export PYTORCH_NPU_ALLOC_CONF="${PYTORCH_NPU_ALLOC_CONF:-expandable_segments:True}" +export VLLM_PLUGINS="${VLLM_PLUGINS:-ascend,afd}" +export GLOO_SOCKET_IFNAME="$NIC_NAME" +export TP_SOCKET_IFNAME="$NIC_NAME" +export HCCL_SOCKET_IFNAME="$NIC_NAME" + +ADDITIONAL_CONFIG="$(printf \ + '{"afd":{"role":"attention","connector":"CAMP2pAFDConnector","host":"%s","port":%s,"num_attention_ranks":%s,"num_ffn_ranks":%s}}' \ + "$AFD_HOST" "$AFD_PORT" "$ATTENTION_RANKS" "$FFN_RANKS")" + +exec env VLLM_USE_V1=1 VLLM_USE_V2_MODEL_RUNNER=0 vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 \ + --port "$SERVER_PORT" \ + --tensor-parallel-size "$TP_SIZE" \ + --data-parallel-size 1 \ + --enable-expert-parallel \ + --enforce-eager \ + --max-model-len "${MAX_MODEL_LEN:-32768}" \ + --max-num-batched-tokens "${MAX_NUM_BATCHED_TOKENS:-2048}" \ + --max-num-seqs "${MAX_NUM_SEQS:-16}" \ + --gpu-memory-utilization "${MEMORY_UTILIZATION:-0.9}" \ + --block-size 128 \ + --quantization ascend \ + --tokenizer-mode deepseek_v4 \ + --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --reasoning-parser deepseek_v4 \ + --no-enable-prefix-caching \ + --trust-remote-code \ + --served-model-name deepseek_v4_afd \ + --additional-config "$ADDITIONAL_CONFIG" diff --git a/recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_ffn.sh b/recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_ffn.sh new file mode 100644 index 00000000..d1e4abf1 --- /dev/null +++ b/recipe/npu/CAMP2pAFDConnector/deepseek_v4/afd_ffn.sh @@ -0,0 +1,49 @@ +#!/usr/bin/env bash +set -euo pipefail + +MODEL_PATH="${MODEL_PATH:-/mnt/share/weight/DeepSeek-V4-Flash}" +: "${AFD_HOST:?Set AFD_HOST to this FFN node IP}" +: "${NIC_NAME:?Set NIC_NAME to the HCCL/Gloo network interface}" + +AFD_PORT="${AFD_PORT:-29666}" +SERVER_PORT="${SERVER_PORT:-8901}" +TP_SIZE="${TP_SIZE:-8}" +ATTENTION_RANKS="${ATTENTION_RANKS:-8}" +FFN_RANKS="${FFN_RANKS:-8}" + +export ASCEND_RT_VISIBLE_DEVICES="${ASCEND_RT_VISIBLE_DEVICES:-0,1,2,3,4,5,6,7}" +export HCCL_BUFFSIZE="${HCCL_BUFFSIZE:-1024}" +export HCCL_OP_EXPANSION_MODE="${HCCL_OP_EXPANSION_MODE:-AIV}" +export OMP_PROC_BIND="${OMP_PROC_BIND:-false}" +export OMP_NUM_THREADS="${OMP_NUM_THREADS:-10}" +export PYTORCH_NPU_ALLOC_CONF="${PYTORCH_NPU_ALLOC_CONF:-expandable_segments:True}" +export VLLM_PLUGINS="${VLLM_PLUGINS:-ascend,afd}" +export GLOO_SOCKET_IFNAME="$NIC_NAME" +export TP_SOCKET_IFNAME="$NIC_NAME" +export HCCL_SOCKET_IFNAME="$NIC_NAME" + +ADDITIONAL_CONFIG="$(printf \ + '{"afd":{"role":"ffn","connector":"CAMP2pAFDConnector","host":"%s","port":%s,"num_attention_ranks":%s,"num_ffn_ranks":%s}}' \ + "$AFD_HOST" "$AFD_PORT" "$ATTENTION_RANKS" "$FFN_RANKS")" + +exec env VLLM_USE_V1=1 VLLM_USE_V2_MODEL_RUNNER=0 vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 \ + --port "$SERVER_PORT" \ + --tensor-parallel-size "$TP_SIZE" \ + --data-parallel-size 1 \ + --enable-expert-parallel \ + --enforce-eager \ + --max-model-len "${MAX_MODEL_LEN:-32768}" \ + --max-num-batched-tokens "${MAX_NUM_BATCHED_TOKENS:-2048}" \ + --max-num-seqs "${MAX_NUM_SEQS:-16}" \ + --gpu-memory-utilization "${MEMORY_UTILIZATION:-0.9}" \ + --block-size 128 \ + --quantization ascend \ + --tokenizer-mode deepseek_v4 \ + --tool-call-parser deepseek_v4 \ + --enable-auto-tool-choice \ + --reasoning-parser deepseek_v4 \ + --no-enable-prefix-caching \ + --trust-remote-code \ + --served-model-name deepseek_v4_afd_ffn \ + --additional-config "$ADDITIONAL_CONFIG" diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 016cef0e..ecf8e955 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -195,6 +195,117 @@ def test_camp2p_a5_attention_token_counts_missing_metadata_returns_none(): assert attention_token_counts({0: None}, 0, 4) is None +def test_camp2p_a5_sends_hidden_states_then_v4_input_ids(monkeypatch): + torch = pytest.importorskip("torch") + connector = CAMP2pAFDConnector( + 0, + 0, + _vllm_config(), + _afd_config(role="attention"), + 0, + ) + connector._initialized = True + connector.afd_pg_list = [object()] + hidden_states = torch.ones((3, 16), dtype=torch.float16) + input_ids = torch.tensor([11, 13, 17], dtype=torch.int32) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=2, + stage_idx=0, + seq_len=3, + ), + ) + sent = [] + + monkeypatch.setattr(camp2p_module, "is_a5", lambda: True) + monkeypatch.setattr( + camp2p_module, + "get_forward_context", + lambda: SimpleNamespace(), + ) + monkeypatch.setattr( + camp2p_module, + "p2p_send", + lambda pg, tensor, dst: sent.append((pg, tensor, dst)), + ) + + connector.send_attn_output( + hidden_states, + context, + input_ids=input_ids, + ) + + assert sent[0][1] is hidden_states + assert sent[1][1] is input_ids + assert [item[2] for item in sent] == [0, 0] + + +def test_camp2p_a5_receives_token_aligned_v4_input_ids(monkeypatch): + torch = pytest.importorskip("torch") + connector = _init_ffn_connector(0, _vllm_config()) + connector.afd_pg_list = [object()] + connector.dp_metadata_list = {0: _FakeDPMetadata([2, 3, 5, 7])} + receives = [] + + monkeypatch.setattr(camp2p_module, "is_a5", lambda: True) + + def fake_recv(pg, shape, dtype, device, src_rank): + receives.append((pg, shape, dtype, device, src_rank)) + if len(shape) == 1: + return torch.arange(shape[0], dtype=torch.int32) + src_rank * 10 + return torch.full(shape, float(src_rank), dtype=torch.float16) + + monkeypatch.setattr(camp2p_module, "p2p_recv", fake_recv) + + payload = connector.recv_attn_output( + ubatch_idx=0, + layer_idx=4, + recv_input_ids=True, + ) + + assert payload.hidden_states.shape == (5, 16) + assert payload.input_ids is not None + assert payload.input_ids.dtype is torch.int32 + assert payload.input_ids.tolist() == [20, 21, 30, 31, 32] + assert payload.context.metadata.seq_lens == [2, 3] + assert [(item[1], item[4]) for item in receives] == [ + ((2, 16), 2), + ((2,), 2), + ((3, 16), 3), + ((3,), 3), + ] + + +def test_camp2p_custom_a2e_explicitly_rejects_v4_input_ids(monkeypatch): + torch = pytest.importorskip("torch") + connector = CAMP2pAFDConnector( + 0, + 0, + _vllm_config(), + _afd_config(role="attention"), + 0, + ) + connector._initialized = True + hidden_states = torch.ones((2, 16), dtype=torch.float16) + input_ids = torch.tensor([1, 2], dtype=torch.int32) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=2, + ), + ) + + monkeypatch.setattr(camp2p_module, "is_a5", lambda: False) + + with pytest.raises(NotImplementedError, match="A2E input_ids transport"): + connector.send_attn_output( + hidden_states, + context, + input_ids=input_ids, + ) + + def test_camp2p_extra_info_rejects_unknown_mix_placement(): with pytest.raises(ValueError, match="unknown CAMP2P connector_extra_config"): CAMP2PExtraInfo.from_mapping({"mix_placement": True}) diff --git a/tests/unit/model_executor/models/test_deepseek_v4_common.py b/tests/unit/model_executor/models/test_deepseek_v4_common.py new file mode 100644 index 00000000..5f066124 --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_common.py @@ -0,0 +1,25 @@ +import pytest + +pytest.importorskip("vllm") + +from afd_plugin.model_executor.models.deepseek_v4_common import ( + _checkpoint_weight_roles, +) + + +def test_v4_common_weight_policy_supports_gpu_and_ascend_names(): + assert _checkpoint_weight_roles("model.layers.0.ffn.gate.weight") == frozenset( + ("ffn",), + ) + assert _checkpoint_weight_roles( + "model.layers.0.mlp.experts.0.down_proj.weight", + ) == frozenset(("ffn",)) + assert _checkpoint_weight_roles( + "model.layers.0.self_attn.q_proj.weight", + ) == frozenset(("attention",)) + assert _checkpoint_weight_roles("model.hc_head_fn") == frozenset( + ("attention",), + ) + assert _checkpoint_weight_roles("model.embed_tokens.weight") == frozenset( + ("attention", "ffn"), + ) diff --git a/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py b/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py index fa39bd24..1ebbd300 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py @@ -9,6 +9,8 @@ from afd_plugin.model_executor.models.deepseek_v4 import ( # noqa: E402 AFDDeepseekV4ForCausalLM, +) +from afd_plugin.model_executor.models.deepseek_v4_common import ( # noqa: E402 _checkpoint_weight_roles, ) @@ -31,6 +33,7 @@ def __iter__(self): "layers.0.ffn.gate.weight", "layers.1.ffn.experts.0.w1.weight", "model.layers.2.ffn.shared_experts.w2.weight", + "model.layers.3.mlp.experts.0.down_proj.weight", ], ) def test_v4_raw_checkpoint_ffn_paths_are_ffn_owned(name): diff --git a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py new file mode 100644 index 00000000..3862bf31 --- /dev/null +++ b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py @@ -0,0 +1,105 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("torch_npu") +pytest.importorskip("vllm_ascend") + +from vllm_ascend.models.deepseek_v4 import model as native # noqa: E402 + +from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 + + +class _FakeConnector: + def __init__(self) -> None: + self.sent = None + + def send_attn_output(self, hidden_states, context, **kwargs) -> None: + self.sent = (hidden_states, context, kwargs) + + def recv_ffn_output(self, *, ref_tensor, ubatch_idx): + return ref_tensor * 0.5 + + +def test_npu_v4_wrapper_uses_ascend_native_classes(): + assert issubclass( + adapter.AFDNPUDeepseekV4DecoderLayer, + native.DeepseekV2DecoderLayer, + ) + assert issubclass( + adapter.AFDNPUDeepseekV4ForCausalLM, + native.AscendDeepseekV4ForCausalLM, + ) + registered_model_cls = adapter.AFDNPUDeepseekV4ForCausalLM.model_cls + assert registered_model_cls is adapter.AFDNPUDeepseekV4Model + assert adapter.AFDNPUDeepseekV4ForCausalLM.afd_requires_input_ids + + +def test_npu_v4_remote_ffn_sends_input_ids(monkeypatch): + connector = _FakeConnector() + afd_metadata = SimpleNamespace(connector=connector, stage_idx=0) + monkeypatch.setattr( + adapter, + "get_afd_metadata_from_forward_context", + lambda: afd_metadata, + ) + monkeypatch.setattr( + adapter, + "get_forward_context", + lambda: SimpleNamespace(ubatch_idx=1), + ) + monkeypatch.setattr( + adapter, + "maybe_apply_dbo_yield", + lambda hidden_states, *, role: hidden_states, + ) + proxy = adapter.RemoteNPUDeepseekV4FFN(layer_idx=3) + hidden_states = torch.ones((2, 4), dtype=torch.float16) + input_ids = torch.tensor([7, 11], dtype=torch.int32) + + output = proxy(hidden_states, input_ids) + + assert connector.sent is not None + assert connector.sent[1].metadata.layer_idx == 3 + assert connector.sent[1].metadata.stage_idx == 1 + assert connector.sent[2]["input_ids"] is input_ids + assert torch.equal(output, hidden_states * 0.5) + + +def test_npu_v4_model_requires_eager_a5_p2p(monkeypatch): + afd_config = SimpleNamespace( + compute_gate_on_attention=False, + connector="CAMP2pAFDConnector", + role="attention", + ) + vllm_config = SimpleNamespace( + model_config=SimpleNamespace(enforce_eager=False), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_elastic_ep=False, + enable_eplb=False, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + use_sequence_parallel_moe=False, + ), + lora_config=None, + speculative_config=None, + ) + monkeypatch.setattr( + adapter.native, + "current_platform", + SimpleNamespace(device_type="npu"), + ) + monkeypatch.setattr(adapter, "is_a5", lambda: True) + monkeypatch.setattr( + adapter, + "parse_afd_config", + lambda *_args, **_kwargs: afd_config, + ) + + with pytest.raises(RuntimeError, match="requires --enforce-eager"): + adapter.AFDNPUDeepseekV4Model(vllm_config=vllm_config) diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index 035131d1..0b2e8e80 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -37,6 +37,39 @@ def test_deepseek_afd_model_registration_paths_are_lazy_strings(): ) +def test_deepseek_v4_registration_selects_ascend_wrapper_on_npu(): + gpu_registration = afd_plugin._DEEPSEEK_MODEL_REGISTRATIONS["DeepseekV4ForCausalLM"] + + assert afd_plugin._model_registration_for_device( + "DeepseekV4ForCausalLM", + gpu_registration, + "npu", + ) == ( + "afd_plugin.model_executor.models.npu.deepseek_v4:AFDNPUDeepseekV4ForCausalLM" + ) + assert ( + afd_plugin._model_registration_for_device( + "DeepseekV4ForCausalLM", + gpu_registration, + "cuda", + ) + == gpu_registration + ) + + +def test_backend_selection_does_not_change_other_model_registrations(): + registration = afd_plugin._DEEPSEEK_MODEL_REGISTRATIONS["DeepseekV3ForCausalLM"] + + assert ( + afd_plugin._model_registration_for_device( + "DeepseekV3ForCausalLM", + registration, + "npu", + ) + == registration + ) + + def test_qwen3_moe_afd_model_registration_path_is_lazy_string(): registrations = afd_plugin._QWEN_MODEL_REGISTRATIONS diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 82f1e317..b67bb797 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -55,13 +55,14 @@ def _temporarily_reimport_module(module_name: str) -> Iterator[ModuleType]: package.__dict__[module_attribute] = original_package_attribute -def _ffn_payload(hidden_states, metadata, states=None): +def _ffn_payload(hidden_states, metadata, states=None, input_ids=None): return AFDA2FTransferPayload( hidden_states=hidden_states, context=AFDTransferContext( metadata=metadata, states=states if states is not None else AFDTransferState(), ), + input_ids=input_ids, ) @@ -135,6 +136,7 @@ def __init__(self, *, attn_size=1, ffn_size=1, role_rank=0, world_rank=0): AFDA2FTransferPayload | tuple[object, AFDTransferMetadata] ] = deque() self.ffn_outputs = [] + self.recv_calls = [] self.updates = [] self.attn_size = attn_size self.ffn_size = ffn_size @@ -158,6 +160,7 @@ def update_state_from_dp_metadata(self, payload): ) def recv_attn_output(self, ubatch_idx=None, **kwargs): + self.recv_calls.append((ubatch_idx, kwargs)) for item in tuple(self.attn_outputs): payload = ( item @@ -190,6 +193,10 @@ def compute_ffn_output(self, hidden_states, layer_idx, **kwargs): return f"npu-ffn({hidden_states}, layer={layer_idx})" +class _InputIdRecordingFakeModel(_RecordingFakeModel): + afd_requires_input_ids = True + + class _FakeStructuredFFNModel: def compute_ffn_output(self, hidden_states, layer_idx, **_kwargs): return AFDF2ATransferPayload( @@ -1410,6 +1417,42 @@ def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): ] +def test_npu_ffn_runner_requests_and_forwards_v4_input_ids(monkeypatch): + _patch_ffn_forward_context(monkeypatch) + runner = _new_ffn_runner() + runner.vllm_config = _vllm_config(role="ffn") + runner.connector = _FakeFFNConnector() + runner.model = _InputIdRecordingFakeModel() + runner.num_layers = 1 + runner.max_num_tokens = 2 + runner.use_aclgraph = False + runner._acl_graphs = {} + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=2, + ) + runner.connector.attn_outputs.append( + _ffn_payload("hidden", metadata, input_ids="token-ids"), + ) + + runner.execute_model(dp_metadata_list={0: _FakeDPMetadata([2])}) + + assert runner.connector.recv_calls == [ + ( + 0, + { + "layer_idx": 0, + "max_num_tokens": 2, + "recv_input_ids": True, + }, + ), + ] + assert runner.model.calls == [ + ("hidden", 0, {"input_ids": "token-ids"}), + ] + + def test_npu_ffn_runner_builds_forward_context_for_each_dbo_stage(monkeypatch): _require_npu_runtime() from afd_plugin.v1.worker.npu import ffn_model_runner From 37ba794bc1cd41a24008eb2d0397d8ff4e5638ee Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=8B=E5=85=83=E7=9B=9F?= Date: Mon, 7 Sep 2026 14:40:21 +0800 Subject: [PATCH 3/5] fix(npu): support v0.26 DeepSeek-V4 module layout --- .../models/test_npu_deepseek_v4_contract.py | 28 ++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py index 3862bf31..c11f1566 100644 --- a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py +++ b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py @@ -1,6 +1,6 @@ from __future__ import annotations -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import pytest @@ -39,6 +39,32 @@ def test_npu_v4_wrapper_uses_ascend_native_classes(): assert adapter.AFDNPUDeepseekV4ForCausalLM.afd_requires_input_ids +def test_npu_v4_native_import_supports_flat_module(monkeypatch): + flat_module = ModuleType("vllm_ascend.models.deepseek_v4") + monkeypatch.setattr(adapter, "import_module", lambda _name: flat_module) + + assert adapter._import_native_deepseek_v4() is flat_module + + +def test_npu_v4_native_import_supports_package_module(monkeypatch): + package_module = ModuleType("vllm_ascend.models.deepseek_v4") + package_module.__path__ = [] + model_module = ModuleType("vllm_ascend.models.deepseek_v4.model") + imported_names = [] + + def fake_import_module(name): + imported_names.append(name) + return model_module if name.endswith(".model") else package_module + + monkeypatch.setattr(adapter, "import_module", fake_import_module) + + assert adapter._import_native_deepseek_v4() is model_module + assert imported_names == [ + "vllm_ascend.models.deepseek_v4", + "vllm_ascend.models.deepseek_v4.model", + ] + + def test_npu_v4_remote_ffn_sends_input_ids(monkeypatch): connector = _FakeConnector() afd_metadata = SimpleNamespace(connector=connector, stage_idx=0) From 160c37d59e6d74b4e8f3a1a4d9668454b9079cc6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=8B=E5=85=83=E7=9B=9F?= Date: Mon, 7 Sep 2026 19:45:45 +0800 Subject: [PATCH 4/5] fix_graph --- .../models/test_npu_deepseek_v4_contract.py | 84 ++++++++++++++++++- 1 file changed, 82 insertions(+), 2 deletions(-) diff --git a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py index c11f1566..e8e817a5 100644 --- a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py +++ b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py @@ -9,10 +9,10 @@ pytest.importorskip("torch_npu") pytest.importorskip("vllm_ascend") -from vllm_ascend.models.deepseek_v4 import model as native # noqa: E402 - from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 +native = adapter.native + class _FakeConnector: def __init__(self) -> None: @@ -129,3 +129,83 @@ def test_npu_v4_model_requires_eager_a5_p2p(monkeypatch): with pytest.raises(RuntimeError, match="requires --enforce-eager"): adapter.AFDNPUDeepseekV4Model(vllm_config=vllm_config) + + +def test_npu_v4_model_preserves_eager_compile_contract(monkeypatch): + afd_config = SimpleNamespace( + compute_gate_on_attention=False, + connector="CAMP2pAFDConnector", + role="ffn", + ) + compilation_config = SimpleNamespace(mode="none") + hf_config = SimpleNamespace( + hc_eps=1e-5, + hc_mult=1, + hidden_size=8, + num_hidden_layers=1, + rms_norm_eps=1e-5, + vocab_size=16, + ) + vllm_config = SimpleNamespace( + compilation_config=compilation_config, + lora_config=None, + model_config=SimpleNamespace( + enforce_eager=True, + hf_config=hf_config, + ), + parallel_config=SimpleNamespace( + decode_context_parallel_size=1, + enable_elastic_ep=False, + enable_eplb=False, + pipeline_parallel_size=1, + prefill_context_parallel_size=1, + use_sequence_parallel_moe=False, + ), + quant_config=None, + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + speculative_config=None, + ) + monkeypatch.setattr( + adapter.native, + "current_platform", + SimpleNamespace(device_type="npu"), + ) + monkeypatch.setattr(adapter, "is_a5", lambda: True) + monkeypatch.setattr( + adapter, + "parse_afd_config", + lambda *_args, **_kwargs: afd_config, + ) + monkeypatch.setattr( + adapter.native, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=False, is_last_rank=False), + ) + monkeypatch.setattr( + adapter.native, + "make_layers", + lambda *_args, **_kwargs: (0, 0, torch.nn.ModuleList()), + ) + monkeypatch.setattr( + adapter.native, + "make_pp_empty_intermediate_tensors", + lambda _model, factory: factory, + raising=False, + ) + monkeypatch.setattr( + adapter.native, + "PPMissingLayer", + torch.nn.Identity, + ) + monkeypatch.setattr( + adapter.AFDNPUDeepseekV4Model, + "forward", + lambda _self: "eager-forward", + ) + + model = adapter.AFDNPUDeepseekV4Model(vllm_config=vllm_config) + + assert model.vllm_config is vllm_config + assert model.compilation_config is compilation_config + assert model.do_not_compile is True + assert model() == "eager-forward" From 9299ffa2d7ded830cc4c1c8a9751bf0a66f2989e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E5=AE=8B=E5=85=83=E7=9B=9F?= Date: Tue, 8 Sep 2026 09:18:49 +0800 Subject: [PATCH 5/5] fix(npu): preserve DeepSeek-V4 input ids across AFD --- afd_plugin/compat/npu/feature_validation.py | 9 +- .../model_executor/models/npu/deepseek_v4.py | 44 ++ .../models/npu/deepseek_v4_p2p.py | 575 ++++++++++++++++++ afd_plugin/v1/worker/npu/ffn_model_runner.py | 5 + docs/design/module/model_integration.md | 8 +- .../compat/npu/test_dsv4_async_validation.py | 15 +- .../models/test_npu_deepseek_v4_contract.py | 106 +++- tests/unit/v1/worker/test_npu_runtime.py | 20 +- 8 files changed, 772 insertions(+), 10 deletions(-) create mode 100644 afd_plugin/model_executor/models/npu/deepseek_v4_p2p.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index a8278e93..569e5fba 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/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( diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index daa01905..f112f786 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -568,8 +568,52 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return super().load_weights(_iter_role_weights(weights, role=self.afd_role)) +class AFDNPUDeepseekV4ForCausalLM(AFDDeepseekV4ForCausalLM): + """Select the Async CAM or A5 P2P DSV4 model for the active connector.""" + + model_cls = AFDDeepseekV4Model + afd_requires_input_ids = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + afd_config = parse_afd_config(vllm_config, validate=False) + if afd_config.connector == "CAMP2pAFDConnector": + from afd_plugin.model_executor.models.npu.deepseek_v4_p2p import ( + AFDNPUDeepseekV4Model, + ) + + self.model_cls = AFDNPUDeepseekV4Model + self.afd_requires_input_ids = True + else: + self.model_cls = AFDDeepseekV4Model + self.afd_requires_input_ids = False + super().__init__(vllm_config=vllm_config, prefix=prefix) + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + """Return no native expert mappings on a P2P Attention rank.""" + if ( + self.afd_config.connector == "CAMP2pAFDConnector" + and self.afd_role == _ATTENTION_ROLE + ): + return [] + return super().get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + if self.afd_config.connector != "CAMP2pAFDConnector": + return super().load_weights(weights) + + from afd_plugin.model_executor.models.deepseek_v4_common import ( + _iter_role_weights as iter_p2p_role_weights, + ) + + return native.AscendDeepseekV4ForCausalLM.load_weights( + self, + iter_p2p_role_weights(weights, role=self.afd_role), + ) + + __all__ = [ "AFDDeepseekV4DecoderLayer", "AFDDeepseekV4ForCausalLM", "AFDDeepseekV4Model", + "AFDNPUDeepseekV4ForCausalLM", ] diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_p2p.py b/afd_plugin/model_executor/models/npu/deepseek_v4_p2p.py new file mode 100644 index 00000000..d0d493f8 --- /dev/null +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_p2p.py @@ -0,0 +1,575 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Ascend A5 AFD wrapper for the native vLLM-Ascend DeepSeek-V4 model. + +The Attention worker retains DSA attention, normalization, and the complete +mHC residual stream. The FFN worker owns the native Ascend V4 MoE. Only the +normalized two-dimensional FFN activation and token-aligned input IDs cross +the synchronous A5 HCCL P2P boundary. + +The connector API intentionally represents input IDs as an optional payload. +That same contract is the integration seam for the planned A5 A2E/E2A +operators; they can replace the P2P transport without changing this model. +""" + +from collections.abc import Iterable +from importlib import import_module +from itertools import islice +from types import ModuleType +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context + +from afd_plugin.config import parse_afd_config +from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata +from afd_plugin.connectors.npu.camp2p_a5 import is_a5 +from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.models.deepseek_v4_common import _iter_role_weights +from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield + + +def _import_native_deepseek_v4() -> ModuleType: + """Load DeepSeek-V4 across the flat and package Ascend layouts.""" + module_name = "vllm_ascend.models.deepseek_v4" + deepseek_v4 = import_module(module_name) + if hasattr(deepseek_v4, "__path__"): + return import_module(f"{module_name}.model") + return deepseek_v4 + + +native = _import_native_deepseek_v4() + + +class RemoteNPUDeepseekV4FFN(nn.Module): + """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" + + def __init__(self, *, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + + def forward( + self, + hidden_states: torch.Tensor, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + if input_ids is None: + raise RuntimeError("DeepSeek-V4 remote FFN requires input_ids") + if input_ids.ndim != 1 or input_ids.shape[0] != hidden_states.shape[0]: + raise ValueError( + "DeepSeek-V4 input_ids must be one-dimensional and token-aligned", + ) + + afd_metadata = get_afd_metadata_from_forward_context() + if afd_metadata is None: + raise RuntimeError("RemoteNPUDeepseekV4FFN requires AFD metadata") + forward_context = get_forward_context() + stage_idx = int( + getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), + ) + afd_metadata.stage_idx = stage_idx + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=self.layer_idx, + stage_idx=stage_idx, + seq_len=int(hidden_states.shape[0]), + ) + context = AFDTransferContext(metadata=metadata) + afd_metadata.connector.send_attn_output( + hidden_states, + context, + input_ids=input_ids, + ) + hidden_states = maybe_apply_dbo_yield( + hidden_states, + role="attention", + ) + return afd_metadata.connector.recv_ffn_output( + ref_tensor=hidden_states, + ubatch_idx=stage_idx, + ) + + +class AFDNPUDeepseekV4DecoderLayer(native.DeepseekV2DecoderLayer): + """Ascend DeepSeek-V4 decoder with an FFN-boundary AFD split.""" + + # Patch reason: the native Ascend layer always allocates Attention and FFN. + # Patch functionality: allocate only the stage owned by the active AFD role. + # Signature: matches upstream; no added parameters. + # Upstream: vllm-ascend/vllm_ascend/models/deepseek_v4/model.py + # Commit: 4fe7ddbf94bc28bcd2b9f3d2d93f0fe1f0499cf5 + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config=None, + topk_indices_buffer: torch.Tensor | None = None, + is_draft_layer: bool = False, + ) -> None: + # ### PATCH START: require a role before allocating native stages. + nn.Module.__init__(self) + afd_config = parse_afd_config(vllm_config, validate=False) + # ### PATCH END + + if config is None: + config = vllm_config.model_config.hf_config + cache_config = vllm_config.cache_config + quant_config = vllm_config.quant_config + parallel_config = vllm_config.parallel_config + + self.hidden_size = config.hidden_size + max_position_embeddings = config.rope_parameters[ + "original_max_position_embeddings" + ] + layer_idx = int(prefix.split(sep=".")[-1]) + self.layer_idx = layer_idx + self.norm_eps = config.rms_norm_eps + + # ### PATCH START: replace the remote stage with a parameter-free proxy. + if afd_config.role == "attention": + self.self_attn = native.DeepseekV4Attention( + vllm_config=vllm_config, + config=config, + max_position_embeddings=max_position_embeddings, + cache_config=cache_config, + quant_config=quant_config, + prefix=f"{prefix}.self_attn", + topk_indices_buffer=topk_indices_buffer, + ) + self.mlp = RemoteNPUDeepseekV4FFN(layer_idx=layer_idx) + elif afd_config.role == "ffn": + self.self_attn = native.PPMissingLayer() + self.mlp = native.DeepseekV4MoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + is_draft_layer=is_draft_layer, + ) + else: + raise ValueError(f"unsupported AFD role {afd_config.role!r}") + # ### PATCH END + + # ### PATCH START: mHC state and normalization are Attention-owned. + if afd_config.role == "ffn": + return + # ### PATCH END + self.input_layernorm = native.RMSNorm( + config.hidden_size, + eps=self.norm_eps, + ) + self.post_attention_layernorm = native.RMSNorm( + config.hidden_size, + eps=self.norm_eps, + ) + self.routed_scaling_factor = getattr( + config, + "routed_scaling_factor", + 1.0, + ) + self.hc_mult = hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + mix_hc = (2 + hc_mult) * hc_mult + hc_dim = hc_mult * config.hidden_size + self.hc_attn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32), + ) + self.hc_ffn_fn = nn.Parameter( + torch.empty(mix_hc, hc_dim, dtype=torch.float32), + ) + self.hc_attn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + ) + self.hc_ffn_base = nn.Parameter( + torch.empty(mix_hc, dtype=torch.float32), + ) + self.hc_attn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + self.hc_ffn_scale = nn.Parameter(torch.empty(3, dtype=torch.float32)) + + # Patch reason: native forward directly invokes its locally allocated FFN. + # Patch functionality: retain native NPU mHC locally while the proxy sends + # only normalized FFN activations and token IDs to the FFN worker. + # Signature: matches upstream; no added parameters. + # Upstream: vllm-ascend/vllm_ascend/models/deepseek_v4/model.py + # Commit: 4fe7ddbf94bc28bcd2b9f3d2d93f0fe1f0499cf5 + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + llama_4_scaling: torch.Tensor | None = None, + input_ids: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + # ### PATCH START: prohibit execution on the FFN worker. + if isinstance(self.self_attn, native.PPMissingLayer): + raise RuntimeError("DeepSeek-V4 decoder forward is Attention-owned") + # ### PATCH END + residual = hidden_states.clone() + hidden_states, post, comb = self.hc_pre( + hidden_states, + self.hc_attn_fn, + self.hc_attn_scale, + self.hc_attn_base, + ) + hidden_states = self.input_layernorm(hidden_states) + attn_kwargs = { + "positions": positions, + "hidden_states": hidden_states, + "llama_4_scaling": llama_4_scaling, + } + hidden_states = self.self_attn(**attn_kwargs) + hidden_states = self.hc_post(hidden_states, residual, post, comb) + residual = hidden_states.clone() + hidden_states, post, comb = self.hc_pre( + hidden_states, + self.hc_ffn_fn, + self.hc_ffn_scale, + self.hc_ffn_base, + ) + hidden_states = self.post_attention_layernorm(hidden_states) + # ### PATCH START: this call enters the synchronous remote FFN proxy. + hidden_states = self.mlp(hidden_states, input_ids) + # ### PATCH END + hidden_states = self.hc_post(hidden_states, residual, post, comb) + return hidden_states, residual + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + *, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + """Execute the native Ascend V4 MoE, including hash routing.""" + if not isinstance(self.mlp, native.DeepseekV4MoE): + raise RuntimeError("DeepSeek-V4 FFN compute is FFN-role only") + if input_ids is None: + raise RuntimeError("DeepSeek-V4 FFN compute requires input_ids") + return self.mlp(hidden_states, input_ids) + + +class AFDNPUDeepseekV4Model(native.DeepseekV4Model): + """Role-aware Ascend DeepSeek-V4 model for the initial A5 P2P route.""" + + # Patch reason: native Ascend V4 allocates all decoder stages on every rank. + # Patch functionality: construct role-aware layers and preserve the + # inherited eager compilation contract. + # Signature: matches upstream; no added parameters. + # Upstream: vllm-ascend/vllm_ascend/models/deepseek_v4/model.py + # Commit: 4fe7ddbf94bc28bcd2b9f3d2d93f0fe1f0499cf5 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + # ### PATCH START: validate the deliberately narrow first release. + nn.Module.__init__(self) + self.afd_config = parse_afd_config(vllm_config, validate=False) + if native.current_platform.device_type != "npu": + raise RuntimeError("AFD NPU DeepSeek-V4 requires an Ascend platform") + if not is_a5(): + raise RuntimeError("AFD NPU DeepSeek-V4 currently supports A5 only") + if self.afd_config.connector != "CAMP2pAFDConnector": + raise RuntimeError( + "AFD NPU DeepSeek-V4 requires CAMP2pAFDConnector", + ) + if self.afd_config.compute_gate_on_attention: + raise RuntimeError( + "AFD NPU DeepSeek-V4 requires gate computation on FFN", + ) + if not vllm_config.model_config.enforce_eager: + raise RuntimeError( + "AFD NPU DeepSeek-V4 A5 P2P requires --enforce-eager", + ) + if vllm_config.speculative_config is not None: + raise RuntimeError( + "AFD NPU DeepSeek-V4 does not yet support speculative decoding", + ) + if vllm_config.lora_config is not None: + raise RuntimeError("AFD NPU DeepSeek-V4 does not yet support LoRA") + parallel_config = vllm_config.parallel_config + if parallel_config.pipeline_parallel_size != 1: + raise RuntimeError("AFD NPU DeepSeek-V4 does not support PP") + if ( + parallel_config.prefill_context_parallel_size != 1 + or parallel_config.decode_context_parallel_size != 1 + ): + raise RuntimeError("AFD NPU DeepSeek-V4 does not support CP") + if parallel_config.use_sequence_parallel_moe: + raise RuntimeError("AFD NPU DeepSeek-V4 does not support SP MoE") + if parallel_config.enable_eplb or parallel_config.enable_elastic_ep: + raise RuntimeError( + "AFD NPU DeepSeek-V4 does not support EPLB or elastic EP", + ) + # The native model is decorated with support_torch_compile. Calling + # nn.Module.__init__ above avoids constructing the unsplit native + # stages, but also bypasses the decorator's initialization wrapper. + # This route requires eager execution, so restore the wrapper state + # with compilation explicitly disabled. + self.vllm_config = vllm_config + self.compilation_config = vllm_config.compilation_config + self.do_not_compile = True + # ### PATCH END + + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.device = native.current_platform.device_type + self.vocab_size = config.vocab_size + self.is_v32 = hasattr(config, "index_topk") + + # ### PATCH START: sparse-index storage is Attention-owned. + if self.afd_config.role == "attention" and self.is_v32: + topk_indices_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + config.index_topk, + dtype=torch.int32, + device=self.device, + ) + else: + topk_indices_buffer = None + self.topk_indices_buffer = topk_indices_buffer + # ### PATCH END + + if native.get_pp_group().is_first_rank: + self.embed_tokens = native.VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = native.PPMissingLayer() + + # ### PATCH START: use the role-aware Ascend decoder constructor. + self.start_layer, self.end_layer, self.layers = native.make_layers( + config.num_hidden_layers, + lambda prefix: AFDNPUDeepseekV4DecoderLayer( + vllm_config, + prefix, + topk_indices_buffer=topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + # ### PATCH END + + if native.get_pp_group().is_last_rank: + self.norm = native.RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + ) + else: + self.norm = native.PPMissingLayer() + + def make_empty_intermediate_tensors( + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> native.IntermediateTensors: + return native.IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, config.hidden_size), + dtype=dtype, + device=device, + ), + }, + ) + + # ### PATCH START: preserve the pre-module-split Ascend PP contract. + make_pp_empty = getattr( + native, + "make_pp_empty_intermediate_tensors", + None, + ) + if make_pp_empty is None: + # vLLM-Ascend before the DSV4 module split exposes the factory + # directly instead of wrapping it for pipeline-parallel models. + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors + else: + self.make_empty_intermediate_tensors = make_pp_empty( + self, + make_empty_intermediate_tensors, + ) + # ### PATCH END + + self.norm_eps = config.rms_norm_eps + self.hc_eps = config.hc_eps + self.hc_mult = hc_mult = config.hc_mult + hc_dim = hc_mult * config.hidden_size + + # ### PATCH START: final mHC state is constructed only on Attention. + if self.afd_config.role == "attention": + self.hc_head_fn = nn.Parameter( + torch.empty(hc_mult, hc_dim, dtype=torch.float32), + ) + self.hc_head_base = nn.Parameter( + torch.empty(hc_mult, dtype=torch.float32), + ) + self.hc_head_scale = nn.Parameter( + torch.empty(1, dtype=torch.float32), + ) + self.hc_norm = native.RMSNorm( + hc_dim, + eps=config.rms_norm_eps, + has_weight=False, + dtype=torch.float32, + ) + else: + self.hc_head_fn = None + self.hc_head_base = None + self.hc_head_scale = None + self.hc_norm = native.PPMissingLayer() + self._mtp_hidden_buffer = None + # ### PATCH END + + # Patch reason: the v0.26 Ascend model forward drops input_ids when it + # invokes decoder layers, so the remote AFD FFN cannot perform hash routing. + # Patch functionality: preserve the pinned native forward while forwarding + # token IDs to each role-aware layer and guarding the disabled MTP buffer. + # Signature: matches upstream; no added parameters. + # Upstream: vllm-ascend/vllm_ascend/models/deepseek_v4.py + # Commit: 80d8c194f7584b17fe08065ea99a130916f6b0e7 + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: native.IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | native.IntermediateTensors: + if native.get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = None + + llama_4_scaling_config = None + llama_4_scaling: torch.Tensor | None + if llama_4_scaling_config is not None: + llama_4_scaling = native._get_llama_4_scaling( + original_max_position_embeddings=llama_4_scaling_config[ + "original_max_position_embeddings" + ], + scaling_beta=llama_4_scaling_config["beta"], + positions=positions, + ) + else: + llama_4_scaling = None + + if native.get_pp_group().is_first_rank: + hidden_states = hidden_states.unsqueeze(1).repeat( + 1, + self.hc_mult, + 1, + ) + aux_hidden_states: list[torch.Tensor] = [] + for layer in islice(self.layers, self.start_layer, self.end_layer): + # ### PATCH START: retain token IDs across the AFD layer boundary. + hidden_states, residual = layer( + positions, + hidden_states, + residual, + llama_4_scaling, + input_ids=input_ids, + ) + # ### PATCH END + if layer.layer_idx + 1 in self.aux_hidden_state_layers: + aux_hidden_states.append(hidden_states.mean(dim=1)) + + # ### PATCH START: AFD rejects speculative decoding and owns no MTP buffer. + if self._mtp_hidden_buffer is not None: + forward_context = get_forward_context() + if forward_context is not None and forward_context.flash_comm_v1_enabled: + h_states_flat = native.tensor_model_parallel_all_gather( + hidden_states.flatten(1), + dim=0, + ) + pad_size = forward_context.pad_size + if pad_size > 0: + h_states_flat = h_states_flat[:-pad_size] + num_tokens = h_states_flat.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_(h_states_flat) + else: + num_tokens = hidden_states.shape[0] + self._mtp_hidden_buffer[:num_tokens].copy_( + hidden_states.flatten(1), + ) + # ### PATCH END + + if not native.get_pp_group().is_last_rank: + return native.IntermediateTensors( + { + "hidden_states": hidden_states, + }, + ) + + hidden_states = self.hc_head( + hidden_states, + self.hc_head_fn, + self.hc_head_scale, + self.hc_head_base, + ) + hidden_states = self.norm(hidden_states) + if aux_hidden_states: + return hidden_states, aux_hidden_states + return hidden_states + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + *, + input_ids: torch.Tensor | None, + ) -> torch.Tensor: + return self.layers[layer_idx].compute_ffn_output( + hidden_states, + input_ids=input_ids, + ) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return tuple(range(int(self.config.num_hidden_layers))) + + +class AFDNPUDeepseekV4ForCausalLM(native.AscendDeepseekV4ForCausalLM): + """Ascend V4 causal LM exposing the NPU FFN-runner model contract.""" + + model_cls = AFDNPUDeepseekV4Model + afd_requires_input_ids = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + self.afd_config = parse_afd_config(vllm_config, validate=False) + self.afd_role = self.afd_config.role + super().__init__(vllm_config=vllm_config, prefix=prefix) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + *, + input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> torch.Tensor: + return self.model.compute_ffn_output( + hidden_states, + layer_idx, + input_ids=input_ids, + ) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return self.model.get_experts_layer_indices() + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + """Return native expert mappings only where real experts are owned.""" + if self.afd_role == "attention": + return [] + return super().get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + return super().load_weights( + _iter_role_weights(weights, role=self.afd_role), + ) + + +__all__ = ["AFDNPUDeepseekV4ForCausalLM"] diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 5e6f1beb..792c8587 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -294,6 +294,11 @@ def _ffn_forward( "AFD model requires input_ids but the connector " "did not return them", ) + # DeepSeek-V4 hash routing in vLLM-Ascend v0.26 reads + # token IDs from ForwardContext rather than the MoE + # forward argument. Publish the received IDs in both + # places so the model contract remains version-stable. + forward_context.input_ids = payload.input_ids compute_kwargs["input_ids"] = payload.input_ids rank_ffn_output = self.model.compute_ffn_output( hidden_states=hidden_states, diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index bc8f90d1..3cd7d403 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -143,9 +143,11 @@ focused unit coverage but no repository model or accuracy E2E case. ### DeepSeek V4 Ascend A5 boundary On Ascend, the same `AFDDeepseekV4ForCausalLM` registry alias resolves to the -backend-local `AFDNPUDeepseekV4ForCausalLM`. It subclasses vLLM-Ascend's -native DeepSeek-V4 implementation, retaining DSA Attention, NPU mHC operators, -and the native `DeepseekV4MoE`. Attention owns all residual-stream and +backend-local `AFDNPUDeepseekV4ForCausalLM`. That wrapper dispatches by +connector: the upstream `deepseek_v4.py` implementation remains responsible +for Async CAM, while `deepseek_v4_p2p.py` owns the A5 P2P model split. Both +paths retain vLLM-Ascend's DSA Attention, NPU mHC operators, and native +`DeepseekV4MoE`. In the P2P path, Attention owns all residual-stream and normalization state and uses `RemoteNPUDeepseekV4FFN`; FFN owns gate, hash router, shared experts, and routed experts. diff --git a/tests/unit/compat/npu/test_dsv4_async_validation.py b/tests/unit/compat/npu/test_dsv4_async_validation.py index 1d74027c..6fb67c02 100644 --- a/tests/unit/compat/npu/test_dsv4_async_validation.py +++ b/tests/unit/compat/npu/test_dsv4_async_validation.py @@ -24,12 +24,21 @@ def _afd_config( ) -def test_dsv4_rejects_camp2p_connector() -> None: - with pytest.raises(RuntimeError, match="only CAMAsyncAFDConnector"): +def test_dsv4_accepts_camp2p_connector() -> None: + _fail_if_unsupported_dsv4_connector( + _afd_config( + compute_gate_on_attention=False, + connector="CAMP2pAFDConnector", + ), + ) + + +def test_dsv4_rejects_unknown_connector() -> None: + with pytest.raises(RuntimeError, match="CAMAsyncAFDConnector or CAMP2p"): _fail_if_unsupported_dsv4_connector( _afd_config( compute_gate_on_attention=False, - connector="CAMP2pAFDConnector", + connector="UnknownConnector", ), ) diff --git a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py index e8e817a5..11bb4980 100644 --- a/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py +++ b/tests/unit/model_executor/models/test_npu_deepseek_v4_contract.py @@ -9,7 +9,12 @@ pytest.importorskip("torch_npu") pytest.importorskip("vllm_ascend") -from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 +from afd_plugin.model_executor.models.npu import ( # noqa: E402 + deepseek_v4 as router, +) +from afd_plugin.model_executor.models.npu import ( # noqa: E402 + deepseek_v4_p2p as adapter, +) native = adapter.native @@ -25,6 +30,24 @@ def recv_ffn_output(self, *, ref_tensor, ubatch_idx): return ref_tensor * 0.5 +class _RecordingDecoderLayer(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.layer_idx = 0 + self.input_ids = None + + def forward( + self, + positions, + hidden_states, + residual, + llama_4_scaling=None, + input_ids=None, + ): + self.input_ids = input_ids + return hidden_states, residual + + def test_npu_v4_wrapper_uses_ascend_native_classes(): assert issubclass( adapter.AFDNPUDeepseekV4DecoderLayer, @@ -39,6 +62,46 @@ def test_npu_v4_wrapper_uses_ascend_native_classes(): assert adapter.AFDNPUDeepseekV4ForCausalLM.afd_requires_input_ids +@pytest.mark.parametrize( + ("connector", "expected_model_cls", "requires_input_ids"), + [ + ("CAMAsyncAFDConnector", router.AFDDeepseekV4Model, False), + ("CAMP2pAFDConnector", adapter.AFDNPUDeepseekV4Model, True), + ], +) +def test_npu_v4_router_selects_model_for_connector( + monkeypatch, + connector, + expected_model_cls, + requires_input_ids, +): + afd_config = SimpleNamespace(connector=connector) + monkeypatch.setattr( + router, + "parse_afd_config", + lambda *_args, **_kwargs: afd_config, + ) + + selected = {} + + def fake_async_init(self, *, vllm_config, prefix=""): + selected["model_cls"] = self.model_cls + selected["requires_input_ids"] = self.afd_requires_input_ids + + monkeypatch.setattr( + router.AFDDeepseekV4ForCausalLM, + "__init__", + fake_async_init, + ) + + router.AFDNPUDeepseekV4ForCausalLM(vllm_config=SimpleNamespace()) + + assert selected == { + "model_cls": expected_model_cls, + "requires_input_ids": requires_input_ids, + } + + def test_npu_v4_native_import_supports_flat_module(monkeypatch): flat_module = ModuleType("vllm_ascend.models.deepseek_v4") monkeypatch.setattr(adapter, "import_module", lambda _name: flat_module) @@ -96,6 +159,47 @@ def test_npu_v4_remote_ffn_sends_input_ids(monkeypatch): assert torch.equal(output, hidden_states * 0.5) +def test_npu_v4_model_forward_preserves_input_ids_without_mtp(monkeypatch): + pp_group = SimpleNamespace(is_first_rank=True, is_last_rank=True) + monkeypatch.setattr(adapter.native, "get_pp_group", lambda: pp_group) + monkeypatch.setattr( + adapter.AFDNPUDeepseekV4Model, + "hc_head", + lambda _self, hidden_states, *_args: hidden_states[:, 0, :], + ) + + model = adapter.AFDNPUDeepseekV4Model.__new__( + adapter.AFDNPUDeepseekV4Model, + ) + torch.nn.Module.__init__(model) + layer = _RecordingDecoderLayer() + model.hc_mult = 1 + model.start_layer = 0 + model.end_layer = 1 + model.layers = torch.nn.ModuleList([layer]) + model.aux_hidden_state_layers = () + model._mtp_hidden_buffer = None + model.hc_head_fn = None + model.hc_head_scale = None + model.hc_head_base = None + model.norm = torch.nn.Identity() + + input_ids = torch.tensor([7, 11], dtype=torch.int32) + positions = torch.tensor([0, 1], dtype=torch.int64) + inputs_embeds = torch.ones((2, 4), dtype=torch.float16) + + output = adapter.AFDNPUDeepseekV4Model.forward( + model, + input_ids, + positions, + intermediate_tensors=None, + inputs_embeds=inputs_embeds, + ) + + assert layer.input_ids is input_ids + assert torch.equal(output, inputs_embeds) + + def test_npu_v4_model_requires_eager_a5_p2p(monkeypatch): afd_config = SimpleNamespace( compute_gate_on_attention=False, diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index b67bb797..0b082c19 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -1418,7 +1418,24 @@ def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): def test_npu_ffn_runner_requests_and_forwards_v4_input_ids(monkeypatch): - _patch_ffn_forward_context(monkeypatch) + _require_npu_runtime() + from afd_plugin.v1.worker.npu import ffn_model_runner + + forward_context = SimpleNamespace( + additional_kwargs={}, + dp_metadata=None, + all_moe_layers={}, + ) + + @contextmanager + def fake_ascend_forward_context(**_kwargs): + yield forward_context + + monkeypatch.setattr( + ffn_model_runner, + "ascend_forward_context", + fake_ascend_forward_context, + ) runner = _new_ffn_runner() runner.vllm_config = _vllm_config(role="ffn") runner.connector = _FakeFFNConnector() @@ -1438,6 +1455,7 @@ def test_npu_ffn_runner_requests_and_forwards_v4_input_ids(monkeypatch): runner.execute_model(dp_metadata_list={0: _FakeDPMetadata([2])}) + assert forward_context.input_ids == "token-ids" assert runner.connector.recv_calls == [ ( 0,