feat(npu): support token-balanced Async CAM MoE ubatching

afd_plugin/async_moe.py

@@ -0,0 +1,44 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Backend-independent Async CAM MoE stage contracts."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+ASYNC_MOE_NUM_STAGES = 2
+ASYNC_MOE_REQUEST_SPLIT = "request"
+ASYNC_MOE_TOKEN_SPLIT = "token"
+
+
+@dataclass(frozen=True)
+class AsyncMoeStage:
+    """One ordered stage in the flattened Attention token layout.
+
+    ``token_slice`` describes the stage's ordered real-token range in the
+    parent batch. ``input_tokens`` includes only the minimum stage-local
+    padding required by the Attention TP/SP layout.
+    """
+
+    request_slice: slice
+    token_slice: slice
+    input_tokens: int
+
+    @property
+    def num_tokens(self) -> int:
+        return self.input_tokens
+
+    @property
+    def actual_tokens(self) -> int:
+        return int(self.token_slice.stop) - int(self.token_slice.start)
+
+    def is_empty(self) -> bool:
+        return self.actual_tokens <= 0 or self.input_tokens < self.actual_tokens
+
+
+__all__ = [
+    "ASYNC_MOE_NUM_STAGES",
+    "ASYNC_MOE_REQUEST_SPLIT",
+    "ASYNC_MOE_TOKEN_SPLIT",
+    "AsyncMoeStage",
+]

afd_plugin/compat/npu/feature_validation.py

@@ -87,7 +87,7 @@ def _fail_if_unsupported_npu_afd_async_features(
         raise RuntimeError(
             "CAMAsyncAFDConnector supports only eager Attention/FFN execution",
         )
-    if bool(parallel_config.use_ubatching):
+    if bool(parallel_config.enable_dbo) or bool(parallel_config.use_ubatching):
         raise RuntimeError(
             "CAMAsyncAFDConnector does not support vLLM native ubatching/DBO",
         )
@@ -111,27 +111,43 @@ def _fail_if_unsupported_npu_async_moe_ubatching_features(
     num_ubatches: int,
     split: str,
 ) -> None:
-    from afd_plugin.connectors.npu.async_cam import ASYNC_MOE_REQUEST_SPLIT
+    from afd_plugin.async_moe import (
+        ASYNC_MOE_NUM_STAGES,
+        ASYNC_MOE_REQUEST_SPLIT,
+        ASYNC_MOE_TOKEN_SPLIT,
+    )
 
     parallel_config = vllm_config.parallel_config
     if not afd_config.compute_gate_on_attention:
         raise RuntimeError(
             "async_moe_ubatching requires compute_gate_on_attention=true",
         )
-    if num_ubatches != 2:
+    if num_ubatches != ASYNC_MOE_NUM_STAGES:
         raise RuntimeError(
             "async_moe_ubatching currently supports exactly two stages; "
             f"got async_moe_num_ubatches={num_ubatches}",
         )
-    if split != ASYNC_MOE_REQUEST_SPLIT:
+    if split not in (ASYNC_MOE_REQUEST_SPLIT, ASYNC_MOE_TOKEN_SPLIT):
         raise RuntimeError(
-            "async_moe_ubatching currently supports only request-boundary split; "
+            "async_moe_split must be 'request' or 'token'; "
             f"got async_moe_split={split!r}",
         )
+    # Attention owns stage planning and SP layout conversion. FFN workers
+    # consume CAM work items and may use an independent TP/EP topology.
+    if afd_config.is_ffn_server:
+        return
     if int(parallel_config.decode_context_parallel_size) > 1:
         raise RuntimeError(
             "async_moe_ubatching does not support decode context parallel metadata yet",
         )
+    if split == ASYNC_MOE_TOKEN_SPLIT and (
+        int(parallel_config.tensor_parallel_size) <= 1
+        or int(parallel_config.prefill_context_parallel_size) > 1
+    ):
+        raise RuntimeError(
+            "async_moe_split='token' requires a non-PCP Attention DP+TP/SP "
+            "topology with tensor_parallel_size > 1",
+        )
 
 
 __all__ = ["fail_if_unsupported_npu_afd_features"]

afd_plugin/connectors/npu/async_cam.py

@@ -18,9 +18,10 @@
 The supported deployment requires ``async=true``, eager execution, Ascend CAM
 operator packages, and matching topology/configuration on every rank. vLLM
 native DBO, ACL graph execution, and decode are not supported.
-Optional AFD-managed MoE ubatching is a separate two-stage request-boundary
-pipeline. See ``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration,
-rank derivation, launch guidance, and the full limitations.
+Optional AFD-managed MoE ubatching is a separate two-stage pipeline using
+request boundaries for PCP or token-balanced stages for non-PCP DP+TP/SP. See
+``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration, rank
+derivation, launch guidance, and the full limitations.
 """
 
 from __future__ import annotations
@@ -36,6 +37,11 @@
 from torch import Tensor
 from vllm.logger import init_logger
 
+from afd_plugin.async_moe import (
+    ASYNC_MOE_NUM_STAGES,
+    ASYNC_MOE_REQUEST_SPLIT,
+    ASYNC_MOE_TOKEN_SPLIT,
+)
 from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available
 from afd_plugin.config import AFDConfig
 from afd_plugin.config_utils import (
@@ -64,7 +70,6 @@
 AFD_ASYNC_CAM_GROUP_NAME = "afd_async_cam"
 CAM_COMM_ID = 0
 ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp"
-ASYNC_MOE_REQUEST_SPLIT = "request"
 
 _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset(
     {
@@ -86,15 +91,16 @@ class AFDAsyncExtraInfo(ConnectorExtraInfo):
     Attributes:
         dynamic_quant: Dynamic quantization mode accepted by CAM operators.
         attn_ranks_per_dp: Number of Attention ranks in each data-parallel group.
-        async_moe_ubatching: Whether request-boundary async MoE ubatching is used.
+        async_moe_ubatching: Whether two-stage async MoE ubatching is used.
         async_moe_num_ubatches: Number of stages used by async MoE ubatching.
-        async_moe_split: Boundary at which async MoE work is split.
+        async_moe_split: ``"request"`` for request boundaries or ``"token"``
+            for token-balanced non-PCP DP+TP/SP stages.
     """
 
     dynamic_quant: int = 0
     attn_ranks_per_dp: int = 1
     async_moe_ubatching: bool = False
-    async_moe_num_ubatches: int = 2
+    async_moe_num_ubatches: int = ASYNC_MOE_NUM_STAGES
     async_moe_split: str = ASYNC_MOE_REQUEST_SPLIT
 
     @classmethod
@@ -129,7 +135,7 @@ def from_mapping(cls, raw: Mapping[str, Any] | None) -> AFDAsyncExtraInfo:
                 field_name="async_moe_ubatching",
             ),
             async_moe_num_ubatches=coerce_extra_positive_int(
-                raw.get("async_moe_num_ubatches", 2),
+                raw.get("async_moe_num_ubatches", ASYNC_MOE_NUM_STAGES),
                 field_name="async_moe_num_ubatches",
             ),
             async_moe_split=coerce_extra_str(
@@ -927,6 +933,9 @@ def _validate_topk_payload(
     "AFDAsyncFFNWorkItem",
     "AFDAsyncTopology",
     "ATTN_RANKS_PER_DP_CONFIG_KEY",
+    "ASYNC_MOE_NUM_STAGES",
+    "ASYNC_MOE_REQUEST_SPLIT",
+    "ASYNC_MOE_TOKEN_SPLIT",
     "CAM_COMM_ID",
     "build_async_topology",
 ]

afd_plugin/model_executor/models/forward_context.py

@@ -6,21 +6,81 @@
 
 from collections.abc import Iterator
 from contextlib import contextmanager
+from dataclasses import dataclass
 from functools import wraps
-from typing import Any, Final, TypedDict
+from typing import Any, Final
 
 import vllm.forward_context as forward_context_module
 from vllm.forward_context import ForwardContext, get_forward_context
-from vllm.v1.worker.ubatch_utils import UBatchSlices
 
+from afd_plugin.async_moe import AsyncMoeStage
 from afd_plugin.connectors import AFDForwardContextMetadata
 
 ASYNC_MOE_UBATCH_METADATA_KEY: Final[str] = "afd_async_moe_ubatch_metadata"
 
 
-class AsyncMoeUbatchMetadata(TypedDict):
-    attn_metadata: object
-    ubatch_slices: UBatchSlices
+@dataclass(frozen=True)
+class AsyncMoeUbatchMetadata:
+    """Immutable execution plan for one two-stage Async CAM model forward.
+
+    Each stage's token slice uses the parent real-token coordinate space, while
+    ``stage.num_tokens`` includes its minimum physical padding. Under sequence
+    parallelism each physical stage is divided evenly across TP ranks by the
+    model-side layout helper. ``parent_input_tokens`` records the original
+    padded layout restored after both stages.
+    """
+
+    attn_metadata: list[dict[str, object]]
+    ubatch_slices: tuple[AsyncMoeStage, ...]
+    parent_input_tokens: int
+    use_sequence_parallel: bool
+
+    @property
+    def stage_actual_token_counts(self) -> tuple[int, ...]:
+        return tuple(stage.actual_tokens for stage in self.ubatch_slices)
+
+    @property
+    def num_parent_input_tokens(self) -> int:
+        return self.parent_input_tokens
+
+    def __post_init__(self) -> None:
+        num_stages = len(self.ubatch_slices)
+        if not num_stages or len(self.attn_metadata) != num_stages:
+            raise ValueError(
+                "Async CAM execution-plan fields must describe the same "
+                f"non-empty stage count: attention={len(self.attn_metadata)}, "
+                f"slices={num_stages}",
+            )
+
+        expected_token_start = 0
+        for stage_slice, actual_tokens in zip(
+            self.ubatch_slices,
+            self.stage_actual_token_counts,
+            strict=True,
+        ):
+            token_slice = stage_slice.token_slice
+            token_start = int(token_slice.start)
+            token_stop = int(token_slice.stop)
+            actual_token_extent = token_stop - token_start
+            input_tokens = int(stage_slice.num_tokens)
+            if token_start != expected_token_start or actual_token_extent <= 0:
+                raise ValueError(
+                    "Async CAM stage token slices must be contiguous, ordered, "
+                    f"and non-empty: token_slice={token_slice}, "
+                    f"expected_start={expected_token_start}",
+                )
+            if not 0 < int(actual_tokens) <= input_tokens:
+                raise ValueError(
+                    "Async CAM stage actual-token count must fit its physical "
+                    f"extent: actual={actual_tokens}, input={input_tokens}",
+                )
+            expected_token_start = token_stop
+        if self.num_parent_input_tokens < expected_token_start:
+            raise ValueError(
+                "Async CAM parent input extent must cover every real token: "
+                f"parent_input_tokens={self.num_parent_input_tokens}, "
+                f"actual_tokens={expected_token_start}",
+            )
 
 
 def get_afd_metadata_from_forward_context(
@@ -90,6 +150,7 @@ def create_forward_context_with_afd(*args: Any, **kwargs: Any) -> ForwardContext
 
 __all__ = [
     "ASYNC_MOE_UBATCH_METADATA_KEY",
+    "AsyncMoeUbatchMetadata",
     "get_afd_metadata_from_forward_context",
     "get_async_moe_ubatch_metadata_from_forward_context",
     "use_afd_metadata_provider",