[NPU] Enable native target and DSpark draft Attention graphs

afd_plugin/compat/npu/feature_validation.py

@@ -14,7 +14,11 @@
     is_afd_async_dp,
     parse_afd_config,
 )
-from afd_plugin.validation import cam_receive_capacity, requested_ffn_mode
+from afd_plugin.validation import (
+    cam_receive_capacity,
+    requested_ffn_mode,
+    validate_async_cam_attention_graph,
+)
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -231,14 +235,13 @@ def _fail_if_unsupported_npu_afd_async_features(
         )
     if not bool(vllm_config.model_config.enforce_eager):
         if afd_config.is_attention_server:
-            raise RuntimeError(
-                "CAMAsyncAFDConnector supports only eager Attention execution"
-            )
-        requested_ffn_mode(vllm_config)
-        if not vllm_config.use_v2_model_runner:
-            raise RuntimeError("Async FFN FULL requires MRV2")
-        if parallel_config.tensor_parallel_size != 1:
-            raise RuntimeError("Async FFN FULL requires F TP=1")
+            validate_async_cam_attention_graph(vllm_config)
+        else:
+            requested_ffn_mode(vllm_config)
+            if not vllm_config.use_v2_model_runner:
+                raise RuntimeError("Async FFN FULL requires MRV2")
+            if parallel_config.tensor_parallel_size != 1:
+                raise RuntimeError("Async FFN FULL requires F TP=1")
     if bool(parallel_config.enable_dbo) or bool(parallel_config.use_ubatching):
         if not afd_config.is_attention_server:
             raise RuntimeError("Async CAM DBO/ubatching requires the Attention role")
@@ -262,26 +265,32 @@ def _fail_if_unsupported_npu_afd_async_features(
                 raise RuntimeError("AFD DSpark supports only DeepSeek-V4 on NPU")
             # The shared MRV2 validator owns PP/CP and static EP constraints.
             # Keep only the additional target/draft restrictions here.
-            # F may replay its expert transaction in FULL; Attention and native
-            # draft proposal execution stay eager.
+            # Both graphs use the native decode path. Draft capture is owned
+            # by native target capture, so draft-only graph mode is unsupported.
+            # F keeps its independent TP1 expert execution.
             if (
-                (
-                    afd_config.is_attention_server
-                    and not vllm_config.model_config.enforce_eager
+                speculative_config.draft_parallel_config.tensor_parallel_size
+                != parallel_config.tensor_parallel_size
+                or (
+                    afd_config.is_ffn_server
+                    and parallel_config.tensor_parallel_size != 1
                 )
-                or not speculative_config.enforce_eager
-                or parallel_config.tensor_parallel_size != 1
-                or speculative_config.draft_parallel_config.tensor_parallel_size != 1
             ):
                 raise RuntimeError(
-                    "AFD DSpark requires eager Attention/draft and target/draft TP=1"
+                    "AFD DSpark requires matching target/draft TP and F TP=1"
                 )
+            if (
+                afd_config.is_attention_server
+                and not speculative_config.enforce_eager
+                and vllm_config.model_config.enforce_eager
+            ):
+                raise RuntimeError("DSpark draft Graph requires target Attention graph")
         elif afd_config.is_attention_server and (
             speculative_config.use_eagle() or speculative_config.uses_draft_model()
         ):
             raise RuntimeError(
                 "CAMAsyncAFDConnector does not support native draft models or MTP; "
-                "only DSV4 DSpark eager is supported",
+                "only DSV4 DSpark is supported",
             )
 
     factor_value = os.environ.get("BATCH_SIZE_FACTOR", "1")

afd_plugin/v1/worker/attention_model_runner_v2.py

@@ -197,8 +197,8 @@ def _use_afd_capture_input_preparation(
 
     # Patch reason: native ModelCudaGraphManager.capture calls this exact
     # module symbol once before each warmup/formal-capture forward.
-    # Patch functionality: preserve native preparation and publish the
-    # matching AFD event before the forward starts.
+    # Patch functionality: preserve native preparation and stage the matching
+    # AFD metadata before the forward; synchronous connectors also send control.
     # Signature: matches vLLM v0.32.0 prepare_inputs_to_capture exactly,
     # including the trailing pcp_manager parameter.
     # Upstream source: vllm/v1/worker/gpu/cudagraph_utils.py,
@@ -237,10 +237,11 @@ def prepare_inputs_to_capture(
             None,
             int(desc.num_tokens),
         )
-        runner.send_dp_metadata(
-            runner.build_capture_dp_metadata(int(desc.num_tokens)),
-            None,
-        )
+        if runner.connector.control_plane is not None:
+            runner.send_dp_metadata(
+                runner.build_capture_dp_metadata(int(desc.num_tokens)),
+                None,
+            )
         runner._afd_suppress_metadata_send = True
         return attention_state
         # ### PATCH END: stage one exact AFD capture event.

afd_plugin/v1/worker/npu/attention_model_runner_v2.py

@@ -82,10 +82,10 @@ def _use_afd_fullgraph_replay_hook(
     original_run_fullgraph = manager.run_fullgraph
     manager_state[_AFD_FULLGRAPH_HOOK_MARKER] = True
 
-    # Patch reason: native FULL replay bypasses ForwardContext creation, so the
-    # execute-scoped AFD provider cannot publish runtime control.
-    # Patch functionality: wrap only this manager instance and publish one
-    # ordinary padded control payload immediately before each native replay.
+    # Patch reason: native FULL replay bypasses the ordinary forward, so its
+    # attention-parameter update needs fresh AFD real/padded token metadata.
+    # Patch functionality: stage that metadata on this manager's replay scope;
+    # synchronous connectors also publish one padded control payload.
     # Signature: matches vLLM v0.32.0 ModelCudaGraphManager.run_fullgraph exactly.
     # Upstream source: vllm/v1/worker/gpu/cudagraph_utils.py,
     # ModelCudaGraphManager.run_fullgraph; commit
@@ -119,10 +119,11 @@ def run_fullgraph(
             runner._is_warmup = False
             runner._afd_is_graph_capturing = False
             runner._afd_is_graph_replaying = True
-            runner.send_dp_metadata(
-                runner.build_capture_dp_metadata(padded_tokens),
-                None,
-            )
+            if runner.connector.control_plane is not None:
+                runner.send_dp_metadata(
+                    runner.build_capture_dp_metadata(padded_tokens),
+                    None,
+                )
             result = original_run_fullgraph(desc)
         finally:
             runner._afd_is_graph_replaying = previous_is_graph_replaying
@@ -269,6 +270,7 @@ def __init__(
                     self.vllm_config,
                     self.device,
                     afd_context=self._afd_draft_context,
+                    afd_runner=self,
                 )
                 self.speculator.update_stream = self.update_stream
         except BaseException:

afd_plugin/v1/worker/npu/dspark_graph.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project
+"""Install AFD metadata around native draft graph capture and replay."""
+
+from __future__ import annotations
+
+from collections.abc import Iterator, Mapping
+from contextlib import contextmanager
+from types import MethodType
+from typing import TYPE_CHECKING
+
+import torch
+from vllm.config import CUDAGraphMode
+from vllm.v1.kv_cache_interface import KVCacheConfig
+from vllm.v1.worker.gpu.block_table import BlockTables
+from vllm.v1.worker.gpu.cudagraph_utils import (
+    AttentionState,
+    BatchExecutionDescriptor,
+)
+from vllm.v1.worker.gpu.input_batch import InputBuffers
+from vllm.v1.worker.gpu.spec_decode.dflash import cudagraph as dflash_cudagraph
+from vllm.v1.worker.utils import AttentionGroup
+from vllm_ascend.worker.v2.spec_decode.dflash.aclgraph import DFlashAclGraphManager
+
+from afd_plugin.v1.worker.attention_model_runner_v2 import _AFDCaptureEventTracker
+
+if TYPE_CHECKING:
+    from afd_plugin.v1.worker.npu.attention_model_runner_v2 import (
+        AFDNPUAttentionModelRunnerV2,
+    )
+
+
+@contextmanager
+def use_afd_draft_capture(
+    runner: AFDNPUAttentionModelRunnerV2,
+    manager: DFlashAclGraphManager,
+) -> Iterator[None]:
+    """Associate every native draft capture forward with its own AFD batch."""
+    descs = manager._capture_descs.get(CUDAGraphMode.FULL, [])[
+        : manager._max_full_descs_to_capture
+    ]
+    if not runner.speculative_config.enforce_eager and not descs:
+        raise RuntimeError("AFD DSpark draft FULL has no native capture descriptors")
+    tracker = _AFDCaptureEventTracker(descs)
+    original_prepare = dflash_cudagraph._prepare_dflash_inputs_to_capture
+
+    # Upstream: vLLM 5df63a72, dflash/cudagraph.py; Ascend bd42eea0 delegates
+    # draft capture here. Signature matches the native preparation function.
+    # Native input/KV preparation remains in the saved function; only AFD
+    # metadata is added. Remove when draft capture exposes a plugin callback.
+    def prepare_inputs(
+        num_reqs: int,
+        num_tokens: int,
+        input_buffers: InputBuffers,
+        block_tables: BlockTables,
+        attn_groups: list[list[AttentionGroup]],
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        skip_attn: bool,
+        causal: bool | Mapping[int, bool],
+    ) -> AttentionState:
+        # ### PATCH START: stage descriptor-local draft metadata.
+        state = original_prepare(
+            num_reqs,
+            num_tokens,
+            input_buffers,
+            block_tables,
+            attn_groups,
+            kv_cache_config,
+            max_model_len,
+            skip_attn,
+            causal,
+        )
+        desc, is_warmup = tracker.consume(num_reqs, num_tokens)
+        runner._afd_pending_metadata = runner.build_afd_metadata(None, desc.num_tokens)
+        runner._afd_suppress_metadata_send = True
+        runner._is_warmup = is_warmup
+        runner._afd_is_graph_capturing = not is_warmup
+        runner._afd_is_graph_replaying = False
+        return state
+        # ### PATCH END: stage descriptor-local draft metadata.
+
+    previous_capturing = runner._afd_is_graph_capturing
+    previous_replaying = runner._afd_is_graph_replaying
+    dflash_cudagraph._prepare_dflash_inputs_to_capture = prepare_inputs
+    try:
+        yield
+        tracker.assert_complete()
+    finally:
+        dflash_cudagraph._prepare_dflash_inputs_to_capture = original_prepare
+        runner._afd_is_graph_capturing = previous_capturing
+        runner._afd_is_graph_replaying = previous_replaying
+
+
+@contextmanager
+def use_afd_draft_replay(
+    runner: AFDNPUAttentionModelRunnerV2,
+    manager: DFlashAclGraphManager,
+    real_tokens: int,
+) -> Iterator[None]:
+    """Preserve native replay while staging the draft's real and padded rows."""
+    state = manager.__dict__
+    had_override = "run_fullgraph" in state
+    previous_override = state.get("run_fullgraph")
+    original_replay = manager.run_fullgraph
+
+    # Upstream: Ascend bd42eea0 DFlashAclGraphManager.run_fullgraph.
+    # Signature matches upstream. Delegate graph replay and attention updates;
+    # add only AFD's model-local metadata. Remove when the native manager has
+    # an instance-local pre-replay callback.
+    def run_fullgraph(
+        self,
+        desc: BatchExecutionDescriptor,
+    ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]:
+        # ### PATCH START: install draft-local replay metadata.
+        previous = (
+            runner._afd_pending_metadata,
+            runner._afd_suppress_metadata_send,
+            runner._is_warmup,
+            runner._afd_is_graph_capturing,
+            runner._afd_is_graph_replaying,
+        )
+        try:
+            metadata = runner.build_afd_metadata(None, real_tokens)
+            metadata.tokens_lens = [desc.num_tokens]
+            runner._afd_pending_metadata = metadata
+            runner._afd_suppress_metadata_send = True
+            runner._is_warmup = False
+            runner._afd_is_graph_capturing = False
+            runner._afd_is_graph_replaying = True
+            return original_replay(desc)
+        finally:
+            (
+                runner._afd_pending_metadata,
+                runner._afd_suppress_metadata_send,
+                runner._is_warmup,
+                runner._afd_is_graph_capturing,
+                runner._afd_is_graph_replaying,
+            ) = previous
+        # ### PATCH END: install draft-local replay metadata.
+
+    try:
+        manager.run_fullgraph = MethodType(run_fullgraph, manager)
+        yield
+    finally:
+        if had_override:
+            manager.run_fullgraph = previous_override
+        else:
+            del manager.run_fullgraph