feat(npu): token-balanced async MoE microbatch split for DP+TP/SP

afd_plugin/compat/npu/feature_validation.py

@@ -4,6 +4,7 @@
 
 from __future__ import annotations
 
+import logging
 from typing import TYPE_CHECKING
 
 from afd_plugin.config import (
@@ -18,6 +19,8 @@
 
     from afd_plugin.connectors.base import ConnectorExtraInfo
 
+logger = logging.getLogger(__name__)
+
 
 def fail_if_unsupported_npu_afd_features(
     vllm_config: VllmConfig,
@@ -111,7 +114,13 @@ 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.connectors.npu.async_cam import (
+        ASYNC_MOE_REQUEST_SPLIT,
+        ASYNC_MOE_TOKEN_SPLIT,
+    )
+    from afd_plugin.v1.worker.npu.ubatch_utils import (
+        enable_token_balanced_async_moe_split,
+    )
 
     parallel_config = vllm_config.parallel_config
     if not afd_config.compute_gate_on_attention:
@@ -123,15 +132,35 @@ def _fail_if_unsupported_npu_async_moe_ubatching_features(
             "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}",
         )
+    # The Attention process owns stage planning and SP layout conversion.
+    # FFN processes consume the resulting stage payloads, so their independent
+    # TP/CP topology must not be used to validate the Attention split policy.
+    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",
         )
+    token_split_capable = enable_token_balanced_async_moe_split(vllm_config)
+    if split == ASYNC_MOE_TOKEN_SPLIT:
+        if not token_split_capable:
+            raise RuntimeError(
+                "async_moe_split='token' requires a non-PCP Attention DP+TP/SP "
+                "topology (Attention tensor_parallel_size > 1, no "
+                "prefill/decode context parallel)",
+            )
+    elif split == ASYNC_MOE_REQUEST_SPLIT and token_split_capable:
+        logger.warning(
+            "async_moe_ubatching runs on a non-PCP DP+TP/SP topology with "
+            "async_moe_split='request'; request lengths can be skewed, "
+            "consider async_moe_split='token' for token-balanced "
+            "microbatches",
+        )
 
 
 __all__ = ["fail_if_unsupported_npu_afd_features"]

afd_plugin/connectors/npu/async_cam.py

@@ -18,9 +18,11 @@
 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. It uses
+request boundaries without sequence parallelism and TP-aligned token stages
+when sequence parallelism is active. See
+``docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md`` for configuration, rank
+derivation, launch guidance, and the full limitations.
 """
 
 from __future__ import annotations
@@ -65,6 +67,7 @@
 CAM_COMM_ID = 0
 ATTN_RANKS_PER_DP_CONFIG_KEY = "attn_ranks_per_dp"
 ASYNC_MOE_REQUEST_SPLIT = "request"
+ASYNC_MOE_TOKEN_SPLIT = "token"
 
 _AFD_ASYNC_EXTRA_CONFIG_FIELDS: Final[frozenset[str]] = frozenset(
     {
@@ -85,10 +88,14 @@ 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.
+        attn_ranks_per_dp: Number of Attention ranks in each data-parallel
+            replica. This is the CAM Attention grouping width and is
+            independent of the FFN process's local TP size.
+        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: Boundary at which async MoE work is split: "request"
+            for request boundaries, "token" for TP-aligned token-balanced
+            split points (non-PCP DP+TP/SP topologies only).
     """
 
     dynamic_quant: int = 0
@@ -233,8 +240,9 @@ def __init__(
 
         Communication resources are created collectively by
         ``init_afd_connector``. ``afd_role_rank`` must already contain the
-        process's role-local DP/PCP-derived offset; ``attn_ranks_per_dp`` is
-        used as the CAM Attention TP width.
+        process's role-local distributed-rank-derived offset;
+        ``attn_ranks_per_dp`` is used as the CAM Attention grouping width and
+        does not describe the FFN process's local TP topology.
         """
         super().__init__(rank, local_rank, vllm_config, afd_config)
         self._initialized = False

afd_plugin/model_executor/models/forward_context.py

@@ -18,7 +18,44 @@
 ASYNC_MOE_UBATCH_METADATA_KEY: Final[str] = "afd_async_moe_ubatch_metadata"
 
 
-class AsyncMoeUbatchMetadata(TypedDict):
+class _AsyncMoeUbatchMetadataOptional(TypedDict, total=False):
+    """Optional SP-local fields for ``AsyncMoeUbatchMetadata``.
+
+    A separate ``total=False`` base class is used instead of
+    ``typing.NotRequired`` because the plugin supports Python 3.10, where
+    ``NotRequired`` is only available via the external ``typing_extensions``
+    package. Keeping the base class avoids adding a new runtime dependency.
+    """
+
+    # SP-local stage layout: when SP shards the full batch across TP ranks,
+    # transpose it into equal per-stage rank shards before attention and
+    # restore the original layout after the async MoE pipeline.
+    use_sp_stage_resharding: bool
+    sp_local_stage_slices: UBatchSlices
+    stage_actual_token_counts: list[int]
+    sp_local_stage_actual_token_counts: list[int]
+
+
+class AsyncMoeUbatchMetadata(_AsyncMoeUbatchMetadataOptional):
+    """Async MoE ubatch sidecar metadata carried by the forward context.
+
+    The required fields (``attn_metadata``, ``ubatch_slices``) are always
+    populated by the attention model runner before the model forward starts.
+    The optional layout fields describe padded stage inputs and their real
+    token coverage:
+
+    - ``use_sp_stage_resharding``: set to ``True`` by the attention model
+      runner when Ascend sequence parallelism is active and the applied split
+      is TP-aligned.
+    - ``sp_local_stage_slices``: set by
+      ``build_async_moe_stage_inputs`` during model forward; its lengths
+      describe the equal per-rank stage shards, not ranges into the original
+      full-batch local shard.
+    - ``stage_actual_token_counts`` and
+      ``sp_local_stage_actual_token_counts`` distinguish real token rows from
+      DP/SP padding globally and on the current TP rank.
+    """
+
     attn_metadata: object
     ubatch_slices: UBatchSlices