bugfix:DBO ubatch NPUGraph capture passes per-stage metadata to the standard full-graph updater

[yujuancao07]

Purpose

Add NPU support for MLA native DBO with Ascend ACL FULL graph replay,
including dense MLA models such as DeepSeek-V2-Lite.

With DBO enabled, the two child ubatches capture independent MLA graph
parameters. The existing vLLM-Ascend process-global graph parameter registry
cannot represent both child graphs safely. This change creates one registry
per ubatch, merges the captured records before replay, and preserves the
required eNPU/non-eNPU update ordering.

Issue

• Related issue(s): [Bug]: DBO ubatch NPUGraph capture passes per-stage metadata to the standard full-graph updater #177
• Closing keyword, only if fully resolved: Closes [Bug]: DBO ubatch NPUGraph capture passes per-stage metadata to the standard full-graph updater #177

Scope

• In scope:

  • Support MLA DBO FULL graph capture and replay on NPU.
  • Maintain independent MLA graph parameter registries for two ubatches.
  • Share the aggregate FIA workspace between the two child registries.
  • Merge graph records in layer-major, ubatch-minor order before replay.
  • Enable the existing DeepSeek-V2-Lite NPU DBO E2E cases.

• Out of scope:

  • Changing the existing sparse MLA path.
  • Changes to the vLLM source checkout.
  • Changes to CAMAsync scheduling behavior.

Test Plan

Focused unit and runner tests

python -m pytest -vv \
  tests/unit/compat/test_runtime.py \
  tests/unit/v1/worker/test_npu_mla_graph.py \
  tests/unit/v1/worker/test_npu_runtime.py \
  tests/e2e/test_runner.py

[jiangkuaixue123]

Thanks for the investigation. The root cause is not the shape of attn_metadata, but that MLA on Ascend has not yet been adapted to DBO + FULL graph.

In the previous AFD v0.13-based implementation, MLA had a dedicated mechanism to update the full-graph parameters. That mechanism should be ported to the current implementation. Skipping _update_full_graph_params_if_needed() for ubatches only avoids the immediate exception; it does not provide the MLA full-graph update semantics that are required.

This is why DeepSeek-V2-Lite cannot currently run with DBO + FULL graph. DeepSeek-V3.2 can run with the same configuration because it uses DSA rather than MLA, so it does not hit this MLA-specific gap.

Therefore, we will not accept the current workaround as the fix. Please port the previous MLA-specific full-graph update mechanism and validate DeepSeek-V2-Lite with DBO + FULL graph.

[jiangkuaixue123]

Reviewed commit 602ad8a6e01d18ac7355806099a9e8e8a028373d with InferMatrixCopilot Direct.

Requesting changes for the inline P1 below.

Validation performed against this snapshot:

• compileall passed for every changed Python file.
• tests/unit/v1/worker/test_npu_mla_graph.py: 25 passed on CPython 3.12 / macOS arm64 with PyTorch.
• The broader NPU runtime test set could not be collected on this host because the pinned vllm / vllm_ascend / torch_npu runtime is unavailable; GitHub currently exposes no combined CI statuses for this SHA.

Cleanup: please remove the unrelated docs/assets/WeChat.jpg binary replacement (272,151 → 412,617 bytes). It is outside the MLA graph fix and adds repository bloat/review noise.

[jiangkuaixue123]

[P1] Reject or support speculative MLA before enabling DBO FULL graph

With use_mla=True, DBO enabled, a FULL graph mode, and a non-None speculative_config, this condition makes supports_mla_dbo_full_graph false, so the configuration is accepted. _install_ascend_ubatch_wrapper() then also sets mla_full_graph_enabled=False; once the DBO threshold creates two ubatches, _model_forward() leaves the update to the upstream MLA path. The parent forward_context.attn_metadata is a list[dict] from _build_attention_metadata_with_ubatches(), while vLLM-Ascend 0.19.1 MLAImpl.update_graph_params() takes the main-model branch and calls attn_metadata.keys(), so replay fails instead of using the new per-stage registries. Either fail fast for speculative MLA + DBO FULL graph, or extend the per-ubatch registry/update path to this case; add a regression that reaches _model_forward() rather than asserting this configuration is allowed.

[jiangkuaixue123]

The per-ubatch registry direction is reasonable, but I do not think the current ownership boundary is ready to merge.

The existing P1 about speculative MLA remains unresolved: validation accepts speculative MLA + DBO + FULL graph while the per-ubatch registry path is disabled, so execution falls back to the upstream single-batch updater with list[dict] metadata.

More broadly, the generic Ascend ubatch wrapper now owns MLA feature policy, GraphParams internals, FIA workspace selection, layer-record merging, a bound runner updater, and eNPU/non-eNPU update ordering. Please extract these MLA-specific responsibilities into a dedicated graph coordinator/adapter. The wrapper should retain only generic ubatch graph orchestration: an opaque capture state plus small capture/replay lifecycle hooks. Completely avoiding wrapper changes is not necessary, but the wrapper should not become the owner of MLA graph semantics.

Before approval, please also provide real NPU E2E evidence for DeepSeek-V2-Lite with DBO + FULL_DECODE_ONLY. The stub-heavy unit tests cannot validate the external-event, two-thread forward-context switching, and ACL graph update/replay ordering exercised by this change. If both eNPU and non-eNPU paths are supported, please validate both orderings.

Correction to my earlier review: the current PR diff does not contain a docs/assets/WeChat.jpg change, so please disregard that cleanup note.

[jiangkuaixue123]

[P1] Do not silently mark a required MLA patch as applied

apply_afd_mla_graph_patch() returns None both when installation succeeds and when its broad import except Exception path silently returns. The next line then sets _PATCHES_APPLIED = True. If installation fails, the wrapper can still claim ownership of MLA full-graph updates, but the child forward contexts will continue using the upstream process-global registry instead of their per-ubatch registries; the guard also prevents a later retry. Please make installation status explicit and fail fast when MLA DBO FULL graph requires this patch (while preserving the optional no-Ascend path). Add a regression covering a failed or unavailable resolver installation.

[jiangkuaixue123]

[P2] Keep MLA graph policy out of the generic ubatch wrapper

These constructor additions are the start of a much broader ownership inversion: AscendUBatchWrapper now decides MLA support, imports and constructs GraphParams, selects the aggregate FIA workspace, merges layer records, calls a bound private runner updater, and enforces speculative/PCP policy. The duplicated policy has already diverged from feature validation in the speculative-MLA case. Please move this behavior into an MLA-specific coordinator/adapter owned by the NPU attention runner, and let the wrapper interact with it through opaque capture state and small capture/replay hooks. The generic wrapper should continue to own splitting, scheduling, capture, replay, and stream fencing—not MLA parameter semantics.

[jiangkuaixue123]

Please append fully e2e test on NPU. And Dsv2-lite profiling.

[jiangkuaixue123]

fix

[jiangkuaixue123]

Add comment and docstring