diff --git a/afd_plugin/__init__.py b/afd_plugin/__init__.py index f7d81451..42b3781d 100644 --- a/afd_plugin/__init__.py +++ b/afd_plugin/__init__.py @@ -6,6 +6,8 @@ import importlib.util import logging +import multiprocessing +import os from importlib.metadata import PackageNotFoundError, version from pathlib import Path from types import MappingProxyType @@ -45,6 +47,51 @@ def __getattr__(name: str): _logger = logging.getLogger(__name__) _registered = False + +def _force_spawn_multiprocessing_if_requested() -> None: + """Pin Python's default context for the A3 Python 3.12 runtime. + + ``VLLM_WORKER_MULTIPROC_METHOD`` controls vLLM's explicit context only. + Some NPU runtime helpers use Python's default context and otherwise retain + the task's inherited ``forkserver`` setting, which cannot restore the + signal-handler state in this environment. Keep the override opt-in so it + remains isolated to the affected CAM async recipe. + """ + if os.environ.get("AFD_FORCE_SPAWN_MULTIPROCESSING") != "1": + return + multiprocessing.set_start_method("spawn", force=True) + + # A few optional runtime helpers request ``forkserver`` explicitly instead + # of consulting the default start method. On the A3 Python 3.12 image the + # forkserver cannot restore its inherited signal-handler state, so every + # child it starts exits before running user code. Keep this narrowly + # opt-in with the recipe environment variable above: callers requesting a + # forkserver receive the equivalent spawn context instead. + contexts = multiprocessing.context._concrete_contexts + contexts["forkserver"] = contexts["spawn"] + + # TE Fusion keeps a module reference to ``multiprocessing`` and calls its + # public ``get_context("forkserver")`` API directly inside each model + # worker. Replacing the registry alone is not sufficient for all Python + # 3.12 context instances, so redirect that explicit request as well. + if not getattr(multiprocessing, "_afd_spawn_context_redirect", False): + original_get_context = multiprocessing.get_context + + def get_spawn_context(method: str | None = None): + if method == "forkserver": + method = "spawn" + return original_get_context(method) + + multiprocessing.get_context = get_spawn_context + multiprocessing._afd_spawn_context_redirect = True + + +# vLLM model workers import the configured worker class directly in spawned +# child interpreters; they do not invoke the general-plugin entry point below. +# Apply the opt-in setting at package import time so it is also in effect before +# Ascend TE Fusion initializes its compilation workers. +_force_spawn_multiprocessing_if_requested() + _DEEPSEEK_MODEL_REGISTRATIONS = { "DeepseekForCausalLM": ( "afd_plugin.model_executor.models.deepseek_v2:AFDDeepseekForCausalLM" @@ -59,7 +106,9 @@ def __getattr__(name: str): "afd_plugin.model_executor.models.deepseek_v2:AFDDeepseekV3ForCausalLM" ), "DeepseekV4ForCausalLM": ( - "afd_plugin.model_executor.models.deepseek_v4:AFDDeepseekV4ForCausalLM" + "afd_plugin.model_executor.models.npu.deepseek_v4:AFDDeepseekV4ForCausalLM" + if importlib.util.find_spec("torch_npu") is not None + else "afd_plugin.model_executor.models.deepseek_v4:AFDDeepseekV4ForCausalLM" ), "GlmMoeDsaForCausalLM": ( "afd_plugin.model_executor.models.deepseek_v2:AFDGlmMoeDsaForCausalLM" @@ -102,6 +151,7 @@ def register_afd() -> None: _logger.debug("AFD plugin: register_afd() already completed") return + _force_spawn_multiprocessing_if_requested() _logger.debug("AFD plugin: register_afd() called") if importlib.util.find_spec("vllm") is None: _logger.debug("AFD plugin: vLLM not found, skipping runtime registration") diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 823144b7..d95592d9 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -40,6 +40,8 @@ def fail_if_unsupported_npu_afd_features( afd_config, extra_info, ) + if _is_dsv4_target(vllm_config): + _fail_if_unsupported_dsv4_async_features(afd_config, extra_info) return if afd_config.compute_gate_on_attention: @@ -84,6 +86,45 @@ def fail_if_unsupported_npu_afd_features( ) +def _is_dsv4_target(vllm_config: VllmConfig) -> bool: + """Return whether the target model is DeepSeek V4.""" + model_config = vllm_config.model_config + hf_config = getattr(model_config, "hf_config", None) + if hf_config is None: + hf_config = getattr(model_config, "hf_text_config", None) + if hf_config is None: + return False + if getattr(hf_config, "model_type", None) == "deepseek_v4": + return True + architectures = getattr(hf_config, "architectures", ()) or () + return any("DeepseekV4" in str(architecture) for architecture in architectures) + + +def _fail_if_unsupported_dsv4_async_features( + afd_config: AFDConfig, + extra_info: ConnectorExtraInfo, +) -> None: + """Keep the initial DSV4 Async CAM target path deliberately narrow.""" + from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + + if not afd_config.compute_gate_on_attention: + raise RuntimeError( + "DSV4 CAMAsyncAFDConnector requires compute_gate_on_attention=true", + ) + if not isinstance(extra_info, AFDAsyncExtraInfo): + raise TypeError( + "DSV4 CAMAsyncAFDConnector requires AFDAsyncExtraInfo, got " + f"{type(extra_info).__name__}", + ) + if extra_info.dynamic_quant != 1: + raise RuntimeError( + "DSV4 Flash-INT8 CAMAsyncAFDConnector requires dynamicQuant=1", + ) + # async_moe_ubatching is validated by + # _fail_if_unsupported_npu_async_moe_ubatching_features for all async CAM + # targets, so DSV4 no longer needs a bespoke rejection here. + + def _fail_if_unsupported_npu_afd_async_features( vllm_config: VllmConfig, afd_config: AFDConfig, @@ -123,6 +164,43 @@ def _fail_if_unsupported_npu_afd_async_features( raise RuntimeError( "CAMAsyncAFDConnector currently supports only dynamicQuant 0 or 1", ) + _validate_cam_world_topology(vllm_config, afd_config, extra_info) + + +def _validate_cam_world_topology( + vllm_config: VllmConfig, + afd_config: AFDConfig, + extra_info: ConnectorExtraInfo, +) -> None: + """Require each role's local layout to fill the one CAM world.""" + from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + + if not isinstance(extra_info, AFDAsyncExtraInfo): + return + parallel_config = vllm_config.parallel_config + attn_ranks_per_dp = int(extra_info.attn_ranks_per_dp) + if afd_config.role == "attention": + if int(parallel_config.tensor_parallel_size) != attn_ranks_per_dp: + raise RuntimeError( + "CAMAsyncAFDConnector Attention tensor_parallel_size must equal " + "attn_ranks_per_dp", + ) + local_world_size = ( + int(parallel_config.data_parallel_size) * attn_ranks_per_dp + ) + expected_world_size = afd_config.num_attention_ranks + else: + local_world_size = ( + int(parallel_config.data_parallel_size) + * int(parallel_config.tensor_parallel_size) + ) + expected_world_size = afd_config.num_ffn_ranks + if local_world_size != expected_world_size: + raise RuntimeError( + "CAMAsyncAFDConnector " + f"{afd_config.role} DPxTP world size must equal its configured " + f"role size, got {local_world_size} and {expected_world_size}", + ) def _fail_if_unsupported_npu_async_moe_ubatching_features( diff --git a/afd_plugin/compat/npu/ops.py b/afd_plugin/compat/npu/ops.py index f2ea057b..1437a853 100644 --- a/afd_plugin/compat/npu/ops.py +++ b/afd_plugin/compat/npu/ops.py @@ -4,6 +4,7 @@ from __future__ import annotations +import ctypes import os from functools import lru_cache from pathlib import Path @@ -16,6 +17,10 @@ CAM_DISPATCH_RECV = "async_dispatch_recv" CAM_COMBINE_SEND = "async_combine_send" CAM_COMBINE_RECV = "async_combine_recv" +CAM_CUST_OPAPI_ENV = "CAM_CUST_OPAPI_LIB_PATH" +CAM_CUST_OPAPI_DEFAULT = Path( + "/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so", +) def get_afd_cann_vendor_path() -> Path: @@ -68,6 +73,24 @@ def _assert_cam_namespace_registered(torch: object) -> None: ) +def _preload_cam_cust_opapi() -> None: + """Make CAM's custom op-api symbols visible to the Python extension. + + ``umdk_cam_op_lib`` resolves the aclnn CAM entry points by name while its + operators are first invoked. CANN's stock ``libopapi.so`` does not export + those symbols; they live in CAM's ``libopapi.so``. Load that library + globally before importing the extension so the real CAM operators bind to + the vendor implementation instead of the stock op-api library. + """ + + library = Path(os.environ.get(CAM_CUST_OPAPI_ENV, CAM_CUST_OPAPI_DEFAULT)) + if not library.is_file(): + raise RuntimeError( + "CAMAsyncAFDConnector requires CAM libopapi.so; expected " + f"{library}. Install the CAM operator package first.", + ) + ctypes.CDLL(str(library), mode=ctypes.RTLD_GLOBAL) + @lru_cache(maxsize=1) def ensure_cam_p2p_ops_available() -> None: """Import the custom operators used by ``CAMP2pAFDConnector``. @@ -107,6 +130,7 @@ def ensure_cam_async_ops_available() -> None: """Ensure the runtime exposes the real CAM async operator namespace.""" try: + _preload_cam_cust_opapi() import torch import torch_npu # noqa: F401 import umdk_cam_op_lib # noqa: F401 diff --git a/afd_plugin/compat/npu/runtime.py b/afd_plugin/compat/npu/runtime.py index 30297064..8b6c11c7 100644 --- a/afd_plugin/compat/npu/runtime.py +++ b/afd_plugin/compat/npu/runtime.py @@ -69,7 +69,6 @@ def apply_afd_ascend_patches_if_needed() -> None: from afd_plugin.compat.patches.npu.mla_graph import ( apply_afd_mla_graph_patch, ) - apply_afd_ascend_config_patch_if_needed() if not apply_afd_mla_graph_patch(): raise RuntimeError( diff --git a/afd_plugin/compat/patches/engine_core.py b/afd_plugin/compat/patches/engine_core.py index 2e689728..70f93f33 100644 --- a/afd_plugin/compat/patches/engine_core.py +++ b/afd_plugin/compat/patches/engine_core.py @@ -240,6 +240,8 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig: return _AFDFFNKVCacheConfig() # ### PATCH END: AFD FFN late-loaded KV cache bypass + import vllm_ascend.patch.platform.patch_kv_cache_utils # noqa: F401 + start = time.time() core_module.register_all_kvcache_specs(vllm_config) diff --git a/afd_plugin/config.py b/afd_plugin/config.py index 256590ab..79ae4713 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -328,7 +328,6 @@ def validate_afd_config( f"num_ffn_ranks must be positive, got {config.num_ffn_ranks}", ) - __all__ = [ "AFDConfig", "AFD_ASYNC_CONNECTOR", diff --git a/afd_plugin/connectors/metadata.py b/afd_plugin/connectors/metadata.py index 6c982d71..b5bfb079 100644 --- a/afd_plugin/connectors/metadata.py +++ b/afd_plugin/connectors/metadata.py @@ -128,6 +128,7 @@ class AFDControlPayload: dp_metadata_list: dict[int, AFDDPMetadata] is_graph_capturing: bool is_warmup: bool + input_ids_by_stage: dict[int, list[int]] = field(default_factory=dict) def __post_init__(self) -> None: self.dp_metadata_list = { @@ -139,6 +140,14 @@ def __post_init__(self) -> None: else dp_metadata for stage_idx, dp_metadata in self.dp_metadata_list.items() } + self.input_ids_by_stage = { + int(stage_idx): torch.as_tensor( + input_ids, + dtype=torch.int64, + device="cpu", + ).flatten().tolist() + for stage_idx, input_ids in self.input_ids_by_stage.items() + } @dataclass(slots=True) @@ -319,10 +328,14 @@ def encode_control_payload(payload: AFDControlPayload) -> bytes: """ metadata_payload: dict[str, dict[str, int | list[int]]] = {} for stage_idx, dp_metadata in payload.dp_metadata_list.items(): - metadata_payload[str(int(stage_idx))] = { + stage_payload: dict[str, int | list[int]] = { "num_tokens_across_dp_cpu": dp_metadata.num_tokens_across_dp_cpu.tolist(), "max_tokens_across_dp_cpu": _to_int(dp_metadata.max_tokens_across_dp_cpu), } + input_ids = payload.input_ids_by_stage.get(int(stage_idx)) + if input_ids is not None: + stage_payload["input_ids"] = list(input_ids) + metadata_payload[str(int(stage_idx))] = stage_payload wire_payload = { "dp_metadata_list": metadata_payload, @@ -352,10 +365,16 @@ def decode_control_payload(payload_bytes: bytes) -> AFDControlPayload: ) for stage_idx, metadata in payload["dp_metadata_list"].items() } + input_ids_by_stage = { + int(stage_idx): [int(value) for value in metadata["input_ids"]] + for stage_idx, metadata in payload["dp_metadata_list"].items() + if "input_ids" in metadata + } return AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=bool(payload.get("is_graph_capturing", False)), is_warmup=bool(payload.get("is_warmup", False)), + input_ids_by_stage=input_ids_by_stage, ) diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index 00dd128a..28109618 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -349,7 +349,24 @@ def recv_ffn_work_item( "TokenNums_Rankid_Layeridx from async_dispatch_recv", ) total_num_tokens = max(1, int(token_nums_rankid_layeridx[0].item())) - layer_idx = int(token_nums_rankid_layeridx[2].item()) + # A zero in the header's leading count is valid for an FFN rank that + # received no routed tokens in this dispatch. Such ranks must still + # enter combine-send (the zero-token fallback below supplies its + # placeholder) so every participant completes the CAM collective. + # CAM 209 currently returns zero in the documented layer-index slot + # for every dispatch-recv result. The connector-driven FFN runner + # already has the authoritative decoder-layer order, so consuming this + # stale field makes every DSV4 work item execute FFN layer 0. CAM + # combine-send also uses this header to match the completion to its + # original dispatch, so merely fixing the local model call is not + # sufficient: all later layers would still be combined as layer 0. + # CAM returns the true decoder-layer index in the dispatch-recv header + # (field 2 of token_nums_rankid_layeridx). Use it for both the local + # FFN layer computation and the combine-send header so the completion + # matches the attention-side combine-recv. + layer_idx = max(0, int(token_nums_rankid_layeridx[2].item())) + token_nums_rankid_layeridx = token_nums_rankid_layeridx.clone() + states.token_nums_rankid_layeridx = token_nums_rankid_layeridx expert_token_nums_shared = states.expert_token_nums_shared if expert_token_nums_shared is None: @@ -802,7 +819,6 @@ def send_ffn_output( self.tp_size, self.group_name, ) - def _require_initialized(self) -> None: if not self._initialized: raise RuntimeError("CAMAsyncAFDConnector is not initialized") @@ -856,7 +872,7 @@ def build_async_topology( *, num_routed_experts: int | None = None, ) -> AFDAsyncTopology: - """Validate role-local rank settings and derive the CAM HCCL world rank. + """Validate role-local rank settings and derive HCCL world rank. The world is Attention-first: Attention role rank ``i`` maps to world rank ``i`` and FFN role rank ``j`` maps to diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index 32f25424..40ea5c4c 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -181,9 +181,9 @@ class _CAMP2PTopology: """Describe where one connector process sits in the communication groups. ``world_rank`` is the process number in the complete AFD group. - ``p2p_rank`` is its number in the smaller Gloo group used to exchange metadata. - For an Attention rank, ``dp_metadata_destinations`` lists the FFN ranks - that should receive its metadata. + ``p2p_rank`` is its number in the Gloo group used to exchange metadata. + Every Attention rank participates in that group so an FFN rank can receive + the token IDs for every Attention hidden-state slice it aggregates. """ role: str @@ -194,21 +194,27 @@ class _CAMP2PTopology: ffn_size: int min_size: int dp_metadata_destinations: tuple[int, ...] + dp_metadata_sources: tuple[int, ...] @property def p2p_world_size(self) -> int: - """Return the number of FFN and participating Attention metadata ranks.""" - return self.ffn_size + self.min_size + """Return the number of FFN and Attention metadata ranks.""" + return self.ffn_size + self.attention_size @property def participates_in_p2p_group(self) -> bool: """Return whether this process joins the Gloo DP-metadata group.""" - return self.world_rank < self.ffn_size or self.is_attn_top_min_size_rank + return True @property def is_attn_top_min_size_rank(self) -> bool: - """Return whether this is an Attention metadata-sender rank.""" - return self.ffn_size <= self.world_rank < self.ffn_size + self.min_size + """Return whether this is an Attention metadata-sender rank. + + The property name is retained for connector compatibility. Every + Attention rank sends one payload to the FFN rank that receives its + hidden-state slice. + """ + return self.role == "attention" class CAMP2pAFDConnector(AFDConnectorBase): @@ -265,6 +271,7 @@ def __init__( self.ratio = self.attn_size // self.ffn_size self.dst_list = list(self.topology.dp_metadata_destinations) self.dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata] = {} + self.input_ids_by_stage: dict[int, list[int]] = {} self.is_graph_capturing = False self.is_warmup = False self.scheduler_config = vllm_config.scheduler_config @@ -383,6 +390,7 @@ def close(self) -> None: dist.destroy_process_group(group) self.p2p_pg = None self.ffn_pg = None + self.input_ids_by_stage = {} self.afd_pg = None self.afd_pg_list = [] self.hccl_comm_name = "" @@ -636,6 +644,7 @@ def update_state_from_dp_metadata( ) -> None: connector = self.connector connector.dp_metadata_list = payload.dp_metadata_list + connector.input_ids_by_stage = payload.input_ids_by_stage connector.is_graph_capturing = payload.is_graph_capturing connector.is_warmup = payload.is_warmup @@ -662,12 +671,68 @@ def recv_dp_metadata_list(self) -> AFDControlPayload: connector = self.connector if connector.p2p_pg is None: raise RuntimeError("CAMP2P metadata process group is not initialized") - src = connector.p2p_rank % connector.min_size + connector.ffn_size - return recv_control_payload( - src=src, - group=connector.p2p_pg, - device=torch.device("cpu"), + payloads = tuple( + recv_control_payload( + src=src, + group=connector.p2p_pg, + device=torch.device("cpu"), + ) + for src in connector.topology.dp_metadata_sources ) + return _aggregate_camp2p_control_payloads(payloads) + + +def _aggregate_camp2p_control_payloads( + payloads: tuple[AFDControlPayload, ...], +) -> AFDControlPayload: + """Concatenate peer token IDs in CAMP2P hidden-state receive order. + + CAMP2P aggregates consecutive Attention role ranks on each FFN rank. The + control payload follows the same order, so DSV4 hash routing sees one + token ID per received hidden-state row, including every DBO ubatch stage. + """ + if not payloads: + raise RuntimeError("CAMP2P FFN control plane received no Attention payloads") + + reference = payloads[0] + stage_ids = tuple(sorted(reference.dp_metadata_list)) + for payload in payloads[1:]: + if tuple(sorted(payload.dp_metadata_list)) != stage_ids: + raise RuntimeError( + "CAMP2P Attention peers sent different control-plane stages" + ) + if ( + payload.is_graph_capturing != reference.is_graph_capturing + or payload.is_warmup != reference.is_warmup + ): + raise RuntimeError( + "CAMP2P Attention peers sent inconsistent control-plane flags" + ) + + input_ids_by_stage: dict[int, list[int]] = {} + for stage_idx in stage_ids: + stage_input_ids = tuple( + payload.input_ids_by_stage.get(stage_idx) for payload in payloads + ) + if all(input_ids is None for input_ids in stage_input_ids): + continue + if any(input_ids is None for input_ids in stage_input_ids): + raise RuntimeError( + "CAMP2P Attention peers must either all provide input_ids or " + f"all omit them for stage {stage_idx}" + ) + input_ids_by_stage[stage_idx] = [ + token_id + for input_ids in stage_input_ids + for token_id in input_ids + ] + + return AFDControlPayload( + dp_metadata_list=reference.dp_metadata_list, + is_graph_capturing=reference.is_graph_capturing, + is_warmup=reference.is_warmup, + input_ids_by_stage=input_ids_by_stage, + ) def build_camp2p_topology( @@ -677,8 +742,9 @@ def build_camp2p_topology( """Calculate the communication rank numbers for one process. FFN processes come first in the main AFD group, followed by Attention - processes. All FFN ranks and the first ``min(A, F)`` Attention ranks also - join the smaller Gloo group that exchanges token counts and batch details. + processes. The Gloo control group contains every process: an FFN rank + receives one payload from each Attention rank whose hidden states CAMP2P + aggregates on that FFN rank. Args: afd_config: Process role and total Attention/FFN rank counts. @@ -698,6 +764,11 @@ def build_camp2p_topology( "CAMP2P requires attention_size >= ffn_size, got " f"{attention_size} < {ffn_size}", ) + if attention_size % ffn_size != 0: + raise ValueError( + "CAMP2P requires attention_size to be an integer multiple of " + f"ffn_size, got attention_size={attention_size}, ffn_size={ffn_size}", + ) if role_rank < 0: raise ValueError(f"CAMP2P role rank must be non-negative, got {role_rank}") @@ -708,7 +779,7 @@ def build_camp2p_topology( f"(rank={role_rank}, size={attention_size})", ) world_rank = ffn_size + role_rank - p2p_rank = role_rank + min(ffn_size, attention_size) + p2p_rank = world_rank elif afd_config.role == "ffn": if role_rank >= ffn_size: raise ValueError( @@ -721,13 +792,20 @@ def build_camp2p_topology( raise ValueError(f"unknown AFD role {afd_config.role!r}") min_size = min(attention_size, ffn_size) - destinations: list[int] = [] - if ffn_size <= world_rank < ffn_size + min_size: - local_attention_rank = world_rank - ffn_size - dst = local_attention_rank - while dst < ffn_size: - destinations.append(dst) - dst += min_size + attention_ranks_per_ffn = attention_size // ffn_size + destinations: tuple[int, ...] = () + sources: tuple[int, ...] = () + if afd_config.role == "attention": + destinations = (role_rank // attention_ranks_per_ffn,) + else: + attention_rank_start = role_rank * attention_ranks_per_ffn + sources = tuple( + ffn_size + attention_rank + for attention_rank in range( + attention_rank_start, + attention_rank_start + attention_ranks_per_ffn, + ) + ) return _CAMP2PTopology( role=afd_config.role, @@ -737,7 +815,8 @@ def build_camp2p_topology( attention_size=attention_size, ffn_size=ffn_size, min_size=min_size, - dp_metadata_destinations=tuple(destinations), + dp_metadata_destinations=destinations, + dp_metadata_sources=sources, ) diff --git a/afd_plugin/distributed/afd_process_group.py b/afd_plugin/distributed/afd_process_group.py index 24102973..cbbcbd01 100644 --- a/afd_plugin/distributed/afd_process_group.py +++ b/afd_plugin/distributed/afd_process_group.py @@ -42,12 +42,13 @@ def __exit__(self, exc_type: object, exc_value: object, tb: object) -> None: def init_afd_process_group( *, backend: str, - init_method: str, + init_method: str | None = None, world_size: int, rank: int, group_name: str, timeout: timedelta, pg_options: Any | None = None, + store: Any | None = None, ) -> ProcessGroup: """Create a plugin-owned process group without patching vLLM source. @@ -56,13 +57,16 @@ def init_afd_process_group( unavailable. """ - rendezvous_iterator = rendezvous( - init_method, - rank, - world_size, - timeout=timeout, - ) - store, rank, world_size = next(rendezvous_iterator) + if store is None: + if init_method is None: + raise ValueError("init_method is required when store is not provided") + rendezvous_iterator = rendezvous( + init_method, + rank, + world_size, + timeout=timeout, + ) + store, rank, world_size = next(rendezvous_iterator) store.set_timeout(timeout) prefixed_store = PrefixStore(group_name, store) backend_value = Backend(backend) if backend else Backend("undefined") diff --git a/afd_plugin/model_executor/models/model_utils.py b/afd_plugin/model_executor/models/model_utils.py index 17e10443..48b329ae 100644 --- a/afd_plugin/model_executor/models/model_utils.py +++ b/afd_plugin/model_executor/models/model_utils.py @@ -19,6 +19,8 @@ def get_afd_model_config( device_type: Literal["cuda", "npu"], ) -> ModelConfig: """Return a model config that resolves to an AFD model implementation.""" + from afd_plugin import register_afd + register_afd() for model_arch in model_config.hf_config.architectures: if model_arch in _MODEL_REGISTRATIONS: diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py new file mode 100644 index 00000000..23901e77 --- /dev/null +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -0,0 +1,853 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek V4 AFD model wrapper for the Ascend native implementation. + +The Ascend DSV4 implementation is not a DeepSeek V2 layer with different +dimensions. Its decoder layer owns the DSA attention path and the +hyper-connection (HC) state transitions around both attention and FFN. This +adapter therefore keeps the native layer forward path on the Attention role +and replaces only the FFN module with the AFD remote proxy. The FFN role +constructs the native MoE module and exposes it through the runner-facing +``compute_ffn_output`` hook. +""" + +from collections.abc import Iterable, Iterator +from copy import copy +from itertools import islice +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.distributed.parallel_state import get_tp_group +from vllm.forward_context import get_forward_context, override_forward_context +from vllm.model_executor.layers import fused_moe +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.sequence import IntermediateTensors + +from afd_plugin.config import AFD_ASYNC_CONNECTOR, parse_afd_config +from afd_plugin.connectors import ( + AFDExpertRoutingSpec, + AFDF2ATransferPayload, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.models.deepseek_v2 import RemoteFFNProxy +from afd_plugin.model_executor.models.npu.async_cam_layout import ( + AsyncMoeUbatchMetadata, + CAMDispatchLayout, + build_async_moe_stage_inputs, + get_async_moe_ubatch_metadata_from_forward_context, + prepare_cam_dispatch_payload, + restore_async_moe_stage_outputs, + restore_cam_dispatch_output, +) + +try: + from vllm_ascend.models import deepseek_v4 as native +except ImportError as exc: # pragma: no cover - only reachable off Ascend. + raise ImportError( + "DSV4 AFD support requires the vLLM-Ascend native DSV4 model" + ) from exc + + +_ATTENTION_ROLE = "attention" +_FFN_ROLE = "ffn" +_BOTH_ROLES = frozenset((_ATTENTION_ROLE, _FFN_ROLE)) + +def _refresh_ascend_fused_moe() -> None: + """Bind native DSV4 MoE construction to the Ascend implementation.""" + # vLLM-Ascend applies this replacement during platform initialization, + # while the DSV4 module keeps a module-level FusedMoE binding. Refresh it + # before either AFD role constructs its local model. + if native.current_platform.device_type == "npu": + native.FusedMoE = fused_moe.FusedMoE + + +def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: + 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], tuple(parts[marker_idx + 3 :]) + return None + + +def _checkpoint_weight_roles(name: str) -> frozenset[str]: + """Return the AFD owner for a DSV4 checkpoint path. + + DSV4 checkpoints use ``attn``/``ffn`` names while the Ascend runtime + model exposes ``self_attn``/``mlp``. The native loader performs that name + conversion later, so filtering must understand both spellings here. + """ + + layer_path = _weight_layer_path(name) + if layer_path is None: + return _BOTH_ROLES + + _, stage, remainder = layer_path + if stage in ("attn", "self_attn"): + return frozenset((_ATTENTION_ROLE,)) + if stage in ("ffn", "mlp"): + if remainder and remainder[0] == "gate": + return _BOTH_ROLES + return frozenset((_FFN_ROLE,)) + # HC parameters and any future shared layer parameters are required by + # both role-local model instances. + return _BOTH_ROLES + + +def _iter_role_weights( + weights: Iterable[tuple[str, torch.Tensor]], + *, + role: str, +) -> Iterator[tuple[str, torch.Tensor]]: + for name, loaded_weight in weights: + if role in _checkpoint_weight_roles(name): + yield name, loaded_weight + + +class AFDDeepseekV4AttentionGateRemoteMoE(RemoteFFNProxy): + """DSV4 gate shell that routes local Attention tokens through Async CAM.""" + + def __init__( + self, + *, + config: Any, + layer_idx: int, + prefix: str, + ) -> None: + super().__init__(layer_idx=layer_idx) + self.top_k = int(config.num_experts_per_tok) + self.n_routed_experts = int(config.n_routed_experts) + self.renormalize = bool(config.norm_topk_prob) + self.scoring_func = getattr(config, "scoring_func", "softmax") + self.num_expert_group = int(getattr(config, "n_group", 1)) + self.topk_group = int(getattr(config, "topk_group", 1)) + self.routed_scaling_factor = float( + getattr(config, "routed_scaling_factor", 1.5), + ) + self.gate = ReplicatedLinear( + config.hidden_size, + config.n_routed_experts, + bias=False, + quant_config=None, + prefix=f"{prefix}.gate", + ) + self.gate.precast_fp32_weight = True + if layer_idx < config.num_hash_layers: + self.gate.tid2eid = nn.Parameter( + torch.zeros( + config.vocab_size, + config.num_experts_per_tok, + dtype=torch.int32, + ), + requires_grad=False, + ) + self.gate.e_score_correction_bias = None + else: + self.gate.tid2eid = None + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=torch.float32), + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Select DSV4 experts on Attention and exchange routed work via CAM.""" + from afd_plugin.model_executor.models.npu import deepseek_v4_attention_gate + + topk_weights, topk_ids = deepseek_v4_attention_gate.compute_attention_gate_topk( + self, + hidden_states, + ) + dispatch_payload = prepare_cam_dispatch_payload( + hidden_states, + topk_weights, + topk_ids, + None, + use_sequence_parallel=get_forward_context().flash_comm_v1_enabled, + ) + output = self._send_and_receive( + dispatch_payload.hidden_states, + topk_weights=dispatch_payload.topk_weights, + topk_ids=dispatch_payload.topk_ids, + ) + return restore_cam_dispatch_output(output, dispatch_payload.layout) + + +class AFDDeepseekV4DecoderLayer(native.DeepseekV2DecoderLayer): + """Role-local DSV4 decoder layer. + + The inherited native ``forward`` is intentionally retained. On the + Attention role it executes the full native HC/DSA sequence and the + ``RemoteFFNProxy`` makes the AFD transfer at exactly the native FFN + boundary. The FFN role is connector-driven and does not call this full + forward method; it invokes ``compute_ffn_output`` instead. + """ + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str, + config=None, + topk_indices_buffer: torch.Tensor | None = None, + is_draft_layer: bool = False, + ) -> None: + if is_draft_layer: + raise ValueError("AFD DSV4 decoder layers do not support draft layers") + afd_config = parse_afd_config(vllm_config, validate=False) + nn.Module.__init__(self) + 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 + layer_idx = int(prefix.split(sep=".")[-1]) + + self.vllm_config = vllm_config + self.config = config + self.afd_role = afd_config.role + self.layer_idx = layer_idx + self.hidden_size = config.hidden_size + self.norm_eps = config.rms_norm_eps + self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0) + self.hc_mult = config.hc_mult + self.hc_sinkhorn_iters = config.hc_sinkhorn_iters + self.hc_eps = config.hc_eps + self.compute_gate_on_attention = bool(afd_config.compute_gate_on_attention) + self.is_moe_layer = True + self.use_sequence_parallel_moe = False + + max_position_embeddings = config.rope_parameters[ + "original_max_position_embeddings" + ] + if afd_config.role == _ATTENTION_ROLE: + 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, + ) + if self.compute_gate_on_attention: + self.mlp = AFDDeepseekV4AttentionGateRemoteMoE( + config=config, + layer_idx=layer_idx, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = RemoteFFNProxy(layer_idx=layer_idx) + elif afd_config.role == _FFN_ROLE: + self.self_attn = native.PPMissingLayer() + _refresh_ascend_fused_moe() + self.mlp = native.DeepseekV4MoE( + config=config, + parallel_config=parallel_config, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + is_draft_layer=False, + ) + else: # pragma: no cover - parse_afd_config validates role values. + raise ValueError(f"Unsupported AFD role: {afd_config.role!r}") + + 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, + ) + mix_hc = (2 + self.hc_mult) * self.hc_mult + hc_dim = self.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)) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + *, + group_list: torch.Tensor | None = None, + dynamic_scales: torch.Tensor | None = None, + expand_x_shared: torch.Tensor | None = None, + dynamic_scales_shared: torch.Tensor | None = None, + topk_scales: torch.Tensor | None = None, + group_list_type: int = 1, + **_: Any, + ) -> torch.Tensor | AFDF2ATransferPayload: + if not isinstance(self.mlp, native.DeepseekV4MoE): + raise RuntimeError( + "DSV4 compute_ffn_output requires the FFN role to own the MoE" + ) + if self.compute_gate_on_attention: + if group_list is None: + raise RuntimeError( + "DSV4 Attention-side routing requires CAM group_list on FFN", + ) + from afd_plugin.model_executor.models.npu import ( + deepseek_v2_attention_gate, + ) + + return deepseek_v2_attention_gate.compute_attention_gate_moe_ffn( + self, + hidden_states=hidden_states, + group_list=group_list, + dynamic_scales=dynamic_scales, + expand_x_shared=expand_x_shared, + dynamic_scales_shared=dynamic_scales_shared, + topk_scales=topk_scales, + group_list_type=group_list_type, + ) + if self.mlp.hash: + input_ids = get_forward_context().input_ids + num_tokens = hidden_states.reshape(-1, hidden_states.shape[-1]).shape[0] + if input_ids is None: + raise RuntimeError( + "DSV4 hash routing requires input_ids from the CAMP2P " + "control plane" + ) + if input_ids.numel() != num_tokens: + raise RuntimeError( + "DSV4 hash routing input_ids do not align with FFN hidden " + f"states: got {input_ids.numel()} IDs for {num_tokens} rows" + ) + return self.mlp(hidden_states) + + +@native.support_torch_compile +class AFDDeepseekV4Model(native.DeepseekV4Model): + """DSV4 model with role-local Attention/FFN allocations.""" + + fall_back_to_pt_during_load = False + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + afd_config = parse_afd_config(vllm_config, validate=False) + if ( + afd_config.connector == AFD_ASYNC_CONNECTOR + and not afd_config.compute_gate_on_attention + ): + raise ValueError( + "DSV4 CAMAsyncAFDConnector requires " + "compute_gate_on_attention=true", + ) + if vllm_config.parallel_config.use_sequence_parallel_moe: + raise RuntimeError("AFD DSV4 does not support sequence-parallel MoE") + + _refresh_ascend_fused_moe() + nn.Module.__init__(self) + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.vllm_config = vllm_config + self.compilation_config = vllm_config.compilation_config + self.afd_config = afd_config + self.config = config + self.device = native.current_platform.device_type + self.hidden_size = config.hidden_size + self.vocab_size = config.vocab_size + self.hc_mult = config.hc_mult + self.norm_eps = config.rms_norm_eps + self.hc_eps = config.hc_eps + + self.is_v32 = hasattr(config, "index_topk") + if self.is_v32 and afd_config.role == _ATTENTION_ROLE: + 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 + + 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() + + self.start_layer, self.end_layer, self.layers = native.make_layers( + config.num_hidden_layers, + lambda prefix, **_: AFDDeepseekV4DecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + topk_indices_buffer=topk_indices_buffer, + ), + prefix=f"{prefix}.layers", + ) + + 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() + + hc_dim = self.hc_mult * config.hidden_size + self.hc_head_fn = nn.Parameter( + torch.empty(self.hc_mult, hc_dim, dtype=torch.float32) + ) + self.hc_head_base = nn.Parameter( + torch.empty(self.hc_mult, dtype=torch.float32) + ) + self.hc_head_scale = nn.Parameter(torch.empty(1, dtype=torch.float32)) + self._mtp_hidden_buffer = torch.empty( + vllm_config.scheduler_config.max_num_batched_tokens, + hc_dim, + dtype=vllm_config.model_config.dtype, + device=self.device, + ) + + def make_empty_intermediate_tensors( + batch_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> IntermediateTensors: + return IntermediateTensors( + { + "hidden_states": torch.zeros( + (batch_size, self.hc_mult, config.hidden_size), + dtype=dtype, + device=device, + ) + } + ) + + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors + self.aux_hidden_state_layers: tuple[int, ...] = () + self.num_redundant_experts = ( + vllm_config.parallel_config.eplb_config.num_redundant_experts + ) + + def compute_ffn_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + **kwargs: Any, + ) -> torch.Tensor: + return self.layers[layer_idx].compute_ffn_output(hidden_states, **kwargs) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + async_moe_metadata = get_async_moe_ubatch_metadata_from_forward_context() + if async_moe_metadata is None: + return super().forward( + input_ids, + positions, + intermediate_tensors, + inputs_embeds, + ) + if input_ids is None or inputs_embeds is not None: + raise NotImplementedError( + "DSV4 AFD async MoE ubatching requires token input_ids for " + "hash routing and does not support inputs_embeds" + ) + return _run_async_moe_ubatch_forward( + self, + input_ids, + positions, + intermediate_tensors, + async_moe_metadata, + inputs_embeds, + ) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return tuple( + layer.layer_idx + for layer in self.layers + if isinstance(layer, AFDDeepseekV4DecoderLayer) + and isinstance(layer.mlp, native.DeepseekV4MoE) + ) + + def get_experts_routing_spec( + self, + layer_idx: int, + ) -> AFDExpertRoutingSpec: + layer = self.layers[layer_idx] + if not isinstance(layer.mlp, native.DeepseekV4MoE): + raise RuntimeError("DSV4 layer does not own a native MoE") + gate = layer.mlp.gate + return AFDExpertRoutingSpec( + router_logits_width=int(layer.mlp.n_routed_experts), + router_logits_dtype=gate.out_dtype or gate.weight.dtype, + ) + + def compute_experts_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + router_logits: torch.Tensor, + ) -> torch.Tensor: + raise NotImplementedError( + "DSV4 AFD does not support compute_gate_on_attention yet" + ) + + +def _pad_stage_input_ids( + input_ids: torch.Tensor, + metadata: AsyncMoeUbatchMetadata, +) -> list[torch.Tensor]: + """Build the token IDs that correspond to each physical AFD stage.""" + + flat_input_ids = input_ids.reshape(-1) + actual_parent_tokens = max(int(stage.token_slice.stop) for stage in metadata.stages) + if int(flat_input_ids.numel()) < actual_parent_tokens: + raise ValueError( + "DSV4 async MoE input_ids do not cover the staged tokens: " + f"input_ids={int(flat_input_ids.numel())}, " + f"staged_tokens={actual_parent_tokens}", + ) + tp_group = get_tp_group() + tp_rank = int(tp_group.rank_in_group) + tp_size = int(tp_group.world_size) + stage_input_ids: list[torch.Tensor] = [] + for stage in metadata.stages: + ids = flat_input_ids[stage.token_slice] + physical_tokens = int(stage.input_tokens) + if int(ids.numel()) < physical_tokens: + ids = torch.nn.functional.pad( + ids, + (0, physical_tokens - int(ids.numel())), + value=-1, + ) + if metadata.use_sequence_parallel: + if physical_tokens % tp_size != 0: + raise ValueError( + "DSV4 async MoE stage is not TP divisible: " + f"tokens={physical_tokens}, tp_size={tp_size}", + ) + local_tokens = physical_tokens // tp_size + local_start = tp_rank * local_tokens + ids = ids[local_start : local_start + local_tokens] + stage_input_ids.append(ids) + return stage_input_ids + + +def _run_async_moe_ubatch_forward( + model: AFDDeepseekV4Model, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + metadata: AsyncMoeUbatchMetadata, + inputs_embeds: torch.Tensor | None, +) -> torch.Tensor | IntermediateTensors: + """Run DSV4's two AFD-owned CAM stages in one model invocation.""" + + if len(metadata.stages) != 2: + raise ValueError( + "DSV4 async MoE currently requires exactly two AFD stages, got " + f"{len(metadata.stages)}", + ) + pp_group = get_pp_group() + if pp_group.is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = model.embed_input_ids(input_ids) + hidden_states = hidden_states.unsqueeze(1).repeat(1, model.hc_mult, 1) + else: + if intermediate_tensors is None: + raise ValueError("DSV4 pipeline stage requires intermediate tensors") + hidden_states = intermediate_tensors["hidden_states"] + + parent_context = get_forward_context() + if bool(parent_context.flash_comm_v1_enabled) != metadata.use_sequence_parallel: + raise RuntimeError( + "DSV4 async MoE stage layout does not match FlashComm1: " + f"layout_sequence_parallel={metadata.use_sequence_parallel}, " + f"flash_comm_v1_enabled={bool(parent_context.flash_comm_v1_enabled)}", + ) + afd_metadata = get_afd_metadata_from_forward_context(parent_context) + if afd_metadata is None: + raise RuntimeError("DSV4 async MoE requires AFD forward metadata") + + stage_inputs = build_async_moe_stage_inputs( + hidden_states, + None, + positions, + None, + metadata, + ) + stage_hidden_states = stage_inputs.hidden_states + stage_positions = stage_inputs.positions + stage_input_ids = _pad_stage_input_ids(input_ids, metadata) + stage_dispatch_layouts: list[CAMDispatchLayout | None] = [None, None] + stage_dispatch_refs: list[torch.Tensor | None] = [None, None] + stage_pending_dispatches: list[Any | None] = [None, None] + stage_ffn_state: list[ + tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None + ] = [None, None] + + def stage_context(stage_idx: int): + stage = metadata.stages[stage_idx] + context = copy(parent_context) + context.attn_metadata = metadata.attn_metadata[stage_idx] + context.additional_kwargs = dict(parent_context.additional_kwargs or {}) + context.ubatch_idx = stage_idx + context.num_ubatches = len(metadata.stages) + context.dbo_enabled = False + context.input_ids = stage_input_ids[stage_idx] + if metadata.use_sequence_parallel: + context.num_tokens = stage.actual_tokens + context.pad_size = int(stage.input_tokens) - stage.actual_tokens + else: + context.num_tokens = int(stage.input_tokens) + context.pad_size = 0 + return context + + def compute_stage_attention( + layer: AFDDeepseekV4DecoderLayer, + stage_idx: int, + ) -> None: + if not isinstance(layer.mlp, AFDDeepseekV4AttentionGateRemoteMoE): + raise RuntimeError( + "DSV4 async MoE requires Attention-side expert routing", + ) + with override_forward_context(stage_context(stage_idx)): + stage_hidden = stage_hidden_states[stage_idx] + attn_residual = stage_hidden.clone() + stage_hidden, attn_post, attn_comb = layer.hc_pre( + stage_hidden, + layer.hc_attn_fn, + layer.hc_attn_scale, + layer.hc_attn_base, + ) + stage_hidden = layer.input_layernorm(stage_hidden) + stage_hidden = layer.self_attn( + positions=stage_positions[stage_idx], + hidden_states=stage_hidden, + llama_4_scaling=None, + ) + stage_hidden = layer.hc_post( + stage_hidden, + attn_residual, + attn_post, + attn_comb, + ) + ffn_residual = stage_hidden.clone() + stage_hidden, ffn_post, ffn_comb = layer.hc_pre( + stage_hidden, + layer.hc_ffn_fn, + layer.hc_ffn_scale, + layer.hc_ffn_base, + ) + stage_hidden = layer.post_attention_layernorm(stage_hidden) + from afd_plugin.model_executor.models.npu import ( + deepseek_v4_attention_gate, + ) + + topk_weights, topk_ids = ( + deepseek_v4_attention_gate.compute_attention_gate_topk( + layer.mlp, + stage_hidden, + ) + ) + dispatch = prepare_cam_dispatch_payload( + stage_hidden, + topk_weights, + topk_ids, + None, + use_sequence_parallel=metadata.use_sequence_parallel, + ) + stage_pending_dispatches[stage_idx] = dispatch + stage_ffn_state[stage_idx] = (ffn_residual, ffn_post, ffn_comb) + + def send_stage_attention( + layer: AFDDeepseekV4DecoderLayer, + stage_idx: int, + ) -> None: + dispatch = stage_pending_dispatches[stage_idx] + if dispatch is None: + raise RuntimeError( + f"DSV4 async MoE stage {stage_idx} has no computed Attention", + ) + transfer_metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=layer.layer_idx, + stage_idx=stage_idx, + seq_len=int(dispatch.hidden_states.shape[0]), + ) + afd_metadata.connector.send_attn_output( + dispatch.hidden_states, + AFDTransferContext(metadata=transfer_metadata), + topk_weights=dispatch.topk_weights, + topk_ids=dispatch.topk_ids, + ) + stage_dispatch_layouts[stage_idx] = dispatch.layout + stage_dispatch_refs[stage_idx] = dispatch.hidden_states + stage_pending_dispatches[stage_idx] = None + + def receive_and_complete( + layer: AFDDeepseekV4DecoderLayer, + stage_idx: int, + ) -> None: + layout = stage_dispatch_layouts[stage_idx] + dispatch_ref = stage_dispatch_refs[stage_idx] + ffn_state = stage_ffn_state[stage_idx] + if layout is None or dispatch_ref is None or ffn_state is None: + raise RuntimeError( + f"DSV4 async MoE stage {stage_idx} has no pending FFN work", + ) + local_output = afd_metadata.connector.recv_ffn_output( + ref_tensor=dispatch_ref, + ubatch_idx=stage_idx, + ) + ffn_output = restore_cam_dispatch_output(local_output, layout) + ffn_residual, ffn_post, ffn_comb = ffn_state + with override_forward_context(stage_context(stage_idx)): + stage_hidden_states[stage_idx] = layer.hc_post( + ffn_output, + ffn_residual, + ffn_post, + ffn_comb, + ) + stage_dispatch_layouts[stage_idx] = None + stage_dispatch_refs[stage_idx] = None + stage_ffn_state[stage_idx] = None + + layers = list(islice(model.layers, model.start_layer, model.end_layer)) + if not layers: + restored_hidden_states = restore_async_moe_stage_outputs( + stage_hidden_states, + metadata, + ) + else: + _run_two_stage_async_moe_schedule( + layers, + compute_stage_attention, + send_stage_attention, + receive_and_complete, + ) + restored_hidden_states = restore_async_moe_stage_outputs( + stage_hidden_states, + metadata, + ) + + if parent_context.flash_comm_v1_enabled: + hidden_flat = native.tensor_model_parallel_all_gather( + restored_hidden_states.flatten(1), + dim=0, + ) + if parent_context.pad_size > 0: + hidden_flat = hidden_flat[: -parent_context.pad_size] + else: + hidden_flat = restored_hidden_states.flatten(1) + model._mtp_hidden_buffer[: hidden_flat.shape[0]].copy_(hidden_flat) + + if not pp_group.is_last_rank: + return IntermediateTensors({"hidden_states": restored_hidden_states}) + output = model.hc_head( + restored_hidden_states, + model.hc_head_fn, + model.hc_head_scale, + model.hc_head_base, + ) + return model.norm(output) + + +def _run_two_stage_async_moe_schedule( + layers: list[AFDDeepseekV4DecoderLayer], + compute_stage_attention: Any, + send_stage_attention: Any, + receive_and_complete: Any, +) -> None: + """Pipeline two AFD-owned stages through Attention and remote FFN. + + Once a send succeeds, any later exception is intentionally fatal: the + EngineCore propagates model-forward failures and terminates the worker. + Blindly draining CAM here would be unsafe because the remote receive point + is unknown after an asynchronous operator failure. + """ + + if not layers: + return + compute_stage_attention(layers[0], 0) + send_stage_attention(layers[0], 0) + for layer_idx in range(len(layers) - 1): + current_layer = layers[layer_idx] + next_layer = layers[layer_idx + 1] + compute_stage_attention(current_layer, 1) + receive_and_complete(current_layer, 0) + send_stage_attention(current_layer, 1) + compute_stage_attention(next_layer, 0) + receive_and_complete(current_layer, 1) + send_stage_attention(next_layer, 0) + last_layer = layers[-1] + compute_stage_attention(last_layer, 1) + receive_and_complete(last_layer, 0) + send_stage_attention(last_layer, 1) + receive_and_complete(last_layer, 1) + + +class AFDDeepseekV4ForCausalLM(native.AscendDeepseekV4ForCausalLM): + """Ascend DSV4 causal LM wrapper for AFD.""" + + model_cls = AFDDeepseekV4Model + + 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, + **kwargs: Any, + ) -> torch.Tensor: + return self.model.compute_ffn_output(hidden_states, layer_idx, **kwargs) + + def get_experts_layer_indices(self) -> tuple[int, ...]: + return self.model.get_experts_layer_indices() + + def get_experts_routing_spec( + self, + layer_idx: int, + ) -> AFDExpertRoutingSpec: + return self.model.get_experts_routing_spec(layer_idx) + + def compute_experts_output( + self, + hidden_states: torch.Tensor, + layer_idx: int, + router_logits: torch.Tensor, + ) -> torch.Tensor: + return self.model.compute_experts_output( + hidden_states, + layer_idx, + router_logits, + ) + + 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__ = [ + "AFDDeepseekV4DecoderLayer", + "AFDDeepseekV4ForCausalLM", + "AFDDeepseekV4Model", +] diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py b/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py new file mode 100644 index 00000000..cd0d5bb9 --- /dev/null +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""DeepSeek V4 Attention-side routing helpers for Async CAM.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from afd_plugin.model_executor.models.npu.deepseek_v4 import ( + AFDDeepseekV4AttentionGateRemoteMoE, + ) + + +def compute_attention_gate_topk( + moe: AFDDeepseekV4AttentionGateRemoteMoE, + hidden_states: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DSV4 routing without entering vLLM's native MoE communicator. + + AFD Async CAM owns the cross-role dispatch. The vLLM-Ascend fused + selector's hash path instead assumes native EP/SP communication and calls + ``forward_context.moe_comm_method.pad_and_split_input_ids``. That object + is intentionally absent on Attention ranks, including the KV-cache profile + forward. Use the same CANN routing operators directly on local Attention + tokens, then hand their IDs and weights to CAM dispatch. + """ + + router_logits, _ = moe.gate(hidden_states) + if moe.scoring_func == "sqrtsoftplus": + topk_weights, topk_ids = _compute_sqrtsoftplus_topk(moe, router_logits) + else: + topk_weights, topk_ids = _compute_standard_topk(moe, router_logits) + return topk_weights.to(torch.float32), topk_ids + + +def _compute_sqrtsoftplus_topk( + moe: AFDDeepseekV4AttentionGateRemoteMoE, + router_logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run DSV4's sqrtsoftplus CANN router without native MoE communication.""" + + if moe.scoring_func != "sqrtsoftplus": + raise RuntimeError( + "DSV4 Hash routing requires scoring_func='sqrtsoftplus', got " + f"{moe.scoring_func!r}", + ) + + tid2eid = moe.gate.tid2eid + input_ids = None + if tid2eid is not None: + from vllm.forward_context import get_forward_context + + forward_context = get_forward_context() + input_ids = getattr(forward_context, "input_ids", None) + if input_ids is None: + raise RuntimeError( + "DSV4 Hash routing requires local input_ids in the forward " + "context", + ) + input_ids = input_ids.reshape(-1).to(torch.int64) + if input_ids.numel() != router_logits.shape[0]: + raise RuntimeError( + "DSV4 Hash routing input_ids/token count mismatch on Attention: " + f"input_ids={input_ids.numel()} router_tokens={router_logits.shape[0]}", + ) + input_ids = torch.where(input_ids == -1, 0, input_ids) + tid2eid = tid2eid.to(torch.int32) + correction_bias = moe.gate.e_score_correction_bias + if correction_bias is not None and correction_bias.dtype != router_logits.dtype: + correction_bias = correction_bias.to(router_logits.dtype) + topk_weights, topk_ids, _ = torch.ops._C_ascend.moe_gating_top_k_hash( + x=router_logits, + k=moe.top_k, + bias=correction_bias, + input_ids=input_ids, + tid2eid=tid2eid, + k_group=moe.topk_group, + group_count=moe.num_expert_group, + routed_scaling_factor=moe.routed_scaling_factor, + eps=1e-20, + group_select_mode=1, + renorm=0, + norm_type=2, + out_flag=False, + ) + return topk_weights, topk_ids + + +def _compute_standard_topk( + moe: AFDDeepseekV4AttentionGateRemoteMoE, + router_logits: torch.Tensor, +) -> tuple[torch.Tensor, torch.Tensor]: + """Run the non-Hash CANN selector without native MoE communication.""" + + from vllm_ascend.device.device_op import DeviceOperator + + norm_type_by_scoring_func = {"softmax": 0, "sigmoid": 1} + try: + norm_type = norm_type_by_scoring_func[moe.scoring_func] + except KeyError as exc: + raise RuntimeError( + "Unsupported non-Hash DSV4 routing scoring function: " + f"{moe.scoring_func!r}", + ) from exc + correction_bias = moe.gate.e_score_correction_bias + if correction_bias is not None and correction_bias.dtype != router_logits.dtype: + correction_bias = correction_bias.to(router_logits.dtype) + topk_weights, topk_ids, _ = DeviceOperator.moe_gating_top_k( + router_logits, + k=moe.top_k, + k_group=moe.topk_group, + group_count=moe.num_expert_group, + group_select_mode=1, + renorm=int(moe.renormalize), + norm_type=norm_type, + out_flag=False, + routed_scaling_factor=moe.routed_scaling_factor, + eps=1e-20, + bias_opt=correction_bias, + ) + return topk_weights, topk_ids diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index e5fd8789..c2d10f39 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -49,7 +49,6 @@ from vllm_ascend.ops.rotary_embedding import update_cos_sin from vllm_ascend.spec_decode.dflash_proposer import AscendDflashProposer from vllm_ascend.spec_decode.draft_proposer import AscendDraftModelProposer -from vllm_ascend.spec_decode.dspark_proposer import AscendDSparkProposer from vllm_ascend.spec_decode.eagle_proposer import AscendEagleProposer from vllm_ascend.spec_decode.step3p5 import AscendStep3p5MTPProposer from vllm_ascend.utils import ( @@ -205,6 +204,7 @@ def _model_forward( ): forward_context = get_forward_context() # ### PATCH START: AFD forward-context metadata + forward_context.input_ids = input_ids if self.ubatch_slices is not None: forward_context.ubatch_slices = self.ubatch_slices forward_context.dbo_enabled = False @@ -237,14 +237,14 @@ def _model_forward( num_tokens_padded, ) - # ### PATCH START: AFD defers FlashComm gather to the ubatch wrapper + # ### PATCH START: AFD defers FlashComm gather to model execution if ( forward_context.flash_comm_v1_enabled and not forward_context.dbo_enabled and not isinstance(hidden_states, IntermediateTensors) ): hidden_states = self._all_gather_hidden_states_and_aux(hidden_states) - # ### PATCH END: AFD defers FlashComm gather to the ubatch wrapper + # ### PATCH END: AFD defers FlashComm gather to model execution return hidden_states # Upstream source: vllm-ascend commit 80d8c194f, @@ -276,8 +276,7 @@ def _build_attention_metadata( num_reqs_padded, ) if self.afd_async_extra_info.async_moe_ubatching: - self.ubatch_slices = None - return self._build_attention_metadata_with_async_moe_ubatches( + result = self._build_attention_metadata_with_async_moe_ubatches( num_tokens=num_tokens, num_reqs=num_reqs, max_query_len=max_query_len, @@ -291,6 +290,8 @@ def _build_attention_metadata( num_scheduled_tokens_np=num_scheduled_tokens_np, cascade_attn_prefix_lens=cascade_attn_prefix_lens, ) + self.ubatch_slices = None + return result self._afd_pending_metadata = self._build_afd_metadata( ubatch_slices, num_tokens, @@ -391,22 +392,6 @@ def _build_attention_metadata_with_async_moe_ubatches( UBatchSlice(stage.request_slice, stage.token_slice) for stage in stages ] - logger.debug( - "AFD NPU async MoE ubatch split; num_reqs=%s num_tokens=%s " - "num_scheduled_tokens=%s split=%s sequence_parallel=%s " - "request_slices=%s token_slices=%s stage_input_tokens=%s " - "stage_actual_tokens=%s", - len(num_scheduled_tokens_np), - num_tokens, - num_scheduled_tokens_np.tolist(), - self.afd_async_extra_info.async_moe_split, - use_sequence_parallel, - [(stage.request_slice.start, stage.request_slice.stop) for stage in stages], - [(stage.token_slice.start, stage.token_slice.stop) for stage in stages], - [int(stage.input_tokens) for stage in stages], - [stage.actual_tokens for stage in stages], - ) - stage_attn_metadata, _ = self._build_attention_metadata_with_ubatches( num_tokens=num_tokens, num_reqs=num_reqs, @@ -428,6 +413,10 @@ def _build_attention_metadata_with_async_moe_ubatches( use_sequence_parallel=use_sequence_parallel, parent_input_tokens=num_tokens_padded, ) + self._afd_pending_metadata = self._build_afd_metadata( + stage_slices, + num_tokens, + ) return full_metadata # Upstream source: vllm-ascend commit 80d8c194f, @@ -754,7 +743,7 @@ def _build_attn_group_metadata( ) if self.speculative_config and isinstance( self.drafter, - AscendStep3p5MTPProposer | AscendDSparkProposer, + AscendStep3p5MTPProposer, ): self.drafter.set_per_group_attn_metadata( kv_cache_gid, @@ -766,8 +755,7 @@ def _build_attn_group_metadata( self.drafter, AscendEagleProposer | AscendDraftModelProposer - | AscendDflashProposer - | AscendDSparkProposer, + | AscendDflashProposer, ): if self.drafter.attn_layer_names[0] in kv_cache_group.layer_names: spec_decode_common_attn_metadata = cm @@ -1438,10 +1426,19 @@ def _install_afd_metadata_on_forward_context( return dp_metadata = forward_context.dp_metadata ubatch_slices = forward_context.ubatch_slices + if self._afd_async_moe_ubatch_metadata is not None: + ubatch_slices = [ + UBatchSlice(stage.request_slice, stage.token_slice) + for stage in self._afd_async_moe_ubatch_metadata.stages + ] padded_graph_tokens = _full_cudagraph_padded_tokens(forward_context) if padded_graph_tokens is not None and not ubatch_slices: dp_metadata = self._build_capture_dp_metadata(padded_graph_tokens) - self._send_dp_metadata(dp_metadata, ubatch_slices) + self._send_dp_metadata( + dp_metadata, + ubatch_slices, + input_ids=getattr(forward_context, "input_ids", None), + ) def _install_async_moe_ubatch_metadata_on_forward_context( self, @@ -1482,6 +1479,8 @@ def _send_dp_metadata( self, dp_metadata: DPMetadata | AFDDPMetadata | None, ubatch_slices: UBatchSlices | None, + *, + input_ids: torch.Tensor | None = None, ) -> None: assert self.connector.control_plane is not None, ( "_send_dp_metadata needs control plane driven connectors" @@ -1497,22 +1496,33 @@ def _send_dp_metadata( else: dp_metadata = self._ensure_dp_metadata(dp_metadata) dp_metadata_list = {0: dp_metadata} + input_ids_by_stage: dict[int, list[int]] = {} + if input_ids is not None: + if ubatch_slices and len(ubatch_slices) > 1: + for idx, ubatch in enumerate(ubatch_slices): + input_ids_by_stage[idx] = ( + input_ids[ubatch.token_slice] + .detach() + .to(device="cpu", dtype=torch.int64) + .flatten() + .tolist() + ) + else: + input_ids_by_stage[0] = ( + input_ids.detach() + .to(device="cpu", dtype=torch.int64) + .flatten() + .tolist() + ) is_warmup = bool(self._is_warmup) is_graph_capturing = bool(self._afd_is_graph_capturing) payload = AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + input_ids_by_stage=input_ids_by_stage, ) self.connector.control_plane.update_state_from_dp_metadata(payload) - logger.warning( - "AFD NPU Attention send_dp_metadata decision; world_rank=%d " - "key=%s is_graph_capturing=%s is_warmup=%s", - self.connector.world_rank, - _dp_metadata_debug_key(dp_metadata_list), - is_graph_capturing, - is_warmup, - ) self.connector.control_plane.send_dp_metadata_list(payload) def _ensure_dp_metadata( @@ -1920,18 +1930,6 @@ def _make_uniform_dp_metadata(dp_size: int, num_tokens: int) -> AFDDPMetadata: return AFDDPMetadata(num_tokens_across_dp_cpu=num_tokens_across_dp_cpu) -def _dp_metadata_debug_key( - dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], -) -> tuple[tuple[int, tuple]]: - key_parts: list[tuple[int, tuple]] = [] - for stage_idx, metadata in sorted(dp_metadata_list.items()): - values_tuple = tuple( - int(value) for value in metadata.num_tokens_across_dp_cpu.tolist() - ) - key_parts.append((int(stage_idx), values_tuple)) - return tuple(key_parts) - - def _normalize_metadata_ubatch_slices( ubatch_slices: UBatchSlices | None, num_tokens_padded: int | None, diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 83e81284..d36b3481 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -36,6 +36,7 @@ AFDAsyncTransferState, CAMAsyncAFDConnector, ) +from afd_plugin.connectors.npu.camp2p import CAMP2pAFDConnector from afd_plugin.v1.worker.attention_model_runner import ( _resolve_world_ranks, ) @@ -56,7 +57,6 @@ logger = init_logger(__name__) - class AFDNPUFFNModelRunner(NPUModelRunner): """Connector-driven NPU FFN runner for AFD execution.""" @@ -208,6 +208,7 @@ def _ffn_forward( _make_dp_metadata_payload( dp_metadata_list, is_graph_capturing=is_graph_capturing, + input_ids_by_stage=_camp2p_input_ids_by_stage(self.connector), ), ) num_stages = max(len(dp_metadata_list), 1) @@ -263,6 +264,12 @@ def _ffn_forward( forward_context.dp_metadata = dp_metadata_list.get(stage_idx) forward_context.additional_kwargs["afd_metadata"] = metadata assert states, "Context.states must not be None" + _install_ffn_input_ids( + forward_context, + _camp2p_input_ids_by_stage(self.connector).get(stage_idx), + num_tokens=int(hidden_states.shape[0]), + device=hidden_states.device, + ) _set_moe_layer_index(forward_context, layer_idx) rank_ffn_output = self.model.compute_ffn_output( @@ -330,6 +337,13 @@ def _ffn_forward_connector_driven( work_item, rank_ffn_output, ) + # ``async_combine_send`` returns after enqueueing its NPU work. + # Without an acknowledgement/credit path, immediately calling + # dispatch-recv for all following MoE layers can reuse CAM's + # bounded buffers while the prior combine is still in flight. + # The target baseline is deliberately serialized per layer; + # async MoE ubatching will introduce explicit stage credits. + torch.npu.synchronize() return rank_ffn_output def capture_model( @@ -478,9 +492,12 @@ def _ffn_layer_indices(runner: AFDNPUFFNModelRunner) -> range | list[int]: def _is_moe_layer(hf_config: object, layer_idx: int) -> bool: moe_layer_freq = getattr(hf_config, "moe_layer_freq", 1) + # DeepSeek V4 has routed experts in every decoder layer and does not expose + # the V2/V3 ``first_k_dense_replace`` compatibility field. + first_moe_layer = getattr(hf_config, "first_k_dense_replace", 0) return ( - hf_config.n_routed_experts is not None - and layer_idx >= hf_config.first_k_dense_replace + getattr(hf_config, "n_routed_experts", None) is not None + and layer_idx >= first_moe_layer and layer_idx % moe_layer_freq == 0 ) @@ -490,14 +507,44 @@ def _make_dp_metadata_payload( *, is_graph_capturing: bool = False, is_warmup: bool = False, + input_ids_by_stage: dict[int, list[int]] | None = None, ) -> AFDControlPayload: return AFDControlPayload( dp_metadata_list=dp_metadata_list, is_graph_capturing=is_graph_capturing, is_warmup=is_warmup, + input_ids_by_stage={} if input_ids_by_stage is None else input_ids_by_stage, ) +def _camp2p_input_ids_by_stage( + connector: AFDConnectorBase, +) -> dict[int, list[int]]: + if isinstance(connector, CAMP2pAFDConnector): + return connector.input_ids_by_stage + return {} + + +def _install_ffn_input_ids( + forward_context: Any, + input_ids: list[int] | None, + *, + num_tokens: int, + device: torch.device, +) -> None: + """Install token ids required by native DSV4 hash routing on FFN.""" + if input_ids is None: + forward_context.input_ids = None + return + ids = torch.as_tensor(input_ids, dtype=torch.int64, device=device).flatten() + if ids.numel() != num_tokens: + raise RuntimeError( + "CAMP2P input_ids must align with received hidden states: " + f"got {ids.numel()} token IDs for {num_tokens} hidden-state rows" + ) + forward_context.input_ids = ids + + def _ffn_token_counts_across_ranks( connector: AFDConnectorBase, dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index 54b8120a..665c308f 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -141,6 +141,9 @@ def _run_ffn_server_loop(self) -> None: continue payload = self.model_runner.connector.control_plane.recv_dp_metadata_list() + self.model_runner.connector.control_plane.update_state_from_dp_metadata( + payload + ) dp_metadata_list = payload.dp_metadata_list is_attn_graph_capturing = payload.is_graph_capturing is_warmup = payload.is_warmup diff --git a/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index e7ca5d2e..2b0e2984 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/afd_plugin/v1/worker/npu/forward_context.py @@ -109,7 +109,12 @@ def create_ascend_forward_context( cur_forward_context.max_tokens_across_pcp ) new_forward_context.sinks = cur_forward_context.sinks - new_forward_context.input_ids = cur_forward_context.input_ids + if cur_forward_context.input_ids is not None: + new_forward_context.input_ids = cur_forward_context.input_ids[ + ubatch_slice.token_slice + ] + else: + new_forward_context.input_ids = None new_forward_context.eplb_heat_collection_status = ( cur_forward_context.eplb_heat_collection_status ) diff --git a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py index 93f30653..c2c9d845 100644 --- a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py +++ b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py @@ -155,9 +155,9 @@ def __init__( assert not enable_enpu, "AscendUBatchWrapper does not support ENPU" self.runnable = runnable self.vllm_config = vllm_config + assert self.vllm_config.parallel_config.num_ubatches == AFD_NPU_NUM_UBATCHES self.compilation_config = vllm_config.compilation_config self.comm_stream = torch.npu.Stream(device=device) - assert self.vllm_config.parallel_config.num_ubatches == AFD_NPU_NUM_UBATCHES self.ready_barrier = threading.Barrier(_READY_BARRIER_PARTIES) self.cudagraphs: dict[AscendNPUGraphKey, AscendNPUGraphMetaData] = {} self.cudagraph_wrapper = None diff --git a/tests/unit/compat/npu/test_dsv4_async_validation.py b/tests/unit/compat/npu/test_dsv4_async_validation.py new file mode 100644 index 00000000..45b7a5fc --- /dev/null +++ b/tests/unit/compat/npu/test_dsv4_async_validation.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Unit coverage for the narrow DSV4 Async CAM feature contract.""" + +from types import SimpleNamespace + +import pytest + +from afd_plugin.compat.npu.feature_validation import ( + _fail_if_unsupported_dsv4_async_features, +) +from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo + + +def _afd_config(*, compute_gate_on_attention: bool) -> SimpleNamespace: + return SimpleNamespace(compute_gate_on_attention=compute_gate_on_attention) + + +def test_dsv4_async_requires_attention_side_gate() -> None: + with pytest.raises(RuntimeError, match="compute_gate_on_attention"): + _fail_if_unsupported_dsv4_async_features( + _afd_config(compute_gate_on_attention=False), + AFDAsyncExtraInfo(dynamic_quant=1), + ) + + +@pytest.mark.parametrize( + ("extra_info", "message"), + [ + ( + AFDAsyncExtraInfo(dynamic_quant=0), + "dynamicQuant=1", + ), + ], +) +def test_dsv4_async_rejects_deferred_features( + extra_info: AFDAsyncExtraInfo, + message: str, +) -> None: + with pytest.raises(RuntimeError, match=message): + _fail_if_unsupported_dsv4_async_features( + _afd_config(compute_gate_on_attention=True), + extra_info, + ) + + +def test_dsv4_async_accepts_target_baseline_and_async_moe_ubatching() -> None: + _fail_if_unsupported_dsv4_async_features( + _afd_config(compute_gate_on_attention=True), + AFDAsyncExtraInfo( + dynamic_quant=1, + attn_ranks_per_dp=2, + async_moe_ubatching=True, + ), + ) diff --git a/tests/unit/connectors/test_async_cam_connector.py b/tests/unit/connectors/test_async_cam_connector.py index 0e64b530..eb11831e 100644 --- a/tests/unit/connectors/test_async_cam_connector.py +++ b/tests/unit/connectors/test_async_cam_connector.py @@ -138,7 +138,13 @@ def zeros(self, shape, *, dtype, device): def _vllm_config(*, tp_size: int = 1, pcp_size: int = 1, extra_config=None): return SimpleNamespace( - additional_config={"afd": {"connector_extra_config": extra_config or {}}}, + additional_config={ + "afd": { + "connector_extra_config": extra_config + if extra_config is not None + else {"attn_ranks_per_dp": 4}, + }, + }, parallel_config=SimpleNamespace( data_parallel_size=1, data_parallel_rank=0, @@ -165,6 +171,15 @@ def _afd_config(*, role: str): ) +def _dp2_afd_config(*, role: str): + return AFDConfig( + connector="CAMAsyncAFDConnector", + role=role, + num_attention_ranks=8, + num_ffn_ranks=8, + ) + + def _topk_payload(batch_size: int, topk: int = 2): return { "topk_ids": _FakeTensor((batch_size, topk), dtype="int32"), @@ -227,13 +242,13 @@ def test_async_connector_uses_attn_ranks_per_dp_for_cam_tp_size(): _vllm_config( tp_size=4, pcp_size=2, - extra_config={"attn_ranks_per_dp": "3"}, + extra_config={"attn_ranks_per_dp": "4"}, ), _afd_config(role="attention"), 0, ) - assert connector.tp_size == 3 + assert connector.tp_size == 4 @pytest.mark.parametrize("value", [True, "bad"]) @@ -273,6 +288,38 @@ def test_async_topology_uses_cam_attention_first_rank_layout(): assert ffn.expert_per_rank == 4 +@pytest.mark.parametrize( + ("role", "role_rank", "expected_world_rank"), + [ + ("attention", 3, 3), + ("attention", 4, 4), + ("ffn", 3, 11), + ("ffn", 4, 12), + ], +) +def test_async_topology_uses_one_world_for_all_dp_replicas( + role, + role_rank, + expected_world_rank, +): + topology = build_async_topology( + _dp2_afd_config(role=role), + role_rank, + num_routed_experts=32, + ) + + assert topology.world_rank == expected_world_rank + assert topology.attn_size == 8 + assert topology.ffn_size == 8 + assert topology.world_size == 16 + assert topology.expert_per_rank == 4 + + +def test_async_extra_info_rejects_removed_shared_ffn_pool_option(): + with pytest.raises(ValueError, match="unknown AFD async connector_extra_config"): + AFDAsyncExtraInfo.from_mapping({"shared_ffn_pool": "true"}) + + def test_async_connector_init_creates_attention_first_hccl_group(monkeypatch): calls = [] fake_torch = _FakeTorch() @@ -311,11 +358,11 @@ def fake_init_afd_process_group(**kwargs): assert calls == [ { "backend": "hccl", - "init_method": "tcp://127.0.0.1:1239", "world_size": 6, "rank": 5, "group_name": AFD_ASYNC_CAM_GROUP_NAME, "timeout": calls[0]["timeout"], + "init_method": "tcp://127.0.0.1:1239", }, ] assert connector.cam_pg is not None @@ -325,6 +372,45 @@ def fake_init_afd_process_group(**kwargs): assert connector._placeholder.shape == (1,) +def test_async_connector_init_uses_one_hccl_group_for_all_dp(monkeypatch): + calls = [] + fake_torch = _FakeTorch() + monkeypatch.setattr(async_cam_module, "torch", fake_torch) + monkeypatch.setattr( + async_cam_module, + "ensure_cam_async_ops_available", + lambda: None, + ) + + def fake_init_afd_process_group(**kwargs): + calls.append(kwargs) + backend = SimpleNamespace( + get_hccl_comm_name=lambda rank: f"hccl:{kwargs['group_name']}:{rank}", + ) + return SimpleNamespace(_get_backend=lambda device: backend) + + monkeypatch.setattr( + async_cam_module, + "init_afd_process_group", + fake_init_afd_process_group, + ) + connector = CAMAsyncAFDConnector( + 0, + 0, + _vllm_config(extra_config={"attn_ranks_per_dp": 4}), + _dp2_afd_config(role="ffn"), + 4, + ) + + connector.init_afd_connector() + + assert calls[0]["world_size"] == 16 + assert calls[0]["rank"] == 12 + assert calls[0]["group_name"] == AFD_ASYNC_CAM_GROUP_NAME + assert calls[0]["init_method"] == "tcp://127.0.0.1:1239" + assert connector.group_name == "hccl:afd_async_cam:12" + + def test_async_connector_disables_dp_metadata_control_plane(): connector = CAMAsyncAFDConnector( 0, @@ -345,7 +431,7 @@ def test_async_connector_calls_cam_shaped_ops(monkeypatch): 0, _vllm_config( pcp_size=3, - extra_config={"attn_ranks_per_dp": 3}, + extra_config={"attn_ranks_per_dp": 4}, ), _afd_config(role="attention"), 0, @@ -379,7 +465,7 @@ def test_async_connector_calls_cam_shaped_ops(monkeypatch): assert fake_torch.ops.umdk_cam_op_lib.calls[1][1][4] == CAM_COMM_ID assert fake_torch.ops.umdk_cam_op_lib.calls[0][1][5:11] == (3, 16, 2, 2, 4, 4) assert fake_torch.ops.umdk_cam_op_lib.calls[1][1][5:11] == (3, 16, 2, 2, 4, 4) - assert fake_torch.ops.umdk_cam_op_lib.calls[0][1][14] == 3 + assert fake_torch.ops.umdk_cam_op_lib.calls[0][1][14] == 4 assert isinstance(context.states, AFDAsyncTransferState) assert isinstance(context.states, AFDTransferState) @@ -409,7 +495,7 @@ def test_async_ffn_side_dispatch_recv_and_combine_send(monkeypatch): assert states.dynamic_scales.shape == (4,) assert states.expand_x_shared.shape == (2, 16) assert states.dynamic_scales_shared.shape == (2,) - assert states.group_list.shape == (4,) + assert states.group_list.shape == (connector.topology.expert_per_rank,) assert states.expert_token_nums_shared.shape == (1,) assert fake_torch.ops.umdk_cam_op_lib.calls[0][0] == "dispatch_recv" assert fake_torch.ops.umdk_cam_op_lib.calls[1][0] == "combine_send" @@ -473,11 +559,7 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): hidden_size=connector.hidden_size, topk=connector.topk, layer_idx=layer_idx, - token_nums_rankid_layeridx=[ - _FakeScalar(7), - _FakeScalar(0), - _FakeScalar(11), - ], + token_nums_rankid_layeridx=torch.tensor([7, 0, 11], dtype=torch.int64), expert_token_nums_shared=[_FakeScalar(2)], group_list=torch.tensor([2, 3], dtype=torch.int64), dynamic_scales=_FakeTensorLike("scales"), @@ -491,7 +573,10 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): monkeypatch.setattr(connector, "recv_attn_output", fake_recv_attn_output) - work_item = connector.recv_ffn_work_item(stage_idx=0, max_num_tokens=16) + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=16, + ) states = work_item.context.states assert work_item.layer_idx == 11 @@ -537,11 +622,7 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): hidden_size=connector.hidden_size, topk=connector.topk, layer_idx=layer_idx, - token_nums_rankid_layeridx=[ - _FakeScalar(6), - _FakeScalar(0), - _FakeScalar(23), - ], + token_nums_rankid_layeridx=torch.tensor([6, 0, 23], dtype=torch.int64), expert_token_nums_shared=[_FakeScalar(0)], group_list=group_list, expand_x_shared=_FakeTensorLike("shared-hidden"), @@ -554,7 +635,10 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): monkeypatch.setattr(connector, "recv_attn_output", fake_recv_attn_output) - work_item = connector.recv_ffn_work_item(stage_idx=0, max_num_tokens=16) + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=16, + ) states = work_item.context.states assert work_item.layer_idx == 23 @@ -597,11 +681,7 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): hidden_size=connector.hidden_size, topk=connector.topk, layer_idx=layer_idx, - token_nums_rankid_layeridx=[ - _FakeScalar(5), - _FakeScalar(0), - _FakeScalar(7), - ], + token_nums_rankid_layeridx=torch.tensor([5, 0, 7], dtype=torch.int64), expert_token_nums_shared=[_FakeScalar(5)], group_list=_FakeIntVector([0] * 8), expand_x_shared=_FakeTensorLike("shared-hidden"), @@ -614,7 +694,10 @@ def fake_recv_attn_output(*, stage_idx, layer_idx, batch_size, ubatch_idx): monkeypatch.setattr(connector, "recv_attn_output", fake_recv_attn_output) - work_item = connector.recv_ffn_work_item(stage_idx=0, max_num_tokens=16) + work_item = connector.recv_ffn_work_item( + stage_idx=0, + max_num_tokens=16, + ) sent_output = connector.send_ffn_work_item_output( work_item, AFDF2ATransferPayload( diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 1abc3acb..2c2ad76a 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -12,6 +12,8 @@ from afd_plugin.config import AFDConfig from afd_plugin.connectors import ( AFDConnectorFactory, + AFDControlPayload, + AFDDPMetadata, AFDTransferContext, AFDTransferMetadata, AFDTransferState, @@ -84,7 +86,7 @@ def test_camp2p_factory_creates_connector(): assert connector.extra_info.core_num == 12 -def test_camp2p_topology_matches_original_rank_layout(): +def test_camp2p_topology_maps_every_attention_rank_to_its_ffn_group(): attn0 = build_camp2p_topology(_afd_config(role="attention"), 0) attn1 = build_camp2p_topology(_afd_config(role="attention"), 1) attn2 = build_camp2p_topology(_afd_config(role="attention"), 2) @@ -98,10 +100,50 @@ def test_camp2p_topology_matches_original_rank_layout(): assert (attn1.world_rank, attn1.p2p_rank, attn1.dp_metadata_destinations) == ( 3, 3, + (0,), + ) + assert (attn2.world_rank, attn2.p2p_rank, attn2.dp_metadata_destinations) == ( + 4, + 4, (1,), ) - assert not attn2.participates_in_p2p_group + assert attn2.participates_in_p2p_group assert (ffn1.world_rank, ffn1.p2p_rank) == (1, 1) + assert ffn1.dp_metadata_sources == (4, 5) + assert ffn1.p2p_world_size == 6 + + +def test_camp2p_control_payload_aggregates_input_ids_per_stage(): + torch = pytest.importorskip("torch") + dp_metadata_list = { + 0: AFDDPMetadata(torch.tensor([2, 3, 5, 7], dtype=torch.int32)), + 1: AFDDPMetadata(torch.tensor([11, 13, 17, 19], dtype=torch.int32)), + } + payload0 = AFDControlPayload( + dp_metadata_list=dp_metadata_list, + is_graph_capturing=False, + is_warmup=False, + input_ids_by_stage={0: [10, 11], 1: [12]}, + ) + payload1 = AFDControlPayload( + dp_metadata_list=dp_metadata_list, + is_graph_capturing=False, + is_warmup=False, + input_ids_by_stage={0: [20, 21, 22], 1: [23, 24]}, + ) + + merged = camp2p_module._aggregate_camp2p_control_payloads( + (payload0, payload1) + ) + + assert merged.input_ids_by_stage == { + 0: [10, 11, 20, 21, 22], + 1: [12, 23, 24], + } + assert torch.equal( + merged.dp_metadata_list[0].num_tokens_across_dp_cpu, + payload0.dp_metadata_list[0].num_tokens_across_dp_cpu, + ) def _init_ffn_connector(rank, vllm_config): diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index b650f480..2dceb924 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -198,6 +198,22 @@ def test_remote_proxy_requires_forward_metadata(monkeypatch): adapter.RemoteFFNProxy(layer_idx=0)(torch.ones(1, 4)) +def test_remote_proxy_exchanges_cam_during_profile(monkeypatch): + events = [] + _install_fake_forward_context(monkeypatch, events) + monkeypatch.setattr( + adapter, + "get_forward_context", + lambda: SimpleNamespace(in_profile_run=True, ubatch_idx=0), + ) + hidden_states = torch.ones(2, 4) + + output = adapter.RemoteFFNProxy(layer_idx=0)(hidden_states) + + assert torch.equal(output, hidden_states * 0.25) + assert [event[0] for event in events] == ["send", "yield", "recv"] + + def test_synchronous_model_forward_delegates_to_native(monkeypatch): calls = [] expected = torch.ones(1, 4) diff --git a/tests/unit/model_executor/models/test_deepseek_v4_async_schedule_contract.py b/tests/unit/model_executor/models/test_deepseek_v4_async_schedule_contract.py new file mode 100644 index 00000000..d64574cb --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_async_schedule_contract.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + + +native = pytest.importorskip("vllm_ascend.models.deepseek_v4") + +from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 + + +def test_dsv4_afd_schedule_executes_two_layers_and_two_stages(): + events = [] + pending = {} + stage_context = { + 0: SimpleNamespace( + ubatch_idx=0, + input_ids=(11, 13), + attn_metadata="metadata-0", + ), + 1: SimpleNamespace( + ubatch_idx=1, + input_ids=(17, 19), + attn_metadata="metadata-1", + ), + } + layers = [SimpleNamespace(layer_idx=0), SimpleNamespace(layer_idx=1)] + + class FakeConnector: + def send(self, layer_idx, stage_idx, context, hc_post, hc_comb): + assert context.ubatch_idx == stage_idx + pending[layer_idx, stage_idx] = (hc_post, hc_comb, context) + events.append(("send", layer_idx, stage_idx)) + + def recv(self, layer_idx, stage_idx): + hc_post, hc_comb, context = pending.pop((layer_idx, stage_idx)) + assert hc_post == f"hc-post-{layer_idx}-{stage_idx}" + assert hc_comb == f"hc-comb-{layer_idx}-{stage_idx}" + assert context is stage_context[stage_idx] + events.append(("recv", layer_idx, stage_idx)) + + connector = FakeConnector() + + def compute(layer, stage_idx): + context = stage_context[stage_idx] + assert context.ubatch_idx == stage_idx + assert context.input_ids == ((11, 13), (17, 19))[stage_idx] + assert context.attn_metadata == f"metadata-{stage_idx}" + pending[layer.layer_idx, stage_idx] = ( + f"hc-post-{layer.layer_idx}-{stage_idx}", + f"hc-comb-{layer.layer_idx}-{stage_idx}", + context, + ) + events.append(("compute", layer.layer_idx, stage_idx)) + + def send(layer, stage_idx): + hc_post, hc_comb, context = pending[layer.layer_idx, stage_idx] + connector.send(layer.layer_idx, stage_idx, context, hc_post, hc_comb) + + def recv(layer, stage_idx): + connector.recv(layer.layer_idx, stage_idx) + + adapter._run_two_stage_async_moe_schedule(layers, compute, send, recv) + + assert events == [ + ("compute", 0, 0), + ("send", 0, 0), + ("compute", 0, 1), + ("recv", 0, 0), + ("send", 0, 1), + ("compute", 1, 0), + ("recv", 0, 1), + ("send", 1, 0), + ("compute", 1, 1), + ("recv", 1, 0), + ("send", 1, 1), + ("recv", 1, 1), + ] + assert pending == {} + + +def test_dsv4_afd_schedule_propagates_post_send_failure(): + events = [] + layer = SimpleNamespace(layer_idx=0) + + def compute(_layer, stage_idx): + events.append(("compute", stage_idx)) + if stage_idx == 1: + raise RuntimeError("stage-one failed") + + def send(_layer, stage_idx): + events.append(("send", stage_idx)) + + with pytest.raises(RuntimeError, match="stage-one failed"): + adapter._run_two_stage_async_moe_schedule( + [layer], + compute, + send, + lambda *_args: events.append(("recv",)), + ) + + assert events == [("compute", 0), ("send", 0), ("compute", 1)] + + +@pytest.mark.parametrize( + ("input_ids", "inputs_embeds"), + ((None, None), (object(), object())), +) +def test_dsv4_afd_ubatch_rejects_inputs_without_token_ids( + monkeypatch, + input_ids, + inputs_embeds, +): + model = object.__new__(adapter.AFDDeepseekV4Model) + monkeypatch.setattr( + adapter, + "get_async_moe_ubatch_metadata_from_forward_context", + lambda: object(), + ) + + with pytest.raises(NotImplementedError, match="requires token input_ids"): + model.forward(input_ids, object(), None, inputs_embeds) + + +def test_dsv4_model_uses_native_forward_without_afd_stage_plan(monkeypatch): + model = object.__new__(adapter.AFDDeepseekV4Model) + sentinel = object() + monkeypatch.setattr( + adapter, + "get_async_moe_ubatch_metadata_from_forward_context", + lambda: None, + ) + monkeypatch.setattr(native.DeepseekV4Model, "forward", lambda *_args: sentinel) + + assert model.forward(None, object(), None, object()) is sentinel diff --git a/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py b/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py new file mode 100644 index 00000000..3565e05a --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py @@ -0,0 +1,35 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +from __future__ import annotations + +from pathlib import Path + + +def test_dsv4_async_gate_bypasses_native_moe_communicator() -> None: + """Async CAM must not invoke Ascend's EP/SP selector wrapper. + + That wrapper calls ``forward_context.moe_comm_method`` during DSV4 Hash + routing. AFD owns this communication boundary, so the Attention helper + must invoke the CANN routing operators directly on its local tokens. + """ + + source = Path( + "afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py", + ).read_text() + + assert "torch.ops._C_ascend.moe_gating_top_k_hash(" in source + assert "DeviceOperator.moe_gating_top_k(" in source + assert "correction_bias = correction_bias.to(router_logits.dtype)" in source + assert "from vllm_ascend.ops.fused_moe.experts_selector" not in source + assert "forward_context.moe_comm_method.pad_and_split_input_ids(" not in source + assert "connector.select_experts" not in source + + +def test_dsv4_async_gate_validates_local_hash_token_alignment() -> None: + source = Path( + "afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py", + ).read_text() + + assert "DSV4 Hash routing input_ids/token count mismatch on Attention" in source + assert "input_ids = input_ids.reshape(-1).to(torch.int64)" in source diff --git a/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py b/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py new file mode 100644 index 00000000..c87ab408 --- /dev/null +++ b/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from types import MethodType, SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("vllm_ascend.models.deepseek_v4") + +from afd_plugin.model_executor.models.npu import async_cam_layout # noqa: E402 +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_attention_gate, +) + + +def _make_proxy(monkeypatch, *, use_sequence_parallel: bool, in_profile_run: bool): + proxy = object.__new__(adapter.AFDDeepseekV4AttentionGateRemoteMoE) + torch.nn.Module.__init__(proxy) + topk_weights = torch.arange(32, dtype=torch.float32).reshape(16, 2) + topk_ids = torch.arange(32, dtype=torch.int32).reshape(16, 2) + monkeypatch.setattr( + deepseek_v4_attention_gate, + "compute_attention_gate_topk", + lambda _proxy, _hidden_states: (topk_weights, topk_ids), + ) + monkeypatch.setattr( + adapter, + "get_forward_context", + lambda: SimpleNamespace( + flash_comm_v1_enabled=use_sequence_parallel, + in_profile_run=in_profile_run, + ), + ) + return proxy, topk_weights, topk_ids + + +def test_dsv4_profile_gate_shards_plain_tp8_and_restores(monkeypatch): + proxy, topk_weights, topk_ids = _make_proxy( + monkeypatch, + use_sequence_parallel=False, + in_profile_run=True, + ) + tp_group = SimpleNamespace(world_size=8, rank_in_group=3) + monkeypatch.setattr(async_cam_layout, "get_tp_group", lambda: tp_group) + hidden_states = torch.arange(64, dtype=torch.float32).reshape(16, 4) + gathered_output = hidden_states + 100 + monkeypatch.setattr( + async_cam_layout, + "tensor_model_parallel_all_gather", + lambda _tensor, token_dim: gathered_output, + ) + sends = [] + + def send_and_receive(_self, local_hidden_states, **kwargs): + sends.append((local_hidden_states, kwargs)) + return local_hidden_states + 100 + + proxy._send_and_receive = MethodType(send_and_receive, proxy) + + output = proxy(hidden_states) + + sent_hidden_states, send_kwargs = sends[0] + assert torch.equal(sent_hidden_states, hidden_states[6:8]) + assert torch.equal(send_kwargs["topk_weights"], topk_weights[6:8]) + assert torch.equal(send_kwargs["topk_ids"], topk_ids[6:8]) + assert torch.equal(output, gathered_output) + + +def test_dsv4_gate_keeps_sequence_parallel_tokens_rank_local(monkeypatch): + proxy, topk_weights, topk_ids = _make_proxy( + monkeypatch, + use_sequence_parallel=True, + in_profile_run=False, + ) + monkeypatch.setattr( + async_cam_layout, + "get_tp_group", + lambda: SimpleNamespace(world_size=8, rank_in_group=3), + ) + hidden_states = torch.arange(64, dtype=torch.float32).reshape(16, 4) + sends = [] + + def send_and_receive(_self, local_hidden_states, **kwargs): + sends.append((local_hidden_states, kwargs)) + return local_hidden_states + 1 + + proxy._send_and_receive = MethodType(send_and_receive, proxy) + + output = proxy(hidden_states) + + sent_hidden_states, send_kwargs = sends[0] + assert sent_hidden_states is hidden_states + assert send_kwargs["topk_weights"] is topk_weights + assert send_kwargs["topk_ids"] is topk_ids + assert torch.equal(output, hidden_states + 1) diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_forward_context.py index bff67d8d..3376717b 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_forward_context.py @@ -189,25 +189,15 @@ def test_async_cam_profile_forward_runs_matched_connector_io(monkeypatch): flash_comm_v1_enabled=True, ) monkeypatch.setattr(async_forward, "get_forward_context", lambda: forward_context) - monkeypatch.setattr( - async_forward, - "maybe_apply_dbo_yield", - lambda hidden_states, **_kwargs: hidden_states, - ) - connector_calls = [] - pending_outputs = [] + connector_calls: list[str] = [] - def send_attn_output(hidden_states, context, **kwargs): - layer_idx = context.metadata.layer_idx - connector_calls.append(("send", layer_idx, context.metadata.stage_idx)) - pending_outputs.append((layer_idx, hidden_states.clone())) + def send_attn_output(*args, **kwargs): + connector_calls.append("send") - def recv_ffn_output(*, ref_tensor, ubatch_idx): - layer_idx, dispatched_hidden_states = pending_outputs.pop(0) - assert torch.equal(ref_tensor, dispatched_hidden_states) - connector_calls.append(("recv", layer_idx, ubatch_idx)) - return dispatched_hidden_states + 10 + def recv_ffn_output(ref_tensor, ubatch_idx): + connector_calls.append("recv") + return ref_tensor connector = SimpleNamespace( send_attn_output=send_attn_output, @@ -215,49 +205,10 @@ def recv_ffn_output(*, ref_tensor, ubatch_idx): ) afd_metadata = SimpleNamespace(connector=connector, stage_idx=0) - dispatch_layouts = [] - - def prepare_dispatch_payload( - hidden_states, - topk_weights, - topk_ids, - router_logits, - *, - use_sequence_parallel, - ): - assert use_sequence_parallel is True - layout = object() - dispatch_layouts.append(layout) - return SimpleNamespace( - hidden_states=hidden_states, - topk_weights=topk_weights, - topk_ids=topk_ids, - router_logits=router_logits, - layout=layout, - ) - - restored_layouts = [] - - def restore_dispatch_output(local_output, layout): - restored_layouts.append(layout) - return local_output - - monkeypatch.setattr( - async_forward, - "prepare_cam_dispatch_payload", - prepare_dispatch_payload, - ) - monkeypatch.setattr( - async_forward, - "restore_cam_dispatch_output", - restore_dispatch_output, - ) - class _ProfileMoELayer: is_moe_layer = True - def __init__(self, layer_idx): - self.layer_idx = layer_idx + layer_idx = 0 def compute_attn_output( self, @@ -275,7 +226,7 @@ def compute_attn_output( ) model = SimpleNamespace( - layers=[_ProfileMoELayer(0), _ProfileMoELayer(1)], + layers=[_ProfileMoELayer(), _ProfileMoELayer()], start_layer=0, end_layer=2, ) @@ -289,16 +240,9 @@ def compute_attn_output( afd_metadata, ) - assert torch.equal(output, hidden_states + 22) + assert torch.equal(output, hidden_states + 2) assert residual is None - assert connector_calls == [ - ("send", 0, 0), - ("recv", 0, 0), - ("send", 1, 0), - ("recv", 1, 0), - ] - assert restored_layouts == dispatch_layouts - assert pending_outputs == [] + assert connector_calls == ["send", "recv", "send", "recv"] def test_deepseek_afd_wrapper_keeps_full_model_compile_enabled(): diff --git a/tests/unit/v1/worker/test_npu_mla_graph.py b/tests/unit/v1/worker/test_npu_mla_graph.py index 567f3ddf..a1372583 100644 --- a/tests/unit/v1/worker/test_npu_mla_graph.py +++ b/tests/unit/v1/worker/test_npu_mla_graph.py @@ -382,6 +382,22 @@ def test_ubatch_wrapper_rejects_enpu(monkeypatch): ) +def test_ubatch_wrapper_requires_native_two_ubatch_config(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + config = SimpleNamespace( + parallel_config=SimpleNamespace(num_ubatches=0), + compilation_config=object(), + ) + + with pytest.raises(AssertionError): + wrapper_module.AscendUBatchWrapper( + lambda: None, + config, + wrapper_module.CUDAGraphMode.NONE, + torch.device("cpu"), + ) + + def test_npu_graph_key_separates_stage_shapes_and_lora(monkeypatch): wrapper_module = _load_ubatch_wrapper_module(monkeypatch) diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 79674247..502c6ddc 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -89,12 +89,15 @@ class _RecordingConnector: def __init__(self): self.dp_metadata_updates = [] self.sent_dp_metadata_lists = [] + self.dp_metadata_payloads = [] + self.sent_dp_metadata_payloads = [] # The runners reach the control plane through connector.control_plane; # the fake serves as both. self.control_plane = self def update_state_from_dp_metadata(self, payload): assert isinstance(payload, AFDControlPayload) + self.dp_metadata_payloads.append(payload) self.dp_metadata_updates.append( ( payload.dp_metadata_list, @@ -105,6 +108,7 @@ def update_state_from_dp_metadata(self, payload): def send_dp_metadata_list(self, payload): assert isinstance(payload, AFDControlPayload) + self.sent_dp_metadata_payloads.append(payload) self.sent_dp_metadata_lists.append( ( payload.dp_metadata_list, @@ -477,9 +481,10 @@ def __call__(self, **_model_inputs): updates = [] runner._update_full_graph_params_if_needed = lambda *args: updates.append(args) + input_ids = object() result = runner._model_forward( 8, - input_ids=None, + input_ids=input_ids, positions=object(), intermediate_tensors=None, inputs_embeds=None, @@ -487,6 +492,7 @@ def __call__(self, **_model_inputs): assert result == "hidden_states" assert len(updates) == expected_updates + assert forward_context.input_ids is input_ids def test_npu_attention_runner_installs_mla_graph_wrapper(monkeypatch): @@ -537,6 +543,7 @@ def test_npu_attention_runner_builds_and_sets_metadata(): runner._afd_is_graph_capturing = False runner._afd_pending_metadata = None runner._afd_transaction_counter = 0 + runner.ubatch_slices = None runner._afd_suppress_metadata_send = False forward_context = SimpleNamespace( additional_kwargs={}, @@ -641,7 +648,12 @@ def test_npu_attention_runner_sends_per_ubatch_dp_metadata(): ), ] - runner._send_dp_metadata(None, ubatch_slices) + torch = pytest.importorskip("torch") + runner._send_dp_metadata( + None, + ubatch_slices, + input_ids=torch.tensor([10, 11, 12, 13, 20, 21, 22]), + ) dp_metadata_list = runner.connector.dp_metadata_updates[0][0] assert sorted(dp_metadata_list) == [0, 1] @@ -651,6 +663,10 @@ def test_npu_attention_runner_sends_per_ubatch_dp_metadata(): assert sorted(sent_dp_metadata_list) == [0, 1] assert _tokens(sent_dp_metadata_list[0]) == [4] assert _tokens(sent_dp_metadata_list[1]) == [3] + assert runner.connector.sent_dp_metadata_payloads[0].input_ids_by_stage == { + 0: [10, 11, 12, 13], + 1: [20, 21, 22], + } def test_npu_attention_capture_microbatch_also_captures_single_stage(): @@ -959,6 +975,10 @@ def build_stage_metadata(self, *args, **kwargs): slice(550, 1099), ] assert materialized_full_metadata == [(full_attn_metadata, runner.positions)] + assert runner.ubatch_slices is None + assert runner._afd_pending_metadata.num_stages == 2 + assert runner._afd_pending_metadata.tokens_start_loc == [0, 550] + assert runner._afd_pending_metadata.tokens_lens == [550, 549] def test_npu_attention_runner_isolates_dsa_caches_per_stage(monkeypatch): @@ -1349,6 +1369,37 @@ def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): ] +def test_npu_ffn_input_ids_require_exact_hidden_state_alignment(): + _require_npu_runtime() + from afd_plugin.v1.worker.npu.ffn_model_runner import _install_ffn_input_ids + + torch = pytest.importorskip("torch") + forward_context = SimpleNamespace() + _install_ffn_input_ids( + forward_context, + [10, 11], + num_tokens=2, + device=torch.device("cpu"), + ) + assert forward_context.input_ids.tolist() == [10, 11] + + with pytest.raises(RuntimeError, match="must align"): + _install_ffn_input_ids( + forward_context, + [12], + num_tokens=2, + device=torch.device("cpu"), + ) + + _install_ffn_input_ids( + forward_context, + None, + num_tokens=2, + device=torch.device("cpu"), + ) + assert forward_context.input_ids is None + + 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 @@ -2459,3 +2510,35 @@ def test_npu_attention_runner_load_model_initializes_connector_after_weights( expected.append("connector_init") assert events == expected assert connector.is_initialized is True + + +def test_npu_attention_runner_afd_ubatching_does_not_install_native_wrapper( + monkeypatch, +): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import attention_model_runner + + events = [] + connector = _LifecycleConnector(events) + runner = object.__new__(attention_model_runner.AFDNPUAttentionModelRunner) + runner.connector = connector + runner.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace( + use_ubatching=False, + num_ubatches=0, + ), + ) + monkeypatch.setattr( + attention_model_runner.NPUModelRunner, + "load_model", + lambda self: events.append("model_load"), + ) + monkeypatch.setattr( + attention_model_runner.AFDNPUAttentionModelRunner, + "_install_ascend_ubatch_wrapper", + lambda self: events.append("wrapper_install"), + ) + + runner.load_model() + + assert events == ["model_load", "connector_init"] diff --git a/tools/itask/launch_dsv4_afd_cross_node.sh b/tools/itask/launch_dsv4_afd_cross_node.sh new file mode 100644 index 00000000..994bac92 --- /dev/null +++ b/tools/itask/launch_dsv4_afd_cross_node.sh @@ -0,0 +1,134 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Two A3-task DSV4 async CAM deployment. +# +# Attention is one global DP3TP8 vLLM process group, started once on each +# node: node1 owns DP0-1 (NPUs 0-15) and hosts the API; node2 owns DP2 +# (NPUs 0-7) and joins headlessly. Both Attention invocations use node1 as +# the vLLM DP coordinator. All 24 Attention ranks and 8 FFN ranks join one +# CAM communicator through node1; node2 hosts the FFN DP8/EP8 ranks on NPUs 8-15. +: "${ROLE:?ROLE must be attention or ffn}" +: "${NODE_IP:?NODE_IP is required}" +: "${ATTENTION_NODE_ID:=1}" # 1 = node1 (DP0-1), 2 = node2 (DP2) +: "${API_PORT:=8900}" +: "${DP_RPC_PORT:=29550}" +: "${NIC_NAME:=eth0}" +: "${VISIBLE_DEVICES:=0,1,2,3,4,5,6,7}" +: "${PROFILE_VARIANT:=none}" +: "${PROFILE_ROOT:=/tmp/dsv4_afd_profiles}" +: "${GPU_MEMORY_UTILIZATION:=0.70}" +: "${HCCL_BUFFSIZE:=512}" +: "${MAX_NUM_BATCHED_TOKENS:=1024}" + +case "$ROLE" in + attention|ffn) ;; + *) echo "ROLE must be attention or ffn" >&2; exit 2 ;; +esac + +if [[ "$ROLE" == attention ]]; then + case "$ATTENTION_NODE_ID" in + 1) + ATTN_DP_SIZE_LOCAL=2 + ATTN_DP_START_RANK=0 + ATTN_HEADLESS_ARGS=() + API_SERVER_ARGS=(--api-server-count 1) + ;; + 2) + ATTN_DP_SIZE_LOCAL=1 + ATTN_DP_START_RANK=2 + ATTN_HEADLESS_ARGS=(--headless) + API_SERVER_ARGS=() + ;; + *) + echo "ATTENTION_NODE_ID must be 1 or 2 for ROLE=attention" >&2 + exit 2 + ;; + esac +fi + +case "$PROFILE_VARIANT" in + none) + PROFILER_ARGS=() + ;; + full) + PROFILE_DIR="$PROFILE_ROOT/afd_${ROLE}_node${ATTENTION_NODE_ID}_full" + PROFILER_ARGS=(--profiler-config "{\"profiler\":\"torch\",\"torch_profiler_dir\":\"$PROFILE_DIR\",\"torch_profiler_with_stack\":true,\"torch_profiler_record_shapes\":true,\"torch_profiler_with_memory\":false,\"torch_profiler_use_gzip\":false,\"ignore_frontend\":true,\"delay_iterations\":10,\"max_iterations\":9,\"warmup_iterations\":0,\"active_iterations\":10,\"wait_iterations\":0}") + ;; + ops) + PROFILE_DIR="$PROFILE_ROOT/afd_${ROLE}_node${ATTENTION_NODE_ID}_ops" + PROFILER_ARGS=(--profiler-config "{\"profiler\":\"torch\",\"torch_profiler_dir\":\"$PROFILE_DIR\",\"torch_profiler_with_stack\":false,\"torch_profiler_record_shapes\":false,\"torch_profiler_with_memory\":false,\"torch_profiler_use_gzip\":false,\"ignore_frontend\":true,\"delay_iterations\":10,\"max_iterations\":9,\"warmup_iterations\":0,\"active_iterations\":10,\"wait_iterations\":0}") + ;; + *) + echo "PROFILE_VARIANT must be none, full, or ops" >&2 + exit 2 + ;; +esac + +MODEL_PATH=/mnt/sfs_turbo/models/DeepSeek-V4-Flash-w8a8-mtp +# Resolve the plugin from this launcher instead of a task-specific workdir. +# ``itask sync`` may assign each task a different workspace name. +SCRIPT_DIR=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +PLUGIN_ROOT=${PLUGIN_ROOT:-$(cd -- "$SCRIPT_DIR/../.." && pwd)} +: "${NODE1_IP:?NODE1_IP must be the validated node1 Pod IP}" +: "${NODE2_IP:?NODE2_IP must be the validated node2 Pod IP}" + +source /usr/local/Ascend/cann-9.0.1/set_env.sh +export ASCEND_RT_VISIBLE_DEVICES="$VISIBLE_DEVICES" +export PYTHONPATH="$PLUGIN_ROOT:/vllm-workspace/vllm-ascend:/vllm-workspace/vllm:${PYTHONPATH:-}" +export VLLM_PLUGINS=ascend,afd +export VLLM_WORKER_MULTIPROC_METHOD=spawn +export AFD_FORCE_SPAWN_MULTIPROCESSING=1 +CAM_VENDOR=/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM +CAM_OPAPI_DIR="$CAM_VENDOR/op_api/lib" +CAM_OPAPI="$CAM_OPAPI_DIR/libopapi.so" +TORCH_NPU_LIB=/usr/local/python3.12.13/lib/python3.12/site-packages/torch_npu/lib +TORCH_LIB=/usr/local/python3.12.13/lib/python3.12/site-packages/torch/lib +export ASCEND_CUSTOM_OPP_PATH="$CAM_VENDOR:${ASCEND_CUSTOM_OPP_PATH:-}" +# umdk_cam_op_lib resolves libopapi.so by name on first request. Put CAM's +# vendor implementation before CANN's stock libopapi.so, which lacks CAM ops. +export LD_LIBRARY_PATH="$TORCH_LIB":"$TORCH_NPU_LIB":/usr/local/Ascend/driver/lib64/driver:/usr/local/Ascend/driver/lib64:"$CAM_OPAPI_DIR":"$CAM_VENDOR/op_api":/usr/local/Ascend/cann-9.0.1/aarch64-linux/lib64:/usr/local/Ascend/cann-9.0.1/runtime/lib64:${LD_LIBRARY_PATH:-} +export CAM_CUST_OPAPI_LIB_PATH="$CAM_OPAPI" +export LD_PRELOAD="$CAM_OPAPI${LD_PRELOAD:+:$LD_PRELOAD}" +export HCCL_IF_IP="$NODE_IP" +export HCCL_SOCKET_IFNAME="$NIC_NAME" +export GLOO_SOCKET_IFNAME="$NIC_NAME" +export TP_SOCKET_IFNAME="$NIC_NAME" +export HCCL_BUFFSIZE HCCL_OP_EXPANSION_MODE=AIV +export HCCL_CONNECT_TIMEOUT=1800 HCCL_EXEC_TIMEOUT=1800 +export OMP_PROC_BIND=false OMP_NUM_THREADS=10 +export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True +export VLLM_ASCEND_ENABLE_FLASHCOMM1=0 + +if [[ "$ROLE" == attention ]]; then + # Keep all Attention ranks in one vLLM DP3TP8 process group. The connector + # derives AFD role rank from vLLM's global DP rank and TP rank, so no manual + # per-node role-rank offset is used here. + PARALLEL_ARGS=( + --data-parallel-size 3 + --data-parallel-size-local "$ATTN_DP_SIZE_LOCAL" + --data-parallel-start-rank "$ATTN_DP_START_RANK" + --data-parallel-address "$NODE1_IP" + --data-parallel-rpc-port "$DP_RPC_PORT" + --tensor-parallel-size 8 + "${ATTN_HEADLESS_ARGS[@]}" + ) + WORKER=afd_plugin.v1.worker.npu.AFDNPUAttentionWorker + MODEL_NAME=dsv4-afd-attention +else + PARALLEL_ARGS=(--data-parallel-size 8 --tensor-parallel-size 1) + API_SERVER_ARGS=(--api-server-count 1) + WORKER=afd_plugin.v1.worker.npu.AFDNPUFFNWorker + MODEL_NAME=dsv4-afd-ffn +fi + +ADDITIONAL_CONFIG="{\"enable_force_load_balance\":false,\"afd\":{\"role\":\"$ROLE\",\"connector\":\"CAMAsyncAFDConnector\",\"async\":true,\"host\":\"$NODE1_IP\",\"port\":1239,\"num_attention_ranks\":24,\"num_ffn_ranks\":8,\"compute_gate_on_attention\":true,\"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":8,\"async_moe_ubatching\":true}}}" + +exec env VLLM_USE_V1=1 /usr/local/python3.12.13/bin/vllm serve "$MODEL_PATH" \ + --host 0.0.0.0 --port "$API_PORT" "${API_SERVER_ARGS[@]}" --served-model-name "$MODEL_NAME" \ + --worker-cls "$WORKER" "${PARALLEL_ARGS[@]}" --enable-expert-parallel \ + --enforce-eager --quantization ascend --tokenizer-mode deepseek_v4 \ + --block-size 128 --max-model-len 8192 --max-num-batched-tokens "$MAX_NUM_BATCHED_TOKENS" --max-num-seqs 2 \ + --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 16}' \ + --trust-remote-code --no-enable-prefix-caching --enable-chunked-prefill \ + --additional-config "$ADDITIONAL_CONFIG" "${PROFILER_ARGS[@]}"