From 7f4ac5a0b6ca3b29a0faac5231a7673e500730b1 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 14:20:15 +0800 Subject: [PATCH 1/8] feat(npu): capture AsyncCAM layered FFN in a reusable graph Signed-off-by: jiangkuaixue123 --- afd_plugin/compat/npu/feature_validation.py | 18 +- afd_plugin/connectors/npu/async_cam.py | 295 ++++++++++++++++- afd_plugin/distributed/afd_process_group.py | 5 + .../v1/worker/npu/attention_model_runner.py | 22 +- afd_plugin/v1/worker/npu/ffn_model_runner.py | 205 +++++++++++- afd_plugin/v1/worker/npu/ffn_worker.py | 108 ++++++- .../op_kernel/afd_async_combine_send.h | 8 + .../op_host/afd_async_dispatch_recv.cpp | 15 +- .../op_kernel/afd_async_dispatch_recv.h | 165 +++++++++- .../afd_async_dispatch_recv_tiling.h | 1 + csrc/npu/pybind/torch_binding_cam_async.cpp | 9 +- .../async_cam_communication_graph_probe.py | 306 ++++++++++++++++++ .../npu/async_cam_ffn_graph_service_smoke.py | 125 +++++++ tests/npu/async_cam_graph_bootstrap_probe.py | 97 ++++++ .../npu/test_async_cam_layered_graph_probe.py | 121 +++++++ tests/unit/compat/npu/test_cam_async_ops.py | 29 +- .../connectors/test_async_cam_connector.py | 146 +++++++++ tests/unit/v1/worker/test_npu_runtime.py | 255 ++++++++++++++- 18 files changed, 1886 insertions(+), 44 deletions(-) create mode 100644 tests/npu/async_cam_communication_graph_probe.py create mode 100644 tests/npu/async_cam_ffn_graph_service_smoke.py create mode 100644 tests/npu/async_cam_graph_bootstrap_probe.py create mode 100644 tests/npu/test_async_cam_layered_graph_probe.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index ffc54f53d..653f3c301 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -184,9 +184,21 @@ def _fail_if_unsupported_npu_afd_async_features( "with async=true and connector='CAMAsyncAFDConnector'", ) if not bool(vllm_config.model_config.enforce_eager): - raise RuntimeError( - "CAMAsyncAFDConnector supports only eager Attention/FFN execution", - ) + graph_mode = vllm_config.compilation_config.cudagraph_mode.name + if not ( + afd_config.role == "ffn" + and async_cam_layered_gmm_enabled() + and _is_dsv4_target(vllm_config) + and afd_config.compute_gate_on_attention + and extra_info.dynamic_quant == 1 + and not vllm_config.use_v2_model_runner + and graph_mode == "FULL" + ): + raise RuntimeError( + "CAMAsyncAFDConnector graph execution requires FFN FULL, " + "DeepSeek V4, layered W4A8 GMM, Attention-side gate, " + "dynamicQuant=1, and ModelRunnerV1; Attention remains eager" + ) if bool(parallel_config.enable_dbo) or bool(parallel_config.use_ubatching): raise RuntimeError( "CAMAsyncAFDConnector does not support vLLM native ubatching/DBO", diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index f79c33070..f5904284f 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """Ascend CAM asynchronous connector for Attention/FFN disaggregation. -``CAMAsyncAFDConnector`` is the eager-only Ascend inference data path. +``CAMAsyncAFDConnector`` is the Ascend inference data path. Attention ranks run MoE routing, submit activations with CAM async dispatch-send, and receive combined expert output with combine-recv. FFN ranks receive routed-expert activations with dispatch-recv, execute their @@ -15,10 +15,10 @@ the CAM operator payload and does not create a separate Gloo DP-metadata control plane. -The supported deployment requires ``async=true``, eager execution, Ascend CAM -operator packages, and matching topology/configuration on every rank. Regular -prefill and autoregressive decode steps use the same connector; vLLM native DBO -and ACL graph execution are not supported. +The supported deployment requires ``async=true``, Ascend CAM operator packages, +and matching topology/configuration on every rank. Attention runs eager; the +layered W4A8 FFN path can use one FULL graph. Regular prefill and autoregressive +decode steps use the same connector; vLLM native DBO is not supported. Optional AFD-managed MoE ubatching is a separate two-stage pipeline using request boundaries or token-balanced stages for DP+TP/SP. See ``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration, rank @@ -28,13 +28,16 @@ from __future__ import annotations import os +import time from collections.abc import Mapping from dataclasses import dataclass from datetime import timedelta from typing import TYPE_CHECKING, Any, Final, cast import torch +import torch.distributed as dist from torch import Tensor +from torch.distributed.distributed_c10d import PrefixStore, Store from vllm.logger import init_logger from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available @@ -72,6 +75,12 @@ ASYNC_MOE_NUM_STAGES = 2 ASYNC_MOE_REQUEST_SPLIT = "request" ASYNC_MOE_TOKEN_SPLIT = "token" +FFN_STARTUP_TIMEOUT_SECONDS = 300 +FFN_STARTUP_POLL_SECONDS = 0.25 +FFN_STARTUP_NONCE_BYTES = 16 +STARTUP_WARMUP_TOKEN_COUNT = 1 +STARTUP_WARMUP_HIDDEN_VALUE = 0.25 +FFN_GRAPH_STOP_BUFFER_INT64_VALUES = 4 _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( { @@ -262,6 +271,10 @@ def __init__( self.comm_id = CAM_COMM_ID self.tp_size = extra_info.attn_ranks_per_dp self.cam_pg: ProcessGroup | None = None + self._rendezvous_store: Store | None = None + self._startup_store: Store | None = None + self._ffn_graph_stop_stream: torch.npu.Stream | None = None + self._ffn_graph_stop_host: Tensor | None = None self.topology = build_async_topology( afd_config, role_rank, @@ -304,6 +317,7 @@ def init_afd_connector(self) -> None: pg_options=create_hccl_process_group_options( self.hccl_buffer_size_mb, ), + on_rendezvous=self._retain_rendezvous_store, ) backend = self.cam_pg._get_backend(torch.device("npu")) self.group_name = str(backend.get_hccl_comm_name(self.world_rank)) @@ -316,17 +330,280 @@ def init_afd_connector(self) -> None: ) self._initialized = True - def close(self) -> None: + def _retain_rendezvous_store(self, store: Store) -> None: + self._rendezvous_store = store + + def _get_startup_store(self) -> Store: + """Make one rank-shared namespace tied to this CAM process group.""" + if self._startup_store is not None: + return self._startup_store + if self.cam_pg is None or self._rendezvous_store is None: + raise RuntimeError("CAM startup coordination requires an initialized group") + nonce = torch.zeros( + FFN_STARTUP_NONCE_BYTES, + dtype=torch.uint8, + device=f"npu:{self.local_rank}", + ) + if self.world_rank == 0: + nonce.copy_( + torch.tensor( + list(os.urandom(FFN_STARTUP_NONCE_BYTES)), + dtype=torch.uint8, + device=nonce.device, + ) + ) + logger.info("CAM startup nonce broadcast start rank=%d", self.world_rank) + dist.broadcast(nonce, src=0, group=self.cam_pg) + logger.info("CAM startup nonce broadcast done rank=%d", self.world_rank) + torch.npu.synchronize() + logger.info("CAM startup nonce synchronize done rank=%d", self.world_rank) + namespace = bytes(nonce.cpu().tolist()).hex() + self._startup_store = PrefixStore( + f"{AFD_ASYNC_CAM_GROUP_NAME}/startup/{namespace}", + self._rendezvous_store, + ) + return self._startup_store + + def prepare_startup_coordination(self) -> None: + self._get_startup_store() + + def enable_ffn_graph_stop(self, first_layer_idx: int) -> None: + """Keep a graph-owned 32-byte stop input and an independent copy stream.""" + if self.topology.role != "ffn" or not self.is_initialized: + raise RuntimeError("CAM graph stop input requires initialized FFN rank") + if self._ffn_graph_stop_stream is not None: + return + initial = torch.zeros(FFN_GRAPH_STOP_BUFFER_INT64_VALUES, dtype=torch.int64) + initial[1] = first_layer_idx + stop = initial.clone().pin_memory() + stop[0] = 1 + self._placeholder = initial.view(self.activation_dtype).to( + device=f"npu:{self.local_rank}" + ) + self._ffn_graph_stop_host = stop.view(self.activation_dtype) + self._ffn_graph_stop_stream = torch.npu.Stream(device=self.local_rank) + + def signal_ffn_graph_stop(self) -> None: + """Write the stop bit while the graph stream may be blocked in DR.""" + stream = self._ffn_graph_stop_stream + source = self._ffn_graph_stop_host + destination = self._placeholder + if stream is None or source is None or destination is None: + raise RuntimeError("CAM FFN graph stop input was not initialized") + with torch.npu.stream(stream): + destination.copy_(source, non_blocking=True) + stream.synchronize() + + def _wait_for_startup_entries( + self, keys: list[str], failure_keys: list[str] + ) -> list[str]: + store = self._get_startup_store() + deadline = time.monotonic() + FFN_STARTUP_TIMEOUT_SECONDS + while True: + for key in failure_keys: + if store.check([key]): + status = store.get(key).decode() + if key.startswith("attn/failure/") or status.startswith("failed:"): + raise RuntimeError(f"CAM startup {key} {status}") + values = [] + pending = [] + for key in keys: + if store.check([key]): + values.append(store.get(key).decode()) + else: + pending.append(key) + if not pending: + return values + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for CAM startup keys: {sorted(pending)}" + ) + time.sleep(FFN_STARTUP_POLL_SECONDS) + + def publish_ffn_mode(self, layer_idx: int | None) -> None: + store = self._get_startup_store() + if self.topology.role != "ffn": + raise RuntimeError("Only CAM FFN ranks may publish their mode") + mode = "eager" if layer_idx is None else f"graph:{layer_idx}" + store.set(f"ffn/mode/{self.world_rank}", mode) + + def wait_for_ffn_modes(self) -> int | None: + ffn_ranks = range(self.attn_size, self.topology.world_size) + modes = self._wait_for_startup_entries( + [f"ffn/mode/{rank}" for rank in ffn_ranks], + [f"ffn/{rank}" for rank in ffn_ranks], + ) + if len(set(modes)) != 1: + raise RuntimeError(f"CAM FFN startup modes disagree: {modes}") + mode = modes[0] + if mode == "eager": + return None + if not mode.startswith("graph:"): + raise RuntimeError(f"Invalid CAM FFN startup mode: {mode}") + return int(mode.removeprefix("graph:")) + + def publish_attention_startup_failure(self, error: str) -> None: + self._get_startup_store().set(f"attn/failure/{self.world_rank}", error) + + def _attention_failure_keys(self) -> list[str]: + return [f"attn/failure/{rank}" for rank in range(self.attn_size)] + + def _ffn_failure_keys(self) -> list[str]: + return [ + f"ffn/{rank}" for rank in range(self.attn_size, self.topology.world_size) + ] + + def wait_for_attention_warmup_prepared(self) -> None: + self._wait_for_startup_entries( + [f"attn/prepared/{rank}" for rank in range(self.attn_size)], + self._attention_failure_keys(), + ) + + def publish_ffn_warmup_started(self) -> None: + self._get_startup_store().set(f"ffn/warmup/start/{self.world_rank}", "ready") + + def publish_ffn_warmup_done(self) -> None: + self._get_startup_store().set(f"ffn/warmup/done/{self.world_rank}", "ready") + + def wait_for_communication_warmup_done(self) -> None: + self._wait_for_startup_entries( + [f"attn/done/{rank}" for rank in range(self.attn_size)] + + [ + f"ffn/warmup/done/{rank}" + for rank in range(self.attn_size, self.topology.world_size) + ], + self._attention_failure_keys() + self._ffn_failure_keys(), + ) + + def run_attention_startup_warmup(self, layer_idx: int) -> None: + """Complete one isolated CAM transaction before FFN graph capture.""" + if self.topology.role != "attention": + raise RuntimeError("CAM Attention warmup requires Attention role") + device = f"npu:{self.local_rank}" + hidden = torch.full( + (STARTUP_WARMUP_TOKEN_COUNT, self.hidden_size), + STARTUP_WARMUP_HIDDEN_VALUE, + dtype=self.activation_dtype, + device=device, + ) + tp_rank = self.world_rank % self.tp_size + expert_ids = [ + ((tp_rank * self.topk + index) % self.ffn_size) * self.expert_per_rank + + ((tp_rank * self.topk + index) // self.ffn_size) % self.expert_per_rank + for index in range(self.topk) + ] + ids = torch.tensor([expert_ids], dtype=torch.int32, device=device) + weights = torch.full( + (STARTUP_WARMUP_TOKEN_COUNT, self.topk), + 1.0 / self.topk, + dtype=torch.float32, + device=device, + ) + store = self._get_startup_store() + store.set(f"attn/prepared/{self.world_rank}", "ready") + self._wait_for_startup_entries( + [ + f"ffn/warmup/start/{rank}" + for rank in range(self.attn_size, self.topology.world_size) + ], + self._ffn_failure_keys(), + ) + logger.info("CAM Attention startup warmup DS rank=%d", self.world_rank) + torch.ops.afd_ascend.afd_async_dispatch_send( + hidden, + ids, + self.comm_args, + self.comm_id, + self.max_num_batched_tokens, + STARTUP_WARMUP_TOKEN_COUNT, + self.hidden_size, + self.topk, + self.ffn_size, + self.attn_size, + self.expert_per_rank, + self.world_rank, + self.topology.world_size, + layer_idx, + self.tp_size, + self.dynamic_quant, + self.group_name, + ) + warmup_output = torch.ops.afd_ascend.afd_async_combine_recv( + self._placeholder, + ids, + weights, + self.comm_args, + self.comm_id, + STARTUP_WARMUP_TOKEN_COUNT, + self.hidden_size, + self.topk, + self.ffn_size, + self.attn_size, + self.expert_per_rank, + self.world_rank, + self.topology.world_size, + self.group_name, + ) + torch.npu.synchronize() + del warmup_output + logger.info("CAM Attention startup warmup CR done rank=%d", self.world_rank) + store.set(f"attn/done/{self.world_rank}", "ready") + self.wait_for_communication_warmup_done() + + def publish_ffn_startup(self, error: str | None = None) -> None: + store = self._get_startup_store() + if self.topology.role != "ffn": + raise RuntimeError("Only CAM FFN ranks may publish startup status") + key = f"ffn/{self.world_rank}" + if error is not None: + store.set(key, f"failed:{error}") + return + status = store.compare_set(key, "", "ready").decode() + if status != "ready": + raise RuntimeError(f"CAM FFN startup {key} {status}") + + def wait_for_ffn_ready(self) -> None: + store = self._get_startup_store() + keys = [ + f"ffn/{rank}" for rank in range(self.attn_size, self.topology.world_size) + ] + deadline = time.monotonic() + FFN_STARTUP_TIMEOUT_SECONDS + while True: + pending = [] + for key in keys: + if not store.check([key]): + pending.append(key) + continue + status = store.get(key).decode() + if status != "ready": + raise RuntimeError(f"CAM FFN startup {key} {status}") + if not pending: + return + if pending and time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for CAM FFN ranks: {sorted(pending)}" + ) + time.sleep(FFN_STARTUP_POLL_SECONDS) + + def close(self, *, release_buffers: bool = True) -> None: """Destroy the HCCL process group and clear pending transfer states.""" if self.cam_pg is not None: - import torch.distributed as dist - dist.destroy_process_group(self.cam_pg) self.cam_pg = None + self._initialized = False + if release_buffers: + self.release_closed_buffers() + + def release_closed_buffers(self) -> None: + if self._initialized: + raise RuntimeError("Cannot release active CAM operator buffers") self.comm_args = None self._placeholder = None + self._ffn_graph_stop_stream = None + self._ffn_graph_stop_host = None + self._startup_store = None + self._rendezvous_store = None self._pending_attention_payloads.clear() - self._initialized = False def select_experts(self, **kwargs: Any) -> tuple[Tensor, Tensor]: """Run the pinned vLLM-Ascend expert selector on Attention.""" diff --git a/afd_plugin/distributed/afd_process_group.py b/afd_plugin/distributed/afd_process_group.py index 5dbfab507..519aa3439 100644 --- a/afd_plugin/distributed/afd_process_group.py +++ b/afd_plugin/distributed/afd_process_group.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections.abc import Callable from datetime import timedelta from typing import Any @@ -12,6 +13,7 @@ from torch.distributed.distributed_c10d import ( PrefixStore, ProcessGroup, + Store, _new_process_group_helper, _update_default_pg, _world, @@ -67,6 +69,7 @@ def init_afd_process_group( group_name: str, timeout: timedelta, pg_options: Any | None = None, + on_rendezvous: Callable[[Store], None] | None = None, ) -> ProcessGroup: """Create a plugin-owned process group without patching vLLM source. @@ -83,6 +86,8 @@ def init_afd_process_group( ) store, rank, world_size = next(rendezvous_iterator) store.set_timeout(timeout) + if on_rendezvous is not None: + on_rendezvous(store) prefixed_store = PrefixStore(group_name, store) backend_value = Backend(backend) if backend else Backend("undefined") pg_options_param_name = ( diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index 8f586d04c..e50b0deb1 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -85,7 +85,10 @@ AFDDPMetadata, AFDForwardContextMetadata, ) -from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo +from afd_plugin.connectors.npu.async_cam import ( + AFDAsyncExtraInfo, + CAMAsyncAFDConnector, +) from afd_plugin.model_executor.models.npu.async_cam_layout import ( ASYNC_MOE_UBATCH_METADATA_KEY, AsyncMoeUbatchMetadata, @@ -1557,7 +1560,10 @@ def _build_capture_dp_metadata(self, num_tokens: int) -> DPMetadata | AFDDPMetad # Attention and FFN model loading overlap across roles. Connector # initialization is skipped when it already succeeded (the FFN side uses # the same guard), and the rendezvous completes before memory profiling or - # the first model forward. The connector itself is idempotent. + # the first model forward. Async CAM then waits for every FFN service to + # complete a controlled CAM communication warmup, finish capture, and start + # its receive loop before Attention can profile. + # The connector itself is idempotent. # Signature: matches upstream; no added parameters. def load_model(self) -> None: super().load_model() @@ -1569,6 +1575,18 @@ def load_model(self) -> None: # weights" barrier before memory profiling. if not self.connector.is_initialized: self.connector.init_afd_connector() + # ### PATCH START: CAM startup coordination + if isinstance(self.connector, CAMAsyncAFDConnector): + try: + layer_idx = self.connector.wait_for_ffn_modes() + if layer_idx is not None: + self.connector.run_attention_startup_warmup(layer_idx) + self.connector.wait_for_ffn_ready() + except Exception as exc: + if self.connector._startup_store is not None: + self.connector.publish_attention_startup_failure(str(exc)) + raise + # ### PATCH END: CAM startup coordination def _install_ascend_ubatch_wrapper(self) -> None: if isinstance(self.model, AscendUBatchWrapper): diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 7e6ab1d39..4c3ba734a 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -40,7 +40,10 @@ CAMAsyncAFDConnector, ) from afd_plugin.envs import async_cam_layered_gmm_enabled -from afd_plugin.model_executor.npu.async_cam_w4a8 import AsyncCAMW4A8Executor +from afd_plugin.model_executor.npu.async_cam_w4a8 import ( + CAM_LAYER_INDEX, + AsyncCAMW4A8Executor, +) from afd_plugin.v1.worker.attention_metadata import ( _resolve_world_ranks, ) @@ -60,6 +63,7 @@ from afd_plugin.connectors import AFDConnectorBase logger = init_logger(__name__) +FFN_GRAPH_REPLAY_LOG_INTERVAL = 1000 class AFDNPUFFNModelRunner(NPUModelRunner): @@ -102,13 +106,22 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: self._is_shutdown = False self._layered_gmm_requested = async_cam_layered_gmm_enabled() self._layered_executor: AsyncCAMW4A8Executor | None = None + self._async_cam_ffn_graph_enabled = ( + isinstance(self.connector, CAMAsyncAFDConnector) + and not self.model_config.enforce_eager + and self.compilation_config.cudagraph_mode.name == "FULL" + ) + self._async_cam_ffn_graph: torch.npu.NPUGraph | None = None + self._async_cam_ffn_output: torch.Tensor | None = None + self._async_cam_ffn_replays = 0 @staticmethod def parse_config(vllm_config: VllmConfig) -> AFDConfig: return parse_afd_config(vllm_config, expected_role="ffn") def initialize_afd_connector(self) -> None: - self._initialize_layered_executor() + if self._layered_executor is None: + self._initialize_layered_executor() self.connector.init_afd_connector() def _initialize_layered_executor(self) -> None: @@ -125,9 +138,11 @@ def _initialize_layered_executor(self) -> None: raise ValueError("layered W4A8 GMM requires Ascend 910C/A3") if self.connector.dynamic_quant != 1: raise ValueError("layered W4A8 GMM requires dynamicQuant=1") - if self.use_aclgraph or self.vllm_config.use_v2_model_runner: + if self.vllm_config.use_v2_model_runner: + raise ValueError("layered W4A8 GMM requires ModelRunnerV1") + if self.use_aclgraph and not self._async_cam_ffn_graph_enabled: raise ValueError( - "layered W4A8 GMM requires eager ModelRunnerV1" + "layered W4A8 GMM graph requires FFN FULL mode" ) if tuple(sorted(layer.layer_idx for layer in layers)) != tuple( _ffn_layer_indices(self) @@ -149,6 +164,8 @@ def _initialize_layered_executor(self) -> None: layers[0].swiglu_limit != 0.0, ) return + if self._async_cam_ffn_graph_enabled: + raise ValueError(f"CAM async FFN FULL requires layered W4A8 GMM: {reason}") logger.info( "AFD_ASYNC_CAM_LAYERED_GMM requested=%s actual=legacy reason=%s", self._layered_gmm_requested, @@ -203,9 +220,178 @@ def execute_connector_driven_step(self) -> None: "AFD connector", ) step_afd_npu_profiler(self.prof) + if self._async_cam_ffn_graph_enabled: + graph = self._async_cam_ffn_graph + if graph is None: + raise RuntimeError( + "CAM async FFN graph was not captured before dispatch" + ) + layer_indices = _ffn_layer_indices(self) + previous_replays = self._async_cam_ffn_replays + for _ in layer_indices: + graph.replay() + self._async_cam_ffn_replays += 1 + if previous_replays == 0: + logger.info( + "CAM async FFN graph replay rank=%d count=%d", + self.connector.world_rank, + self._async_cam_ffn_replays, + ) + elif ( + self._async_cam_ffn_replays // FFN_GRAPH_REPLAY_LOG_INTERVAL + > previous_replays // FFN_GRAPH_REPLAY_LOG_INTERVAL + ): + logger.debug( + "CAM async FFN graph replay rank=%d count=%d", + self.connector.world_rank, + self._async_cam_ffn_replays, + ) + return None self._ffn_forward_connector_driven() return None + def warmup_async_cam_ffn_graph(self) -> None: + """Prepare layered GMM locally before collective CAM initialization.""" + if ( + not self._async_cam_ffn_graph_enabled + or self._async_cam_ffn_graph is not None + ): + return + if self._layered_executor is None: + self._initialize_layered_executor() + executor = self._layered_executor + if executor is None: + raise RuntimeError("CAM async FFN FULL requires layered GMM") + connector = cast(CAMAsyncAFDConnector, self.connector) + # Meta dispatch uses the very same C++ allocator and float32 + # BATCH_SIZE_FACTOR parser as the real receive, without communication. + meta_outputs = torch.ops.afd_ascend.afd_async_dispatch_recv( + torch.empty((1,), dtype=connector.activation_dtype, device="meta"), + torch.empty((1,), dtype=torch.float16, device="meta"), + connector.comm_id, + connector.max_num_batched_tokens, + connector.hidden_size, + connector.topk, + connector.ffn_size, + connector.attn_size, + connector.expert_per_rank, + connector.world_rank, + connector.topology.world_size, + connector.tp_size, + connector.dynamic_quant, + "warmup", + ) + capacity = meta_outputs[0].shape[0] + hidden_states = torch.empty( + meta_outputs[0].shape, + dtype=meta_outputs[0].dtype, + device=self.device, + ) + hidden_states[0].zero_() + dynamic_scales = torch.empty( + meta_outputs[1].shape, + dtype=meta_outputs[1].dtype, + device=self.device, + ) + dynamic_scales[0] = 1.0 + group_list = torch.zeros( + meta_outputs[3].shape, + dtype=meta_outputs[3].dtype, + device=self.device, + ) + group_list[0] = 1 + batch_info = torch.zeros( + meta_outputs[2].shape, + dtype=meta_outputs[2].dtype, + device=self.device, + ) + batch_info[CAM_LAYER_INDEX] = executor.layer_ids[0] + executor(hidden_states, dynamic_scales, group_list, batch_info) + torch.npu.synchronize() + logger.info( + "CAM async FFN local layered warmup complete rank=%d capacity=%d", + connector.world_rank, + capacity, + ) + + def warmup_async_cam_ffn_communication(self) -> None: + """Consume one controlled CAM work item from each Attention DP group.""" + if not self._async_cam_ffn_graph_enabled: + return + connector = cast(CAMAsyncAFDConnector, self.connector) + attention_dp_size = connector.attn_size // connector.tp_size + for _ in range(attention_dp_size): + self._execute_layered_work_item() + torch.npu.synchronize() + logger.info( + "CAM async FFN communication warmup done rank=%d items=%d", + connector.world_rank, + attention_dp_size, + ) + + def capture_async_cam_ffn_graph(self) -> None: + """Capture one DR, layered GMM, CS transaction for every MoE layer.""" + if ( + not self._async_cam_ffn_graph_enabled + or self._async_cam_ffn_graph is not None + ): + return + if self._layered_executor is None or not self.connector.is_initialized: + raise RuntimeError("CAM async FFN graph requires weights and HCCL group") + graph = torch.npu.NPUGraph() + logger.info( + "CAM async FFN graph capture start rank=%d", self.connector.world_rank + ) + with torch.npu.graph(graph, pool=self.graph_pool): + logger.info( + "CAM async FFN graph context entered rank=%d", + self.connector.world_rank, + ) + output = self._execute_layered_work_item(capture_trace=True) + self._async_cam_ffn_output = output + self._async_cam_ffn_graph = graph + logger.info( + "CAM async FFN graph capture complete rank=%d graphs=1", + self.connector.world_rank, + ) + + def release_async_cam_ffn_graph(self) -> None: + if self._async_cam_ffn_graph is None: + return + self._async_cam_ffn_graph = None + self._async_cam_ffn_output = None + logger.info( + "CAM async FFN graph released rank=%d replays=%d", + self.connector.world_rank, + self._async_cam_ffn_replays, + ) + + def _execute_layered_work_item( + self, *, capture_trace: bool = False + ) -> torch.Tensor: + connector = cast(CAMAsyncAFDConnector, self.connector) + executor = self._layered_executor + if executor is None: + raise RuntimeError("CAM async FFN layered executor is unavailable") + payload = connector.recv_attn_output() + if capture_trace: + logger.info("CAM async FFN graph DR captured rank=%d", connector.world_rank) + states = cast(AFDAsyncTransferState, payload.context.states) + output = executor( + payload.hidden_states, + states.dynamic_scales, + states.group_list, + states.token_nums_rankid_layeridx, + ) + if capture_trace: + logger.info( + "CAM async FFN graph GMM captured rank=%d", connector.world_rank + ) + connector.send_ffn_output(output, payload.context) + if capture_trace: + logger.info("CAM async FFN graph CS captured rank=%d", connector.world_rank) + return output + def execute_model( self, scheduler_output: SchedulerOutput | None = None, @@ -372,15 +558,7 @@ def _ffn_forward_connector_driven( connector = cast(CAMAsyncAFDConnector, self.connector) if self._layered_executor is not None: for _ in _ffn_layer_indices(self): - payload = connector.recv_attn_output() - states = cast(AFDAsyncTransferState, payload.context.states) - rank_ffn_output = self._layered_executor( - payload.hidden_states, - states.dynamic_scales, - states.group_list, - states.token_nums_rankid_layeridx, - ) - connector.send_ffn_output(rank_ffn_output, payload.context) + rank_ffn_output = self._execute_layered_work_item() return rank_ffn_output for _ in _ffn_layer_indices(self): @@ -533,6 +711,7 @@ def shutdown(self) -> None: stop_afd_npu_profiler(self.prof) if self.connector.is_initialized: self.connector.close() + self.release_async_cam_ffn_graph() super().shutdown() self._is_shutdown = True diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index dded65435..f5aa5e335 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -22,6 +22,7 @@ fix_all2all_backend_for_afd, npu_afd_num_ubatches, ) +from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector from afd_plugin.model_executor.models.model_utils import get_afd_model_config from afd_plugin.v1.worker.npu.ffn_model_runner import AFDNPUFFNModelRunner from afd_plugin.validation import ( @@ -38,6 +39,8 @@ logger = logging.getLogger(__name__) FFN_SHUTDOWN_TIMEOUT_SECONDS = 5 +FFN_GRAPH_SHUTDOWN_TIMEOUT_SECONDS = 30 +FFN_STARTUP_THREAD_TIMEOUT_SECONDS = 30 class AFDNPUFFNWorker(NPUWorker): @@ -56,6 +59,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._ffn_thread: threading.Thread | None = None self._ffn_shutdown_event: threading.Event | None = None self._ffn_loop_error: BaseException | None = None + self._ffn_loop_started_event: threading.Event | None = None self._cpu_binding_attempted = False def init_device(self) -> None: @@ -91,7 +95,6 @@ def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]: def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None: self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks self.model_runner.initialize_kv_cache(kv_cache_config) - self.model_runner.initialize_afd_connector() self.start_ffn_server_loop() def compile_or_warm_up_model(self) -> CompilationTimes: @@ -113,11 +116,46 @@ def start_ffn_server_loop(self) -> None: self.raise_ffn_loop_error_if_any() connector = self.model_runner.connector - if not connector.is_initialized: - self.model_runner.initialize_afd_connector() + is_async_cam = isinstance(connector, CAMAsyncAFDConnector) + if is_async_cam: + logger.info( + "CAM FFN startup entered rank=%d graph=%s", + connector.world_rank, + self.model_runner._async_cam_ffn_graph_enabled, + ) + try: + if is_async_cam: + logger.info("CAM FFN local warmup start rank=%d", connector.world_rank) + self.model_runner.warmup_async_cam_ffn_graph() + logger.info("CAM FFN local warmup done rank=%d", connector.world_rank) + if not connector.is_initialized: + self.model_runner.initialize_afd_connector() + if is_async_cam: + logger.info("CAM FFN startup store start rank=%d", connector.world_rank) + connector.prepare_startup_coordination() + logger.info("CAM FFN startup store done rank=%d", connector.world_rank) + if self.model_runner._async_cam_ffn_graph_enabled: + executor = self.model_runner._layered_executor + if executor is None: + raise RuntimeError("CAM FFN graph has no layered executor") + connector.enable_ffn_graph_stop(executor.layer_ids[0]) + connector.publish_ffn_mode(executor.layer_ids[0]) + connector.wait_for_attention_warmup_prepared() + connector.publish_ffn_warmup_started() + self.model_runner.warmup_async_cam_ffn_communication() + connector.publish_ffn_warmup_done() + connector.wait_for_communication_warmup_done() + self.model_runner.capture_async_cam_ffn_graph() + else: + connector.publish_ffn_mode(None) + except Exception as exc: + if is_async_cam and connector._startup_store is not None: + connector.publish_ffn_startup(error=str(exc)) + raise self._bind_cpus_once() self._ffn_shutdown_event = threading.Event() + self._ffn_loop_started_event = threading.Event() self._ffn_loop_error = None def ffn_worker_loop() -> None: @@ -132,14 +170,35 @@ def ffn_worker_loop() -> None: ) return self._ffn_loop_error = exc + if self._ffn_loop_started_event is not None: + self._ffn_loop_started_event.set() logger.exception("AFD NPU FFN worker loop failed") + if is_async_cam and connector._startup_store is not None: + connector.publish_ffn_startup(error=str(exc)) self._ffn_thread = threading.Thread( target=ffn_worker_loop, name="afd-npu-ffn-worker-loop", daemon=True, ) - self._ffn_thread.start() + try: + self._ffn_thread.start() + if is_async_cam: + if not self._ffn_loop_started_event.wait( + timeout=FFN_STARTUP_THREAD_TIMEOUT_SECONDS + ): + self.raise_ffn_loop_error_if_any() + raise TimeoutError("AFD NPU FFN service thread did not start") + self.raise_ffn_loop_error_if_any() + if not self._ffn_thread.is_alive(): + raise RuntimeError( + "AFD NPU FFN service thread stopped during startup" + ) + connector.publish_ffn_startup() + except Exception as exc: + if is_async_cam and connector._startup_store is not None: + connector.publish_ffn_startup(error=str(exc)) + raise def _bind_cpus_once(self) -> None: if self._cpu_binding_attempted: @@ -167,6 +226,8 @@ def _run_ffn_server_loop(self) -> None: return torch.npu.set_device(self.device) + if isinstance(self.model_runner.connector, CAMAsyncAFDConnector): + self._ffn_loop_started_event.set() while not event.is_set(): if self.model_runner.connector.control_plane is None: self.model_runner.execute_connector_driven_step() @@ -200,20 +261,51 @@ def stop_ffn_server_loop(self) -> None: if event is not None: event.set() - # CAM recv blocks in the connector operator. Release the communicator - # first so the daemon can observe the shutdown event, then wait for it - # before the parent runner releases model tensors. - self.model_runner.connector.close() + # The graph DR can block inside the NPU kernel. Signal its graph-owned + # input from an independent stream, then wait for every queued replay + # before destroying the communicator or releasing graph anchors. + connector = self.model_runner.connector + graph_cam = ( + isinstance(connector, CAMAsyncAFDConnector) + and self.model_runner._async_cam_ffn_graph_enabled + ) thread = self._ffn_thread + if graph_cam and thread is not None: + logger.info("CAM FFN graph shutdown entered rank=%d", connector.world_rank) + if graph_cam and thread is not None and thread.is_alive(): + connector.signal_ffn_graph_stop() + logger.info("CAM FFN graph stop copied rank=%d", connector.world_rank) + thread.join(timeout=FFN_GRAPH_SHUTDOWN_TIMEOUT_SECONDS) + if thread.is_alive(): + raise RuntimeError("AFD NPU FFN graph loop did not stop after stop bit") + logger.info("CAM FFN graph thread joined rank=%d", connector.world_rank) + # The stopped graph must release its captured HCCL resources while the + # communicator is still valid. Keep operator anchors until both are + # released; the legacy path still closes before its bounded join. + if graph_cam: + self.model_runner.release_async_cam_ffn_graph() + logger.info( + "CAM FFN graph resources released rank=%d", connector.world_rank + ) + connector.close(release_buffers=False) + logger.info("CAM FFN communicator closed rank=%d", connector.world_rank) + else: + connector.close() if thread is not None: thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): raise RuntimeError( "AFD NPU FFN worker loop did not stop after connector close", ) + if graph_cam: + connector.release_closed_buffers() + logger.info("CAM FFN graph anchors released rank=%d", connector.world_rank) self._ffn_thread = None self._ffn_shutdown_event = None + self._ffn_loop_started_event = None self.raise_ffn_loop_error_if_any() + if graph_cam and thread is not None: + logger.info("CAM FFN graph shutdown complete rank=%d", connector.world_rank) def shutdown(self) -> None: # Stop the connector-driven daemon before NPUWorker releases the model diff --git a/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h b/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h index 02c57971c..6d9ec5f00 100644 --- a/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h +++ b/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h @@ -22,6 +22,8 @@ constexpr uint64_t CAM_MAX_RANK_SIZE = 384; // max NPUs supported by the Cam co constexpr uint32_t UB_ALIGN = 32; // UB aligned to 32 bytes constexpr uint32_t MAX_AIV_NUM = 48; // max AIV core count constexpr uint32_t INFO_NUM = 5; // number of valid batch-info fields; also start/end expert per chunk +constexpr uint32_t ATTN_RANK_ID_INDEX = 1; +constexpr int64_t STOP_ATTENTION_RANK = -1; template __aicore__ inline void SyncFunc() @@ -178,6 +180,12 @@ __aicore__ inline void AfdAsyncCombineSend::Process() SyncFunc(); SyncFunc(); + // FFN graph shutdown emits a local synthetic work item. Its negative + // Attention rank must never be used to address a remote CAM window. + if (tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) == STOP_ATTENTION_RANK) { + return; + } + CombineSend(); } diff --git a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp index fd54729d0..333d998a7 100644 --- a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp +++ b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp @@ -57,6 +57,7 @@ constexpr static int LIMIT_EXPERT_RANK_SIZE_MIN = 1; constexpr static int BATCH_INFO_VAL_NUM = 5; constexpr static int UB_ALIGN = 32; constexpr static int MAX_AIV_NUM = 48; +constexpr static int GRAPH_STOP_BUFFER_ELEMENTS = 16; constexpr static float EP_BALANCE_FACTOR = 1.2; @@ -134,14 +135,20 @@ static ge::graphStatus TilingFunc(gert::TilingContext* context) OPS_ERR_IF(dynamicQuant != 0 && dynamicQuant != 1, OPS_LOG_E(nodeName, "dynamicQuant is invalid, only support 0 or 1, but got dynamicQuant=%ld.", - dynamicQuant), return ge::GRAPH_FAILED); + dynamicQuant), return ge::GRAPH_FAILED); + + OPS_ERR_IF(xShape != nullptr && xShape->GetStorageShape().GetDimNum() == ONE_DIM && + xShape->GetStorageShape().GetDim(0) == GRAPH_STOP_BUFFER_ELEMENTS && dynamicQuant != 1, + OPS_LOG_E(nodeName, "graph stop buffer requires dynamicQuant=1"), return ge::GRAPH_FAILED); OPS_ERR_IF(xShape == nullptr, OPS_LOG_E(nodeName, "xShape is null."), return ge::GRAPH_FAILED); OPS_ERR_IF(xShape->GetStorageShape().GetDimNum() != ONE_DIM, OPS_LOG_E(nodeName, "xShape dim is invalid, must be %d, but got dimNum=%u.", ONE_DIM, xShape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); - OPS_ERR_IF(xShape->GetStorageShape().GetDim(0) != 1, - OPS_LOG_E(nodeName, "xShape dim0 is invalid, must be 1, but got dim0=%u.", + OPS_ERR_IF(xShape->GetStorageShape().GetDim(0) != 1 && + xShape->GetStorageShape().GetDim(0) != GRAPH_STOP_BUFFER_ELEMENTS, + OPS_LOG_E(nodeName, "xShape dim0 is invalid, must be 1 or %d, but got dim0=%u.", + GRAPH_STOP_BUFFER_ELEMENTS, xShape->GetStorageShape().GetDim(0)), return ge::GRAPH_FAILED); int64_t limitMaxSeqLenPerRank = maxSeqLen / tpSize; @@ -188,6 +195,8 @@ static ge::graphStatus TilingFunc(gert::TilingContext* context) tilingData->moeDistributeDispatchInfo.tpSize = tpSize; tilingData->moeDistributeDispatchInfo.maxTokenNum = maxTokenNum; tilingData->moeDistributeDispatchInfo.dynamicQuant = dynamicQuant; + tilingData->moeDistributeDispatchInfo.shutdownEnabled = + xShape->GetStorageShape().GetDim(0) == GRAPH_STOP_BUFFER_ELEMENTS; uint32_t blockDim = 1U; auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); diff --git a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h index b5bbee659..4cbf0ad75 100644 --- a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h +++ b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h @@ -30,6 +30,10 @@ constexpr uint32_t LAYER_INDEX_INDEX = 2; constexpr uint32_t START_EXPERT_INDEX = 3; constexpr uint32_t END_EXPERT_INDEX = 4; constexpr uint32_t MOE_PRE_TOKEN_NUM_INDEX = 3; +constexpr uint32_t STOP_BUFFER_ELEMENTS = 4; // four int64 values in the 32-byte x input +constexpr uint32_t STOP_ATTENTION_RANK = 0xFFFFFFFFU; +constexpr int64_t PENDING_ATTENTION_RANK = -2; +constexpr int64_t STOP_ATTENTION_RANK_SIGNED = -1; template __aicore__ inline void SyncFunc() { @@ -77,6 +81,9 @@ class AfdAsyncDispatchRecv { uint32_t startExpert, uint32_t endExpert); __aicore__ inline uint32_t WaitFromAttn(); + __aicore__ inline uint32_t WaitFromAttnWithStop(); + __aicore__ inline void PublishStopDecision(int64_t rankId); + __aicore__ inline void EmitStopWorkItem(); __aicore__ inline void ClearFlags(uint32_t attnRankId); @@ -124,6 +131,9 @@ class AfdAsyncDispatchRecv { uint32_t hiddenSizeQuant_{0}; bool isError_{false}; + bool shutdownEnabled_{false}; + uint32_t stopLayerId_{0}; + GlobalTensor stopBufferGMTensor_; }; template @@ -147,6 +157,10 @@ __aicore__ inline void AfdAsyncDispatchRecv::Init( totalWorkspaceSize_ = tilingData->moeDistributeDispatchInfo.totalWorkspaceSize; // total shared mem tpSize_ = tilingData->moeDistributeDispatchInfo.tpSize; maxTokenNum_ = tilingData->moeDistributeDispatchInfo.maxTokenNum; + shutdownEnabled_ = tilingData->moeDistributeDispatchInfo.shutdownEnabled != 0; + if (shutdownEnabled_) { + stopBufferGMTensor_.SetGlobalBuffer((__gm__ int64_t *)x); + } uint32_t maxSeqLenPerRank = maxSeqLen_ / tpSize_; uint64_t batchInfoWorkspaceSize = MathCeil(sizeof(int64_t) * BATCH_INFO_VAL_NUM, UB_ALIGN); @@ -242,7 +256,11 @@ __aicore__ inline void AfdAsyncDispatchRecv::Process() template __aicore__ inline void AfdAsyncDispatchRecv::DispatchRecv() { - uint32_t tpAttnRankId = WaitFromAttn(); + uint32_t tpAttnRankId = shutdownEnabled_ ? WaitFromAttnWithStop() : WaitFromAttn(); + if (tpAttnRankId == STOP_ATTENTION_RANK) { + EmitStopWorkItem(); + return; + } tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) = tpAttnRankId; // attnRankId tpBatchInfoTensor_(LAYER_INDEX_INDEX) = batchInfoTensor_(LAYER_INDEX_INDEX); // layerIndex @@ -496,6 +514,151 @@ __aicore__ inline void AfdAsyncDispatchRecv::DispatchRecvRo } } +template +__aicore__ inline void AfdAsyncDispatchRecv::PublishStopDecision(int64_t rankId) +{ + for (uint32_t i = 0; i < UB_ALIGN / sizeof(int64_t); ++i) { + tpBatchInfoTensor_(i) = 0; + } + tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) = rankId; + SyncFunc(); + DataCopy(batchInfoGMTensor_, tpBatchInfoTensor_, UB_ALIGN / sizeof(int64_t)); + SyncFunc(); +} + +template +__aicore__ inline uint32_t AfdAsyncDispatchRecv::WaitFromAttnWithStop() +{ + // Only AIV0 chooses between a normal DP group and a stop work item. The + // output's rank field is a per-launch decision slot; no other op can read + // it until this dispatch-recv kernel completes. + if (aivId_ == 0) { + PublishStopDecision(PENDING_ATTENTION_RANK); + } + SyncAll(); + + int64_t selectedRank = PENDING_ATTENTION_RANK; + GM_ADDR dstGM = GetPeerAddrByRankId(moeRankId_); + GlobalTensor peerBatchInfoGMTensor; + if (aivId_ == 0) { + GlobalTensor lastDpGMTensor; + lastDpGMTensor.SetGlobalBuffer((__gm__ int64_t *)(dstGM + (totalWorkspaceSize_ - UB_ALIGN))); + DataCopy(batchInfoTensor_, lastDpGMTensor, UB_ALIGN / sizeof(int64_t)); + SyncFunc(); + uint32_t nextAttnRankId = batchInfoTensor_(0); + SyncFunc(); + + while (selectedRank == PENDING_ATTENTION_RANK) { + // x is a 32-byte, graph-owned control input in this opt-in path. + // A separate H2D stream writes its first int64 on shutdown. + DataCopy(batchInfoTensor_, stopBufferGMTensor_, STOP_BUFFER_ELEMENTS); + SyncFunc(); + if (batchInfoTensor_(0) != 0) { + stopLayerId_ = static_cast(batchInfoTensor_(1)); + selectedRank = STOP_ATTENTION_RANK_SIGNED; + } + SyncFunc(); + if (selectedRank == STOP_ATTENTION_RANK_SIGNED) { + PublishStopDecision(selectedRank); + return STOP_ATTENTION_RANK; + } + + for (uint32_t i = 0; i < (attnRankNum_ / tpSize_); ++i) { + bool allReceived = true; + for (uint32_t j = tpSize_; j > 0; --j) { + uint32_t srcRankId = (nextAttnRankId + (tpSize_ * i + j - 1)) % attnRankNum_; + peerBatchInfoGMTensor.SetGlobalBuffer((__gm__ int64_t *)(dstGM + + workspaceSizePerAttn_ * srcRankId)); + DataCopyPad(batchInfoTensor_, peerBatchInfoGMTensor, + {1U, (uint32_t)(sizeof(int64_t) * BATCH_INFO_VAL_NUM), 0U, 0U, 0U}, + {false, 0U, 0U, 0U}); + SyncFunc(); + SyncFunc(); + if (batchInfoTensor_(TOKEN_NUM_INDEX) == 0) { + allReceived = false; + break; + } + } + if (allReceived) { + selectedRank = (nextAttnRankId + (tpSize_ * i)) % attnRankNum_; + // Preserve the legacy shared-window marker for normal work. + SyncFunc(); + DataCopyPad(peerBatchInfoGMTensor[BATCH_INFO_VAL_NUM - 1], + batchInfoTensor_, {1U, sizeof(int64_t), 0U, 0U, 0U}); + SyncFunc(); + PublishStopDecision(selectedRank); + break; + } + } + } + } else { + while (selectedRank == PENDING_ATTENTION_RANK) { + DataCopy(batchInfoTensor_, batchInfoGMTensor_, UB_ALIGN / sizeof(int64_t)); + SyncFunc(); + selectedRank = batchInfoTensor_(ATTN_RANK_ID_INDEX); + SyncFunc(); + } + if (selectedRank == STOP_ATTENTION_RANK_SIGNED) { + return STOP_ATTENTION_RANK; + } + } + + // All AIVs need the real layer/token fields from the selected DP group. + // The existing DispatchRecv SyncAll keeps AIV0 from clearing them early. + peerBatchInfoGMTensor.SetGlobalBuffer((__gm__ int64_t *)(dstGM + + workspaceSizePerAttn_ * static_cast(selectedRank))); + DataCopyPad(batchInfoTensor_, peerBatchInfoGMTensor, + {1U, (uint32_t)(sizeof(int64_t) * BATCH_INFO_VAL_NUM), 0U, 0U, 0U}, + {false, 0U, 0U, 0U}); + SyncFunc(); + SyncFunc(); + return static_cast(selectedRank); +} + +template +__aicore__ inline void AfdAsyncDispatchRecv::EmitStopWorkItem() +{ + if (aivId_ != 0) { + return; + } + // The layered W4A8 GMM sees one valid zero token on expert 0. Combine-send + // recognizes the negative Attention rank and skips all remote writes. + // AscendC Duplicate does not support the INT8 output type here. + for (uint32_t i = 0; i < hiddenSize_; ++i) { + xOutTensor_(i) = static_cast(0); + } + SyncFunc(); + DataCopy(xOutGMTensor_, xOutTensor_, hiddenSize_); + SyncFunc(); + + xOutFloatTensor_(0) = 1.0F; + SyncFunc(); + DataCopyPad(quantScaleGMTensor_, xOutFloatTensor_, {1U, sizeof(float), 0U, 0U, 0U}); + SyncFunc(); + + for (uint32_t i = 0; i < routeExpertNumPerMoe_; ++i) { + epRecvCountOutTensor_(i) = i == 0 ? 1 : 0; + } + SyncFunc(); + DataCopyPad(epRecvCountRoutedGMTensor_, epRecvCountOutTensor_, + {1U, (uint32_t)(sizeof(int64_t) * routeExpertNumPerMoe_), 0U, 0U, 0U}); + SyncFunc(); + + uint32_t infoCount = INFO_NUM + tpSize_ + routeExpertNumPerMoe_ * tpSize_; + for (uint32_t i = 0; i < infoCount; ++i) { + tpBatchInfoTensor_(i) = 0; + } + tpBatchInfoTensor_(TOKEN_NUM_INDEX) = 1; + tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) = STOP_ATTENTION_RANK_SIGNED; + tpBatchInfoTensor_(LAYER_INDEX_INDEX) = stopLayerId_; + tpBatchInfoTensor_(START_EXPERT_INDEX) = 0; + tpBatchInfoTensor_(END_EXPERT_INDEX) = 0; + SyncFunc(); + DataCopyPad(batchInfoGMTensor_, tpBatchInfoTensor_, + {1U, (uint32_t)(sizeof(int64_t) * infoCount), 0U, 0U, 0U}); + SyncFunc(); +} + template __aicore__ inline uint32_t AfdAsyncDispatchRecv::WaitFromAttn() { diff --git a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h index 9f3b4af6f..2afe60c5a 100644 --- a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h +++ b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h @@ -30,6 +30,7 @@ struct AfdAsyncDispatchRecvInfo { uint64_t totalWorkspaceSize; uint64_t maxTokenNum; int dynamicQuant; + int shutdownEnabled; }; struct AfdAsyncDispatchRecvTilingData { diff --git a/csrc/npu/pybind/torch_binding_cam_async.cpp b/csrc/npu/pybind/torch_binding_cam_async.cpp index f4dfb6058..8b5bc66ac 100644 --- a/csrc/npu/pybind/torch_binding_cam_async.cpp +++ b/csrc/npu/pybind/torch_binding_cam_async.cpp @@ -155,7 +155,14 @@ std::vector dispatch_recv( check_tp(max_seq_len, tp_size, attn_rank_num); TORCH_CHECK(moe_rank_id >= attn_rank_num && moe_rank_id < world_size, "moe_rank_id must be the global rank of a MoE participant"); - TORCH_CHECK(x.dim() == 1 && x.size(0) == 1, "x must be a one-element anchor"); + // x[16] is an opt-in graph stop ABI. Old one-element eager anchors retain + // their existing behavior; an old binary rejects the new shape visibly. + constexpr int64_t kGraphStopBufferElements = 16; + TORCH_CHECK(x.dim() == 1 && + (x.size(0) == 1 || x.size(0) == kGraphStopBufferElements), + "x must be a one-element anchor or a 16-element FFN graph stop buffer"); + TORCH_CHECK(x.size(0) == 1 || dynamic_quant == 1, + "FFN graph stop buffer requires dynamic_quant=1"); TORCH_CHECK(comm_args.scalar_type() == at::kHalf, "comm_args must have dtype float16"); TORCH_CHECK(dynamic_quant == 0 || dynamic_quant == 1, diff --git a/tests/npu/async_cam_communication_graph_probe.py b/tests/npu/async_cam_communication_graph_probe.py new file mode 100644 index 000000000..04a70bab5 --- /dev/null +++ b/tests/npu/async_cam_communication_graph_probe.py @@ -0,0 +1,306 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Two-rank CAM DR -> simple compute -> CS graph capture probe. + +Run with torchrun --standalone --nproc-per-node=2. Rank 0 is Attention and +rank 1 is FFN. ``cold`` captures immediately; ``warmup`` completes one eager +CAM transaction on both ranks before capture. Both modes send a matched task +for graph replay only after capture completes. ``shutdown`` first follows the +successful ``warmup`` path, then starts an unmatched FFN replay to test whether +destroying the CAM process group releases a blocked dispatch-recv kernel. +``stop`` instead uses the opt-in 32-byte FFN input to drain that replay before +destroying the communicator. +""" + +from __future__ import annotations + +import argparse +import os +import threading +import time +from datetime import timedelta + +import torch +import torch.distributed as dist +import torch_npu # noqa: F401 - register torch.npu + +from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available +from afd_plugin.distributed.afd_process_group import init_afd_process_group + +HIDDEN_SIZE = 512 +MAX_SEQUENCE_LENGTH = 262144 +MAX_SEQ_LEN = 1 +TOP_K = 1 +ATTN_RANKS = 1 +FFN_RANKS = 1 +EXPERTS_PER_RANK = 1 +DYNAMIC_QUANT = 1 +COMM_ID = 0 +RENDEZVOUS_TIMEOUT_SECONDS = 120 +SHUTDOWN_IDLE_SECONDS = 2 +SHUTDOWN_JOIN_SECONDS = 5 +GRAPH_STOP_BUFFER_INT64_VALUES = 4 +UNMATCHED_REPLAY_COUNT = 43 + + +def log(rank: int, phase: str) -> None: + print(f"rank={rank} {phase}", flush=True) + + +def attention_transaction( + *, rank: int, group_name: str, comm_args: torch.Tensor, anchor: torch.Tensor +) -> None: + ops = torch.ops.afd_ascend + hidden = torch.full( + (MAX_SEQ_LEN, HIDDEN_SIZE), 0.25, device="npu", dtype=torch.bfloat16 + ) + expert_ids = torch.zeros((MAX_SEQ_LEN, TOP_K), device="npu", dtype=torch.int32) + expert_weights = torch.ones((MAX_SEQ_LEN, TOP_K), device="npu", dtype=torch.float32) + log(rank, "DS start") + ops.afd_async_dispatch_send( + hidden, + expert_ids, + comm_args, + COMM_ID, + MAX_SEQ_LEN, + MAX_SEQ_LEN, + HIDDEN_SIZE, + TOP_K, + FFN_RANKS, + ATTN_RANKS, + EXPERTS_PER_RANK, + rank, + ATTN_RANKS + FFN_RANKS, + 0, + ATTN_RANKS, + DYNAMIC_QUANT, + group_name, + ) + log(rank, "DS returned; CR start") + output = ops.afd_async_combine_recv( + anchor, + expert_ids, + expert_weights, + comm_args, + COMM_ID, + MAX_SEQ_LEN, + HIDDEN_SIZE, + TOP_K, + FFN_RANKS, + ATTN_RANKS, + EXPERTS_PER_RANK, + rank, + ATTN_RANKS + FFN_RANKS, + group_name, + ) + torch.npu.synchronize() + torch.testing.assert_close(output, hidden, atol=0.01, rtol=0.01) + log(rank, "CR complete") + + +def ffn_transaction( + *, rank: int, group_name: str, comm_args: torch.Tensor, anchor: torch.Tensor +) -> tuple[torch.Tensor, ...]: + ops = torch.ops.afd_ascend + log(rank, "DR start") + expanded, scales, batch_info, counts = ops.afd_async_dispatch_recv( + anchor, + comm_args, + COMM_ID, + MAX_SEQ_LEN, + HIDDEN_SIZE, + TOP_K, + FFN_RANKS, + ATTN_RANKS, + EXPERTS_PER_RANK, + rank, + ATTN_RANKS + FFN_RANKS, + ATTN_RANKS, + DYNAMIC_QUANT, + group_name, + ) + log(rank, "DR returned; compute start") + output = (expanded.float() * scales.unsqueeze(1)).to(torch.bfloat16) + log(rank, "compute returned; CS start") + ops.afd_async_combine_send( + output, + comm_args, + batch_info, + COMM_ID, + MAX_SEQ_LEN, + HIDDEN_SIZE, + TOP_K, + FFN_RANKS, + ATTN_RANKS, + EXPERTS_PER_RANK, + rank, + ATTN_RANKS + FFN_RANKS, + ATTN_RANKS, + group_name, + ) + log(rank, "CS returned") + # Keep the graph's intermediate output and the receive metadata alive. + return output, expanded, scales, batch_info, counts + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--mode", choices=("cold", "warmup", "shutdown", "stop"), required=True + ) + parser.add_argument("--cam-port", type=int, required=True) + args = parser.parse_args() + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + if int(os.environ["WORLD_SIZE"]) != ATTN_RANKS + FFN_RANKS: + raise ValueError("Probe requires exactly two ranks") + os.environ["BATCH_SIZE_FACTOR"] = str(1 / MAX_SEQUENCE_LENGTH) + torch.npu.set_device(local_rank) + ensure_cam_async_ops_available() + dist.init_process_group( + "gloo", timeout=timedelta(seconds=RENDEZVOUS_TIMEOUT_SECONDS) + ) + # torchrun sets this for its own agent store. The CAM tcp:// rendezvous + # needs its own rank-0 server on --cam-port. + os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False" + cam_group = init_afd_process_group( + backend="hccl", + init_method=f"tcp://127.0.0.1:{args.cam_port}", + world_size=ATTN_RANKS + FFN_RANKS, + rank=rank, + group_name="afd_async_cam_communication_graph_probe", + timeout=timedelta(seconds=RENDEZVOUS_TIMEOUT_SECONDS), + ) + group_name = str( + cam_group._get_backend(torch.device("npu")).get_hccl_comm_name(rank) + ) + comm_args = torch.empty((1,), dtype=torch.float16, device="npu") + anchor = torch.empty((1,), dtype=torch.bfloat16, device="npu") + stop_host = None + stop_stream = None + if args.mode == "stop" and rank == ATTN_RANKS: + control = torch.zeros(GRAPH_STOP_BUFFER_INT64_VALUES, dtype=torch.int64) + anchor = control.view(torch.bfloat16).to("npu") + control[0] = 1 + stop_host = control.pin_memory().view(torch.bfloat16) + stop_stream = torch.npu.Stream() + log(rank, "CAM HCCL group initialized") + + with torch.inference_mode(): + if args.mode in ("warmup", "shutdown", "stop"): + log(rank, "eager warmup start") + if rank == 0: + attention_transaction( + rank=rank, group_name=group_name, comm_args=comm_args, anchor=anchor + ) + else: + ffn_transaction( + rank=rank, group_name=group_name, comm_args=comm_args, anchor=anchor + ) + torch.npu.synchronize() + log(rank, "eager warmup complete") + dist.barrier() + + graph = None + references = None + if rank == ATTN_RANKS: + graph = torch.npu.NPUGraph() + log(rank, "graph capture start") + with torch.npu.graph(graph): + log(rank, "graph context entered") + references = ffn_transaction( + rank=rank, group_name=group_name, comm_args=comm_args, anchor=anchor + ) + log(rank, "graph capture complete") + dist.barrier() + + log(rank, "matched replay transaction start") + if rank == 0: + attention_transaction( + rank=rank, group_name=group_name, comm_args=comm_args, anchor=anchor + ) + else: + assert graph is not None and references is not None + graph.replay() + torch.npu.synchronize() + log(rank, "graph replay complete") + dist.barrier() + + if args.mode in ("shutdown", "stop"): + results: dict[str, str] = {} + + def synchronize_unmatched_replay() -> None: + try: + torch.npu.set_device(local_rank) + torch.npu.synchronize() + results["sync"] = "complete" + except Exception as exc: + results["sync"] = f"failed: {exc}" + + def close_cam_group() -> None: + try: + torch.npu.set_device(local_rank) + dist.destroy_process_group(cam_group) + results["close"] = "complete" + except Exception as exc: + results["close"] = f"failed: {exc}" + + def copy_stop_bit() -> None: + assert stop_host is not None and stop_stream is not None + try: + torch.npu.set_device(local_rank) + with torch.npu.stream(stop_stream): + anchor.copy_(stop_host, non_blocking=True) + stop_stream.synchronize() + results["stop_copy"] = "complete" + except Exception as exc: + results["stop_copy"] = f"failed: {exc}" + + sync_thread = None + if rank == ATTN_RANKS: + assert graph is not None and references is not None + log(rank, "unmatched graph replay start") + for _ in range(UNMATCHED_REPLAY_COUNT): + graph.replay() + log( + rank, + f"unmatched graph replay submitted count={UNMATCHED_REPLAY_COUNT}; " + "anchors retained", + ) + sync_thread = threading.Thread( + target=synchronize_unmatched_replay, daemon=True + ) + sync_thread.start() + dist.barrier() + time.sleep(SHUTDOWN_IDLE_SECONDS) + if sync_thread is not None: + log(rank, f"unmatched synchronize blocked={sync_thread.is_alive()}") + if args.mode == "stop" and rank == ATTN_RANKS: + stop_thread = threading.Thread(target=copy_stop_bit, daemon=True) + stop_thread.start() + stop_thread.join(timeout=SHUTDOWN_JOIN_SECONDS) + log(rank, f"graph stop bit copy={results.get('stop_copy', 'blocked')}") + if results.get("stop_copy") != "complete": + log(rank, "shutdown probe failed; forcing this probe process exit") + os._exit(2) + close_thread = threading.Thread(target=close_cam_group, daemon=True) + close_thread.start() + close_thread.join(timeout=SHUTDOWN_JOIN_SECONDS) + if sync_thread is not None: + sync_thread.join(timeout=SHUTDOWN_JOIN_SECONDS) + log(rank, f"post-close sync={results.get('sync', 'blocked')}") + log(rank, f"CAM group close={results.get('close', 'blocked')}") + if results.get("close") != "complete" or ( + sync_thread is not None and results.get("sync") != "complete" + ): + log(rank, "shutdown probe failed; forcing this probe process exit") + os._exit(2) + dist.destroy_process_group() + return + + dist.destroy_process_group(cam_group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/npu/async_cam_ffn_graph_service_smoke.py b/tests/npu/async_cam_ffn_graph_service_smoke.py new file mode 100644 index 000000000..b0733548c --- /dev/null +++ b/tests/npu/async_cam_ffn_graph_service_smoke.py @@ -0,0 +1,125 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Reproducible prefill-only API smoke for the AsyncCam FFN graph deployment.""" + +from __future__ import annotations + +import argparse +import json +import time +from concurrent.futures import ThreadPoolExecutor, as_completed +from datetime import datetime, timezone +from pathlib import Path +from threading import Event +from urllib.error import HTTPError, URLError +from urllib.request import Request, urlopen + + +DEFAULT_TIMEOUT_SECONDS = 120 +CHUNK_SIZE = 8192 +CONCURRENT_REQUESTS = 4 + + +def send_completion( + endpoint: str, model: str, name: str, prompt: str, timeout: int +) -> dict: + payload = json.dumps( + {"model": model, "prompt": prompt, "max_tokens": 1, "temperature": 0} + ).encode() + request = Request( + endpoint, + data=payload, + headers={"Content-Type": "application/json"}, + ) + started_at = datetime.now(timezone.utc).isoformat() + start = time.monotonic() + result = {"name": name, "started_at": started_at} + try: + with urlopen(request, timeout=timeout) as response: + body = json.loads(response.read()) + usage = body.get("usage") or {} + result.update( + { + "ok": bool(body.get("choices")) + and usage.get("completion_tokens") == 1, + "request_id": body.get("id"), + "usage": usage, + "completion": [choice.get("text") for choice in body.get("choices", [])], + } + ) + except (HTTPError, URLError, TimeoutError, ValueError) as exc: + result.update({"ok": False, "error": str(exc)}) + result["elapsed_seconds"] = round(time.monotonic() - start, 3) + return result + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--endpoint", default="http://127.0.0.1:17100/v1/completions") + parser.add_argument("--model", default="dsv4") + parser.add_argument("--timeout", type=int, default=DEFAULT_TIMEOUT_SECONDS) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + + # Unique numbers avoid tokenizer compression of a repeated word and cross + # the deployed 8192-token prefill chunk boundary. Verify from API usage. + long_prompt = "List the next number after: " + " ".join( + f"{number:05d}" for number in range(8500) + ) + medium_prompt = "Continue this sequence: " + " ".join( + f"{number:05d}" for number in range(1000) + ) + cases = [ + ("short", "The capital of France is"), + ("medium", medium_prompt), + ("over_chunk", long_prompt), + ] + args.output.parent.mkdir(parents=True, exist_ok=True) + results = [] + with args.output.open("w") as output_file: + + def record(result: dict) -> None: + results.append(result) + line = json.dumps(result, ensure_ascii=False, sort_keys=True) + print(line, flush=True) + output_file.write(line + "\n") + output_file.flush() + + for name, prompt in cases: + result = send_completion(args.endpoint, args.model, name, prompt, args.timeout) + record(result) + if not result["ok"]: + return 1 + if name == "over_chunk" and ( + (result.get("usage") or {}).get("prompt_tokens", 0) <= CHUNK_SIZE + ): + return 1 + + start_concurrent = Event() + + def concurrent_case(index: int) -> dict: + start_concurrent.wait() + return send_completion( + args.endpoint, + args.model, + f"concurrent_{index}", + f"Continue: {index}. " + medium_prompt, + args.timeout, + ) + + with ThreadPoolExecutor(max_workers=CONCURRENT_REQUESTS) as executor: + futures = [ + executor.submit(concurrent_case, index) + for index in range(CONCURRENT_REQUESTS) + ] + start_concurrent.set() + for future in as_completed(futures): + record(future.result()) + + over_chunk = next(result for result in results if result["name"] == "over_chunk") + prompt_tokens = (over_chunk.get("usage") or {}).get("prompt_tokens", 0) + return 0 if all(result["ok"] for result in results) and prompt_tokens > CHUNK_SIZE else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/tests/npu/async_cam_graph_bootstrap_probe.py b/tests/npu/async_cam_graph_bootstrap_probe.py new file mode 100644 index 000000000..89eaaa81b --- /dev/null +++ b/tests/npu/async_cam_graph_bootstrap_probe.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Two-rank CAM process-group graph bootstrap probe without CAM work items. + +Launch with ``torchrun --nproc-per-node=2`` and pass a free ``--cam-port``. +The default torchrun rendezvous port must be different from ``--cam-port``. +This isolates HCCL startup work from dispatch-recv capture behavior. +""" + +from __future__ import annotations + +import argparse +import os +from datetime import timedelta + +import torch +import torch.distributed as dist +import torch_npu # noqa: F401 - register torch.npu before the probe uses it +from torch.distributed.distributed_c10d import Store + +from afd_plugin.distributed.afd_process_group import init_afd_process_group + +PROBE_TIMEOUT_SECONDS = 120 +NONCE_BYTES = 16 + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--cam-port", type=int, required=True) + parser.add_argument("--skip-nonce", action="store_true") + args = parser.parse_args() + + rank = int(os.environ["RANK"]) + local_rank = int(os.environ["LOCAL_RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + if world_size != 2: + raise ValueError("This probe requires exactly two ranks") + + torch.npu.set_device(local_rank) + dist.init_process_group( + backend="gloo", timeout=timedelta(seconds=PROBE_TIMEOUT_SECONDS) + ) + # This second rendezvous owns its TCPStore; torchrun's agent serves only + # the default group port, so rank zero must create the CAM store here. + os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False" + rendezvous_store: Store | None = None + + def retain_store(store: Store) -> None: + nonlocal rendezvous_store + rendezvous_store = store + + cam_group = init_afd_process_group( + backend="hccl", + init_method=f"tcp://127.0.0.1:{args.cam_port}", + world_size=world_size, + rank=rank, + group_name="afd_async_cam_graph_bootstrap_probe", + timeout=timedelta(seconds=PROBE_TIMEOUT_SECONDS), + on_rendezvous=retain_store, + ) + if rendezvous_store is None: + raise RuntimeError("CAM rendezvous store was not retained") + print(f"rank={rank} HCCL group initialized", flush=True) + + if not args.skip_nonce: + nonce = torch.zeros(NONCE_BYTES, dtype=torch.uint8, device=f"npu:{local_rank}") + if rank == 0: + nonce.copy_( + torch.tensor( + list(os.urandom(NONCE_BYTES)), + dtype=torch.uint8, + device=nonce.device, + ) + ) + dist.broadcast(nonce, src=0, group=cam_group) + print(f"rank={rank} nonce broadcast returned", flush=True) + torch.npu.synchronize() + print(f"rank={rank} nonce NPU synchronize returned", flush=True) + + graph = torch.npu.NPUGraph() + value = torch.ones((1,), device=f"npu:{local_rank}") + print(f"rank={rank} graph capture start", flush=True) + with torch.npu.graph(graph): + print(f"rank={rank} graph context entered", flush=True) + result = value + 1 + print(f"rank={rank} graph capture complete", flush=True) + graph.replay() + torch.npu.synchronize() + if result.item() != 2: + raise AssertionError(f"Unexpected graph result on rank {rank}") + print(f"rank={rank} graph replay complete", flush=True) + dist.destroy_process_group(cam_group) + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/npu/test_async_cam_layered_graph_probe.py b/tests/npu/test_async_cam_layered_graph_probe.py new file mode 100644 index 000000000..8546330a3 --- /dev/null +++ b/tests/npu/test_async_cam_layered_graph_probe.py @@ -0,0 +1,121 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Opt-in device probe for reusing one layered W4A8 computation graph. + +This isolates the device-controlled layer and row-count inputs. Full AsyncCam +DR/GMM/CS graph capture requires a running Attention/FFN service and is tested +separately on the complete deployment. +""" + +from __future__ import annotations + +import os + +import pytest + +pytestmark = pytest.mark.npu + + +def test_layered_w4a8_graph_reuses_device_metadata(): + if os.environ.get("AFD_RUN_ASCEND_OP_RUNTIME") != "1": + pytest.skip("requires opt-in 910C runtime") + + torch = pytest.importorskip("torch") + torch_npu = pytest.importorskip("torch_npu") + from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available + from afd_plugin.model_executor.npu.async_cam_w4a8 import ( + AsyncCAMW4A8Executor, + W4A8LayerWeights, + ) + + ensure_cam_async_ops_available() + torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) + torch.npu.config.allow_internal_format = True + torch.manual_seed(19) + + experts = 2 + hidden = 256 + intermediate = 256 + capacity = 32 + + def make_weight(k: int, n: int, factor: float, squeeze: bool): + values = torch.randint(-7, 8, (experts, k, n), dtype=torch.int8) + pairs = values.reshape(-1, 2) + packed = ( + torch.bitwise_or( + torch.bitwise_left_shift(pairs[:, 1], 4), + torch.bitwise_and(pairs[:, 0], 0x0F), + ) + .reshape(experts, k, n // 2) + .clone() + ) + weight = torch_npu.npu_format_cast(packed.npu(), 29).view(torch.int32) + scale = torch.full((experts, 1, n), factor, dtype=torch.float16).float() + compensation = 8 * (values.float() * scale).sum(dim=1) + encoded_scale = scale.view(torch.int32).to(torch.int64) + if squeeze: + encoded_scale = encoded_scale.squeeze(1) + return weight.contiguous(), encoded_scale.npu(), compensation.npu() + + layers = [] + for layer_idx, factor in ((2, 0.01), (5, 0.04)): + w13, s13, b13 = make_weight(hidden, 2 * intermediate, factor, True) + w2, s2, b2 = make_weight(intermediate, hidden, factor, False) + layers.append( + W4A8LayerWeights( + layer_idx, w13, w2, s13, s2, b13, b2, True, 0.0, 1.0 + ) + ) + executor = AsyncCAMW4A8Executor(layers) + + static_hidden = torch.zeros((capacity, hidden), dtype=torch.int8, device="npu") + static_scales = torch.full((capacity,), 0.005, device="npu") + static_counts = torch.tensor([1, 0], dtype=torch.int64, device="npu") + static_info = torch.tensor([capacity * 2, 0, 2, 1], dtype=torch.int64, device="npu") + + # Initialize the custom operators before graph capture without changing + # the addresses used by the graph itself. + executor(static_hidden, static_scales, static_counts, static_info) + torch.npu.synchronize() + graph = torch.npu.NPUGraph() + with torch.npu.graph(graph): + captured_output = executor( + static_hidden, static_scales, static_counts, static_info + ) + + shared_hidden = torch.randint(-100, 101, (capacity, hidden), dtype=torch.int8) + layer_outputs = {} + for layer_idx, counts in ( + (2, (1, 0)), + (5, (1, 0)), + (2, (0, 7)), + (5, (16, 16)), + (2, (0, 0)), + ): + valid_rows = sum(counts) + static_hidden.copy_(shared_hidden.npu()) + static_counts.copy_(torch.tensor(counts, dtype=torch.int64).npu()) + static_info.copy_( + torch.tensor([capacity * 2, 0, layer_idx, valid_rows], dtype=torch.int64).npu() + ) + graph.replay() + torch.npu.synchronize() + assert captured_output.shape == (capacity, hidden) + graphed = captured_output[:valid_rows].cpu().float() + eager = executor( + static_hidden, static_scales, static_counts, static_info + )[:valid_rows].cpu().float() + if valid_rows: + difference = (graphed - eager).abs() + print( + f"layer={layer_idx} counts={counts} " + f"max_abs={difference.max().item():.6f} " + f"mean_abs={difference.mean().item():.6f}" + ) + torch.testing.assert_close(graphed, eager, rtol=0.04, atol=0.05) + if counts == (1, 0): + layer_outputs[layer_idx] = eager + + assert not torch.allclose( + layer_outputs[2], layer_outputs[5], rtol=0.04, atol=0.05 + ), "synthetic layers must have distinguishable outputs for the same input" diff --git a/tests/unit/compat/npu/test_cam_async_ops.py b/tests/unit/compat/npu/test_cam_async_ops.py index 10a2bc561..8d1364ff8 100644 --- a/tests/unit/compat/npu/test_cam_async_ops.py +++ b/tests/unit/compat/npu/test_cam_async_ops.py @@ -38,9 +38,15 @@ def cam_runtime() -> ModuleType: return torch -def _dispatch_recv(torch: ModuleType, *, dynamic_quant: int = 0, moe_rank_id: int = 2): +def _dispatch_recv( + torch: ModuleType, + *, + dynamic_quant: int = 0, + moe_rank_id: int = 2, + anchor_size: int = 1, +): return torch.ops.afd_ascend.afd_async_dispatch_recv( - torch.empty(1, device="meta", dtype=torch.bfloat16), + torch.empty(anchor_size, device="meta", dtype=torch.bfloat16), torch.empty(1, device="meta", dtype=torch.float16), comm_id=0, max_seq_len=64, @@ -57,6 +63,25 @@ def _dispatch_recv(torch: ModuleType, *, dynamic_quant: int = 0, moe_rank_id: in ) +@pytest.mark.parametrize( + ("anchor_size", "dynamic_quant", "accepted"), + [(1, 0, True), (1, 1, True), (16, 1, True), (16, 0, False)], +) +def test_dispatch_recv_graph_stop_buffer_meta_contract( + cam_runtime, anchor_size, dynamic_quant, accepted +): + if accepted: + outputs = _dispatch_recv( + cam_runtime, anchor_size=anchor_size, dynamic_quant=dynamic_quant + ) + assert len(outputs) == 4 + else: + with pytest.raises(RuntimeError, match="requires dynamic_quant=1"): + _dispatch_recv( + cam_runtime, anchor_size=anchor_size, dynamic_quant=dynamic_quant + ) + + @pytest.mark.parametrize( "name", ( diff --git a/tests/unit/connectors/test_async_cam_connector.py b/tests/unit/connectors/test_async_cam_connector.py index fea5658d3..46693a726 100644 --- a/tests/unit/connectors/test_async_cam_connector.py +++ b/tests/unit/connectors/test_async_cam_connector.py @@ -353,6 +353,7 @@ def fake_init_afd_process_group(**kwargs): "timeout": calls[0]["timeout"], "init_method": "tcp://127.0.0.1:1239", "pg_options": pg_options, + "on_rendezvous": connector._retain_rendezvous_store, }, ] assert connector.cam_pg is not None @@ -851,3 +852,148 @@ def test_close_releases_every_stage_routing(monkeypatch): assert set(connector._pending_attention_payloads) == {0, 1} connector.close() assert connector._pending_attention_payloads == {} + + +def test_ffn_startup_failure_reaches_attention_waiter(): + class FakeStore: + def __init__(self): + self.values = {} + + def set(self, key, value): + self.values[key] = value.encode() + + def check(self, keys): + return all(key in self.values for key in keys) + + def get(self, key): + return self.values[key] + + store = FakeStore() + ffn = object.__new__(CAMAsyncAFDConnector) + ffn._startup_store = store + ffn.topology = SimpleNamespace(role="ffn", world_size=3) + ffn.world_rank = 2 + ffn.publish_ffn_startup(error="graph capture failed") + + attention = object.__new__(CAMAsyncAFDConnector) + attention._startup_store = store + attention.attn_size = 2 + attention.topology = SimpleNamespace(world_size=3) + with pytest.raises(RuntimeError, match="graph capture failed"): + attention.wait_for_ffn_ready() + + +def test_attention_wait_rechecks_previously_ready_ffn(monkeypatch): + class ChangingStore: + def __init__(self): + self.values = {"ffn/8": b"ready"} + + def check(self, keys): + return all(key in self.values for key in keys) + + def get(self, key): + return self.values[key] + + attention = object.__new__(CAMAsyncAFDConnector) + store = ChangingStore() + attention._startup_store = store + attention.attn_size = 8 + attention.topology = SimpleNamespace(world_size=10) + + def fail_first_rank(_: float) -> None: + store.values["ffn/8"] = b"failed:loop stopped" + store.values["ffn/9"] = b"ready" + + monkeypatch.setattr(async_cam_module.time, "sleep", fail_first_rank) + + with pytest.raises(RuntimeError, match="loop stopped"): + attention.wait_for_ffn_ready() + + +def test_ffn_ready_cannot_overwrite_failed_status(): + class FakeStore: + def __init__(self): + self.values = {"ffn/8": b"failed:startup crashed"} + + def compare_set(self, key, expected, desired): + current = self.values.get(key, b"") + if current == expected.encode(): + current = desired.encode() + self.values[key] = current + return current + + ffn = object.__new__(CAMAsyncAFDConnector) + ffn._startup_store = FakeStore() + ffn.topology = SimpleNamespace(role="ffn") + ffn.world_rank = 8 + with pytest.raises(RuntimeError, match="startup crashed"): + ffn.publish_ffn_startup() + + +def test_graph_warmup_wait_propagates_attention_failure(): + class FakeStore: + def __init__(self): + self.values = {"attn/failure/0": b"failed before dispatch"} + + def check(self, keys): + return all(key in self.values for key in keys) + + def get(self, key): + return self.values[key] + + ffn = object.__new__(CAMAsyncAFDConnector) + ffn._startup_store = FakeStore() + ffn.attn_size = 2 + with pytest.raises(RuntimeError, match="failed before dispatch"): + ffn.wait_for_attention_warmup_prepared() + + +def test_graph_capture_waits_for_all_attention_and_ffn_warmup_done(monkeypatch): + class FakeStore: + def __init__(self): + self.values = { + "attn/done/0": b"ready", + "ffn/warmup/done/1": b"ready", + } + + def check(self, keys): + return all(key in self.values for key in keys) + + def get(self, key): + return self.values[key] + + store = FakeStore() + ffn = object.__new__(CAMAsyncAFDConnector) + ffn._startup_store = store + ffn.attn_size = 1 + ffn.topology = SimpleNamespace(world_size=3) + polls = [] + + def finish_second_ffn(_: float) -> None: + polls.append(True) + store.values["ffn/warmup/done/2"] = b"ready" + + monkeypatch.setattr(async_cam_module.time, "sleep", finish_second_ffn) + ffn.wait_for_communication_warmup_done() + assert polls == [True] + + +def test_close_can_retain_graph_anchors_until_worker_stops(): + connector = object.__new__(CAMAsyncAFDConnector) + connector.cam_pg = None + connector._initialized = True + connector.comm_args = object() + connector._placeholder = object() + connector._startup_store = object() + connector._rendezvous_store = object() + connector._pending_attention_payloads = {} + comm_args = connector.comm_args + placeholder = connector._placeholder + + connector.close(release_buffers=False) + + assert connector.comm_args is comm_args + assert connector._placeholder is placeholder + connector.release_closed_buffers() + assert connector.comm_args is None + assert connector._placeholder is None diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index d63cfefe8..847b8035f 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -446,6 +446,10 @@ def _new_ffn_runner(): runner = object.__new__(AFDNPUFFNModelRunner) runner._layered_executor = None runner._layered_gmm_requested = False + runner._async_cam_ffn_graph_enabled = False + runner._async_cam_ffn_graph = None + runner._async_cam_ffn_output = None + runner._async_cam_ffn_replays = 0 runner.prof = None runner.device = SimpleNamespace(type="npu") runner._is_shutdown = False @@ -2440,7 +2444,33 @@ def test_npu_async_feature_validation_requires_async_config_and_eager(): config = _vllm_config(connector="CAMAsyncAFDConnector", async_dp=True) config.model_config.enforce_eager = False - with pytest.raises(RuntimeError, match="only eager"): + with pytest.raises(RuntimeError, match="requires FFN FULL"): + fail_if_unsupported_npu_afd_features(config) + + +def test_npu_async_feature_validation_allows_only_layered_dsv4_ffn_full( + monkeypatch, +): + _require_npu_runtime() + from afd_plugin.compat.npu import feature_validation + + config = _vllm_config( + role="ffn", + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + extra_config={"dynamicQuant": "1"}, + ) + config.model_config.enforce_eager = False + config.model_config.hf_config = SimpleNamespace(model_type="deepseek_v4") + config.use_v2_model_runner = False + monkeypatch.setattr( + feature_validation, "async_cam_layered_gmm_enabled", lambda: True + ) + fail_if_unsupported_npu_afd_features(config) + + config.compilation_config.cudagraph_mode.name = "FULL_DECODE_ONLY" + with pytest.raises(RuntimeError, match="requires FFN FULL"): fail_if_unsupported_npu_afd_features(config) @@ -2695,7 +2725,8 @@ def test_npu_ffn_runner_disables_mask_only_for_cam_mrv1( def fake_native_init(self, vllm_config, device): self.model_config = SimpleNamespace( - hf_config=SimpleNamespace(num_hidden_layers=1) + hf_config=SimpleNamespace(num_hidden_layers=1), + enforce_eager=True, ) module.ascend_context._reserved_mc2_mask = mask @@ -2849,3 +2880,223 @@ def test_npu_attention_runner_afd_ubatching_does_not_install_native_wrapper( runner.load_model() assert events == ["model_load", "connector_init"] + + +def test_npu_async_cam_ffn_graph_replays_one_item_per_moe_layer(monkeypatch): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import ffn_model_runner + + class FakeGraph: + def __init__(self): + self.replays = 0 + + def replay(self): + self.replays += 1 + + graph = FakeGraph() + runner = object.__new__(ffn_model_runner.AFDNPUFFNModelRunner) + runner.connector = SimpleNamespace(control_plane=None, world_rank=8) + runner.prof = None + runner._async_cam_ffn_graph_enabled = True + runner._async_cam_ffn_graph = graph + runner._async_cam_ffn_replays = 0 + monkeypatch.setattr(ffn_model_runner, "step_afd_npu_profiler", lambda _: None) + monkeypatch.setattr(ffn_model_runner, "_ffn_layer_indices", lambda _: [1, 4, 7]) + + runner.execute_connector_driven_step() + + assert graph.replays == 3 + assert runner._async_cam_ffn_replays == 3 + + +def test_npu_async_cam_ffn_graph_capture_is_idempotent_and_failure_is_atomic( + monkeypatch, +): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import ffn_model_runner + + @contextmanager + def fake_graph_context(graph, *, pool): + yield graph + + graph_creations = [] + + def make_graph(): + graph = SimpleNamespace() + graph_creations.append(graph) + return graph + + monkeypatch.setattr(ffn_model_runner.torch.npu, "NPUGraph", make_graph) + monkeypatch.setattr(ffn_model_runner.torch.npu, "graph", fake_graph_context) + runner = object.__new__(ffn_model_runner.AFDNPUFFNModelRunner) + runner.connector = SimpleNamespace(is_initialized=True, world_rank=8) + runner._layered_executor = object() + runner._async_cam_ffn_graph_enabled = True + runner._async_cam_ffn_graph = None + runner._async_cam_ffn_output = None + runner.graph_pool = None + runner._execute_layered_work_item = lambda **_: "captured-output" + + runner.capture_async_cam_ffn_graph() + runner.capture_async_cam_ffn_graph() + assert len(graph_creations) == 1 + assert runner._async_cam_ffn_output == "captured-output" + + runner._async_cam_ffn_graph = None + runner._async_cam_ffn_output = None + runner._execute_layered_work_item = lambda **_: (_ for _ in ()).throw( + RuntimeError("capture failed") + ) + with pytest.raises(RuntimeError, match="capture failed"): + runner.capture_async_cam_ffn_graph() + assert runner._async_cam_ffn_graph is None + assert runner._async_cam_ffn_output is None + + +def test_npu_async_cam_ffn_startup_orders_warmup_capture_and_ready(monkeypatch): + _require_npu_runtime() + from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector + from afd_plugin.v1.worker.npu import ffn_worker + + events = [] + connector = object.__new__(CAMAsyncAFDConnector) + connector._initialized = False + connector._startup_store = object() + connector.world_rank = 8 + connector.prepare_startup_coordination = lambda: events.append("coordination") + connector.enable_ffn_graph_stop = lambda layer_idx: events.append( + ("stop input", layer_idx) + ) + connector.publish_ffn_mode = lambda layer_idx: events.append(("mode", layer_idx)) + connector.wait_for_attention_warmup_prepared = lambda: events.append("prepared") + connector.publish_ffn_warmup_started = lambda: events.append("warmup start") + connector.publish_ffn_warmup_done = lambda: events.append("warmup done") + connector.wait_for_communication_warmup_done = lambda: events.append("all done") + connector.publish_ffn_startup = lambda **_: events.append("ready") + worker = _new_ffn_worker() + worker._ffn_thread = None + worker._ffn_shutdown_event = None + worker._bind_cpus_once = lambda: events.append("bind") + + def initialize_connector(): + events.append("connector") + connector._initialized = True + + worker.model_runner = SimpleNamespace( + connector=connector, + _async_cam_ffn_graph_enabled=True, + _layered_executor=SimpleNamespace(layer_ids=(0,)), + warmup_async_cam_ffn_graph=lambda: events.append("warmup"), + warmup_async_cam_ffn_communication=lambda: events.append("communication"), + initialize_afd_connector=initialize_connector, + capture_async_cam_ffn_graph=lambda: events.append("capture"), + ) + + class FakeThread: + def __init__(self, *, target, name, daemon): + events.append("thread created") + + def start(self): + events.append("thread started") + worker._ffn_loop_started_event.set() + + def is_alive(self): + return True + + monkeypatch.setattr(ffn_worker.threading, "Thread", FakeThread) + + worker.start_ffn_server_loop() + + assert events == [ + "warmup", + "connector", + "coordination", + ("stop input", 0), + ("mode", 0), + "prepared", + "warmup start", + "communication", + "warmup done", + "all done", + "capture", + "bind", + "thread created", + "thread started", + "ready", + ] + + +def test_npu_async_cam_graph_shutdown_drains_before_destroying_group(): + _require_npu_runtime() + from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector + + events = [] + connector = object.__new__(CAMAsyncAFDConnector) + connector.world_rank = 8 + connector.signal_ffn_graph_stop = lambda: events.append("stop bit") + connector.close = lambda **_: events.append("close group") + connector.release_closed_buffers = lambda: events.append("release anchors") + worker = _new_ffn_worker() + worker._ffn_shutdown_event = threading.Event() + worker._ffn_loop_started_event = threading.Event() + + class FakeThread: + active = True + + def is_alive(self): + return self.active + + def join(self, *, timeout): + events.append("join") + self.active = False + + worker._ffn_thread = FakeThread() + worker.model_runner = SimpleNamespace( + connector=connector, + _async_cam_ffn_graph_enabled=True, + release_async_cam_ffn_graph=lambda: events.append("release graph"), + ) + + worker.stop_ffn_server_loop() + + assert events == [ + "stop bit", + "join", + "release graph", + "close group", + "join", + "release anchors", + ] + + +def test_npu_async_cam_graph_shutdown_retains_anchors_if_join_times_out(): + _require_npu_runtime() + from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector + + events = [] + connector = object.__new__(CAMAsyncAFDConnector) + connector.world_rank = 8 + connector.signal_ffn_graph_stop = lambda: events.append("stop bit") + connector.close = lambda **_: events.append("close group") + connector.release_closed_buffers = lambda: events.append("release anchors") + worker = _new_ffn_worker() + worker._ffn_shutdown_event = threading.Event() + + class BlockedThread: + def is_alive(self): + return True + + def join(self, *, timeout): + events.append("join timeout") + + worker._ffn_thread = BlockedThread() + worker.model_runner = SimpleNamespace( + connector=connector, + _async_cam_ffn_graph_enabled=True, + release_async_cam_ffn_graph=lambda: events.append("release graph"), + ) + + with pytest.raises(RuntimeError, match="did not stop after stop bit"): + worker.stop_ffn_server_loop() + + assert events == ["stop bit", "join timeout"] From 272b158baaa26409449675b6b8a61a81d93ed9e4 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 15:05:59 +0800 Subject: [PATCH 2/8] refactor(npu): retain existing CAM kernels for FFN graphs Signed-off-by: jiangkuaixue123 --- afd_plugin/connectors/npu/async_cam.py | 32 ---- afd_plugin/v1/worker/npu/ffn_worker.py | 34 ++-- .../op_kernel/afd_async_combine_send.h | 8 - .../op_host/afd_async_dispatch_recv.cpp | 15 +- .../op_kernel/afd_async_dispatch_recv.h | 165 +----------------- .../afd_async_dispatch_recv_tiling.h | 1 - csrc/npu/pybind/torch_binding_cam_async.cpp | 9 +- .../async_cam_communication_graph_probe.py | 98 +---------- tests/unit/compat/npu/test_cam_async_ops.py | 29 +-- tests/unit/v1/worker/test_npu_runtime.py | 14 +- 10 files changed, 24 insertions(+), 381 deletions(-) diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index f5904284f..20c05e0ac 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -80,7 +80,6 @@ FFN_STARTUP_NONCE_BYTES = 16 STARTUP_WARMUP_TOKEN_COUNT = 1 STARTUP_WARMUP_HIDDEN_VALUE = 0.25 -FFN_GRAPH_STOP_BUFFER_INT64_VALUES = 4 _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( { @@ -273,8 +272,6 @@ def __init__( self.cam_pg: ProcessGroup | None = None self._rendezvous_store: Store | None = None self._startup_store: Store | None = None - self._ffn_graph_stop_stream: torch.npu.Stream | None = None - self._ffn_graph_stop_host: Tensor | None = None self.topology = build_async_topology( afd_config, role_rank, @@ -367,33 +364,6 @@ def _get_startup_store(self) -> Store: def prepare_startup_coordination(self) -> None: self._get_startup_store() - def enable_ffn_graph_stop(self, first_layer_idx: int) -> None: - """Keep a graph-owned 32-byte stop input and an independent copy stream.""" - if self.topology.role != "ffn" or not self.is_initialized: - raise RuntimeError("CAM graph stop input requires initialized FFN rank") - if self._ffn_graph_stop_stream is not None: - return - initial = torch.zeros(FFN_GRAPH_STOP_BUFFER_INT64_VALUES, dtype=torch.int64) - initial[1] = first_layer_idx - stop = initial.clone().pin_memory() - stop[0] = 1 - self._placeholder = initial.view(self.activation_dtype).to( - device=f"npu:{self.local_rank}" - ) - self._ffn_graph_stop_host = stop.view(self.activation_dtype) - self._ffn_graph_stop_stream = torch.npu.Stream(device=self.local_rank) - - def signal_ffn_graph_stop(self) -> None: - """Write the stop bit while the graph stream may be blocked in DR.""" - stream = self._ffn_graph_stop_stream - source = self._ffn_graph_stop_host - destination = self._placeholder - if stream is None or source is None or destination is None: - raise RuntimeError("CAM FFN graph stop input was not initialized") - with torch.npu.stream(stream): - destination.copy_(source, non_blocking=True) - stream.synchronize() - def _wait_for_startup_entries( self, keys: list[str], failure_keys: list[str] ) -> list[str]: @@ -599,8 +569,6 @@ def release_closed_buffers(self) -> None: raise RuntimeError("Cannot release active CAM operator buffers") self.comm_args = None self._placeholder = None - self._ffn_graph_stop_stream = None - self._ffn_graph_stop_host = None self._startup_store = None self._rendezvous_store = None self._pending_attention_payloads.clear() diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index f5aa5e335..3baa4259d 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -39,7 +39,6 @@ logger = logging.getLogger(__name__) FFN_SHUTDOWN_TIMEOUT_SECONDS = 5 -FFN_GRAPH_SHUTDOWN_TIMEOUT_SECONDS = 30 FFN_STARTUP_THREAD_TIMEOUT_SECONDS = 30 @@ -138,7 +137,6 @@ def start_ffn_server_loop(self) -> None: executor = self.model_runner._layered_executor if executor is None: raise RuntimeError("CAM FFN graph has no layered executor") - connector.enable_ffn_graph_stop(executor.layer_ids[0]) connector.publish_ffn_mode(executor.layer_ids[0]) connector.wait_for_attention_warmup_prepared() connector.publish_ffn_warmup_started() @@ -261,9 +259,9 @@ def stop_ffn_server_loop(self) -> None: if event is not None: event.set() - # The graph DR can block inside the NPU kernel. Signal its graph-owned - # input from an independent stream, then wait for every queued replay - # before destroying the communicator or releasing graph anchors. + # A graph replay may be blocked in DR until Attention sends another + # work item. Do not destroy its communicator or graph-owned tensors + # while the worker thread is still using them. connector = self.model_runner.connector graph_cam = ( isinstance(connector, CAMAsyncAFDConnector) @@ -271,27 +269,20 @@ def stop_ffn_server_loop(self) -> None: ) thread = self._ffn_thread if graph_cam and thread is not None: - logger.info("CAM FFN graph shutdown entered rank=%d", connector.world_rank) - if graph_cam and thread is not None and thread.is_alive(): - connector.signal_ffn_graph_stop() - logger.info("CAM FFN graph stop copied rank=%d", connector.world_rank) - thread.join(timeout=FFN_GRAPH_SHUTDOWN_TIMEOUT_SECONDS) + thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): - raise RuntimeError("AFD NPU FFN graph loop did not stop after stop bit") - logger.info("CAM FFN graph thread joined rank=%d", connector.world_rank) - # The stopped graph must release its captured HCCL resources while the - # communicator is still valid. Keep operator anchors until both are - # released; the legacy path still closes before its bounded join. + raise RuntimeError( + "AFD NPU FFN graph loop is still active; retain graph and " + "connector resources until the worker process exits" + ) + # Release a completed graph before its communicator. The legacy path + # still closes the connector before joining its receive thread. if graph_cam: self.model_runner.release_async_cam_ffn_graph() - logger.info( - "CAM FFN graph resources released rank=%d", connector.world_rank - ) connector.close(release_buffers=False) - logger.info("CAM FFN communicator closed rank=%d", connector.world_rank) else: connector.close() - if thread is not None: + if thread is not None and not graph_cam: thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): raise RuntimeError( @@ -299,13 +290,10 @@ def stop_ffn_server_loop(self) -> None: ) if graph_cam: connector.release_closed_buffers() - logger.info("CAM FFN graph anchors released rank=%d", connector.world_rank) self._ffn_thread = None self._ffn_shutdown_event = None self._ffn_loop_started_event = None self.raise_ffn_loop_error_if_any() - if graph_cam and thread is not None: - logger.info("CAM FFN graph shutdown complete rank=%d", connector.world_rank) def shutdown(self) -> None: # Stop the connector-driven daemon before NPUWorker releases the model diff --git a/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h b/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h index 6d9ec5f00..02c57971c 100644 --- a/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h +++ b/csrc/npu/ascend_kernels/afd_async_combine_send/op_kernel/afd_async_combine_send.h @@ -22,8 +22,6 @@ constexpr uint64_t CAM_MAX_RANK_SIZE = 384; // max NPUs supported by the Cam co constexpr uint32_t UB_ALIGN = 32; // UB aligned to 32 bytes constexpr uint32_t MAX_AIV_NUM = 48; // max AIV core count constexpr uint32_t INFO_NUM = 5; // number of valid batch-info fields; also start/end expert per chunk -constexpr uint32_t ATTN_RANK_ID_INDEX = 1; -constexpr int64_t STOP_ATTENTION_RANK = -1; template __aicore__ inline void SyncFunc() @@ -180,12 +178,6 @@ __aicore__ inline void AfdAsyncCombineSend::Process() SyncFunc(); SyncFunc(); - // FFN graph shutdown emits a local synthetic work item. Its negative - // Attention rank must never be used to address a remote CAM window. - if (tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) == STOP_ATTENTION_RANK) { - return; - } - CombineSend(); } diff --git a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp index 333d998a7..fd54729d0 100644 --- a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp +++ b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_host/afd_async_dispatch_recv.cpp @@ -57,7 +57,6 @@ constexpr static int LIMIT_EXPERT_RANK_SIZE_MIN = 1; constexpr static int BATCH_INFO_VAL_NUM = 5; constexpr static int UB_ALIGN = 32; constexpr static int MAX_AIV_NUM = 48; -constexpr static int GRAPH_STOP_BUFFER_ELEMENTS = 16; constexpr static float EP_BALANCE_FACTOR = 1.2; @@ -135,20 +134,14 @@ static ge::graphStatus TilingFunc(gert::TilingContext* context) OPS_ERR_IF(dynamicQuant != 0 && dynamicQuant != 1, OPS_LOG_E(nodeName, "dynamicQuant is invalid, only support 0 or 1, but got dynamicQuant=%ld.", - dynamicQuant), return ge::GRAPH_FAILED); - - OPS_ERR_IF(xShape != nullptr && xShape->GetStorageShape().GetDimNum() == ONE_DIM && - xShape->GetStorageShape().GetDim(0) == GRAPH_STOP_BUFFER_ELEMENTS && dynamicQuant != 1, - OPS_LOG_E(nodeName, "graph stop buffer requires dynamicQuant=1"), return ge::GRAPH_FAILED); + dynamicQuant), return ge::GRAPH_FAILED); OPS_ERR_IF(xShape == nullptr, OPS_LOG_E(nodeName, "xShape is null."), return ge::GRAPH_FAILED); OPS_ERR_IF(xShape->GetStorageShape().GetDimNum() != ONE_DIM, OPS_LOG_E(nodeName, "xShape dim is invalid, must be %d, but got dimNum=%u.", ONE_DIM, xShape->GetStorageShape().GetDimNum()), return ge::GRAPH_FAILED); - OPS_ERR_IF(xShape->GetStorageShape().GetDim(0) != 1 && - xShape->GetStorageShape().GetDim(0) != GRAPH_STOP_BUFFER_ELEMENTS, - OPS_LOG_E(nodeName, "xShape dim0 is invalid, must be 1 or %d, but got dim0=%u.", - GRAPH_STOP_BUFFER_ELEMENTS, + OPS_ERR_IF(xShape->GetStorageShape().GetDim(0) != 1, + OPS_LOG_E(nodeName, "xShape dim0 is invalid, must be 1, but got dim0=%u.", xShape->GetStorageShape().GetDim(0)), return ge::GRAPH_FAILED); int64_t limitMaxSeqLenPerRank = maxSeqLen / tpSize; @@ -195,8 +188,6 @@ static ge::graphStatus TilingFunc(gert::TilingContext* context) tilingData->moeDistributeDispatchInfo.tpSize = tpSize; tilingData->moeDistributeDispatchInfo.maxTokenNum = maxTokenNum; tilingData->moeDistributeDispatchInfo.dynamicQuant = dynamicQuant; - tilingData->moeDistributeDispatchInfo.shutdownEnabled = - xShape->GetStorageShape().GetDim(0) == GRAPH_STOP_BUFFER_ELEMENTS; uint32_t blockDim = 1U; auto ascendcPlatform = platform_ascendc::PlatformAscendC(context->GetPlatformInfo()); diff --git a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h index 4cbf0ad75..b5bbee659 100644 --- a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h +++ b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv.h @@ -30,10 +30,6 @@ constexpr uint32_t LAYER_INDEX_INDEX = 2; constexpr uint32_t START_EXPERT_INDEX = 3; constexpr uint32_t END_EXPERT_INDEX = 4; constexpr uint32_t MOE_PRE_TOKEN_NUM_INDEX = 3; -constexpr uint32_t STOP_BUFFER_ELEMENTS = 4; // four int64 values in the 32-byte x input -constexpr uint32_t STOP_ATTENTION_RANK = 0xFFFFFFFFU; -constexpr int64_t PENDING_ATTENTION_RANK = -2; -constexpr int64_t STOP_ATTENTION_RANK_SIGNED = -1; template __aicore__ inline void SyncFunc() { @@ -81,9 +77,6 @@ class AfdAsyncDispatchRecv { uint32_t startExpert, uint32_t endExpert); __aicore__ inline uint32_t WaitFromAttn(); - __aicore__ inline uint32_t WaitFromAttnWithStop(); - __aicore__ inline void PublishStopDecision(int64_t rankId); - __aicore__ inline void EmitStopWorkItem(); __aicore__ inline void ClearFlags(uint32_t attnRankId); @@ -131,9 +124,6 @@ class AfdAsyncDispatchRecv { uint32_t hiddenSizeQuant_{0}; bool isError_{false}; - bool shutdownEnabled_{false}; - uint32_t stopLayerId_{0}; - GlobalTensor stopBufferGMTensor_; }; template @@ -157,10 +147,6 @@ __aicore__ inline void AfdAsyncDispatchRecv::Init( totalWorkspaceSize_ = tilingData->moeDistributeDispatchInfo.totalWorkspaceSize; // total shared mem tpSize_ = tilingData->moeDistributeDispatchInfo.tpSize; maxTokenNum_ = tilingData->moeDistributeDispatchInfo.maxTokenNum; - shutdownEnabled_ = tilingData->moeDistributeDispatchInfo.shutdownEnabled != 0; - if (shutdownEnabled_) { - stopBufferGMTensor_.SetGlobalBuffer((__gm__ int64_t *)x); - } uint32_t maxSeqLenPerRank = maxSeqLen_ / tpSize_; uint64_t batchInfoWorkspaceSize = MathCeil(sizeof(int64_t) * BATCH_INFO_VAL_NUM, UB_ALIGN); @@ -256,11 +242,7 @@ __aicore__ inline void AfdAsyncDispatchRecv::Process() template __aicore__ inline void AfdAsyncDispatchRecv::DispatchRecv() { - uint32_t tpAttnRankId = shutdownEnabled_ ? WaitFromAttnWithStop() : WaitFromAttn(); - if (tpAttnRankId == STOP_ATTENTION_RANK) { - EmitStopWorkItem(); - return; - } + uint32_t tpAttnRankId = WaitFromAttn(); tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) = tpAttnRankId; // attnRankId tpBatchInfoTensor_(LAYER_INDEX_INDEX) = batchInfoTensor_(LAYER_INDEX_INDEX); // layerIndex @@ -514,151 +496,6 @@ __aicore__ inline void AfdAsyncDispatchRecv::DispatchRecvRo } } -template -__aicore__ inline void AfdAsyncDispatchRecv::PublishStopDecision(int64_t rankId) -{ - for (uint32_t i = 0; i < UB_ALIGN / sizeof(int64_t); ++i) { - tpBatchInfoTensor_(i) = 0; - } - tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) = rankId; - SyncFunc(); - DataCopy(batchInfoGMTensor_, tpBatchInfoTensor_, UB_ALIGN / sizeof(int64_t)); - SyncFunc(); -} - -template -__aicore__ inline uint32_t AfdAsyncDispatchRecv::WaitFromAttnWithStop() -{ - // Only AIV0 chooses between a normal DP group and a stop work item. The - // output's rank field is a per-launch decision slot; no other op can read - // it until this dispatch-recv kernel completes. - if (aivId_ == 0) { - PublishStopDecision(PENDING_ATTENTION_RANK); - } - SyncAll(); - - int64_t selectedRank = PENDING_ATTENTION_RANK; - GM_ADDR dstGM = GetPeerAddrByRankId(moeRankId_); - GlobalTensor peerBatchInfoGMTensor; - if (aivId_ == 0) { - GlobalTensor lastDpGMTensor; - lastDpGMTensor.SetGlobalBuffer((__gm__ int64_t *)(dstGM + (totalWorkspaceSize_ - UB_ALIGN))); - DataCopy(batchInfoTensor_, lastDpGMTensor, UB_ALIGN / sizeof(int64_t)); - SyncFunc(); - uint32_t nextAttnRankId = batchInfoTensor_(0); - SyncFunc(); - - while (selectedRank == PENDING_ATTENTION_RANK) { - // x is a 32-byte, graph-owned control input in this opt-in path. - // A separate H2D stream writes its first int64 on shutdown. - DataCopy(batchInfoTensor_, stopBufferGMTensor_, STOP_BUFFER_ELEMENTS); - SyncFunc(); - if (batchInfoTensor_(0) != 0) { - stopLayerId_ = static_cast(batchInfoTensor_(1)); - selectedRank = STOP_ATTENTION_RANK_SIGNED; - } - SyncFunc(); - if (selectedRank == STOP_ATTENTION_RANK_SIGNED) { - PublishStopDecision(selectedRank); - return STOP_ATTENTION_RANK; - } - - for (uint32_t i = 0; i < (attnRankNum_ / tpSize_); ++i) { - bool allReceived = true; - for (uint32_t j = tpSize_; j > 0; --j) { - uint32_t srcRankId = (nextAttnRankId + (tpSize_ * i + j - 1)) % attnRankNum_; - peerBatchInfoGMTensor.SetGlobalBuffer((__gm__ int64_t *)(dstGM + - workspaceSizePerAttn_ * srcRankId)); - DataCopyPad(batchInfoTensor_, peerBatchInfoGMTensor, - {1U, (uint32_t)(sizeof(int64_t) * BATCH_INFO_VAL_NUM), 0U, 0U, 0U}, - {false, 0U, 0U, 0U}); - SyncFunc(); - SyncFunc(); - if (batchInfoTensor_(TOKEN_NUM_INDEX) == 0) { - allReceived = false; - break; - } - } - if (allReceived) { - selectedRank = (nextAttnRankId + (tpSize_ * i)) % attnRankNum_; - // Preserve the legacy shared-window marker for normal work. - SyncFunc(); - DataCopyPad(peerBatchInfoGMTensor[BATCH_INFO_VAL_NUM - 1], - batchInfoTensor_, {1U, sizeof(int64_t), 0U, 0U, 0U}); - SyncFunc(); - PublishStopDecision(selectedRank); - break; - } - } - } - } else { - while (selectedRank == PENDING_ATTENTION_RANK) { - DataCopy(batchInfoTensor_, batchInfoGMTensor_, UB_ALIGN / sizeof(int64_t)); - SyncFunc(); - selectedRank = batchInfoTensor_(ATTN_RANK_ID_INDEX); - SyncFunc(); - } - if (selectedRank == STOP_ATTENTION_RANK_SIGNED) { - return STOP_ATTENTION_RANK; - } - } - - // All AIVs need the real layer/token fields from the selected DP group. - // The existing DispatchRecv SyncAll keeps AIV0 from clearing them early. - peerBatchInfoGMTensor.SetGlobalBuffer((__gm__ int64_t *)(dstGM + - workspaceSizePerAttn_ * static_cast(selectedRank))); - DataCopyPad(batchInfoTensor_, peerBatchInfoGMTensor, - {1U, (uint32_t)(sizeof(int64_t) * BATCH_INFO_VAL_NUM), 0U, 0U, 0U}, - {false, 0U, 0U, 0U}); - SyncFunc(); - SyncFunc(); - return static_cast(selectedRank); -} - -template -__aicore__ inline void AfdAsyncDispatchRecv::EmitStopWorkItem() -{ - if (aivId_ != 0) { - return; - } - // The layered W4A8 GMM sees one valid zero token on expert 0. Combine-send - // recognizes the negative Attention rank and skips all remote writes. - // AscendC Duplicate does not support the INT8 output type here. - for (uint32_t i = 0; i < hiddenSize_; ++i) { - xOutTensor_(i) = static_cast(0); - } - SyncFunc(); - DataCopy(xOutGMTensor_, xOutTensor_, hiddenSize_); - SyncFunc(); - - xOutFloatTensor_(0) = 1.0F; - SyncFunc(); - DataCopyPad(quantScaleGMTensor_, xOutFloatTensor_, {1U, sizeof(float), 0U, 0U, 0U}); - SyncFunc(); - - for (uint32_t i = 0; i < routeExpertNumPerMoe_; ++i) { - epRecvCountOutTensor_(i) = i == 0 ? 1 : 0; - } - SyncFunc(); - DataCopyPad(epRecvCountRoutedGMTensor_, epRecvCountOutTensor_, - {1U, (uint32_t)(sizeof(int64_t) * routeExpertNumPerMoe_), 0U, 0U, 0U}); - SyncFunc(); - - uint32_t infoCount = INFO_NUM + tpSize_ + routeExpertNumPerMoe_ * tpSize_; - for (uint32_t i = 0; i < infoCount; ++i) { - tpBatchInfoTensor_(i) = 0; - } - tpBatchInfoTensor_(TOKEN_NUM_INDEX) = 1; - tpBatchInfoTensor_(ATTN_RANK_ID_INDEX) = STOP_ATTENTION_RANK_SIGNED; - tpBatchInfoTensor_(LAYER_INDEX_INDEX) = stopLayerId_; - tpBatchInfoTensor_(START_EXPERT_INDEX) = 0; - tpBatchInfoTensor_(END_EXPERT_INDEX) = 0; - SyncFunc(); - DataCopyPad(batchInfoGMTensor_, tpBatchInfoTensor_, - {1U, (uint32_t)(sizeof(int64_t) * infoCount), 0U, 0U, 0U}); - SyncFunc(); -} - template __aicore__ inline uint32_t AfdAsyncDispatchRecv::WaitFromAttn() { diff --git a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h index 2afe60c5a..9f3b4af6f 100644 --- a/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h +++ b/csrc/npu/ascend_kernels/afd_async_dispatch_recv/op_kernel/afd_async_dispatch_recv_tiling.h @@ -30,7 +30,6 @@ struct AfdAsyncDispatchRecvInfo { uint64_t totalWorkspaceSize; uint64_t maxTokenNum; int dynamicQuant; - int shutdownEnabled; }; struct AfdAsyncDispatchRecvTilingData { diff --git a/csrc/npu/pybind/torch_binding_cam_async.cpp b/csrc/npu/pybind/torch_binding_cam_async.cpp index 8b5bc66ac..f4dfb6058 100644 --- a/csrc/npu/pybind/torch_binding_cam_async.cpp +++ b/csrc/npu/pybind/torch_binding_cam_async.cpp @@ -155,14 +155,7 @@ std::vector dispatch_recv( check_tp(max_seq_len, tp_size, attn_rank_num); TORCH_CHECK(moe_rank_id >= attn_rank_num && moe_rank_id < world_size, "moe_rank_id must be the global rank of a MoE participant"); - // x[16] is an opt-in graph stop ABI. Old one-element eager anchors retain - // their existing behavior; an old binary rejects the new shape visibly. - constexpr int64_t kGraphStopBufferElements = 16; - TORCH_CHECK(x.dim() == 1 && - (x.size(0) == 1 || x.size(0) == kGraphStopBufferElements), - "x must be a one-element anchor or a 16-element FFN graph stop buffer"); - TORCH_CHECK(x.size(0) == 1 || dynamic_quant == 1, - "FFN graph stop buffer requires dynamic_quant=1"); + TORCH_CHECK(x.dim() == 1 && x.size(0) == 1, "x must be a one-element anchor"); TORCH_CHECK(comm_args.scalar_type() == at::kHalf, "comm_args must have dtype float16"); TORCH_CHECK(dynamic_quant == 0 || dynamic_quant == 1, diff --git a/tests/npu/async_cam_communication_graph_probe.py b/tests/npu/async_cam_communication_graph_probe.py index 04a70bab5..89ab282de 100644 --- a/tests/npu/async_cam_communication_graph_probe.py +++ b/tests/npu/async_cam_communication_graph_probe.py @@ -5,19 +5,13 @@ Run with torchrun --standalone --nproc-per-node=2. Rank 0 is Attention and rank 1 is FFN. ``cold`` captures immediately; ``warmup`` completes one eager CAM transaction on both ranks before capture. Both modes send a matched task -for graph replay only after capture completes. ``shutdown`` first follows the -successful ``warmup`` path, then starts an unmatched FFN replay to test whether -destroying the CAM process group releases a blocked dispatch-recv kernel. -``stop`` instead uses the opt-in 32-byte FFN input to drain that replay before -destroying the communicator. +for graph replay only after capture completes. """ from __future__ import annotations import argparse import os -import threading -import time from datetime import timedelta import torch @@ -37,10 +31,6 @@ DYNAMIC_QUANT = 1 COMM_ID = 0 RENDEZVOUS_TIMEOUT_SECONDS = 120 -SHUTDOWN_IDLE_SECONDS = 2 -SHUTDOWN_JOIN_SECONDS = 5 -GRAPH_STOP_BUFFER_INT64_VALUES = 4 -UNMATCHED_REPLAY_COUNT = 43 def log(rank: int, phase: str) -> None: @@ -145,9 +135,7 @@ def ffn_transaction( def main() -> None: parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument( - "--mode", choices=("cold", "warmup", "shutdown", "stop"), required=True - ) + parser.add_argument("--mode", choices=("cold", "warmup"), required=True) parser.add_argument("--cam-port", type=int, required=True) args = parser.parse_args() rank = int(os.environ["RANK"]) @@ -176,18 +164,10 @@ def main() -> None: ) comm_args = torch.empty((1,), dtype=torch.float16, device="npu") anchor = torch.empty((1,), dtype=torch.bfloat16, device="npu") - stop_host = None - stop_stream = None - if args.mode == "stop" and rank == ATTN_RANKS: - control = torch.zeros(GRAPH_STOP_BUFFER_INT64_VALUES, dtype=torch.int64) - anchor = control.view(torch.bfloat16).to("npu") - control[0] = 1 - stop_host = control.pin_memory().view(torch.bfloat16) - stop_stream = torch.npu.Stream() log(rank, "CAM HCCL group initialized") with torch.inference_mode(): - if args.mode in ("warmup", "shutdown", "stop"): + if args.mode == "warmup": log(rank, "eager warmup start") if rank == 0: attention_transaction( @@ -226,78 +206,6 @@ def main() -> None: log(rank, "graph replay complete") dist.barrier() - if args.mode in ("shutdown", "stop"): - results: dict[str, str] = {} - - def synchronize_unmatched_replay() -> None: - try: - torch.npu.set_device(local_rank) - torch.npu.synchronize() - results["sync"] = "complete" - except Exception as exc: - results["sync"] = f"failed: {exc}" - - def close_cam_group() -> None: - try: - torch.npu.set_device(local_rank) - dist.destroy_process_group(cam_group) - results["close"] = "complete" - except Exception as exc: - results["close"] = f"failed: {exc}" - - def copy_stop_bit() -> None: - assert stop_host is not None and stop_stream is not None - try: - torch.npu.set_device(local_rank) - with torch.npu.stream(stop_stream): - anchor.copy_(stop_host, non_blocking=True) - stop_stream.synchronize() - results["stop_copy"] = "complete" - except Exception as exc: - results["stop_copy"] = f"failed: {exc}" - - sync_thread = None - if rank == ATTN_RANKS: - assert graph is not None and references is not None - log(rank, "unmatched graph replay start") - for _ in range(UNMATCHED_REPLAY_COUNT): - graph.replay() - log( - rank, - f"unmatched graph replay submitted count={UNMATCHED_REPLAY_COUNT}; " - "anchors retained", - ) - sync_thread = threading.Thread( - target=synchronize_unmatched_replay, daemon=True - ) - sync_thread.start() - dist.barrier() - time.sleep(SHUTDOWN_IDLE_SECONDS) - if sync_thread is not None: - log(rank, f"unmatched synchronize blocked={sync_thread.is_alive()}") - if args.mode == "stop" and rank == ATTN_RANKS: - stop_thread = threading.Thread(target=copy_stop_bit, daemon=True) - stop_thread.start() - stop_thread.join(timeout=SHUTDOWN_JOIN_SECONDS) - log(rank, f"graph stop bit copy={results.get('stop_copy', 'blocked')}") - if results.get("stop_copy") != "complete": - log(rank, "shutdown probe failed; forcing this probe process exit") - os._exit(2) - close_thread = threading.Thread(target=close_cam_group, daemon=True) - close_thread.start() - close_thread.join(timeout=SHUTDOWN_JOIN_SECONDS) - if sync_thread is not None: - sync_thread.join(timeout=SHUTDOWN_JOIN_SECONDS) - log(rank, f"post-close sync={results.get('sync', 'blocked')}") - log(rank, f"CAM group close={results.get('close', 'blocked')}") - if results.get("close") != "complete" or ( - sync_thread is not None and results.get("sync") != "complete" - ): - log(rank, "shutdown probe failed; forcing this probe process exit") - os._exit(2) - dist.destroy_process_group() - return - dist.destroy_process_group(cam_group) dist.destroy_process_group() diff --git a/tests/unit/compat/npu/test_cam_async_ops.py b/tests/unit/compat/npu/test_cam_async_ops.py index 8d1364ff8..10a2bc561 100644 --- a/tests/unit/compat/npu/test_cam_async_ops.py +++ b/tests/unit/compat/npu/test_cam_async_ops.py @@ -38,15 +38,9 @@ def cam_runtime() -> ModuleType: return torch -def _dispatch_recv( - torch: ModuleType, - *, - dynamic_quant: int = 0, - moe_rank_id: int = 2, - anchor_size: int = 1, -): +def _dispatch_recv(torch: ModuleType, *, dynamic_quant: int = 0, moe_rank_id: int = 2): return torch.ops.afd_ascend.afd_async_dispatch_recv( - torch.empty(anchor_size, device="meta", dtype=torch.bfloat16), + torch.empty(1, device="meta", dtype=torch.bfloat16), torch.empty(1, device="meta", dtype=torch.float16), comm_id=0, max_seq_len=64, @@ -63,25 +57,6 @@ def _dispatch_recv( ) -@pytest.mark.parametrize( - ("anchor_size", "dynamic_quant", "accepted"), - [(1, 0, True), (1, 1, True), (16, 1, True), (16, 0, False)], -) -def test_dispatch_recv_graph_stop_buffer_meta_contract( - cam_runtime, anchor_size, dynamic_quant, accepted -): - if accepted: - outputs = _dispatch_recv( - cam_runtime, anchor_size=anchor_size, dynamic_quant=dynamic_quant - ) - assert len(outputs) == 4 - else: - with pytest.raises(RuntimeError, match="requires dynamic_quant=1"): - _dispatch_recv( - cam_runtime, anchor_size=anchor_size, dynamic_quant=dynamic_quant - ) - - @pytest.mark.parametrize( "name", ( diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 847b8035f..a8ae3559a 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -2964,9 +2964,6 @@ def test_npu_async_cam_ffn_startup_orders_warmup_capture_and_ready(monkeypatch): connector._startup_store = object() connector.world_rank = 8 connector.prepare_startup_coordination = lambda: events.append("coordination") - connector.enable_ffn_graph_stop = lambda layer_idx: events.append( - ("stop input", layer_idx) - ) connector.publish_ffn_mode = lambda layer_idx: events.append(("mode", layer_idx)) connector.wait_for_attention_warmup_prepared = lambda: events.append("prepared") connector.publish_ffn_warmup_started = lambda: events.append("warmup start") @@ -3011,7 +3008,6 @@ def is_alive(self): "warmup", "connector", "coordination", - ("stop input", 0), ("mode", 0), "prepared", "warmup start", @@ -3026,14 +3022,13 @@ def is_alive(self): ] -def test_npu_async_cam_graph_shutdown_drains_before_destroying_group(): +def test_npu_async_cam_graph_shutdown_releases_after_thread_stops(): _require_npu_runtime() from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector events = [] connector = object.__new__(CAMAsyncAFDConnector) connector.world_rank = 8 - connector.signal_ffn_graph_stop = lambda: events.append("stop bit") connector.close = lambda **_: events.append("close group") connector.release_closed_buffers = lambda: events.append("release anchors") worker = _new_ffn_worker() @@ -3060,11 +3055,9 @@ def join(self, *, timeout): worker.stop_ffn_server_loop() assert events == [ - "stop bit", "join", "release graph", "close group", - "join", "release anchors", ] @@ -3076,7 +3069,6 @@ def test_npu_async_cam_graph_shutdown_retains_anchors_if_join_times_out(): events = [] connector = object.__new__(CAMAsyncAFDConnector) connector.world_rank = 8 - connector.signal_ffn_graph_stop = lambda: events.append("stop bit") connector.close = lambda **_: events.append("close group") connector.release_closed_buffers = lambda: events.append("release anchors") worker = _new_ffn_worker() @@ -3096,7 +3088,7 @@ def join(self, *, timeout): release_async_cam_ffn_graph=lambda: events.append("release graph"), ) - with pytest.raises(RuntimeError, match="did not stop after stop bit"): + with pytest.raises(RuntimeError, match="still active"): worker.stop_ffn_server_loop() - assert events == ["stop bit", "join timeout"] + assert events == ["join timeout"] From 57d145cbe24f8e7e773c310b71355bc41be78913 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 15:18:50 +0800 Subject: [PATCH 3/8] fix(npu): clarify FFN graph connector types and format probes Signed-off-by: jiangkuaixue123 --- afd_plugin/v1/worker/npu/ffn_model_runner.py | 19 ++-- afd_plugin/v1/worker/npu/ffn_worker.py | 93 +++++++++++-------- .../npu/async_cam_ffn_graph_service_smoke.py | 27 ++++-- .../npu/test_async_cam_layered_graph_probe.py | 18 ++-- tests/unit/v1/worker/test_npu_runtime.py | 2 +- 5 files changed, 93 insertions(+), 66 deletions(-) diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 4c3ba734a..94b67c53a 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -221,6 +221,7 @@ def execute_connector_driven_step(self) -> None: ) step_afd_npu_profiler(self.prof) if self._async_cam_ffn_graph_enabled: + connector = cast(CAMAsyncAFDConnector, self.connector) graph = self._async_cam_ffn_graph if graph is None: raise RuntimeError( @@ -234,7 +235,7 @@ def execute_connector_driven_step(self) -> None: if previous_replays == 0: logger.info( "CAM async FFN graph replay rank=%d count=%d", - self.connector.world_rank, + connector.world_rank, self._async_cam_ffn_replays, ) elif ( @@ -243,7 +244,7 @@ def execute_connector_driven_step(self) -> None: ): logger.debug( "CAM async FFN graph replay rank=%d count=%d", - self.connector.world_rank, + connector.world_rank, self._async_cam_ffn_replays, ) return None @@ -336,33 +337,33 @@ def capture_async_cam_ffn_graph(self) -> None: or self._async_cam_ffn_graph is not None ): return - if self._layered_executor is None or not self.connector.is_initialized: + connector = cast(CAMAsyncAFDConnector, self.connector) + if self._layered_executor is None or not connector.is_initialized: raise RuntimeError("CAM async FFN graph requires weights and HCCL group") graph = torch.npu.NPUGraph() - logger.info( - "CAM async FFN graph capture start rank=%d", self.connector.world_rank - ) + logger.info("CAM async FFN graph capture start rank=%d", connector.world_rank) with torch.npu.graph(graph, pool=self.graph_pool): logger.info( "CAM async FFN graph context entered rank=%d", - self.connector.world_rank, + connector.world_rank, ) output = self._execute_layered_work_item(capture_trace=True) self._async_cam_ffn_output = output self._async_cam_ffn_graph = graph logger.info( "CAM async FFN graph capture complete rank=%d graphs=1", - self.connector.world_rank, + connector.world_rank, ) def release_async_cam_ffn_graph(self) -> None: if self._async_cam_ffn_graph is None: return + connector = cast(CAMAsyncAFDConnector, self.connector) self._async_cam_ffn_graph = None self._async_cam_ffn_output = None logger.info( "CAM async FFN graph released rank=%d replays=%d", - self.connector.world_rank, + connector.world_rank, self._async_cam_ffn_replays, ) diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index 3baa4259d..1f9d6b794 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -115,45 +115,56 @@ def start_ffn_server_loop(self) -> None: self.raise_ffn_loop_error_if_any() connector = self.model_runner.connector - is_async_cam = isinstance(connector, CAMAsyncAFDConnector) - if is_async_cam: + cam_connector = ( + connector if isinstance(connector, CAMAsyncAFDConnector) else None + ) + if cam_connector is not None: logger.info( "CAM FFN startup entered rank=%d graph=%s", - connector.world_rank, + cam_connector.world_rank, self.model_runner._async_cam_ffn_graph_enabled, ) try: - if is_async_cam: - logger.info("CAM FFN local warmup start rank=%d", connector.world_rank) + if cam_connector is not None: + logger.info( + "CAM FFN local warmup start rank=%d", cam_connector.world_rank + ) self.model_runner.warmup_async_cam_ffn_graph() - logger.info("CAM FFN local warmup done rank=%d", connector.world_rank) + logger.info( + "CAM FFN local warmup done rank=%d", cam_connector.world_rank + ) if not connector.is_initialized: self.model_runner.initialize_afd_connector() - if is_async_cam: - logger.info("CAM FFN startup store start rank=%d", connector.world_rank) - connector.prepare_startup_coordination() - logger.info("CAM FFN startup store done rank=%d", connector.world_rank) + if cam_connector is not None: + logger.info( + "CAM FFN startup store start rank=%d", cam_connector.world_rank + ) + cam_connector.prepare_startup_coordination() + logger.info( + "CAM FFN startup store done rank=%d", cam_connector.world_rank + ) if self.model_runner._async_cam_ffn_graph_enabled: executor = self.model_runner._layered_executor if executor is None: raise RuntimeError("CAM FFN graph has no layered executor") - connector.publish_ffn_mode(executor.layer_ids[0]) - connector.wait_for_attention_warmup_prepared() - connector.publish_ffn_warmup_started() + cam_connector.publish_ffn_mode(executor.layer_ids[0]) + cam_connector.wait_for_attention_warmup_prepared() + cam_connector.publish_ffn_warmup_started() self.model_runner.warmup_async_cam_ffn_communication() - connector.publish_ffn_warmup_done() - connector.wait_for_communication_warmup_done() + cam_connector.publish_ffn_warmup_done() + cam_connector.wait_for_communication_warmup_done() self.model_runner.capture_async_cam_ffn_graph() else: - connector.publish_ffn_mode(None) + cam_connector.publish_ffn_mode(None) except Exception as exc: - if is_async_cam and connector._startup_store is not None: - connector.publish_ffn_startup(error=str(exc)) + if cam_connector is not None and cam_connector._startup_store is not None: + cam_connector.publish_ffn_startup(error=str(exc)) raise self._bind_cpus_once() self._ffn_shutdown_event = threading.Event() - self._ffn_loop_started_event = threading.Event() + started_event = threading.Event() + self._ffn_loop_started_event = started_event self._ffn_loop_error = None def ffn_worker_loop() -> None: @@ -168,11 +179,13 @@ def ffn_worker_loop() -> None: ) return self._ffn_loop_error = exc - if self._ffn_loop_started_event is not None: - self._ffn_loop_started_event.set() + started_event.set() logger.exception("AFD NPU FFN worker loop failed") - if is_async_cam and connector._startup_store is not None: - connector.publish_ffn_startup(error=str(exc)) + if ( + cam_connector is not None + and cam_connector._startup_store is not None + ): + cam_connector.publish_ffn_startup(error=str(exc)) self._ffn_thread = threading.Thread( target=ffn_worker_loop, @@ -181,10 +194,8 @@ def ffn_worker_loop() -> None: ) try: self._ffn_thread.start() - if is_async_cam: - if not self._ffn_loop_started_event.wait( - timeout=FFN_STARTUP_THREAD_TIMEOUT_SECONDS - ): + if cam_connector is not None: + if not started_event.wait(timeout=FFN_STARTUP_THREAD_TIMEOUT_SECONDS): self.raise_ffn_loop_error_if_any() raise TimeoutError("AFD NPU FFN service thread did not start") self.raise_ffn_loop_error_if_any() @@ -192,10 +203,10 @@ def ffn_worker_loop() -> None: raise RuntimeError( "AFD NPU FFN service thread stopped during startup" ) - connector.publish_ffn_startup() + cam_connector.publish_ffn_startup() except Exception as exc: - if is_async_cam and connector._startup_store is not None: - connector.publish_ffn_startup(error=str(exc)) + if cam_connector is not None and cam_connector._startup_store is not None: + cam_connector.publish_ffn_startup(error=str(exc)) raise def _bind_cpus_once(self) -> None: @@ -225,7 +236,9 @@ def _run_ffn_server_loop(self) -> None: torch.npu.set_device(self.device) if isinstance(self.model_runner.connector, CAMAsyncAFDConnector): - self._ffn_loop_started_event.set() + started_event = self._ffn_loop_started_event + if started_event is not None: + started_event.set() while not event.is_set(): if self.model_runner.connector.control_plane is None: self.model_runner.execute_connector_driven_step() @@ -263,12 +276,14 @@ def stop_ffn_server_loop(self) -> None: # work item. Do not destroy its communicator or graph-owned tensors # while the worker thread is still using them. connector = self.model_runner.connector - graph_cam = ( - isinstance(connector, CAMAsyncAFDConnector) + graph_connector = ( + connector + if isinstance(connector, CAMAsyncAFDConnector) and self.model_runner._async_cam_ffn_graph_enabled + else None ) thread = self._ffn_thread - if graph_cam and thread is not None: + if graph_connector is not None and thread is not None: thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): raise RuntimeError( @@ -277,19 +292,19 @@ def stop_ffn_server_loop(self) -> None: ) # Release a completed graph before its communicator. The legacy path # still closes the connector before joining its receive thread. - if graph_cam: + if graph_connector is not None: self.model_runner.release_async_cam_ffn_graph() - connector.close(release_buffers=False) + graph_connector.close(release_buffers=False) else: connector.close() - if thread is not None and not graph_cam: + if thread is not None and graph_connector is None: thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): raise RuntimeError( "AFD NPU FFN worker loop did not stop after connector close", ) - if graph_cam: - connector.release_closed_buffers() + if graph_connector is not None: + graph_connector.release_closed_buffers() self._ffn_thread = None self._ffn_shutdown_event = None self._ffn_loop_started_event = None diff --git a/tests/npu/async_cam_ffn_graph_service_smoke.py b/tests/npu/async_cam_ffn_graph_service_smoke.py index b0733548c..bcf0ea3be 100644 --- a/tests/npu/async_cam_ffn_graph_service_smoke.py +++ b/tests/npu/async_cam_ffn_graph_service_smoke.py @@ -11,10 +11,10 @@ from datetime import datetime, timezone from pathlib import Path from threading import Event +from typing import Any from urllib.error import HTTPError, URLError from urllib.request import Request, urlopen - DEFAULT_TIMEOUT_SECONDS = 120 CHUNK_SIZE = 8192 CONCURRENT_REQUESTS = 4 @@ -22,7 +22,7 @@ def send_completion( endpoint: str, model: str, name: str, prompt: str, timeout: int -) -> dict: +) -> dict[str, Any]: payload = json.dumps( {"model": model, "prompt": prompt, "max_tokens": 1, "temperature": 0} ).encode() @@ -33,18 +33,19 @@ def send_completion( ) started_at = datetime.now(timezone.utc).isoformat() start = time.monotonic() - result = {"name": name, "started_at": started_at} + result: dict[str, Any] = {"name": name, "started_at": started_at} try: with urlopen(request, timeout=timeout) as response: body = json.loads(response.read()) usage = body.get("usage") or {} result.update( { - "ok": bool(body.get("choices")) - and usage.get("completion_tokens") == 1, + "ok": bool(body.get("choices")) and usage.get("completion_tokens") == 1, "request_id": body.get("id"), "usage": usage, - "completion": [choice.get("text") for choice in body.get("choices", [])], + "completion": [ + choice.get("text") for choice in body.get("choices", []) + ], } ) except (HTTPError, URLError, TimeoutError, ValueError) as exc: @@ -78,7 +79,7 @@ def main() -> int: results = [] with args.output.open("w") as output_file: - def record(result: dict) -> None: + def record(result: dict[str, Any]) -> None: results.append(result) line = json.dumps(result, ensure_ascii=False, sort_keys=True) print(line, flush=True) @@ -86,7 +87,9 @@ def record(result: dict) -> None: output_file.flush() for name, prompt in cases: - result = send_completion(args.endpoint, args.model, name, prompt, args.timeout) + result = send_completion( + args.endpoint, args.model, name, prompt, args.timeout + ) record(result) if not result["ok"]: return 1 @@ -97,7 +100,7 @@ def record(result: dict) -> None: start_concurrent = Event() - def concurrent_case(index: int) -> dict: + def concurrent_case(index: int) -> dict[str, Any]: start_concurrent.wait() return send_completion( args.endpoint, @@ -118,7 +121,11 @@ def concurrent_case(index: int) -> dict: over_chunk = next(result for result in results if result["name"] == "over_chunk") prompt_tokens = (over_chunk.get("usage") or {}).get("prompt_tokens", 0) - return 0 if all(result["ok"] for result in results) and prompt_tokens > CHUNK_SIZE else 1 + return ( + 0 + if all(result["ok"] for result in results) and prompt_tokens > CHUNK_SIZE + else 1 + ) if __name__ == "__main__": diff --git a/tests/npu/test_async_cam_layered_graph_probe.py b/tests/npu/test_async_cam_layered_graph_probe.py index 8546330a3..4d6619941 100644 --- a/tests/npu/test_async_cam_layered_graph_probe.py +++ b/tests/npu/test_async_cam_layered_graph_probe.py @@ -62,9 +62,7 @@ def make_weight(k: int, n: int, factor: float, squeeze: bool): w13, s13, b13 = make_weight(hidden, 2 * intermediate, factor, True) w2, s2, b2 = make_weight(intermediate, hidden, factor, False) layers.append( - W4A8LayerWeights( - layer_idx, w13, w2, s13, s2, b13, b2, True, 0.0, 1.0 - ) + W4A8LayerWeights(layer_idx, w13, w2, s13, s2, b13, b2, True, 0.0, 1.0) ) executor = AsyncCAMW4A8Executor(layers) @@ -96,15 +94,21 @@ def make_weight(k: int, n: int, factor: float, squeeze: bool): static_hidden.copy_(shared_hidden.npu()) static_counts.copy_(torch.tensor(counts, dtype=torch.int64).npu()) static_info.copy_( - torch.tensor([capacity * 2, 0, layer_idx, valid_rows], dtype=torch.int64).npu() + torch.tensor( + [capacity * 2, 0, layer_idx, valid_rows], dtype=torch.int64 + ).npu() ) graph.replay() torch.npu.synchronize() assert captured_output.shape == (capacity, hidden) graphed = captured_output[:valid_rows].cpu().float() - eager = executor( - static_hidden, static_scales, static_counts, static_info - )[:valid_rows].cpu().float() + eager = ( + executor(static_hidden, static_scales, static_counts, static_info)[ + :valid_rows + ] + .cpu() + .float() + ) if valid_rows: difference = (graphed - eager).abs() print( diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index a8ae3559a..819f688c2 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -2958,7 +2958,7 @@ def test_npu_async_cam_ffn_startup_orders_warmup_capture_and_ready(monkeypatch): from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector from afd_plugin.v1.worker.npu import ffn_worker - events = [] + events: list[str | tuple[str, int | None]] = [] connector = object.__new__(CAMAsyncAFDConnector) connector._initialized = False connector._startup_store = object() From 44d088569f90d28fe574d193294556c7564450b4 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 15:24:54 +0800 Subject: [PATCH 4/8] refactor(npu): require FULL and rely on paired FFN warmup Signed-off-by: jiangkuaixue123 --- afd_plugin/compat/npu/feature_validation.py | 13 +++++ afd_plugin/v1/worker/npu/ffn_model_runner.py | 59 +------------------- afd_plugin/v1/worker/npu/ffn_worker.py | 8 ++- tests/unit/v1/worker/test_npu_runtime.py | 25 +++++++-- 4 files changed, 41 insertions(+), 64 deletions(-) diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 653f3c301..b480747dc 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -183,6 +183,19 @@ def _fail_if_unsupported_npu_afd_async_features( "CAMAsyncAFDConnector requires additional_config['afd'] " "with async=true and connector='CAMAsyncAFDConnector'", ) + if ( + afd_config.role == "ffn" + and _is_dsv4_target(vllm_config) + and async_cam_layered_gmm_enabled() + and ( + bool(vllm_config.model_config.enforce_eager) + or vllm_config.compilation_config.cudagraph_mode.name != "FULL" + ) + ): + raise RuntimeError( + "DSV4 layered Async CAM FFN requires enforce_eager=false " + "and cudagraph_mode=FULL" + ) if not bool(vllm_config.model_config.enforce_eager): graph_mode = vllm_config.compilation_config.cudagraph_mode.name if not ( diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 94b67c53a..a6380aa3a 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -41,7 +41,6 @@ ) from afd_plugin.envs import async_cam_layered_gmm_enabled from afd_plugin.model_executor.npu.async_cam_w4a8 import ( - CAM_LAYER_INDEX, AsyncCAMW4A8Executor, ) from afd_plugin.v1.worker.attention_metadata import ( @@ -251,8 +250,8 @@ def execute_connector_driven_step(self) -> None: self._ffn_forward_connector_driven() return None - def warmup_async_cam_ffn_graph(self) -> None: - """Prepare layered GMM locally before collective CAM initialization.""" + def prepare_async_cam_ffn_graph(self) -> None: + """Initialize the layered executor before collective CAM warmup.""" if ( not self._async_cam_ffn_graph_enabled or self._async_cam_ffn_graph is not None @@ -260,60 +259,8 @@ def warmup_async_cam_ffn_graph(self) -> None: return if self._layered_executor is None: self._initialize_layered_executor() - executor = self._layered_executor - if executor is None: + if self._layered_executor is None: raise RuntimeError("CAM async FFN FULL requires layered GMM") - connector = cast(CAMAsyncAFDConnector, self.connector) - # Meta dispatch uses the very same C++ allocator and float32 - # BATCH_SIZE_FACTOR parser as the real receive, without communication. - meta_outputs = torch.ops.afd_ascend.afd_async_dispatch_recv( - torch.empty((1,), dtype=connector.activation_dtype, device="meta"), - torch.empty((1,), dtype=torch.float16, device="meta"), - connector.comm_id, - connector.max_num_batched_tokens, - connector.hidden_size, - connector.topk, - connector.ffn_size, - connector.attn_size, - connector.expert_per_rank, - connector.world_rank, - connector.topology.world_size, - connector.tp_size, - connector.dynamic_quant, - "warmup", - ) - capacity = meta_outputs[0].shape[0] - hidden_states = torch.empty( - meta_outputs[0].shape, - dtype=meta_outputs[0].dtype, - device=self.device, - ) - hidden_states[0].zero_() - dynamic_scales = torch.empty( - meta_outputs[1].shape, - dtype=meta_outputs[1].dtype, - device=self.device, - ) - dynamic_scales[0] = 1.0 - group_list = torch.zeros( - meta_outputs[3].shape, - dtype=meta_outputs[3].dtype, - device=self.device, - ) - group_list[0] = 1 - batch_info = torch.zeros( - meta_outputs[2].shape, - dtype=meta_outputs[2].dtype, - device=self.device, - ) - batch_info[CAM_LAYER_INDEX] = executor.layer_ids[0] - executor(hidden_states, dynamic_scales, group_list, batch_info) - torch.npu.synchronize() - logger.info( - "CAM async FFN local layered warmup complete rank=%d capacity=%d", - connector.world_rank, - capacity, - ) def warmup_async_cam_ffn_communication(self) -> None: """Consume one controlled CAM work item from each Attention DP group.""" diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index 1f9d6b794..0113fd5a4 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -127,11 +127,13 @@ def start_ffn_server_loop(self) -> None: try: if cam_connector is not None: logger.info( - "CAM FFN local warmup start rank=%d", cam_connector.world_rank + "CAM FFN layered executor prepare start rank=%d", + cam_connector.world_rank, ) - self.model_runner.warmup_async_cam_ffn_graph() + self.model_runner.prepare_async_cam_ffn_graph() logger.info( - "CAM FFN local warmup done rank=%d", cam_connector.world_rank + "CAM FFN layered executor prepare done rank=%d", + cam_connector.world_rank, ) if not connector.is_initialized: self.model_runner.initialize_afd_connector() diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 819f688c2..c5d58555d 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -2469,9 +2469,24 @@ def test_npu_async_feature_validation_allows_only_layered_dsv4_ffn_full( ) fail_if_unsupported_npu_afd_features(config) - config.compilation_config.cudagraph_mode.name = "FULL_DECODE_ONLY" - with pytest.raises(RuntimeError, match="requires FFN FULL"): - fail_if_unsupported_npu_afd_features(config) + for enforce_eager, graph_mode in ( + (True, "FULL"), + (False, "NONE"), + (False, "PIECEWISE"), + (False, "FULL_AND_PIECEWISE"), + (False, "FULL_DECODE_ONLY"), + ): + config.model_config.enforce_eager = enforce_eager + config.compilation_config.cudagraph_mode.name = graph_mode + with pytest.raises(RuntimeError, match="requires enforce_eager=false"): + fail_if_unsupported_npu_afd_features(config) + + monkeypatch.setattr( + feature_validation, "async_cam_layered_gmm_enabled", lambda: False + ) + config.model_config.enforce_eager = True + config.compilation_config.cudagraph_mode.name = "NONE" + fail_if_unsupported_npu_afd_features(config) @pytest.mark.parametrize( @@ -2983,7 +2998,7 @@ def initialize_connector(): connector=connector, _async_cam_ffn_graph_enabled=True, _layered_executor=SimpleNamespace(layer_ids=(0,)), - warmup_async_cam_ffn_graph=lambda: events.append("warmup"), + prepare_async_cam_ffn_graph=lambda: events.append("prepare"), warmup_async_cam_ffn_communication=lambda: events.append("communication"), initialize_afd_connector=initialize_connector, capture_async_cam_ffn_graph=lambda: events.append("capture"), @@ -3005,7 +3020,7 @@ def is_alive(self): worker.start_ffn_server_loop() assert events == [ - "warmup", + "prepare", "connector", "coordination", ("mode", 0), From 04211aa683dd24b360b6f14c5ae8b16d97d00e68 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 15:26:06 +0800 Subject: [PATCH 5/8] fix(npu): preserve eager FFN alongside FULL graphs Signed-off-by: jiangkuaixue123 --- afd_plugin/compat/npu/feature_validation.py | 13 ------------- tests/unit/v1/worker/test_npu_runtime.py | 17 ++++++----------- 2 files changed, 6 insertions(+), 24 deletions(-) diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index b480747dc..653f3c301 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -183,19 +183,6 @@ def _fail_if_unsupported_npu_afd_async_features( "CAMAsyncAFDConnector requires additional_config['afd'] " "with async=true and connector='CAMAsyncAFDConnector'", ) - if ( - afd_config.role == "ffn" - and _is_dsv4_target(vllm_config) - and async_cam_layered_gmm_enabled() - and ( - bool(vllm_config.model_config.enforce_eager) - or vllm_config.compilation_config.cudagraph_mode.name != "FULL" - ) - ): - raise RuntimeError( - "DSV4 layered Async CAM FFN requires enforce_eager=false " - "and cudagraph_mode=FULL" - ) if not bool(vllm_config.model_config.enforce_eager): graph_mode = vllm_config.compilation_config.cudagraph_mode.name if not ( diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index c5d58555d..e3521f17f 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -2469,21 +2469,16 @@ def test_npu_async_feature_validation_allows_only_layered_dsv4_ffn_full( ) fail_if_unsupported_npu_afd_features(config) - for enforce_eager, graph_mode in ( - (True, "FULL"), - (False, "NONE"), - (False, "PIECEWISE"), - (False, "FULL_AND_PIECEWISE"), - (False, "FULL_DECODE_ONLY"), + for graph_mode in ( + "NONE", + "PIECEWISE", + "FULL_AND_PIECEWISE", + "FULL_DECODE_ONLY", ): - config.model_config.enforce_eager = enforce_eager config.compilation_config.cudagraph_mode.name = graph_mode - with pytest.raises(RuntimeError, match="requires enforce_eager=false"): + with pytest.raises(RuntimeError, match="requires FFN FULL"): fail_if_unsupported_npu_afd_features(config) - monkeypatch.setattr( - feature_validation, "async_cam_layered_gmm_enabled", lambda: False - ) config.model_config.enforce_eager = True config.compilation_config.cudagraph_mode.name = "NONE" fail_if_unsupported_npu_afd_features(config) From 81b09366d62c4a1587e21ac78e1a9eed1d28c270 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 16:40:23 +0800 Subject: [PATCH 6/8] refactor(npu): centralize AsyncCAM graph startup coordination Signed-off-by: jiangkuaixue123 --- afd_plugin/connectors/factory.py | 21 + afd_plugin/connectors/npu/async_cam.py | 263 +----------- afd_plugin/distributed/__init__.py | 2 + afd_plugin/distributed/afd_process_group.py | 30 ++ afd_plugin/v1/worker/npu/async_cam_startup.py | 377 ++++++++++++++++++ .../v1/worker/npu/attention_model_runner.py | 66 ++- afd_plugin/v1/worker/npu/ffn_model_runner.py | 73 +++- afd_plugin/v1/worker/npu/ffn_worker.py | 123 +++--- .../connectors/test_async_cam_connector.py | 280 ++++++------- .../unit/v1/worker/test_async_cam_startup.py | 237 +++++++++++ tests/unit/v1/worker/test_npu_runtime.py | 117 +++--- 11 files changed, 1030 insertions(+), 559 deletions(-) create mode 100644 afd_plugin/v1/worker/npu/async_cam_startup.py create mode 100644 tests/unit/v1/worker/test_async_cam_startup.py diff --git a/afd_plugin/connectors/factory.py b/afd_plugin/connectors/factory.py index 6e94e1ac1..9a47da9ec 100644 --- a/afd_plugin/connectors/factory.py +++ b/afd_plugin/connectors/factory.py @@ -19,6 +19,10 @@ if TYPE_CHECKING: from vllm.config import VllmConfig + from afd_plugin.distributed.afd_process_group import ( + ProcessGroupRendezvousContext, + ) + class AFDConnectorFactory: _registry: dict[str, Callable[[], type[AFDConnectorBase]]] = {} @@ -53,12 +57,29 @@ def create_connector( local_rank: int, vllm_config: VllmConfig, afd_config: AFDConfig | None = None, + *, + rendezvous_context: ProcessGroupRendezvousContext | None = None, ) -> AFDConnectorBase: config = afd_config or parse_afd_config(vllm_config) if config.connector not in cls._registry: raise ValueError(f"unsupported AFD connector type: {config.connector}") connector_cls = cls._registry[config.connector]() role_rank = resolve_role_rank(vllm_config, config) + if rendezvous_context is not None: + from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector + + if not issubclass(connector_cls, CAMAsyncAFDConnector): + raise TypeError( + "rendezvous_context requires a CAMAsyncAFDConnector class" + ) + return connector_cls( + rank, + local_rank, + vllm_config, + config, + role_rank, + rendezvous_context=rendezvous_context, + ) return connector_cls( rank, local_rank, diff --git a/afd_plugin/connectors/npu/async_cam.py b/afd_plugin/connectors/npu/async_cam.py index 20c05e0ac..fb6834b22 100644 --- a/afd_plugin/connectors/npu/async_cam.py +++ b/afd_plugin/connectors/npu/async_cam.py @@ -28,7 +28,6 @@ from __future__ import annotations import os -import time from collections.abc import Mapping from dataclasses import dataclass from datetime import timedelta @@ -37,7 +36,6 @@ import torch import torch.distributed as dist from torch import Tensor -from torch.distributed.distributed_c10d import PrefixStore, Store from vllm.logger import init_logger from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available @@ -61,6 +59,7 @@ AFDTransferState, ) from afd_plugin.distributed import ( + ProcessGroupRendezvousContext, create_hccl_process_group_options, init_afd_process_group, ) @@ -75,11 +74,6 @@ ASYNC_MOE_NUM_STAGES = 2 ASYNC_MOE_REQUEST_SPLIT = "request" ASYNC_MOE_TOKEN_SPLIT = "token" -FFN_STARTUP_TIMEOUT_SECONDS = 300 -FFN_STARTUP_POLL_SECONDS = 0.25 -FFN_STARTUP_NONCE_BYTES = 16 -STARTUP_WARMUP_TOKEN_COUNT = 1 -STARTUP_WARMUP_HIDDEN_VALUE = 0.25 _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset( { @@ -245,6 +239,8 @@ def __init__( vllm_config: VllmConfig, afd_config: AFDConfig, role_rank: int, + *, + rendezvous_context: ProcessGroupRendezvousContext | None = None, ) -> None: """Derive CAM topology, tensor dimensions, and connector state. @@ -270,8 +266,7 @@ def __init__( self.comm_id = CAM_COMM_ID self.tp_size = extra_info.attn_ranks_per_dp self.cam_pg: ProcessGroup | None = None - self._rendezvous_store: Store | None = None - self._startup_store: Store | None = None + self._rendezvous_context = rendezvous_context self.topology = build_async_topology( afd_config, role_rank, @@ -314,7 +309,11 @@ def init_afd_connector(self) -> None: pg_options=create_hccl_process_group_options( self.hccl_buffer_size_mb, ), - on_rendezvous=self._retain_rendezvous_store, + on_rendezvous=( + self._rendezvous_context.retain_store + if self._rendezvous_context is not None + else None + ), ) backend = self.cam_pg._get_backend(torch.device("npu")) self.group_name = str(backend.get_hccl_comm_name(self.world_rank)) @@ -325,253 +324,21 @@ def init_afd_connector(self) -> None: dtype=self.activation_dtype, device=device, ) + if self._rendezvous_context is not None: + self._rendezvous_context.bind(self.cam_pg) self._initialized = True - def _retain_rendezvous_store(self, store: Store) -> None: - self._rendezvous_store = store - - def _get_startup_store(self) -> Store: - """Make one rank-shared namespace tied to this CAM process group.""" - if self._startup_store is not None: - return self._startup_store - if self.cam_pg is None or self._rendezvous_store is None: - raise RuntimeError("CAM startup coordination requires an initialized group") - nonce = torch.zeros( - FFN_STARTUP_NONCE_BYTES, - dtype=torch.uint8, - device=f"npu:{self.local_rank}", - ) - if self.world_rank == 0: - nonce.copy_( - torch.tensor( - list(os.urandom(FFN_STARTUP_NONCE_BYTES)), - dtype=torch.uint8, - device=nonce.device, - ) - ) - logger.info("CAM startup nonce broadcast start rank=%d", self.world_rank) - dist.broadcast(nonce, src=0, group=self.cam_pg) - logger.info("CAM startup nonce broadcast done rank=%d", self.world_rank) - torch.npu.synchronize() - logger.info("CAM startup nonce synchronize done rank=%d", self.world_rank) - namespace = bytes(nonce.cpu().tolist()).hex() - self._startup_store = PrefixStore( - f"{AFD_ASYNC_CAM_GROUP_NAME}/startup/{namespace}", - self._rendezvous_store, - ) - return self._startup_store - - def prepare_startup_coordination(self) -> None: - self._get_startup_store() - - def _wait_for_startup_entries( - self, keys: list[str], failure_keys: list[str] - ) -> list[str]: - store = self._get_startup_store() - deadline = time.monotonic() + FFN_STARTUP_TIMEOUT_SECONDS - while True: - for key in failure_keys: - if store.check([key]): - status = store.get(key).decode() - if key.startswith("attn/failure/") or status.startswith("failed:"): - raise RuntimeError(f"CAM startup {key} {status}") - values = [] - pending = [] - for key in keys: - if store.check([key]): - values.append(store.get(key).decode()) - else: - pending.append(key) - if not pending: - return values - if time.monotonic() >= deadline: - raise TimeoutError( - f"Timed out waiting for CAM startup keys: {sorted(pending)}" - ) - time.sleep(FFN_STARTUP_POLL_SECONDS) - - def publish_ffn_mode(self, layer_idx: int | None) -> None: - store = self._get_startup_store() - if self.topology.role != "ffn": - raise RuntimeError("Only CAM FFN ranks may publish their mode") - mode = "eager" if layer_idx is None else f"graph:{layer_idx}" - store.set(f"ffn/mode/{self.world_rank}", mode) - - def wait_for_ffn_modes(self) -> int | None: - ffn_ranks = range(self.attn_size, self.topology.world_size) - modes = self._wait_for_startup_entries( - [f"ffn/mode/{rank}" for rank in ffn_ranks], - [f"ffn/{rank}" for rank in ffn_ranks], - ) - if len(set(modes)) != 1: - raise RuntimeError(f"CAM FFN startup modes disagree: {modes}") - mode = modes[0] - if mode == "eager": - return None - if not mode.startswith("graph:"): - raise RuntimeError(f"Invalid CAM FFN startup mode: {mode}") - return int(mode.removeprefix("graph:")) - - def publish_attention_startup_failure(self, error: str) -> None: - self._get_startup_store().set(f"attn/failure/{self.world_rank}", error) - - def _attention_failure_keys(self) -> list[str]: - return [f"attn/failure/{rank}" for rank in range(self.attn_size)] - - def _ffn_failure_keys(self) -> list[str]: - return [ - f"ffn/{rank}" for rank in range(self.attn_size, self.topology.world_size) - ] - - def wait_for_attention_warmup_prepared(self) -> None: - self._wait_for_startup_entries( - [f"attn/prepared/{rank}" for rank in range(self.attn_size)], - self._attention_failure_keys(), - ) - - def publish_ffn_warmup_started(self) -> None: - self._get_startup_store().set(f"ffn/warmup/start/{self.world_rank}", "ready") - - def publish_ffn_warmup_done(self) -> None: - self._get_startup_store().set(f"ffn/warmup/done/{self.world_rank}", "ready") - - def wait_for_communication_warmup_done(self) -> None: - self._wait_for_startup_entries( - [f"attn/done/{rank}" for rank in range(self.attn_size)] - + [ - f"ffn/warmup/done/{rank}" - for rank in range(self.attn_size, self.topology.world_size) - ], - self._attention_failure_keys() + self._ffn_failure_keys(), - ) - - def run_attention_startup_warmup(self, layer_idx: int) -> None: - """Complete one isolated CAM transaction before FFN graph capture.""" - if self.topology.role != "attention": - raise RuntimeError("CAM Attention warmup requires Attention role") - device = f"npu:{self.local_rank}" - hidden = torch.full( - (STARTUP_WARMUP_TOKEN_COUNT, self.hidden_size), - STARTUP_WARMUP_HIDDEN_VALUE, - dtype=self.activation_dtype, - device=device, - ) - tp_rank = self.world_rank % self.tp_size - expert_ids = [ - ((tp_rank * self.topk + index) % self.ffn_size) * self.expert_per_rank - + ((tp_rank * self.topk + index) // self.ffn_size) % self.expert_per_rank - for index in range(self.topk) - ] - ids = torch.tensor([expert_ids], dtype=torch.int32, device=device) - weights = torch.full( - (STARTUP_WARMUP_TOKEN_COUNT, self.topk), - 1.0 / self.topk, - dtype=torch.float32, - device=device, - ) - store = self._get_startup_store() - store.set(f"attn/prepared/{self.world_rank}", "ready") - self._wait_for_startup_entries( - [ - f"ffn/warmup/start/{rank}" - for rank in range(self.attn_size, self.topology.world_size) - ], - self._ffn_failure_keys(), - ) - logger.info("CAM Attention startup warmup DS rank=%d", self.world_rank) - torch.ops.afd_ascend.afd_async_dispatch_send( - hidden, - ids, - self.comm_args, - self.comm_id, - self.max_num_batched_tokens, - STARTUP_WARMUP_TOKEN_COUNT, - self.hidden_size, - self.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, - layer_idx, - self.tp_size, - self.dynamic_quant, - self.group_name, - ) - warmup_output = torch.ops.afd_ascend.afd_async_combine_recv( - self._placeholder, - ids, - weights, - self.comm_args, - self.comm_id, - STARTUP_WARMUP_TOKEN_COUNT, - self.hidden_size, - self.topk, - self.ffn_size, - self.attn_size, - self.expert_per_rank, - self.world_rank, - self.topology.world_size, - self.group_name, - ) - torch.npu.synchronize() - del warmup_output - logger.info("CAM Attention startup warmup CR done rank=%d", self.world_rank) - store.set(f"attn/done/{self.world_rank}", "ready") - self.wait_for_communication_warmup_done() - - def publish_ffn_startup(self, error: str | None = None) -> None: - store = self._get_startup_store() - if self.topology.role != "ffn": - raise RuntimeError("Only CAM FFN ranks may publish startup status") - key = f"ffn/{self.world_rank}" - if error is not None: - store.set(key, f"failed:{error}") - return - status = store.compare_set(key, "", "ready").decode() - if status != "ready": - raise RuntimeError(f"CAM FFN startup {key} {status}") - - def wait_for_ffn_ready(self) -> None: - store = self._get_startup_store() - keys = [ - f"ffn/{rank}" for rank in range(self.attn_size, self.topology.world_size) - ] - deadline = time.monotonic() + FFN_STARTUP_TIMEOUT_SECONDS - while True: - pending = [] - for key in keys: - if not store.check([key]): - pending.append(key) - continue - status = store.get(key).decode() - if status != "ready": - raise RuntimeError(f"CAM FFN startup {key} {status}") - if not pending: - return - if pending and time.monotonic() >= deadline: - raise TimeoutError( - f"Timed out waiting for CAM FFN ranks: {sorted(pending)}" - ) - time.sleep(FFN_STARTUP_POLL_SECONDS) - - def close(self, *, release_buffers: bool = True) -> None: - """Destroy the HCCL process group and clear pending transfer states.""" + def close(self) -> None: + """Destroy the communicator and release its operator buffers.""" if self.cam_pg is not None: dist.destroy_process_group(self.cam_pg) self.cam_pg = None self._initialized = False - if release_buffers: - self.release_closed_buffers() - - def release_closed_buffers(self) -> None: - if self._initialized: - raise RuntimeError("Cannot release active CAM operator buffers") self.comm_args = None self._placeholder = None - self._startup_store = None - self._rendezvous_store = None self._pending_attention_payloads.clear() + if self._rendezvous_context is not None: + self._rendezvous_context.invalidate() def select_experts(self, **kwargs: Any) -> tuple[Tensor, Tensor]: """Run the pinned vLLM-Ascend expert selector on Attention.""" diff --git a/afd_plugin/distributed/__init__.py b/afd_plugin/distributed/__init__.py index b7b28853b..995ae585c 100644 --- a/afd_plugin/distributed/__init__.py +++ b/afd_plugin/distributed/__init__.py @@ -15,6 +15,7 @@ def __getattr__(name: str): if name in { "DefaultProcessGroupSwitcher", + "ProcessGroupRendezvousContext", "create_hccl_process_group_options", "init_afd_process_group", }: @@ -29,6 +30,7 @@ def __getattr__(name: str): __all__ = [ "AFDRankMapping", "DefaultProcessGroupSwitcher", + "ProcessGroupRendezvousContext", "build_rank_mapping", "create_hccl_process_group_options", "init_afd_process_group", diff --git a/afd_plugin/distributed/afd_process_group.py b/afd_plugin/distributed/afd_process_group.py index 519aa3439..2ea33db20 100644 --- a/afd_plugin/distributed/afd_process_group.py +++ b/afd_plugin/distributed/afd_process_group.py @@ -23,6 +23,35 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer +class ProcessGroupRendezvousContext: + """Borrow the Store and process group created by one CAM rendezvous.""" + + def __init__(self) -> None: + self._store: Store | None = None + self._process_group: ProcessGroup | None = None + self._closed = False + + def retain_store(self, store: Store) -> None: + if self._closed or self._store is not None: + raise RuntimeError("CAM rendezvous context cannot accept another Store") + self._store = store + + def bind(self, process_group: ProcessGroup) -> None: + if self._closed or self._store is None or self._process_group is not None: + raise RuntimeError("CAM rendezvous context is not ready to bind") + self._process_group = process_group + + def borrow(self) -> tuple[Store, ProcessGroup]: + if self._closed or self._store is None or self._process_group is None: + raise RuntimeError("CAM rendezvous context is not bound") + return self._store, self._process_group + + def invalidate(self) -> None: + self._closed = True + self._store = None + self._process_group = None + + class DefaultProcessGroupSwitcher: """Temporarily switch PyTorch's default process group.""" @@ -123,6 +152,7 @@ def init_afd_process_group( __all__ = [ "DefaultProcessGroupSwitcher", + "ProcessGroupRendezvousContext", "create_hccl_process_group_options", "init_afd_process_group", ] diff --git a/afd_plugin/v1/worker/npu/async_cam_startup.py b/afd_plugin/v1/worker/npu/async_cam_startup.py new file mode 100644 index 000000000..8988f3fd6 --- /dev/null +++ b/afd_plugin/v1/worker/npu/async_cam_startup.py @@ -0,0 +1,377 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Host-side startup coordination for the Async CAM Attention/FFN world.""" + +from __future__ import annotations + +import os +import time +from collections.abc import Callable +from dataclasses import dataclass + +import torch +import torch.distributed as dist +from torch.distributed.distributed_c10d import PrefixStore, Store +from vllm.logger import init_logger + +from afd_plugin.connectors.base import AFDConnectorBase +from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata +from afd_plugin.connectors.npu.async_cam import ( + AFD_ASYNC_CAM_GROUP_NAME, + AFDAsyncTopology, +) +from afd_plugin.distributed.afd_process_group import ProcessGroupRendezvousContext + +STARTUP_TIMEOUT_SECONDS = 300 +STARTUP_POLL_SECONDS = 0.25 +STARTUP_NONCE_BYTES = 16 +STARTUP_WARMUP_STAGE = 0 +STARTUP_WARMUP_TOKEN_COUNT = 1 +STARTUP_WARMUP_HIDDEN_VALUE = 0.25 + +logger = init_logger(__name__) + + +@dataclass(frozen=True, slots=True) +class AsyncCamStartupSpec: + """Config-derived, immutable shape and topology of one CAM participant.""" + + topology: AFDAsyncTopology + local_rank: int + tp_size: int + hidden_size: int + topk: int + activation_dtype: torch.dtype + + @property + def attention_dp_size(self) -> int: + if self.topology.attn_size % self.tp_size: + raise ValueError("CAM Attention size must be divisible by TP size") + return self.topology.attn_size // self.tp_size + + +@dataclass(frozen=True, slots=True) +class FFNStartupPlan: + """Selected FFN startup mode and the layer used by graph warmup.""" + + use_graph: bool + first_layer_idx: int | None = None + + def __post_init__(self) -> None: + if self.use_graph and ( + self.first_layer_idx is None or self.first_layer_idx < 0 + ): + raise ValueError("CAM FFN graph startup requires a valid first layer") + if not self.use_graph and self.first_layer_idx is not None: + raise ValueError("CAM FFN eager startup cannot select a graph layer") + + +class AsyncCamStartupCoordinator: + """Run one rank's complete CAM startup without owning its communicator.""" + + def __init__( + self, + connector: AFDConnectorBase, + rendezvous_context: ProcessGroupRendezvousContext, + spec: AsyncCamStartupSpec, + ) -> None: + self._connector = connector + self._rendezvous_context = rendezvous_context + self._spec = spec + self._store: Store | None = None + self._started = False + self._starting = False + self._failure: Exception | None = None + + @property + def failed(self) -> bool: + return self._failure is not None + + @property + def started(self) -> bool: + return self._started + + def _begin(self, role: str) -> bool: + if self._spec.topology.role != role: + raise RuntimeError(f"CAM {role} startup requires {role} rank") + if self._failure is not None: + raise RuntimeError("CAM startup previously failed") from self._failure + if self._started: + self._rendezvous_context.borrow() + return False + if self._starting: + raise RuntimeError("CAM startup is already in progress") + self._starting = True + return True + + def _get_store(self) -> Store: + if self._store is not None: + self._rendezvous_context.borrow() + return self._store + store, process_group = self._rendezvous_context.borrow() + topology = self._spec.topology + nonce = torch.zeros( + STARTUP_NONCE_BYTES, + dtype=torch.uint8, + device=f"npu:{self._spec.local_rank}", + ) + if topology.world_rank == 0: + nonce.copy_( + torch.tensor( + list(os.urandom(STARTUP_NONCE_BYTES)), + dtype=torch.uint8, + device=nonce.device, + ) + ) + logger.info("CAM startup nonce broadcast start rank=%d", topology.world_rank) + dist.broadcast(nonce, src=0, group=process_group) + logger.info("CAM startup nonce broadcast done rank=%d", topology.world_rank) + torch.npu.synchronize() + logger.info("CAM startup nonce synchronize done rank=%d", topology.world_rank) + namespace = bytes(nonce.cpu().tolist()).hex() + self._store = PrefixStore( + f"{AFD_ASYNC_CAM_GROUP_NAME}/startup/{namespace}", store + ) + return self._store + + def _failure_keys(self) -> list[str]: + topology = self._spec.topology + return [f"attn/failure/{rank}" for rank in range(topology.attn_size)] + [ + f"ffn/{rank}" for rank in range(topology.attn_size, topology.world_size) + ] + + def _check_failures(self, store: Store, stage: str) -> None: + for key in self._failure_keys(): + if not store.check([key]): + continue + status = store.get(key).decode() + if key.startswith("attn/failure/") or status.startswith("failed:"): + raise RuntimeError( + f"CAM startup rank={self._spec.topology.world_rank} " + f"stage={stage} {key} {status}" + ) + + def _wait_for_entries(self, keys: list[str], *, stage: str) -> list[str]: + store = self._get_store() + deadline = time.monotonic() + STARTUP_TIMEOUT_SECONDS + while True: + self._check_failures(store, stage) + values = [] + pending = [] + for key in keys: + if store.check([key]): + values.append(store.get(key).decode()) + else: + pending.append(key) + if not pending: + self._check_failures(store, stage) + return values + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for CAM startup rank=" + f"{self._spec.topology.world_rank} stage={stage} " + f"keys={sorted(pending)}" + ) + time.sleep(STARTUP_POLL_SECONDS) + + def _wait_for_ffn_modes(self) -> int | None: + topology = self._spec.topology + modes = self._wait_for_entries( + [ + f"ffn/mode/{rank}" + for rank in range(topology.attn_size, topology.world_size) + ], + stage="mode", + ) + if len(set(modes)) != 1: + raise RuntimeError(f"CAM FFN startup modes disagree: {modes}") + mode = modes[0] + if mode == "eager": + return None + if not mode.startswith("graph:"): + raise RuntimeError(f"Invalid CAM FFN startup mode: {mode}") + layer_idx = int(mode.removeprefix("graph:")) + if layer_idx < 0: + raise RuntimeError(f"Invalid CAM FFN startup layer: {layer_idx}") + return layer_idx + + def _run_attention_warmup(self, layer_idx: int) -> None: + spec = self._spec + topology = spec.topology + device = f"npu:{spec.local_rank}" + hidden = torch.full( + (STARTUP_WARMUP_TOKEN_COUNT, spec.hidden_size), + STARTUP_WARMUP_HIDDEN_VALUE, + dtype=spec.activation_dtype, + device=device, + ) + tp_rank = topology.world_rank % spec.tp_size + expert_ids = [ + ((tp_rank * spec.topk + index) % topology.ffn_size) + * topology.expert_per_rank + + ((tp_rank * spec.topk + index) // topology.ffn_size) + % topology.expert_per_rank + for index in range(spec.topk) + ] + ids = torch.tensor([expert_ids], dtype=torch.int32, device=device) + weights = torch.full( + (STARTUP_WARMUP_TOKEN_COUNT, spec.topk), + 1.0 / spec.topk, + dtype=torch.float32, + device=device, + ) + store = self._get_store() + store.set(f"attn/prepared/{topology.world_rank}", "ready") + self._wait_for_entries( + [ + f"ffn/warmup/start/{rank}" + for rank in range(topology.attn_size, topology.world_size) + ], + stage="warmup-start", + ) + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=layer_idx, + stage_idx=STARTUP_WARMUP_STAGE, + seq_len=STARTUP_WARMUP_TOKEN_COUNT, + ) + context = AFDTransferContext(metadata=metadata) + logger.info("CAM Attention startup warmup DS rank=%d", topology.world_rank) + self._connector.send_attn_output( + hidden, context, topk_ids=ids, topk_weights=weights + ) + # Omit context and routing kwargs so the connector consumes its FIFO. + warmup_output = self._connector.recv_ffn_output( + hidden, ubatch_idx=STARTUP_WARMUP_STAGE + ) + torch.npu.synchronize() + del warmup_output + logger.info("CAM Attention startup warmup CR done rank=%d", topology.world_rank) + store.set(f"attn/done/{topology.world_rank}", "ready") + self._wait_for_communication_warmup_done() + + def _wait_for_communication_warmup_done(self) -> None: + topology = self._spec.topology + self._wait_for_entries( + [f"attn/done/{rank}" for rank in range(topology.attn_size)] + + [ + f"ffn/warmup/done/{rank}" + for rank in range(topology.attn_size, topology.world_size) + ], + stage="warmup-done", + ) + + def _wait_for_ffn_ready(self) -> None: + store = self._get_store() + topology = self._spec.topology + keys = [ + f"ffn/{rank}" for rank in range(topology.attn_size, topology.world_size) + ] + deadline = time.monotonic() + STARTUP_TIMEOUT_SECONDS + while True: + self._check_failures(store, "ready") + pending = [] + for key in keys: + if not store.check([key]): + pending.append(key) + continue + status = store.get(key).decode() + if status != "ready": + raise RuntimeError(f"CAM FFN startup {key} {status}") + if not pending: + self._check_failures(store, "ready") + return + if time.monotonic() >= deadline: + raise TimeoutError( + f"Timed out waiting for CAM FFN ranks: {sorted(pending)}" + ) + time.sleep(STARTUP_POLL_SECONDS) + + def report_failure(self, error: Exception) -> None: + """Publish a sticky failure only if the nonce Store already exists.""" + if self._failure is not None: + return + self._failure = error + store = self._store + if store is None: + return + topology = self._spec.topology + key = ( + f"attn/failure/{topology.world_rank}" + if topology.role == "attention" + else f"ffn/{topology.world_rank}" + ) + try: + store.set(key, f"failed:{error}") + except Exception: + logger.warning("CAM startup failure publication failed: %s", key) + + def start_attention(self) -> None: + if not self._begin("attention"): + return + try: + if not self._connector.is_initialized: + self._connector.init_afd_connector() + self._get_store() + layer_idx = self._wait_for_ffn_modes() + if layer_idx is not None: + self._run_attention_warmup(layer_idx) + self._wait_for_ffn_ready() + self._started = True + except Exception as exc: + self.report_failure(exc) + raise + finally: + self._starting = False + + def start_ffn( + self, + *, + prepare: Callable[[], FFNStartupPlan], + consume_warmup: Callable[[int], None], + capture: Callable[[], None], + start_receiver: Callable[[], None], + ) -> None: + if not self._begin("ffn"): + return + try: + plan = prepare() + if not self._connector.is_initialized: + self._connector.init_afd_connector() + store = self._get_store() + topology = self._spec.topology + mode = f"graph:{plan.first_layer_idx}" if plan.use_graph else "eager" + store.set(f"ffn/mode/{topology.world_rank}", mode) + observed_layer = self._wait_for_ffn_modes() + if observed_layer != plan.first_layer_idx: + raise RuntimeError( + f"CAM FFN startup rank={topology.world_rank} mode mismatch: " + f"local={mode}, all=graph:{observed_layer}" + ) + if plan.use_graph: + self._wait_for_entries( + [f"attn/prepared/{rank}" for rank in range(topology.attn_size)], + stage="warmup-prepared", + ) + store.set(f"ffn/warmup/start/{topology.world_rank}", "ready") + consume_warmup(self._spec.attention_dp_size) + store.set(f"ffn/warmup/done/{topology.world_rank}", "ready") + self._wait_for_communication_warmup_done() + capture() + start_receiver() + if self._failure is not None: + raise RuntimeError("CAM FFN receiver failed during startup") from ( + self._failure + ) + key = f"ffn/{topology.world_rank}" + status = store.compare_set(key, "", "ready").decode() + if status != "ready": + raise RuntimeError(f"CAM FFN startup {key} {status}") + self._started = True + except Exception as exc: + self.report_failure(exc) + raise + finally: + self._starting = False + + +__all__ = ["AsyncCamStartupCoordinator", "AsyncCamStartupSpec", "FFNStartupPlan"] diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index e50b0deb1..7efd8038b 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -89,6 +89,7 @@ AFDAsyncExtraInfo, CAMAsyncAFDConnector, ) +from afd_plugin.distributed.afd_process_group import ProcessGroupRendezvousContext from afd_plugin.model_executor.models.npu.async_cam_layout import ( ASYNC_MOE_UBATCH_METADATA_KEY, AsyncMoeUbatchMetadata, @@ -108,6 +109,10 @@ _resolve_world_ranks, build_ubatch_dp_metadata_list, ) +from afd_plugin.v1.worker.npu.async_cam_startup import ( + AsyncCamStartupCoordinator, + AsyncCamStartupSpec, +) from afd_plugin.v1.worker.npu.npu_ubatch_wrapper import AscendUBatchWrapper from afd_plugin.v1.worker.npu.ubatch_utils import ( check_enable_ubatch, @@ -142,12 +147,38 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) rank, _ = _resolve_world_ranks() local_rank = int(device.index) - self.connector = AFDConnectorFactory.create_connector( - rank, - local_rank, - vllm_config, - self.afd_config, + rendezvous_context = ( + ProcessGroupRendezvousContext() + if afd_config.connector == AFD_ASYNC_CONNECTOR + else None ) + if rendezvous_context is None: + self.connector = AFDConnectorFactory.create_connector( + rank, local_rank, vllm_config, self.afd_config + ) + else: + self.connector = AFDConnectorFactory.create_connector( + rank, + local_rank, + vllm_config, + self.afd_config, + rendezvous_context=rendezvous_context, + ) + self._async_cam_startup: AsyncCamStartupCoordinator | None = None + if rendezvous_context is not None: + assert isinstance(self.connector, CAMAsyncAFDConnector) + self._async_cam_startup = AsyncCamStartupCoordinator( + self.connector, + rendezvous_context, + AsyncCamStartupSpec( + topology=self.connector.topology, + local_rank=local_rank, + tp_size=self.connector.tp_size, + hidden_size=vllm_config.model_config.hf_config.hidden_size, + topk=vllm_config.model_config.hf_config.num_experts_per_tok, + activation_dtype=vllm_config.model_config.dtype, + ), + ) self.afd_async_extra_info = AFDAsyncExtraInfo() if afd_config.connector == AFD_ASYNC_CONNECTOR: connector_extra_info = self.connector.extra_info @@ -1573,19 +1604,12 @@ def load_model(self) -> None: # rendezvous is the blocking cross-role collective, so it is # deliberately last: it doubles as the "both roles finished loading # weights" barrier before memory profiling. - if not self.connector.is_initialized: - self.connector.init_afd_connector() # ### PATCH START: CAM startup coordination - if isinstance(self.connector, CAMAsyncAFDConnector): - try: - layer_idx = self.connector.wait_for_ffn_modes() - if layer_idx is not None: - self.connector.run_attention_startup_warmup(layer_idx) - self.connector.wait_for_ffn_ready() - except Exception as exc: - if self.connector._startup_store is not None: - self.connector.publish_attention_startup_failure(str(exc)) - raise + startup = self._async_cam_startup + if startup is not None: + startup.start_attention() + elif not self.connector.is_initialized: + self.connector.init_afd_connector() # ### PATCH END: CAM startup coordination def _install_ascend_ubatch_wrapper(self) -> None: @@ -1948,6 +1972,14 @@ def sync_and_slice_intermediate_tensors( def shutdown(self) -> None: stop_afd_npu_profiler(self.prof) + startup = self._async_cam_startup + if startup is not None and self.connector.is_initialized: + if not startup.started or startup.failed: + raise RuntimeError( + "CAM Attention startup did not complete safely; " + "retain communicator resources" + ) + torch.npu.synchronize() self.connector.close() super().shutdown() diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index a6380aa3a..ecbfe4127 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -26,7 +26,7 @@ step_afd_npu_profiler, stop_afd_npu_profiler, ) -from afd_plugin.config import AFDConfig, parse_afd_config +from afd_plugin.config import AFD_ASYNC_CONNECTOR, AFDConfig, parse_afd_config from afd_plugin.connectors import ( AFDConnectorFactory, AFDControlPayload, @@ -39,6 +39,7 @@ AFDAsyncTransferState, CAMAsyncAFDConnector, ) +from afd_plugin.distributed.afd_process_group import ProcessGroupRendezvousContext from afd_plugin.envs import async_cam_layered_gmm_enabled from afd_plugin.model_executor.npu.async_cam_w4a8 import ( AsyncCAMW4A8Executor, @@ -52,6 +53,11 @@ make_ffn_graph_key, ) from afd_plugin.v1.worker.ffn_model_runner import _set_moe_layer_index +from afd_plugin.v1.worker.npu.async_cam_startup import ( + AsyncCamStartupCoordinator, + AsyncCamStartupSpec, + FFNStartupPlan, +) if TYPE_CHECKING: from vllm.sequence import IntermediateTensors @@ -81,12 +87,38 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device) -> None: ) rank, _ = _resolve_world_ranks() local_rank = int(device.index) - self.connector = AFDConnectorFactory.create_connector( - rank, - local_rank, - vllm_config, - self.afd_config, + rendezvous_context = ( + ProcessGroupRendezvousContext() + if afd_config.connector == AFD_ASYNC_CONNECTOR + else None ) + if rendezvous_context is None: + self.connector = AFDConnectorFactory.create_connector( + rank, local_rank, vllm_config, self.afd_config + ) + else: + self.connector = AFDConnectorFactory.create_connector( + rank, + local_rank, + vllm_config, + self.afd_config, + rendezvous_context=rendezvous_context, + ) + self._async_cam_startup: AsyncCamStartupCoordinator | None = None + if rendezvous_context is not None: + assert isinstance(self.connector, CAMAsyncAFDConnector) + self._async_cam_startup = AsyncCamStartupCoordinator( + self.connector, + rendezvous_context, + AsyncCamStartupSpec( + topology=self.connector.topology, + local_rank=local_rank, + tp_size=self.connector.tp_size, + hidden_size=vllm_config.model_config.hf_config.hidden_size, + topk=vllm_config.model_config.hf_config.num_experts_per_tok, + activation_dtype=vllm_config.model_config.dtype, + ), + ) if ( isinstance(self.connector, CAMAsyncAFDConnector) and self.afd_config.compute_gate_on_attention @@ -250,31 +282,29 @@ def execute_connector_driven_step(self) -> None: self._ffn_forward_connector_driven() return None - def prepare_async_cam_ffn_graph(self) -> None: - """Initialize the layered executor before collective CAM warmup.""" - if ( - not self._async_cam_ffn_graph_enabled - or self._async_cam_ffn_graph is not None - ): - return + def prepare_async_cam_ffn_startup(self) -> FFNStartupPlan: + """Prepare the executor and select eager or one FULL graph.""" if self._layered_executor is None: self._initialize_layered_executor() - if self._layered_executor is None: - raise RuntimeError("CAM async FFN FULL requires layered GMM") + if not self._async_cam_ffn_graph_enabled: + return FFNStartupPlan(use_graph=False) + executor = self._layered_executor + if executor is None or not executor.layer_ids: + raise RuntimeError("CAM async FFN FULL requires layered GMM layers") + return FFNStartupPlan(use_graph=True, first_layer_idx=executor.layer_ids[0]) - def warmup_async_cam_ffn_communication(self) -> None: + def warmup_async_cam_ffn_communication(self, count: int) -> None: """Consume one controlled CAM work item from each Attention DP group.""" if not self._async_cam_ffn_graph_enabled: return connector = cast(CAMAsyncAFDConnector, self.connector) - attention_dp_size = connector.attn_size // connector.tp_size - for _ in range(attention_dp_size): + for _ in range(count): self._execute_layered_work_item() torch.npu.synchronize() logger.info( "CAM async FFN communication warmup done rank=%d items=%d", connector.world_rank, - attention_dp_size, + count, ) def capture_async_cam_ffn_graph(self) -> None: @@ -657,6 +687,11 @@ def shutdown(self) -> None: if self._is_shutdown: return stop_afd_npu_profiler(self.prof) + if self._async_cam_startup is not None and self.connector.is_initialized: + raise RuntimeError( + "CAM FFN receiver must drain and close its connector before " + "model runner shutdown" + ) if self.connector.is_initialized: self.connector.close() self.release_async_cam_ffn_graph() diff --git a/afd_plugin/v1/worker/npu/ffn_worker.py b/afd_plugin/v1/worker/npu/ffn_worker.py index 0113fd5a4..69b818750 100644 --- a/afd_plugin/v1/worker/npu/ffn_worker.py +++ b/afd_plugin/v1/worker/npu/ffn_worker.py @@ -59,6 +59,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._ffn_shutdown_event: threading.Event | None = None self._ffn_loop_error: BaseException | None = None self._ffn_loop_started_event: threading.Event | None = None + self._ffn_receiver_drained = False self._cpu_binding_attempted = False def init_device(self) -> None: @@ -109,72 +110,47 @@ def execute_model( ) def start_ffn_server_loop(self) -> None: + startup = self.model_runner._async_cam_startup if self._ffn_thread is not None and self._ffn_thread.is_alive(): self.raise_ffn_loop_error_if_any() + if startup is not None and startup.failed: + raise RuntimeError("CAM FFN startup previously failed") return - self.raise_ffn_loop_error_if_any() - connector = self.model_runner.connector - cam_connector = ( - connector if isinstance(connector, CAMAsyncAFDConnector) else None - ) - if cam_connector is not None: - logger.info( - "CAM FFN startup entered rank=%d graph=%s", - cam_connector.world_rank, - self.model_runner._async_cam_ffn_graph_enabled, + if startup is not None and self._ffn_thread is not None: + raise RuntimeError("CAM FFN receiver stopped; startup cannot be reused") + if startup is not None: + startup.start_ffn( + prepare=self.model_runner.prepare_async_cam_ffn_startup, + consume_warmup=self.model_runner.warmup_async_cam_ffn_communication, + capture=self.model_runner.capture_async_cam_ffn_graph, + start_receiver=self._start_ffn_receiver, ) - try: - if cam_connector is not None: - logger.info( - "CAM FFN layered executor prepare start rank=%d", - cam_connector.world_rank, - ) - self.model_runner.prepare_async_cam_ffn_graph() - logger.info( - "CAM FFN layered executor prepare done rank=%d", - cam_connector.world_rank, - ) - if not connector.is_initialized: - self.model_runner.initialize_afd_connector() - if cam_connector is not None: - logger.info( - "CAM FFN startup store start rank=%d", cam_connector.world_rank - ) - cam_connector.prepare_startup_coordination() - logger.info( - "CAM FFN startup store done rank=%d", cam_connector.world_rank - ) - if self.model_runner._async_cam_ffn_graph_enabled: - executor = self.model_runner._layered_executor - if executor is None: - raise RuntimeError("CAM FFN graph has no layered executor") - cam_connector.publish_ffn_mode(executor.layer_ids[0]) - cam_connector.wait_for_attention_warmup_prepared() - cam_connector.publish_ffn_warmup_started() - self.model_runner.warmup_async_cam_ffn_communication() - cam_connector.publish_ffn_warmup_done() - cam_connector.wait_for_communication_warmup_done() - self.model_runner.capture_async_cam_ffn_graph() - else: - cam_connector.publish_ffn_mode(None) - except Exception as exc: - if cam_connector is not None and cam_connector._startup_store is not None: - cam_connector.publish_ffn_startup(error=str(exc)) - raise + return + if not self.model_runner.connector.is_initialized: + self.model_runner.initialize_afd_connector() + self._start_ffn_receiver() + def _start_ffn_receiver(self) -> None: + """Start the receive thread and confirm that its NPU device is set.""" self._bind_cpus_once() self._ffn_shutdown_event = threading.Event() started_event = threading.Event() self._ffn_loop_started_event = started_event self._ffn_loop_error = None + self._ffn_receiver_drained = False + startup = self.model_runner._async_cam_startup def ffn_worker_loop() -> None: try: self._run_ffn_server_loop() except Exception as exc: shutdown_event = self._ffn_shutdown_event - if shutdown_event is not None and shutdown_event.is_set(): + if ( + startup is None + and shutdown_event is not None + and shutdown_event.is_set() + ): logger.debug( "AFD NPU FFN receive loop stopped during shutdown", exc_info=True, @@ -183,11 +159,8 @@ def ffn_worker_loop() -> None: self._ffn_loop_error = exc started_event.set() logger.exception("AFD NPU FFN worker loop failed") - if ( - cam_connector is not None - and cam_connector._startup_store is not None - ): - cam_connector.publish_ffn_startup(error=str(exc)) + if startup is not None: + startup.report_failure(exc) self._ffn_thread = threading.Thread( target=ffn_worker_loop, @@ -196,7 +169,7 @@ def ffn_worker_loop() -> None: ) try: self._ffn_thread.start() - if cam_connector is not None: + if startup is not None: if not started_event.wait(timeout=FFN_STARTUP_THREAD_TIMEOUT_SECONDS): self.raise_ffn_loop_error_if_any() raise TimeoutError("AFD NPU FFN service thread did not start") @@ -205,10 +178,9 @@ def ffn_worker_loop() -> None: raise RuntimeError( "AFD NPU FFN service thread stopped during startup" ) - cam_connector.publish_ffn_startup() except Exception as exc: - if cam_connector is not None and cam_connector._startup_store is not None: - cam_connector.publish_ffn_startup(error=str(exc)) + if startup is not None: + startup.report_failure(exc) raise def _bind_cpus_once(self) -> None: @@ -262,11 +234,11 @@ def _run_ffn_server_loop(self) -> None: is_graph_replaying=is_graph_replaying, ) torch.npu.synchronize() + self._ffn_receiver_drained = True def raise_ffn_loop_error_if_any(self) -> None: error = self._ffn_loop_error if error is not None: - self._ffn_loop_error = None raise RuntimeError("AFD NPU FFN worker loop failed") from error def stop_ffn_server_loop(self) -> None: @@ -278,35 +250,40 @@ def stop_ffn_server_loop(self) -> None: # work item. Do not destroy its communicator or graph-owned tensors # while the worker thread is still using them. connector = self.model_runner.connector - graph_connector = ( - connector - if isinstance(connector, CAMAsyncAFDConnector) - and self.model_runner._async_cam_ffn_graph_enabled - else None + cam_connector = ( + connector if isinstance(connector, CAMAsyncAFDConnector) else None ) thread = self._ffn_thread - if graph_connector is not None and thread is not None: + if cam_connector is not None and thread is not None: thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): raise RuntimeError( - "AFD NPU FFN graph loop is still active; retain graph and " + "AFD NPU FFN receive loop is still active; retain graph and " "connector resources until the worker process exits" ) - # Release a completed graph before its communicator. The legacy path - # still closes the connector before joining its receive thread. - if graph_connector is not None: + self.raise_ffn_loop_error_if_any() + if not self._ffn_receiver_drained: + raise RuntimeError( + "AFD NPU FFN receiver did not confirm device completion; " + "retain graph and connector resources" + ) + startup = self.model_runner._async_cam_startup + if cam_connector is not None and startup is not None and startup.failed: + raise RuntimeError( + "CAM FFN startup failed; retain communicator and graph resources" + ) + # Only a normally drained CAM receiver may release captured resources. + if cam_connector is not None: self.model_runner.release_async_cam_ffn_graph() - graph_connector.close(release_buffers=False) + cam_connector.close() else: connector.close() - if thread is not None and graph_connector is None: + if thread is not None and cam_connector is None: thread.join(timeout=FFN_SHUTDOWN_TIMEOUT_SECONDS) if thread.is_alive(): raise RuntimeError( "AFD NPU FFN worker loop did not stop after connector close", ) - if graph_connector is not None: - graph_connector.release_closed_buffers() self._ffn_thread = None self._ffn_shutdown_event = None self._ffn_loop_started_event = None diff --git a/tests/unit/connectors/test_async_cam_connector.py b/tests/unit/connectors/test_async_cam_connector.py index 46693a726..71fbf4317 100644 --- a/tests/unit/connectors/test_async_cam_connector.py +++ b/tests/unit/connectors/test_async_cam_connector.py @@ -33,6 +33,9 @@ CAMAsyncAFDConnector, build_async_topology, ) +from afd_plugin.distributed.afd_process_group import ( # noqa: E402 + ProcessGroupRendezvousContext, +) class _FakeTensorLike: @@ -353,7 +356,7 @@ def fake_init_afd_process_group(**kwargs): "timeout": calls[0]["timeout"], "init_method": "tcp://127.0.0.1:1239", "pg_options": pg_options, - "on_rendezvous": connector._retain_rendezvous_store, + "on_rendezvous": None, }, ] assert connector.cam_pg is not None @@ -512,6 +515,84 @@ def test_async_ffn_side_dispatch_recv_and_combine_send(monkeypatch): assert logs[index][2]["max_seq_len"] == connector.max_num_batched_tokens +def test_startup_warmup_uses_real_connector_fifo_then_formal_task(monkeypatch): + from afd_plugin.v1.worker.npu import async_cam_startup + + fake_torch = _FakeTorch() + routed_ids = [] + synchronizations = [] + fake_torch.full = lambda shape, value, *, dtype, device: _FakeTensor( + shape, dtype=dtype, device=device + ) + + def make_ids(values, *, dtype, device): + routed_ids.append(values) + return _FakeTensor((len(values), len(values[0])), dtype=dtype, device=device) + + fake_torch.tensor = make_ids + fake_torch.npu = SimpleNamespace( + synchronize=lambda: synchronizations.append("done") + ) + monkeypatch.setattr(async_cam_module, "torch", fake_torch) + monkeypatch.setattr(async_cam_startup, "torch", fake_torch) + connector = CAMAsyncAFDConnector( + 0, 0, _vllm_config(), _afd_config(role="attention"), 0 + ) + connector._initialized = True + connector.comm_args = _FakeTensor((1,), dtype="fp16") + + class ReadyStore: + values = { + **{f"ffn/warmup/start/{rank}": b"ready" for rank in (4, 5)}, + **{f"ffn/warmup/done/{rank}": b"ready" for rank in (4, 5)}, + **{f"attn/done/{rank}": b"ready" for rank in (1, 2, 3)}, + } + + def set(self, key, value): + self.values[key] = value.encode() + + def check(self, keys): + return all(key in self.values for key in keys) + + def get(self, key): + return self.values[key] + + context = ProcessGroupRendezvousContext() + store = ReadyStore() + context.retain_store(store) + context.bind(SimpleNamespace()) + startup = async_cam_startup.AsyncCamStartupCoordinator( + connector, + context, + async_cam_startup.AsyncCamStartupSpec( + topology=connector.topology, + local_rank=0, + tp_size=connector.tp_size, + hidden_size=connector.hidden_size, + topk=connector.topk, + activation_dtype=fake_torch.bfloat16, + ), + ) + startup._store = store + startup._run_attention_warmup(layer_idx=3) + + assert routed_ids == [[[0, 4]]] + assert synchronizations == ["done"] + assert connector._pending_attention_payloads == {} + assert store.values["attn/done/0"] == b"ready" + + formal_context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=5, stage_idx=0, seq_len=1 + ) + ) + formal_hidden = _FakeTensor((1, 16)) + connector.send_attn_output(formal_hidden, formal_context, **_topk_payload(1)) + assert connector._pending_attention_payloads[0][0][0] is formal_context + connector.recv_ffn_output(formal_hidden, ubatch_idx=0) + assert connector._pending_attention_payloads == {} + + def test_async_combine_send_requires_dispatch_recv_token_metadata(monkeypatch): fake_torch = _FakeTorch() monkeypatch.setattr(async_cam_module, "torch", fake_torch) @@ -854,146 +935,69 @@ def test_close_releases_every_stage_routing(monkeypatch): assert connector._pending_attention_payloads == {} -def test_ffn_startup_failure_reaches_attention_waiter(): - class FakeStore: - def __init__(self): - self.values = {} - - def set(self, key, value): - self.values[key] = value.encode() - - def check(self, keys): - return all(key in self.values for key in keys) - - def get(self, key): - return self.values[key] - - store = FakeStore() - ffn = object.__new__(CAMAsyncAFDConnector) - ffn._startup_store = store - ffn.topology = SimpleNamespace(role="ffn", world_size=3) - ffn.world_rank = 2 - ffn.publish_ffn_startup(error="graph capture failed") - - attention = object.__new__(CAMAsyncAFDConnector) - attention._startup_store = store - attention.attn_size = 2 - attention.topology = SimpleNamespace(world_size=3) - with pytest.raises(RuntimeError, match="graph capture failed"): - attention.wait_for_ffn_ready() - - -def test_attention_wait_rechecks_previously_ready_ffn(monkeypatch): - class ChangingStore: - def __init__(self): - self.values = {"ffn/8": b"ready"} - - def check(self, keys): - return all(key in self.values for key in keys) - - def get(self, key): - return self.values[key] - - attention = object.__new__(CAMAsyncAFDConnector) - store = ChangingStore() - attention._startup_store = store - attention.attn_size = 8 - attention.topology = SimpleNamespace(world_size=10) - - def fail_first_rank(_: float) -> None: - store.values["ffn/8"] = b"failed:loop stopped" - store.values["ffn/9"] = b"ready" - - monkeypatch.setattr(async_cam_module.time, "sleep", fail_first_rank) - - with pytest.raises(RuntimeError, match="loop stopped"): - attention.wait_for_ffn_ready() - - -def test_ffn_ready_cannot_overwrite_failed_status(): - class FakeStore: - def __init__(self): - self.values = {"ffn/8": b"failed:startup crashed"} - - def compare_set(self, key, expected, desired): - current = self.values.get(key, b"") - if current == expected.encode(): - current = desired.encode() - self.values[key] = current - return current - - ffn = object.__new__(CAMAsyncAFDConnector) - ffn._startup_store = FakeStore() - ffn.topology = SimpleNamespace(role="ffn") - ffn.world_rank = 8 - with pytest.raises(RuntimeError, match="startup crashed"): - ffn.publish_ffn_startup() - - -def test_graph_warmup_wait_propagates_attention_failure(): - class FakeStore: - def __init__(self): - self.values = {"attn/failure/0": b"failed before dispatch"} - - def check(self, keys): - return all(key in self.values for key in keys) - - def get(self, key): - return self.values[key] - - ffn = object.__new__(CAMAsyncAFDConnector) - ffn._startup_store = FakeStore() - ffn.attn_size = 2 - with pytest.raises(RuntimeError, match="failed before dispatch"): - ffn.wait_for_attention_warmup_prepared() - - -def test_graph_capture_waits_for_all_attention_and_ffn_warmup_done(monkeypatch): - class FakeStore: - def __init__(self): - self.values = { - "attn/done/0": b"ready", - "ffn/warmup/done/1": b"ready", - } - - def check(self, keys): - return all(key in self.values for key in keys) +def test_async_factory_injects_context_only_into_cam_connector(): + context = ProcessGroupRendezvousContext() + connector = AFDConnectorFactory.create_connector( + 0, + 0, + _vllm_config(), + _afd_config(role="attention"), + rendezvous_context=context, + ) + assert isinstance(connector, CAMAsyncAFDConnector) + assert connector._rendezvous_context is context + with pytest.raises(RuntimeError, match="not bound"): + context.borrow() - def get(self, key): - return self.values[key] + with pytest.raises(TypeError, match="requires a CAMAsyncAFDConnector"): + AFDConnectorFactory.create_connector( + 0, + 0, + _vllm_config(), + AFDConfig(connector="CAMP2pAFDConnector", role="attention"), + rendezvous_context=context, + ) - store = FakeStore() - ffn = object.__new__(CAMAsyncAFDConnector) - ffn._startup_store = store - ffn.attn_size = 1 - ffn.topology = SimpleNamespace(world_size=3) - polls = [] - def finish_second_ffn(_: float) -> None: - polls.append(True) - store.values["ffn/warmup/done/2"] = b"ready" +def test_async_connector_binds_context_only_after_anchor_init(monkeypatch): + context = ProcessGroupRendezvousContext() + store = object() + backend = SimpleNamespace(get_hccl_comm_name=lambda rank: "cam-test") + process_group = SimpleNamespace(_get_backend=lambda device: backend) + fake_torch = _FakeTorch() + monkeypatch.setattr(async_cam_module, "torch", fake_torch) + monkeypatch.setattr( + async_cam_module, "ensure_cam_async_ops_available", lambda: None + ) - monkeypatch.setattr(async_cam_module.time, "sleep", finish_second_ffn) - ffn.wait_for_communication_warmup_done() - assert polls == [True] + def fake_init_afd_process_group(**kwargs): + kwargs["on_rendezvous"](store) + with pytest.raises(RuntimeError, match="not bound"): + context.borrow() + return process_group + monkeypatch.setattr( + async_cam_module, "init_afd_process_group", fake_init_afd_process_group + ) + monkeypatch.setattr( + async_cam_module, "create_hccl_process_group_options", lambda _: None + ) + monkeypatch.setattr( + async_cam_module.dist, + "destroy_process_group", + lambda group: None, + ) + connector = CAMAsyncAFDConnector( + 0, + 0, + _vllm_config(), + _afd_config(role="attention"), + 0, + rendezvous_context=context, + ) -def test_close_can_retain_graph_anchors_until_worker_stops(): - connector = object.__new__(CAMAsyncAFDConnector) - connector.cam_pg = None - connector._initialized = True - connector.comm_args = object() - connector._placeholder = object() - connector._startup_store = object() - connector._rendezvous_store = object() - connector._pending_attention_payloads = {} - comm_args = connector.comm_args - placeholder = connector._placeholder - - connector.close(release_buffers=False) - - assert connector.comm_args is comm_args - assert connector._placeholder is placeholder - connector.release_closed_buffers() - assert connector.comm_args is None - assert connector._placeholder is None + connector.init_afd_connector() + assert context.borrow() == (store, process_group) + connector.close() + with pytest.raises(RuntimeError, match="not bound"): + context.borrow() diff --git a/tests/unit/v1/worker/test_async_cam_startup.py b/tests/unit/v1/worker/test_async_cam_startup.py new file mode 100644 index 000000000..a2cb84b48 --- /dev/null +++ b/tests/unit/v1/worker/test_async_cam_startup.py @@ -0,0 +1,237 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Host-side Async CAM startup ordering and failure tests.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch_npu") + +from afd_plugin.connectors.npu.async_cam import AFDAsyncTopology # noqa: E402 +from afd_plugin.distributed.afd_process_group import ( # noqa: E402 + ProcessGroupRendezvousContext, +) +from afd_plugin.v1.worker.npu import async_cam_startup # noqa: E402 +from afd_plugin.v1.worker.npu.async_cam_startup import ( # noqa: E402 + AsyncCamStartupCoordinator, + AsyncCamStartupSpec, + FFNStartupPlan, +) + + +class FakeStore: + def __init__(self, values: dict[str, bytes] | None = None): + self.values = dict(values or {}) + + def set(self, key: str, value: str) -> None: + self.values[key] = value.encode() + + def check(self, keys: list[str]) -> bool: + return all(key in self.values for key in keys) + + def get(self, key: str) -> bytes: + return self.values[key] + + def compare_set(self, key: str, expected: str, desired: str) -> bytes: + current = self.values.get(key, b"") + if current == expected.encode(): + current = desired.encode() + self.values[key] = current + return current + + +def make_coordinator( + *, role: str, attn_size: int = 1, ffn_size: int = 1, store: FakeStore +) -> tuple[AsyncCamStartupCoordinator, SimpleNamespace]: + world_rank = 0 if role == "attention" else attn_size + topology = AFDAsyncTopology( + role=role, + role_rank=0, + world_rank=world_rank, + attn_size=attn_size, + ffn_size=ffn_size, + expert_per_rank=1, + ) + connector = SimpleNamespace(is_initialized=True) + context = ProcessGroupRendezvousContext() + context.retain_store(store) + context.bind(SimpleNamespace()) + coordinator = AsyncCamStartupCoordinator( + connector, + context, + AsyncCamStartupSpec( + topology=topology, + local_rank=0, + tp_size=1, + hidden_size=16, + topk=1, + activation_dtype=async_cam_startup.torch.bfloat16, + ), + ) + coordinator._store = store + return coordinator, connector + + +def test_mode_disagreement_prevents_receiver_and_capture(): + store = FakeStore({"ffn/mode/2": b"eager"}) + coordinator, _ = make_coordinator(role="ffn", attn_size=1, ffn_size=2, store=store) + calls = [] + with pytest.raises(RuntimeError, match="modes disagree"): + coordinator.start_ffn( + prepare=lambda: FFNStartupPlan(use_graph=True, first_layer_idx=3), + consume_warmup=lambda count: calls.append("warmup"), + capture=lambda: calls.append("capture"), + start_receiver=lambda: calls.append("receiver"), + ) + assert calls == [] + assert store.values["ffn/1"].startswith(b"failed:") + + +def test_all_attention_and_ffn_warmup_must_finish_before_capture(monkeypatch): + store = FakeStore( + { + "attn/prepared/0": b"ready", + "attn/done/0": b"ready", + "ffn/mode/2": b"graph:3", + } + ) + coordinator, _ = make_coordinator(role="ffn", attn_size=1, ffn_size=2, store=store) + events = [] + + def finish_peer(_: float) -> None: + assert events == ["prepare", "warmup:1"] + store.set("ffn/warmup/done/2", "ready") + + monkeypatch.setattr(async_cam_startup.time, "sleep", finish_peer) + coordinator.start_ffn( + prepare=lambda: ( + events.append("prepare") + or FFNStartupPlan(use_graph=True, first_layer_idx=3) + ), + consume_warmup=lambda count: events.append(f"warmup:{count}"), + capture=lambda: events.append("capture"), + start_receiver=lambda: events.append("receiver"), + ) + assert events == ["prepare", "warmup:1", "capture", "receiver"] + assert store.values["ffn/1"] == b"ready" + assert coordinator.started + + +@pytest.mark.parametrize("failure_phase", ["capture", "receiver"]) +def test_capture_or_receiver_failure_is_sticky(failure_phase): + store = FakeStore({"attn/prepared/0": b"ready", "attn/done/0": b"ready"}) + coordinator, _ = make_coordinator(role="ffn", store=store) + + def fail_if_selected(phase: str) -> None: + if phase == failure_phase: + raise RuntimeError(f"{phase} failed") + + with pytest.raises(RuntimeError, match=f"{failure_phase} failed"): + coordinator.start_ffn( + prepare=lambda: FFNStartupPlan(use_graph=True, first_layer_idx=3), + consume_warmup=lambda count: None, + capture=lambda: fail_if_selected("capture"), + start_receiver=lambda: fail_if_selected("receiver"), + ) + assert store.values["ffn/1"].startswith(b"failed:") + assert not coordinator.started + with pytest.raises(RuntimeError, match="previously failed"): + coordinator.start_ffn( + prepare=lambda: FFNStartupPlan(use_graph=True, first_layer_idx=3), + consume_warmup=lambda count: None, + capture=lambda: None, + start_receiver=lambda: None, + ) + + +def test_ready_cannot_replace_failed_status(): + class RacingStore(FakeStore): + def compare_set(self, key: str, expected: str, desired: str) -> bytes: + self.set(key, "failed:loop stopped") + return self.values[key] + + store = RacingStore() + coordinator, _ = make_coordinator(role="ffn", store=store) + with pytest.raises(RuntimeError, match="loop stopped"): + coordinator.start_ffn( + prepare=lambda: FFNStartupPlan(use_graph=False), + consume_warmup=lambda count: None, + capture=lambda: None, + start_receiver=lambda: None, + ) + assert store.values["ffn/1"].startswith(b"failed:") + + +def test_attention_rereads_ready_ranks_for_later_failure(monkeypatch): + store = FakeStore({"ffn/1": b"ready"}) + coordinator, _ = make_coordinator( + role="attention", attn_size=1, ffn_size=2, store=store + ) + + def fail_after_first_poll(_: float) -> None: + store.set("ffn/1", "failed:receiver stopped") + store.set("ffn/2", "ready") + + monkeypatch.setattr(async_cam_startup.time, "sleep", fail_after_first_poll) + with pytest.raises(RuntimeError, match="receiver stopped"): + coordinator._wait_for_ffn_ready() + + +def test_attention_start_does_not_return_before_ffn_ready(monkeypatch): + store = FakeStore({"ffn/mode/1": b"eager"}) + coordinator, _ = make_coordinator(role="attention", store=store) + polls = [] + + def publish_after_wait(_: float) -> None: + assert not coordinator.started + polls.append(True) + store.set("ffn/1", "ready") + + monkeypatch.setattr(async_cam_startup.time, "sleep", publish_after_wait) + coordinator.start_attention() + assert polls == [True] + assert coordinator.started + + +def test_prepare_failure_before_store_does_not_start_nonce_collective(monkeypatch): + store = FakeStore() + coordinator, connector = make_coordinator(role="ffn", store=store) + connector.is_initialized = False + connector.init_afd_connector = lambda: pytest.fail("group was initialized") + coordinator._store = None + monkeypatch.setattr( + coordinator, + "_get_store", + lambda: pytest.fail("nonce collective was started"), + ) + with pytest.raises(RuntimeError, match="weights unavailable"): + coordinator.start_ffn( + prepare=lambda: (_ for _ in ()).throw(RuntimeError("weights unavailable")), + consume_warmup=lambda count: None, + capture=lambda: None, + start_receiver=lambda: None, + ) + assert coordinator.failed + assert store.values == {} + + +def test_repeated_eager_start_is_idempotent_and_closed_context_is_rejected(): + store = FakeStore() + coordinator, connector = make_coordinator(role="ffn", store=store) + events = [] + connector.init_afd_connector = lambda: events.append("init") + callbacks = dict( + prepare=lambda: events.append("prepare") or FFNStartupPlan(False), + consume_warmup=lambda count: events.append("warmup"), + capture=lambda: events.append("capture"), + start_receiver=lambda: events.append("receiver"), + ) + coordinator.start_ffn(**callbacks) + coordinator.start_ffn(**callbacks) + assert events == ["prepare", "receiver"] + coordinator._rendezvous_context.invalidate() + with pytest.raises(RuntimeError, match="not bound"): + coordinator.start_ffn(**callbacks) diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index e3521f17f..723e43fa9 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -450,6 +450,7 @@ def _new_ffn_runner(): runner._async_cam_ffn_graph = None runner._async_cam_ffn_output = None runner._async_cam_ffn_replays = 0 + runner._async_cam_startup = None runner.prof = None runner.device = SimpleNamespace(type="npu") runner._is_shutdown = False @@ -472,6 +473,7 @@ def _new_ffn_worker(): worker = object.__new__(AFDNPUFFNWorker) worker._ffn_loop_error = None + worker._ffn_receiver_drained = False # Most tests exercise the daemon loop rather than CPU placement. Tests for # the startup binding path explicitly reset this guard. worker._cpu_binding_attempted = True @@ -1965,6 +1967,23 @@ def test_npu_ffn_runner_shutdown_is_idempotent(monkeypatch): assert parent_calls == [runner] +def test_npu_cam_eager_runner_rejects_direct_close_with_live_connector(monkeypatch): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import ffn_model_runner + + runner = _new_ffn_runner() + closes = [] + runner._async_cam_startup = SimpleNamespace(started=True, failed=False) + runner.connector = SimpleNamespace( + is_initialized=True, close=lambda: closes.append(True) + ) + monkeypatch.setattr(ffn_model_runner, "stop_afd_npu_profiler", lambda _: None) + + with pytest.raises(RuntimeError, match="receiver must drain"): + runner.shutdown() + assert closes == [] + + def test_npu_ffn_worker_scheduler_execute_model_fails_fast(): worker = _new_ffn_worker() @@ -1995,6 +2014,7 @@ def test_npu_ffn_worker_start_binds_physical_npu_once_before_daemon( worker.local_rank = 3 worker.model_runner = SimpleNamespace( connector=SimpleNamespace(is_initialized=True), + _async_cam_startup=None, ) events: list[tuple[str, object]] = [] monkeypatch.setattr( @@ -2060,6 +2080,7 @@ def test_npu_ffn_worker_cpu_binding_failure_does_not_abort_daemon_start( worker.local_rank = 5 worker.model_runner = SimpleNamespace( connector=SimpleNamespace(is_initialized=True), + _async_cam_startup=None, ) thread_starts: list[bool] = [] monkeypatch.setattr( @@ -2105,6 +2126,7 @@ def test_npu_ffn_worker_loop_error_is_propagated(caplog): worker._ffn_loop_error = None worker.model_runner = SimpleNamespace( connector=SimpleNamespace(is_initialized=True), + _async_cam_startup=None, ) expected_error = RuntimeError("boom") @@ -2136,6 +2158,7 @@ def test_npu_ffn_worker_ignores_receive_error_during_shutdown(caplog): worker._ffn_loop_error = None worker.model_runner = SimpleNamespace( connector=SimpleNamespace(is_initialized=True), + _async_cam_startup=None, ) def stop_while_receiving(): @@ -2173,7 +2196,7 @@ def is_alive(self): connector = SimpleNamespace(close=lambda: calls.append(("close", None))) worker._ffn_thread = _StoppingThread() - worker.model_runner = SimpleNamespace(connector=connector) + worker.model_runner = SimpleNamespace(connector=connector, _async_cam_startup=None) monkeypatch.setattr( ffn_worker.NPUWorker, "shutdown", @@ -2211,6 +2234,7 @@ def is_alive(self): worker._ffn_thread = thread worker.model_runner = SimpleNamespace( connector=SimpleNamespace(close=lambda: calls.append(("close", None))), + _async_cam_startup=None, ) monkeypatch.setattr( ffn_worker.NPUWorker, @@ -2833,6 +2857,7 @@ def test_npu_attention_runner_load_model_initializes_connector_after_weights( connector = _LifecycleConnector(events) runner = object.__new__(attention_model_runner.AFDNPUAttentionModelRunner) runner.connector = connector + runner._async_cam_startup = None runner.vllm_config = SimpleNamespace( parallel_config=SimpleNamespace(use_ubatching=use_ubatching), ) @@ -2963,73 +2988,37 @@ def make_graph(): assert runner._async_cam_ffn_output is None -def test_npu_async_cam_ffn_startup_orders_warmup_capture_and_ready(monkeypatch): +def test_npu_async_cam_ffn_worker_delegates_one_complete_startup(): _require_npu_runtime() - from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector - from afd_plugin.v1.worker.npu import ffn_worker - events: list[str | tuple[str, int | None]] = [] - connector = object.__new__(CAMAsyncAFDConnector) - connector._initialized = False - connector._startup_store = object() - connector.world_rank = 8 - connector.prepare_startup_coordination = lambda: events.append("coordination") - connector.publish_ffn_mode = lambda layer_idx: events.append(("mode", layer_idx)) - connector.wait_for_attention_warmup_prepared = lambda: events.append("prepared") - connector.publish_ffn_warmup_started = lambda: events.append("warmup start") - connector.publish_ffn_warmup_done = lambda: events.append("warmup done") - connector.wait_for_communication_warmup_done = lambda: events.append("all done") - connector.publish_ffn_startup = lambda **_: events.append("ready") + events = [] worker = _new_ffn_worker() worker._ffn_thread = None - worker._ffn_shutdown_event = None - worker._bind_cpus_once = lambda: events.append("bind") + worker._start_ffn_receiver = lambda: events.append("receiver") - def initialize_connector(): - events.append("connector") - connector._initialized = True + class FakeStartup: + failed = False + + def start_ffn(self, *, prepare, consume_warmup, capture, start_receiver): + events.append("coordinator") + plan = prepare() + assert plan.use_graph and plan.first_layer_idx == 3 + consume_warmup(2) + capture() + start_receiver() worker.model_runner = SimpleNamespace( - connector=connector, - _async_cam_ffn_graph_enabled=True, - _layered_executor=SimpleNamespace(layer_ids=(0,)), - prepare_async_cam_ffn_graph=lambda: events.append("prepare"), - warmup_async_cam_ffn_communication=lambda: events.append("communication"), - initialize_afd_connector=initialize_connector, + _async_cam_startup=FakeStartup(), + prepare_async_cam_ffn_startup=lambda: SimpleNamespace( + use_graph=True, first_layer_idx=3 + ), + warmup_async_cam_ffn_communication=lambda count: events.append( + ("warmup", count) + ), capture_async_cam_ffn_graph=lambda: events.append("capture"), ) - - class FakeThread: - def __init__(self, *, target, name, daemon): - events.append("thread created") - - def start(self): - events.append("thread started") - worker._ffn_loop_started_event.set() - - def is_alive(self): - return True - - monkeypatch.setattr(ffn_worker.threading, "Thread", FakeThread) - worker.start_ffn_server_loop() - - assert events == [ - "prepare", - "connector", - "coordination", - ("mode", 0), - "prepared", - "warmup start", - "communication", - "warmup done", - "all done", - "capture", - "bind", - "thread created", - "thread started", - "ready", - ] + assert events == ["coordinator", ("warmup", 2), "capture", "receiver"] def test_npu_async_cam_graph_shutdown_releases_after_thread_stops(): @@ -3039,11 +3028,11 @@ def test_npu_async_cam_graph_shutdown_releases_after_thread_stops(): events = [] connector = object.__new__(CAMAsyncAFDConnector) connector.world_rank = 8 - connector.close = lambda **_: events.append("close group") - connector.release_closed_buffers = lambda: events.append("release anchors") + connector.close = lambda: events.append("close group and anchors") worker = _new_ffn_worker() worker._ffn_shutdown_event = threading.Event() worker._ffn_loop_started_event = threading.Event() + worker._ffn_receiver_drained = True class FakeThread: active = True @@ -3059,6 +3048,7 @@ def join(self, *, timeout): worker.model_runner = SimpleNamespace( connector=connector, _async_cam_ffn_graph_enabled=True, + _async_cam_startup=SimpleNamespace(failed=False, started=True), release_async_cam_ffn_graph=lambda: events.append("release graph"), ) @@ -3067,8 +3057,7 @@ def join(self, *, timeout): assert events == [ "join", "release graph", - "close group", - "release anchors", + "close group and anchors", ] @@ -3079,8 +3068,7 @@ def test_npu_async_cam_graph_shutdown_retains_anchors_if_join_times_out(): events = [] connector = object.__new__(CAMAsyncAFDConnector) connector.world_rank = 8 - connector.close = lambda **_: events.append("close group") - connector.release_closed_buffers = lambda: events.append("release anchors") + connector.close = lambda: events.append("close group and anchors") worker = _new_ffn_worker() worker._ffn_shutdown_event = threading.Event() @@ -3095,6 +3083,7 @@ def join(self, *, timeout): worker.model_runner = SimpleNamespace( connector=connector, _async_cam_ffn_graph_enabled=True, + _async_cam_startup=SimpleNamespace(failed=False, started=True), release_async_cam_ffn_graph=lambda: events.append("release graph"), ) From 412f69d37a5ef38c07ce0994b2a4d7a8c756fb08 Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 16:46:12 +0800 Subject: [PATCH 7/8] test(npu): align runner fixtures with startup assembly Signed-off-by: jiangkuaixue123 --- tests/unit/v1/worker/test_npu_runtime.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 723e43fa9..e74897c6f 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -2779,7 +2779,10 @@ def fake_native_init(self, vllm_config, device): module.AFDNPUFFNModelRunner, "parse_config", staticmethod( - lambda config: SimpleNamespace(compute_gate_on_attention=compute_gate) + lambda config: SimpleNamespace( + compute_gate_on_attention=compute_gate, + connector="P2pNcclAFDConnector", + ) ), ) @@ -2895,6 +2898,7 @@ def test_npu_attention_runner_afd_ubatching_does_not_install_native_wrapper( connector = _LifecycleConnector(events) runner = object.__new__(attention_model_runner.AFDNPUAttentionModelRunner) runner.connector = connector + runner._async_cam_startup = None runner.vllm_config = SimpleNamespace( parallel_config=SimpleNamespace( use_ubatching=False, From 4e86c7182e570c650f98d6c6b45910fe084f99eb Mon Sep 17 00:00:00 2001 From: jiangkuaixue123 Date: Wed, 30 Sep 2026 16:46:54 +0800 Subject: [PATCH 8/8] chore(npu): log completed AsyncCAM startup readiness Signed-off-by: jiangkuaixue123 --- afd_plugin/v1/worker/npu/async_cam_startup.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/afd_plugin/v1/worker/npu/async_cam_startup.py b/afd_plugin/v1/worker/npu/async_cam_startup.py index 8988f3fd6..42631cd16 100644 --- a/afd_plugin/v1/worker/npu/async_cam_startup.py +++ b/afd_plugin/v1/worker/npu/async_cam_startup.py @@ -317,6 +317,11 @@ def start_attention(self) -> None: self._run_attention_warmup(layer_idx) self._wait_for_ffn_ready() self._started = True + logger.info( + "CAM Attention startup all FFN ready rank=%d mode=%s", + self._spec.topology.world_rank, + "eager" if layer_idx is None else f"graph:{layer_idx}", + ) except Exception as exc: self.report_failure(exc) raise @@ -367,6 +372,9 @@ def start_ffn( if status != "ready": raise RuntimeError(f"CAM FFN startup {key} {status}") self._started = True + logger.info( + "CAM FFN startup ready rank=%d mode=%s", topology.world_rank, mode + ) except Exception as exc: self.report_failure(exc) raise