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

[ShwStone]

Summary

Add token-balanced, AFD-managed asynchronous MoE ubatching for
CAMAsyncAFDConnector on non-PCP DP+TP/SP Attention topologies.

The new token policy divides the real flattened token workload into exactly two
approximately balanced stages while preserving token order, stage-local
Attention metadata, and CAM dispatch/combine pairing. Request-boundary splitting
remains available and continues to be used by the existing PCP-oriented policy.
Native vLLM DBO remains unsupported for Ascend Async CAM.

Closes #149.

This PR supersedes #152 and #154 and combines their feature and E2E work on the
latest main. Both earlier PRs were developed while the model-configuration
identity regression described in #172 was present. That unrelated regression
caused TP2 FlashComm1 MLA/RoPE failures and materially distorted the debugging
and validation of those changes. #172 was fixed by #174; this replacement was
rebuilt and revalidated on the corrected baseline.

What changed

• Add a plugin-owned two-stage planner for token-balanced and request-boundary
  Async CAM MoE ubatching.
• Distinguish real scheduled tokens, parent DP padding, and stage-local TP/SP
  padding.
• Rebuild stage-local request/token offsets and Ascend Attention metadata,
  including sequence lengths, positions, slot mappings, and routing state.
• Preserve stage identity across asynchronous CAM dispatch/combine operations.
• Convert FlashComm1/SP full-batch TP shards into stage-local shards and restore
  the parent layout after the MoE pipeline.
• For plain TP, shard replicated tokens only at the CAM boundary and all-gather
  FFN output before the next Attention layer.
• Keep Attention and FFN parallel topology validation independent. The tested
  deployment uses Attention DP3TP2 and FFN DP2TP1 with EP2.
• Fail early for unsupported configurations, including more than two stages,
  native DBO, decode context parallel, and token splitting on unsupported
  Attention topologies.
• Add opt-in shape/layout diagnostics through
  AFD_ASYNC_MOE_LAYOUT_LOG=1; the diagnostics do not inspect tensor values or
  introduce device synchronization.
• Update the Async CAM user guide and runtime/connector design documentation.

No vLLM or vLLM-Ascend source-tree changes or model monkey patches are added.

Supported and validated configurations

The Ascend DP3TP2 Attention + DP2TP1/EP2 FFN GSM8K E2E matrix passed in all six
distinct configurations:

Split policy	Attention FlashComm1/SP	Async MoE ubatching	Result
token	enabled	enabled	PASS
token	disabled	enabled	PASS
request	enabled	enabled	PASS
request	disabled	enabled	PASS
request	enabled	disabled	PASS
request	disabled	disabled	PASS

When async MoE ubatching is disabled, token and request split settings are
operationally equivalent because no stage split is performed; the matrix uses
request as the representative no-ubatch configuration.

Validation

Unit tests passed, including the planner, stage metadata reconstruction,
forward-context/state-machine behavior, SP/plain-TP layout conversion, and
startup validation coverage.

Focused unit command:

python -m pytest -vv \
  tests/unit/v1/worker/test_async_moe_ubatch_planner.py \
  tests/unit/model_executor/models/test_forward_context.py \
  tests/unit/v1/worker/test_npu_runtime.py

Six-configuration Ascend E2E matrix:

AFD_NPU_ASYNC_CAM_RUN_DP3TP2=1 \
python -m pytest -svv \
  tests/e2e/accuracy/test_gsm8k_npu_async_cam.py::test_gsm8k_lm_eval_async_cam_dp3tp2_ep2

Because of the available test time, the accuracy run was limited to one
representative configuration:

• Configuration: token-ubatch-sp
• Dataset: GSM8K
• Samples: 100
• Score: 22.00%
• Reported std: 0.0416

The remaining five configurations completed their E2E runs successfully, but
were not all repeated with the 100-sample accuracy limit.

Scope and compatibility

In scope:

• Ascend CAMAsyncAFDConnector prefill execution.
• Exactly two AFD-managed asynchronous MoE stages.
• Token and request split policies.
• FlashComm1/SP and plain-TP Attention layouts.
• Independent Attention DP3TP2 and FFN DP2TP1/EP2 topology.
• Single-request, odd-token, skewed-length, and padded batches.

Out of scope:

• Native vLLM DBO for Async CAM.
• Decode context parallel.
• More than two asynchronous MoE stages.
• GPU-specific implementations.
• E2E process-lifecycle cleanup changes.

Review request

Even excluding tests and documentation, this PR changes a substantial amount of
runtime code (9 production files, approximately +1306/-186). Review is
requested, with particular attention to:

1. async_moe_sp.py: SP/plain-TP token layout conversion and padding rules.
2. deepseek_v2_async_cam_forward.py: two-stage dispatch/receive ordering and
   pending state pairing.
3. attention_model_runner.py and ubatch_utils.py: stage metadata rebuilding
   and real-versus-physical token counts.
4. feature_validation.py: Attention/FFN topology constraints and preservation
   of the existing request-boundary PCP policy.

Documentation

• docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md
• docs/design/module/attention_runtime.md
• docs/design/module/connector_contracts.md
• docs/design/module/execution_platforms.md

[Copilot]

Pull request overview

Adds plugin-owned, token-balanced two-stage Async CAM MoE microbatching for Ascend NPU, enabling non-PCP DP+TP/SP Attention topologies while preserving stage-local metadata and CAM dispatch/combine pairing.

Changes:

• Introduces a CPU-safe two-stage planner supporting async_moe_split="token" (token-balanced) and "request" (request-boundary) for Async CAM.
• Rebuilds stage-local Attention metadata and forward-context state, including distinct handling for real tokens vs parent DP padding vs stage-local TP/SP padding.
• Adds SP/plain-TP CAM-boundary layout conversions, expands validation rules, and extends unit + opt-in GSM8K E2E coverage with updated documentation.

Reviewed changes

Copilot reviewed 18 out of 18 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
afd_plugin/async_moe.py	Defines backend-independent Async CAM stage contracts/constants.
afd_plugin/v1/worker/npu/async_moe_ubatch.py	Adds the two-stage planner for request/token split policies.
afd_plugin/v1/worker/npu/attention_model_runner.py	Integrates stage planning, stage metadata building, and isolates metadata builders for Async CAM.
afd_plugin/v1/worker/npu/ubatch_utils.py	Adds Async CAM-specific attention-metadata slicing that distinguishes real vs stage-physical tokens.
afd_plugin/model_executor/models/forward_context.py	Replaces TypedDict with an immutable Async MoE execution-plan dataclass.
afd_plugin/model_executor/models/npu/async_moe_sp.py	Implements SP/plain-TP stage tensor layout conversion and CAM dispatch/restore layout handling.
afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py	Updates Async CAM forward orchestration to use stage inputs + CAM dispatch layouts and restores outputs.
afd_plugin/connectors/npu/async_cam.py	Documents/exports token vs request split constants and updates connector extra-info defaults.
afd_plugin/compat/npu/feature_validation.py	Tightens feature validation (native DBO/ubatching rejected; token split constraints enforced on Attention).
tests/unit/v1/worker/test_async_moe_ubatch_planner.py	Adds unit tests for the stage planner split policies and edge cases.
tests/unit/v1/worker/test_npu_runtime.py	Extends runtime validation + metadata tests; adds module reimport helper for isolation.
tests/unit/model_executor/models/test_forward_context.py	Adds tests for execution-plan validation and stage-context behavior/order/layout conversion.
tests/e2e/conftest.py	Enhances E2E launcher to support connector overrides, per-role TP sizing, async mode, and per-role env vars.
tests/e2e/accuracy/test_gsm8k_npu_async_cam.py	Adds opt-in GSM8K accuracy matrix for DP3TP2 Attention + DP2TP1/EP2 FFN Async CAM modes.
docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md	Updates user guide for token-balanced ubatching, env var guidance, and E2E matrix description.
docs/design/module/attention_runtime.md	Documents token-balanced stages and topology semantics (Attention-local SP, FFN independence).
docs/design/module/connector_contracts.md	Updates CAMAsync connector contract text to reflect request/token stage pipeline.
docs/design/module/execution_platforms.md	Updates platform matrix + behavior notes for token-balanced Async CAM pipeline.

---

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

[ShwStone]

Follow-up validation: repeatable token-SP vs request-SP score difference

Additional GSM8K limit=100 runs show the following results:

Async MoE split	FlashComm1/SP	Runs	Mean
token	enabled	0.22, 0.20, 0.19	0.203
request	enabled	0.25, 0.24	0.245

This qualifies the single 0.22 token-SP result currently mentioned in the PR description. The observed mean gap is about 4.2 percentage points.

I do not yet consider this evidence of a confirmed token-split correctness regression:

• At 100 examples, the binomial standard error of each aggregate score is roughly 4 percentage points.
• The same configuration is not bitwise/reproduction deterministic across runs (token-SP varies by 3 points; request-SP by 1 point).
• --limit 100 selects the same dataset prefix when the task/dataset version is unchanged, so the right analysis is paired per-example comparison rather than comparing only aggregate scores.

Static inspection so far:

• The main semantic difference is that token splitting can divide one request across the two stages, while request splitting cannot.
• The stage-2 context accounting appears intentional: vLLM-Ascend reconstructs context as seq_lens - query_lens; keeping the full sequence length while shortening stage-2 query length includes the tokens processed by stage 1.
• Positions, slot mapping, block table rows, query ranges, SP gather/slice order, and stage padding are rebuilt/sliced consistently in the paths inspected.
• Existing unit tests cover the single-request mid-request slice shapes, but they do not prove numerical equivalence between a split MLA forward and a non-split forward.

Next checks:

1. Compare the five samples_*.jsonl files and verify identical document IDs, prompts/few-shot examples, and generation settings.
2. Build per-document correctness vectors and count:
   • flips between repeated runs of the same mode;
   • request-correct/token-wrong versus token-correct/request-wrong cases.
3. Compare raw completions for disagreements (generation divergence vs answer-extraction differences).
4. If prompts are identical, isolate attention-only and MoE-stage numerical differences for a request split inside one long prefill.

Please treat this as an open validation finding, not yet a pass/fail conclusion. The per-sample lm-eval artifacts are needed to distinguish a strategy bug from batch/DP scheduling and low-precision execution variability.

[ShwStone]

Important refinement from the code and the previously captured runtime shapes: when an Attention DP has only one scheduled long prefill, request cannot form two stages and plan_async_moe_stages returns None; token still splits that request in the middle. In that common case the reported comparison is therefore split-prefill versus request-policy fallback, not two simultaneously active ubatch policies. The second token stage is intentionally represented as a chunked prefill with the first stage in KV context. This makes split-prefill numerical/correctness equivalence the primary investigation target. We should also make future validation prove whether each requested policy actually activated, rather than infer activation from config alone.

[ShwStone]

Follow-up: batch-size-1 accuracy result

With lm-eval batch size set to 1, both SP configurations produced the same GSM8K limit=100 accuracy:

Async MoE split	FlashComm1/SP	Batch size	Accuracy
token	enabled	1	0.28
request	enabled	1	0.28

This controlled single-request result does not reproduce the earlier token-SP versus request-SP score difference. In particular, it provides no evidence that splitting a single request with the token policy causes an accuracy loss relative to the request-policy fallback in this run.

We will therefore stop investigating batch-size effects for now. The earlier batch-size-8 aggregate difference remains documented above, but it is not currently treated as a confirmed split-policy correctness regression. No code change is proposed from these accuracy observations.

[jiangkuaixue123]

Why is this a patch function？

[ShwStone]

It is not a runtime patch and it does not replace or wrap an upstream function. This ubatch_utils.py module is a plugin-owned copy of the pinned vLLM-Ascend implementation, so the PATCH START/END markers were intended only to delimit the AFD-specific delta inside copied upstream code, as required by the repository guidelines. The helper itself is imported and called explicitly by the AFD NPU runner.

I agree that the Patch reason/functionality/signature wording is misleading for a plugin-owned helper. It would be clearer to describe this as an AFD-specific extension while retaining only the copied-upstream delta markers.