[Feat]: Add NPU ModelRunnerV2 DBO with Eager and FULL_DECODE_ONLY ACL Graph Support

[lirx-pd]

Purpose

This PR adds temporary Dual Batch Overlap (DBO) support to the AFD NPU
ModelRunnerV2 path. The pinned vLLM v0.26.0 ABI does not yet include the
ModelRunnerV2 DBO implementation it needs.

It adapts ModelRunnerV2 DBO behavior from specture724/vllm branch
feat/v2/dbo-fullcg at commit 626fee7831 to vLLM ABI 568afb3a13. The
missing upstream pieces live in an isolated compatibility backport. When the
pinned vLLM release gains native ModelRunnerV2 DBO support, this backport can
be removed.

The AFD NPU ModelRunnerV2 path supports:

• eager execution without DBO;
• eager execution with two DBO microbatches;
• ACL graph execution without DBO using FULL or FULL_DECODE_ONLY;
• DBO ACL graph execution using FULL_DECODE_ONLY.

The change does not modify the vLLM or vLLM-Ascend source trees. NPU
ModelRunnerV1 keeps its existing execution path.

Architecture diagram

flowchart TB
    Config["AFD NPU V2 configuration"]
    Validation["Scoped v0.26 validation shim"]
    Runner["AFDNPUAttentionModelRunnerV2"]
    Dispatch["DP synchronization and DBO dispatch"]
    Single["Native single-batch V2 path"]
    Split["Two request-boundary uBatches"]
    Eager["AFDAscendUBatchRunnerV2<br/>eager execution"]
    Graph["AFDModelAclGraphManagerV2<br/>DBO graph capture/replay"]
    Yield["Per-layer DBO handoff"]
    Connector["CAMP2pAFDConnector"]
    FFN["Existing AFD NPU FFN runner"]
    Output["Merged model output"]
    Guard["Upstream capability guard<br/>remove backport when native DBO lands"]

    Config --> Validation --> Runner
    Runner --> Dispatch
    Dispatch -->|threshold not met| Single --> Connector
    Dispatch -->|eager DBO| Split --> Eager --> Yield
    Dispatch -->|captured DBO shape| Split --> Graph --> Yield
    Yield --> Connector --> FFN --> Connector --> Output
    Guard -.-> Validation

    classDef new fill:#fff7ed,stroke:#f59e0b,color:#92400e,stroke-width:2px;
    classDef reused fill:#ecfdf5,stroke:#16a34a,color:#166534,stroke-width:2px;
    classDef upstream fill:#eff6ff,stroke:#3b82f6,color:#1e40af,stroke-width:2px;

    class Validation,Dispatch,Split,Eager,Graph,Guard new;
    class Runner,Yield,Connector,FFN,Output reused;
    class Single upstream;

Loading

Issue

• Related issue(s): [Feature]: Add DBO support for Ascend NPU ModelRunnerV2 #265
• Closing keyword, only if fully resolved: [Feature]: Add DBO support for Ascend NPU ModelRunnerV2 #265

Scope

In scope

• Backport the ModelRunnerV2 DBO execution pieces missing from the pinned vLLM
  v0.26.0 ABI into afd-plugin.
• Enable two-microbatch eager DBO for the AFD NPU ModelRunnerV2 Attention path.
• Enable two-microbatch FULL_DECODE_ONLY ACL graph capture and replay.
• Synchronize token counts, graph mode, uniform decode state, and DBO selection
  across DP ranks before execution.
• Apply the configured decode/prefill DBO thresholds using the minimum real
  token count across DP ranks.
• Pad DBO execution to the maximum selected DP shape while falling back to a
  single batch when graph padding would leave the second microbatch empty.
• Split ModelRunnerV2 inputs at request boundaries and preserve Ascend-specific
  input, attention, block-table, slot-mapping, and forward-context state for
  each microbatch.
• Run the two eager microbatches through isolated forward contexts and merge
  tensor, auxiliary-hidden-state, and intermediate outputs in request order.
• Capture plugin-owned twin DBO graphs separately from native single-batch ACL
  graphs and publish matching AFD metadata for warmup, capture, and replay.
• Reuse the upstream FIA workspace while maintaining per-microbatch MLA graph
  parameters and updating the captured task-group handles during replay.
• Allow the existing connector-driven FFN runner to replay the matching graph
  during repeated Attention-side capture events.
• Narrowly bypass vLLM's native DBO/all-to-all validation only for an active AFD
  NPU ModelRunnerV2 configuration.
• Preserve native execution when DBO is disabled or the runtime threshold is
  not met.
• Preserve the existing NPU ModelRunnerV1 DBO and graph paths.

Out of scope

• Native or generally reusable vLLM ModelRunnerV2 DBO support.
• More than two microbatches.
• DBO with DP size one.
• DBO ACL graph modes other than FULL_DECODE_ONLY.
• Speculative decoding, LoRA, multimodal models, or encoder-decoder models with
  NPU ModelRunnerV2 DBO.
• PP, PCP, or DCP greater than one, elastic expert parallelism, EPLB, sequence
  parallel MoE, or other topologies already unsupported by AFD NPU
  ModelRunnerV2.
• Asynchronous CAMP or asynchronous AFD connector lifecycles.
• Performance tuning or a claim that DBO improves every model or workload.

Implementation notes

Temporary vLLM v0.26 ABI backport

afd_plugin/compat/backports/vllm_v026_mrv2_dbo contains the copied and
adapted behavior missing from the pinned upstream ABI:

• an AFD batch descriptor carrying num_ubatches;
• DP synchronization and DBO/graph dispatch;
• request-boundary microbatch slicing;
• ModelRunnerV2 input and model-input slicing;
• per-microbatch attention metadata preparation;
• output merging for tensors and Ascend auxiliary output structures;
• the temporary ModelRunnerV2 execute path.

At startup, the backport checks the upstream BatchExecutionDescriptor shape.
If upstream already provides num_ubatches, initialization fails with an
instruction to remove the compatibility layer rather than shadowing the native
implementation.

Eager DBO execution

AFDNPUAttentionModelRunnerV2.execute_model() uses the backported execute path
only when vLLM reports use_ubatching; otherwise it uses the native
vLLM-Ascend runner.

For a selected DBO batch, it:

1. prepares the native ModelRunnerV2 input batch and model state;
2. splits the padded batch into exactly two request-boundary slices;
3. builds independent Ascend attention metadata, slot mappings, DP metadata,
   and forward contexts for both stages;
4. runs both stages under the existing DBO handoff used by the AFD model proxy;
5. merges the ordered outputs before native sampling and request-state update.

ModelRunnerV2 stores stage-specific Ascend state in
ForwardContext.additional_kwargs. ModelRunnerV1 continues to use its existing
attributes, so this does not remove or redirect V1 state.

DP dispatch and fallback

All DP ranks participate in one CPU-group reduction for token count,
uniform-token state, local DBO eligibility, and requested graph mode. This
keeps the eager or graph and single- or dual-batch decision consistent across
ranks.

DBO is used only when every rank allows it, the configured threshold is met,
and both microbatches still contain real work after final graph padding.
Otherwise, the step uses the native single-batch descriptor. If a rank cannot
use a graph, all ranks run the two microbatches eagerly for that step.

Full ACL graph DBO

During ModelRunnerV2 initialization, a scoped wrapper replaces only the graph
manager factory used by the current AFD runner. finally blocks restore the
original vLLM and vLLM-Ascend symbols.

AFDModelAclGraphManagerV2 keeps the native single-batch graph descriptors and
owns separate two-microbatch twins for eligible, evenly divisible capture
shapes. Each DBO graph includes:

• a warmup execution and a formal capture execution;
• two stage-specific forward contexts and metadata builders;
• per-stage MLA graph parameters backed by the native aggregate FIA workspace;
• staged replay state containing the real runtime token counts;
• AFD control metadata published before Attention and FFN capture/replay;
• the same FULL replay ordering fence required by vLLM-Ascend.

An uncaptured DBO shape uses eager DBO. A captured single-batch shape continues
through the native graph manager.

Compatibility and isolation

• The validation relaxation applies only to active AFD NPU ModelRunnerV2 DBO
  configurations and restores the original config values after upstream
  validation.
• Graph-manager substitution is limited to one runner initialization scope and
  is always restored.
• Full-graph replay hooks are installed on one manager instance and restored
  after each execution scope.
• Non-AFD configurations preserve upstream validation and runner identity.
• DBO-disabled AFD ModelRunnerV2 execution continues to use the native runner.
• ModelRunnerV1 keeps its existing execution, forward-context, DBO, and ACL
  graph ownership.
• Plugin-owned DBO graphs and staged replay state are released during cleanup.

Configuration boundary

NPU ModelRunnerV2 DBO requires:

• vLLM v0.26.0 without native ModelRunnerV2 DBO descriptors;
• the synchronous CAMP2pAFDConnector;
• compute_gate_on_attention=false;
• DP size greater than one;
• exactly two microbatches;
• matching configured AFD ranks and DP * TP;
• PP, PCP, and DCP equal to one;
• static expert parallelism;
• a registered AFD model;
• eager execution, or ACL graph mode FULL_DECODE_ONLY.

Unsupported combinations fail during validation before model execution.

Test plan

Unit and compatibility coverage

• DBO threshold selection and uniform DP padding.
• Single-batch fallback when a rank is below threshold.
• Single-batch fallback when graph padding would create an empty second stage.
• Eager DBO fallback when DP ranks disagree on graph availability.
• Compatibility with the pinned upstream graph-manager dispatch signature.
• Request-boundary slicing, all-padding trailing slices, metadata-builder
  selection, and structured output merging.
• AFD-only configuration-validation relaxation and restoration of original
  values.
• Rejection of unsupported DBO topology, graph mode, and model features.
• ModelRunnerV2 forward-context storage without changing ModelRunnerV1 state.
• Graph-manager/replay-hook scoping, restoration, failure cleanup, and
  non-AFD/native-runner isolation.
• FFN graph replay for repeated capture keys.

Ascend E2E coverage

The hardware E2E matrix includes:

• ModelRunnerV2 eager DBO off versus on;
• ModelRunnerV2 FULL_DECODE_ONLY DBO off versus on;
• ModelRunnerV1 DBO regression coverage;
• deterministic output consistency for the corresponding DBO on/off pairs.

Test result

• The Ascend E2E consistency suite passed for NPU ModelRunnerV2 DBO, including
  eager and FULL_DECODE_ONLY graph execution.
• The NPU ModelRunnerV1 DBO regression check passed on its existing
  implementation path.
• DeepSeek-V2-Lite performance validation showed no DBO performance benefit.

Limitations and follow-up

• This compatibility backport is not the long-term home for generic
  ModelRunnerV2 DBO behavior.
• Current NPU support is restricted to two microbatches and the validated AFD
  topology above.
• Once the pinned vLLM release provides native ModelRunnerV2 DBO descriptors,
  dispatch, slicing, and execution, remove the compatibility package and scoped
  validation and graph-manager patches. Retain only NPU-specific execution and
  ACL graph integration that is still needed.

---

Essential PR checklist

• Eager and FULL_DECODE_ONLY DBO behavior is documented.
• The temporary vLLM v0.26 backport and its removal condition are explicit.
• DP synchronization, fallback, slicing, and graph ownership are described.
• Unsupported topology and model-feature boundaries are stated.
• ModelRunnerV1 and non-AFD isolation are documented.
• Ascend consistency and performance results are stated.
• No upstream vLLM or vLLM-Ascend files are modified.

[jiangkuaixue123]

Thanks for the detailed design and test coverage. I found two blockers before this can merge:

1. The temporary backport copies/adapts upstream execution code but does not follow this repository's patch-marking requirements; see the inline comment.
2. The current head fails the pre-commit workflow. The failing job reports mypy errors in the changed source/test files and the SPDX hook modifies four touched test files. Please run the full pre-commit checks and make the workflow green.

[jiangkuaixue123]

This function is a copied/adapted upstream execution path, but the AFD-specific differences are not marked. The repository guidelines require copied or wrapped upstream functions to have the patch reason/functionality/signature comments immediately above them and to surround only the AFD-specific deltas with # ### PATCH START: ... / # ### PATCH END: .... Please re-copy from the exact pinned source and mark the adaptations so future vLLM upgrades can mechanically compare and reapply this backport. The same applies to the copied helpers in runtime.py.

[lirx-pd]

All the related files have been refactored

[jiangkuaixue123]

A follow-up pass focused specifically on abstraction and defensive programming. Most of the larger boundaries (AFDAscendUBatchRunnerV2, graph-entry state, and the scoped graph-manager context managers) are justified because they isolate substantial lifecycle/state. The overengineering is concentrated in compatibility sentries and optional attribute probing. assert_backport_required() is also a single-use helper that duplicates the module ABI sentry; I suggest removing both and relying on the pinned ABI plus direct accesses/tests.

[jiangkuaixue123]

This module-to-module ABI handshake looks over-defensive. These modules are shipped from the same package/checkout, and there is already a second proactive guard in assert_backport_required(). The getattr(..., 1) fallback also masks the actual missing-symbol failure. Please remove this private ABI protocol (and the single-use descriptor-field sentry) and let the pinned vLLM contract plus normal import/static-test failures expose drift.

[jiangkuaixue123]

Please avoid probing this contract with getattr and replacing the original failure with a custom RuntimeError. The DBO initialization path installs AFDModelAclGraphManagerV2, which defines dispatch_ubatches; call cudagraph_manager.dispatch_ubatches(...) directly (and preferably give the manager a concrete protocol/type). If that contract changes, the original attribute/type failure should remain visible, per the repository's upstream-compatibility guidance.

[jiangkuaixue123]

These hasattr branches silently turn an ABI mismatch into partially initialized graph state. Because this backport targets pinned vLLM/vLLM-Ascend versions, access _get_workspace_buffer() / set_workspace_buffer() directly with the expected builder type and let an upstream incompatibility fail at its source. This removes defensive branching and makes static checking useful.

[hsliuustc0106]

Four actionable findings from review of the current head.

[hsliuustc0106]

[P1] Package the new backport modules

Package discovery sets namespaces = false, but afd_plugin/compat/backports has no __init__.py. I verified that setuptools.find_packages() excludes afd_plugin.compat.backports.vllm_v026_mrv2_dbo; checkout tests pass through implicit namespace imports, while an installed wheel will omit these modules and fail the new imports. Please add the parent package marker or enable namespace discovery, plus a wheel-import smoke test.

[lirx-pd]

Fix this problem

[hsliuustc0106]

[P1] Mark DBO graph replay in control metadata

This early return bypasses the code that sets _afd_is_graph_replaying = True. AFDModelAclGraphManagerV2.run_fullgraph() subsequently calls send_dp_metadata(), so the FFN receives is_graph_replaying=False and selects eager execution rather than its matching captured graph. Please wrap this call with the same replay-flag save/set/restore logic, or set the flag in the plugin manager before it sends control.

[lirx-pd]

Fix this problem

[hsliuustc0106]

[P1] Cancel event waits after a stage failure

Aborting ready_barrier only releases threads still entering their contexts. Once execution begins, stages wait on cpu_wait_event; if one stage fails, its sibling can resume, reach the next DBO yield, and wait forever because the failed stage can no longer signal it. close_execution() then blocks indefinitely in thread.join(). Please add shared cancellation checked by the event-wait loop and wake every stage on failure, with a regression test that fails one stage mid-forward.

[lirx-pd]

Fix this problem

[hsliuustc0106]

[P2] Release ubatch ownership before native shutdown

clear_afd_graphs() only clears graph entries. Both this runner and the graph manager still retain ubatch_runner, which retains model_state; ModelState owns the model itself. Therefore super().shutdown() deleting self.model_state and self.model does not release the weights, defeating same-process memory reclamation. Please break these plugin-owned references before native shutdown.

[lirx-pd]

Fix this problem

[hsliuustc0106]

Independent review found two blocking correctness gaps in the NPU ModelRunnerV2 DBO graph path.

[hsliuustc0106]

[P1] Preserve replay state for AFD DBO descriptors. execute_model resets _afd_is_graph_replaying to false, and this branch delegates before setting it true. AFDModelAclGraphManagerV2.run_fullgraph then sends the control payload with that false value, while the FFN graph policy selects REPLAY only when is_graph_replaying is true. As a result, the Attention DBO graph replays but the matching FFN graph runs eagerly. Please set and restore the replay flag around this delegation and add a cross-role payload/replay regression test.

[lirx-pd]

Fix this problem

[hsliuustc0106]

[P1] Reject unsupported non-MLA DBO graph configurations. A registered Qwen3MoeForCausalLM configuration with use_mla=false, DP=2, two ubatches, and FULL_DECODE_ONLY passes this validator. The new AFDModelAclGraphManagerV2 capture path then unconditionally requires an MLA/FIA workspace and raises. Please either require MLA here for DBO graph mode or provide a backend-generic capture path, with a non-MLA regression test.

[lirx-pd]

Fix this problem