diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 73edc9990..80f8de270 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/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") diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 82eefc156..3ff495685 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/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. diff --git a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py index 1c2b8214d..8f791cca7 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/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: diff --git a/afd_plugin/v1/worker/npu/dspark_graph.py b/afd_plugin/v1/worker/npu/dspark_graph.py new file mode 100644 index 000000000..e17e1990e --- /dev/null +++ b/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 diff --git a/afd_plugin/v1/worker/npu/dspark_speculator.py b/afd_plugin/v1/worker/npu/dspark_speculator.py index 0de54a59b..41940c87e 100644 --- a/afd_plugin/v1/worker/npu/dspark_speculator.py +++ b/afd_plugin/v1/worker/npu/dspark_speculator.py @@ -4,19 +4,30 @@ from __future__ import annotations -from collections.abc import Callable -from contextlib import AbstractContextManager -from typing import Any +from collections.abc import Callable, Iterator +from contextlib import AbstractContextManager, contextmanager, nullcontext +from typing import TYPE_CHECKING, Any, cast import torch from vllm.config import VllmConfig from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.spec_decode.dflash import speculator as dflash_speculator +from vllm_ascend.worker.v2.spec_decode.dflash.aclgraph import DFlashAclGraphManager from vllm_ascend.worker.v2.spec_decode.dflash.speculator import ( prepare_dflash_inputs_factory, ) from vllm_ascend.worker.v2.spec_decode.dspark.speculator import AscendDSparkSpeculator +from afd_plugin.v1.worker.npu.dspark_graph import ( + use_afd_draft_capture, + use_afd_draft_replay, +) + +if TYPE_CHECKING: + from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( + AFDNPUAttentionModelRunnerV2, + ) + class AFDDSparkSpeculator(AscendDSparkSpeculator): def __init__( @@ -25,9 +36,37 @@ def __init__( device: torch.device, *, afd_context: Callable[[bool, bool], AbstractContextManager[None]], + afd_runner: AFDNPUAttentionModelRunnerV2, ) -> None: super().__init__(vllm_config, device) self.afd_context = afd_context + self.afd_runner = afd_runner + + @contextmanager + def _physical_input_preparation(self, *, enabled: bool = True) -> Iterator[None]: + original_prepare = dflash_speculator.prepare_dflash_inputs + if enabled: + dflash_speculator.prepare_dflash_inputs = self._afd_prepare_inputs + try: + yield + finally: + dflash_speculator.prepare_dflash_inputs = original_prepare + + # Upstream: vLLM 5df63a72 DFlashSpeculator.capture, inherited by Ascend + # bd42eea0 DSpark. Signature matches. Delegate native context/KV and Markov + # capture, adding draft-local AFD metadata and the physical block size. + # Remove when native capture exposes metadata and physical-size hooks. + def capture(self) -> None: + # ### PATCH START: scope native draft capture. + assert self.query_cudagraph_manager is not None + manager = cast(DFlashAclGraphManager, self.query_cudagraph_manager) + with ( + self.afd_context(True, False), + self._physical_input_preparation(), + use_afd_draft_capture(self.afd_runner, manager), + ): + super().capture() + # ### PATCH END: scope native draft capture. # Upstream: Ascend bd42eea0, AscendDSparkSpeculator.set_attn. # Its input factory captures the global cache_config.block_size, which can @@ -88,30 +127,41 @@ def propose( num_speculative_tokens: int | None = None, ) -> torch.Tensor: # ### PATCH START: scope only the native draft proposal - prepare_inputs = dflash_speculator.prepare_dflash_inputs - if not (dummy_run and skip_attn_for_dummy_run): - dflash_speculator.prepare_dflash_inputs = self._afd_prepare_inputs - try: - with self.afd_context(dummy_run, is_profile): - return super().propose( - input_batch, - attn_metadata, - slot_mappings, - last_hidden_states, - aux_hidden_states, - num_sampled, - num_rejected, - last_sampled, - next_prefill_tokens, - temperature, - seeds, - None, - dummy_run, - skip_attn_for_dummy_run, - mm_inputs, - is_profile=is_profile, - num_speculative_tokens=num_speculative_tokens, - ) - finally: - dflash_speculator.prepare_dflash_inputs = prepare_inputs + manager = self.query_cudagraph_manager + # Native memory profiling runs before KV/graph-manager initialization. + replay_scope = ( + use_afd_draft_replay( + self.afd_runner, + cast(DFlashAclGraphManager, manager), + input_batch.num_reqs * self.num_query_per_req, + ) + if manager is not None + else nullcontext() + ) + with ( + self.afd_context(dummy_run, is_profile), + self._physical_input_preparation( + enabled=not (dummy_run and skip_attn_for_dummy_run) + ), + replay_scope, + ): + return super().propose( + input_batch, + attn_metadata, + slot_mappings, + last_hidden_states, + aux_hidden_states, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + temperature, + seeds, + None, + dummy_run, + skip_attn_for_dummy_run, + mm_inputs, + is_profile=is_profile, + num_speculative_tokens=num_speculative_tokens, + ) # ### PATCH END: scope only the native draft proposal diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index 7d586b3a5..b2ef9b0eb 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -89,6 +89,25 @@ def validate_graph_capacity( raise ValueError("Async CAM graph source rows exceed the uint16 protocol bound") +def validate_async_cam_attention_graph(vllm_config: VllmConfig) -> None: + """Validate AFD dependencies for native Async CAM Attention graphs.""" + model = vllm_config.model_config + parallel = vllm_config.parallel_config + if ( + not vllm_config.use_v2_model_runner + or model.hf_text_config.model_type != "deepseek_v4" + ): + raise RuntimeError("Async CAM Attention graph requires MRV2 DeepSeek-V4") + if parallel.enable_dbo or parallel.use_ubatching: + raise RuntimeError("Async CAM Attention graph does not support DBO/ubatching") + if ( + parallel.pipeline_parallel_size, + parallel.prefill_context_parallel_size, + parallel.decode_context_parallel_size, + ) != (1, 1, 1): + raise RuntimeError("Async CAM Attention graph does not support PP or CP") + + def validate_gpu_model_runner_v2_config( vllm_config: VllmConfig, *, @@ -207,11 +226,9 @@ def validate_npu_model_runner_v2_config( ) if not vllm_config.model_config.enforce_eager: if expected_role == "attention": - raise RuntimeError( - "AFD NPU ModelRunnerV2 Async CAM requires eager execution " - "on Attention" - ) - requested_ffn_mode(vllm_config) + validate_async_cam_attention_graph(vllm_config) + else: + requested_ffn_mode(vllm_config) elif afd_config.connector != CAMP2P_CONNECTOR: raise RuntimeError( "AFD NPU ModelRunnerV2 requires either synchronous CAMP2pAFDConnector " diff --git a/docs/design/module/attention_runtime.md b/docs/design/module/attention_runtime.md index ae6b48f28..57188eeda 100644 --- a/docs/design/module/attention_runtime.md +++ b/docs/design/module/attention_runtime.md @@ -13,6 +13,7 @@ primary_code_paths: - "afd_plugin/v1/worker/npu/attention_model_runner_v2.py" - "afd_plugin/v1/worker/npu/attention_worker.py" - "afd_plugin/v1/worker/npu/dspark_speculator.py" + - "afd_plugin/v1/worker/npu/dspark_graph.py" related_code_paths: - "afd_plugin/connectors/**" - "afd_plugin/model_executor/**" @@ -24,6 +25,7 @@ depends_on: - "execution_platforms.md" - "compatibility_and_patches.md" validation_paths: + - "tests/unit/v1/worker/npu/test_dspark_graph.py" - "tests/unit/v1/worker/test_attention_model_runner.py" - "tests/unit/v1/worker/test_model_runner_v2.py" - "tests/unit/v1/worker/test_npu_runtime.py" @@ -166,7 +168,7 @@ Both platforms require V2. The paired deployment constraints are: | Excluded features | Elastic EP, EPLB, sequence-parallel MoE, and compile SP | Elastic EP and EPLB; model-specific SP restrictions | Elastic EP, EPLB, native draft models/MTP except DSV4 DSpark, finegrained TP, and explicit shared-expert DP | | Sequence parallelism | Unsupported | Model-specific; synchronous DeepSeek-V2/V3 SP is rejected | Attention-local FlashComm1/SP supported; disabled on FFN; no PCP/DCP | | Native DBO | Exactly two microbatches, Attention DP1 or DP > 1 | Eager or graph; exactly two microbatches, Attention DP > 1 | Eager only; exactly two microbatches at Attention DP1 or DP > 1 | -| Graph execution | Eager or `FULL_DECODE_ONLY` | Eager, `FULL`, or `FULL_DECODE_ONLY`; plain MLA graph DBO requires `FULL_DECODE_ONLY` | Eager Attention; independent FFN also supports the narrow W4A8/W8A8 FULL contract | +| Graph execution | Eager or `FULL_DECODE_ONLY` | Eager, `FULL`, or `FULL_DECODE_ONLY`; plain MLA graph DBO requires `FULL_DECODE_ONLY` | Eager Attention or DSV4 target native graph without DBO; native `FULL`/`FULL_DECODE_ONLY` resolution; Attention-local SP and native PD may remain enabled; DSpark draft may be eager or native Graph; independent FFN supports the narrow W4A8/W8A8 FULL contract | | Model | Registered AFD architecture; DeepEP M2N is limited to Qwen3 MoE or DeepSeek-V2 | Registered AFD architecture | DeepSeek-V4 with layered W4A8/W8A8 FFN experts and `dynamicQuant=1` | Both roles validate the paired V2 topology, but only Attention uses the native @@ -335,12 +337,14 @@ shutdown; Ascend and other models keep native preparation and profiling. Ascend V2 supports exactly two native microbatches: synchronous CAMP2P requires Attention DP > 1 and supports graph DBO, while Async CAM uses eager -Attention and also supports Attention DP1. Plain MLA graph DBO requires +Attention for DBO and also supports Attention DP1. Its separate target +DSV4 decode graph profile excludes DBO. Plain MLA graph DBO requires `FULL_DECODE_ONLY`. -For full-graph capture, the V2 runner -wraps native input preparation and publishes one warmup and one capture -payload per descriptor. CUDA two-stage preparation installs separate AFD +For full-graph capture, the V2 runner wraps native input preparation and +stages descriptor-matched AFD metadata. Connectors with a control plane also +publish one warmup and one capture payload per descriptor; Async CAM does not. +CUDA two-stage preparation installs separate AFD metadata on both native stage contexts and sends their complete layout before eager execution or FULL replay. Single-stage FULL replay publishes its padded control shape through an execute-scoped manager hook; two-stage replay reuses diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md index b52d2ade5..1992482ca 100644 --- a/docs/design/module/execution_platforms.md +++ b/docs/design/module/execution_platforms.md @@ -285,11 +285,12 @@ groups and also keeps Attention TP peers aligned. ### ACL Graph and NPU Graph NPU Attention follows the upstream ACL graph dispatcher while adding AFD -metadata and control-plane coordination. Native MRV2 owns single/two-stage -graph descriptors, warmup/capture, persistent outputs and replay. AFD retains -formal per-stage MLA registries by the complete native descriptor, including -LoRA specialization, and publishes one matching FFN control payload before -each warmup, capture or live replay. +metadata and, for synchronous CAMP2P, control-plane coordination. Native MRV2 +owns single/two-stage graph descriptors, warmup/capture, persistent outputs and +replay. AFD retains formal per-stage MLA registries by the complete native descriptor, including +LoRA specialization. Synchronous CAMP2P publishes one matching FFN control +payload before each warmup, capture or live replay; Async CAM stages local +AFD metadata without sending DP control messages. For MLA DBO full graphs, each stage records its own upstream `GraphParams`. `merge_mla_graph_params()` validates identical layer order and record counts, @@ -305,12 +306,16 @@ register no FIA workspace, so those models, DeepSeek V4 among them, take the generic two-stage path with no MLA registries. With synchronous CAMP2P, the NPU V2 runner supports eager, `FULL`, and -`FULL_DECODE_ONLY`; Async CAM Attention is eager-only. Like CUDA V2, -it publishes descriptor-matched warmup/capture control outside formal graph -capture and installs an instance-scoped pre-replay hook because native full -replay creates no `ForwardContext`. Native `NPUUBatchRunner` owns preparation, -execution and merge for eager and graph DBO. Plain MLA graph DBO requires -`FULL_DECODE_ONLY`; SFA/DSA +`FULL_DECODE_ONLY`; Async CAM additionally admits DSV4 target native +`FULL` or `FULL_DECODE_ONLY`, with native mode resolution and no DBO. DSpark +draft may be eager or native Graph. Attention-local FlashComm1/SP and native PD +can remain enabled. Device qualification uses compilation mode 0. +Synchronous CAMP2P publishes descriptor-matched warmup/capture control outside +formal graph capture. Async CAM has no control plane and only stages local +AFD metadata. Both install an instance-scoped pre-replay hook because native +full replay creates no `ForwardContext`. Native `NPUUBatchRunner` owns +preparation, execution and merge for eager and graph DBO. Plain MLA graph DBO +requires `FULL_DECODE_ONLY`; SFA/DSA accepted modes still require model-specific qualification. Native two-stage capture shares the ordinary descriptor's FIA workspace; replay merges captured records with current metadata for Ascend updates. Each child carries its actual stage ID for the opaque CAMP2P HCCL @@ -318,13 +323,13 @@ group lookup. Async CAM enables eager native DBO at DP1; it has no separate MoE stage pipeline. Async FFN FULL accepts eager Attention DBO; FFN stays outside Attention's native DBO thread schedule. -The FFN runner owns a separate ACL graph cache keyed by stage token counts and -A/F topology. Warmup runs the eager FFN path. Formal capture updates connector -state through `connector.control_plane` before entering `torch.npu.graph(...)`, +The synchronous FFN runner owns a separate ACL graph cache keyed by stage token +counts and A/F topology. Warmup runs the eager FFN path. Formal capture updates +connector state through `connector.control_plane` before entering `torch.npu.graph(...)`, so replay contains only model and data-plane operations. Replay additionally requires the control payload's `is_graph_replaying=true`; an unknown key or eager payload executes eagerly. -CAM async has no control plane and uses the independent -[Async FFN runtime](../npu/async_ffn_runtime.md). +Async CAM has no control plane and uses the separate complete-transaction +graph described in the [Async FFN runtime guide](../npu/async_ffn_runtime.md). ### Ascend native operators and packaging diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 236f49611..7af50ca4a 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -85,7 +85,7 @@ registers lazy AFD wrapper paths under `AFD`-prefixed aliases. | `DeepseekV3ForCausalLM` | `AFDDeepseekV3ForCausalLM` | `AFDDeepseekV3ForCausalLM` | | `DeepseekV32ForCausalLM` | `AFDDeepseekV32ForCausalLM` | `AFDDeepseekV3ForCausalLM` | | `DeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | `AFDDeepseekV4ForCausalLM` | -| `DSparkDraftModel` | Unsupported | `AFDDSparkDeepseekV4ForCausalLM` (Async CAM: eager A, eager/FULL F) | +| `DSparkDraftModel` | Unsupported | `AFDDSparkDeepseekV4ForCausalLM` (Async CAM: eager/native decode Graph A, eager/FULL F) | | `GlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | `AFDGlmMoeDsaForCausalLM` | | `Qwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | `AFDQwen3MoeForCausalLM` | | `Qwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | `AFDQwen3_5MoeForConditionalGeneration` | @@ -176,8 +176,12 @@ installs its own AFD metadata and opts out of cross-DP synchronization; native sampling, rejection, heads and context-KV updates are inherited. The [0.32 recipe](../../../recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) describes A/F serving. In PD deployments, Prefill uses the upstream complete -draft model and KV connector. Attention and native draft proposal execution stay -eager. F may use the existing FULL transaction graph with one homogeneous W4A8 or W8A8 +draft model and KV connector. Attention may use native graphs with or without +PD and Attention-local FlashComm1/SP, with matching target/draft TP and DBO +disabled. The draft can remain eager or use its own native graph alongside the +target graph. Native resolution selects the effective mode for `FULL` or +`FULL_DECODE_ONLY`; device qualification uses compilation mode 0. +F may use the existing FULL transaction graph with one homogeneous W4A8 or W8A8 table covering target and draft; selection remains on device. DSpark excludes DBO. diff --git a/docs/design/npu/async_ffn_runtime.md b/docs/design/npu/async_ffn_runtime.md index 299dab9c9..33eb5d9bb 100644 --- a/docs/design/npu/async_ffn_runtime.md +++ b/docs/design/npu/async_ffn_runtime.md @@ -1,7 +1,7 @@ # MRV2 Async CAM FFN eager/FULL runtime FULL captures one routed-expert transaction: dispatch receive (DR), -W4A8/W8A8 FFN, and combine send (CS). Attention remains MRV2 eager. +W4A8/W8A8 FFN, and combine send (CS). Attention uses MRV2 eager or native target decode graphs. `AsyncNPUFFNModelRunner` owns Async CAM model loading, eager/FULL transactions and receive-safe resource release independently of the synchronous CAMP2P @@ -12,7 +12,13 @@ and receive-safe resource release independently of the synchronous CAMP2P The target is the pinned vLLM 0.32 / vLLM-Ascend stack on A3/910C, with homogeneous DeepSeek-V4 W4A8 or W8A8 layers, dynamicQuant=1, static expert placement, F TP=1 and Attention-side gating. Both roles select -`VLLM_USE_V2_MODEL_RUNNER=1`. Attention keeps `--enforce-eager`. +`VLLM_USE_V2_MODEL_RUNNER=1`. Attention normally keeps `--enforce-eager`; the +[target Attention graph](../../npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#deepseek-v4-target-attention-graph) +profile accepts `FULL` or `FULL_DECODE_ONLY`, with native graph-mode resolution. +Attention-local FlashComm1/SP and native PD can be combined with the graphs. +DSpark draft may remain eager or use its native graph alongside target graph. +Device qualification uses compilation mode 0; other compilation modes have not +been validated. Other model-based drafts and MTP remain unsupported. Attention may enable native DBO under the existing two-microbatch constraints. Enable native DBO only on Attention; FFN consumes CAM work items without joining Attention's DBO thread schedule. @@ -28,8 +34,8 @@ decode-only graph modes are rejected. DeepSeek-V4 DSpark uses this same F graph for target and draft routed experts. The combined layer table and startup agreement cover target0–42/draft43–45; all layers must share the same W4A8 or W8A8 layout. Native P loads the complete -draft model. Target/draft Attention and the native speculative configuration -remain eager. DSpark excludes DBO; see the +draft model. Target Attention may use the native decode graph; the draft can +remain eager or use its native query Graph. DSpark excludes DBO; see the [DSpark recipe](../../../recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md). W8A8 reuses the loaded INT8 NZ weights, FP32 W13 scales and BF16 W2 scales. diff --git a/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md b/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md index df23e1314..543b137a7 100644 --- a/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md +++ b/docs/npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md @@ -17,11 +17,12 @@ Current runtime constraints and validation evidence are summarized in the ## When to use this connector -Use this connector for eager asynchronous Ascend inference with these requirements: +Use this connector for asynchronous Ascend inference with these requirements: - plugin-owned Ascend operators are built and installed on every node; - Attention performs MoE gating before dispatch to FFN ranks; -- Attention execution is eager and AFD async-DP is enabled; **`async_dp=true` is required**; +- Attention uses eager execution or the DeepSeek-V4 graph profile below; + **`async_dp=true` is required**; - regular prefill and autoregressive decode steps use the same CAM path; - optional ubatching uses exactly two native MRV2 DBO microbatches, including Attention DP1. @@ -33,8 +34,8 @@ Attention DP replicas make that decision independently; FFN is connector-driven. FFN also supports the experimental complete-transaction FULL mode described in [Async FFN runtime](../design/npu/async_ffn_runtime.md#selection-and-startup). -FULL requires MRV2 on both roles and accepts Attention-side DBO under the -existing constraints. Attention graph execution remains unsupported. +FULL requires MRV2 on both roles and accepts eager Attention-side DBO under +the existing constraints. The target Attention graph profile is described below. ## ModelRunnerV2 deployment @@ -45,7 +46,7 @@ Attention uses native Ascend ModelRunnerV2; FFN uses the independent The V2 Async CAM contract is: - DeepSeek-V4 with homogeneous W4A8 or W8A8 routed experts on Ascend 910C/A3; -- `--enforce-eager` on Attention and `--enable-expert-parallel` on both roles, with +- eager Attention or the target graph profile below, and `--enable-expert-parallel` on both roles, with `async_dp=true` and `compute_gate_on_attention=true`; - on FFN, select eager with `--enforce-eager`; for FULL, omit `--enforce-eager` and pass `--compilation-config '{"cudagraph_mode":"FULL"}'`, following the narrow @@ -67,6 +68,41 @@ per-ubatch model forwards, and output merge. Native splitting can divide one request across microbatches. Native draft-model/MTP paths other than DeepSeek-V4 DSpark are unsupported. +### DeepSeek-V4 target Attention graph + +For a DeepSeek-V4 deployment, replace Attention's `--enforce-eager` +with native ACL decode graph selection: + +```bash +--compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY"}' +``` + +The profile requires PP=PCP=DCP=1, static expert placement, +and DBO/ubatching disabled. Attention-local FlashComm1/SP and native PD may remain +enabled. DSpark may stay enabled: `speculative_config.enforce_eager=true` keeps the +draft eager; `false` enables its native query Graph alongside target Graph. Target and draft use matching +TP widths. Other model-based drafts and MTP remain unsupported. PD uses the +existing upstream KV connector; the W4A8 device comparison covers +`MooncakeHybridConnector` with a native full-draft P. Draft Graph requires target Graph because +native target capture owns the draft capture lifecycle. +Attention DPs remain independent. Each Attention TP group uses the existing +CAM token split and output gather; set `attn_ranks_per_dp` to Attention TP +on both roles. + +Native Ascend owns graph inputs, batch buckets, attention-parameter updates and +replay. Ineligible batches, including prefill, retain native eager fallback. +`FULL` is also accepted and resolved by the native attention backend; the +current DSV4 backend resolves it to `FULL_DECODE_ONLY`. The device qualification +uses compilation mode 0, which disables Dynamo while retaining ACL graph +capture. Other compilation modes have not been validated by this qualification. +AFD stages local real/padded token metadata before capture/replay; Async CAM +does not send synchronous DP control messages. Existing synchronous connectors +retain their control-plane behavior. + +Keep the existing FFN execution settings. Transport warmup and F +service readiness use the existing startup protocol. This option introduces no +new graph owner, native kernel or shutdown protocol. + ## CAM async data flow CAM async removes synchronization between data-parallel replicas and avoids the @@ -276,7 +312,8 @@ actual vendor library path. Missing source operators fail startup. ## Current limitations -- Attention is eager. FFN FULL targets homogeneous W4A8 or W8A8 layers on A3/910C; +- Attention is eager except for the DeepSeek-V4 native decode graph profile above. + FFN FULL targets homogeneous W4A8 or W8A8 layers on A3/910C; see [supported configuration](../design/npu/async_ffn_runtime.md#selection-and-startup) and [validation commands](../design/npu/async_ffn_runtime.md#validation). Broader models and topologies require separate validation. diff --git a/recipe/README.md b/recipe/README.md index 3e7cb9fd3..694418b0d 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -28,7 +28,7 @@ Recipes are organized by hardware backend, connector, and model: | GPU | `P2pNcclAFDConnector` | Qwen3.5-122B-A10B-FP8 | Decode | Three-pod graph recipe (`2A0F,2A0F,0A2F`); other layouts remain configurable through `LAYOUT` | [Graph](gpu/P2pNcclAFDConnector/qwen3_5_122b_a10b_fp8/prefill_decode_colocation/multipod_2a_2a_2f_graph.sh) | | Ascend NPU | `CAMP2pAFDConnector` | DeepSeek-V3.2 | Decode | Historical v0.19.1rc1 recipe; not qualified for the 0.32 target | [Synchronous decode](npu/CAMP2pAFDConnector/deepseek_v3_2/README.md) | | Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V3.2 | Prefill / decode | Historical v0.19.1rc1/v0.26 recipes; standard V3.2 is unsupported by the current Async CAM path | [Async CAM](npu/CAMAsyncAFDConnector/deepseek_v3_2/README.md) | -| Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V4-Flash W4A8/W8A8 DSpark | Separate P / decode, or colocated | v0.32; eager A TP1, F layered eager/FULL. W4A8 PD: native full-draft P TP4, A DP2, F EP4. W8A8 A/F: A DP2, F EP8; PD not qualified | [DSpark A/F](npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) | +| Ascend NPU | `CAMAsyncAFDConnector` | DeepSeek-V4-Flash W4A8/W8A8 DSpark | Separate P / decode, or colocated | v0.32; F layered eager/FULL. W4A8 target/draft native decode Graph: A DP1/TP4, F EP4, with Attention-local SP and with/without PD; native full-draft P TP4 for PD. Earlier eager W4A8 PD: A DP2/TP1, F EP4. W8A8 eager A/F: A DP2/TP1, F EP8; A Graph and PD not qualified | [DSpark A/F](npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md) | Open the model-level README before running a recipe. It documents the required hardware and runtime baseline, topology, environment variables, launch order, diff --git a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md index 48d51b85a..fec550967 100644 --- a/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md +++ b/recipe/npu/CAMAsyncAFDConnector/deepseek_v4_flash/README.md @@ -34,11 +34,14 @@ and network interface. A TP1/DP1 and F EP4 use five distinct DIEs. MODEL=/absolute/path/to/dsv4-flash-dspark-w4a8 LOCAL_IP=REPLACE_WITH_COMMUNICATION_IP INTERFACE=eth0 -NUM_A=1 +A_DP=1 +A_TP=1 +NUM_A=$((A_DP * A_TP)) NUM_F=4 A_DEVICES=0 F_DEVICES=4,5,6,7 FFN_EXEC_ARGS=(--enforce-eager) +A_EXEC_ARGS=(--enforce-eager) export VLLM_USE_V2_MODEL_RUNNER=1 VLLM_ASCEND_ENABLE_FLASHCOMM1=0 export BATCH_SIZE_FACTOR=0.00390625 @@ -61,14 +64,36 @@ COMMON_ARGS=( --no-disable-hybrid-kv-cache-manager ) DSPARK_CONFIG="{\"method\":\"dspark\",\"model\":\"$MODEL\",\"num_speculative_tokens\":5,\"enforce_eager\":true}" +A_DSPARK_CONFIG="$DSPARK_CONFIG" ``` -For DP2 set `NUM_A=2` and `A_DEVICES=0,1`. The W4A8 example uses F EP4. +For DP2/TP1 set `A_DP=2`, recompute `NUM_A`, and use `A_DEVICES=0,1`. The W4A8 example uses F EP4. For a W8A8 checkpoint, size F for its larger weights; the validated W8A8 configuration uses `NUM_F=8` and `F_DEVICES=8,9,10,11,12,13,14,15`. -Target and draft Attention use eager execution. F uses layered eager or layered -FULL, with one combined target/draft weight table. Layered eager can be slower -for small decode batches, so measure the chosen deployment. +Target and draft use matching Attention TP widths; F remains TP1 with EP. +Target Attention may use eager execution or the native decode graph. +`A_DSPARK_CONFIG` selects native draft eager/Graph execution on A; F's +`DSPARK_CONFIG` and expert execution settings remain independent. To enable both +Attention graphs, set the following before launching. Keep `enforce_eager=true` +in `A_DSPARK_CONFIG` when only the target graph is wanted. + +```bash +A_EXEC_ARGS=(--compilation-config '{"mode":0,"cudagraph_mode":"FULL_DECODE_ONLY"}') +A_DSPARK_CONFIG="{\"method\":\"dspark\",\"model\":\"$MODEL\",\"num_speculative_tokens\":5,\"enforce_eager\":false}" +``` + +Attention-local FlashComm1/SP and native PD may remain enabled with both graphs. +For PD, keep the native full-draft P and KV-transfer configuration described +below. The W4A8 graph comparison covers A DP1/TP4, F EP4, and native P TP4 when +using `MooncakeHybridConnector`. Compilation mode 0 is the qualified example; +other compilation modes have not been validated here. + +Native capture buckets are normalized separately for target verification and +draft queries; choose buckets covering both query widths and `max-num-seqs`. +For A TP4, set `A_TP=4`, `A_DP=1`, recompute `NUM_A`, and use four distinct +A devices. Both roles use the same A TP width for CAM topology. F uses layered +eager or layered FULL, with one combined target/draft weight table. Layered eager can be slower for +small decode batches, so measure the chosen deployment. For F FULL, replace the F settings above before launching: @@ -78,8 +103,8 @@ FFN_EXEC_ARGS=(--compilation-config '{"cudagraph_mode":"FULL"}') F captures one DR → layered FFN → CS transaction and repeatedly replays it for both target and draft layers. The bounded replay group is independent of request -and target/draft phase boundaries. A and the native draft configuration retain -eager execution. FULL requires matching W4A8 or W8A8 layouts on all layers; unsupported +and target/draft phase boundaries. Attention target/draft graph choices are +independent of F's execution mode. FULL requires matching W4A8 or W8A8 layouts on all layers; unsupported graph modes fail startup. DSpark does not support DBO. Layered W4A8 forwards the model's `swiglu_limit` for both target and draft @@ -106,17 +131,17 @@ vllm serve "$MODEL" "${COMMON_ARGS[@]}" "${FFN_EXEC_ARGS[@]}" \ \"host\":\"$LOCAL_IP\",\"port\":12499, \"num_attention_ranks\":$NUM_A,\"num_ffn_ranks\":$NUM_F, \"compute_gate_on_attention\":true, - \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":1} + \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":$A_TP} } }" >ffn.log 2>&1 & ``` ```bash ASCEND_RT_VISIBLE_DEVICES="$A_DEVICES" VLLM_PLUGINS=ascend,afd \ -vllm serve "$MODEL" "${COMMON_ARGS[@]}" --enforce-eager \ +vllm serve "$MODEL" "${COMMON_ARGS[@]}" "${A_EXEC_ARGS[@]}" \ --port 7100 --data-parallel-rpc-port 17100 \ - --tensor-parallel-size 1 --data-parallel-size "$NUM_A" \ - --all2all-backend allgather_reducescatter --speculative-config "$DSPARK_CONFIG" \ + --tensor-parallel-size "$A_TP" --data-parallel-size "$A_DP" \ + --all2all-backend allgather_reducescatter --speculative-config "$A_DSPARK_CONFIG" \ --additional-config "{ \"enable_cpu_binding\":false,\"enable_force_eplb\":false, \"enable_dsa_cp\":false,\"multistream_dsv4_dsa_overlap\":false, @@ -125,7 +150,7 @@ vllm serve "$MODEL" "${COMMON_ARGS[@]}" --enforce-eager \ \"host\":\"$LOCAL_IP\",\"port\":12499, \"num_attention_ranks\":$NUM_A,\"num_ffn_ranks\":$NUM_F, \"compute_gate_on_attention\":true, - \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":1} + \"connector_extra_config\":{\"dynamicQuant\":1,\"attn_ranks_per_dp\":$A_TP} } }" >attention.log 2>&1 & ``` diff --git a/tests/unit/compat/npu/test_feature_validation.py b/tests/unit/compat/npu/test_feature_validation.py index 93ed73baf..cda824302 100644 --- a/tests/unit/compat/npu/test_feature_validation.py +++ b/tests/unit/compat/npu/test_feature_validation.py @@ -19,7 +19,10 @@ from tests.unit.fakes.vllm_config import ( # noqa: E402 async_moe_config as _async_moe_config, ) -from tests.unit.fakes.vllm_config import npu_dspark_config # noqa: E402 +from tests.unit.fakes.vllm_config import ( # noqa: E402 + npu_async_cam_target_graph_config, + npu_dspark_config, +) from tests.unit.fakes.vllm_config import ( # noqa: E402 npu_vllm_config as _vllm_config, ) @@ -35,7 +38,7 @@ [ ("attention", True, True, "deepseek_v4", None), ("ffn", True, True, "deepseek_v4", None), - ("attention", False, True, "deepseek_v4", "DSpark requires eager"), + ("attention", False, True, "deepseek_v4", "draft Graph requires target"), ( "attention", True, @@ -43,13 +46,19 @@ "deepseek_v2", "requires DeepSeek V4", ), + ("ffn", False, True, "deepseek_v4", None), ("ffn", True, False, "deepseek_v4", None), - ("attention", True, False, "deepseek_v4", "eager Attention"), + ("ffn", False, False, "deepseek_v4", None), + ("attention", True, False, "deepseek_v4", None), + ("attention", False, False, "deepseek_v4", None), ], ) def test_dspark_feature_boundaries(role, draft_eager, target_eager, model_type, error): config = npu_dspark_config(role=role, draft_eager=draft_eager) config.model_config.enforce_eager = target_eager + if role == "attention" and not target_eager: + config.compilation_config.mode = 0 + config.compilation_config.cudagraph_mode.name = "FULL_DECODE_ONLY" config.model_config.hf_text_config.model_type = model_type if error: with pytest.raises(RuntimeError, match=error): @@ -58,6 +67,16 @@ def test_dspark_feature_boundaries(role, draft_eager, target_eager, model_type, fail_if_unsupported_npu_afd_features(config) +@pytest.mark.parametrize("draft_eager", [True, False]) +def test_dspark_target_graph_reuses_tp4(monkeypatch, draft_eager): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_dspark_config(tp_size=4, draft_eager=draft_eager) + config.model_config.enforce_eager = False + config.compilation_config.mode = 0 + config.compilation_config.cudagraph_mode.name = "FULL_DECODE_ONLY" + fail_if_unsupported_npu_afd_features(config) + + @pytest.mark.parametrize( "mode", ["FULL_DECODE_ONLY", "PIECEWISE", "FULL_AND_PIECEWISE"] ) @@ -173,7 +192,7 @@ def _mla_dbo(cudagraph_mode, **model_fields): ), pytest.param( _config_with(model={"enforce_eager": False}, **ASYNC_CAM), - "only eager", + "FULL_DECODE_ONLY", id="async-graph", ), pytest.param( @@ -277,6 +296,41 @@ def test_async_cam_rejects_models_without_layered_ffn(role, model_type): fail_if_unsupported_npu_afd_features(config) +@pytest.mark.parametrize("dp_size", [1, 2]) +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_async_cam_target_graph_feature_admission(monkeypatch, dp_size, tp_size): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "1") + config = npu_async_cam_target_graph_config( + data_parallel_size=dp_size, + tensor_parallel_size=tp_size, + num_attention_ranks=dp_size * tp_size, + ) + config.kv_transfer_config = object() + config.parallel_config.use_sequence_parallel_moe = True + fail_if_unsupported_npu_afd_features(config) + + +@pytest.mark.parametrize( + "method,uses_draft,error", + [ + ("mtp", True, "native draft models or MTP"), + ("ngram", False, None), + ], +) +def test_async_cam_target_graph_feature_preserves_draft_model_boundary( + monkeypatch, method, uses_draft, error +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config() + config.speculative_config = SimpleNamespace( + method=method, + use_eagle=lambda: False, + uses_draft_model=lambda: uses_draft, + ) + with pytest.raises(RuntimeError, match=error) if error else nullcontext(): + fail_if_unsupported_npu_afd_features(config) + + def test_npu_afd_num_ubatches_reads_two_ubatch_dbo(): assert npu_afd_num_ubatches(_vllm_config(**TWO_UBATCH_DBO)) == 2 diff --git a/tests/unit/compat/patches/test_signatures.py b/tests/unit/compat/patches/test_signatures.py index 0b7bd68c4..d9ebad4bb 100644 --- a/tests/unit/compat/patches/test_signatures.py +++ b/tests/unit/compat/patches/test_signatures.py @@ -237,6 +237,22 @@ def _use_installed_patch(binding, symbol, marker_name, install): "fullgraph-replay", id="fullgraph-replay", ), + pytest.param( + "vllm.v1.worker.gpu.spec_decode.dflash.cudagraph", + None, + "_prepare_dflash_inputs_to_capture", + None, + "npu-draft-capture", + id="npu-draft-capture", + ), + pytest.param( + "vllm_ascend.worker.v2.spec_decode.dflash.aclgraph", + "DFlashAclGraphManager", + "run_fullgraph", + None, + "npu-draft-replay", + id="npu-draft-replay", + ), pytest.param( "vllm_ascend.ascend_config", None, @@ -352,6 +368,26 @@ def test_patch_preserves_upstream_signature( tracker = adapter._AFDCaptureEventTracker([]) scope = adapter._use_afd_capture_input_preparation(None, tracker) + elif scope_name in {"npu-draft-capture", "npu-draft-replay"}: + from vllm.config import CUDAGraphMode + + from afd_plugin.v1.worker.npu import dspark_graph as adapter + + if scope_name == "npu-draft-capture": + runner = SimpleNamespace( + speculative_config=SimpleNamespace(enforce_eager=True), + _afd_is_graph_capturing=False, + _afd_is_graph_replaying=False, + ) + manager = SimpleNamespace( + _capture_descs={CUDAGraphMode.FULL: []}, + _max_full_descs_to_capture=None, + ) + scope = adapter.use_afd_draft_capture(runner, manager) + else: + native_class = getattr(target, owner) + binding = native_class.__new__(native_class) + scope = adapter.use_afd_draft_replay(SimpleNamespace(), binding, 1) elif scope_name == "camp2p-dispatch": from afd_plugin.v1.worker.npu import attention_model_runner_v2 as adapter diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py index 6d8a53e0c..f61a7f7df 100644 --- a/tests/unit/config/test_validation.py +++ b/tests/unit/config/test_validation.py @@ -18,7 +18,11 @@ validate_gpu_model_runner_v2_config, validate_npu_model_runner_v2_config, ) -from tests.unit.fakes.vllm_config import mrv2_vllm_config +from tests.unit.fakes.vllm_config import ( + mrv2_vllm_config, + npu_async_cam_target_graph_config, + npu_dspark_config, +) MODEL_UTILS_MODULE = "afd_plugin.model_executor.models.model_utils" @@ -156,6 +160,97 @@ def test_gpu_validator_accepts_target_config_without_removed_enable_sp(mrv2_conf ) +@pytest.mark.parametrize("dp_size", [1, 2]) +@pytest.mark.parametrize("tp_size", [1, 2]) +def test_async_cam_target_graph_accepts_independent_attention_dps( + monkeypatch, mrv2_config, dp_size, tp_size +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config( + data_parallel_size=dp_size, + tensor_parallel_size=tp_size, + num_attention_ranks=dp_size * tp_size, + ) + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + +@pytest.mark.parametrize("draft_eager", [True, False]) +def test_async_cam_target_graph_accepts_dspark_tp4( + monkeypatch, mrv2_config, draft_eager +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config( + tensor_parallel_size=4, num_attention_ranks=4 + ) + config.speculative_config = npu_dspark_config( + tp_size=4, draft_eager=draft_eager + ).speculative_config + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + +@pytest.mark.parametrize( + "path,value,error", + [ + ("use_v2_model_runner", False, "ModelRunnerV2"), + ("model_config.hf_text_config.model_type", "deepseek_v2", "DeepSeek-V4"), + ("parallel_config.enable_dbo", True, "DBO/ubatching"), + ("parallel_config.use_ubatching", True, "DBO/ubatching"), + ("parallel_config.pipeline_parallel_size", 2, "PP or CP"), + ("parallel_config.prefill_context_parallel_size", 2, "PP or CP"), + ("parallel_config.decode_context_parallel_size", 2, "PP or CP"), + ("compilation_config.cudagraph_mode.name", "NONE", "supports ACL graph modes"), + ( + "compilation_config.cudagraph_mode.name", + "PIECEWISE", + "supports ACL graph modes", + ), + ], +) +def test_async_cam_target_graph_rejects_unsupported_configuration( + monkeypatch, mrv2_config, path, value, error +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "0") + config = npu_async_cam_target_graph_config() + *parents, name = path.split(".") + target = config + for parent in parents: + target = getattr(target, parent) + setattr(target, name, value) + with pytest.raises(RuntimeError, match=error): + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + +@pytest.mark.parametrize( + "path,value", + [ + ("kv_transfer_config", object()), + ("parallel_config.use_sequence_parallel_moe", True), + ("speculative_config", SimpleNamespace(method="ngram")), + ("compilation_config.cudagraph_mode.name", "FULL"), + ("compilation_config.mode", 3), + ], +) +def test_async_cam_target_graph_preserves_native_options( + monkeypatch, mrv2_config, path, value +): + monkeypatch.setenv("VLLM_ASCEND_ENABLE_FLASHCOMM1", "1") + config = npu_async_cam_target_graph_config() + *parents, name = path.split(".") + target = config + for parent in parents: + target = getattr(target, parent) + setattr(target, name, value) + validate_npu_model_runner_v2_config( + config, expected_role="attention", device_type="npu" + ) + + @pytest.mark.parametrize( "role,tp,error", [ @@ -655,8 +750,8 @@ def test_gpu_v2_validator_outcome(mrv2_config, monkeypatch, overrides, error): ), pytest.param( {**_NPU_ASYNC, "model.enforce_eager": False}, - _err("eager"), - id="reject-async-cam-graph", + _err("MRV2 DeepSeek-V4"), + id="reject-async-cam-non-dsv4-graph", ), pytest.param( {**_NPU_ASYNC, "par.pipeline_parallel_size": 2}, @@ -697,6 +792,7 @@ def test_gpu_v2_validator_outcome(mrv2_config, monkeypatch, overrides, error): ) def test_npu_v2_validator_outcome(mrv2_config, overrides, error): config, call = _validator_call(overrides) + config.model_config.hf_text_config = config.model_config.hf_config with _expect(error): validate_npu_model_runner_v2_config( config, diff --git a/tests/unit/fakes/npu_attention.py b/tests/unit/fakes/npu_attention.py new file mode 100644 index 000000000..80e8ae9d0 --- /dev/null +++ b/tests/unit/fakes/npu_attention.py @@ -0,0 +1,54 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Native Ascend metadata seams without device allocation.""" + +from types import SimpleNamespace + +from tests.unit.fakes.vllm_config import npu_vllm_config + + +def metadata_runner(): + from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( + AFDNPUAttentionModelRunnerV2, + ) + + runner = object.__new__(AFDNPUAttentionModelRunnerV2) + runner.vllm_config = npu_vllm_config() + runner.vllm_config.kv_transfer_config = None + runner.speculative_config = SimpleNamespace(enforce_eager=False) + runner.connector = SimpleNamespace(control_plane=None) + runner._afd_pending_metadata = None + runner._afd_suppress_metadata_send = False + runner._afd_transaction_counter = 0 + runner._is_warmup = runner._afd_is_profile = False + runner._afd_is_graph_capturing = runner._afd_is_graph_replaying = False + return runner + + +def forward_context(num_tokens): + return SimpleNamespace( + batch_descriptor=SimpleNamespace(num_tokens=num_tokens), + attn_metadata={}, + additional_kwargs={}, + ubatch_slices=None, + ) + + +class DraftGraphManager: + def __init__(self, runner, descs=(), *, fail=False): + from vllm.config import CUDAGraphMode + + self.runner = runner + self._capture_descs = {CUDAGraphMode.FULL: list(descs)} + self._max_full_descs_to_capture = None + self.fail = fail + self.emitted = [] + + def run_fullgraph(self, desc): + context = forward_context(desc.num_tokens) + self.runner.install_afd_metadata_on_forward_context(context) + metadata = context.additional_kwargs["afd_metadata"] + self.emitted.append(metadata) + if self.fail: + raise RuntimeError("native replay failed") + return metadata diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index e67a400da..c374efebd 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -45,6 +45,7 @@ def npu_parallel_config(**overrides): "pipeline_parallel_size": 1, "use_sequence_parallel_moe": False, "enable_eplb": False, + "enable_elastic_ep": False, "prefill_context_parallel_size": 1, "decode_context_parallel_size": 1, "dbo_decode_token_threshold": 1, @@ -101,10 +102,34 @@ def npu_vllm_config( fast_moe_cold_start=False, ), speculative_config=speculative_config, + kv_transfer_config=None, scheduler_config=SimpleNamespace(max_num_batched_tokens=16), ) +def npu_async_cam_target_graph_config(**parallel_overrides): + config = npu_vllm_config( + connector="CAMAsyncAFDConnector", + async_dp=True, + compute_gate_on_attention=True, + cudagraph_mode="FULL_DECODE_ONLY", + extra_config={ + "dynamicQuant": 1, + "attn_ranks_per_dp": parallel_overrides.get("tensor_parallel_size", 1), + }, + enable_expert_parallel=True, + **parallel_overrides, + ) + config.model_config.enforce_eager = False + config.model_config.hf_config = config.model_config.hf_text_config = ( + SimpleNamespace( + model_type="deepseek_v4", architectures=["DeepseekV4ForCausalLM"] + ) + ) + config.compilation_config.mode = 0 + return config + + def async_moe_config( *, role="attention", @@ -253,7 +278,7 @@ def m2n_vllm_config( ) -def npu_dspark_config(*, role="attention", draft_eager=True): +def npu_dspark_config(*, role="attention", draft_eager=True, tp_size=1): hf = SimpleNamespace( architectures=["DSparkDraftModel"], vocab_size=16, @@ -266,12 +291,14 @@ def npu_dspark_config(*, role="attention", draft_eager=True): connector="CAMAsyncAFDConnector", async_dp=True, compute_gate_on_attention=True, - extra_config={"dynamicQuant": 1}, + tensor_parallel_size=tp_size, + num_attention_ranks=tp_size, + extra_config={"dynamicQuant": 1, "attn_ranks_per_dp": tp_size}, speculative_config=SimpleNamespace( method="dspark", enforce_eager=draft_eager, num_speculative_tokens=5, - draft_parallel_config=SimpleNamespace(tensor_parallel_size=1), + draft_parallel_config=SimpleNamespace(tensor_parallel_size=tp_size), draft_model_config=SimpleNamespace(hf_config=hf, hf_text_config=hf), ), ) diff --git a/tests/unit/v1/worker/npu/test_dspark_graph.py b/tests/unit/v1/worker/npu/test_dspark_graph.py new file mode 100644 index 000000000..42884a3b0 --- /dev/null +++ b/tests/unit/v1/worker/npu/test_dspark_graph.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Draft graph scopes preserve independent AFD batches and native preparation.""" + +from contextlib import nullcontext + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") +pytest.importorskip("vllm_ascend") + +from vllm.config import CUDAGraphMode # noqa: E402 +from vllm.v1.worker.gpu.cudagraph_utils import ( # noqa: E402 + AttentionState, + BatchExecutionDescriptor, +) + +from afd_plugin.v1.worker.npu import dspark_graph # noqa: E402 +from tests.unit.fakes.npu_attention import ( # noqa: E402 + DraftGraphManager, + forward_context, + metadata_runner, +) + + +@pytest.mark.parametrize("fail", [False, True]) +def test_draft_replay_isolates_padding_and_restores_target_context(fail): + runner = metadata_runner() + target = forward_context(12) + runner.install_afd_metadata_on_forward_context(target) + target_metadata = target.additional_kwargs["afd_metadata"] + manager = DraftGraphManager(runner, fail=fail) + desc = BatchExecutionDescriptor(CUDAGraphMode.FULL, num_tokens=15, num_reqs=3) + expected = ( + pytest.raises(RuntimeError, match="native replay failed") + if fail + else nullcontext() + ) + with ( + expected, + runner._afd_draft_context(False, False), + dspark_graph.use_afd_draft_replay(runner, manager, real_tokens=10), + ): + manager.run_fullgraph(desc) + draft = manager.emitted[0] + assert draft.tokens_lens == [15] + assert draft.tokens_unpadded_lens == [10] + assert draft.transaction_id != target_metadata.transaction_id + following = forward_context(12) + runner.install_afd_metadata_on_forward_context(following) + assert following.additional_kwargs["afd_metadata"] is target_metadata + manager.fail = False + assert manager.run_fullgraph(desc) is target_metadata + + +@pytest.mark.parametrize("fail", [False, True]) +def test_draft_capture_creates_distinct_batches_and_restores_native_prepare( + monkeypatch, fail +): + runner = metadata_runner() + desc = BatchExecutionDescriptor(CUDAGraphMode.FULL, num_tokens=10, num_reqs=2) + manager = DraftGraphManager(runner, [desc]) + native_result = AttentionState({}, {}) + monkeypatch.setattr( + dspark_graph.dflash_cudagraph, + "_prepare_dflash_inputs_to_capture", + lambda *args, **kwargs: native_result, + ) + emitted = [] + expected = ( + pytest.raises(RuntimeError, match="native capture failed") + if fail + else nullcontext() + ) + with ( + expected, + runner._afd_draft_context(True, False), + dspark_graph.use_afd_draft_capture(runner, manager), + ): + for _ in range(2): + result = dspark_graph.dflash_cudagraph._prepare_dflash_inputs_to_capture( + 2, 10, None, None, None, None, 4096, False, False + ) + assert result is native_result + context = forward_context(10) + runner.install_afd_metadata_on_forward_context(context) + emitted.append(context.additional_kwargs["afd_metadata"]) + if fail: + raise RuntimeError("native capture failed") + assert [item.tokens_lens for item in emitted] == [[10], [10]] + assert emitted[0].transaction_id != emitted[1].transaction_id + assert ( + dspark_graph.dflash_cudagraph._prepare_dflash_inputs_to_capture() + is native_result + ) + + +def test_draft_full_rejects_missing_native_capture_descriptors(): + runner = metadata_runner() + with ( + pytest.raises(RuntimeError, match="no native capture descriptors"), + dspark_graph.use_afd_draft_capture(runner, DraftGraphManager(runner)), + ): + pass diff --git a/tests/unit/v1/worker/test_npu_model_runner_v2.py b/tests/unit/v1/worker/test_npu_model_runner_v2.py index 5c60e7c57..7cb0283c8 100644 --- a/tests/unit/v1/worker/test_npu_model_runner_v2.py +++ b/tests/unit/v1/worker/test_npu_model_runner_v2.py @@ -57,7 +57,9 @@ def _runner(*, async_dp=False): async_dp=async_dp, num_attention_ranks=2, num_ffn_ranks=2 ) runner.connector = SimpleNamespace( - extra_info=SimpleNamespace(), attention_step=nullcontext + extra_info=SimpleNamespace(), + control_plane=None if async_dp else object(), + attention_step=nullcontext, ) runner.cudagraph_manager = None runner.ubatch_runner = None @@ -200,10 +202,11 @@ def update(backend, stream, forward, tokens, config, speculative): @pytest.mark.parametrize("profile_only", [False, True]) @pytest.mark.parametrize("raise_in_native", [False, True]) -def test_capture_forwards_new_inputs_and_restores( - monkeypatch, profile_only, raise_in_native -): # long-test: capture lifecycle - runner = _runner() +@pytest.mark.parametrize("async_dp", [False, True]) +def test_capture_forwards_new_inputs_and_restores( # long-test: capture lifecycle + monkeypatch, profile_only, raise_in_native, async_dp +): + runner = _runner(async_dp=async_dp) descs = [ cudagraph_utils.BatchExecutionDescriptor( cg_mode=CUDAGraphMode.FULL, num_reqs=1, num_tokens=8 @@ -311,11 +314,72 @@ def native_capture(self, *, profile_only=False): ] if not profile_only: expected_events.extend([(32, True, False), (32, False, True)]) - assert events == expected_events + assert events == ([] if async_dp else expected_events) assert cudagraph_utils.prepare_inputs_to_capture is native_prepare assert original_prepare is not native_prepare +@pytest.mark.parametrize("async_dp", [False, True]) +@pytest.mark.parametrize("fail_replay", [False, True]) +def test_replay_real_and_padded_metadata( # long-test: native replay/context seam + monkeypatch, async_dp, fail_replay +): + runner = _runner(async_dp=async_dp) + original_metadata = runner._afd_pending_metadata + runner._afd_transaction_counter = 0 + runner.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_DECODE_ONLY + runner.build_capture_dp_metadata = lambda tokens: tokens + sent = [] + runner.send_dp_metadata = lambda tokens, _slices: sent.append(tokens) + emitted = [] + desc = cudagraph_utils.BatchExecutionDescriptor( + CUDAGraphMode.FULL, num_tokens=8, num_reqs=8 + ) + + def context(): + return SimpleNamespace( + batch_descriptor=None, + attn_metadata={}, + additional_kwargs={}, + ubatch_slices=None, + ) + + class Manager: + def run_fullgraph(self, descriptor): + update_context = context() + runner.install_afd_metadata_on_forward_context(update_context) + emitted.append(update_context.additional_kwargs["afd_metadata"]) + if fail_replay: + raise RuntimeError("native replay failed") + return "native-output" + + runner.cudagraph_manager = Manager() + + def native_execute(self, *args, **kwargs): + assert should_skip_dp_coordination() is async_dp + return self.cudagraph_manager.run_fullgraph(desc) + + monkeypatch.setattr(NPUModelRunner, "execute_model", native_execute) + monkeypatch.setattr(npu_v2, "step_afd_npu_profiler", lambda _prof: None) + expected = ( + pytest.raises(RuntimeError, match="native replay failed") + if fail_replay + else nullcontext() + ) + with expected: + assert ( + runner.execute_model(SimpleNamespace(total_num_scheduled_tokens=5)) + == "native-output" + ) + assert emitted[0].tokens_lens == [8] + assert emitted[0].tokens_unpadded_lens == [5] + assert sent == ([] if async_dp else [8]) + following = context() + runner.install_afd_metadata_on_forward_context(following) + assert following.additional_kwargs["afd_metadata"] is original_metadata + assert not should_skip_dp_coordination() + + @pytest.mark.parametrize("num_tokens", [0, 1, 7]) @pytest.mark.parametrize("dp_rank", [0, 1]) def test_async_execute_uses_native_local_dp_dispatch(monkeypatch, num_tokens, dp_rank):