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

[ShwStone]

CAMAsyncAFDConnector previously split async MoE prefill work into two microbatches only at request boundaries. This works for PCP deployments but produces heavily imbalanced stages in non-PCP DP+TP/SP topologies where request lengths differ.

Add a token-count-based split policy:

• Introduce async_moe_split connector_extra_config with values "request" (default, preserves PCP behavior) and "token".
• On non-PCP DP+TP/SP topologies, split the padded prefill token workload into two TP-aligned stages of approximately equal size.
• Rebuild stage-local request/token offsets, sequence lengths, and Ascend attention metadata for each stage.
• Add SP-local ubatch slices (sp_local_ubatch_slices) so model-side stage slicing works when sequence parallelism shards tokens across TP ranks.
• Move stage-input slicing into ubatch_sp.py and keep the forward orchestration in deepseek_v2_async_cam_forward.py minimal.

Startup validation fail-fast:

• Reject async_moe_split="token" on unsupported topologies.
• Warn when running "request" on a topology that supports "token".

Tests cover request-boundary vs token-balanced selection, TP-size alignment, padded/odd token counts, SP-local token mapping, and topology validation.

Purpose

Implement the token-balanced async MoE microbatch split for non-PCP DP+TP/SP topologies, as described in #149. The goal is to avoid heavily imbalanced pipeline stages when request lengths are skewed, while keeping the existing PCP request-boundary behavior unchanged.

Issue

• Related issue(s): [Feature]: Support token-based CAMAsync microbatch splitting for non-PCP DP+TP/SP topologies #149
• Closing keyword, only if fully resolved: Closes [Feature]: Support token-based CAMAsync microbatch splitting for non-PCP DP+TP/SP topologies #149

Scope

• In scope:
  • CAMAsync async MoE ubatch planning for prefill with exactly two microbatches.
  • Token-balanced split on non-PCP DP+TP/SP topologies.
  • TP-aligned split points and padded batch handling.
  • Stage-local attention metadata reconstruction.
  • SP-local ubatch slices for sequence-parallel tensor slicing.
  • Startup validation and user-facing async_moe_split config contract.
  • CPU-safe unit tests for planner and SP-local slicing.
• Out of scope:
  • Native vLLM DBO behavior changes.
  • Decode context parallel support.
  • More than two microbatches.
  • GPU-specific implementations.

Implementation Notes

• The split policy is selected in plugin-owned afd_plugin.v1.worker.npu.ubatch_utils.create_async_moe_ubatch_slices, which returns both the slices and the applied policy string (request_boundary or token_balanced_tp).
• Topology capability is shared between runtime planning and startup validation via enable_token_balanced_async_moe_split.
• SP-local slice mapping is isolated in afd_plugin.model_executor.models.npu.ubatch_sp.build_async_moe_stage_inputs; the forward orchestration only consumes the resulting stage tensors and slice metadata.
• No changes are made to the vLLM or vLLM-Ascend source checkout; all extension points are plugin-owned classes/functions.
• All new imports are CPU-safe; NPU-specific modules are stubbed in unit tests.

Test Plan

• CPU-only unit tests:
  • tests/unit/v1/worker/test_npu_runtime.py — ubatch planning, SP-local slice mapping, validation warnings/errors.
  • tests/unit/model_executor/models/test_forward_context.py — forward orchestration structure assertions.
• GPU/NPU-gated E2E (planned, requires hardware):
  • CAMAsync prefill with DP+TP and DP+SP, async_moe_ubatching=true, two microbatches.
  • Output accuracy comparison against non-ubatched CAMAsync baseline.
  • Skewed-length workload balance check.
  • Existing PCP CAMAsync recipe regression check.

Test Result

• CPU unit tests pass locally: py_compile, ruff check, ruff format.
• Full pytest suite was not executed; NPU E2E validation is pending hardware access.

Docs Impact

• Files updated: inline code comments and docstrings in modified modules.
• No standalone documentation changes; the feature is controlled through connector_extra_config.async_moe_split.

---

Essential PR Checklist

• Purpose is clear and linked to public context when possible.
• Scope is bounded.
• Compatibility with vLLM v0.19.1 is considered.
• No changes are made to the vLLM source checkout.
• Plugin-owned classes or explicit dotted class paths are preferred over monkey patches.
• Any compat shim or monkey patch is isolated, idempotent, version-guarded, documented, and tested.
• Imports remain CPU-safe; CUDA-heavy work is delayed or GPU-gated.
• Validation evidence is included, including skipped GPU tests when applicable.
• Documentation impact is stated.

[jiangkuaixue123]

Test plan: Let’s first use DSV2-Lite to validate accuracy with TP/SP u-batch splitting.

[Copilot]

Pull request overview

This PR adds a token-balanced async MoE microbatch split policy for CAMAsync on non-PCP DP+TP/SP topologies, aiming to reduce pipeline stage imbalance when request lengths are skewed, while keeping the existing request-boundary behavior as the default.

Changes:

• Introduces async_moe_split selection ("request" default, "token" token-balanced TP-aligned) and topology capability detection/validation.
• Implements token-balanced split planning (including padded token handling) and stage-local SP-aware slicing via a new ubatch_sp helper.
• Expands CPU-safe unit tests to cover split selection, alignment, padding edge cases, and SP-local stage slicing behavior.

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/unit/v1/worker/test_npu_runtime.py	Adds CPU-safe stubs and unit tests for token-balanced split behavior, padding edge cases, SP-local slice mapping, and validation messaging.
tests/unit/model_executor/models/test_forward_context.py	Updates forward-orchestration assertions to reflect the new stage-input slicing helper.
afd_plugin/v1/worker/npu/ubatch_utils.py	Adds token-balanced split policy, split-point computation, async ubatch slice creation, and SP-local ubatch slice mapping.
afd_plugin/v1/worker/npu/attention_model_runner.py	Switches async MoE ubatch planning to the new policy-aware helper and records whether SP-local slicing should be used.
afd_plugin/model_executor/models/npu/ubatch_sp.py	New module implementing SP-aware stage-input slicing and mapping global ubatch splits to per-rank token shards.
afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py	Replaces inline slicing with build_async_moe_stage_inputs and threads SP-local slices through stage metadata/context.
afd_plugin/model_executor/models/forward_context.py	Extends async MoE ubatch metadata typing to include SP-local slicing fields.
afd_plugin/connectors/npu/async_cam.py	Documents and exposes "token" as a valid async_moe_split config value.
afd_plugin/compat/npu/feature_validation.py	Adds fail-fast validation for "token" split on unsupported topologies and warning guidance when "request" is used on capable topologies.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 9 out of 9 changed files in this pull request and generated no new comments.

Comments suppressed due to low confidence (4)

afd_plugin/v1/worker/npu/ubatch_utils.py:165

• AGENTS.md asks to avoid getattr/hasattr on vLLM/vLLM-Ascend structures so upstream compatibility issues fail loudly (AGENTS.md:64-69). In enable_token_balanced_async_moe_split, using getattr with defaults can silently mask missing/renamed fields and change behavior; access parallel_config attributes directly instead.

    if int(getattr(parallel_config, "tensor_parallel_size", 1)) <= 1:
        return False
    if int(getattr(parallel_config, "prefill_context_parallel_size", 1)) != 1:
        return False
    return int(getattr(parallel_config, "decode_context_parallel_size", 1)) == 1

afd_plugin/v1/worker/npu/ubatch_utils.py:286

• build_sp_local_ubatch_slices_for_current_rank currently swallows all exceptions and silently falls back to global slices. This can hide real upstream/runtime breakages (e.g., get_tp_group/enable_sp API changes) and make SP slicing wrong without an obvious failure. Prefer only treating missing optional imports as a fallback, and let other exceptions surface.

        if not bool(enable_sp()) or tp_size <= 1:
            return ubatch_slices
    except Exception:
        return ubatch_slices

afd_plugin/v1/worker/npu/ubatch_utils.py:228

• The create_async_moe_ubatch_slices docstring says it raises AssertionError when num_ubatches != 2, but the function currently has no such assertion and will attempt to compute split points for any num_ubatches. Either update the docstring or enforce the 2-ubatch contract to match the stated behavior (and async_moe_ubatching validation).

    Returns ``(None, "request_boundary")`` when the batch is too small to
    split (fewer than ``num_ubatches`` requests or tokens); callers should
    run such steps unubatched. Raises AssertionError if ``num_ubatches``
    is not 2.
    """

afd_plugin/model_executor/models/npu/ubatch_sp.py:243

• _slice_sequence_tensor and _pad_sequence_tensor are currently unused (only definitions exist) and they implement the heuristic slicing/padding behavior that was previously identified as buggy for tensors whose token dimension is not inferred by shape[0] <= 4. Removing these unused helpers will reduce the risk of accidentally reintroducing that bug.

def _slice_sequence_tensor(tensor: torch.Tensor, token_slice: slice) -> torch.Tensor:
    if tensor.dim() > 1 and int(tensor.shape[0]) <= 4:
        return tensor[:, token_slice]
    return tensor[token_slice]

[jiangkuaixue123]

Please fix DCO and CI.

[jiangkuaixue123]

Scan:

Category	Result
Tests / verification	3 finding(s) below
Security	no finding reported
Docs / comments	2 finding(s) below
Behavior / compatibility	no finding reported
Correctness	no finding reported

Validated:

• [sweep] Checked ASYNC_MOE_REQUEST_SPLIT consumers via grep: feature_validation.py (updated), async_cam.py (constant unchanged, docstring updated). No stale consumers.
• [sweep] Checked create_request_boundary_ubatch_slices callers via grep: attention_model_runner.py now calls create_async_moe_ubatch_slices; create_request_boundary_ubatch_slices still exported and used as fallback. Tests updated.
• [sweep] Checked AsyncMoeUbatchMetadata consumers via grep: forward_context.py (definition extended with optional fields), deepseek_v2_async_cam_forward.py (reads new sp_local_ubatch_slices via .get with fallback), deepseek_v2.py (reads via get_async_moe_ubatch_metadata_from_forward_context). Backward-compatible.
• [sweep] Checked removed _slice_optional_first_dim, _slice_positions, _slice_llama_4_scaling: all were private to deepseek_v2_async_cam_forward.py. No external consumers found — safe removal.
• [sweep] Checked async_moe_split default: ASYNC_MOE_REQUEST_SPLIT = "request" unchanged. New ASYNC_MOE_TOKEN_SPLIT = "token" constant added. AFDAsyncExtraInfo.async_moe_split default remains "request".
• [sweep] Checked docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md:154 — says 'request-boundary splitting only', stale after PR. Line in 'Current limitations' also says 'two request-boundary stages'. Must document "token".
• [sweep] Checked docs/design/module/connector_contracts.md:94 — says 'optional request-boundary pipeline', stale. docs/design_zh/module/connector_contracts.md:77 same.
• [sweep] Checked recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md:72 — describes 'request-level splitting' for this recipe, still accurate but doesn't acknowledge "token" option.

afd_plugin/v1/worker/npu/attention_model_runner.py:234 [major] — The diff changes _build_attention_metadata_with_async_moe_ubatches to call create_async_moe_ubatch_slices and store use_sp_local_ubatch_slices in the metadata dict. There are no unit tests for this path. Add a CPU unit test that instantiates AFDNPUAttentionModelRunner with async_moe_ubatching=True, split='token', and a suitable batch; verify the metadata contains use_sp_local_ubatch_slices and the correct ubatch_slices. (evidence: grep pattern '_build_attention_metadata_with_async_moe_ubatches|async_moe_ubatching' in tests/ returned zero matches; read_file of tests/unit/v1/worker/test_npu_runtime.py confirmed no test exercises this method.)

afd_plugin/model_executor/models/npu/ubatch_sp.py:190 [major] — _slice_sequence_tensor_for_sp_stage has branches for llama_4_scaling and 2D positions that are not exercised by test_ubatch_sp_stage_inputs_use_sp_local_ranges (which uses 1D positions and llama_4_scaling=None). Add test cases with 2D positions and nonzero llama_4_scaling to cover the local_token_dim==1 and global_token_dim==1 paths. (evidence: read_file of test_ubatch_sp_stage_inputs_use_sp_local_ranges shows only 1D positions and llama_4_scaling=None; read_file of _slice_sequence_tensor_for_sp_stage shows branches at lines 198, 200, 204, 207, 209 that are not exercised.)

docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md:154 [minor] — The diff adds async_moe_split="token" support, but docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md:154 still states 'The current async connector supports request-boundary splitting only.' Update this line, the 'AFD-managed asynchronous MoE ubatching' paragraph, and the 'Current limitations' bullet to document the token split, its topology preconditions (non-PCP, TP>1, DP+TP/SP), and use cases; similarly update recipe/npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md:72 and the design docs. (evidence: grep for async_moe_split across docs/ found 5 locations restricting the field to request-boundary only; diff adds ASYNC_MOE_TOKEN_SPLIT = "token" in async_cam.py:68 and "token" is accepted in feature_validation.py:141-155.)

afd_plugin/model_executor/models/npu/ubatch_sp.py:198 [minor] — _slice_sequence_tensor_for_sp_stage handles local_token_dim 0 and 1 but returns the tensor unchanged when the token dimension is >1 (and not detected by global_token_dim). The sibling _slice_positions uses positions[..., token_slice] for higher dims. If a sequence tensor with tokens on a higher axis reaches this path, the input may be passed through incorrectly. Add an ellipsis fallback or assert that the dimensionality is ≤2, and document the assumption. (evidence: Read ubatch_sp.py lines 185-227 (_slice_sequence_tensor_for_sp_stage): only local_token_dim == 0 → tensor[tensor_token_slice] and local_token_dim == 1 → tensor[:, tensor_token_slice]. Compare with _slice_positions at line 240: positions[..., token_slice] for dim>1.)

afd_plugin/model_executor/models/npu/ubatch_sp.py:131 [minor] — In _build_async_moe_stage_inputs_with_slices, stage_hidden_states uses bare indexing hidden_states[slice] while stage_residual uses _slice_and_pad_first_dim. Both slices are exact (non‑SP path), so the asymmetry is functionally correct but confusing—and the failure modes differ: out‑of‑bounds on hidden_states would crash while residual would silently pad. Use the same slicing helper (e.g., _slice_and_pad_first_dim) for both, or add a comment explaining the intent. (evidence: Read ubatch_sp.py lines 125-140: stage_hidden_states = [hidden_states[sp_local_ubatch_slice.token_slice] ...] vs stage_residual = [_slice_and_pad_first_dim(residual, sp_local_ubatch_slice.token_slice) ...]. The old code in deepseek_v2_async_cam_forward.py used _slice_optional_first_dim (just a bare slice) for both.)

Verdict: REQUEST CHANGES

[jiangkuaixue123]

Two suggestions on the async MoE ubatch metadata type definition.

[jiangkuaixue123]

Could you add a class docstring explaining why the SP-local microbatch metadata is optional and when each field is populated and consumed?

[ShwStone]

Added a class docstring on AsyncMoeUbatchMetadata explaining required vs optional fields and when each is populated/consumed.

[jiangkuaixue123]

One compatibility concern for the updated ubatch split path.

[jiangkuaixue123]

Overall, the structure LGTM. Please address the review comments and validate the accuracy. Once those are complete, this should be ready to merge.

[ShwStone]

Pushed a follow-up fix commit to resolve a test collection failure on CPU-only development paths.

Root cause: The previous commit added a module-level from vllm.logger import init_logger in afd_plugin/compat/npu/feature_validation.py. tests/unit/v1/worker/test_npu_runtime.py imports afd_plugin.compat.npu at module scope, so in CI environments that only have pytest installed (without the vLLM runtime), test collection threw ModuleNotFoundError: No module named 'vllm' instead of skipping the NPU runtime tests gracefully.

Fix: Replaced from vllm.logger import init_logger with the standard library logging.getLogger(__name__); logger.warning(...) calls remain unchanged.

[jiangkuaixue123]

Please fix DCO

[ShwStone]

Thanks for the thorough scan. All findings have been addressed in the latest push:

1. [major] attention_model_runner.py:234 — Missing unit tests for _build_attention_metadata_with_async_moe_ubatches

Added two tests:

• test_npu_attention_runner_builds_async_moe_ubatch_metadata (token split)
• test_npu_attention_runner_builds_async_moe_ubatch_metadata_request_split (request split)

2. [major] ubatch_sp.py:190 — Missing 2D positions / llama_4_scaling tests

Added test_ubatch_sp_stage_inputs_2d_positions_and_llama4_scaling covering 2D positions (dim=1) and non-None llama_4_scaling (dim=1).

3. [minor] Docs stale

Updated CAM_ASYNC_CONNECTOR_USER_GUIDE.md, connector_contracts.md, and the DeepSeek v3.2 recipe README to document async_moe_split="token" and its topology preconditions. Note: docs/design_zh/module/connector_contracts.md does not exist in the current tree (checked origin/main and all branches), so no change needed there.

4. [minor] ubatch_sp.py:198 — Silent pass-through for >1 dim token axis

_slice_sequence_tensor_for_sp_stage now raises ValueError when the token dimension is not on axis 0 or 1.

5. [minor] ubatch_sp.py:131 — Asymmetric slicing

Unified _build_async_moe_stage_inputs_with_slices to use _slice_and_pad_first_dim for both hidden_states and residual, with @overload to keep type inference precise.

DCO sign-off and GPG signature are also included in the latest commit.

[ShwStone]

DCO sign-off added. The squashed commit now includes Signed-off-by and a GPG signature.