From 559812d4286cafc41e30c68af5d367ad44b0c114 Mon Sep 17 00:00:00 2001 From: yujuancao07 Date: Tue, 4 Aug 2026 16:29:19 +0800 Subject: [PATCH] feat(npu): support MLA DBO full-graph replay Signed-off-by: yujuancao07 --- afd_plugin/compat/npu/feature_validation.py | 26 +- afd_plugin/compat/npu/runtime.py | 14 +- afd_plugin/compat/patches/npu/mla_graph.py | 58 ++ .../v1/worker/npu/attention_model_runner.py | 15 +- afd_plugin/v1/worker/npu/forward_context.py | 16 +- afd_plugin/v1/worker/npu/mla_graph.py | 127 +++ .../v1/worker/npu/npu_ubatch_wrapper.py | 182 +++- .../models/deepseek_v2_lite/test_e2e_npu.py | 26 +- tests/e2e/runner.py | 1 + tests/e2e/test_runner.py | 34 +- .../compat/patches/test_config_validation.py | 14 +- tests/unit/compat/test_runtime.py | 117 ++- tests/unit/v1/worker/test_npu_mla_graph.py | 932 ++++++++++++++++++ tests/unit/v1/worker/test_npu_runtime.py | 247 +++-- 14 files changed, 1685 insertions(+), 124 deletions(-) create mode 100644 afd_plugin/compat/patches/npu/mla_graph.py create mode 100644 afd_plugin/v1/worker/npu/mla_graph.py create mode 100644 tests/unit/v1/worker/test_npu_mla_graph.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 32cd5f9b..ea03c385 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/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( diff --git a/afd_plugin/compat/npu/runtime.py b/afd_plugin/compat/npu/runtime.py index e360096e..460dcf7e 100644 --- a/afd_plugin/compat/npu/runtime.py +++ b/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__ = [ diff --git a/afd_plugin/compat/patches/npu/mla_graph.py b/afd_plugin/compat/patches/npu/mla_graph.py new file mode 100644 index 00000000..b195c06f --- /dev/null +++ b/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", +] diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index dcc71974..022bc029 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/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: diff --git a/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index 522acaea..98518ea3 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/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) diff --git a/afd_plugin/v1/worker/npu/mla_graph.py b/afd_plugin/v1/worker/npu/mla_graph.py new file mode 100644 index 00000000..6c337a09 --- /dev/null +++ b/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", +] diff --git a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py index d0632511..5f06471b 100644 --- a/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py +++ b/afd_plugin/v1/worker/npu/npu_ubatch_wrapper.py @@ -20,17 +20,27 @@ ) from vllm.forward_context import ( DPMetadata, + ForwardContext, get_forward_context, override_forward_context, ) from vllm.sequence import IntermediateTensors from vllm.v1.worker.gpu_ubatch_wrapper import UbatchMetadata, UBatchWrapper -from vllm_ascend.compilation.acl_graph import ACLGraphWrapper +from vllm_ascend.compilation.acl_graph import ( + ACLGraphWrapper, + GraphParams, + get_graph_params, +) from vllm_ascend.utils import enable_sp from afd_plugin.v1.worker.npu.forward_context import ( create_ascend_forward_context, ) +from afd_plugin.v1.worker.npu.mla_graph import ( + merge_mla_graph_params, + new_mla_graph_params, + override_mla_graph_params, +) from afd_plugin.v1.worker.npu.ubatching import ( AscendUBatchContext, make_ubatch_contexts, @@ -48,6 +58,20 @@ class AscendNPUGraphMetaData: aclgraph: torch.npu.NPUGraph ubatch_metadata: list[AscendUbatchMetadata] outputs: torch.Tensor | IntermediateTensors | None = None + mla_graph_params: tuple[GraphParams, GraphParams] | None = None + + +@dataclass(frozen=True) +class AscendNPUGraphKey: + stage_num_tokens: tuple[int, int] + has_lora: bool + num_active_loras: int + + +FullGraphParamsUpdater = Callable[ + [ForwardContext, int, torch.Tensor | None], + None, +] class AscendUBatchWrapper(UBatchWrapper): @@ -59,13 +83,18 @@ def __init__( vllm_config: VllmConfig, runtime_mode: CUDAGraphMode, device: torch.device, + *, + mla_full_graph_enabled: bool = False, + full_graph_params_updater: FullGraphParamsUpdater | None = None, + enable_enpu: bool = False, ): + assert not enable_enpu, "AscendUBatchWrapper does not support ENPU" self.runnable = runnable self.vllm_config = vllm_config self.compilation_config = vllm_config.compilation_config self.comm_stream = torch.npu.Stream(device=device) self.ready_barrier = threading.Barrier(3) - self.cudagraphs: dict[int, AscendNPUGraphMetaData] = {} + self.cudagraphs: dict[AscendNPUGraphKey, AscendNPUGraphMetaData] = {} self.cudagraph_wrapper = None if runtime_mode is not CUDAGraphMode.NONE: self.cudagraph_wrapper = ACLGraphWrapper( @@ -74,6 +103,8 @@ def __init__( runtime_mode=runtime_mode, ) self.device = device + self.mla_full_graph_enabled = mla_full_graph_enabled + self.full_graph_params_updater = full_graph_params_updater @property def graph_pool(self): @@ -96,6 +127,23 @@ def __getattr__(self, key: str): def unwrap(self) -> Callable: return self.runnable + def owns_full_graph_update( + self, + forward_context: ForwardContext, + ) -> bool: + uses_mla_ubatch_full_graph = ( + self.mla_full_graph_enabled + and forward_context.ubatch_slices is not None + and forward_context.cudagraph_runtime_mode is CUDAGraphMode.FULL + ) + if not uses_mla_ubatch_full_graph: + return False + if forward_context.max_tokens_across_pcp not in (None, 0): + raise RuntimeError( + "MLA DBO FULL graph does not support PCP execution", + ) + return True + def __call__(self, *args, **kwargs): forward_context = get_forward_context() batch_descriptor = forward_context.batch_descriptor @@ -103,17 +151,27 @@ def __call__(self, *args, **kwargs): cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode if ubatch_slices is None: - if cudagraph_runtime_mode is CUDAGraphMode.FULL: - assert batch_descriptor is not None - if batch_descriptor.num_tokens in self.cudagraphs: - cudagraph_runtime_mode = CUDAGraphMode.NONE if cudagraph_runtime_mode in (CUDAGraphMode.NONE, CUDAGraphMode.PIECEWISE): return self.runnable(*args, **kwargs) assert self.cudagraph_wrapper is not None return self.cudagraph_wrapper(*args, **kwargs) + mla_full_graph_active = self.owns_full_graph_update(forward_context) attn_metadata = forward_context.attn_metadata - num_tokens = sum(ubatch_slice.num_tokens for ubatch_slice in ubatch_slices) + if len(ubatch_slices) != 2: + raise RuntimeError( + "Ascend FULL graph requires exactly two ubatches; " + f"got {len(ubatch_slices)}", + ) + stage_num_tokens = ( + ubatch_slices[0].num_tokens, + ubatch_slices[1].num_tokens, + ) + graph_key = AscendNPUGraphKey( + stage_num_tokens, + batch_descriptor.has_lora, + batch_descriptor.num_active_loras, + ) input_ids = kwargs["input_ids"] positions = kwargs["positions"] intermediate_tensors = kwargs["intermediate_tensors"] @@ -140,9 +198,14 @@ def __call__(self, *args, **kwargs): ubatch_dp_metadata.append(None) if ( - num_tokens not in self.cudagraphs + graph_key not in self.cudagraphs and cudagraph_runtime_mode is CUDAGraphMode.FULL ): + mla_graph_params = ( + self._new_mla_capture_params(stage_num_tokens) + if mla_full_graph_active + else None + ) ubatch_metadata = self._make_ubatch_metadata( ubatch_slices, attn_metadata, @@ -154,15 +217,30 @@ def __call__(self, *args, **kwargs): ubatch_dp_metadata, batch_descriptor, CUDAGraphMode.NONE, + mla_graph_params=mla_graph_params, + ) + return self._capture_ubatches( + ubatch_metadata, + self.runnable, + graph_key=graph_key, + mla_graph_params=mla_graph_params, ) - return self._capture_ubatches(ubatch_metadata, self.runnable) if ( - num_tokens in self.cudagraphs + graph_key in self.cudagraphs and cudagraph_runtime_mode is CUDAGraphMode.FULL ): - cudagraph_metadata = self.cudagraphs[num_tokens] - cudagraph_metadata.aclgraph.replay() - get_forward_context().dbo_enabled = True + cudagraph_metadata = self.cudagraphs[graph_key] + if mla_full_graph_active: + self._replay_mla_graph( + cudagraph_metadata, + forward_context, + stage_num_tokens[0], + positions, + ) + else: + torch.npu.current_stream().synchronize() + cudagraph_metadata.aclgraph.replay() + forward_context.dbo_enabled = True return cudagraph_metadata.outputs ubatch_metadata = self._make_ubatch_metadata( @@ -179,6 +257,71 @@ def __call__(self, *args, **kwargs): ) return self._run_ubatches(ubatch_metadata, self.runnable) + def _new_mla_capture_params( + self, + stage_num_tokens: tuple[int, int], + ) -> tuple[GraphParams, GraphParams]: + if stage_num_tokens[0] != stage_num_tokens[1]: + raise RuntimeError( + "MLA DBO FULL graph requires equal padded token counts; " + f"got {stage_num_tokens}", + ) + + aggregate_num_tokens = sum(stage_num_tokens) + graph_params = get_graph_params() + if ( + graph_params is None + or graph_params.workspaces.get(aggregate_num_tokens) is None + ): + raise RuntimeError( + "MLA DBO FULL graph requires the single-batch FIA workspace " + f"for {aggregate_num_tokens} tokens", + ) + workspace = graph_params.workspaces[aggregate_num_tokens] + child_num_tokens = stage_num_tokens[0] + return ( + new_mla_graph_params(child_num_tokens, workspace), + new_mla_graph_params(child_num_tokens, workspace), + ) + + def _replay_mla_graph( + self, + graph_metadata: AscendNPUGraphMetaData, + forward_context: ForwardContext, + num_tokens: int, + positions: torch.Tensor | None, + ) -> None: + if graph_metadata.mla_graph_params is None: + raise RuntimeError( + "MLA DBO FULL graph cache entry has no capture registry", + ) + if self.full_graph_params_updater is None: + raise RuntimeError( + "MLA DBO FULL graph cache entry has no parameter updater", + ) + + merged_metadata, merged_params = merge_mla_graph_params( + forward_context.attn_metadata, + graph_metadata.mla_graph_params, + num_tokens, + ) + + def update_params() -> None: + with override_mla_graph_params( + forward_context, + merged_metadata, + merged_params, + ): + self.full_graph_params_updater( + forward_context, + num_tokens, + positions, + ) + + torch.npu.current_stream().synchronize() + graph_metadata.aclgraph.replay() + update_params() + def _make_ubatch_metadata( self, ubatch_slices, @@ -191,6 +334,8 @@ def _make_ubatch_metadata( dp_metadata, batch_descriptor, cudagraph_runtime_mode, + *, + mla_graph_params: tuple[GraphParams, GraphParams] | None = None, ) -> list[AscendUbatchMetadata]: cur_forward_context = get_forward_context() forward_contexts = [] @@ -208,6 +353,9 @@ def _make_ubatch_metadata( cudagraph_runtime_mode=cudagraph_runtime_mode, ubatch_num=i, skip_compiled=cur_forward_context.skip_compiled, + mla_graph_params=( + mla_graph_params[i] if mla_graph_params is not None else None + ), ) ) @@ -354,10 +502,12 @@ def _capture_ubatches( self, ubatch_metadata: list[AscendUbatchMetadata], model, + *, + graph_key: AscendNPUGraphKey, + mla_graph_params: tuple[GraphParams, GraphParams] | None, ) -> torch.Tensor | IntermediateTensors: results: list[tuple[int, torch.Tensor | IntermediateTensors]] = [] compute_stream = ubatch_metadata[0].context.compute_stream - num_tokens = sum(metadata.num_tokens for metadata in ubatch_metadata) with override_forward_context(None): ubatch_threads = [] @@ -373,6 +523,7 @@ def _capture_ubatches( cudagraph_metadata = AscendNPUGraphMetaData( aclgraph=torch.npu.NPUGraph(), ubatch_metadata=ubatch_metadata, + mla_graph_params=mla_graph_params, ) with torch.npu.graph( cudagraph_metadata.aclgraph, @@ -387,12 +538,13 @@ def _capture_ubatches( sorted_results, ubatch_metadata, ) - self.cudagraphs[num_tokens] = cudagraph_metadata + self.cudagraphs[graph_key] = cudagraph_metadata get_forward_context().dbo_enabled = True return cudagraph_metadata.outputs __all__ = [ + "AscendNPUGraphKey", "AscendNPUGraphMetaData", "AscendUBatchWrapper", "AscendUbatchMetadata", diff --git a/tests/e2e/models/deepseek_v2_lite/test_e2e_npu.py b/tests/e2e/models/deepseek_v2_lite/test_e2e_npu.py index 2052c90e..ac70a97f 100644 --- a/tests/e2e/models/deepseek_v2_lite/test_e2e_npu.py +++ b/tests/e2e/models/deepseek_v2_lite/test_e2e_npu.py @@ -8,7 +8,7 @@ base / +TP / +DBO / +profile / +TP+DBO+profile × {eager, graph} -= 10 tests. DBO variants self-skip on NPU (DBO is not supported there yet). += 10 tests. Profiler is enabled purely through AFD_NPU_{ATTENTION,FFN}_PROFILER_* env vars, which leak through runner.py's os.environ.copy() into the vllm worker — no runner/source change required. @@ -27,6 +27,7 @@ REPO_ROOT = Path(__file__).resolve().parents[4] RUNNER = REPO_ROOT / "tests" / "e2e" / "runner.py" +EAGER_DBO_REQUEST_COUNT = 4 def _npu_list() -> list[str]: @@ -44,11 +45,6 @@ def _model_path() -> str: return model -def _skip_dbo_on_npu() -> None: - """DBO is not supported on NPU yet — skip cleanly instead of failing.""" - pytest.skip("DBO is not supported on NPU yet") - - def _graph_capture_size() -> int: return int(os.environ.get("AFD_NPU_E2E_GRAPH_CAPTURE_SIZE", "8")) @@ -126,6 +122,16 @@ def _run_e2e( ), ], ) + if not graph: + request_count = str(EAGER_DBO_REQUEST_COUNT) + command.extend( + [ + "--num-requests", + request_count, + "--request-concurrency", + request_count, + ], + ) subprocess.run(command, cwd=REPO_ROOT, check=True) @@ -217,13 +223,12 @@ def test_deepseek_v2_2a2f_tp_graph(): # --------------------------------------------------------------------------- -# 2A2F + DBO: eager + graph (NPU self-skips — DBO unsupported) +# 2A2F + DBO: eager + FULL graph # --------------------------------------------------------------------------- @pytest.mark.npu def test_deepseek_v2_2a2f_dbo_eager(): - _skip_dbo_on_npu() _run_e2e( npus=_npu_list(), api_port_base=int( @@ -237,7 +242,6 @@ def test_deepseek_v2_2a2f_dbo_eager(): @pytest.mark.npu @pytest.mark.slow def test_deepseek_v2_2a2f_dbo_graph(): - _skip_dbo_on_npu() _run_e2e( npus=_npu_list(), api_port_base=int( @@ -289,7 +293,7 @@ def test_deepseek_v2_2a2f_profile_graph( # --------------------------------------------------------------------------- -# 2A2F + TP=2 + DBO + profiler: eager + graph (NPU self-skips — DBO unsupported) +# 2A2F + TP=2 + DBO + profiler: eager + graph # --------------------------------------------------------------------------- @@ -298,7 +302,6 @@ def test_deepseek_v2_2a2f_tp_dbo_profile_eager( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ): - _skip_dbo_on_npu() attn_dir, ffn_dir = _enable_profiler(tmp_path, monkeypatch) _run_e2e( npus=_npu_list(), @@ -320,7 +323,6 @@ def test_deepseek_v2_2a2f_tp_dbo_profile_graph( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ): - _skip_dbo_on_npu() attn_dir, ffn_dir = _enable_profiler(tmp_path, monkeypatch) _run_e2e( npus=_npu_list(), diff --git a/tests/e2e/runner.py b/tests/e2e/runner.py index ee86a51e..9547fda0 100644 --- a/tests/e2e/runner.py +++ b/tests/e2e/runner.py @@ -471,6 +471,7 @@ def build_env( role: str | None = None, ) -> dict[str, str]: env = os.environ.copy() + env.setdefault("VLLM_ENGINE_READY_TIMEOUT_S", "18000") if args.device_backend == "npu": env["ASCEND_RT_VISIBLE_DEVICES"] = visible_devices else: diff --git a/tests/e2e/test_runner.py b/tests/e2e/test_runner.py index f6f66ee2..9faa2951 100644 --- a/tests/e2e/test_runner.py +++ b/tests/e2e/test_runner.py @@ -8,6 +8,7 @@ import pytest from tests.e2e import runner +from tests.e2e.models.deepseek_v2_lite import test_e2e_npu as npu_e2e from tests.e2e.runner import build_vllm_command @@ -97,6 +98,16 @@ def test_runner_uses_plugin_decode_bench_connector(): ) +@pytest.mark.parametrize("role", ["attention", "ffn"]) +def test_runner_uses_auto_worker_selection_for_npu(role): + args = _args() + args.device_backend = "npu" + + command = build_vllm_command(args, role=role) + + assert "--worker-cls" not in command + + def test_runner_builds_npu_async_cam_role_specific_topology(): args = _args() args.device_backend = "npu" @@ -122,8 +133,6 @@ def test_runner_builds_npu_async_cam_role_specific_topology(): assert _arg_value(attention_command, "--tensor-parallel-size") == "2" assert _arg_value(ffn_command, "--data-parallel-size") == "2" assert _arg_value(ffn_command, "--tensor-parallel-size") == "1" - assert "--worker-cls" not in attention_command - assert "--worker-cls" not in ffn_command additional_config = json.loads(_arg_value(attention_command, "--additional-config")) afd_config = additional_config["afd"] @@ -182,6 +191,27 @@ def fake_request_completion(_args): assert len(responses) == 2 +def test_npu_eager_dbo_sends_enough_requests_for_two_ubatches(monkeypatch): + captured = {} + + def fake_run(command, **_kwargs): + captured["command"] = command + + monkeypatch.setattr(npu_e2e, "_model_path", lambda: "/models/DeepSeek-V2-Lite") + monkeypatch.setattr(npu_e2e.subprocess, "run", fake_run) + + npu_e2e._run_e2e( + npus=["0", "1", "2", "3"], + api_port_base=18400, + afd_port=6279, + dbo=True, + ) + + command = captured["command"] + assert _arg_value(command, "--num-requests") == "4" + assert _arg_value(command, "--request-concurrency") == "4" + + def test_request_completion_includes_http_error_body(monkeypatch): args = _args() error = urllib.error.HTTPError( diff --git a/tests/unit/compat/patches/test_config_validation.py b/tests/unit/compat/patches/test_config_validation.py index e7928830..f9b66d55 100644 --- a/tests/unit/compat/patches/test_config_validation.py +++ b/tests/unit/compat/patches/test_config_validation.py @@ -7,7 +7,9 @@ import pytest +from afd_plugin.compat import npu as npu_compat from afd_plugin.compat.npu import runtime as ascend_runtime +from afd_plugin.compat.patches.npu import mla_graph from afd_plugin.validation import ( ATTENTION_WORKER_FQCN, FFN_WORKER_FQCN, @@ -205,6 +207,7 @@ def create_engine_config(engine_args, usage_context=None, headless=False): monkeypatch.setitem(sys.modules, "vllm_ascend", fake_package) monkeypatch.setitem(sys.modules, "vllm_ascend.platform", fake_platform) sys.modules["vllm.platforms"].current_platform = NPUPlatform + monkeypatch.setattr(mla_graph, "apply_afd_mla_graph_patch", lambda: True) monkeypatch.setattr(ascend_runtime, "_PATCHES_APPLIED", False) return arg_utils_module, NPUPlatform, events @@ -434,6 +437,11 @@ def test_config_validation_patch_auto_selects_afd_worker( expected_worker_cls, ): arg_utils_module, config_module = _install_fake_vllm_config(monkeypatch) + monkeypatch.setattr( + npu_compat, + "apply_afd_ascend_patches_if_needed", + lambda: None, + ) config_module.VllmConfig.platform_worker_cls = platform_worker_cls _set_fake_platform(is_cuda=is_cuda, device_type=device_type) _load_patch_module() @@ -473,7 +481,6 @@ def test_config_validation_patch_auto_selects_without_ubatching(monkeypatch): def test_config_validation_installs_ascend_patch_only_on_npu(monkeypatch): arg_utils_module, config_module = _install_fake_vllm_config(monkeypatch) - import afd_plugin.compat.npu as npu_compat calls = [] monkeypatch.setattr( @@ -523,6 +530,11 @@ def test_config_validation_patch_rejects_unsupported_auto_platform( device_type, ): arg_utils_module, config_module = _install_fake_vllm_config(monkeypatch) + monkeypatch.setattr( + npu_compat, + "apply_afd_ascend_patches_if_needed", + lambda: None, + ) config_module.VllmConfig.platform_worker_cls = platform_worker_cls _set_fake_platform(is_cuda=is_cuda, device_type=device_type) _load_patch_module() diff --git a/tests/unit/compat/test_runtime.py b/tests/unit/compat/test_runtime.py index c2a334e3..7123426e 100644 --- a/tests/unit/compat/test_runtime.py +++ b/tests/unit/compat/test_runtime.py @@ -130,6 +130,8 @@ def set_ascend_forward_context( def test_npu_afd_config_patch_restores_dbo_for_afd(monkeypatch): + from afd_plugin.compat.patches.npu import mla_graph + fake_package = ModuleType("vllm_ascend") fake_package.__path__ = [] fake_platform = ModuleType("vllm_ascend.platform") @@ -176,6 +178,7 @@ def afd_vllm_config(*, active=True): return config fake_platform.NPUPlatform = NPUPlatform + monkeypatch.setattr(mla_graph, "apply_afd_mla_graph_patch", lambda: True) monkeypatch.setitem(sys.modules, "vllm_ascend", fake_package) monkeypatch.setitem(sys.modules, "vllm_ascend.platform", fake_platform) monkeypatch.setattr(ascend_runtime, "_PATCHES_APPLIED", False) @@ -204,14 +207,18 @@ def afd_vllm_config(*, active=True): assert inactive_config.parallel_config.all2all_backend == "flashinfer_all2allv" -def test_npu_afd_config_patch_retries_after_initial_import_error(monkeypatch): +def test_npu_afd_config_patch_raises_and_retries_after_import_error(monkeypatch): + from afd_plugin.compat.patches.npu import mla_graph + fake_package = ModuleType("vllm_ascend") fake_package.__path__ = [] monkeypatch.setitem(sys.modules, "vllm_ascend", fake_package) monkeypatch.setitem(sys.modules, "vllm_ascend.platform", None) + monkeypatch.setattr(mla_graph, "apply_afd_mla_graph_patch", lambda: True) monkeypatch.setattr(ascend_runtime, "_PATCHES_APPLIED", False) - ascend_runtime.apply_afd_ascend_patches_if_needed() + with pytest.raises(RuntimeError, match="DBO config patch"): + ascend_runtime.apply_afd_ascend_patches_if_needed() assert ascend_runtime._PATCHES_APPLIED is False @@ -229,3 +236,109 @@ def check_and_update_config(cls, vllm_config): assert ascend_runtime._PATCHES_APPLIED is True assert hasattr(NPUPlatform, "_afd_plugin_ascend_platform_patch_state") + + +def test_npu_patches_reject_missing_mla_resolver(monkeypatch): + fake_vllm = ModuleType("vllm") + fake_vllm.__path__ = [] + fake_forward_context = ModuleType("vllm.forward_context") + fake_forward_context.get_forward_context = lambda: None + fake_forward_context.is_forward_context_available = lambda: False + fake_ascend = ModuleType("vllm_ascend") + fake_ascend.__path__ = [] + fake_platform = ModuleType("vllm_ascend.platform") + fake_attention = ModuleType("vllm_ascend.attention") + fake_attention.__path__ = [] + + class NPUPlatform: + @classmethod + def check_and_update_config(cls, vllm_config): + del cls, vllm_config + + fake_platform.NPUPlatform = NPUPlatform + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + monkeypatch.setitem(sys.modules, "vllm.forward_context", fake_forward_context) + monkeypatch.setitem(sys.modules, "vllm_ascend", fake_ascend) + monkeypatch.setitem(sys.modules, "vllm_ascend.platform", fake_platform) + monkeypatch.setitem(sys.modules, "vllm_ascend.attention", fake_attention) + monkeypatch.delitem( + sys.modules, + "vllm_ascend.attention.mla_v1", + raising=False, + ) + monkeypatch.setattr(ascend_runtime, "_PATCHES_APPLIED", False) + + with pytest.raises(RuntimeError, match="MLA graph patch"): + ascend_runtime.apply_afd_ascend_patches_if_needed() + assert ascend_runtime._PATCHES_APPLIED is False + + +def test_npu_patches_route_mla_graph_params_from_forward_context(monkeypatch): + fake_vllm = ModuleType("vllm") + fake_vllm.__path__ = [] + fake_forward_context = ModuleType("vllm.forward_context") + fake_ascend = ModuleType("vllm_ascend") + fake_ascend.__path__ = [] + fake_platform = ModuleType("vllm_ascend.platform") + fake_attention = ModuleType("vllm_ascend.attention") + fake_attention.__path__ = [] + fake_mla = ModuleType("vllm_ascend.attention.mla_v1") + + class NPUPlatform: + @classmethod + def check_and_update_config(cls, vllm_config): + del cls, vllm_config + + upstream_registry = object() + afd_registry = object() + forward_context = SimpleNamespace( + additional_kwargs={"afd_mla_graph_params": afd_registry}, + ) + context_available = True + + def get_forward_context(): + return forward_context + + def is_forward_context_available(): + return context_available + + def get_graph_params(): + return upstream_registry + + fake_forward_context.get_forward_context = get_forward_context + fake_forward_context.is_forward_context_available = is_forward_context_available + fake_platform.NPUPlatform = NPUPlatform + monkeypatch.setitem(sys.modules, "vllm", fake_vllm) + monkeypatch.setitem( + sys.modules, + "vllm.forward_context", + fake_forward_context, + ) + monkeypatch.setitem(sys.modules, "vllm_ascend", fake_ascend) + monkeypatch.setitem(sys.modules, "vllm_ascend.platform", fake_platform) + monkeypatch.setitem(sys.modules, "vllm_ascend.attention", fake_attention) + monkeypatch.setitem( + sys.modules, + "vllm_ascend.attention.mla_v1", + fake_mla, + ) + monkeypatch.setattr(ascend_runtime, "_PATCHES_APPLIED", False) + + with pytest.raises(AttributeError, match="get_graph_params"): + ascend_runtime.apply_afd_ascend_patches_if_needed() + assert ascend_runtime._PATCHES_APPLIED is False + + fake_mla.get_graph_params = get_graph_params + ascend_runtime.apply_afd_ascend_patches_if_needed() + patched_get_graph_params = fake_mla.get_graph_params + + assert fake_mla.get_graph_params() is afd_registry + + forward_context.additional_kwargs = {} + assert fake_mla.get_graph_params() is upstream_registry + + context_available = False + assert fake_mla.get_graph_params() is upstream_registry + + ascend_runtime.apply_afd_ascend_patches_if_needed() + assert fake_mla.get_graph_params is patched_get_graph_params diff --git a/tests/unit/v1/worker/test_npu_mla_graph.py b/tests/unit/v1/worker/test_npu_mla_graph.py new file mode 100644 index 00000000..a0de74aa --- /dev/null +++ b/tests/unit/v1/worker/test_npu_mla_graph.py @@ -0,0 +1,932 @@ +from __future__ import annotations + +import importlib +import sys +from contextlib import contextmanager +from types import ModuleType, SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") + + +def _reload_module( + monkeypatch: pytest.MonkeyPatch, + module_name: str, +) -> ModuleType: + """Reload a module and restore both import caches after the test.""" + package_name, module_attribute = module_name.rsplit(".", 1) + package = importlib.import_module(package_name) + monkeypatch.setattr(package, module_attribute, None, raising=False) + monkeypatch.delitem(sys.modules, module_name, raising=False) + return importlib.import_module(module_name) + + +def _load_mla_graph_module(monkeypatch): + fake_ascend = ModuleType("vllm_ascend") + fake_ascend.__path__ = [] + fake_compilation = ModuleType("vllm_ascend.compilation") + fake_compilation.__path__ = [] + fake_acl_graph = ModuleType("vllm_ascend.compilation.acl_graph") + + class GraphParams: + def __init__(self, events, workspaces, handles, attn_params): + self.events = events + self.workspaces = workspaces + self.handles = handles + self.attn_params = attn_params + + fake_acl_graph.GraphParams = GraphParams + monkeypatch.setitem(sys.modules, "vllm_ascend", fake_ascend) + monkeypatch.setitem( + sys.modules, + "vllm_ascend.compilation", + fake_compilation, + ) + monkeypatch.setitem( + sys.modules, + "vllm_ascend.compilation.acl_graph", + fake_acl_graph, + ) + + module_name = "afd_plugin.v1.worker.npu.mla_graph" + return _reload_module(monkeypatch, module_name) + + +def _graph_params( + num_tokens, + workspace, + events, + handles, + attn_params, +): + return SimpleNamespace( + events={num_tokens: list(events)}, + workspaces={num_tokens: workspace}, + handles={num_tokens: list(handles)}, + attn_params={num_tokens: list(attn_params)}, + ) + + +def _load_forward_context_module(monkeypatch): + fake_vllm = ModuleType("vllm") + fake_vllm.__path__ = [] + fake_config = ModuleType("vllm.config") + + class CUDAGraphMode: + NONE = object() + FULL = object() + + fake_config.CUDAGraphMode = CUDAGraphMode + fake_config.VllmConfig = object + fake_distributed = ModuleType("vllm.distributed") + fake_distributed.get_dp_group = lambda: SimpleNamespace(world_size=1) + fake_distributed.get_tensor_model_parallel_world_size = lambda: 1 + fake_forward_context = ModuleType("vllm.forward_context") + + class ForwardContext(SimpleNamespace): + def __init__(self, **kwargs): + super().__init__(**kwargs) + + fake_forward_context.BatchDescriptor = object + fake_forward_context.DPMetadata = object + fake_forward_context.ForwardContext = ForwardContext + fake_vllm_v1 = ModuleType("vllm.v1") + fake_vllm_v1.__path__ = [] + fake_vllm_worker = ModuleType("vllm.v1.worker") + fake_vllm_worker.__path__ = [] + fake_ubatch_utils = ModuleType("vllm.v1.worker.ubatch_utils") + fake_ubatch_utils.UBatchSlices = list + + fake_ascend = ModuleType("vllm_ascend") + fake_ascend.__path__ = [] + fake_ops = ModuleType("vllm_ascend.ops") + fake_ops.__path__ = [] + fake_fused_moe = ModuleType("vllm_ascend.ops.fused_moe") + fake_fused_moe.__path__ = [] + fake_comm_method = ModuleType( + "vllm_ascend.ops.fused_moe.moe_comm_method", + ) + fake_comm_method.get_moe_comm_method = lambda value: value + fake_afd_ubatch = ModuleType("afd_plugin.v1.worker.ubatch_wrapper") + fake_afd_ubatch.build_ubatch_additional_kwargs = lambda kwargs, metadata: { + **kwargs, + "afd_metadata": metadata, + } + fake_afd_ubatch.build_ubatch_afd_metadata = lambda metadata, _slices, _index: ( + metadata + ) + + modules = { + "vllm": fake_vllm, + "vllm.config": fake_config, + "vllm.distributed": fake_distributed, + "vllm.forward_context": fake_forward_context, + "vllm.v1": fake_vllm_v1, + "vllm.v1.worker": fake_vllm_worker, + "vllm.v1.worker.ubatch_utils": fake_ubatch_utils, + "vllm_ascend": fake_ascend, + "vllm_ascend.ops": fake_ops, + "vllm_ascend.ops.fused_moe": fake_fused_moe, + "vllm_ascend.ops.fused_moe.moe_comm_method": fake_comm_method, + "afd_plugin.v1.worker.ubatch_wrapper": fake_afd_ubatch, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + module_name = "afd_plugin.v1.worker.npu.forward_context" + return _reload_module(monkeypatch, module_name) + + +def _load_ubatch_wrapper_module(monkeypatch): + class FakeStream: + def __init__(self, device=None): + self.device = device + + def synchronize(self): + return None + + class FakeNPUGraph: + def replay(self): + return None + + fake_npu = SimpleNamespace( + Stream=FakeStream, + NPUGraph=FakeNPUGraph, + current_device=lambda: 0, + current_stream=lambda: FakeStream(), + ) + monkeypatch.setattr(torch, "npu", fake_npu, raising=False) + monkeypatch.setitem(sys.modules, "torch_npu", ModuleType("torch_npu")) + + fake_vllm = ModuleType("vllm") + fake_vllm.__path__ = [] + fake_config = ModuleType("vllm.config") + + class CUDAGraphMode: + NONE = object() + FULL = object() + PIECEWISE = object() + + fake_config.CUDAGraphMode = CUDAGraphMode + fake_config.VllmConfig = object + fake_distributed = ModuleType("vllm.distributed") + fake_distributed.get_pp_group = lambda: SimpleNamespace(is_last_rank=True) + fake_distributed.get_tensor_model_parallel_world_size = lambda: 1 + fake_distributed.tensor_model_parallel_all_gather = lambda value, dim=0: value + fake_forward_context = ModuleType("vllm.forward_context") + + class DPMetadata: + @staticmethod + def make(*_args, **_kwargs): + return None + + @contextmanager + def override_forward_context(_context): + yield + + fake_forward_context.DPMetadata = DPMetadata + fake_forward_context.ForwardContext = object + fake_forward_context.get_forward_context = lambda: None + fake_forward_context.override_forward_context = override_forward_context + fake_sequence = ModuleType("vllm.sequence") + fake_sequence.IntermediateTensors = type("IntermediateTensors", (), {}) + fake_vllm_v1 = ModuleType("vllm.v1") + fake_vllm_v1.__path__ = [] + fake_vllm_worker = ModuleType("vllm.v1.worker") + fake_vllm_worker.__path__ = [] + fake_gpu_wrapper = ModuleType("vllm.v1.worker.gpu_ubatch_wrapper") + fake_gpu_wrapper.UbatchMetadata = object + fake_gpu_wrapper.UBatchWrapper = object + + fake_ascend = ModuleType("vllm_ascend") + fake_ascend.__path__ = [] + fake_compilation = ModuleType("vllm_ascend.compilation") + fake_compilation.__path__ = [] + fake_acl_graph = ModuleType("vllm_ascend.compilation.acl_graph") + + class ACLGraphWrapper: + pass + + class GraphParams: + def __init__(self, events, workspaces, handles, attn_params): + self.events = events + self.workspaces = workspaces + self.handles = handles + self.attn_params = attn_params + + fake_acl_graph.ACLGraphWrapper = ACLGraphWrapper + fake_acl_graph.GraphParams = GraphParams + fake_acl_graph.get_graph_params = lambda: None + fake_ascend_utils = ModuleType("vllm_ascend.utils") + fake_ascend_utils.enable_sp = lambda: False + fake_child_context = ModuleType( + "afd_plugin.v1.worker.npu.forward_context", + ) + fake_child_context.create_ascend_forward_context = lambda *_args, **_kwargs: ( + SimpleNamespace() + ) + fake_ubatching = ModuleType("afd_plugin.v1.worker.npu.ubatching") + fake_ubatching.AscendUBatchContext = object + fake_ubatching.make_ubatch_contexts = lambda **_kwargs: [] + + modules = { + "vllm": fake_vllm, + "vllm.config": fake_config, + "vllm.distributed": fake_distributed, + "vllm.forward_context": fake_forward_context, + "vllm.sequence": fake_sequence, + "vllm.v1": fake_vllm_v1, + "vllm.v1.worker": fake_vllm_worker, + "vllm.v1.worker.gpu_ubatch_wrapper": fake_gpu_wrapper, + "vllm_ascend": fake_ascend, + "vllm_ascend.compilation": fake_compilation, + "vllm_ascend.compilation.acl_graph": fake_acl_graph, + "vllm_ascend.utils": fake_ascend_utils, + "afd_plugin.v1.worker.npu.forward_context": fake_child_context, + "afd_plugin.v1.worker.npu.ubatching": fake_ubatching, + } + for name, module in modules.items(): + monkeypatch.setitem(sys.modules, name, module) + + mla_graph_name = "afd_plugin.v1.worker.npu.mla_graph" + _reload_module(monkeypatch, mla_graph_name) + wrapper_name = "afd_plugin.v1.worker.npu.npu_ubatch_wrapper" + return _reload_module(monkeypatch, wrapper_name) + + +def _parent_forward_context(): + return SimpleNamespace( + additional_kwargs={}, + all_moe_layers={}, + moe_comm_type="mc2", + in_profile_run=False, + capturing=False, + mmrs_fusion=False, + flash_comm_v1_enabled=False, + flashcomm_v2_enabled=False, + is_first_layer=True, + layer_idx=0, + prefetch_mlp_gate_up_proj=False, + prefetch_mlp_down_proj=False, + model_instance=None, + is_draft_model=False, + is_draft_model_prefill=False, + draft_attn_metadatas=None, + max_tokens_across_pcp=None, + mc2_mask=None, + ) + + +def _two_slices(first, second): + return [ + SimpleNamespace( + request_slice=slice(0, 1), + token_slice=slice(0, first), + num_tokens=first, + ), + SimpleNamespace( + request_slice=slice(1, 2), + token_slice=slice(first, first + second), + num_tokens=second, + ), + ] + + +def _batch_descriptor( + num_tokens=8, + *, + has_lora=False, + num_active_loras=0, +): + return SimpleNamespace( + num_tokens=num_tokens, + has_lora=has_lora, + num_active_loras=num_active_loras, + ) + + +def _new_wrapper_for_unit_test(wrapper_module, *, mla_full_graph_enabled): + wrapper = object.__new__(wrapper_module.AscendUBatchWrapper) + wrapper.mla_full_graph_enabled = mla_full_graph_enabled + wrapper.cudagraphs = {} + wrapper.cudagraph_wrapper = None + wrapper.runnable = lambda **_kwargs: "eager" + wrapper.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(data_parallel_size=1), + ) + return wrapper + + +class _RecordingCallable: + def __init__(self, result): + self.result = result + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + return self.result + + +def test_ubatch_wrapper_rejects_enpu(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + + with pytest.raises(AssertionError, match="does not support ENPU"): + wrapper_module.AscendUBatchWrapper( + lambda: None, + SimpleNamespace(compilation_config=object()), + wrapper_module.CUDAGraphMode.NONE, + torch.device("cpu"), + enable_enpu=True, + ) + + +def test_npu_graph_key_separates_stage_shapes_and_lora(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + + keys = { + wrapper_module.AscendNPUGraphKey((4, 4), False, 0), + wrapper_module.AscendNPUGraphKey((3, 5), False, 0), + wrapper_module.AscendNPUGraphKey((4, 4), True, 1), + wrapper_module.AscendNPUGraphKey((4, 4), True, 2), + } + + assert len(keys) == 4 + + +def test_merge_mla_graph_params_is_layer_major_ubatch_minor(monkeypatch): + mla_graph = _load_mla_graph_module(monkeypatch) + workspace = object() + stage_params = ( + _graph_params( + 4, + workspace, + ["e00", "e10"], + ["h00", "h10"], + ["p00", "p10"], + ), + _graph_params( + 4, + workspace, + ["e01", "e11"], + ["h01", "h11"], + ["p01", "p11"], + ), + ) + metadata = [ + {"layer0": "m00", "layer1": "m10"}, + {"layer0": "m01", "layer1": "m11"}, + ] + + merged_metadata, merged = mla_graph.merge_mla_graph_params( + metadata, + stage_params, + 4, + ) + + assert list(merged_metadata) == [ + ("layer0", 0), + ("layer0", 1), + ("layer1", 0), + ("layer1", 1), + ] + assert list(merged_metadata.values()) == ["m00", "m01", "m10", "m11"] + assert merged.events[4] == ["e00", "e01", "e10", "e11"] + assert merged.handles[4] == ["h00", "h01", "h10", "h11"] + assert merged.attn_params[4] == ["p00", "p01", "p10", "p11"] + assert merged.workspaces[4] is workspace + + +def test_merge_mla_graph_params_requires_two_metadata_stages(monkeypatch): + mla_graph = _load_mla_graph_module(monkeypatch) + workspace = object() + graph_params = ( + _graph_params(4, workspace, ["e0"], ["h0"], ["p0"]), + _graph_params(4, workspace, ["e1"], ["h1"], ["p1"]), + ) + + with pytest.raises(RuntimeError, match="exactly two metadata stages"): + mla_graph.merge_mla_graph_params( + [{"layer0": "m0"}], + graph_params, + 4, + ) + + +def test_merge_mla_graph_params_requires_matching_layer_order(monkeypatch): + mla_graph = _load_mla_graph_module(monkeypatch) + workspace = object() + graph_params = ( + _graph_params( + 4, + workspace, + ["e00", "e10"], + ["h00", "h10"], + ["p00", "p10"], + ), + _graph_params( + 4, + workspace, + ["e01", "e11"], + ["h01", "h11"], + ["p01", "p11"], + ), + ) + + with pytest.raises(RuntimeError, match="layer order differs"): + mla_graph.merge_mla_graph_params( + [ + {"layer0": "m00", "layer1": "m10"}, + {"layer1": "m11", "layer0": "m01"}, + ], + graph_params, + 4, + ) + + +def test_merge_mla_graph_params_requires_shared_workspace(monkeypatch): + mla_graph = _load_mla_graph_module(monkeypatch) + graph_params = ( + _graph_params(4, object(), ["e0"], ["h0"], ["p0"]), + _graph_params(4, object(), ["e1"], ["h1"], ["p1"]), + ) + + with pytest.raises(RuntimeError, match="shared FIA workspace"): + mla_graph.merge_mla_graph_params( + [{"layer0": "m0"}, {"layer0": "m1"}], + graph_params, + 4, + ) + + +def test_merge_mla_graph_params_rejects_record_count_mismatch(monkeypatch): + mla_graph = _load_mla_graph_module(monkeypatch) + workspace = object() + graph_params = ( + _graph_params(4, workspace, ["e0"], ["h0"], ["p0"]), + _graph_params(4, workspace, [], [], []), + ) + + with pytest.raises(RuntimeError, match="record count mismatch"): + mla_graph.merge_mla_graph_params( + [{"layer0": "m0"}, {"layer0": "m1"}], + graph_params, + 4, + ) + + +def test_override_mla_graph_params_restores_context_after_error(monkeypatch): + mla_graph = _load_mla_graph_module(monkeypatch) + original_metadata = [{"layer0": "old"}] + original_kwargs = {"afd_metadata": object()} + context = SimpleNamespace( + attn_metadata=original_metadata, + additional_kwargs=original_kwargs, + ) + merged_metadata = {("layer0", 0): "new"} + merged_params = object() + + with ( + pytest.raises(RuntimeError, match="update failed"), + mla_graph.override_mla_graph_params( + context, + merged_metadata, + merged_params, + ), + ): + assert context.attn_metadata is merged_metadata + assert context.additional_kwargs is not original_kwargs + assert ( + context.additional_kwargs["afd_metadata"] + is (original_kwargs["afd_metadata"]) + ) + assert context.additional_kwargs["afd_mla_graph_params"] is merged_params + raise RuntimeError("update failed") + + assert context.attn_metadata is original_metadata + assert context.additional_kwargs is original_kwargs + + +def test_child_forward_context_installs_mla_capture_registry(monkeypatch): + forward_context = _load_forward_context_module(monkeypatch) + registry = object() + + child = forward_context.create_ascend_forward_context( + _parent_forward_context(), + attn_metadata=None, + vllm_config=SimpleNamespace( + compilation_config=SimpleNamespace(static_forward_context={}), + ), + ubatch_slices=_two_slices(4, 4), + ubatch_num=1, + mla_graph_params=registry, + ) + + assert child.capturing is True + assert child.additional_kwargs["afd_mla_graph_params"] is registry + assert child.ubatch_idx == 1 + assert child.num_ubatches == 2 + + +def test_mla_capture_params_isolate_records_and_share_workspace(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + workspace = object() + monkeypatch.setattr( + wrapper_module, + "get_graph_params", + lambda: SimpleNamespace(workspaces={8: workspace}), + raising=False, + ) + + registries = wrapper._new_mla_capture_params((4, 4)) + + assert registries[0] is not registries[1] + assert registries[0].events[4] is not registries[1].events[4] + assert registries[0].handles[4] is not registries[1].handles[4] + assert registries[0].attn_params[4] is not registries[1].attn_params[4] + assert registries[0].workspaces[4] is workspace + assert registries[1].workspaces[4] is workspace + + +def test_mla_capture_params_require_equal_stage_tokens(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + monkeypatch.setattr( + wrapper_module, + "get_graph_params", + lambda: SimpleNamespace(workspaces={8: object()}), + ) + + with pytest.raises(RuntimeError, match="equal padded token counts"): + wrapper._new_mla_capture_params((2, 6)) + + +def test_mla_capture_params_require_single_graph_workspace(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + monkeypatch.setattr( + wrapper_module, + "get_graph_params", + lambda: SimpleNamespace(workspaces={8: None}), + ) + + with pytest.raises(RuntimeError, match="single-batch FIA workspace"): + wrapper._new_mla_capture_params((4, 4)) + + +def test_single_full_graph_uses_inner_wrapper_when_ubatch_total_matches( + monkeypatch, +): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + wrapper.cudagraphs[8] = SimpleNamespace( + aclgraph=SimpleNamespace(replay=lambda: None), + outputs="ubatch", + ) + wrapper.cudagraph_wrapper = _RecordingCallable("single") + context = SimpleNamespace( + ubatch_slices=None, + cudagraph_runtime_mode=wrapper_module.CUDAGraphMode.FULL, + batch_descriptor=_batch_descriptor(), + ) + monkeypatch.setattr( + wrapper_module, + "get_forward_context", + lambda: context, + ) + + result = wrapper(input_ids="ids") + + assert result == "single" + assert wrapper.cudagraph_wrapper.calls == [{"input_ids": "ids"}] + + +def test_full_graph_capture_passes_shape_key_and_mla_registries(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + workspace = object() + monkeypatch.setattr( + wrapper_module, + "get_graph_params", + lambda: SimpleNamespace(workspaces={8: workspace}), + ) + context = SimpleNamespace( + ubatch_slices=_two_slices(4, 4), + cudagraph_runtime_mode=wrapper_module.CUDAGraphMode.FULL, + batch_descriptor=_batch_descriptor(), + attn_metadata=[{"layer0": "m0"}, {"layer0": "m1"}], + is_draft_model=False, + max_tokens_across_pcp=0, + ) + monkeypatch.setattr( + wrapper_module, + "get_forward_context", + lambda: context, + ) + captured = {} + + def make_ubatch_metadata(*_args, mla_graph_params=None, **_kwargs): + captured["make_params"] = mla_graph_params + return ["metadata"] + + def capture_ubatches( + ubatch_metadata, + model, + *, + graph_key, + mla_graph_params, + ): + captured["metadata"] = ubatch_metadata + captured["model"] = model + captured["graph_key"] = graph_key + captured["capture_params"] = mla_graph_params + return "captured" + + wrapper._make_ubatch_metadata = make_ubatch_metadata + wrapper._capture_ubatches = capture_ubatches + + result = wrapper( + input_ids=object(), + positions=object(), + intermediate_tensors=None, + inputs_embeds=None, + ) + + assert result == "captured" + assert captured["graph_key"] == wrapper_module.AscendNPUGraphKey( + (4, 4), + False, + 0, + ) + assert captured["make_params"] is captured["capture_params"] + assert captured["capture_params"][0].workspaces[4] is workspace + assert captured["capture_params"][1].workspaces[4] is workspace + + +def test_mla_graph_replay_updates_child_params_each_time_in_runtime_order(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + calls = [] + workspace = object() + graph_params = ( + _graph_params(4, workspace, ["e0"], ["h0"], ["p0"]), + _graph_params(4, workspace, ["e1"], ["h1"], ["p1"]), + ) + graph = SimpleNamespace(replay=lambda: calls.append("replay")) + wrapper.cudagraphs[wrapper_module.AscendNPUGraphKey((4, 4), False, 0)] = ( + wrapper_module.AscendNPUGraphMetaData( + aclgraph=graph, + ubatch_metadata=[], + outputs="output", + mla_graph_params=graph_params, + ) + ) + original_metadata = [{"layer0": "m0"}, {"layer0": "m1"}] + original_kwargs = {"afd_metadata": object()} + context = SimpleNamespace( + ubatch_slices=_two_slices(4, 4), + cudagraph_runtime_mode=wrapper_module.CUDAGraphMode.FULL, + batch_descriptor=_batch_descriptor(), + attn_metadata=original_metadata, + additional_kwargs=original_kwargs, + is_draft_model=False, + max_tokens_across_pcp=None, + ) + monkeypatch.setattr( + wrapper_module, + "get_forward_context", + lambda: context, + ) + monkeypatch.setattr( + wrapper_module.torch.npu, + "current_stream", + lambda: SimpleNamespace( + synchronize=lambda: calls.append("synchronize"), + ), + ) + + def update_params(active_context, num_tokens, positions): + calls.append("update") + assert num_tokens == 4 + assert positions is position_tensor + assert list(active_context.attn_metadata) == [ + ("layer0", 0), + ("layer0", 1), + ] + merged = active_context.additional_kwargs["afd_mla_graph_params"] + assert merged.events[4] == ["e0", "e1"] + assert merged.handles[4] == ["h0", "h1"] + assert merged.attn_params[4] == ["p0", "p1"] + + wrapper.full_graph_params_updater = update_params + position_tensor = object() + + results = [ + wrapper( + input_ids=object(), + positions=position_tensor, + intermediate_tensors=None, + inputs_embeds=None, + ) + for _ in range(2) + ] + + assert results == ["output", "output"] + assert calls == ["synchronize", "replay", "update"] * 2 + assert context.attn_metadata is original_metadata + assert context.additional_kwargs is original_kwargs + assert context.dbo_enabled is True + + +def test_non_mla_graph_replay_keeps_stream_fence(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=False, + ) + calls = [] + wrapper.cudagraphs[wrapper_module.AscendNPUGraphKey((4, 4), False, 0)] = ( + wrapper_module.AscendNPUGraphMetaData( + aclgraph=SimpleNamespace( + replay=lambda: calls.append("replay"), + ), + ubatch_metadata=[], + outputs="output", + ) + ) + context = SimpleNamespace( + ubatch_slices=_two_slices(4, 4), + cudagraph_runtime_mode=wrapper_module.CUDAGraphMode.FULL, + batch_descriptor=_batch_descriptor(), + attn_metadata=None, + ) + monkeypatch.setattr( + wrapper_module, + "get_forward_context", + lambda: context, + ) + monkeypatch.setattr( + wrapper_module.torch.npu, + "current_stream", + lambda: SimpleNamespace( + synchronize=lambda: calls.append("synchronize"), + ), + ) + + result = wrapper( + input_ids=object(), + positions=object(), + intermediate_tensors=None, + inputs_embeds=None, + ) + + assert result == "output" + assert calls == ["synchronize", "replay"] + assert context.dbo_enabled is True + + +@pytest.mark.parametrize( + ( + "mla_enabled", + "ubatch_slices", + "runtime_mode_name", + "max_tokens_across_pcp", + "expected", + ), + [ + (True, _two_slices(4, 4), "FULL", None, True), + (True, _two_slices(4, 4), "FULL", 0, True), + (False, _two_slices(4, 4), "FULL", None, False), + (True, None, "FULL", None, False), + (True, _two_slices(4, 4), "PIECEWISE", None, False), + ], +) +def test_wrapper_owns_only_supported_mla_full_graph_updates( + monkeypatch, + mla_enabled, + ubatch_slices, + runtime_mode_name, + max_tokens_across_pcp, + expected, +): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=mla_enabled, + ) + context = SimpleNamespace( + ubatch_slices=ubatch_slices, + cudagraph_runtime_mode=getattr( + wrapper_module.CUDAGraphMode, + runtime_mode_name, + ), + max_tokens_across_pcp=max_tokens_across_pcp, + ) + + assert wrapper.owns_full_graph_update(context) is expected + + +def test_mla_full_graph_rejects_pcp_before_graph_routing(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + context = SimpleNamespace( + ubatch_slices=_two_slices(4, 4), + cudagraph_runtime_mode=wrapper_module.CUDAGraphMode.FULL, + batch_descriptor=_batch_descriptor(), + attn_metadata=[{"layer0": "m0"}, {"layer0": "m1"}], + max_tokens_across_pcp=8, + ) + monkeypatch.setattr( + wrapper_module, + "get_forward_context", + lambda: context, + ) + + with pytest.raises( + RuntimeError, + match="does not support PCP execution", + ): + wrapper( + input_ids=object(), + positions=object(), + intermediate_tensors=None, + inputs_embeds=None, + ) + + assert wrapper.cudagraphs == {} + + +def test_mla_graph_replay_rejects_missing_capture_registry(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + wrapper.full_graph_params_updater = lambda *_args: None + replay_calls = [] + graph_metadata = wrapper_module.AscendNPUGraphMetaData( + aclgraph=SimpleNamespace( + replay=lambda: replay_calls.append("replay"), + ), + ubatch_metadata=[], + mla_graph_params=None, + ) + context = SimpleNamespace( + attn_metadata=[{"layer0": "m0"}, {"layer0": "m1"}], + additional_kwargs={}, + ) + + with pytest.raises(RuntimeError, match="no capture registry"): + wrapper._replay_mla_graph(graph_metadata, context, 4, object()) + + assert replay_calls == [] + + +def test_mla_graph_replay_rejects_missing_updater_before_replay(monkeypatch): + wrapper_module = _load_ubatch_wrapper_module(monkeypatch) + wrapper = _new_wrapper_for_unit_test( + wrapper_module, + mla_full_graph_enabled=True, + ) + wrapper.full_graph_params_updater = None + workspace = object() + graph_metadata = wrapper_module.AscendNPUGraphMetaData( + aclgraph=SimpleNamespace(replay=lambda: replay_calls.append("replay")), + ubatch_metadata=[], + mla_graph_params=( + _graph_params(4, workspace, ["e0"], ["h0"], ["p0"]), + _graph_params(4, workspace, ["e1"], ["h1"], ["p1"]), + ), + ) + context = SimpleNamespace( + attn_metadata=[{"layer0": "m0"}, {"layer0": "m1"}], + additional_kwargs={}, + ) + replay_calls = [] + + with pytest.raises(RuntimeError, match="no parameter updater"): + wrapper._replay_mla_graph(graph_metadata, context, 4, object()) + + assert replay_calls == [] diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index ad323fd2..9fe90c4c 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -20,7 +20,6 @@ AFDA2FTransferPayload, AFDControlPayload, AFDF2ATransferPayload, - AFDForwardContextMetadata, AFDTransferContext, AFDTransferMetadata, AFDTransferState, @@ -192,6 +191,9 @@ def _vllm_config( role="attention", connector="CAMP2pAFDConnector", extra_config=None, + use_mla=False, + cudagraph_mode="FULL", + speculative_config=None, **parallel_overrides, ): async_dp = bool(parallel_overrides.pop("async_dp", False)) @@ -209,11 +211,21 @@ def _vllm_config( }, }, parallel_config=_parallel_config(**parallel_overrides), - model_config=SimpleNamespace(enforce_eager=True), + model_config=SimpleNamespace( + enforce_eager=True, + hf_text_config=SimpleNamespace(), + use_mla=use_mla, + ), compilation_config=SimpleNamespace( - cudagraph_mode=SimpleNamespace(name="FULL"), + cudagraph_mode=SimpleNamespace( + name=cudagraph_mode, + has_full_cudagraphs=lambda: ( + cudagraph_mode in {"FULL", "FULL_DECODE_ONLY", "FULL_AND_PIECEWISE"} + ), + ), fast_moe_cold_start=False, ), + speculative_config=speculative_config, ) @@ -251,6 +263,102 @@ def _new_ffn_worker(): return object.__new__(AFDNPUFFNWorker) +@pytest.mark.parametrize( + ("wrapper_owns_update", "expected_updates"), + [(True, 0), (False, 1)], +) +def test_npu_attention_runner_skips_outer_update_only_for_owned_graph( + monkeypatch, + wrapper_owns_update, + expected_updates, +): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import attention_model_runner + + class FakeUBatchWrapper: + def owns_full_graph_update(self, _forward_context): + return wrapper_owns_update + + def __call__(self, **_model_inputs): + return "hidden_states" + + forward_context = SimpleNamespace( + dbo_enabled=False, + flash_comm_v1_enabled=False, + ) + monkeypatch.setattr( + attention_model_runner, + "AscendUBatchWrapper", + FakeUBatchWrapper, + ) + monkeypatch.setattr( + attention_model_runner, + "get_forward_context", + lambda: forward_context, + ) + + runner = object.__new__( + attention_model_runner.AFDNPUAttentionModelRunner, + ) + runner.enable_enpu = False + runner.model = FakeUBatchWrapper() + runner.ubatch_slices = None + runner._install_afd_metadata_on_forward_context = lambda _context: None + runner._install_async_moe_ubatch_metadata_on_forward_context = lambda _context: None + updates = [] + runner._update_full_graph_params_if_needed = lambda *args: updates.append(args) + + result = runner._model_forward( + 8, + input_ids=None, + positions=object(), + intermediate_tensors=None, + inputs_embeds=None, + ) + + assert result == "hidden_states" + assert len(updates) == expected_updates + + +def test_npu_attention_runner_installs_mla_graph_wrapper(monkeypatch): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import attention_model_runner + + captured = {} + + class RecordingUBatchWrapper: + def __init__(self, *args, **kwargs): + captured["args"] = args + captured["kwargs"] = kwargs + + monkeypatch.setattr( + attention_model_runner, + "AscendUBatchWrapper", + RecordingUBatchWrapper, + ) + runner = object.__new__( + attention_model_runner.AFDNPUAttentionModelRunner, + ) + runner.model = "model" + runner.device = "npu" + runner.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(use_mla=True), + ) + runner.compilation_config = SimpleNamespace( + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: True), + ) + runner.use_sparse = False + runner.enable_enpu = False + + runner._install_ascend_ubatch_wrapper() + + assert captured["args"][:2] == ("model", runner.vllm_config) + assert captured["kwargs"]["mla_full_graph_enabled"] is True + assert captured["kwargs"]["enable_enpu"] is False + updater = captured["kwargs"]["full_graph_params_updater"] + assert updater.__self__ is runner + + def test_npu_attention_runner_builds_and_sets_metadata(): runner = _new_attention_runner() runner.vllm_config = _vllm_config(role="attention") @@ -564,85 +672,6 @@ def is_empty(self): sys.modules[module_name] = original_module -def test_npu_create_ascend_forward_context_marks_current_ubatch(monkeypatch): - _require_npu_runtime() - from afd_plugin.v1.worker.npu import forward_context as forward_context_module - - monkeypatch.setattr( - forward_context_module, - "get_tensor_model_parallel_world_size", - lambda: 1, - ) - monkeypatch.setattr( - forward_context_module, - "get_dp_group", - lambda: SimpleNamespace(world_size=1), - ) - monkeypatch.setattr( - forward_context_module, - "get_moe_comm_method", - lambda moe_comm_type: f"method:{moe_comm_type}", - ) - afd_metadata = AFDForwardContextMetadata( - tokens_start_loc=[0, 4], - requests_start_loc=[0, 1], - stage_idx=0, - connector=object(), - tokens_lens=[4, 3], - num_stages=2, - tokens_unpadded_lens=[4, 3], - ) - cur_forward_context = SimpleNamespace( - additional_kwargs={"afd_metadata": afd_metadata}, - all_moe_layers={}, - moe_comm_type="mc2", - in_profile_run=False, - capturing=False, - mmrs_fusion=False, - flash_comm_v1_enabled=False, - flashcomm_v2_enabled=False, - is_first_layer=True, - layer_idx=0, - prefetch_mlp_gate_up_proj=False, - prefetch_mlp_down_proj=False, - model_instance=None, - is_draft_model=False, - is_draft_model_prefill=False, - draft_attn_metadatas=None, - max_tokens_across_pcp=None, - mc2_mask=None, - ) - ubatch_slices = [ - SimpleNamespace( - request_slice=slice(0, 1), - token_slice=slice(0, 4), - num_tokens=4, - ), - SimpleNamespace( - request_slice=slice(1, 2), - token_slice=slice(4, 7), - num_tokens=3, - ), - ] - vllm_config = SimpleNamespace( - compilation_config=SimpleNamespace(static_forward_context={}), - ) - - new_forward_context = forward_context_module.create_ascend_forward_context( - cur_forward_context, - attn_metadata=None, - vllm_config=vllm_config, - ubatch_slices=ubatch_slices, - ubatch_num=1, - ) - - child_metadata = new_forward_context.additional_kwargs["afd_metadata"] - assert new_forward_context.ubatch_idx == 1 - assert new_forward_context.num_ubatches == 2 - assert new_forward_context.num_tokens == 3 - assert child_metadata.stage_idx == 1 - - def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): _patch_ffn_forward_context(monkeypatch) runner = _new_ffn_runner() @@ -1139,6 +1168,60 @@ def test_npu_feature_validation_allows_two_ubatches_only(): fail_if_unsupported_npu_afd_features(config) +@pytest.mark.parametrize("cudagraph_mode", ["FULL", "FULL_AND_PIECEWISE"]) +def test_npu_feature_validation_requires_decode_only_full_graph_for_mla_dbo( + cudagraph_mode, +): + config = _vllm_config( + use_mla=True, + cudagraph_mode=cudagraph_mode, + enable_dbo=True, + use_ubatching=True, + num_ubatches=2, + ubatch_size=4, + ) + + with pytest.raises(RuntimeError, match="FULL_DECODE_ONLY"): + fail_if_unsupported_npu_afd_features(config) + + fail_if_unsupported_npu_afd_features( + _vllm_config( + use_mla=True, + cudagraph_mode="FULL_DECODE_ONLY", + enable_dbo=True, + use_ubatching=True, + num_ubatches=2, + ubatch_size=4, + ), + ) + + sparse_config = _vllm_config( + use_mla=True, + cudagraph_mode="FULL", + enable_dbo=True, + use_ubatching=True, + num_ubatches=2, + ubatch_size=4, + ) + sparse_config.model_config.hf_text_config = SimpleNamespace(index_topk=8) + fail_if_unsupported_npu_afd_features(sparse_config) + + +def test_npu_feature_validation_rejects_speculative_mla_dbo_full_graph(): + config = _vllm_config( + use_mla=True, + cudagraph_mode="FULL_DECODE_ONLY", + speculative_config=object(), + enable_dbo=True, + use_ubatching=True, + num_ubatches=2, + ubatch_size=4, + ) + + with pytest.raises(RuntimeError, match="does not support speculative decoding"): + fail_if_unsupported_npu_afd_features(config) + + def test_npu_async_feature_validation_requires_async_config_and_eager(): with pytest.raises(RuntimeError, match="async=true"): fail_if_unsupported_npu_afd_features(