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

afd_plugin/compat/npu/feature_validation.py

@@ -56,12 +56,32 @@ def fail_if_unsupported_npu_afd_features(
             )
         extra_info.validate_supported()
 
-    if bool(vllm_config.parallel_config.use_ubatching) and (
-        int(vllm_config.parallel_config.num_ubatches) != 2
-    ):
+    uses_ubatching = bool(vllm_config.parallel_config.use_ubatching)
+    if uses_ubatching and int(vllm_config.parallel_config.num_ubatches) != 2:
         raise RuntimeError(
             "AFD NPU runtime supports exactly two ubatches when DBO is enabled",
         )
+    model_config = vllm_config.model_config
+    # Match the pinned NPUModelRunner's sparse-attention backend selection.
+    uses_sparse_mla = hasattr(
+        model_config.hf_text_config,
+        "index_topk",
+    )
+    cudagraph_mode = vllm_config.compilation_config.cudagraph_mode
+    uses_mla_dbo_full_graph = (
+        uses_ubatching
+        and model_config.use_mla
+        and not uses_sparse_mla
+        and cudagraph_mode.has_full_cudagraphs()
+    )
+    if uses_mla_dbo_full_graph and vllm_config.speculative_config is not None:
+        raise RuntimeError(
+            "AFD NPU MLA DBO FULL graph does not support speculative decoding",
+        )
+    if uses_mla_dbo_full_graph and cudagraph_mode.name != "FULL_DECODE_ONLY":
+        raise RuntimeError(
+            "AFD NPU MLA DBO graph execution requires FULL_DECODE_ONLY",
+        )
 
 
 def _fail_if_unsupported_npu_afd_async_features(

afd_plugin/compat/npu/runtime.py

@@ -28,9 +28,19 @@ def apply_afd_ascend_patches_if_needed() -> None:
     from afd_plugin.compat.patches.npu.ascend_platform import (
         apply_afd_ascend_dbo_config_patch,
     )
+    from afd_plugin.compat.patches.npu.mla_graph import (
+        apply_afd_mla_graph_patch,
+    )
 
-    if apply_afd_ascend_dbo_config_patch():
-        _PATCHES_APPLIED = True
+    if not apply_afd_ascend_dbo_config_patch():
+        raise RuntimeError(
+            "AFD NPU DBO config patch requires vLLM-Ascend NPUPlatform",
+        )
+    if not apply_afd_mla_graph_patch():
+        raise RuntimeError(
+            "AFD NPU MLA graph patch requires the vLLM-Ascend MLA resolver",
+        )
+    _PATCHES_APPLIED = True
 
 
 __all__ = [

afd_plugin/compat/patches/npu/mla_graph.py

@@ -0,0 +1,58 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Route vLLM-Ascend MLA graph parameters through AFD forward contexts.
+
+Upstream source: ``vllm_ascend/attention/mla_v1.py``.
+"""
+
+from __future__ import annotations
+
+AFD_MLA_GRAPH_PARAMS_KEY = "afd_mla_graph_params"
+_MLA_GRAPH_PATCH_ATTR = "_afd_plugin_mla_graph_patch_state"
+
+
+def apply_afd_mla_graph_patch() -> bool:
+    """Install the MLA resolver, returning whether it is available."""
+
+    try:
+        from vllm.forward_context import (
+            get_forward_context,
+            is_forward_context_available,
+        )
+        from vllm_ascend.attention import mla_v1
+    except ImportError:
+        return False
+
+    if hasattr(mla_v1, _MLA_GRAPH_PATCH_ATTR):
+        return True
+
+    original_get_graph_params = mla_v1.get_graph_params
+
+    # Patch reason: AFD aggregate graph capture owns one GraphParams per ubatch.
+    # Patch functionality: resolve graph params from the active AFD forward
+    # context while preserving upstream process-global behavior otherwise.
+    # Signature: matches upstream; no added parameters.
+    def get_graph_params():
+        # ### PATCH START: AFD MLA graph registry
+        if is_forward_context_available():
+            forward_context = get_forward_context()
+            additional_kwargs = forward_context.additional_kwargs or {}
+            graph_params = additional_kwargs.get(AFD_MLA_GRAPH_PARAMS_KEY)
+            if graph_params is not None:
+                return graph_params
+        # ### PATCH END: AFD MLA graph registry
+        return original_get_graph_params()
+
+    mla_v1.get_graph_params = get_graph_params
+    setattr(
+        mla_v1,
+        _MLA_GRAPH_PATCH_ATTR,
+        original_get_graph_params,
+    )
+    return True
+
+
+__all__ = [
+    "AFD_MLA_GRAPH_PARAMS_KEY",
+    "apply_afd_mla_graph_patch",
+]

afd_plugin/v1/worker/npu/attention_model_runner.py

@@ -179,16 +179,18 @@ def _model_forward(self, *args: Any, **kwargs: Any) -> Any:
             **model_kwargs,
         }
         run_model = partial(self.model, **model_inputs)
+        wrapper_owns_full_graph_update = isinstance(
+            self.model, AscendUBatchWrapper
+        ) and self.model.owns_full_graph_update(forward_context)
 
-        if self.enable_enpu:
+        if self.enable_enpu and not wrapper_owns_full_graph_update:
             self._update_full_graph_params_if_needed(
                 forward_context,
                 num_tokens_padded,
                 positions,
             )
-            hidden_states = run_model()
-        else:
-            hidden_states = run_model()
+        hidden_states = run_model()
+        if not self.enable_enpu and not wrapper_owns_full_graph_update:
             self._update_full_graph_params_if_needed(
                 forward_context,
                 num_tokens_padded,
@@ -1347,6 +1349,11 @@ def _install_ascend_ubatch_wrapper(self) -> None:
             self.vllm_config,
             runtime_mode,
             self.device,
+            mla_full_graph_enabled=(
+                self.vllm_config.model_config.use_mla and not self.use_sparse
+            ),
+            full_graph_params_updater=self._update_full_graph_params_if_needed,
+            enable_enpu=self.enable_enpu,
         )
 
     def get_model(self) -> Any:

afd_plugin/v1/worker/npu/forward_context.py

@@ -2,7 +2,10 @@
 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
 """Forward-context helpers for plugin-owned Ascend ubatching."""
 
+from __future__ import annotations
+
 import math
+from typing import TYPE_CHECKING
 
 import torch
 from vllm.config import CUDAGraphMode, VllmConfig
@@ -11,11 +14,17 @@
 from vllm.v1.worker.ubatch_utils import UBatchSlices
 from vllm_ascend.ops.fused_moe.moe_comm_method import get_moe_comm_method
 
+from afd_plugin.compat.patches.npu.mla_graph import (
+    AFD_MLA_GRAPH_PARAMS_KEY,
+)
 from afd_plugin.v1.worker.ubatch_wrapper import (
     build_ubatch_additional_kwargs,
     build_ubatch_afd_metadata,
 )
 
+if TYPE_CHECKING:
+    from vllm_ascend.compilation.acl_graph import GraphParams
+
 
 def create_ascend_forward_context(
     cur_forward_context: ForwardContext,
@@ -27,6 +36,7 @@ def create_ascend_forward_context(
     cudagraph_runtime_mode: CUDAGraphMode | None = None,
     batch_descriptor: BatchDescriptor | None = None,
     skip_compiled: bool = False,
+    mla_graph_params: GraphParams | None = None,
 ) -> ForwardContext:
     if cudagraph_runtime_mode is None:
         cudagraph_runtime_mode = CUDAGraphMode.NONE
@@ -38,6 +48,8 @@ def create_ascend_forward_context(
             parent_kwargs,
             build_ubatch_afd_metadata(afd_metadata, ubatch_slices, ubatch_num),
         )
+    if mla_graph_params is not None:
+        parent_kwargs[AFD_MLA_GRAPH_PARAMS_KEY] = mla_graph_params
 
     new_forward_context = ForwardContext(
         no_compile_layers=vllm_config.compilation_config.static_forward_context,
@@ -62,7 +74,9 @@ def create_ascend_forward_context(
         new_forward_context.moe_comm_type
     )
     new_forward_context.in_profile_run = cur_forward_context.in_profile_run
-    new_forward_context.capturing = cur_forward_context.capturing
+    new_forward_context.capturing = (
+        mla_graph_params is not None or cur_forward_context.capturing
+    )
     new_forward_context.mmrs_fusion = cur_forward_context.mmrs_fusion
     new_forward_context.num_tokens = num_tokens
     new_forward_context.ubatch_idx = int(ubatch_num)

afd_plugin/v1/worker/npu/mla_graph.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Graph-local MLA parameter helpers for Ascend ubatch execution."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, TypeVar
+
+from afd_plugin.compat.patches.npu.mla_graph import (
+    AFD_MLA_GRAPH_PARAMS_KEY,
+)
+
+if TYPE_CHECKING:
+    import torch
+    from vllm.forward_context import ForwardContext
+    from vllm_ascend.compilation.acl_graph import GraphParams
+
+MetadataT = TypeVar("MetadataT")
+
+
+def new_mla_graph_params(
+    num_tokens: int,
+    workspace: torch.Tensor,
+) -> GraphParams:
+    """Create an empty MLA registry for one token shape and FIA workspace."""
+    # Delay the NPU-only import so this helper remains CPU-import-safe.
+    from vllm_ascend.compilation.acl_graph import GraphParams
+
+    return GraphParams(
+        events={num_tokens: []},
+        workspaces={num_tokens: workspace},
+        handles={num_tokens: []},
+        attn_params={num_tokens: []},
+    )
+
+
+def merge_mla_graph_params(
+    attn_metadata: list[dict[str, MetadataT]],
+    graph_params: tuple[GraphParams, GraphParams],
+    num_tokens: int,
+) -> tuple[dict[tuple[str, int], MetadataT], GraphParams]:
+    """Validate and merge two stage-local MLA registries for graph replay."""
+    if len(attn_metadata) != 2 or len(graph_params) != 2:
+        raise RuntimeError(
+            "MLA DBO FULL graph requires exactly two metadata stages; "
+            f"got metadata={len(attn_metadata)}, params={len(graph_params)}",
+        )
+
+    layer_keys = tuple(attn_metadata[0])
+    if tuple(attn_metadata[1]) != layer_keys:
+        raise RuntimeError("MLA DBO FULL graph layer order differs by stage")
+
+    expected_records = len(layer_keys)
+    for stage_index, params in enumerate(graph_params):
+        record_counts = (
+            len(params.events[num_tokens]),
+            len(params.handles[num_tokens]),
+            len(params.attn_params[num_tokens]),
+        )
+        if record_counts != (
+            expected_records,
+            expected_records,
+            expected_records,
+        ):
+            raise RuntimeError(
+                "MLA DBO FULL graph record count mismatch for "
+                f"stage {stage_index}: metadata={expected_records}, "
+                f"events={record_counts[0]}, handles={record_counts[1]}, "
+                f"params={record_counts[2]}",
+            )
+
+    workspace = graph_params[0].workspaces[num_tokens]
+    if graph_params[1].workspaces[num_tokens] is not workspace:
+        raise RuntimeError(
+            "MLA DBO FULL graph requires one shared FIA workspace",
+        )
+
+    merged_metadata: dict[tuple[str, int], MetadataT] = {}
+    merged = new_mla_graph_params(num_tokens, workspace)
+
+    # Preserve the updater's layer-major, stage-minor record order.
+    for layer_index, layer_key in enumerate(layer_keys):
+        for stage_index in range(2):
+            params = graph_params[stage_index]
+            merged_key = (layer_key, stage_index)
+            merged_metadata[merged_key] = attn_metadata[stage_index][layer_key]
+            merged.events[num_tokens].append(
+                params.events[num_tokens][layer_index],
+            )
+            merged.handles[num_tokens].append(
+                params.handles[num_tokens][layer_index],
+            )
+            merged.attn_params[num_tokens].append(
+                params.attn_params[num_tokens][layer_index],
+            )
+
+    return merged_metadata, merged
+
+
+@contextmanager
+def override_mla_graph_params(
+    forward_context: ForwardContext,
+    attn_metadata: dict[tuple[str, int], MetadataT],
+    graph_params: GraphParams,
+) -> Iterator[None]:
+    """Temporarily expose merged MLA state to the upstream graph updater."""
+    original_metadata = forward_context.attn_metadata
+    original_additional_kwargs = forward_context.additional_kwargs
+    # Do not mutate the parent context's shared additional_kwargs mapping.
+    temporary_kwargs = dict(original_additional_kwargs or {})
+    temporary_kwargs[AFD_MLA_GRAPH_PARAMS_KEY] = graph_params
+    forward_context.attn_metadata = attn_metadata
+    forward_context.additional_kwargs = temporary_kwargs
+    try:
+        yield
+    finally:
+        forward_context.attn_metadata = original_metadata
+        forward_context.additional_kwargs = original_additional_kwargs
+
+
+__all__ = [
+    "merge_mla_graph_params",
+    "new_mla_graph_params",
+    "override_mla_graph_params",
+]