[WIP][NPU]Add DSV4 NPU ubatch async connector

afd_plugin/__init__.py

@@ -6,6 +6,8 @@
 
 import importlib.util
 import logging
+import multiprocessing
+import os
 from importlib.metadata import PackageNotFoundError, version
 from pathlib import Path
 from types import MappingProxyType
@@ -45,6 +47,51 @@ def __getattr__(name: str):
 _logger = logging.getLogger(__name__)
 _registered = False
 
+
+def _force_spawn_multiprocessing_if_requested() -> None:
+    """Pin Python's default context for the A3 Python 3.12 runtime.
+
+    ``VLLM_WORKER_MULTIPROC_METHOD`` controls vLLM's explicit context only.
+    Some NPU runtime helpers use Python's default context and otherwise retain
+    the task's inherited ``forkserver`` setting, which cannot restore the
+    signal-handler state in this environment.  Keep the override opt-in so it
+    remains isolated to the affected CAM async recipe.
+    """
+    if os.environ.get("AFD_FORCE_SPAWN_MULTIPROCESSING") != "1":
+        return
+    multiprocessing.set_start_method("spawn", force=True)
+
+    # A few optional runtime helpers request ``forkserver`` explicitly instead
+    # of consulting the default start method.  On the A3 Python 3.12 image the
+    # forkserver cannot restore its inherited signal-handler state, so every
+    # child it starts exits before running user code.  Keep this narrowly
+    # opt-in with the recipe environment variable above: callers requesting a
+    # forkserver receive the equivalent spawn context instead.
+    contexts = multiprocessing.context._concrete_contexts
+    contexts["forkserver"] = contexts["spawn"]
+
+    # TE Fusion keeps a module reference to ``multiprocessing`` and calls its
+    # public ``get_context("forkserver")`` API directly inside each model
+    # worker.  Replacing the registry alone is not sufficient for all Python
+    # 3.12 context instances, so redirect that explicit request as well.
+    if not getattr(multiprocessing, "_afd_spawn_context_redirect", False):
+        original_get_context = multiprocessing.get_context
+
+        def get_spawn_context(method: str | None = None):
+            if method == "forkserver":
+                method = "spawn"
+            return original_get_context(method)
+
+        multiprocessing.get_context = get_spawn_context
+        multiprocessing._afd_spawn_context_redirect = True
+
+
+# vLLM model workers import the configured worker class directly in spawned
+# child interpreters; they do not invoke the general-plugin entry point below.
+# Apply the opt-in setting at package import time so it is also in effect before
+# Ascend TE Fusion initializes its compilation workers.
+_force_spawn_multiprocessing_if_requested()
+
 _DEEPSEEK_MODEL_REGISTRATIONS = {
     "DeepseekForCausalLM": (
         "afd_plugin.model_executor.models.deepseek_v2:AFDDeepseekForCausalLM"
@@ -59,7 +106,9 @@ def __getattr__(name: str):
         "afd_plugin.model_executor.models.deepseek_v2:AFDDeepseekV3ForCausalLM"
     ),
     "DeepseekV4ForCausalLM": (
-        "afd_plugin.model_executor.models.deepseek_v4:AFDDeepseekV4ForCausalLM"
+        "afd_plugin.model_executor.models.npu.deepseek_v4:AFDDeepseekV4ForCausalLM"
+        if importlib.util.find_spec("torch_npu") is not None
+        else "afd_plugin.model_executor.models.deepseek_v4:AFDDeepseekV4ForCausalLM"
     ),
     "GlmMoeDsaForCausalLM": (
         "afd_plugin.model_executor.models.deepseek_v2:AFDGlmMoeDsaForCausalLM"
@@ -102,6 +151,7 @@ def register_afd() -> None:
         _logger.debug("AFD plugin: register_afd() already completed")
         return
 
+    _force_spawn_multiprocessing_if_requested()
     _logger.debug("AFD plugin: register_afd() called")
     if importlib.util.find_spec("vllm") is None:
         _logger.debug("AFD plugin: vLLM not found, skipping runtime registration")

afd_plugin/compat/npu/feature_validation.py

@@ -40,6 +40,8 @@ def fail_if_unsupported_npu_afd_features(
             afd_config,
             extra_info,
         )
+        if _is_dsv4_target(vllm_config):
+            _fail_if_unsupported_dsv4_async_features(afd_config, extra_info)
         return
 
     if afd_config.compute_gate_on_attention:
@@ -84,6 +86,45 @@ def fail_if_unsupported_npu_afd_features(
         )
 
 
+def _is_dsv4_target(vllm_config: VllmConfig) -> bool:
+    """Return whether the target model is DeepSeek V4."""
+    model_config = vllm_config.model_config
+    hf_config = getattr(model_config, "hf_config", None)
+    if hf_config is None:
+        hf_config = getattr(model_config, "hf_text_config", None)
+    if hf_config is None:
+        return False
+    if getattr(hf_config, "model_type", None) == "deepseek_v4":
+        return True
+    architectures = getattr(hf_config, "architectures", ()) or ()
+    return any("DeepseekV4" in str(architecture) for architecture in architectures)
+
+
+def _fail_if_unsupported_dsv4_async_features(
+    afd_config: AFDConfig,
+    extra_info: ConnectorExtraInfo,
+) -> None:
+    """Keep the initial DSV4 Async CAM target path deliberately narrow."""
+    from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo
+
+    if not afd_config.compute_gate_on_attention:
+        raise RuntimeError(
+            "DSV4 CAMAsyncAFDConnector requires compute_gate_on_attention=true",
+        )
+    if not isinstance(extra_info, AFDAsyncExtraInfo):
+        raise TypeError(
+            "DSV4 CAMAsyncAFDConnector requires AFDAsyncExtraInfo, got "
+            f"{type(extra_info).__name__}",
+        )
+    if extra_info.dynamic_quant != 1:
+        raise RuntimeError(
+            "DSV4 Flash-INT8 CAMAsyncAFDConnector requires dynamicQuant=1",
+        )
+    # async_moe_ubatching is validated by
+    # _fail_if_unsupported_npu_async_moe_ubatching_features for all async CAM
+    # targets, so DSV4 no longer needs a bespoke rejection here.
+
+
 def _fail_if_unsupported_npu_afd_async_features(
     vllm_config: VllmConfig,
     afd_config: AFDConfig,
@@ -123,6 +164,43 @@ def _fail_if_unsupported_npu_afd_async_features(
         raise RuntimeError(
             "CAMAsyncAFDConnector currently supports only dynamicQuant 0 or 1",
         )
+    _validate_cam_world_topology(vllm_config, afd_config, extra_info)
+
+
+def _validate_cam_world_topology(
+    vllm_config: VllmConfig,
+    afd_config: AFDConfig,
+    extra_info: ConnectorExtraInfo,
+) -> None:
+    """Require each role's local layout to fill the one CAM world."""
+    from afd_plugin.connectors.npu.async_cam import AFDAsyncExtraInfo
+
+    if not isinstance(extra_info, AFDAsyncExtraInfo):
+        return
+    parallel_config = vllm_config.parallel_config
+    attn_ranks_per_dp = int(extra_info.attn_ranks_per_dp)
+    if afd_config.role == "attention":
+        if int(parallel_config.tensor_parallel_size) != attn_ranks_per_dp:
+            raise RuntimeError(
+                "CAMAsyncAFDConnector Attention tensor_parallel_size must equal "
+                "attn_ranks_per_dp",
+            )
+        local_world_size = (
+            int(parallel_config.data_parallel_size) * attn_ranks_per_dp
+        )
+        expected_world_size = afd_config.num_attention_ranks
+    else:
+        local_world_size = (
+            int(parallel_config.data_parallel_size)
+            * int(parallel_config.tensor_parallel_size)
+        )
+        expected_world_size = afd_config.num_ffn_ranks
+    if local_world_size != expected_world_size:
+        raise RuntimeError(
+            "CAMAsyncAFDConnector "
+            f"{afd_config.role} DPxTP world size must equal its configured "
+            f"role size, got {local_world_size} and {expected_world_size}",
+        )
 
 
 def _fail_if_unsupported_npu_async_moe_ubatching_features(

afd_plugin/compat/npu/ops.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import ctypes
 import os
 from functools import lru_cache
 from pathlib import Path
@@ -16,6 +17,10 @@
 CAM_DISPATCH_RECV = "async_dispatch_recv"
 CAM_COMBINE_SEND = "async_combine_send"
 CAM_COMBINE_RECV = "async_combine_recv"
+CAM_CUST_OPAPI_ENV = "CAM_CUST_OPAPI_LIB_PATH"
+CAM_CUST_OPAPI_DEFAULT = Path(
+    "/usr/local/Ascend/cann-9.0.1/opp/vendors/CAM/op_api/lib/libopapi.so",
+)
 
 
 def get_afd_cann_vendor_path() -> Path:
@@ -68,6 +73,24 @@ def _assert_cam_namespace_registered(torch: object) -> None:
     )
 
 
+def _preload_cam_cust_opapi() -> None:
+    """Make CAM's custom op-api symbols visible to the Python extension.
+
+    ``umdk_cam_op_lib`` resolves the aclnn CAM entry points by name while its
+    operators are first invoked.  CANN's stock ``libopapi.so`` does not export
+    those symbols; they live in CAM's ``libopapi.so``.  Load that library
+    globally before importing the extension so the real CAM operators bind to
+    the vendor implementation instead of the stock op-api library.
+    """
+
+    library = Path(os.environ.get(CAM_CUST_OPAPI_ENV, CAM_CUST_OPAPI_DEFAULT))
+    if not library.is_file():
+        raise RuntimeError(
+            "CAMAsyncAFDConnector requires CAM libopapi.so; expected "
+            f"{library}. Install the CAM operator package first.",
+        )
+    ctypes.CDLL(str(library), mode=ctypes.RTLD_GLOBAL)
+
 @lru_cache(maxsize=1)
 def ensure_cam_p2p_ops_available() -> None:
     """Import the custom operators used by ``CAMP2pAFDConnector``.
@@ -107,6 +130,7 @@ def ensure_cam_async_ops_available() -> None:
     """Ensure the runtime exposes the real CAM async operator namespace."""
 
     try:
+        _preload_cam_cust_opapi()
         import torch
         import torch_npu  # noqa: F401
         import umdk_cam_op_lib  # noqa: F401

afd_plugin/compat/npu/runtime.py

@@ -69,7 +69,6 @@ def apply_afd_ascend_patches_if_needed() -> None:
     from afd_plugin.compat.patches.npu.mla_graph import (
         apply_afd_mla_graph_patch,
     )
-
     apply_afd_ascend_config_patch_if_needed()
     if not apply_afd_mla_graph_patch():
         raise RuntimeError(

afd_plugin/compat/patches/engine_core.py

@@ -240,6 +240,8 @@ def _initialize_kv_caches(self, vllm_config: VllmConfig) -> KVCacheConfig:
         return _AFDFFNKVCacheConfig()
     # ### PATCH END: AFD FFN late-loaded KV cache bypass
 
+    import vllm_ascend.patch.platform.patch_kv_cache_utils  # noqa: F401
+
     start = time.time()
 
     core_module.register_all_kvcache_specs(vllm_config)

afd_plugin/config.py

@@ -328,7 +328,6 @@ def validate_afd_config(
             f"num_ffn_ranks must be positive, got {config.num_ffn_ranks}",
         )
 
-
 __all__ = [
     "AFDConfig",
     "AFD_ASYNC_CONNECTOR",

afd_plugin/connectors/metadata.py

@@ -128,6 +128,7 @@ class AFDControlPayload:
     dp_metadata_list: dict[int, AFDDPMetadata]
     is_graph_capturing: bool
     is_warmup: bool
+    input_ids_by_stage: dict[int, list[int]] = field(default_factory=dict)
 
     def __post_init__(self) -> None:
         self.dp_metadata_list = {
@@ -139,6 +140,14 @@ def __post_init__(self) -> None:
             else dp_metadata
             for stage_idx, dp_metadata in self.dp_metadata_list.items()
         }
+        self.input_ids_by_stage = {
+            int(stage_idx): torch.as_tensor(
+                input_ids,
+                dtype=torch.int64,
+                device="cpu",
+            ).flatten().tolist()
+            for stage_idx, input_ids in self.input_ids_by_stage.items()
+        }
 
 
 @dataclass(slots=True)
@@ -319,10 +328,14 @@ def encode_control_payload(payload: AFDControlPayload) -> bytes:
     """
     metadata_payload: dict[str, dict[str, int | list[int]]] = {}
     for stage_idx, dp_metadata in payload.dp_metadata_list.items():
-        metadata_payload[str(int(stage_idx))] = {
+        stage_payload: dict[str, int | list[int]] = {
             "num_tokens_across_dp_cpu": dp_metadata.num_tokens_across_dp_cpu.tolist(),
             "max_tokens_across_dp_cpu": _to_int(dp_metadata.max_tokens_across_dp_cpu),
         }
+        input_ids = payload.input_ids_by_stage.get(int(stage_idx))
+        if input_ids is not None:
+            stage_payload["input_ids"] = list(input_ids)
+        metadata_payload[str(int(stage_idx))] = stage_payload
 
     wire_payload = {
         "dp_metadata_list": metadata_payload,
@@ -352,10 +365,16 @@ def decode_control_payload(payload_bytes: bytes) -> AFDControlPayload:
         )
         for stage_idx, metadata in payload["dp_metadata_list"].items()
     }
+    input_ids_by_stage = {
+        int(stage_idx): [int(value) for value in metadata["input_ids"]]
+        for stage_idx, metadata in payload["dp_metadata_list"].items()
+        if "input_ids" in metadata
+    }
     return AFDControlPayload(
         dp_metadata_list=dp_metadata_list,
         is_graph_capturing=bool(payload.get("is_graph_capturing", False)),
         is_warmup=bool(payload.get("is_warmup", False)),
+        input_ids_by_stage=input_ids_by_stage,
     )