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

[lirx-pd]

Purpose

This change adds temporary Dual Batch Overlap (DBO) support for the AFD NPU
ModelRunnerV2 path while the pinned vLLM v0.26.0 ABI does not provide the
required ModelRunnerV2 DBO implementation.

The implementation adapts the ModelRunnerV2 DBO behavior from
specture724/vllm branch feat/v2/dbo-fullcg at commit 626fee7831 to the
current vLLM ABI at 568afb3a13, and keeps all missing upstream behavior under
an isolated compatibility backport. This allows the backport to be removed once
the pinned vLLM version contains native ModelRunnerV2 DBO support.

The supported AFD NPU ModelRunnerV2 paths are:

• 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 vLLM and vLLM-Ascend source trees are not modified. Existing NPU
ModelRunnerV1 behavior remains on its original 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 backport

The compatibility package
afd_plugin/compat/backports/vllm_v026_mrv2_dbo contains the copied/adapted
execution behavior that is absent 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.

The backport checks the upstream BatchExecutionDescriptor shape at startup. If
upstream already provides num_ubatches, initialization fails with an explicit
instruction to remove this compatibility layer instead of silently shadowing
the new native implementation.

Eager DBO execution

AFDNPUAttentionModelRunnerV2.execute_model() delegates to the backported
execute path only when vLLM reports use_ubatching. Otherwise it continues to
delegate to the native vLLM-Ascend runner.

For a selected DBO batch, the implementation:

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.

The stage-specific Ascend state is stored in ForwardContext.additional_kwargs
for ModelRunnerV2. ModelRunnerV1 continues to use its existing attributes, so
the V2 adaptation does not remove or redirect V1 state.

DP dispatch and fallback

All DP ranks participate in one CPU-group reduction containing their token
count, uniform-token state, local DBO eligibility, and requested graph mode.
The result ensures that every rank makes the same eager/graph and
single-/dual-batch decision.

DBO is selected only when all ranks allow it, the configured threshold is met,
and both microbatches contain real work after final graph padding. Otherwise the
step uses the native single-batch descriptor. If any DP rank cannot use a graph,
all ranks execute 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. The original vLLM and
vLLM-Ascend symbols are restored in finally blocks.

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

• 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 falls back to 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 covers:

• 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 current Ascend E2E consistency suite passed for NPU ModelRunnerV2 DBO,
  including eager and FULL_DECODE_ONLY graph execution.
• The NPU ModelRunnerV1 DBO regression check passed and remains on its existing
  implementation path.
• Performance validation on DeepSeek-V2-Lite showed no DBO performance benefit.

Limitations and Follow-up

• This is a compatibility backport, not the long-term owner of 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 the
  scoped validation/graph-manager patches, then adapt only the NPU-specific
  execution and ACL graph seams that remain necessary.

---

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.