From 5457b8e68c68896e317d7d7d1018488400cedd30 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Tue, 25 Aug 2026 13:43:17 +0800 Subject: [PATCH 1/7] feat(npu): backport eager DBO for ModelRunnerV2 Signed-off-by: lirx-pd <616517220@qq.com> --- .../backports/vllm_v026_mrv2_dbo/__init__.py | 27 ++ .../backports/vllm_v026_mrv2_dbo/execute.py | 213 ++++++++++ .../backports/vllm_v026_mrv2_dbo/runtime.py | 376 ++++++++++++++++++ .../compat/patches/config_validation.py | 35 +- .../worker/npu/attention_model_runner_v2.py | 43 ++ afd_plugin/v1/worker/npu/forward_context.py | 174 ++++++-- afd_plugin/v1/worker/npu/ubatch_runner_v2.py | 229 +++++++++++ afd_plugin/validation.py | 23 +- .../backports/test_vllm_v026_mrv2_dbo.py | 162 ++++++++ .../compat/patches/test_config_validation.py | 21 + tests/unit/v1/worker/test_model_runner_v2.py | 59 +++ tests/unit/v1/worker/test_npu_mla_graph.py | 1 + tests/unit/v1/worker/test_npu_runtime.py | 58 +++ 13 files changed, 1372 insertions(+), 49 deletions(-) create mode 100644 afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py create mode 100644 afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py create mode 100644 afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py create mode 100644 afd_plugin/v1/worker/npu/ubatch_runner_v2.py create mode 100644 tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py new file mode 100644 index 00000000..873688ab --- /dev/null +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py @@ -0,0 +1,27 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Temporary vLLM 0.26 ModelRunnerV2 eager DBO backport.""" + +from .runtime import ( + AFDBatchExecutionDescriptor, + assert_backport_required, + create_ubatch_slices, + dispatch_afd_dbo_and_sync_dp, + prepare_attn_for_ubatch, + share_metadata_builder_workspaces, + slice_input_batch, + slice_model_inputs, + use_two_metadata_builders, +) + +__all__ = [ + "AFDBatchExecutionDescriptor", + "assert_backport_required", + "create_ubatch_slices", + "dispatch_afd_dbo_and_sync_dp", + "prepare_attn_for_ubatch", + "share_metadata_builder_workspaces", + "slice_input_batch", + "slice_model_inputs", + "use_two_metadata_builders", +] diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py new file mode 100644 index 00000000..b5715c0d --- /dev/null +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py @@ -0,0 +1,213 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""vLLM 0.26 ModelRunnerV2 execute path with the eager DBO seams added.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Any + +import torch +from vllm.config import CUDAGraphMode +from vllm.forward_context import ( + BatchDescriptor, + get_forward_context, + set_forward_context, +) +from vllm.sequence import IntermediateTensors +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.cudagraph_utils import get_uniform_token_count +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.model_runner import ExecuteModelState +from vllm_ascend.worker.v2.input_batch import AscendInputBatch + +from .runtime import ( + AFDBatchExecutionDescriptor, + create_ubatch_slices, + dispatch_afd_dbo_and_sync_dp, +) + +if TYPE_CHECKING: + from vllm.v1.core.sched.output import SchedulerOutput + from vllm.v1.outputs import ModelRunnerOutput + + +def execute_model_v026_eager_dbo( + runner: Any, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + *, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + is_profile: bool = False, +) -> ModelRunnerOutput | IntermediateTensors | None: + """Execute the supported plain-decoder subset with eager DBO.""" + + if not dummy_run: + runner.update_pp_decode_requests() + runner.finish_requests(scheduler_output) + runner.free_states(scheduler_output) + runner.add_requests(scheduler_output) + runner.update_requests(scheduler_output) + runner.block_tables.apply_staged_writes() + if scheduler_output.total_num_scheduled_tokens == 0: + return runner.kv_connector.no_forward(scheduler_output) + + num_reqs = len(scheduler_output.num_scheduled_tokens) + num_tokens = int(scheduler_output.total_num_scheduled_tokens) + max_query_len = max(scheduler_output.num_scheduled_tokens.values()) + uniform_token_count = get_uniform_token_count( + num_reqs, + num_tokens, + max_query_len, + ) + batch_desc, num_tokens_across_dp = dispatch_afd_dbo_and_sync_dp( + num_reqs=num_reqs, + num_tokens=num_tokens, + uniform_token_count=uniform_token_count, + dp_size=runner.dp_size, + dp_rank=runner.dp_rank, + parallel_config=runner.parallel_config, + decode_query_len=runner.decode_query_len, + allow_ubatching=not skip_attn_for_dummy_run, + ) + if batch_desc.num_tokens == 0: + return runner.kv_connector.no_forward(scheduler_output) + + num_ubatches = ( + batch_desc.num_ubatches + if isinstance(batch_desc, AFDBatchExecutionDescriptor) + else 1 + ) + if not dummy_run: + runner.input_buffers.is_padding[:num_tokens].fill_(False) + runner.input_buffers.is_padding[ + num_tokens : batch_desc.num_tokens + ].fill_(True) + input_batch = runner.prepare_inputs(scheduler_output, batch_desc) + block_tables, slot_mappings = runner.prepare_attn(input_batch) + runner.model_state.preprocess_state( + input_batch, + block_tables, + runner.kv_cache_config, + runner.req_states.num_computed_tokens.gpu, + ) + else: + dummy_batch_cls = AscendInputBatch if num_ubatches > 1 else InputBatch + input_batch = dummy_batch_cls.make_dummy( + batch_desc.num_reqs or num_reqs, + batch_desc.num_tokens, + runner.input_buffers, + ) + if not skip_attn_for_dummy_run: + block_tables, slot_mappings = runner.prepare_dummy_attn(input_batch) + else: + block_tables = None + slot_mappings = None + + attn_metadata = None + slot_mappings_by_layer = None + ubatch_slices = None + if num_ubatches > 1: + assert runner.ubatch_runner is not None + assert block_tables is not None and slot_mappings is not None + ubatch_slices = create_ubatch_slices(input_batch, num_ubatches) + elif not (dummy_run and skip_attn_for_dummy_run): + assert block_tables is not None and slot_mappings is not None + slot_mappings_by_layer = build_slot_mappings_by_layer( + slot_mappings, + runner.kv_cache_config, + ) + attn_metadata = runner.model_state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + block_tables, + slot_mappings, + runner.attn_groups, + runner.kv_cache_config, + ) + + model_inputs = { + "input_ids": input_batch.input_ids, + "positions": input_batch.positions, + "inputs_embeds": None, + "intermediate_tensors": None, + **runner.model_state.prepare_inputs(input_batch, runner.req_states), + } + runner.eplb.prepare_forward( + runner.model_config, + input_batch.num_tokens, + ubatch_slices, + ) + + if ubatch_slices is not None: + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=False, + num_active_loras=0, + ) + with set_forward_context( + None, + runner.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + num_tokens_across_dp=num_tokens_across_dp, + batch_descriptor=batch_descriptor, + ubatch_slices=ubatch_slices, + is_padding=input_batch.is_padding, + ): + parent_context = get_forward_context() + afd_metadata = parent_context.additional_kwargs["afd_metadata"] + afd_metadata.tokens_unpadded_lens = [ + max( + 0, + min(int(stage.token_slice.stop), int(input_batch.num_tokens)) + - int(stage.token_slice.start), + ) + for stage in ubatch_slices + ] + ubatch_state = runner.ubatch_runner.prepare( + input_batch, + block_tables, + slot_mappings, + ubatch_slices, + parent_context, + ) + runner.kv_connector.pre_forward(scheduler_output) + model_output = runner.ubatch_runner.run( + runner.model, + model_inputs, + ubatch_state, + ) + else: + batch_descriptor = BatchDescriptor( + num_tokens=input_batch.num_tokens_after_padding, + has_lora=False, + num_active_loras=0, + ) + with set_forward_context( + attn_metadata, + runner.vllm_config, + num_tokens=input_batch.num_tokens_after_padding, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + num_tokens_across_dp=num_tokens_across_dp, + batch_descriptor=batch_descriptor, + slot_mapping=slot_mappings_by_layer, + is_padding=input_batch.is_padding, + ): + runner.kv_connector.pre_forward(scheduler_output) + model_output = runner.model(**model_inputs) + + assert runner.is_last_pp_rank + assert isinstance(model_output, torch.Tensor) + runner.execute_model_state = ExecuteModelState( + input_batch=input_batch, + attn_metadata=attn_metadata, + slot_mappings_by_layer=slot_mappings_by_layer, + hidden_states=model_output, + aux_hidden_states=None, + finished_req_ids=scheduler_output.finished_req_ids, + ) + return None + + +__all__ = ["execute_model_v026_eager_dbo"] diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py new file mode 100644 index 00000000..c66896c9 --- /dev/null +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py @@ -0,0 +1,376 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Runtime pieces missing from vLLM 0.26 ModelRunnerV2 eager DBO. + +The behavior is adapted from ``specture724/vllm`` branch +``feat/v2/dbo-fullcg`` at ``626fee7831``. The target ABI is vLLM +``568afb3a13``. Keep this module self-contained so it can be deleted once the +pinned vLLM release provides native ModelRunnerV2 DBO. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass, fields, replace +from typing import Any + +import numpy as np +import torch +import torch.distributed as dist +from vllm.config import CUDAGraphMode, ParallelConfig +from vllm.distributed.parallel_state import get_dp_group +from vllm.sequence import IntermediateTensors +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.cudagraph_utils import BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.ubatch_utils import ( + UBatchSlice, + UBatchSlices, + check_ubatch_thresholds, + maybe_create_ubatch_slices, +) +from vllm.v1.worker.utils import AttentionGroup + + +@dataclass(frozen=True) +class AFDBatchExecutionDescriptor(BatchExecutionDescriptor): + """v0.26 batch descriptor extended only with the DBO execution count.""" + + num_ubatches: int = 1 + + +def assert_backport_required() -> None: + """Fail when the pinned vLLM ABI no longer needs this backport.""" + + descriptor_fields = {field.name for field in fields(BatchExecutionDescriptor)} + if "num_ubatches" in descriptor_fields: + raise RuntimeError( + "vLLM already provides ModelRunnerV2 DBO descriptors; remove the " + "temporary afd-plugin v0.26 backport", + ) + + +def dispatch_afd_dbo_and_sync_dp( + *, + num_reqs: int, + num_tokens: int, + uniform_token_count: int | None, + dp_size: int, + dp_rank: int, + parallel_config: ParallelConfig, + decode_query_len: int, + allow_ubatching: bool, +) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: + """Select eager DBO consistently across all DP ranks. + + The minimum rank token count controls the threshold, while all ranks run + the maximum token count when DBO is selected. This matches the final + upstream inference rule and avoids a separate per-rank DBO vote. + """ + + if dp_size == 1: + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + ), + None, + ) + + tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu") + tensor[0][dp_rank] = num_tokens + tensor[1][dp_rank] = uniform_token_count or 0 + tensor[2][dp_rank] = int(allow_ubatching) + dist.all_reduce(tensor, group=get_dp_group().cpu_group) + + num_tokens_across_dp = tensor[0] + if torch.all(num_tokens_across_dp == 0).item(): + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=0, + num_reqs=0, + ), + None, + ) + + uniform_decode = bool( + torch.all(tensor[1] == int(decode_query_len)).item() + ) + should_ubatch = ( + bool(torch.all(tensor[2] == 1).item()) + and int(num_tokens_across_dp.max().item()) + >= int(parallel_config.num_ubatches) + and check_ubatch_thresholds( + parallel_config, + int(num_tokens_across_dp.min().item()), + uniform_decode=uniform_decode, + ) + ) + if not should_ubatch: + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + ), + num_tokens_across_dp, + ) + + padded_tokens = int(num_tokens_across_dp.max().item()) + num_tokens_across_dp.fill_(padded_tokens) + return ( + AFDBatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=padded_tokens, + num_reqs=num_reqs, + num_ubatches=int(parallel_config.num_ubatches), + ), + num_tokens_across_dp, + ) + + +def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: + """Split a DP-padded v0.26 InputBatch into microbatch views.""" + + _, padded_slices = maybe_create_ubatch_slices( + True, + input_batch.num_scheduled_tokens, + input_batch.num_tokens_after_padding, + input_batch.num_reqs_after_padding, + num_ubatches, + ) + assert padded_slices is not None + return [ + UBatchSlice( + slice( + min(stage.request_slice.start, input_batch.num_reqs - 1), + min(stage.request_slice.stop, input_batch.num_reqs_after_padding), + ), + stage.token_slice, + ) + for stage in padded_slices + ] + + +def slice_input_batch( + input_batch: InputBatch, + stage: UBatchSlice, + query_start_loc_buffer: torch.Tensor, + seq_lens_buffer: torch.Tensor, +) -> InputBatch: + """Build one microbatch while retaining the v0.26 InputBatch ABI.""" + + assert not stage.is_empty(), f"Ubatch slice {stage} is empty" + req_start = stage.request_slice.start + req_stop = stage.request_slice.stop + tok_start = stage.token_slice.start + tok_stop = stage.token_slice.stop + + num_reqs_after_padding = req_stop - req_start + num_tokens_after_padding = tok_stop - tok_start + num_reqs = max(0, min(req_stop, input_batch.num_reqs) - req_start) + num_tokens = max(0, min(tok_stop, input_batch.num_tokens) - tok_start) + + query_start_loc = query_start_loc_buffer[: num_reqs_after_padding + 1] + torch.sub( + input_batch.query_start_loc[req_start : req_stop + 1], + tok_start, + out=query_start_loc, + ) + query_start_loc.clamp_(0, num_tokens_after_padding) + query_start_loc_np = np.clip( + input_batch.query_start_loc_np[req_start : req_stop + 1] - tok_start, + 0, + num_tokens_after_padding, + ).astype(np.int32) + + seq_lens = seq_lens_buffer[:num_reqs_after_padding] + seq_lens.copy_(input_batch.seq_lens[req_start:req_stop]) + last = num_reqs_after_padding - 1 + seq_lens[last] -= ( + input_batch.query_start_loc[req_stop] - tok_stop + ).clamp_(min=0) + + seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[ + req_start:req_stop + ].clone() + truncated = max(0, int(input_batch.query_start_loc_np[req_stop]) - tok_stop) + if truncated: + seq_lens_cpu_upper_bound[-1] -= truncated + + dcp_local_seq_lens = input_batch.dcp_local_seq_lens + if dcp_local_seq_lens is not None: + dcp_local_seq_lens = dcp_local_seq_lens[req_start:req_stop] + + return replace( + input_batch, + req_ids=input_batch.req_ids[req_start : min(req_stop, input_batch.num_reqs)], + num_reqs=num_reqs, + num_reqs_after_padding=num_reqs_after_padding, + idx_mapping=input_batch.idx_mapping[req_start:req_stop], + idx_mapping_np=input_batch.idx_mapping_np[req_start:req_stop], + num_scheduled_tokens=np.diff(query_start_loc_np)[:num_reqs], + num_tokens=num_tokens, + num_tokens_after_padding=num_tokens_after_padding, + query_start_loc=query_start_loc, + query_start_loc_np=query_start_loc_np, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, + dcp_local_seq_lens=dcp_local_seq_lens, + num_computed_tokens_np=input_batch.num_computed_tokens_np[ + req_start:req_stop + ], + prefill_len_np=input_batch.prefill_len_np[req_start:req_stop], + num_computed_prefill_tokens_np=input_batch.num_computed_prefill_tokens_np[ + req_start:req_stop + ], + is_prefilling_np=input_batch.is_prefilling_np[req_start:req_stop], + max_seq_len_np=( + None + if input_batch.max_seq_len_np is None + else input_batch.max_seq_len_np[req_start:req_stop] + ), + input_ids=input_batch.input_ids[tok_start:tok_stop], + positions=input_batch.positions[tok_start:tok_stop], + is_padding=input_batch.is_padding[tok_start:tok_stop], + prompt_lens=( + None + if input_batch.prompt_lens is None + else input_batch.prompt_lens[req_start:req_stop] + ), + ) + + +def slice_model_inputs( + model_inputs: dict[str, Any], token_slice: slice +) -> dict[str, Any]: + """Narrow the model's per-token inputs to one microbatch.""" + + sliced = dict(model_inputs) + for key in ("input_ids", "inputs_embeds"): + value = model_inputs.get(key) + if value is not None: + sliced[key] = value[token_slice] + positions = model_inputs["positions"] + sliced["positions"] = ( + positions[:, token_slice] if positions.ndim == 2 else positions[token_slice] + ) + intermediate_tensors = model_inputs.get("intermediate_tensors") + if intermediate_tensors is not None: + sliced["intermediate_tensors"] = intermediate_tensors[token_slice] + return sliced + + +def merge_ubatch_outputs(outputs: list[Any]) -> Any: + """Reassemble the native ModelRunnerV2 output structure.""" + + first = outputs[0] + if isinstance(first, IntermediateTensors): + return IntermediateTensors( + { + key: torch.cat([output.tensors[key] for output in outputs], dim=0) + for key in first.tensors + }, + ) + if isinstance(first, tuple): + hidden_states = torch.cat([output[0] for output in outputs], dim=0) + auxiliary = [ + torch.cat([output[1][index] for output in outputs], dim=0) + for index in range(len(first[1])) + ] + return hidden_states, auxiliary + return torch.cat(outputs, dim=0) + + +@contextmanager +def use_two_metadata_builders() -> Iterator[None]: + """Create two builders during one AFD DBO KV-cache initialization.""" + + original = AttentionGroup.create_metadata_builders + + def create_metadata_builders( + group: AttentionGroup, + vllm_config, + device, + kernel_block_size: int | None = None, + num_metadata_builders: int = 1, + ) -> None: + del num_metadata_builders + original( + group, + vllm_config, + device, + kernel_block_size, + num_metadata_builders=2, + ) + + try: + AttentionGroup.create_metadata_builders = create_metadata_builders + yield + finally: + AttentionGroup.create_metadata_builders = original + + +def share_metadata_builder_workspaces( + attn_groups: list[list[AttentionGroup]], +) -> None: + """Share the backend workspace while retaining independent builders.""" + + workspace = None + for groups in attn_groups: + for group in groups: + for builder in group.metadata_builders: + if workspace is None and hasattr(builder, "_get_workspace_buffer"): + workspace = builder._get_workspace_buffer() + elif workspace is not None and hasattr(builder, "set_workspace_buffer"): + builder.set_workspace_buffer(workspace) + + +def prepare_attn_for_ubatch( + model_state: ModelState, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + ubatch_index: int, +) -> dict[str, Any]: + """Select one of the two builders without changing upstream signatures.""" + + if ubatch_index == 0: + return model_state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + ) + + swapped: list[AttentionGroup] = [] + for groups in attn_groups: + for group in groups: + group.metadata_builders[0], group.metadata_builders[ubatch_index] = ( + group.metadata_builders[ubatch_index], + group.metadata_builders[0], + ) + swapped.append(group) + try: + return model_state.prepare_attn( + input_batch, + CUDAGraphMode.NONE, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + ) + finally: + for group in reversed(swapped): + group.metadata_builders[0], group.metadata_builders[ubatch_index] = ( + group.metadata_builders[ubatch_index], + group.metadata_builders[0], + ) diff --git a/afd_plugin/compat/patches/config_validation.py b/afd_plugin/compat/patches/config_validation.py index 78ec51c0..02776bbf 100644 --- a/afd_plugin/compat/patches/config_validation.py +++ b/afd_plugin/compat/patches/config_validation.py @@ -120,18 +120,29 @@ def __post_init__(self): """Verify configs are valid & consistent with each other.""" assert _original_vllm_config_post_init is not None - if not _should_relax_vllm_config_backend(self): + relax_backend = _should_relax_vllm_config_backend(self) + relax_v2_dbo = _should_relax_npu_v2_dbo_validation(self) + if not relax_backend and not relax_v2_dbo: return _original_vllm_config_post_init(self) - # ### PATCH START: AFD repeated ubatching backend validation + # ### PATCH START: AFD repeated ubatching and MRV2 DBO validation parallel_config = self.parallel_config original_backend = parallel_config.all2all_backend - parallel_config.all2all_backend = _AFD_TEMP_BACKEND + if relax_backend: + parallel_config.all2all_backend = _AFD_TEMP_BACKEND + if relax_v2_dbo: + original_enable_dbo = parallel_config.enable_dbo + original_ubatch_size = parallel_config.ubatch_size + parallel_config.enable_dbo = False + parallel_config.ubatch_size = 0 try: result = _original_vllm_config_post_init(self) finally: parallel_config.all2all_backend = original_backend - # ### PATCH END: AFD repeated ubatching backend validation + if relax_v2_dbo: + parallel_config.enable_dbo = original_enable_dbo + parallel_config.ubatch_size = original_ubatch_size + # ### PATCH END: AFD repeated ubatching and MRV2 DBO validation return result @@ -195,6 +206,22 @@ def _should_relax_vllm_config_backend(vllm_config: VllmConfig) -> bool: } +def _should_relax_npu_v2_dbo_validation(vllm_config: VllmConfig) -> bool: + if not _is_target_vllm_compatible(): + return False + afd_config = parse_optional_afd_config(vllm_config) + if afd_config is None or afd_config.connector != "CAMP2pAFDConnector": + return False + + from vllm.platforms import current_platform + + return bool( + current_platform.device_type == "npu" + and vllm_config.use_v2_model_runner + and vllm_config.parallel_config.use_ubatching + ) + + def _is_target_vllm_compatible() -> bool: try: import vllm 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 656f4f12..40307c14 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -169,6 +169,7 @@ def __init__( self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 + self.ubatch_runner = None self.prof = create_afd_npu_profiler("attention") except BaseException: try: @@ -211,6 +212,35 @@ def load_model( if not self.connector.is_initialized: self.connector.init_afd_connector() + def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: + """Initialize native state and the temporary v0.26 eager DBO runner.""" + + if not self.vllm_config.parallel_config.use_ubatching: + super().initialize_kv_cache(kv_cache_config) + return + + from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( + assert_backport_required, + share_metadata_builder_workspaces, + use_two_metadata_builders, + ) + from afd_plugin.v1.worker.npu.ubatch_runner_v2 import ( + AFDAscendUBatchRunnerV2, + ) + + assert_backport_required() + with use_two_metadata_builders(): + super().initialize_kv_cache(kv_cache_config) + share_metadata_builder_workspaces(self.attn_groups) + self.ubatch_runner = AFDAscendUBatchRunnerV2( + self.vllm_config, + self.device, + self.model_state, + self.attn_groups, + self.kv_cache_config, + self.max_num_reqs, + ) + # Patch reason: vLLM v0.26.0 prepares FULL graph inputs before each warmup # and formal capture forward, outside torch.cuda.graph, but does not expose # that lifecycle to AFD's control plane. @@ -381,6 +411,19 @@ def execute_model( self.install_afd_metadata_on_forward_context, ), ): + if self.vllm_config.parallel_config.use_ubatching: + from afd_plugin.compat.backports.vllm_v026_mrv2_dbo.execute import ( + execute_model_v026_eager_dbo, + ) + + return execute_model_v026_eager_dbo( + self, + scheduler_output, + intermediate_tensors, + dummy_run=dummy_run, + skip_attn_for_dummy_run=skip_attn_for_dummy_run, + is_profile=is_profile, + ) return super().execute_model( scheduler_output, intermediate_tensors, diff --git a/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index e7ca5d2e..43d76e0a 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/afd_plugin/v1/worker/npu/forward_context.py @@ -26,6 +26,28 @@ from vllm_ascend.compilation.acl_graph import GraphParams +def _get_ascend_extra( + forward_context: ForwardContext, + vllm_config: VllmConfig, + name: str, +): + if vllm_config.use_v2_model_runner: + return forward_context.additional_kwargs.get(name) + return getattr(forward_context, name) + + +def _set_ascend_extra( + forward_context: ForwardContext, + vllm_config: VllmConfig, + name: str, + value, +) -> None: + if vllm_config.use_v2_model_runner: + forward_context.additional_kwargs[name] = value + else: + setattr(forward_context, name, value) + + def create_ascend_forward_context( cur_forward_context: ForwardContext, attn_metadata, @@ -75,77 +97,143 @@ def create_ascend_forward_context( tp_world_size = get_tensor_model_parallel_world_size() dp_world_size = get_dp_group().world_size - new_forward_context.moe_comm_type = cur_forward_context.moe_comm_type - new_forward_context.moe_comm_method = get_moe_comm_method( - new_forward_context.moe_comm_type + moe_comm_type = _get_ascend_extra( + cur_forward_context, + vllm_config, + "moe_comm_type", ) - new_forward_context.in_profile_run = cur_forward_context.in_profile_run - new_forward_context.capturing = ( - mla_graph_params is not None or cur_forward_context.capturing + _set_ascend_extra( + new_forward_context, + vllm_config, + "moe_comm_type", + moe_comm_type, ) - 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) - new_forward_context.num_ubatches = len(ubatch_slices) - new_forward_context.flash_comm_v1_enabled = ( - cur_forward_context.flash_comm_v1_enabled + _set_ascend_extra( + new_forward_context, + vllm_config, + "moe_comm_method", + get_moe_comm_method(moe_comm_type), ) - new_forward_context.pad_size = 0 - new_forward_context.is_first_layer = cur_forward_context.is_first_layer - new_forward_context.layer_idx = cur_forward_context.layer_idx - new_forward_context.prefetch_mlp_gate_up_proj = ( - cur_forward_context.prefetch_mlp_gate_up_proj + for name in ( + "in_profile_run", + "mmrs_fusion", + "is_first_layer", + "layer_idx", + "prefetch_mlp_gate_up_proj", + "prefetch_mlp_down_proj", + "model_instance", + "is_draft_model", + "is_draft_model_prefill", + "draft_attn_metadatas", + "max_tokens_across_pcp", + "sinks", + "input_ids", + "eplb_heat_collection_status", + ): + _set_ascend_extra( + new_forward_context, + vllm_config, + name, + _get_ascend_extra(cur_forward_context, vllm_config, name), + ) + _set_ascend_extra( + new_forward_context, + vllm_config, + "capturing", + mla_graph_params is not None + or _get_ascend_extra(cur_forward_context, vllm_config, "capturing"), ) - new_forward_context.prefetch_mlp_down_proj = ( - cur_forward_context.prefetch_mlp_down_proj + _set_ascend_extra( + new_forward_context, + vllm_config, + "num_tokens", + num_tokens, ) - new_forward_context.model_instance = cur_forward_context.model_instance - new_forward_context.is_draft_model = cur_forward_context.is_draft_model - new_forward_context.is_draft_model_prefill = ( - cur_forward_context.is_draft_model_prefill + new_forward_context.ubatch_idx = int(ubatch_num) + new_forward_context.num_ubatches = len(ubatch_slices) + flash_comm_v1_enabled = bool( + _get_ascend_extra( + cur_forward_context, + vllm_config, + "flash_comm_v1_enabled", + ) ) - new_forward_context.draft_attn_metadatas = cur_forward_context.draft_attn_metadatas - new_forward_context.max_tokens_across_pcp = ( - cur_forward_context.max_tokens_across_pcp + _set_ascend_extra( + new_forward_context, + vllm_config, + "flash_comm_v1_enabled", + flash_comm_v1_enabled, ) - new_forward_context.sinks = cur_forward_context.sinks - new_forward_context.input_ids = cur_forward_context.input_ids - new_forward_context.eplb_heat_collection_status = ( - cur_forward_context.eplb_heat_collection_status + _set_ascend_extra( + new_forward_context, + vllm_config, + "pad_size", + 0, ) - if new_forward_context.flash_comm_v1_enabled: - new_forward_context.pad_size = ( - tp_world_size - (num_tokens % tp_world_size) - ) % tp_world_size + if flash_comm_v1_enabled: + _set_ascend_extra( + new_forward_context, + vllm_config, + "pad_size", + (tp_world_size - (num_tokens % tp_world_size)) % tp_world_size, + ) if dp_world_size > 1 and dp_metadata is not None: max_tokens_across_dp = dp_metadata.num_tokens_across_dp_cpu.max().item() - if new_forward_context.flash_comm_v1_enabled: + if flash_comm_v1_enabled: padded_length = ( (max_tokens_across_dp + tp_world_size - 1) // tp_world_size * tp_world_size ) - new_forward_context.padded_length = padded_length - new_forward_context.pad_size = padded_length - num_tokens + _set_ascend_extra( + new_forward_context, + vllm_config, + "padded_length", + padded_length, + ) + _set_ascend_extra( + new_forward_context, + vllm_config, + "pad_size", + padded_length - num_tokens, + ) else: max_tokens_across_dp = num_tokens - new_forward_context.max_tokens_across_dp = max_tokens_across_dp + _set_ascend_extra( + new_forward_context, + vllm_config, + "max_tokens_across_dp", + max_tokens_across_dp, + ) - new_forward_context.padded_num_tokens = ( - math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size + padded_num_tokens = math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size + _set_ascend_extra( + new_forward_context, + vllm_config, + "padded_num_tokens", + padded_num_tokens, + ) + cur_mc2_mask = _get_ascend_extra( + cur_forward_context, + vllm_config, + "mc2_mask", ) - cur_mc2_mask = cur_forward_context.mc2_mask if cur_mc2_mask is not None: mc2_mask = torch.zeros( - (new_forward_context.padded_num_tokens,), + (padded_num_tokens,), dtype=cur_mc2_mask.dtype, device=cur_mc2_mask.device, ) mc2_mask[:num_tokens] = True mc2_mask[num_tokens:] = False - new_forward_context.mc2_mask = mc2_mask + _set_ascend_extra( + new_forward_context, + vllm_config, + "mc2_mask", + mc2_mask, + ) new_forward_context.dbo_enabled = True return new_forward_context diff --git a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py new file mode 100644 index 00000000..8ee0b136 --- /dev/null +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Ascend eager DBO runner for vLLM 0.26 ModelRunnerV2.""" + +from __future__ import annotations + +import threading +from contextlib import ExitStack +from dataclasses import dataclass, replace +from typing import Any + +import torch +import torch_npu # noqa: F401 +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.forward_context import ( + DPMetadata, + ForwardContext, + override_forward_context, +) +from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.ubatch_utils import UBatchSlices +from vllm.v1.worker.utils import AttentionGroup + +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( + prepare_attn_for_ubatch, + slice_input_batch, + slice_model_inputs, +) +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo.runtime import ( + merge_ubatch_outputs, +) +from afd_plugin.v1.worker.npu.forward_context import create_ascend_forward_context +from afd_plugin.v1.worker.npu.npu_ubatch_wrapper import _all_gather_ubatch_output +from afd_plugin.v1.worker.npu.ubatching import make_ubatch_contexts + +AFD_NPU_MRV2_NUM_UBATCHES = 2 + + +@dataclass +class AFDAscendUBatchState: + """Prepared inputs and forward contexts for one eager DBO step.""" + + slices: UBatchSlices + forward_contexts: list[ForwardContext] + + +class AFDAscendUBatchRunnerV2: + """Run exactly two ModelRunnerV2 microbatches on Ascend.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + model_state: ModelState, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + max_num_reqs: int, + ) -> None: + self.vllm_config = vllm_config + self.parallel_config = vllm_config.parallel_config + self.num_ubatches = int(self.parallel_config.num_ubatches) + if self.num_ubatches != AFD_NPU_MRV2_NUM_UBATCHES: + raise RuntimeError("AFD NPU ModelRunnerV2 requires exactly two ubatches") + self.device = device + self.model_state = model_state + self.attn_groups = attn_groups + self.kv_cache_config = kv_cache_config + self.ready_barrier = threading.Barrier(self.num_ubatches + 1) + self.query_start_loc_buffers = [ + torch.zeros(max_num_reqs + 2, dtype=torch.int32, device=device) + for _ in range(self.num_ubatches) + ] + self.seq_lens_buffers = [ + torch.zeros(max_num_reqs, dtype=torch.int32, device=device) + for _ in range(self.num_ubatches) + ] + + def prepare( + self, + input_batch: InputBatch, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + slices: UBatchSlices, + parent_context: ForwardContext, + ) -> AFDAscendUBatchState: + forward_contexts: list[ForwardContext] = [] + dp_size = int(self.parallel_config.data_parallel_size) + for stage_index, stage in enumerate(slices): + child_batch = slice_input_batch( + input_batch, + stage, + self.query_start_loc_buffers[stage_index], + self.seq_lens_buffers[stage_index], + ) + child_batch = self._with_ascend_fields(input_batch, child_batch, stage) + stage_slot_mappings = slot_mappings[:, stage.token_slice] + stage_block_tables = tuple( + block_table[stage.request_slice] for block_table in block_tables + ) + attn_metadata = prepare_attn_for_ubatch( + self.model_state, + child_batch, + stage_block_tables, + stage_slot_mappings, + self.attn_groups, + self.kv_cache_config, + stage_index, + ) + stage_tokens = int(stage.num_tokens) + counts = torch.full( + (dp_size,), stage_tokens, dtype=torch.int32, device="cpu" + ) + dp_metadata = DPMetadata.make( + self.parallel_config, + stage_tokens, + counts, + ) + context = create_ascend_forward_context( + parent_context, + attn_metadata, + self.vllm_config, + slices, + ubatch_num=stage_index, + dp_metadata=dp_metadata, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + skip_compiled=parent_context.skip_compiled, + ) + context.slot_mapping = build_slot_mappings_by_layer( + stage_slot_mappings, + self.kv_cache_config, + ) + forward_contexts.append(context) + + return AFDAscendUBatchState( + slices=slices, + forward_contexts=forward_contexts, + ) + + @staticmethod + def _with_ascend_fields( + parent_batch: InputBatch, + child_batch: InputBatch, + stage, + ) -> InputBatch: + req_start = int(stage.request_slice.start) + req_stop = int(stage.request_slice.stop) + tok_stop = int(stage.token_slice.stop) + seq_lens_np = parent_batch.seq_lens_np[req_start:req_stop].copy() + truncated = max(0, int(parent_batch.query_start_loc_np[req_stop]) - tok_stop) + if truncated: + seq_lens_np[-1] -= truncated + return replace( + child_batch, + seq_lens_np=seq_lens_np, + attn_state=parent_batch.attn_state, + ) + + def run( + self, + model: Any, + model_inputs: dict[str, Any], + state: AFDAscendUBatchState, + ) -> Any: + compute_stream = torch.npu.current_stream() + contexts = make_ubatch_contexts( + self.num_ubatches, + compute_stream, + state.forward_contexts, + self.ready_barrier, + ) + outputs: dict[int, Any] = {} + errors: dict[int, BaseException] = {} + + @torch.inference_mode() + def run_stage(context, inputs: dict[str, Any]) -> None: + try: + torch.npu.set_device(self.device) + with context: + outputs[context.id] = model(**inputs) + except BaseException as error: # noqa: BLE001 + errors[context.id] = error + + stack = ExitStack() + stack.enter_context(override_forward_context(None)) + threads = [] + for context, stage in zip(contexts, state.slices, strict=True): + thread = threading.Thread( + target=run_stage, + args=(context, slice_model_inputs(model_inputs, stage.token_slice)), + ) + threads.append(thread) + thread.start() + self.ready_barrier.wait() + try: + contexts[0].cpu_wait_event.set() + for thread in threads: + thread.join() + finally: + stack.close() + + if errors: + failed_stage = min(errors) + raise RuntimeError( + f"AFD NPU microbatch {failed_stage} failed", + ) from errors[failed_stage] + ordered_outputs = [outputs[index] for index in range(self.num_ubatches)] + if state.forward_contexts[0].additional_kwargs["flash_comm_v1_enabled"]: + ordered_outputs = [ + _all_gather_ubatch_output( + output, + context.additional_kwargs["pad_size"], + ) + for output, context in zip( + ordered_outputs, + state.forward_contexts, + strict=True, + ) + ] + return merge_ubatch_outputs(ordered_outputs) + + +__all__ = [ + "AFD_NPU_MRV2_NUM_UBATCHES", + "AFDAscendUBatchRunnerV2", + "AFDAscendUBatchState", +] diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index 9d5f6642..fb522cdc 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -147,8 +147,27 @@ def validate_npu_model_runner_v2_config( or vllm_config.compilation_config.pass_config.enable_sp ): raise RuntimeError("AFD ModelRunnerV2 requires static expert parallelism") - if parallel.enable_dbo or parallel.use_ubatching: - raise RuntimeError("AFD NPU ModelRunnerV2 does not support DBO or ubatching") + dbo_enabled = bool(parallel.enable_dbo or parallel.use_ubatching) + if dbo_enabled: + if parallel.data_parallel_size <= 1 or int(parallel.num_ubatches) != 2: + raise RuntimeError( + "AFD NPU ModelRunnerV2 eager DBO requires DP > 1 and exactly " + "two ubatches", + ) + if not vllm_config.model_config.enforce_eager: + raise RuntimeError( + "AFD NPU ModelRunnerV2 DBO currently supports eager execution only", + ) + if ( + vllm_config.speculative_config is not None + or vllm_config.lora_config is not None + or vllm_config.model_config.is_multimodal_model + or vllm_config.model_config.is_encoder_decoder + ): + raise RuntimeError( + "AFD NPU ModelRunnerV2 DBO does not support speculative decode, " + "LoRA, multimodal, or encoder models", + ) from afd_plugin.model_executor.models.model_utils import ( has_afd_model_registration, diff --git a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py new file mode 100644 index 00000000..8dcb4ab2 --- /dev/null +++ b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + +from dataclasses import replace +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch", exc_type=ImportError) +pytest.importorskip("vllm", exc_type=ImportError) + +from vllm.config import CUDAGraphMode # noqa: E402 +from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers # noqa: E402 + +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( # noqa: E402 + AFDBatchExecutionDescriptor, + create_ubatch_slices, + dispatch_afd_dbo_and_sync_dp, + prepare_attn_for_ubatch, + runtime, # noqa: E402 + slice_input_batch, +) + + +def _parallel_config(*, decode_threshold=4, prefill_threshold=8): + return SimpleNamespace( + use_ubatching=True, + num_ubatches=2, + dbo_decode_token_threshold=decode_threshold, + dbo_prefill_token_threshold=prefill_threshold, + ) + + +def test_dispatch_selects_two_ubatches_and_uniform_padding(monkeypatch): + def all_reduce(tensor, group): + assert group == "cpu-group" + tensor[0] = torch.tensor([8, 12], dtype=torch.int32) + tensor[1] = torch.tensor([1, 1], dtype=torch.int32) + tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + + monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + runtime, + "get_dp_group", + lambda: SimpleNamespace(cpu_group="cpu-group"), + ) + + descriptor, counts = dispatch_afd_dbo_and_sync_dp( + num_reqs=8, + num_tokens=8, + uniform_token_count=1, + dp_size=2, + dp_rank=0, + parallel_config=_parallel_config(), + decode_query_len=1, + allow_ubatching=True, + ) + + assert isinstance(descriptor, AFDBatchExecutionDescriptor) + assert descriptor.cg_mode is CUDAGraphMode.NONE + assert descriptor.num_tokens == 12 + assert descriptor.num_ubatches == 2 + assert counts.tolist() == [12, 12] + + +def test_dispatch_uses_single_batch_when_one_rank_is_below_threshold(monkeypatch): + def all_reduce(tensor, group): + del group + tensor[0] = torch.tensor([3, 12], dtype=torch.int32) + tensor[1] = torch.tensor([1, 1], dtype=torch.int32) + tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + + monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + runtime, + "get_dp_group", + lambda: SimpleNamespace(cpu_group=None), + ) + + descriptor, counts = dispatch_afd_dbo_and_sync_dp( + num_reqs=3, + num_tokens=3, + uniform_token_count=1, + dp_size=2, + dp_rank=0, + parallel_config=_parallel_config(), + decode_query_len=1, + allow_ubatching=True, + ) + + assert not isinstance(descriptor, AFDBatchExecutionDescriptor) + assert descriptor.num_tokens == 3 + assert counts.tolist() == [3, 12] + + +def test_slice_input_batch_keeps_all_padding_trailing_stage_well_formed(): + buffers = InputBuffers(2, 8, torch.device("cpu")) + batch = InputBatch.make_dummy(1, 3, buffers) + buffers.is_padding[:3].fill_(False) + buffers.is_padding[3:8].fill_(True) + batch = replace( + batch, + num_tokens_after_padding=8, + input_ids=buffers.input_ids[:8], + positions=buffers.positions[:8], + is_padding=buffers.is_padding[:8], + ) + stages = create_ubatch_slices(batch, 2) + + trailing = slice_input_batch( + batch, + stages[1], + torch.zeros(3, dtype=torch.int32), + torch.zeros(2, dtype=torch.int32), + ) + + assert [stage.num_tokens for stage in stages] == [4, 4] + assert trailing.num_reqs == 1 + assert trailing.num_tokens == 0 + assert trailing.num_tokens_after_padding == 4 + assert trailing.num_scheduled_tokens.tolist() == [0] + assert trailing.is_padding.tolist() == [True, True, True, True] + + +def test_prepare_attn_for_second_ubatch_restores_builder_order(): + first_builder = object() + second_builder = object() + + class Group: + metadata_builders = [first_builder, second_builder] + + group = Group() + + class ModelState: + def prepare_attn(self, *_args): + return group.metadata_builders[0] + + selected = prepare_attn_for_ubatch( + ModelState(), + input_batch=object(), + block_tables=(), + slot_mappings=object(), + attn_groups=[[group]], + kv_cache_config=object(), + ubatch_index=1, + ) + + assert selected is second_builder + assert group.metadata_builders == [first_builder, second_builder] + + +def test_merge_ubatch_outputs_preserves_ascend_auxiliary_structure(): + outputs = [ + (torch.tensor([[1]]), [torch.tensor([[2]]), torch.tensor([[3]])]), + (torch.tensor([[4]]), [torch.tensor([[5]]), torch.tensor([[6]])]), + ] + + hidden_states, auxiliary = runtime.merge_ubatch_outputs(outputs) + + assert hidden_states.tolist() == [[1], [4]] + assert auxiliary[0].tolist() == [[2], [5]] + assert auxiliary[1].tolist() == [[3], [6]] diff --git a/tests/unit/compat/patches/test_config_validation.py b/tests/unit/compat/patches/test_config_validation.py index 51777de0..b18335d9 100644 --- a/tests/unit/compat/patches/test_config_validation.py +++ b/tests/unit/compat/patches/test_config_validation.py @@ -58,6 +58,7 @@ def create_engine_config(self, usage_context=None, headless=False): }, "native all2all backend assertion" cfg = VllmConfig() cfg.additional_config = self.additional_config + cfg.use_v2_model_runner = self.use_v2_model_runner cfg.parallel_config = SimpleNamespace( use_ubatching=self.enable_dbo or self.ubatch_size > 1, all2all_backend=self.all2all_backend, @@ -93,6 +94,7 @@ def _engine_args(*, active, role="attention", worker_cls="auto"): args.ubatch_size = 0 args.all2all_backend = "allgather_reducescatter" args.worker_cls = worker_cls + args.use_v2_model_runner = False return args @@ -184,6 +186,7 @@ def create_engine_config(engine_args, usage_context=None, headless=False): del usage_context, headless config = config_module.VllmConfig() config.additional_config = engine_args.additional_config + config.use_v2_model_runner = engine_args.use_v2_model_runner config.parallel_config = FakeParallelConfig( enable_dbo=engine_args.enable_dbo, ubatch_size=engine_args.ubatch_size, @@ -403,6 +406,24 @@ def test_config_validation_preserves_npu_dbo_off_behavior(monkeypatch): assert cfg.parallel_config.worker_cls == NPU_FFN_WORKER_FQCN +def test_config_validation_relaxes_only_afd_npu_v2_dbo_gate(monkeypatch): + arg_utils_module, _npu_platform, events = _install_fake_npu_config(monkeypatch) + _load_patch_module() + args = _engine_args(active=True) + args.additional_config["afd"]["connector"] = "CAMP2pAFDConnector" + args.use_v2_model_runner = True + args.ubatch_size = 2 + args.enable_sp = False + args.fail_update = False + + cfg = arg_utils_module.EngineArgs.create_engine_config(args) + + assert ("fix_incompatible_config", False, 0) in events + assert cfg.parallel_config.enable_dbo is True + assert cfg.parallel_config.ubatch_size == 2 + assert cfg.parallel_config.use_ubatching is True + + @pytest.mark.parametrize( ( "role", diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index eed48913..0d044ddd 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -530,6 +530,65 @@ def test_npu_v2_validator_rejects_non_full_acl_graph(cudagraph_mode): ) +def test_npu_v2_validator_allows_eager_dbo_dp2(): + config = _v2_config( + num_attention_ranks=2, + num_ffn_ranks=2, + data_parallel_size=2, + ) + config.additional_config["afd"]["connector"] = "CAMP2pAFDConnector" + config.parallel_config.enable_dbo = True + config.parallel_config.use_ubatching = True + config.parallel_config.num_ubatches = 2 + + validate_npu_model_runner_v2_config( + config, + expected_role="attention", + device_type="npu", + ) + + +@pytest.mark.parametrize( + ("mutation", "message"), + [ + ( + lambda c: ( + setattr(c.parallel_config, "data_parallel_size", 1), + c.additional_config["afd"].update( + num_attention_ranks=1, + num_ffn_ranks=1, + ), + ), + "DP > 1", + ), + (lambda c: setattr(c.parallel_config, "num_ubatches", 4), "two ubatches"), + (lambda c: setattr(c.model_config, "enforce_eager", False), "eager"), + ( + lambda c: setattr(c, "speculative_config", SimpleNamespace()), + "speculative decode", + ), + ], +) +def test_npu_v2_validator_rejects_unsupported_dbo(mutation, message): + config = _v2_config( + num_attention_ranks=2, + num_ffn_ranks=2, + data_parallel_size=2, + ) + config.additional_config["afd"]["connector"] = "CAMP2pAFDConnector" + config.parallel_config.enable_dbo = True + config.parallel_config.use_ubatching = True + config.parallel_config.num_ubatches = 2 + mutation(config) + + with pytest.raises(RuntimeError, match=message): + validate_npu_model_runner_v2_config( + config, + expected_role="attention", + device_type="npu", + ) + + @pytest.mark.parametrize( ("mutation", "message"), [ diff --git a/tests/unit/v1/worker/test_npu_mla_graph.py b/tests/unit/v1/worker/test_npu_mla_graph.py index 07f88561..3eeaa43a 100644 --- a/tests/unit/v1/worker/test_npu_mla_graph.py +++ b/tests/unit/v1/worker/test_npu_mla_graph.py @@ -557,6 +557,7 @@ def test_child_forward_context_installs_mla_capture_registry(monkeypatch): attn_metadata=None, vllm_config=SimpleNamespace( compilation_config=SimpleNamespace(static_forward_context={}), + use_v2_model_runner=False, ), ubatch_slices=_two_slices(4, 4), ubatch_num=1, diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 3c4ecc4a..ebd8b4e2 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -1337,6 +1337,7 @@ def test_npu_create_ascend_forward_context_marks_current_ubatch(monkeypatch): ] vllm_config = SimpleNamespace( compilation_config=SimpleNamespace(static_forward_context={}), + use_v2_model_runner=False, ) new_forward_context = forward_context_module.create_ascend_forward_context( @@ -1354,6 +1355,63 @@ def test_npu_create_ascend_forward_context_marks_current_ubatch(monkeypatch): assert child_metadata.stage_idx == 1 +def test_npu_create_ascend_forward_context_uses_v2_extra_kwargs(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}", + ) + parent = SimpleNamespace( + additional_kwargs={ + "moe_comm_type": "alltoall", + "in_profile_run": False, + "capturing": False, + "mmrs_fusion": False, + "flash_comm_v1_enabled": False, + "is_draft_model": False, + "is_draft_model_prefill": False, + "sinks": False, + "mc2_mask": None, + }, + all_moe_layers={}, + is_padding=None, + ) + slices = [ + SimpleNamespace(token_slice=slice(0, 4), num_tokens=4), + SimpleNamespace(token_slice=slice(4, 7), num_tokens=3), + ] + config = SimpleNamespace( + compilation_config=SimpleNamespace(static_forward_context={}), + use_v2_model_runner=True, + ) + + child = forward_context_module.create_ascend_forward_context( + parent, + attn_metadata=None, + vllm_config=config, + ubatch_slices=slices, + ubatch_num=1, + ) + + assert child.ubatch_idx == 1 + assert child.additional_kwargs["moe_comm_type"] == "alltoall" + assert child.additional_kwargs["moe_comm_method"] == "method:alltoall" + assert child.additional_kwargs["num_tokens"] == 3 + + def test_npu_ffn_runner_executes_eager_ffn_step(monkeypatch): _patch_ffn_forward_context(monkeypatch) runner = _new_ffn_runner() From 2d4665880af718a02ced8e89c1330b038ce87f05 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Tue, 25 Aug 2026 17:42:32 +0800 Subject: [PATCH 2/7] feat(npu): backport full graph DBO for ModelRunnerV2 Signed-off-by: lirx-pd <616517220@qq.com> --- .../backports/vllm_v026_mrv2_dbo/__init__.py | 2 +- .../backports/vllm_v026_mrv2_dbo/execute.py | 33 +- .../backports/vllm_v026_mrv2_dbo/runtime.py | 106 +++-- .../compat/patches/npu/model_runner_v2_dbo.py | 74 ++++ .../v1/worker/npu/aclgraph_manager_v2.py | 391 ++++++++++++++++++ .../worker/npu/attention_model_runner_v2.py | 36 +- afd_plugin/v1/worker/npu/ffn_model_runner.py | 9 +- afd_plugin/v1/worker/npu/ubatch_runner_v2.py | 118 ++++-- afd_plugin/validation.py | 10 +- .../backports/test_vllm_v026_mrv2_dbo.py | 59 ++- tests/unit/v1/worker/test_model_runner_v2.py | 27 +- tests/unit/v1/worker/test_npu_runtime.py | 23 ++ 12 files changed, 786 insertions(+), 102 deletions(-) create mode 100644 afd_plugin/compat/patches/npu/model_runner_v2_dbo.py create mode 100644 afd_plugin/v1/worker/npu/aclgraph_manager_v2.py diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py index 873688ab..64b75874 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Temporary vLLM 0.26 ModelRunnerV2 eager DBO backport.""" +"""Temporary vLLM 0.26 ModelRunnerV2 DBO backport.""" from .runtime import ( AFDBatchExecutionDescriptor, diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py index b5715c0d..885eb940 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""vLLM 0.26 ModelRunnerV2 execute path with the eager DBO seams added.""" +"""vLLM 0.26 ModelRunnerV2 execute path with DBO dispatch and replay seams.""" from __future__ import annotations @@ -40,7 +40,7 @@ def execute_model_v026_eager_dbo( skip_attn_for_dummy_run: bool = False, is_profile: bool = False, ) -> ModelRunnerOutput | IntermediateTensors | None: - """Execute the supported plain-decoder subset with eager DBO.""" + """Execute the supported plain-decoder subset with eager or FULL DBO.""" if not dummy_run: runner.update_pp_decode_requests() @@ -69,6 +69,8 @@ def execute_model_v026_eager_dbo( parallel_config=runner.parallel_config, decode_query_len=runner.decode_query_len, allow_ubatching=not skip_attn_for_dummy_run, + cudagraph_manager=runner.cudagraph_manager, + need_eager=is_profile or skip_attn_for_dummy_run, ) if batch_desc.num_tokens == 0: return runner.kv_connector.no_forward(scheduler_output) @@ -80,9 +82,7 @@ def execute_model_v026_eager_dbo( ) if not dummy_run: runner.input_buffers.is_padding[:num_tokens].fill_(False) - runner.input_buffers.is_padding[ - num_tokens : batch_desc.num_tokens - ].fill_(True) + runner.input_buffers.is_padding[num_tokens : batch_desc.num_tokens].fill_(True) input_batch = runner.prepare_inputs(scheduler_output, batch_desc) block_tables, slot_mappings = runner.prepare_attn(input_batch) runner.model_state.preprocess_state( @@ -119,7 +119,7 @@ def execute_model_v026_eager_dbo( ) attn_metadata = runner.model_state.prepare_attn( input_batch, - CUDAGraphMode.NONE, + batch_desc.cg_mode, block_tables, slot_mappings, runner.attn_groups, @@ -139,7 +139,20 @@ def execute_model_v026_eager_dbo( ubatch_slices, ) - if ubatch_slices is not None: + if ubatch_slices is not None and batch_desc.cg_mode == CUDAGraphMode.FULL: + assert isinstance(batch_desc, AFDBatchExecutionDescriptor) + ubatch_state = runner.ubatch_runner.prepare( + input_batch, + block_tables, + slot_mappings, + ubatch_slices, + None, + cg_mode=CUDAGraphMode.FULL, + ) + runner.cudagraph_manager.stage_replay(batch_desc, ubatch_state) + runner.kv_connector.pre_forward(scheduler_output) + model_output = runner.cudagraph_manager.run_fullgraph(batch_desc) + elif ubatch_slices is not None: batch_descriptor = BatchDescriptor( num_tokens=input_batch.num_tokens_after_padding, has_lora=False, @@ -178,6 +191,10 @@ def execute_model_v026_eager_dbo( model_inputs, ubatch_state, ) + elif batch_desc.cg_mode == CUDAGraphMode.FULL: + assert runner.cudagraph_manager is not None + runner.kv_connector.pre_forward(scheduler_output) + model_output = runner.cudagraph_manager.run_fullgraph(batch_desc) else: batch_descriptor = BatchDescriptor( num_tokens=input_batch.num_tokens_after_padding, @@ -188,7 +205,7 @@ def execute_model_v026_eager_dbo( attn_metadata, runner.vllm_config, num_tokens=input_batch.num_tokens_after_padding, - cudagraph_runtime_mode=CUDAGraphMode.NONE, + cudagraph_runtime_mode=batch_desc.cg_mode, num_tokens_across_dp=num_tokens_across_dp, batch_descriptor=batch_descriptor, slot_mapping=slot_mappings_by_layer, diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py index c66896c9..c66a428b 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Runtime pieces missing from vLLM 0.26 ModelRunnerV2 eager DBO. +"""Runtime pieces missing from vLLM 0.26 ModelRunnerV2 DBO. The behavior is adapted from ``specture724/vllm`` branch ``feat/v2/dbo-fullcg`` at ``626fee7831``. The target ABI is vLLM @@ -62,28 +62,51 @@ def dispatch_afd_dbo_and_sync_dp( parallel_config: ParallelConfig, decode_query_len: int, allow_ubatching: bool, + cudagraph_manager: Any | None = None, + need_eager: bool = False, ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: - """Select eager DBO consistently across all DP ranks. + """Select DBO and FULL graph execution consistently across DP ranks. The minimum rank token count controls the threshold, while all ranks run the maximum token count when DBO is selected. This matches the final upstream inference rule and avoids a separate per-rank DBO vote. """ - if dp_size == 1: - return ( - BatchExecutionDescriptor( + def dispatch( + tokens: int, + ubatches: int, + uniform_tokens: int | None = uniform_token_count, + ) -> BatchExecutionDescriptor: + if need_eager or cudagraph_manager is None: + if ubatches > 1: + return AFDBatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=tokens, + num_reqs=num_reqs, + num_ubatches=ubatches, + ) + return BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, - num_tokens=num_tokens, + num_tokens=tokens, num_reqs=num_reqs, - ), - None, + ) + return cudagraph_manager.dispatch( + num_reqs, + tokens, + uniform_tokens, + num_active_loras=0, + num_ubatches=ubatches, ) - tensor = torch.zeros(3, dp_size, dtype=torch.int32, device="cpu") + if dp_size == 1: + return dispatch(num_tokens, 1), None + + desired_single = dispatch(num_tokens, 1) + tensor = torch.zeros(4, dp_size, dtype=torch.int32, device="cpu") tensor[0][dp_rank] = num_tokens tensor[1][dp_rank] = uniform_token_count or 0 tensor[2][dp_rank] = int(allow_ubatching) + tensor[3][dp_rank] = int(desired_single.cg_mode.value) dist.all_reduce(tensor, group=get_dp_group().cpu_group) num_tokens_across_dp = tensor[0] @@ -97,13 +120,16 @@ def dispatch_afd_dbo_and_sync_dp( None, ) - uniform_decode = bool( - torch.all(tensor[1] == int(decode_query_len)).item() - ) + uniform_decode = bool(torch.all(tensor[1] == int(decode_query_len)).item()) + synced_uniform_token_count: int | None = int(tensor[1][0].item()) + if ( + synced_uniform_token_count == 0 + or not torch.all(tensor[1] == synced_uniform_token_count).item() + ): + synced_uniform_token_count = None should_ubatch = ( bool(torch.all(tensor[2] == 1).item()) - and int(num_tokens_across_dp.max().item()) - >= int(parallel_config.num_ubatches) + and int(num_tokens_across_dp.max().item()) >= int(parallel_config.num_ubatches) and check_ubatch_thresholds( parallel_config, int(num_tokens_across_dp.min().item()), @@ -111,26 +137,28 @@ def dispatch_afd_dbo_and_sync_dp( ) ) if not should_ubatch: - return ( - BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=num_tokens, - num_reqs=num_reqs, - ), - num_tokens_across_dp, - ) + synced_mode = CUDAGraphMode(int(tensor[3].min().item())) + if synced_mode == CUDAGraphMode.NONE: + return ( + BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + ), + num_tokens_across_dp, + ) + padded_tokens = int(num_tokens_across_dp.max().item()) + synced = dispatch(padded_tokens, 1, synced_uniform_token_count) + num_tokens_across_dp.fill_(synced.num_tokens) + return synced, num_tokens_across_dp padded_tokens = int(num_tokens_across_dp.max().item()) num_tokens_across_dp.fill_(padded_tokens) - return ( - AFDBatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=padded_tokens, - num_reqs=num_reqs, - num_ubatches=int(parallel_config.num_ubatches), - ), - num_tokens_across_dp, - ) + return dispatch( + padded_tokens, + int(parallel_config.num_ubatches), + synced_uniform_token_count, + ), num_tokens_across_dp def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: @@ -191,9 +219,7 @@ def slice_input_batch( seq_lens = seq_lens_buffer[:num_reqs_after_padding] seq_lens.copy_(input_batch.seq_lens[req_start:req_stop]) last = num_reqs_after_padding - 1 - seq_lens[last] -= ( - input_batch.query_start_loc[req_stop] - tok_stop - ).clamp_(min=0) + seq_lens[last] -= (input_batch.query_start_loc[req_stop] - tok_stop).clamp_(min=0) seq_lens_cpu_upper_bound = input_batch.seq_lens_cpu_upper_bound[ req_start:req_stop @@ -221,9 +247,7 @@ def slice_input_batch( seq_lens=seq_lens, seq_lens_cpu_upper_bound=seq_lens_cpu_upper_bound, dcp_local_seq_lens=dcp_local_seq_lens, - num_computed_tokens_np=input_batch.num_computed_tokens_np[ - req_start:req_stop - ], + num_computed_tokens_np=input_batch.num_computed_tokens_np[req_start:req_stop], prefill_len_np=input_batch.prefill_len_np[req_start:req_stop], num_computed_prefill_tokens_np=input_batch.num_computed_prefill_tokens_np[ req_start:req_stop @@ -338,17 +362,20 @@ def prepare_attn_for_ubatch( attn_groups: list[list[AttentionGroup]], kv_cache_config: KVCacheConfig, ubatch_index: int, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + for_capture: bool = False, ) -> dict[str, Any]: """Select one of the two builders without changing upstream signatures.""" if ubatch_index == 0: return model_state.prepare_attn( input_batch, - CUDAGraphMode.NONE, + cg_mode, block_tables, slot_mappings, attn_groups, kv_cache_config, + for_capture=for_capture, ) swapped: list[AttentionGroup] = [] @@ -362,11 +389,12 @@ def prepare_attn_for_ubatch( try: return model_state.prepare_attn( input_batch, - CUDAGraphMode.NONE, + cg_mode, block_tables, slot_mappings, attn_groups, kv_cache_config, + for_capture=for_capture, ) finally: for group in reversed(swapped): diff --git a/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py b/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py new file mode 100644 index 00000000..8b7feb15 --- /dev/null +++ b/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py @@ -0,0 +1,74 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Scoped graph-manager injection for the temporary MRV2 DBO backport.""" + +from __future__ import annotations + +from collections.abc import Iterator +from contextlib import contextmanager + +import torch +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.v1.worker.gpu import model_runner as vllm_model_runner +from vllm_ascend.worker.v2 import model_runner as ascend_model_runner + +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( + share_metadata_builder_workspaces, +) +from afd_plugin.v1.worker.npu.aclgraph_manager_v2 import ( + AFDModelAclGraphManagerV2, +) +from afd_plugin.v1.worker.npu.ubatch_runner_v2 import AFDAscendUBatchRunnerV2 + + +@contextmanager +def use_afd_mrv2_dbo_graph_manager(model_runner) -> Iterator[None]: + """Replace only Ascend's initialization-scoped manager wrapper.""" + + original_wrapper = ascend_model_runner.graph_manager_wrapper + + @contextmanager + def graph_manager_wrapper(runner) -> Iterator[None]: + original_manager = vllm_model_runner.ModelCudaGraphManager + + def factory( + vllm_config: VllmConfig, + device: torch.device, + cudagraph_mode: CUDAGraphMode, + decode_query_len: int, + lora_capture_cases: list[int] | None = None, + ) -> AFDModelAclGraphManagerV2: + share_metadata_builder_workspaces(runner.attn_groups) + ubatch_runner = AFDAscendUBatchRunnerV2( + vllm_config, + device, + runner.model_state, + runner.attn_groups, + runner.kv_cache_config, + runner.max_num_reqs, + ) + runner.ubatch_runner = ubatch_runner + return AFDModelAclGraphManagerV2( + vllm_config, + device, + cudagraph_mode, + decode_query_len, + runner, + ubatch_runner, + lora_capture_cases=lora_capture_cases, + ) + + try: + vllm_model_runner.ModelCudaGraphManager = factory + yield + finally: + vllm_model_runner.ModelCudaGraphManager = original_manager + + try: + ascend_model_runner.graph_manager_wrapper = graph_manager_wrapper + yield + finally: + ascend_model_runner.graph_manager_wrapper = original_wrapper + + +__all__ = ["use_afd_mrv2_dbo_graph_manager"] diff --git a/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py new file mode 100644 index 00000000..94eef98d --- /dev/null +++ b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py @@ -0,0 +1,391 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Temporary MRV2 DBO FULL ACL graph manager for vLLM 0.26.""" + +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn +from vllm.config import CUDAGraphMode, VllmConfig +from vllm.distributed.parallel_state import graph_capture +from vllm.forward_context import ( + BatchDescriptor, + get_forward_context, + set_forward_context, +) +from vllm.sequence import IntermediateTensors +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 BatchExecutionDescriptor +from vllm.v1.worker.gpu.input_batch import InputBuffers +from vllm.v1.worker.gpu.model_states.interface import ModelState +from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds +from vllm.v1.worker.utils import AttentionGroup +from vllm_ascend.compilation.acl_graph import ( + get_graph_params, + update_full_graph_params, +) +from vllm_ascend.worker.v2.aclgraph_utils import ModelAclGraphManager, ModelWithContext +from vllm_ascend.worker.v2.input_batch import AscendInputBatch +from vllm_ascend.worker.v2.utils import communicator_switch + +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( + AFDBatchExecutionDescriptor, + create_ubatch_slices, +) +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.ubatch_runner_v2 import ( + AFDAscendUBatchRunnerV2, + AFDAscendUBatchState, +) + + +@dataclass +class _AFDGraphEntry: + graph: Any + output: Any + graph_params: tuple[Any, Any] + workspace: torch.Tensor + + +@dataclass +class _PendingReplay: + descriptor: AFDBatchExecutionDescriptor + state: AFDAscendUBatchState + + +class AFDModelAclGraphManagerV2(ModelAclGraphManager): + """Keep DBO graph descriptors and storage separate from upstream graphs.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + cudagraph_mode: CUDAGraphMode, + decode_query_len: int, + model_runner: Any, + ubatch_runner: AFDAscendUBatchRunnerV2, + lora_capture_cases: list[int] | None = None, + ) -> None: + super().__init__( + vllm_config, + device, + cudagraph_mode, + decode_query_len, + model_runner, + lora_capture_cases=lora_capture_cases, + ) + self.ubatch_runner = ubatch_runner + self._afd_twins = { + desc: AFDBatchExecutionDescriptor( + cg_mode=desc.cg_mode, + num_tokens=desc.num_tokens, + num_reqs=desc.num_reqs, + uniform_token_count=desc.uniform_token_count, + num_active_loras=desc.num_active_loras, + num_ubatches=2, + ) + for desc in self._capture_descs.get(CUDAGraphMode.FULL, []) + if self._needs_ubatch_twin(desc) + } + self._afd_graphs: dict[AFDBatchExecutionDescriptor, _AFDGraphEntry] = {} + self._afd_pending_replay: _PendingReplay | None = None + + def _needs_ubatch_twin(self, desc: BatchExecutionDescriptor) -> bool: + if desc.num_tokens % 2 or desc.num_tokens < 2: + return False + parallel = self.vllm_config.parallel_config + return any( + check_ubatch_thresholds(parallel, desc.num_tokens, uniform_decode=value) + for value in (True, False) + ) + + def dispatch( + self, + num_reqs: int, + num_tokens: int, + uniform_token_count: int | None, + num_active_loras: int, + num_ubatches: int = 1, + ) -> BatchExecutionDescriptor: + base = super().dispatch( + num_reqs, + num_tokens, + uniform_token_count, + num_active_loras, + ) + if num_ubatches == 1: + return base + if num_ubatches != 2: + raise RuntimeError("AFD NPU ModelRunnerV2 requires exactly two ubatches") + twin = self._afd_twins.get(base) + if twin is not None and twin in self._afd_graphs: + return twin + return AFDBatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + num_active_loras=num_active_loras, + num_ubatches=2, + ) + + def stage_replay( + self, + descriptor: AFDBatchExecutionDescriptor, + state: AFDAscendUBatchState, + ) -> None: + if descriptor not in self._afd_graphs: + raise RuntimeError(f"No AFD DBO ACL graph for {descriptor}") + if self._afd_pending_replay is not None: + raise RuntimeError("An AFD DBO ACL graph replay is already pending") + self._afd_pending_replay = _PendingReplay(descriptor, state) + + def clear_afd_graphs(self) -> None: + """Release plugin-owned graphs and any staged replay state.""" + + self._afd_pending_replay = None + self._afd_graphs.clear() + + def capture( + self, + model: nn.Module, + model_state: ModelState, + input_buffers: InputBuffers, + intermediate_tensors: IntermediateTensors | None, + block_tables: BlockTables, + attn_groups: list[list[AttentionGroup]], + kv_cache_config: KVCacheConfig, + has_lora: bool = False, + use_aux_hidden_state_outputs: bool = False, + lora_capture_hook: Callable[[int, int, int], None] | None = None, + progress_bar_desc: str = "Capturing CUDA graphs", + ) -> None: + super().capture( + model, + model_state, + input_buffers, + intermediate_tensors, + block_tables, + attn_groups, + kv_cache_config, + has_lora=has_lora, + use_aux_hidden_state_outputs=use_aux_hidden_state_outputs, + lora_capture_hook=lora_capture_hook, + progress_bar_desc=progress_bar_desc, + ) + if not self._afd_twins: + return + + wrapped_model = ModelWithContext(model) + try: + with graph_capture(device=self.device), communicator_switch(): + for descriptor in self._afd_twins.values(): + self._capture_afd_graph( + descriptor, + wrapped_model, + model_state, + input_buffers, + intermediate_tensors, + block_tables, + ) + except BaseException: + self.clear_afd_graphs() + raise + + def _capture_afd_graph( + self, + descriptor: AFDBatchExecutionDescriptor, + model: nn.Module, + model_state: ModelState, + input_buffers: InputBuffers, + intermediate_tensors: IntermediateTensors | None, + block_tables: BlockTables, + ) -> None: + num_tokens = descriptor.num_tokens + num_reqs = descriptor.num_reqs or min(num_tokens, self.max_num_reqs) + stage_tokens = num_tokens // 2 + aggregate = get_graph_params() + if aggregate is None or aggregate.workspaces.get(num_tokens) is None: + raise RuntimeError( + "MLA DBO FULL graph requires the upstream FIA workspace for " + f"{num_tokens} tokens" + ) + workspace = aggregate.workspaces[num_tokens] + graph_params = ( + new_mla_graph_params(stage_tokens, workspace), + new_mla_graph_params(stage_tokens, workspace), + ) + input_batch = AscendInputBatch.make_dummy( + num_reqs, + num_tokens, + input_buffers, + ) + slices = create_ubatch_slices(input_batch, 2) + model_inputs = { + "input_ids": input_buffers.input_ids[:num_tokens], + "positions": input_buffers.positions[:num_tokens], + **model_state.prepare_dummy_inputs(num_reqs, num_tokens), + } + if not self.is_first_pp_rank: + model_inputs["input_ids"] = None + model_inputs["inputs_embeds"] = None + assert intermediate_tensors is not None + model_inputs["intermediate_tensors"] = intermediate_tensors[:num_tokens] + input_buffers.is_padding.fill_(True) + dummy_tables = block_tables.get_dummy_block_tables(num_reqs) + dummy_slots = block_tables.get_dummy_slot_mappings(num_tokens) + + warmup_state = self._prepare_capture_state( + descriptor, + input_batch, + dummy_tables, + dummy_slots, + slices, + None, + is_warmup=True, + ) + self.ubatch_runner.run(model, model_inputs, warmup_state) + + capture_state = self._prepare_capture_state( + descriptor, + input_batch, + dummy_tables, + dummy_slots, + slices, + graph_params, + is_warmup=False, + ) + finish = self.ubatch_runner.begin_capturable_run( + model, + model_inputs, + capture_state, + for_capture=True, + ) + graph = torch.npu.NPUGraph() + with torch.npu.graph( + graph, + pool=self.pool, + stream=self.ubatch_runner.capture_stream, + ): + output = finish() + self._afd_graphs[descriptor] = _AFDGraphEntry( + graph=graph, + output=output, + graph_params=graph_params, + workspace=workspace, + ) + + def _prepare_capture_state( + self, + descriptor: AFDBatchExecutionDescriptor, + input_batch: AscendInputBatch, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + slices, + graph_params: tuple[Any, Any] | None, + *, + is_warmup: bool, + ) -> AFDAscendUBatchState: + runner = self.model_runner + runner._is_warmup = is_warmup + runner._afd_is_graph_capturing = not is_warmup + runner._afd_pending_metadata = runner.build_afd_metadata( + slices, descriptor.num_tokens + ) + runner.send_dp_metadata( + runner.build_capture_dp_metadata(descriptor.num_tokens), slices + ) + runner._afd_suppress_metadata_send = True + batch_descriptor = BatchDescriptor( + num_tokens=descriptor.num_tokens, + has_lora=False, + num_active_loras=0, + ) + counts = torch.full((self.dp_size,), descriptor.num_tokens, dtype=torch.int32) + with set_forward_context( + None, + self.vllm_config, + num_tokens=descriptor.num_tokens, + cudagraph_runtime_mode=CUDAGraphMode.NONE, + num_tokens_across_dp=counts, + batch_descriptor=batch_descriptor, + ubatch_slices=slices, + is_padding=input_batch.is_padding, + ): + return self.ubatch_runner.prepare( + input_batch, + block_tables, + slot_mappings, + slices, + get_forward_context(), + cg_mode=CUDAGraphMode.FULL, + context_cg_mode=CUDAGraphMode.NONE, + for_capture=True, + mla_graph_params=graph_params, + ) + + def run_fullgraph(self, desc: BatchExecutionDescriptor) -> Any: + if not isinstance(desc, AFDBatchExecutionDescriptor): + return super().run_fullgraph(desc) + pending = self._afd_pending_replay + self._afd_pending_replay = None + if pending is None or pending.descriptor != desc: + raise RuntimeError( + "AFD DBO ACL graph replay state is missing or mismatched" + ) + entry = self._afd_graphs[desc] + state = pending.state + runner = self.model_runner + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_pending_metadata = runner.build_afd_metadata( + state.slices, + sum(state.real_token_counts), + ) + runner._afd_pending_metadata.tokens_unpadded_lens = state.real_token_counts + runner._afd_suppress_metadata_send = True + runner.send_dp_metadata( + runner.build_capture_dp_metadata(desc.num_tokens), state.slices + ) + + assert self.update_stream is not None + self.update_stream.wait_stream(torch.npu.current_stream()) + entry.graph.replay() + stage_tokens = desc.num_tokens // 2 + merged_metadata, merged_params = merge_mla_graph_params( + state.attn_metadata, + entry.graph_params, + stage_tokens, + ) + counts = torch.full((self.dp_size,), desc.num_tokens, dtype=torch.int32) + with set_forward_context( + state.attn_metadata, + self.vllm_config, + num_tokens=desc.num_tokens, + cudagraph_runtime_mode=CUDAGraphMode.FULL, + num_tokens_across_dp=counts, + batch_descriptor=None, + slot_mapping=None, + ): + context = get_forward_context() + with override_mla_graph_params(context, merged_metadata, merged_params): + update_full_graph_params( + self.model_runner.attn_groups[0][0].backend, + self.update_stream, + context, + stage_tokens, + self.vllm_config, + self.model_runner.speculative_config, + ) + return entry.output + + +__all__ = ["AFDModelAclGraphManagerV2"] 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 40307c14..3b69a2be 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -22,6 +22,9 @@ from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm_ascend.worker.v2.model_runner import NPUModelRunner as NPUModelRunnerV2 +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( + AFDBatchExecutionDescriptor, +) from afd_plugin.compat.npu import fail_if_unsupported_npu_afd_features from afd_plugin.compat.npu.profiler import ( create_afd_npu_profiler, @@ -39,6 +42,9 @@ AFDMetadataProviderMixin, _resolve_world_ranks, ) +from afd_plugin.v1.worker.npu.aclgraph_manager_v2 import ( + AFDModelAclGraphManagerV2, +) from afd_plugin.validation import validate_npu_model_runner_v2_config _AFD_FULLGRAPH_HOOK_MARKER = "_afd_fullgraph_replay_hook_active" @@ -81,6 +87,10 @@ def run_fullgraph( self: v2_cudagraph_utils.ModelCudaGraphManager, desc: v2_cudagraph_utils.BatchExecutionDescriptor, ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: + # ### PATCH START: bypass native replay metadata for AFD DBO graphs. + if isinstance(desc, AFDBatchExecutionDescriptor): + return original_run_fullgraph(desc) + # ### PATCH END: bypass native replay metadata for AFD DBO graphs. # ### PATCH START: publish one AFD pre-replay payload. previous_is_graph_replaying = getattr( runner, @@ -221,25 +231,20 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( assert_backport_required, - share_metadata_builder_workspaces, use_two_metadata_builders, ) - from afd_plugin.v1.worker.npu.ubatch_runner_v2 import ( - AFDAscendUBatchRunnerV2, + from afd_plugin.compat.patches.npu.model_runner_v2_dbo import ( + use_afd_mrv2_dbo_graph_manager, ) assert_backport_required() - with use_two_metadata_builders(): + with ( + use_two_metadata_builders(), + use_afd_mrv2_dbo_graph_manager(self), + ): super().initialize_kv_cache(kv_cache_config) - share_metadata_builder_workspaces(self.attn_groups) - self.ubatch_runner = AFDAscendUBatchRunnerV2( - self.vllm_config, - self.device, - self.model_state, - self.attn_groups, - self.kv_cache_config, - self.max_num_reqs, - ) + if self.ubatch_runner is None: + raise RuntimeError("AFD MRV2 DBO graph manager was not initialized") # Patch reason: vLLM v0.26.0 prepares FULL graph inputs before each warmup # and formal capture forward, outside torch.cuda.graph, but does not expose @@ -458,6 +463,11 @@ def shutdown(self) -> None: stop_afd_npu_profiler(self.prof) finally: try: + if isinstance( + self.cudagraph_manager, + AFDModelAclGraphManagerV2, + ): + self.cudagraph_manager.clear_afd_graphs() super().shutdown() finally: self._afd_pending_metadata = None diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 2b6e9281..634db29e 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -404,9 +404,16 @@ def _capture_graphs( graph_key = self._make_graph_key(dp_metadata_list) if graph_key in self._acl_graphs: logger.debug( - "AFD NPU FFN ACL graph capture skipped for existing key=%s", + "AFD NPU FFN replaying existing ACL graph for capture key=%s", graph_key, ) + self.connector.control_plane.update_state_from_dp_metadata( + _make_dp_metadata_payload( + dp_metadata_list, + is_graph_capturing=is_attn_graph_capturing, + ), + ) + self._acl_graphs[graph_key]["graph"].replay() return logger.debug("AFD NPU FFN capturing ACL graph for key=%s", graph_key) diff --git a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py index 8ee0b136..c3b0f177 100644 --- a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Ascend eager DBO runner for vLLM 0.26 ModelRunnerV2.""" +"""Ascend eager and capture-time DBO runner for vLLM 0.26 ModelRunnerV2.""" from __future__ import annotations import threading +from collections.abc import Callable from contextlib import ExitStack from dataclasses import dataclass, replace from typing import Any @@ -41,10 +42,12 @@ @dataclass class AFDAscendUBatchState: - """Prepared inputs and forward contexts for one eager DBO step.""" + """Prepared inputs for one eager or FULL graph DBO step.""" slices: UBatchSlices - forward_contexts: list[ForwardContext] + attn_metadata: list[dict[str, Any]] + forward_contexts: list[ForwardContext] | None + real_token_counts: list[int] class AFDAscendUBatchRunnerV2: @@ -69,6 +72,7 @@ def __init__( self.attn_groups = attn_groups self.kv_cache_config = kv_cache_config self.ready_barrier = threading.Barrier(self.num_ubatches + 1) + self.capture_stream = torch.npu.Stream(device=device) self.query_start_loc_buffers = [ torch.zeros(max_num_reqs + 2, dtype=torch.int32, device=device) for _ in range(self.num_ubatches) @@ -84,9 +88,18 @@ def prepare( block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, slices: UBatchSlices, - parent_context: ForwardContext, + parent_context: ForwardContext | None, + *, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + context_cg_mode: CUDAGraphMode | None = None, + for_capture: bool = False, + mla_graph_params: tuple[Any, Any] | None = None, ) -> AFDAscendUBatchState: + attn_metadata_list: list[dict[str, Any]] = [] forward_contexts: list[ForwardContext] = [] + real_token_counts: list[int] = [] + if context_cg_mode is None: + context_cg_mode = cg_mode dp_size = int(self.parallel_config.data_parallel_size) for stage_index, stage in enumerate(slices): child_batch = slice_input_batch( @@ -108,7 +121,13 @@ def prepare( self.attn_groups, self.kv_cache_config, stage_index, + cg_mode=cg_mode, + for_capture=for_capture, ) + attn_metadata_list.append(attn_metadata) + real_token_counts.append(int(child_batch.num_tokens)) + if parent_context is None: + continue stage_tokens = int(stage.num_tokens) counts = torch.full( (dp_size,), stage_tokens, dtype=torch.int32, device="cpu" @@ -125,8 +144,11 @@ def prepare( slices, ubatch_num=stage_index, dp_metadata=dp_metadata, - cudagraph_runtime_mode=CUDAGraphMode.NONE, + cudagraph_runtime_mode=context_cg_mode, skip_compiled=parent_context.skip_compiled, + mla_graph_params=( + None if mla_graph_params is None else mla_graph_params[stage_index] + ), ) context.slot_mapping = build_slot_mappings_by_layer( stage_slot_mappings, @@ -136,7 +158,9 @@ def prepare( return AFDAscendUBatchState( slices=slices, - forward_contexts=forward_contexts, + attn_metadata=attn_metadata_list, + forward_contexts=forward_contexts or None, + real_token_counts=real_token_counts, ) @staticmethod @@ -164,11 +188,28 @@ def run( model_inputs: dict[str, Any], state: AFDAscendUBatchState, ) -> Any: - compute_stream = torch.npu.current_stream() + return self.begin_capturable_run(model, model_inputs, state)() + + def begin_capturable_run( + self, + model: Any, + model_inputs: dict[str, Any], + state: AFDAscendUBatchState, + *, + for_capture: bool = False, + ) -> Callable[[], Any]: + """Start both stages outside capture and return a one-shot finisher.""" + + forward_contexts = state.forward_contexts + if forward_contexts is None: + raise RuntimeError("uBatch execution requires prepared forward contexts") + compute_stream = ( + self.capture_stream if for_capture else torch.npu.current_stream() + ) contexts = make_ubatch_contexts( self.num_ubatches, compute_stream, - state.forward_contexts, + forward_contexts, self.ready_barrier, ) outputs: dict[int, Any] = {} @@ -194,32 +235,41 @@ def run_stage(context, inputs: dict[str, Any]) -> None: threads.append(thread) thread.start() self.ready_barrier.wait() - try: - contexts[0].cpu_wait_event.set() - for thread in threads: - thread.join() - finally: - stack.close() - - if errors: - failed_stage = min(errors) - raise RuntimeError( - f"AFD NPU microbatch {failed_stage} failed", - ) from errors[failed_stage] - ordered_outputs = [outputs[index] for index in range(self.num_ubatches)] - if state.forward_contexts[0].additional_kwargs["flash_comm_v1_enabled"]: - ordered_outputs = [ - _all_gather_ubatch_output( - output, - context.additional_kwargs["pad_size"], - ) - for output, context in zip( - ordered_outputs, - state.forward_contexts, - strict=True, - ) - ] - return merge_ubatch_outputs(ordered_outputs) + finished = False + + def finish() -> Any: + nonlocal finished + if finished: + raise RuntimeError("uBatch finisher may only be called once") + finished = True + try: + contexts[0].cpu_wait_event.set() + for thread in threads: + thread.join() + finally: + stack.close() + + if errors: + failed_stage = min(errors) + raise RuntimeError( + f"AFD NPU microbatch {failed_stage} failed", + ) from errors[failed_stage] + ordered_outputs = [outputs[index] for index in range(self.num_ubatches)] + if forward_contexts[0].additional_kwargs["flash_comm_v1_enabled"]: + ordered_outputs = [ + _all_gather_ubatch_output( + output, + context.additional_kwargs["pad_size"], + ) + for output, context in zip( + ordered_outputs, + forward_contexts, + strict=True, + ) + ] + return merge_ubatch_outputs(ordered_outputs) + + return finish __all__ = [ diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index fb522cdc..a4cf6911 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -151,12 +151,14 @@ def validate_npu_model_runner_v2_config( if dbo_enabled: if parallel.data_parallel_size <= 1 or int(parallel.num_ubatches) != 2: raise RuntimeError( - "AFD NPU ModelRunnerV2 eager DBO requires DP > 1 and exactly " - "two ubatches", + "AFD NPU ModelRunnerV2 DBO requires DP > 1 and exactly two ubatches", ) - if not vllm_config.model_config.enforce_eager: + if ( + not vllm_config.model_config.enforce_eager + and cudagraph_mode_name(vllm_config) != "FULL_DECODE_ONLY" + ): raise RuntimeError( - "AFD NPU ModelRunnerV2 DBO currently supports eager execution only", + "AFD NPU ModelRunnerV2 DBO ACL graph requires FULL_DECODE_ONLY", ) if ( vllm_config.speculative_config is not None diff --git a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py index 8dcb4ab2..c7939ffa 100644 --- a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py +++ b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py @@ -10,6 +10,9 @@ pytest.importorskip("vllm", exc_type=ImportError) from vllm.config import CUDAGraphMode # noqa: E402 +from vllm.v1.worker.gpu.cudagraph_utils import ( # noqa: E402 + BatchExecutionDescriptor, +) from vllm.v1.worker.gpu.input_batch import InputBatch, InputBuffers # noqa: E402 from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( # noqa: E402 @@ -132,7 +135,7 @@ class Group: group = Group() class ModelState: - def prepare_attn(self, *_args): + def prepare_attn(self, *_args, **_kwargs): return group.metadata_builders[0] selected = prepare_attn_for_ubatch( @@ -149,6 +152,60 @@ def prepare_attn(self, *_args): assert group.metadata_builders == [first_builder, second_builder] +def test_dispatch_requests_captured_two_ubatch_descriptor(monkeypatch): + descriptor = AFDBatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=12, + num_reqs=8, + num_ubatches=2, + ) + dispatch_calls = [] + + class Manager: + def dispatch(self, *args, **kwargs): + dispatch_calls.append((args, kwargs)) + if kwargs["num_ubatches"] == 1: + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=args[1], + num_reqs=args[0], + uniform_token_count=args[2], + ) + return descriptor + + def all_reduce(tensor, group): + del group + tensor[0] = torch.tensor([8, 12], dtype=torch.int32) + tensor[1] = torch.tensor([1, 1], dtype=torch.int32) + tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + + monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + runtime, + "get_dp_group", + lambda: SimpleNamespace(cpu_group=None), + ) + + selected, counts = dispatch_afd_dbo_and_sync_dp( + num_reqs=8, + num_tokens=8, + uniform_token_count=1, + dp_size=2, + dp_rank=0, + parallel_config=_parallel_config(), + decode_query_len=1, + allow_ubatching=True, + cudagraph_manager=Manager(), + ) + + assert selected is descriptor + assert counts.tolist() == [12, 12] + assert dispatch_calls == [ + ((8, 8, 1), {"num_active_loras": 0, "num_ubatches": 1}), + ((8, 12, 1), {"num_active_loras": 0, "num_ubatches": 2}), + ] + + def test_merge_ubatch_outputs_preserves_ascend_auxiliary_structure(): outputs = [ (torch.tensor([[1]]), [torch.tensor([[2]]), torch.tensor([[3]])]), diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 0d044ddd..4920b866 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -548,6 +548,28 @@ def test_npu_v2_validator_allows_eager_dbo_dp2(): ) +def test_npu_v2_validator_allows_full_decode_only_dbo_dp2(): + config = _v2_config( + num_attention_ranks=2, + num_ffn_ranks=2, + data_parallel_size=2, + enforce_eager=False, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + cudagraph_capture_sizes=[TEST_CUDAGRAPH_CAPTURE_SIZE], + max_cudagraph_capture_size=TEST_CUDAGRAPH_CAPTURE_SIZE, + ) + config.additional_config["afd"]["connector"] = "CAMP2pAFDConnector" + config.parallel_config.enable_dbo = True + config.parallel_config.use_ubatching = True + config.parallel_config.num_ubatches = 2 + + validate_npu_model_runner_v2_config( + config, + expected_role="attention", + device_type="npu", + ) + + @pytest.mark.parametrize( ("mutation", "message"), [ @@ -562,7 +584,10 @@ def test_npu_v2_validator_allows_eager_dbo_dp2(): "DP > 1", ), (lambda c: setattr(c.parallel_config, "num_ubatches", 4), "two ubatches"), - (lambda c: setattr(c.model_config, "enforce_eager", False), "eager"), + ( + lambda c: setattr(c.model_config, "enforce_eager", False), + "FULL_DECODE_ONLY", + ), ( lambda c: setattr(c, "speculative_config", SimpleNamespace()), "speculative decode", diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index ebd8b4e2..91d7fd1a 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -1728,6 +1728,29 @@ def test_npu_ffn_runner_skips_replay_when_attention_is_eager(monkeypatch): ] +def test_npu_ffn_capture_event_replays_existing_acl_graph(): + from vllm.config import CUDAGraphMode + + runner = _new_ffn_runner() + runner.vllm_config = _vllm_config(role="ffn") + runner.connector = _FakeFFNConnector() + runner.max_num_tokens = 1 + dp_metadata = {0: _FakeDPMetadata([1])} + graph = _FakeGraph() + runner._acl_graphs = {runner._make_graph_key(dp_metadata): {"graph": graph}} + + runner._capture_graphs( + aclgraph_runtime_mode=CUDAGraphMode.FULL, + dp_metadata_list=dp_metadata, + ) + + assert graph.replay_count == 1 + assert runner.connector.updates[0][1] == { + "is_graph_capturing": True, + "is_warmup": False, + } + + def test_npu_ffn_runner_graph_key_uses_ffn_aggregated_token_counts(): runner = _new_ffn_runner() runner.connector = _FakeFFNConnector(attn_size=8, ffn_size=4) From f0dec2ecf063da4ce0e04d2f3db30674fab23d57 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Thu, 27 Aug 2026 10:43:14 +0800 Subject: [PATCH 3/7] fix the bugs of TypeError Signed-off-by: lirx-pd <616517220@qq.com> --- .../backports/vllm_v026_mrv2_dbo/execute.py | 16 +- .../backports/vllm_v026_mrv2_dbo/runtime.py | 79 +++++-- .../v1/worker/npu/aclgraph_manager_v2.py | 30 +-- afd_plugin/v1/worker/npu/ubatch_runner_v2.py | 25 +- .../backports/test_vllm_v026_mrv2_dbo.py | 218 +++++++++++++++++- 5 files changed, 304 insertions(+), 64 deletions(-) diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py index 885eb940..937107cd 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py @@ -16,10 +16,10 @@ from vllm.sequence import IntermediateTensors from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.cudagraph_utils import get_uniform_token_count -from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_runner import ExecuteModelState from vllm_ascend.worker.v2.input_batch import AscendInputBatch +from . import runtime as dbo_runtime from .runtime import ( AFDBatchExecutionDescriptor, create_ubatch_slices, @@ -31,6 +31,17 @@ from vllm.v1.outputs import ModelRunnerOutput +_EXPECTED_RUNTIME_ABI = 3 +_loaded_runtime_abi = getattr(dbo_runtime, "AFD_MRV2_DBO_RUNTIME_ABI", 1) +if _loaded_runtime_abi != _EXPECTED_RUNTIME_ABI: + raise ImportError( + "AFD ModelRunnerV2 DBO backport modules are out of sync: " + f"execute expects runtime ABI {_EXPECTED_RUNTIME_ABI}, but loaded " + f"ABI {_loaded_runtime_abi}. Reinstall afd-plugin from one checkout " + "and restart every worker process." + ) + + def execute_model_v026_eager_dbo( runner: Any, scheduler_output: SchedulerOutput, @@ -92,8 +103,7 @@ def execute_model_v026_eager_dbo( runner.req_states.num_computed_tokens.gpu, ) else: - dummy_batch_cls = AscendInputBatch if num_ubatches > 1 else InputBatch - input_batch = dummy_batch_cls.make_dummy( + input_batch = AscendInputBatch.make_dummy( batch_desc.num_reqs or num_reqs, batch_desc.num_tokens, runner.input_buffers, diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py index c66a428b..438abdd9 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py @@ -29,10 +29,13 @@ UBatchSlice, UBatchSlices, check_ubatch_thresholds, + is_last_ubatch_empty, maybe_create_ubatch_slices, ) from vllm.v1.worker.utils import AttentionGroup +AFD_MRV2_DBO_RUNTIME_ABI = 3 + @dataclass(frozen=True) class AFDBatchExecutionDescriptor(BatchExecutionDescriptor): @@ -76,27 +79,42 @@ def dispatch( tokens: int, ubatches: int, uniform_tokens: int | None = uniform_token_count, + force_eager: bool = False, ) -> BatchExecutionDescriptor: - if need_eager or cudagraph_manager is None: - if ubatches > 1: - return AFDBatchExecutionDescriptor( - cg_mode=CUDAGraphMode.NONE, - num_tokens=tokens, - num_reqs=num_reqs, - num_ubatches=ubatches, - ) - return BatchExecutionDescriptor( + if force_eager or need_eager or cudagraph_manager is None: + base = BatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, num_tokens=tokens, num_reqs=num_reqs, ) - return cudagraph_manager.dispatch( - num_reqs, - tokens, - uniform_tokens, - num_active_loras=0, - num_ubatches=ubatches, - ) + else: + base = cudagraph_manager.dispatch( + num_reqs, + tokens, + uniform_tokens, + num_active_loras=0, + ) + + if ubatches == 1: + return base + if base.cg_mode == CUDAGraphMode.NONE: + return AFDBatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=tokens, + num_reqs=num_reqs, + uniform_token_count=uniform_tokens, + num_active_loras=base.num_active_loras, + num_ubatches=ubatches, + ) + + dispatch_ubatches = getattr(cudagraph_manager, "dispatch_ubatches", None) + if dispatch_ubatches is None: + manager_type = type(cudagraph_manager).__name__ + raise RuntimeError( + "AFD NPU ModelRunnerV2 FULL graph DBO requires the plugin " + f"graph manager, got {manager_type}" + ) + return dispatch_ubatches(base, ubatches) if dp_size == 1: return dispatch(num_tokens, 1), None @@ -127,6 +145,7 @@ def dispatch( or not torch.all(tensor[1] == synced_uniform_token_count).item() ): synced_uniform_token_count = None + synced_mode = CUDAGraphMode(int(tensor[3].min().item())) should_ubatch = ( bool(torch.all(tensor[2] == 1).item()) and int(num_tokens_across_dp.max().item()) >= int(parallel_config.num_ubatches) @@ -137,7 +156,6 @@ def dispatch( ) ) if not should_ubatch: - synced_mode = CUDAGraphMode(int(tensor[3].min().item())) if synced_mode == CUDAGraphMode.NONE: return ( BatchExecutionDescriptor( @@ -153,12 +171,31 @@ def dispatch( return synced, num_tokens_across_dp padded_tokens = int(num_tokens_across_dp.max().item()) - num_tokens_across_dp.fill_(padded_tokens) - return dispatch( + num_ubatches = int(parallel_config.num_ubatches) + ubatch_desc = dispatch( + padded_tokens, + num_ubatches, + synced_uniform_token_count, + force_eager=synced_mode == CUDAGraphMode.NONE, + ) + # Graph dispatch can round the DP maximum upward. Check the final shape so + # every rank has real work in its last microbatch. + if not is_last_ubatch_empty( + int(num_tokens_across_dp.min().item()), + ubatch_desc.num_tokens, + num_ubatches, + ): + num_tokens_across_dp.fill_(ubatch_desc.num_tokens) + return ubatch_desc, num_tokens_across_dp + + synced = dispatch( padded_tokens, - int(parallel_config.num_ubatches), + 1, synced_uniform_token_count, - ), num_tokens_across_dp + force_eager=synced_mode == CUDAGraphMode.NONE, + ) + num_tokens_across_dp.fill_(synced.num_tokens) + return synced, num_tokens_across_dp def create_ubatch_slices(input_batch: InputBatch, num_ubatches: int) -> UBatchSlices: diff --git a/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py index 94eef98d..6ee5aa2e 100644 --- a/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py +++ b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py @@ -108,22 +108,11 @@ def _needs_ubatch_twin(self, desc: BatchExecutionDescriptor) -> bool: for value in (True, False) ) - def dispatch( + def dispatch_ubatches( self, - num_reqs: int, - num_tokens: int, - uniform_token_count: int | None, - num_active_loras: int, - num_ubatches: int = 1, + base: BatchExecutionDescriptor, + num_ubatches: int, ) -> BatchExecutionDescriptor: - base = super().dispatch( - num_reqs, - num_tokens, - uniform_token_count, - num_active_loras, - ) - if num_ubatches == 1: - return base if num_ubatches != 2: raise RuntimeError("AFD NPU ModelRunnerV2 requires exactly two ubatches") twin = self._afd_twins.get(base) @@ -131,9 +120,10 @@ def dispatch( return twin return AFDBatchExecutionDescriptor( cg_mode=CUDAGraphMode.NONE, - num_tokens=num_tokens, - num_reqs=num_reqs, - num_active_loras=num_active_loras, + num_tokens=base.num_tokens, + num_reqs=base.num_reqs, + uniform_token_count=base.uniform_token_count, + num_active_loras=base.num_active_loras, num_ubatches=2, ) @@ -357,7 +347,11 @@ def run_fullgraph(self, desc: BatchExecutionDescriptor) -> Any: ) assert self.update_stream is not None - self.update_stream.wait_stream(torch.npu.current_stream()) + current_stream = torch.npu.current_stream() + self.update_stream.wait_stream(current_stream) + # This graph bypasses Ascend's ACLGraphWrapper, so preserve its FULL + # replay fence before updating the captured FIA task-group handles. + current_stream.synchronize() entry.graph.replay() stage_tokens = desc.num_tokens // 2 merged_metadata, merged_params = merge_mla_graph_params( diff --git a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py index c3b0f177..3564a181 100644 --- a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -8,7 +8,7 @@ from collections.abc import Callable from contextlib import ExitStack from dataclasses import dataclass, replace -from typing import Any +from typing import Any, cast import torch import torch_npu # noqa: F401 @@ -20,10 +20,10 @@ ) from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer -from vllm.v1.worker.gpu.input_batch import InputBatch from vllm.v1.worker.gpu.model_states.interface import ModelState from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import AttentionGroup +from vllm_ascend.worker.v2.input_batch import AscendInputBatch from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( prepare_attn_for_ubatch, @@ -84,7 +84,7 @@ def __init__( def prepare( self, - input_batch: InputBatch, + input_batch: AscendInputBatch, block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, slices: UBatchSlices, @@ -102,11 +102,14 @@ def prepare( context_cg_mode = cg_mode dp_size = int(self.parallel_config.data_parallel_size) for stage_index, stage in enumerate(slices): - child_batch = slice_input_batch( - input_batch, - stage, - self.query_start_loc_buffers[stage_index], - self.seq_lens_buffers[stage_index], + child_batch = cast( + AscendInputBatch, + slice_input_batch( + input_batch, + stage, + self.query_start_loc_buffers[stage_index], + self.seq_lens_buffers[stage_index], + ), ) child_batch = self._with_ascend_fields(input_batch, child_batch, stage) stage_slot_mappings = slot_mappings[:, stage.token_slice] @@ -165,10 +168,10 @@ def prepare( @staticmethod def _with_ascend_fields( - parent_batch: InputBatch, - child_batch: InputBatch, + parent_batch: AscendInputBatch, + child_batch: AscendInputBatch, stage, - ) -> InputBatch: + ) -> AscendInputBatch: req_start = int(stage.request_slice.start) req_stop = int(stage.request_slice.stop) tok_stop = int(stage.token_slice.stop) diff --git a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py index c7939ffa..c0257334 100644 --- a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py +++ b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +import inspect from dataclasses import replace from types import SimpleNamespace @@ -34,6 +35,13 @@ def _parallel_config(*, decode_threshold=4, prefill_threshold=8): ) +def test_runtime_abi_matches_full_graph_dispatch_signature(): + parameters = inspect.signature(dispatch_afd_dbo_and_sync_dp).parameters + + assert runtime.AFD_MRV2_DBO_RUNTIME_ABI == 3 + assert {"cudagraph_manager", "need_eager"} <= parameters.keys() + + def test_dispatch_selects_two_ubatches_and_uniform_padding(monkeypatch): def all_reduce(tensor, group): assert group == "cpu-group" @@ -163,14 +171,18 @@ def test_dispatch_requests_captured_two_ubatch_descriptor(monkeypatch): class Manager: def dispatch(self, *args, **kwargs): - dispatch_calls.append((args, kwargs)) - if kwargs["num_ubatches"] == 1: - return BatchExecutionDescriptor( - cg_mode=CUDAGraphMode.FULL, - num_tokens=args[1], - num_reqs=args[0], - uniform_token_count=args[2], - ) + dispatch_calls.append(("dispatch", args, kwargs)) + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=args[1], + num_reqs=args[0], + uniform_token_count=args[2], + ) + + def dispatch_ubatches(self, base, num_ubatches): + dispatch_calls.append(("dispatch_ubatches", base, num_ubatches)) + assert base.num_tokens == 12 + assert num_ubatches == 2 return descriptor def all_reduce(tensor, group): @@ -178,6 +190,10 @@ def all_reduce(tensor, group): tensor[0] = torch.tensor([8, 12], dtype=torch.int32) tensor[1] = torch.tensor([1, 1], dtype=torch.int32) tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + tensor[3] = torch.tensor( + [CUDAGraphMode.FULL.value, CUDAGraphMode.FULL.value], + dtype=torch.int32, + ) monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) monkeypatch.setattr( @@ -200,10 +216,190 @@ def all_reduce(tensor, group): assert selected is descriptor assert counts.tolist() == [12, 12] - assert dispatch_calls == [ - ((8, 8, 1), {"num_active_loras": 0, "num_ubatches": 1}), - ((8, 12, 1), {"num_active_loras": 0, "num_ubatches": 2}), + assert dispatch_calls[:2] == [ + ("dispatch", (8, 8, 1), {"num_active_loras": 0}), + ("dispatch", (8, 12, 1), {"num_active_loras": 0}), ] + assert dispatch_calls[2][0] == "dispatch_ubatches" + + +def test_dispatch_accepts_upstream_manager_without_ubatch_keyword(monkeypatch): + dispatch_calls = [] + + class Manager: + def dispatch( + self, + num_reqs, + num_tokens, + uniform_token_count, + num_active_loras, + ): + dispatch_calls.append( + (num_reqs, num_tokens, uniform_token_count, num_active_loras) + ) + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.NONE, + num_tokens=num_tokens, + num_reqs=num_reqs, + uniform_token_count=uniform_token_count, + num_active_loras=num_active_loras, + ) + + def all_reduce(tensor, group): + del group + tensor[0] = torch.tensor([8, 12], dtype=torch.int32) + tensor[1] = torch.tensor([1, 1], dtype=torch.int32) + tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + + monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + runtime, + "get_dp_group", + lambda: SimpleNamespace(cpu_group=None), + ) + + selected, counts = dispatch_afd_dbo_and_sync_dp( + num_reqs=8, + num_tokens=8, + uniform_token_count=1, + dp_size=2, + dp_rank=0, + parallel_config=_parallel_config(), + decode_query_len=1, + allow_ubatching=True, + cudagraph_manager=Manager(), + ) + + assert isinstance(selected, AFDBatchExecutionDescriptor) + assert selected.cg_mode is CUDAGraphMode.NONE + assert selected.num_tokens == 12 + assert selected.num_ubatches == 2 + assert counts.tolist() == [12, 12] + assert dispatch_calls == [(8, 8, 1, 0)] + + +def test_dispatch_avoids_empty_second_ubatch_after_graph_padding(monkeypatch): + dispatch_calls = [] + + class Manager: + def dispatch( + self, + num_reqs, + num_tokens, + uniform_token_count, + num_active_loras, + ): + dispatch_calls.append((num_reqs, num_tokens, uniform_token_count)) + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=8, + num_reqs=8, + uniform_token_count=1, + num_active_loras=num_active_loras, + ) + + def dispatch_ubatches(self, base, num_ubatches): + return AFDBatchExecutionDescriptor( + **vars(base), + num_ubatches=num_ubatches, + ) + + def all_reduce(tensor, group): + del group + tensor[0] = torch.tensor([4, 7], dtype=torch.int32) + tensor[1] = torch.tensor([1, 1], dtype=torch.int32) + tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + tensor[3] = torch.tensor( + [CUDAGraphMode.FULL.value, CUDAGraphMode.FULL.value], + dtype=torch.int32, + ) + + monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + runtime, + "get_dp_group", + lambda: SimpleNamespace(cpu_group=None), + ) + + selected, counts = dispatch_afd_dbo_and_sync_dp( + num_reqs=4, + num_tokens=4, + uniform_token_count=1, + dp_size=2, + dp_rank=0, + parallel_config=_parallel_config(), + decode_query_len=1, + allow_ubatching=True, + cudagraph_manager=Manager(), + ) + + assert not isinstance(selected, AFDBatchExecutionDescriptor) + assert selected.cg_mode is CUDAGraphMode.FULL + assert selected.num_tokens == 8 + assert counts.tolist() == [8, 8] + assert dispatch_calls == [(4, 4, 1), (4, 7, 1), (4, 7, 1)] + + +def test_dispatch_uses_eager_dbo_when_any_rank_misses_graph(monkeypatch): + dispatch_calls = [] + + class Manager: + def dispatch( + self, + num_reqs, + num_tokens, + uniform_token_count, + num_active_loras, + ): + dispatch_calls.append( + (num_reqs, num_tokens, uniform_token_count, num_active_loras) + ) + return BatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=num_tokens, + num_reqs=num_reqs, + uniform_token_count=uniform_token_count, + num_active_loras=num_active_loras, + ) + + def dispatch_ubatches(self, base, num_ubatches): + raise AssertionError((base, num_ubatches)) + + def all_reduce(tensor, group): + del group + tensor[0] = torch.tensor([8, 12], dtype=torch.int32) + tensor[1] = torch.tensor([1, 1], dtype=torch.int32) + tensor[2] = torch.tensor([1, 1], dtype=torch.int32) + tensor[3] = torch.tensor( + [CUDAGraphMode.FULL.value, CUDAGraphMode.NONE.value], + dtype=torch.int32, + ) + + monkeypatch.setattr(runtime.dist, "all_reduce", all_reduce) + monkeypatch.setattr( + runtime, + "get_dp_group", + lambda: SimpleNamespace(cpu_group=None), + ) + + selected, counts = dispatch_afd_dbo_and_sync_dp( + num_reqs=8, + num_tokens=8, + uniform_token_count=1, + dp_size=2, + dp_rank=0, + parallel_config=_parallel_config(), + decode_query_len=1, + allow_ubatching=True, + cudagraph_manager=Manager(), + ) + + assert isinstance(selected, AFDBatchExecutionDescriptor) + assert selected.cg_mode is CUDAGraphMode.NONE + assert selected.num_tokens == 12 + assert selected.num_ubatches == 2 + assert counts.tolist() == [12, 12] + assert dispatch_calls == [(8, 8, 1, 0)] def test_merge_ubatch_outputs_preserves_ascend_auxiliary_structure(): From 550a5ce8d02d5d20fefe4c672eb7c82d5d868d2b Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Thu, 27 Aug 2026 15:57:48 +0800 Subject: [PATCH 4/7] refactor(npu): normalize MRV2 DBO backport patches. Align upstream-derived code with patch markers and remove redundant defensive logic. Signed-off-by: lirx-pd <616517220@qq.com> --- .../backports/vllm_v026_mrv2_dbo/__init__.py | 4 - .../backports/vllm_v026_mrv2_dbo/execute.py | 37 ++- .../backports/vllm_v026_mrv2_dbo/runtime.py | 92 +++--- .../compat/patches/config_validation.py | 2 +- .../compat/patches/npu/model_runner_v2_dbo.py | 46 ++- .../v1/worker/npu/aclgraph_manager_v2.py | 62 +++- .../worker/npu/attention_model_runner_v2.py | 39 ++- afd_plugin/v1/worker/npu/ffn_model_runner.py | 2 + afd_plugin/v1/worker/npu/forward_context.py | 274 ++++++++---------- afd_plugin/v1/worker/npu/ubatch_runner_v2.py | 105 +++++-- afd_plugin/validation.py | 2 + .../backports/test_vllm_v026_mrv2_dbo.py | 3 +- 12 files changed, 396 insertions(+), 272 deletions(-) diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py index 64b75874..c229d11b 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py @@ -4,11 +4,9 @@ from .runtime import ( AFDBatchExecutionDescriptor, - assert_backport_required, create_ubatch_slices, dispatch_afd_dbo_and_sync_dp, prepare_attn_for_ubatch, - share_metadata_builder_workspaces, slice_input_batch, slice_model_inputs, use_two_metadata_builders, @@ -16,11 +14,9 @@ __all__ = [ "AFDBatchExecutionDescriptor", - "assert_backport_required", "create_ubatch_slices", "dispatch_afd_dbo_and_sync_dp", "prepare_attn_for_ubatch", - "share_metadata_builder_workspaces", "slice_input_batch", "slice_model_inputs", "use_two_metadata_builders", diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py index 937107cd..d8da7424 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""vLLM 0.26 ModelRunnerV2 execute path with DBO dispatch and replay seams.""" +"""vLLM 0.26 ModelRunnerV2 execute path with DBO dispatch and replay seams. + +Upstream source: ``vllm/v1/worker/gpu/model_runner.py`` from +``specture724/vllm`` commit ``626fee7831``. AFD-specific changes are delimited +below so this temporary copy can be dropped when native support is available. +""" from __future__ import annotations @@ -19,7 +24,6 @@ from vllm.v1.worker.gpu.model_runner import ExecuteModelState from vllm_ascend.worker.v2.input_batch import AscendInputBatch -from . import runtime as dbo_runtime from .runtime import ( AFDBatchExecutionDescriptor, create_ubatch_slices, @@ -31,17 +35,16 @@ from vllm.v1.outputs import ModelRunnerOutput -_EXPECTED_RUNTIME_ABI = 3 -_loaded_runtime_abi = getattr(dbo_runtime, "AFD_MRV2_DBO_RUNTIME_ABI", 1) -if _loaded_runtime_abi != _EXPECTED_RUNTIME_ABI: - raise ImportError( - "AFD ModelRunnerV2 DBO backport modules are out of sync: " - f"execute expects runtime ABI {_EXPECTED_RUNTIME_ABI}, but loaded " - f"ABI {_loaded_runtime_abi}. Reinstall afd-plugin from one checkout " - "and restart every worker process." - ) - - +# Upstream source: vllm/v1/worker/gpu/model_runner.py, +# GPUModelRunner.execute_model; specture724 commit 626fee7831. +# Patch reason: pinned vLLM v0.26 lacks the ModelRunnerV2 DBO execute path and +# vLLM-Ascend does not provide its NPU adaptation. +# Patch functionality: retain the supported upstream plain-decoder flow while +# adding AFD DP dispatch, Ascend microbatch contexts, and DBO graph replay. +# Signature: extracted from GPUModelRunner.execute_model; ``runner`` replaces +# the bound ``self`` and the remaining parameters match the pinned method. +# Removal/upstream plan: delete this function when pinned vLLM and vLLM-Ascend +# provide native ModelRunnerV2 DBO execution. def execute_model_v026_eager_dbo( runner: Any, scheduler_output: SchedulerOutput, @@ -71,6 +74,7 @@ def execute_model_v026_eager_dbo( num_tokens, max_query_len, ) + # ### PATCH START: AFD v0.26 DBO dispatch batch_desc, num_tokens_across_dp = dispatch_afd_dbo_and_sync_dp( num_reqs=num_reqs, num_tokens=num_tokens, @@ -91,6 +95,7 @@ def execute_model_v026_eager_dbo( if isinstance(batch_desc, AFDBatchExecutionDescriptor) else 1 ) + # ### PATCH END: AFD v0.26 DBO dispatch if not dummy_run: runner.input_buffers.is_padding[:num_tokens].fill_(False) runner.input_buffers.is_padding[num_tokens : batch_desc.num_tokens].fill_(True) @@ -103,11 +108,13 @@ def execute_model_v026_eager_dbo( runner.req_states.num_computed_tokens.gpu, ) else: + # ### PATCH START: Ascend dummy input batch input_batch = AscendInputBatch.make_dummy( batch_desc.num_reqs or num_reqs, batch_desc.num_tokens, runner.input_buffers, ) + # ### PATCH END: Ascend dummy input batch if not skip_attn_for_dummy_run: block_tables, slot_mappings = runner.prepare_dummy_attn(input_batch) else: @@ -117,10 +124,12 @@ def execute_model_v026_eager_dbo( attn_metadata = None slot_mappings_by_layer = None ubatch_slices = None + # ### PATCH START: AFD microbatch input preparation if num_ubatches > 1: assert runner.ubatch_runner is not None assert block_tables is not None and slot_mappings is not None ubatch_slices = create_ubatch_slices(input_batch, num_ubatches) + # ### PATCH END: AFD microbatch input preparation elif not (dummy_run and skip_attn_for_dummy_run): assert block_tables is not None and slot_mappings is not None slot_mappings_by_layer = build_slot_mappings_by_layer( @@ -149,6 +158,7 @@ def execute_model_v026_eager_dbo( ubatch_slices, ) + # ### PATCH START: AFD eager and FULL microbatch execution if ubatch_slices is not None and batch_desc.cg_mode == CUDAGraphMode.FULL: assert isinstance(batch_desc, AFDBatchExecutionDescriptor) ubatch_state = runner.ubatch_runner.prepare( @@ -201,6 +211,7 @@ def execute_model_v026_eager_dbo( model_inputs, ubatch_state, ) + # ### PATCH END: AFD eager and FULL microbatch execution elif batch_desc.cg_mode == CUDAGraphMode.FULL: assert runner.cudagraph_manager is not None runner.kv_connector.pre_forward(scheduler_output) diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py index 438abdd9..59df0281 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py @@ -12,8 +12,8 @@ from collections.abc import Iterator from contextlib import contextmanager -from dataclasses import dataclass, fields, replace -from typing import Any +from dataclasses import dataclass, replace +from typing import TYPE_CHECKING, Any import numpy as np import torch @@ -34,9 +34,13 @@ ) from vllm.v1.worker.utils import AttentionGroup -AFD_MRV2_DBO_RUNTIME_ABI = 3 +if TYPE_CHECKING: + from afd_plugin.v1.worker.npu.aclgraph_manager_v2 import ( + AFDModelAclGraphManagerV2, + ) +# ### PATCH START: AFD v0.26 DBO descriptor @dataclass(frozen=True) class AFDBatchExecutionDescriptor(BatchExecutionDescriptor): """v0.26 batch descriptor extended only with the DBO execution count.""" @@ -44,17 +48,18 @@ class AFDBatchExecutionDescriptor(BatchExecutionDescriptor): num_ubatches: int = 1 -def assert_backport_required() -> None: - """Fail when the pinned vLLM ABI no longer needs this backport.""" - - descriptor_fields = {field.name for field in fields(BatchExecutionDescriptor)} - if "num_ubatches" in descriptor_fields: - raise RuntimeError( - "vLLM already provides ModelRunnerV2 DBO descriptors; remove the " - "temporary afd-plugin v0.26 backport", - ) +# ### PATCH END: AFD v0.26 DBO descriptor +# Upstream source: ``vllm/v1/worker/gpu/dp_utils.py`` from +# ``specture724/vllm`` commit ``626fee7831``. +# Patch reason: pinned vLLM's descriptor and graph-manager dispatch do not carry +# ModelRunnerV2 microbatch state. +# Patch functionality: preserve upstream DP-wide threshold selection while +# routing DBO graph selection through the plugin-owned Ascend manager. +# Signature: standalone backport helper; graph-manager and eager-control +# parameters replace the newer upstream descriptor inputs. +# Removal/upstream plan: delete this helper with the v0.26 compatibility layer. def dispatch_afd_dbo_and_sync_dp( *, num_reqs: int, @@ -65,7 +70,7 @@ def dispatch_afd_dbo_and_sync_dp( parallel_config: ParallelConfig, decode_query_len: int, allow_ubatching: bool, - cudagraph_manager: Any | None = None, + cudagraph_manager: AFDModelAclGraphManagerV2 | None = None, need_eager: bool = False, ) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: """Select DBO and FULL graph execution consistently across DP ranks. @@ -75,6 +80,7 @@ def dispatch_afd_dbo_and_sync_dp( upstream inference rule and avoids a separate per-rank DBO vote. """ + # ### PATCH START: AFD v0.26 graph dispatch adapter def dispatch( tokens: int, ubatches: int, @@ -107,24 +113,21 @@ def dispatch( num_ubatches=ubatches, ) - dispatch_ubatches = getattr(cudagraph_manager, "dispatch_ubatches", None) - if dispatch_ubatches is None: - manager_type = type(cudagraph_manager).__name__ - raise RuntimeError( - "AFD NPU ModelRunnerV2 FULL graph DBO requires the plugin " - f"graph manager, got {manager_type}" - ) - return dispatch_ubatches(base, ubatches) + return cudagraph_manager.dispatch_ubatches(base, ubatches) + + # ### PATCH END: AFD v0.26 graph dispatch adapter if dp_size == 1: return dispatch(num_tokens, 1), None + # ### PATCH START: AFD DBO graph-mode synchronization desired_single = dispatch(num_tokens, 1) tensor = torch.zeros(4, dp_size, dtype=torch.int32, device="cpu") tensor[0][dp_rank] = num_tokens tensor[1][dp_rank] = uniform_token_count or 0 tensor[2][dp_rank] = int(allow_ubatching) tensor[3][dp_rank] = int(desired_single.cg_mode.value) + # ### PATCH END: AFD DBO graph-mode synchronization dist.all_reduce(tensor, group=get_dp_group().cpu_group) num_tokens_across_dp = tensor[0] @@ -139,6 +142,7 @@ def dispatch( ) uniform_decode = bool(torch.all(tensor[1] == int(decode_query_len)).item()) + # ### PATCH START: AFD v0.26 synchronized graph descriptor synced_uniform_token_count: int | None = int(tensor[1][0].item()) if ( synced_uniform_token_count == 0 @@ -146,6 +150,7 @@ def dispatch( ): synced_uniform_token_count = None synced_mode = CUDAGraphMode(int(tensor[3].min().item())) + # ### PATCH END: AFD v0.26 synchronized graph descriptor should_ubatch = ( bool(torch.all(tensor[2] == 1).item()) and int(num_tokens_across_dp.max().item()) >= int(parallel_config.num_ubatches) @@ -170,6 +175,7 @@ def dispatch( num_tokens_across_dp.fill_(synced.num_tokens) return synced, num_tokens_across_dp + # ### PATCH START: AFD DBO graph dispatch and empty-stage fallback padded_tokens = int(num_tokens_across_dp.max().item()) num_ubatches = int(parallel_config.num_ubatches) ubatch_desc = dispatch( @@ -195,6 +201,7 @@ def dispatch( force_eager=synced_mode == CUDAGraphMode.NONE, ) num_tokens_across_dp.fill_(synced.num_tokens) + # ### PATCH END: AFD DBO graph dispatch and empty-stage fallback return synced, num_tokens_across_dp @@ -337,6 +344,7 @@ def merge_ubatch_outputs(outputs: list[Any]) -> Any: for key in first.tensors }, ) + # ### PATCH START: Ascend auxiliary hidden-state output if isinstance(first, tuple): hidden_states = torch.cat([output[0] for output in outputs], dim=0) auxiliary = [ @@ -344,6 +352,7 @@ def merge_ubatch_outputs(outputs: list[Any]) -> Any: for index in range(len(first[1])) ] return hidden_states, auxiliary + # ### PATCH END: Ascend auxiliary hidden-state output return torch.cat(outputs, dim=0) @@ -353,21 +362,31 @@ def use_two_metadata_builders() -> Iterator[None]: original = AttentionGroup.create_metadata_builders + # Upstream source: vllm/v1/worker/utils.py, + # AttentionGroup.create_metadata_builders; commit 568afb3a13. + # Patch reason: pinned vLLM initializes one metadata builder for MRV2. + # Patch functionality: force exactly two builders only during AFD DBO + # KV-cache initialization. + # Signature: matches AttentionGroup.create_metadata_builders exactly. + # Removal/upstream plan: remove this replacement when native MRV2 DBO + # initializes one metadata builder per microbatch. def create_metadata_builders( - group: AttentionGroup, + self, vllm_config, device, kernel_block_size: int | None = None, num_metadata_builders: int = 1, - ) -> None: + ): + # ### PATCH START: AFD two metadata builders del num_metadata_builders original( - group, + self, vllm_config, device, kernel_block_size, num_metadata_builders=2, ) + # ### PATCH END: AFD two metadata builders try: AttentionGroup.create_metadata_builders = create_metadata_builders @@ -376,21 +395,14 @@ def create_metadata_builders( AttentionGroup.create_metadata_builders = original -def share_metadata_builder_workspaces( - attn_groups: list[list[AttentionGroup]], -) -> None: - """Share the backend workspace while retaining independent builders.""" - - workspace = None - for groups in attn_groups: - for group in groups: - for builder in group.metadata_builders: - if workspace is None and hasattr(builder, "_get_workspace_buffer"): - workspace = builder._get_workspace_buffer() - elif workspace is not None and hasattr(builder, "set_workspace_buffer"): - builder.set_workspace_buffer(workspace) - - +# Upstream source: vllm/v1/worker/gpu/model_states/interface.py, +# ModelState.prepare_attn; commit 568afb3a13. +# Patch reason: pinned vLLM has one active metadata builder and no ubatch index. +# Patch functionality: select the builder owned by the requested microbatch +# while invoking the native attention-preparation path unchanged. +# Signature: standalone adapter; ``ubatch_index`` is the only added input. +# Removal/upstream plan: call native prepare_attn directly when it accepts an +# ubatch index. def prepare_attn_for_ubatch( model_state: ModelState, input_batch: InputBatch, @@ -415,6 +427,7 @@ def prepare_attn_for_ubatch( for_capture=for_capture, ) + # ### PATCH START: AFD per-microbatch metadata builder swapped: list[AttentionGroup] = [] for groups in attn_groups: for group in groups: @@ -439,3 +452,4 @@ def prepare_attn_for_ubatch( group.metadata_builders[ubatch_index], group.metadata_builders[0], ) + # ### PATCH END: AFD per-microbatch metadata builder diff --git a/afd_plugin/compat/patches/config_validation.py b/afd_plugin/compat/patches/config_validation.py index 02776bbf..9e8e375e 100644 --- a/afd_plugin/compat/patches/config_validation.py +++ b/afd_plugin/compat/patches/config_validation.py @@ -120,12 +120,12 @@ def __post_init__(self): """Verify configs are valid & consistent with each other.""" assert _original_vllm_config_post_init is not None + # ### PATCH START: AFD repeated ubatching and MRV2 DBO validation relax_backend = _should_relax_vllm_config_backend(self) relax_v2_dbo = _should_relax_npu_v2_dbo_validation(self) if not relax_backend and not relax_v2_dbo: return _original_vllm_config_post_init(self) - # ### PATCH START: AFD repeated ubatching and MRV2 DBO validation parallel_config = self.parallel_config original_backend = parallel_config.all2all_backend if relax_backend: diff --git a/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py b/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py index 8b7feb15..4a707857 100644 --- a/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py +++ b/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py @@ -1,6 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Scoped graph-manager injection for the temporary MRV2 DBO backport.""" +"""Scoped graph-manager injection for the temporary MRV2 DBO backport. + +Upstream source: ``vllm_ascend/worker/v2/model_runner.py`` at commit +``d543ccee0``, function ``graph_manager_wrapper``. +""" from __future__ import annotations @@ -12,51 +16,67 @@ from vllm.v1.worker.gpu import model_runner as vllm_model_runner from vllm_ascend.worker.v2 import model_runner as ascend_model_runner -from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( - share_metadata_builder_workspaces, -) from afd_plugin.v1.worker.npu.aclgraph_manager_v2 import ( AFDModelAclGraphManagerV2, ) from afd_plugin.v1.worker.npu.ubatch_runner_v2 import AFDAscendUBatchRunnerV2 +# Patch reason: vLLM-Ascend's initialization wrapper always constructs the +# native single-batch ACL graph manager. +# Patch functionality: replace only that initialization scope with the AFD DBO +# manager and restore the upstream wrapper afterwards. +# Signature: plugin-owned context manager; ``model_runner`` is the AFD runner +# whose initialization is being scoped. +# Removal/upstream plan: delete this wrapper when vLLM-Ascend accepts a native +# ModelRunnerV2 DBO graph manager/runner pair. @contextmanager def use_afd_mrv2_dbo_graph_manager(model_runner) -> Iterator[None]: """Replace only Ascend's initialization-scoped manager wrapper.""" original_wrapper = ascend_model_runner.graph_manager_wrapper + # Patch reason: the upstream wrapper factory has no DBO runner dependency. + # Patch functionality: construct the AFD two-stage runner and graph manager + # while retaining upstream's temporary ModelCudaGraphManager substitution. + # Signature: matches vLLM-Ascend graph_manager_wrapper exactly. @contextmanager - def graph_manager_wrapper(runner) -> Iterator[None]: + def graph_manager_wrapper(model_runner): original_manager = vllm_model_runner.ModelCudaGraphManager + # Upstream source: vllm_ascend/worker/v2/model_runner.py, + # graph_manager_wrapper.factory; commit d543ccee0. + # Patch reason: the native factory cannot receive an ubatch runner. + # Patch functionality: construct the plugin-owned runner and manager. + # Signature: matches the native nested factory exactly. + # Removal/upstream plan: delete with this scoped wrapper. def factory( vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, decode_query_len: int, lora_capture_cases: list[int] | None = None, - ) -> AFDModelAclGraphManagerV2: - share_metadata_builder_workspaces(runner.attn_groups) + ): + # ### PATCH START: AFD MRV2 DBO graph manager ubatch_runner = AFDAscendUBatchRunnerV2( vllm_config, device, - runner.model_state, - runner.attn_groups, - runner.kv_cache_config, - runner.max_num_reqs, + model_runner.model_state, + model_runner.attn_groups, + model_runner.kv_cache_config, + model_runner.max_num_reqs, ) - runner.ubatch_runner = ubatch_runner + model_runner.ubatch_runner = ubatch_runner return AFDModelAclGraphManagerV2( vllm_config, device, cudagraph_mode, decode_query_len, - runner, + model_runner, ubatch_runner, lora_capture_cases=lora_capture_cases, ) + # ### PATCH END: AFD MRV2 DBO graph manager try: vllm_model_runner.ModelCudaGraphManager = factory diff --git a/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py index 6ee5aa2e..32762892 100644 --- a/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py +++ b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py @@ -1,12 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Temporary MRV2 DBO FULL ACL graph manager for vLLM 0.26.""" +"""Temporary MRV2 DBO FULL ACL graph manager for vLLM 0.26. + +Upstream sources: ``vllm/v1/worker/gpu/cudagraph_utils.py`` from +``specture724/vllm`` commit ``626fee7831`` and +``vllm_ascend/worker/v2/aclgraph_utils.py`` at commit ``d543ccee0``. +""" from __future__ import annotations from collections.abc import Callable from dataclasses import dataclass -from typing import Any +from typing import TYPE_CHECKING import torch import torch.nn as nn @@ -26,6 +31,7 @@ from vllm.v1.worker.ubatch_utils import check_ubatch_thresholds from vllm.v1.worker.utils import AttentionGroup from vllm_ascend.compilation.acl_graph import ( + GraphParams, get_graph_params, update_full_graph_params, ) @@ -47,12 +53,18 @@ AFDAscendUBatchState, ) +if TYPE_CHECKING: + from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( + AFDNPUAttentionModelRunnerV2, + ) + from afd_plugin.v1.worker.npu.npu_ubatch_wrapper import AscendModelOutput + @dataclass class _AFDGraphEntry: - graph: Any - output: Any - graph_params: tuple[Any, Any] + graph: torch.npu.NPUGraph + output: AscendModelOutput + graph_params: tuple[GraphParams, GraphParams] workspace: torch.Tensor @@ -63,15 +75,26 @@ class _PendingReplay: class AFDModelAclGraphManagerV2(ModelAclGraphManager): - """Keep DBO graph descriptors and storage separate from upstream graphs.""" + """Keep DBO graph descriptors and storage separate from upstream graphs. + ``ubatch_runner`` is the only parameter added to the upstream + ``ModelAclGraphManager`` constructor. + """ + + # Upstream source: vllm_ascend/worker/v2/aclgraph_utils.py, + # ModelAclGraphManager.__init__; commit d543ccee0. + # Patch reason: the native manager has no ModelRunnerV2 DBO executor. + # Patch functionality: retain native initialization and register separate + # two-microbatch capture descriptors, graphs, and replay state. + # Signature: adds only ``ubatch_runner`` to the native constructor. + # Removal/upstream plan: use the native manager when it owns DBO graphs. def __init__( self, vllm_config: VllmConfig, device: torch.device, cudagraph_mode: CUDAGraphMode, decode_query_len: int, - model_runner: Any, + model_runner: AFDNPUAttentionModelRunnerV2, ubatch_runner: AFDAscendUBatchRunnerV2, lora_capture_cases: list[int] | None = None, ) -> None: @@ -83,6 +106,7 @@ def __init__( model_runner, lora_capture_cases=lora_capture_cases, ) + # ### PATCH START: AFD DBO graph registry self.ubatch_runner = ubatch_runner self._afd_twins = { desc: AFDBatchExecutionDescriptor( @@ -98,6 +122,7 @@ def __init__( } self._afd_graphs: dict[AFDBatchExecutionDescriptor, _AFDGraphEntry] = {} self._afd_pending_replay: _PendingReplay | None = None + # ### PATCH END: AFD DBO graph registry def _needs_ubatch_twin(self, desc: BatchExecutionDescriptor) -> bool: if desc.num_tokens % 2 or desc.num_tokens < 2: @@ -144,6 +169,12 @@ def clear_afd_graphs(self) -> None: self._afd_pending_replay = None self._afd_graphs.clear() + # Patch reason: native vLLM-Ascend captures only single-batch descriptors. + # Patch functionality: preserve native capture, then capture eligible AFD + # two-stage twins in a separate registry. + # Signature: matches ModelAclGraphManager.capture exactly. + # Removal/upstream plan: delete this override when native Ascend capture + # accepts DBO descriptors and an ubatch runner. def capture( self, model: nn.Module, @@ -171,6 +202,7 @@ def capture( lora_capture_hook=lora_capture_hook, progress_bar_desc=progress_bar_desc, ) + # ### PATCH START: AFD DBO twin graph capture if not self._afd_twins: return @@ -189,6 +221,7 @@ def capture( except BaseException: self.clear_afd_graphs() raise + # ### PATCH END: AFD DBO twin graph capture def _capture_afd_graph( self, @@ -280,7 +313,7 @@ def _prepare_capture_state( block_tables: tuple[torch.Tensor, ...], slot_mappings: torch.Tensor, slices, - graph_params: tuple[Any, Any] | None, + graph_params: tuple[GraphParams, GraphParams] | None, *, is_warmup: bool, ) -> AFDAscendUBatchState: @@ -322,9 +355,19 @@ def _prepare_capture_state( mla_graph_params=graph_params, ) - def run_fullgraph(self, desc: BatchExecutionDescriptor) -> Any: + # Patch reason: native replay has no staged two-microbatch state or AFD + # control-plane metadata. + # Patch functionality: delegate native descriptors unchanged and replay + # only AFD descriptors from the plugin-owned registry. + # Signature: matches ModelAclGraphManager.run_fullgraph exactly. + # Removal/upstream plan: delete this override with the DBO graph registry. + def run_fullgraph( + self, + desc: BatchExecutionDescriptor, + ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: if not isinstance(desc, AFDBatchExecutionDescriptor): return super().run_fullgraph(desc) + # ### PATCH START: AFD DBO FULL graph replay pending = self._afd_pending_replay self._afd_pending_replay = None if pending is None or pending.descriptor != desc: @@ -379,6 +422,7 @@ def run_fullgraph(self, desc: BatchExecutionDescriptor) -> Any: self.vllm_config, self.model_runner.speculative_config, ) + # ### PATCH END: AFD DBO FULL graph replay return entry.output 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 3b69a2be..fa574949 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -24,6 +24,10 @@ from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( AFDBatchExecutionDescriptor, + use_two_metadata_builders, +) +from afd_plugin.compat.backports.vllm_v026_mrv2_dbo.execute import ( + execute_model_v026_eager_dbo, ) from afd_plugin.compat.npu import fail_if_unsupported_npu_afd_features from afd_plugin.compat.npu.profiler import ( @@ -31,6 +35,9 @@ step_afd_npu_profiler, stop_afd_npu_profiler, ) +from afd_plugin.compat.patches.npu.model_runner_v2_dbo import ( + use_afd_mrv2_dbo_graph_manager, +) from afd_plugin.config import AFDConfig, parse_afd_config from afd_plugin.connectors import ( AFDConnectorBase, @@ -179,7 +186,9 @@ def __init__( self._afd_pending_metadata: AFDForwardContextMetadata | None = None self._afd_suppress_metadata_send = False self._afd_transaction_counter = 0 + # ### PATCH START: AFD MRV2 DBO runner state self.ubatch_runner = None + # ### PATCH END: AFD MRV2 DBO runner state self.prof = create_afd_npu_profiler("attention") except BaseException: try: @@ -222,6 +231,15 @@ def load_model( if not self.connector.is_initialized: self.connector.init_afd_connector() + # Upstream source: vLLM v0.26.0 commit 568afb3a1, + # GPUModelRunner.initialize_kv_cache, with vLLM-Ascend's scoped graph + # manager wrapper from commit d543ccee0. + # Patch reason: neither pinned upstream initializes MRV2 metadata builders, + # an Ascend ubatch runner, or DBO graph descriptors. + # Patch functionality: retain native initialization when DBO is disabled; + # otherwise scope the temporary two-builder and graph-manager replacements. + # Signature: matches NPUModelRunnerV2.initialize_kv_cache exactly. + # Removal/upstream plan: remove the DBO branch with the v0.26 backport. def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: """Initialize native state and the temporary v0.26 eager DBO runner.""" @@ -229,22 +247,13 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: super().initialize_kv_cache(kv_cache_config) return - from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( - assert_backport_required, - use_two_metadata_builders, - ) - from afd_plugin.compat.patches.npu.model_runner_v2_dbo import ( - use_afd_mrv2_dbo_graph_manager, - ) - - assert_backport_required() + # ### PATCH START: AFD MRV2 DBO initialization with ( use_two_metadata_builders(), use_afd_mrv2_dbo_graph_manager(self), ): super().initialize_kv_cache(kv_cache_config) - if self.ubatch_runner is None: - raise RuntimeError("AFD MRV2 DBO graph manager was not initialized") + # ### PATCH END: AFD MRV2 DBO initialization # Patch reason: vLLM v0.26.0 prepares FULL graph inputs before each warmup # and formal capture forward, outside torch.cuda.graph, but does not expose @@ -416,11 +425,8 @@ def execute_model( self.install_afd_metadata_on_forward_context, ), ): + # ### PATCH START: AFD MRV2 DBO execute backport if self.vllm_config.parallel_config.use_ubatching: - from afd_plugin.compat.backports.vllm_v026_mrv2_dbo.execute import ( - execute_model_v026_eager_dbo, - ) - return execute_model_v026_eager_dbo( self, scheduler_output, @@ -429,6 +435,7 @@ def execute_model( skip_attn_for_dummy_run=skip_attn_for_dummy_run, is_profile=is_profile, ) + # ### PATCH END: AFD MRV2 DBO execute backport return super().execute_model( scheduler_output, intermediate_tensors, @@ -463,11 +470,13 @@ def shutdown(self) -> None: stop_afd_npu_profiler(self.prof) finally: try: + # ### PATCH START: AFD MRV2 DBO graph cleanup if isinstance( self.cudagraph_manager, AFDModelAclGraphManagerV2, ): self.cudagraph_manager.clear_afd_graphs() + # ### PATCH END: AFD MRV2 DBO graph cleanup super().shutdown() finally: self._afd_pending_metadata = None diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 634db29e..aa856f37 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -402,6 +402,7 @@ def _capture_graphs( "Only DP metadata control plane supports graph capturing." ) graph_key = self._make_graph_key(dp_metadata_list) + # ### PATCH START: AFD MRV2 repeated capture replay if graph_key in self._acl_graphs: logger.debug( "AFD NPU FFN replaying existing ACL graph for capture key=%s", @@ -415,6 +416,7 @@ def _capture_graphs( ) self._acl_graphs[graph_key]["graph"].replay() return + # ### PATCH END: AFD MRV2 repeated capture replay logger.debug("AFD NPU FFN capturing ACL graph for key=%s", graph_key) graph = torch.npu.NPUGraph() diff --git a/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index 43d76e0a..8da7d157 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/afd_plugin/v1/worker/npu/forward_context.py @@ -26,28 +26,15 @@ from vllm_ascend.compilation.acl_graph import GraphParams -def _get_ascend_extra( - forward_context: ForwardContext, - vllm_config: VllmConfig, - name: str, -): - if vllm_config.use_v2_model_runner: - return forward_context.additional_kwargs.get(name) - return getattr(forward_context, name) - - -def _set_ascend_extra( - forward_context: ForwardContext, - vllm_config: VllmConfig, - name: str, - value, -) -> None: - if vllm_config.use_v2_model_runner: - forward_context.additional_kwargs[name] = value - else: - setattr(forward_context, name, value) - - +# Upstream source: vllm-ascend commit 80d8c194f, +# set_additional_forward_context and the ModelRunnerV1 forward-context fields. +# Patch reason: vLLM-Ascend stores Ascend state as direct attributes for MRV1 +# but as guaranteed additional_kwargs entries for MRV2. +# Patch functionality: retain the MRV1 attribute path and update only the +# per-microbatch values in MRV2's concrete additional_kwargs contract. +# Signature: plugin-owned helper; no upstream symbol is replaced. +# Removal/upstream plan: use native Ascend MRV2 ubatch context construction +# when vLLM-Ascend provides it. def create_ascend_forward_context( cur_forward_context: ForwardContext, attn_metadata, @@ -63,7 +50,7 @@ def create_ascend_forward_context( if cudagraph_runtime_mode is None: cudagraph_runtime_mode = CUDAGraphMode.NONE - parent_kwargs = dict(cur_forward_context.additional_kwargs or {}) + parent_kwargs = dict(cur_forward_context.additional_kwargs) afd_metadata = parent_kwargs.get("afd_metadata") if afd_metadata is not None: parent_kwargs = build_ubatch_additional_kwargs( @@ -97,143 +84,126 @@ def create_ascend_forward_context( tp_world_size = get_tensor_model_parallel_world_size() dp_world_size = get_dp_group().world_size - moe_comm_type = _get_ascend_extra( - cur_forward_context, - vllm_config, - "moe_comm_type", - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "moe_comm_type", - moe_comm_type, - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "moe_comm_method", - get_moe_comm_method(moe_comm_type), - ) - for name in ( - "in_profile_run", - "mmrs_fusion", - "is_first_layer", - "layer_idx", - "prefetch_mlp_gate_up_proj", - "prefetch_mlp_down_proj", - "model_instance", - "is_draft_model", - "is_draft_model_prefill", - "draft_attn_metadatas", - "max_tokens_across_pcp", - "sinks", - "input_ids", - "eplb_heat_collection_status", - ): - _set_ascend_extra( - new_forward_context, - vllm_config, - name, - _get_ascend_extra(cur_forward_context, vllm_config, name), - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "capturing", - mla_graph_params is not None - or _get_ascend_extra(cur_forward_context, vllm_config, "capturing"), - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "num_tokens", - num_tokens, - ) new_forward_context.ubatch_idx = int(ubatch_num) new_forward_context.num_ubatches = len(ubatch_slices) - flash_comm_v1_enabled = bool( - _get_ascend_extra( - cur_forward_context, - vllm_config, - "flash_comm_v1_enabled", - ) - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "flash_comm_v1_enabled", - flash_comm_v1_enabled, - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "pad_size", - 0, - ) - if flash_comm_v1_enabled: - _set_ascend_extra( - new_forward_context, - vllm_config, - "pad_size", - (tp_world_size - (num_tokens % tp_world_size)) % tp_world_size, - ) - - if dp_world_size > 1 and dp_metadata is not None: - max_tokens_across_dp = dp_metadata.num_tokens_across_dp_cpu.max().item() + # ### PATCH START: MRV2 Ascend additional kwargs + if vllm_config.use_v2_model_runner: + source = cur_forward_context.additional_kwargs + target = new_forward_context.additional_kwargs + target["capturing"] = mla_graph_params is not None or source["capturing"] + target["num_tokens"] = num_tokens + + flash_comm_v1_enabled = source["flash_comm_v1_enabled"] + pad_size = 0 + padded_length = None if flash_comm_v1_enabled: - padded_length = ( - (max_tokens_across_dp + tp_world_size - 1) - // tp_world_size - * tp_world_size - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "padded_length", - padded_length, - ) - _set_ascend_extra( - new_forward_context, - vllm_config, - "pad_size", - padded_length - num_tokens, + pad_size = (tp_world_size - (num_tokens % tp_world_size)) % tp_world_size + + if dp_world_size > 1 and dp_metadata is not None: + max_tokens_across_dp = dp_metadata.num_tokens_across_dp_cpu.max().item() + if flash_comm_v1_enabled: + padded_length = ( + (max_tokens_across_dp + tp_world_size - 1) + // tp_world_size + * tp_world_size + ) + pad_size = padded_length - num_tokens + else: + max_tokens_across_dp = num_tokens + + padded_num_tokens = ( + math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size + ) + mc2_mask = None + source_mc2_mask = source["mc2_mask"] + if source_mc2_mask is not None: + mc2_mask = torch.zeros( + (padded_num_tokens,), + dtype=source_mc2_mask.dtype, + device=source_mc2_mask.device, ) + mc2_mask[:num_tokens] = True + mc2_mask[num_tokens:] = False + + target["pad_size"] = pad_size + target["padded_length"] = padded_length + target["max_tokens_across_dp"] = max_tokens_across_dp + target["padded_num_tokens"] = padded_num_tokens + target["mc2_mask"] = mc2_mask + # ### PATCH END: MRV2 Ascend additional kwargs else: - max_tokens_across_dp = num_tokens - _set_ascend_extra( - new_forward_context, - vllm_config, - "max_tokens_across_dp", - max_tokens_across_dp, - ) - - padded_num_tokens = math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size - _set_ascend_extra( - new_forward_context, - vllm_config, - "padded_num_tokens", - padded_num_tokens, - ) - cur_mc2_mask = _get_ascend_extra( - cur_forward_context, - vllm_config, - "mc2_mask", - ) - if cur_mc2_mask is not None: - mc2_mask = torch.zeros( - (padded_num_tokens,), - dtype=cur_mc2_mask.dtype, - device=cur_mc2_mask.device, + new_forward_context.moe_comm_type = cur_forward_context.moe_comm_type + new_forward_context.moe_comm_method = get_moe_comm_method( + new_forward_context.moe_comm_type + ) + new_forward_context.in_profile_run = cur_forward_context.in_profile_run + new_forward_context.capturing = ( + mla_graph_params is not None or cur_forward_context.capturing ) - mc2_mask[:num_tokens] = True - mc2_mask[num_tokens:] = False - _set_ascend_extra( - new_forward_context, - vllm_config, - "mc2_mask", - mc2_mask, + new_forward_context.mmrs_fusion = cur_forward_context.mmrs_fusion + new_forward_context.num_tokens = num_tokens + new_forward_context.flash_comm_v1_enabled = ( + cur_forward_context.flash_comm_v1_enabled ) + new_forward_context.pad_size = 0 + new_forward_context.is_first_layer = cur_forward_context.is_first_layer + new_forward_context.layer_idx = cur_forward_context.layer_idx + new_forward_context.prefetch_mlp_gate_up_proj = ( + cur_forward_context.prefetch_mlp_gate_up_proj + ) + new_forward_context.prefetch_mlp_down_proj = ( + cur_forward_context.prefetch_mlp_down_proj + ) + new_forward_context.model_instance = cur_forward_context.model_instance + new_forward_context.is_draft_model = cur_forward_context.is_draft_model + new_forward_context.is_draft_model_prefill = ( + cur_forward_context.is_draft_model_prefill + ) + new_forward_context.draft_attn_metadatas = ( + cur_forward_context.draft_attn_metadatas + ) + new_forward_context.max_tokens_across_pcp = ( + cur_forward_context.max_tokens_across_pcp + ) + new_forward_context.sinks = cur_forward_context.sinks + new_forward_context.input_ids = cur_forward_context.input_ids + new_forward_context.eplb_heat_collection_status = ( + cur_forward_context.eplb_heat_collection_status + ) + + if new_forward_context.flash_comm_v1_enabled: + new_forward_context.pad_size = ( + tp_world_size - (num_tokens % tp_world_size) + ) % tp_world_size + + if dp_world_size > 1 and dp_metadata is not None: + max_tokens_across_dp = dp_metadata.num_tokens_across_dp_cpu.max().item() + if new_forward_context.flash_comm_v1_enabled: + padded_length = ( + (max_tokens_across_dp + tp_world_size - 1) + // tp_world_size + * tp_world_size + ) + new_forward_context.padded_length = padded_length + new_forward_context.pad_size = padded_length - num_tokens + else: + max_tokens_across_dp = num_tokens + new_forward_context.max_tokens_across_dp = max_tokens_across_dp + + new_forward_context.padded_num_tokens = ( + math.ceil(max_tokens_across_dp / tp_world_size) * tp_world_size + ) + cur_mc2_mask = cur_forward_context.mc2_mask + if cur_mc2_mask is not None: + mc2_mask = torch.zeros( + (new_forward_context.padded_num_tokens,), + dtype=cur_mc2_mask.dtype, + device=cur_mc2_mask.device, + ) + mc2_mask[:num_tokens] = True + mc2_mask[num_tokens:] = False + new_forward_context.mc2_mask = mc2_mask new_forward_context.dbo_enabled = True return new_forward_context diff --git a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py index 3564a181..04985f6b 100644 --- a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -1,6 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Ascend eager and capture-time DBO runner for vLLM 0.26 ModelRunnerV2.""" +"""Ascend eager and capture-time DBO runner for vLLM 0.26 ModelRunnerV2. + +Upstream source: ``vllm/v1/worker/gpu/ubatch_utils.py`` from +``specture724/vllm`` commit ``626fee7831``. PATCH markers identify the Ascend +and AFD adaptations to that runner. +""" from __future__ import annotations @@ -21,8 +26,9 @@ from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import build_slot_mappings_by_layer from vllm.v1.worker.gpu.model_states.interface import ModelState -from vllm.v1.worker.ubatch_utils import UBatchSlices +from vllm.v1.worker.ubatch_utils import UBatchSlice, UBatchSlices from vllm.v1.worker.utils import AttentionGroup +from vllm_ascend.compilation.acl_graph import GraphParams from vllm_ascend.worker.v2.input_batch import AscendInputBatch from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( @@ -34,8 +40,14 @@ merge_ubatch_outputs, ) from afd_plugin.v1.worker.npu.forward_context import create_ascend_forward_context -from afd_plugin.v1.worker.npu.npu_ubatch_wrapper import _all_gather_ubatch_output -from afd_plugin.v1.worker.npu.ubatching import make_ubatch_contexts +from afd_plugin.v1.worker.npu.npu_ubatch_wrapper import ( + AscendModelOutput, + _all_gather_ubatch_output, +) +from afd_plugin.v1.worker.npu.ubatching import ( + AscendUBatchContext, + make_ubatch_contexts, +) AFD_NPU_MRV2_NUM_UBATCHES = 2 @@ -53,6 +65,14 @@ class AFDAscendUBatchState: class AFDAscendUBatchRunnerV2: """Run exactly two ModelRunnerV2 microbatches on Ascend.""" + # Upstream source: vllm/v1/worker/gpu/ubatch_utils.py, + # UBatchRunner.__init__; specture724 commit 626fee7831. + # Patch reason: the GPU runner owns CUDA streams and supports a variable + # microbatch count; the AFD Ascend handoff protocol has exactly two stages. + # Patch functionality: allocate Ascend capture state and stage-local input + # buffers without CUDA SM-control state. + # Signature: matches UBatchRunner.__init__ exactly. + # Removal/upstream plan: replace this class with native Ascend MRV2 DBO. def __init__( self, vllm_config: VllmConfig, @@ -64,15 +84,20 @@ def __init__( ) -> None: self.vllm_config = vllm_config self.parallel_config = vllm_config.parallel_config - self.num_ubatches = int(self.parallel_config.num_ubatches) + self.num_ubatches = self.parallel_config.num_ubatches + # ### PATCH START: AFD NPU two-stage constraint if self.num_ubatches != AFD_NPU_MRV2_NUM_UBATCHES: raise RuntimeError("AFD NPU ModelRunnerV2 requires exactly two ubatches") + # ### PATCH END: AFD NPU two-stage constraint self.device = device self.model_state = model_state self.attn_groups = attn_groups self.kv_cache_config = kv_cache_config self.ready_barrier = threading.Barrier(self.num_ubatches + 1) + # ### PATCH START: Ascend capture stream self.capture_stream = torch.npu.Stream(device=device) + # ### PATCH END: Ascend capture stream + # ### PATCH START: Ascend stage-local input buffers self.query_start_loc_buffers = [ torch.zeros(max_num_reqs + 2, dtype=torch.int32, device=device) for _ in range(self.num_ubatches) @@ -81,7 +106,17 @@ def __init__( torch.zeros(max_num_reqs, dtype=torch.int32, device=device) for _ in range(self.num_ubatches) ] + # ### PATCH END: Ascend stage-local input buffers + # Upstream source: vllm/v1/worker/gpu/ubatch_utils.py, + # UBatchRunner.prepare; specture724 commit 626fee7831. + # Patch reason: Ascend requires additional InputBatch fields, its own + # ForwardContext values, and per-stage MLA graph parameters. + # Patch functionality: prepare the supplied two slices for eager execution, + # graph warmup, or graph capture. + # Signature: adds ``slices``, ``parent_context``, ``context_cg_mode``, and + # ``mla_graph_params``; the return type is the concrete Ascend state. + # Removal/upstream plan: use native Ascend UBatchRunner.prepare when added. def prepare( self, input_batch: AscendInputBatch, @@ -93,15 +128,16 @@ def prepare( cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, context_cg_mode: CUDAGraphMode | None = None, for_capture: bool = False, - mla_graph_params: tuple[Any, Any] | None = None, + mla_graph_params: tuple[GraphParams, GraphParams] | None = None, ) -> AFDAscendUBatchState: attn_metadata_list: list[dict[str, Any]] = [] forward_contexts: list[ForwardContext] = [] real_token_counts: list[int] = [] if context_cg_mode is None: context_cg_mode = cg_mode - dp_size = int(self.parallel_config.data_parallel_size) + dp_size = self.parallel_config.data_parallel_size for stage_index, stage in enumerate(slices): + # ### PATCH START: Ascend input batch fields child_batch = cast( AscendInputBatch, slice_input_batch( @@ -112,10 +148,12 @@ def prepare( ), ) child_batch = self._with_ascend_fields(input_batch, child_batch, stage) + # ### PATCH END: Ascend input batch fields stage_slot_mappings = slot_mappings[:, stage.token_slice] stage_block_tables = tuple( block_table[stage.request_slice] for block_table in block_tables ) + # ### PATCH START: AFD v0.26 metadata builder selection attn_metadata = prepare_attn_for_ubatch( self.model_state, child_batch, @@ -127,6 +165,7 @@ def prepare( cg_mode=cg_mode, for_capture=for_capture, ) + # ### PATCH END: AFD v0.26 metadata builder selection attn_metadata_list.append(attn_metadata) real_token_counts.append(int(child_batch.num_tokens)) if parent_context is None: @@ -140,6 +179,7 @@ def prepare( stage_tokens, counts, ) + # ### PATCH START: Ascend microbatch forward context context = create_ascend_forward_context( parent_context, attn_metadata, @@ -157,6 +197,7 @@ def prepare( stage_slot_mappings, self.kv_cache_config, ) + # ### PATCH END: Ascend microbatch forward context forward_contexts.append(context) return AFDAscendUBatchState( @@ -170,7 +211,7 @@ def prepare( def _with_ascend_fields( parent_batch: AscendInputBatch, child_batch: AscendInputBatch, - stage, + stage: UBatchSlice, ) -> AscendInputBatch: req_start = int(stage.request_slice.start) req_stop = int(stage.request_slice.stop) @@ -185,27 +226,40 @@ def _with_ascend_fields( attn_state=parent_batch.attn_state, ) + # Upstream source: vllm/v1/worker/gpu/ubatch_utils.py, UBatchRunner.run; + # specture724 commit 626fee7831. + # Patch reason: this adapter consumes the concrete Ascend ubatch state. + # Patch functionality: start and finish one two-stage Ascend execution. + # Signature: narrows the state and return types to the Ascend contract. + # Removal/upstream plan: use native Ascend UBatchRunner.run when added. def run( self, model: Any, model_inputs: dict[str, Any], - state: AFDAscendUBatchState, - ) -> Any: - return self.begin_capturable_run(model, model_inputs, state)() + ubatch_state: AFDAscendUBatchState, + ) -> AscendModelOutput: + return self.begin_capturable_run(model, model_inputs, ubatch_state)() + # Upstream source: vllm/v1/worker/gpu/ubatch_utils.py, + # UBatchRunner.begin_capturable_run; specture724 commit 626fee7831. + # Patch reason: GPU stream/SM control must be replaced by the existing + # Ascend two-stage contexts and FlashComm output handling. + # Patch functionality: launch both stages and return their finisher. + # Signature: narrows the state/output types; ``for_capture`` is unchanged. + # Removal/upstream plan: use native Ascend capturable DBO when available. def begin_capturable_run( self, model: Any, model_inputs: dict[str, Any], - state: AFDAscendUBatchState, - *, + ubatch_state: AFDAscendUBatchState, for_capture: bool = False, - ) -> Callable[[], Any]: + ) -> Callable[[], AscendModelOutput]: """Start both stages outside capture and return a one-shot finisher.""" - forward_contexts = state.forward_contexts + forward_contexts = ubatch_state.forward_contexts if forward_contexts is None: raise RuntimeError("uBatch execution requires prepared forward contexts") + # ### PATCH START: Ascend microbatch contexts compute_stream = ( self.capture_stream if for_capture else torch.npu.current_stream() ) @@ -215,13 +269,19 @@ def begin_capturable_run( forward_contexts, self.ready_barrier, ) - outputs: dict[int, Any] = {} + # ### PATCH END: Ascend microbatch contexts + outputs: dict[int, AscendModelOutput] = {} errors: dict[int, BaseException] = {} @torch.inference_mode() - def run_stage(context, inputs: dict[str, Any]) -> None: + def run_stage( + context: AscendUBatchContext, + inputs: dict[str, Any], + ) -> None: try: + # ### PATCH START: Ascend worker device torch.npu.set_device(self.device) + # ### PATCH END: Ascend worker device with context: outputs[context.id] = model(**inputs) except BaseException as error: # noqa: BLE001 @@ -230,7 +290,7 @@ def run_stage(context, inputs: dict[str, Any]) -> None: stack = ExitStack() stack.enter_context(override_forward_context(None)) threads = [] - for context, stage in zip(contexts, state.slices, strict=True): + for context, stage in zip(contexts, ubatch_state.slices, strict=True): thread = threading.Thread( target=run_stage, args=(context, slice_model_inputs(model_inputs, stage.token_slice)), @@ -238,13 +298,8 @@ def run_stage(context, inputs: dict[str, Any]) -> None: threads.append(thread) thread.start() self.ready_barrier.wait() - finished = False - def finish() -> Any: - nonlocal finished - if finished: - raise RuntimeError("uBatch finisher may only be called once") - finished = True + def finish() -> AscendModelOutput: try: contexts[0].cpu_wait_event.set() for thread in threads: @@ -258,6 +313,7 @@ def finish() -> Any: f"AFD NPU microbatch {failed_stage} failed", ) from errors[failed_stage] ordered_outputs = [outputs[index] for index in range(self.num_ubatches)] + # ### PATCH START: Ascend FlashComm output gathering if forward_contexts[0].additional_kwargs["flash_comm_v1_enabled"]: ordered_outputs = [ _all_gather_ubatch_output( @@ -270,6 +326,7 @@ def finish() -> Any: strict=True, ) ] + # ### PATCH END: Ascend FlashComm output gathering return merge_ubatch_outputs(ordered_outputs) return finish diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index a4cf6911..deade7de 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -147,6 +147,7 @@ def validate_npu_model_runner_v2_config( or vllm_config.compilation_config.pass_config.enable_sp ): raise RuntimeError("AFD ModelRunnerV2 requires static expert parallelism") + # ### PATCH START: AFD NPU MRV2 DBO validation dbo_enabled = bool(parallel.enable_dbo or parallel.use_ubatching) if dbo_enabled: if parallel.data_parallel_size <= 1 or int(parallel.num_ubatches) != 2: @@ -170,6 +171,7 @@ def validate_npu_model_runner_v2_config( "AFD NPU ModelRunnerV2 DBO does not support speculative decode, " "LoRA, multimodal, or encoder models", ) + # ### PATCH END: AFD NPU MRV2 DBO validation from afd_plugin.model_executor.models.model_utils import ( has_afd_model_registration, diff --git a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py index c0257334..972ccbec 100644 --- a/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py +++ b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py @@ -35,10 +35,9 @@ def _parallel_config(*, decode_threshold=4, prefill_threshold=8): ) -def test_runtime_abi_matches_full_graph_dispatch_signature(): +def test_full_graph_dispatch_signature(): parameters = inspect.signature(dispatch_afd_dbo_and_sync_dp).parameters - assert runtime.AFD_MRV2_DBO_RUNTIME_ABI == 3 assert {"cudagraph_manager", "need_eager"} <= parameters.keys() From 0af85085b0058f70fda83df6e6f19940df65822b Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Thu, 27 Aug 2026 19:03:39 +0800 Subject: [PATCH 5/7] fix the errors in mypy tests and add SPDX headers Signed-off-by: lirx-pd <616517220@qq.com> --- .../backports/vllm_v026_mrv2_dbo/runtime.py | 1 + afd_plugin/v1/worker/attention_metadata.py | 6 ++++ .../worker/npu/attention_model_runner_v2.py | 3 +- afd_plugin/v1/worker/npu/ffn_model_runner.py | 26 ++++++++++++----- .../compat/patches/test_config_validation.py | 9 ++++-- tests/unit/v1/worker/test_model_runner_v2.py | 29 ++++++++++--------- tests/unit/v1/worker/test_npu_mla_graph.py | 10 +++++-- tests/unit/v1/worker/test_npu_runtime.py | 2 ++ 8 files changed, 58 insertions(+), 28 deletions(-) diff --git a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py index 59df0281..60a3b08f 100644 --- a/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py @@ -113,6 +113,7 @@ def dispatch( num_ubatches=ubatches, ) + assert cudagraph_manager is not None return cudagraph_manager.dispatch_ubatches(base, ubatches) # ### PATCH END: AFD v0.26 graph dispatch adapter diff --git a/afd_plugin/v1/worker/attention_metadata.py b/afd_plugin/v1/worker/attention_metadata.py index f633dc59..f401c04d 100644 --- a/afd_plugin/v1/worker/attention_metadata.py +++ b/afd_plugin/v1/worker/attention_metadata.py @@ -11,6 +11,7 @@ from vllm.v1.worker.ubatch_utils import UBatchSlices from afd_plugin.connectors import ( + AFDConnectorBase, AFDControlPayload, AFDDPMetadata, AFDForwardContextMetadata, @@ -28,6 +29,11 @@ class AFDMetadataProviderMixin: """ _afd_is_profile: bool = False + _afd_pending_metadata: AFDForwardContextMetadata | None + _afd_transaction_counter: int + _is_warmup: bool + connector: AFDConnectorBase + vllm_config: VllmConfig def build_afd_metadata( self, 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 fa574949..f5802f69 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -7,6 +7,7 @@ from collections.abc import Iterator from contextlib import contextmanager, nullcontext from types import MethodType +from typing import Any, cast import torch from vllm.config import CUDAGraphMode, VllmConfig @@ -69,7 +70,7 @@ def _use_afd_fullgraph_replay_hook( raise RuntimeError( "AFD ACL graph replay hook requires an initialized graph manager", ) - manager_state = vars(manager) + manager_state = cast(dict[str, Any], vars(manager)) if _AFD_FULLGRAPH_HOOK_MARKER in manager_state: raise RuntimeError("AFD ACL graph replay hook is already active") diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index aa856f37..8bb4c6ee 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -4,7 +4,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Protocol, cast import torch from vllm.compilation.monitor import set_cudagraph_capturing_enabled @@ -57,6 +57,14 @@ logger = init_logger(__name__) +class _MoEConfig(Protocol): + """DeepSeek-family fields required by the FFN layer splitter.""" + + n_routed_experts: int | None + first_k_dense_replace: int + moe_layer_freq: int + + class AFDNPUFFNModelRunner(NPUModelRunner): """Connector-driven NPU FFN runner for AFD execution.""" @@ -177,6 +185,7 @@ def execute_model( graph_key, len(acl_graphs), ) + assert graph_info is not None graph_info["graph"].replay() return None if run_mode in (AFDGraphRunMode.WARMUP, AFDGraphRunMode.CAPTURE): @@ -495,16 +504,15 @@ def _ffn_layer_indices(runner: AFDNPUFFNModelRunner) -> range | list[int]: return [ layer_idx for layer_idx in range(num_layers) - if _is_moe_layer(hf_config, layer_idx) + if _is_moe_layer(cast(_MoEConfig, hf_config), layer_idx) ] -def _is_moe_layer(hf_config: object, layer_idx: int) -> bool: - moe_layer_freq = getattr(hf_config, "moe_layer_freq", 1) +def _is_moe_layer(hf_config: _MoEConfig, layer_idx: int) -> bool: return ( hf_config.n_routed_experts is not None and layer_idx >= hf_config.first_k_dense_replace - and layer_idx % moe_layer_freq == 0 + and layer_idx % hf_config.moe_layer_freq == 0 ) @@ -566,13 +574,15 @@ def _ffn_token_count_for_rank( num_tokens_across_dp: torch.Tensor, ) -> int: values = _to_int_list(num_tokens_across_dp) - role_rank = int(connector.topology.role_rank) + role_rank = int(connector.role_rank) if role_rank >= len(values): return max(1, values[0] if values else 1) return max(1, int(values[role_rank])) -def _to_int_list(value: object) -> list[int]: +def _to_int_list( + value: torch.Tensor | list[int] | tuple[int, ...] | int | float | None, +) -> list[int]: if value is None: return [] if isinstance(value, (int, float)): @@ -605,7 +615,7 @@ def _to_dp_level_token_counts( return num_tokens_across_dp[indices].contiguous() -def _use_npu_aclgraph(vllm_config: VllmConfig, runner: object) -> bool: +def _use_npu_aclgraph(vllm_config: VllmConfig, runner: NPUModelRunner) -> bool: inherited = bool(runner.use_aclgraph) if bool(vllm_config.model_config.enforce_eager): return False diff --git a/tests/unit/compat/patches/test_config_validation.py b/tests/unit/compat/patches/test_config_validation.py index b18335d9..b851a041 100644 --- a/tests/unit/compat/patches/test_config_validation.py +++ b/tests/unit/compat/patches/test_config_validation.py @@ -1,3 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# mypy: disable-error-code=attr-defined +# This test deliberately builds incomplete ModuleType mocks for vLLM modules. + from __future__ import annotations import importlib @@ -107,7 +112,7 @@ def _set_fake_platform(*, is_cuda, device_type): def _install_fake_npu_config(monkeypatch): arg_utils_module, config_module = _install_fake_vllm_config(monkeypatch) - events = [] + events: list[tuple[str, object] | tuple[str, object, object]] = [] class FakeParallelConfig: def __init__( @@ -537,7 +542,7 @@ def test_config_validation_finalizes_async_attention_patch_after_ascend(monkeypa _set_fake_platform(is_cuda=False, device_type="npu") config_patch_calls = [] - engine_patch_configs = [] + engine_patch_configs: list[object] = [] monkeypatch.setattr( npu_compat, "apply_afd_ascend_config_patch_if_needed", diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 4920b866..c2f05747 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """Unit tests for AFD CUDA GPU ModelRunnerV2 support.""" from __future__ import annotations @@ -85,7 +87,7 @@ def stop(self): class _RunnerRecorder: - instances = [] + instances: list[tuple[object, object]] = [] def __init__(self, vllm_config, device): type(self).instances.append((vllm_config, device)) @@ -383,6 +385,14 @@ def test_v2_fullgraph_replay_hook_rejects_same_manager_reentry(): assert runner.cudagraph_manager.__dict__["run_fullgraph"] is original +def _set_single_dp_rank(config): + config.parallel_config.data_parallel_size = 1 + config.additional_config["afd"].update( + num_attention_ranks=1, + num_ffn_ranks=1, + ) + + @pytest.mark.parametrize( ( "role", @@ -573,16 +583,7 @@ def test_npu_v2_validator_allows_full_decode_only_dbo_dp2(): @pytest.mark.parametrize( ("mutation", "message"), [ - ( - lambda c: ( - setattr(c.parallel_config, "data_parallel_size", 1), - c.additional_config["afd"].update( - num_attention_ranks=1, - num_ffn_ranks=1, - ), - ), - "DP > 1", - ), + (_set_single_dp_rank, "DP > 1"), (lambda c: setattr(c.parallel_config, "num_ubatches", 4), "two ubatches"), ( lambda c: setattr(c.model_config, "enforce_eager", False), @@ -1010,7 +1011,7 @@ def test_v2_dp2_repeated_fullgraph_replay_sends_local_real_and_padded_tokens( descriptor = SimpleNamespace(num_tokens=8) payloads = [] metadata_seen = [] - replay_returns = [] + replay_returns: list[str] = [] class ReplayConnector(_RecordingConnector): def send_dp_metadata_list(self, payload): @@ -1143,7 +1144,7 @@ def test_v2_graph_miss_uses_provider_once_without_replay_control(monkeypatch): runner.vllm_config.compilation_config.cudagraph_mode = ( CUDAGraphMode.FULL_DECODE_ONLY ) - replay_calls = [] + replay_calls: list[object] = [] context = ForwardContext( no_compile_layers={}, attn_metadata={}, @@ -1631,7 +1632,7 @@ def close_connector(): if failure == "connector": raise RuntimeError("connector failed") - runner.connector.close = close_connector + monkeypatch.setattr(runner.connector, "close", close_connector) monkeypatch.setattr( "afd_plugin.v1.worker.attention_model_runner_v2.stop_afd_gpu_profiler", stop_profiler, diff --git a/tests/unit/v1/worker/test_npu_mla_graph.py b/tests/unit/v1/worker/test_npu_mla_graph.py index 3eeaa43a..049a7b5d 100644 --- a/tests/unit/v1/worker/test_npu_mla_graph.py +++ b/tests/unit/v1/worker/test_npu_mla_graph.py @@ -1,3 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# mypy: disable-error-code=attr-defined +# This test deliberately builds incomplete ModuleType mocks for NPU imports. + from __future__ import annotations import importlib @@ -926,7 +931,7 @@ def test_mla_graph_replay_rejects_missing_capture_registry(monkeypatch): mla_full_graph_enabled=True, ) wrapper.full_graph_params_updater = lambda *_args: None - replay_calls = [] + replay_calls: list[str] = [] graph_metadata = wrapper_module.AscendNPUGraphMetaData( aclgraph=SimpleNamespace( replay=lambda: replay_calls.append("replay"), @@ -953,6 +958,7 @@ def test_mla_graph_replay_rejects_missing_updater_before_replay(monkeypatch): ) wrapper.full_graph_params_updater = None workspace = object() + replay_calls: list[str] = [] graph_metadata = wrapper_module.AscendNPUGraphMetaData( aclgraph=SimpleNamespace(replay=lambda: replay_calls.append("replay")), ubatch_metadata=[], @@ -965,8 +971,6 @@ def test_mla_graph_replay_rejects_missing_updater_before_replay(monkeypatch): 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) diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 91d7fd1a..263bc92e 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +# mypy: disable-error-code=attr-defined +# This test deliberately builds incomplete ModuleType mocks for NPU imports. from __future__ import annotations From 884ac6444d184f1a01a6021fdc8ff0f1d70c26e8 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Fri, 28 Aug 2026 16:29:59 +0800 Subject: [PATCH 6/7] fix(npu): prevent MRV2 DBO thread initialization deadlock Signed-off-by: lirx-pd <616517220@qq.com> --- afd_plugin/v1/worker/npu/ubatch_runner_v2.py | 46 +++++++++++++++----- afd_plugin/v1/worker/npu/ubatching.py | 24 +++++++--- 2 files changed, 53 insertions(+), 17 deletions(-) diff --git a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py index 04985f6b..96cc213f 100644 --- a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -286,32 +286,56 @@ def run_stage( outputs[context.id] = model(**inputs) except BaseException as error: # noqa: BLE001 errors[context.id] = error + # ### PATCH START: Failed worker release + self.ready_barrier.abort() + # ### PATCH END: Failed worker release stack = ExitStack() stack.enter_context(override_forward_context(None)) - threads = [] - for context, stage in zip(contexts, ubatch_state.slices, strict=True): - thread = threading.Thread( - target=run_stage, - args=(context, slice_model_inputs(model_inputs, stage.token_slice)), - ) - threads.append(thread) - thread.start() - self.ready_barrier.wait() + threads: list[threading.Thread] = [] - def finish() -> AscendModelOutput: + # ### PATCH START: Failed execution cleanup + def close_execution() -> None: try: - contexts[0].cpu_wait_event.set() for thread in threads: thread.join() finally: stack.close() + if self.ready_barrier.broken: + self.ready_barrier.reset() + def raise_stage_error() -> None: if errors: failed_stage = min(errors) raise RuntimeError( f"AFD NPU microbatch {failed_stage} failed", ) from errors[failed_stage] + + try: + for context, stage in zip(contexts, ubatch_state.slices, strict=True): + thread = threading.Thread( + target=run_stage, + args=( + context, + slice_model_inputs(model_inputs, stage.token_slice), + ), + ) + thread.start() + threads.append(thread) + self.ready_barrier.wait() + except BaseException: # noqa: BLE001 + self.ready_barrier.abort() + close_execution() + raise_stage_error() + raise + # ### PATCH END: Failed execution cleanup + + def finish() -> AscendModelOutput: + # ### PATCH START: Failed execution cleanup + contexts[0].cpu_wait_event.set() + close_execution() + raise_stage_error() + # ### PATCH END: Failed execution cleanup ordered_outputs = [outputs[index] for index in range(self.num_ubatches)] # ### PATCH START: Ascend FlashComm output gathering if forward_contexts[0].additional_kwargs["flash_comm_v1_enabled"]: diff --git a/afd_plugin/v1/worker/npu/ubatching.py b/afd_plugin/v1/worker/npu/ubatching.py index fad0de28..13fed783 100644 --- a/afd_plugin/v1/worker/npu/ubatching.py +++ b/afd_plugin/v1/worker/npu/ubatching.py @@ -37,19 +37,31 @@ def __init__( def __enter__(self): _THREAD_ID_TO_CONTEXT[threading.get_ident()] = self.id _CURRENT_CONTEXTS[self.id] = self - self.ready_barrier.wait() - self.cpu_wait_event.wait() - self.cpu_wait_event.clear() - self._restore_context() - self.update_stream(self.compute_stream) + # ### PATCH START: Context failure cleanup + try: + self.ready_barrier.wait() + self.cpu_wait_event.wait() + self.cpu_wait_event.clear() + self._restore_context() + self.update_stream(self.compute_stream) + except BaseException: # noqa: BLE001 + self._release() + raise + # ### PATCH END: Context failure cleanup return self + # ### PATCH START: Shared context cleanup def __exit__(self, exc_type, exc_val, exc_tb): + self._release() + return False + + def _release(self): _CURRENT_CONTEXTS[self.id] = None del _THREAD_ID_TO_CONTEXT[threading.get_ident()] self.cpu_signal_event.set() self.cpu_wait_event.clear() - return False + + # ### PATCH END: Shared context cleanup def _restore_context(self): forward_context._forward_context = self.forward_context From d209945cc5bf375e39d6c9c46a813f1169ac4412 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Sun, 30 Aug 2026 18:12:15 +0800 Subject: [PATCH 7/7] fix: address review findings Signed-off-by: lirx-pd <616517220@qq.com> --- afd_plugin/compat/backports/__init__.py | 3 + .../worker/npu/attention_model_runner_v2.py | 33 ++-- afd_plugin/v1/worker/npu/ubatch_runner_v2.py | 21 ++- afd_plugin/v1/worker/npu/ubatching.py | 20 +- afd_plugin/validation.py | 16 +- tests/unit/package/test_package.py | 53 ++++++ tests/unit/v1/worker/test_model_runner_v2.py | 28 +++ tests/unit/v1/worker/test_npu_runtime.py | 176 ++++++++++++++++++ 8 files changed, 319 insertions(+), 31 deletions(-) create mode 100644 afd_plugin/compat/backports/__init__.py diff --git a/afd_plugin/compat/backports/__init__.py b/afd_plugin/compat/backports/__init__.py new file mode 100644 index 00000000..9cbec707 --- /dev/null +++ b/afd_plugin/compat/backports/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Compatibility backports for the target vLLM runtime.""" 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 f5802f69..0da5dc03 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -95,17 +95,17 @@ def run_fullgraph( self: v2_cudagraph_utils.ModelCudaGraphManager, desc: v2_cudagraph_utils.BatchExecutionDescriptor, ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]] | IntermediateTensors: - # ### PATCH START: bypass native replay metadata for AFD DBO graphs. - if isinstance(desc, AFDBatchExecutionDescriptor): - return original_run_fullgraph(desc) - # ### PATCH END: bypass native replay metadata for AFD DBO graphs. - # ### PATCH START: publish one AFD pre-replay payload. + # ### PATCH START: publish AFD replay control metadata. previous_is_graph_replaying = getattr( runner, "_afd_is_graph_replaying", False, ) try: + runner._afd_is_graph_replaying = True + if isinstance(desc, AFDBatchExecutionDescriptor): + return original_run_fullgraph(desc) + padded_tokens = int(desc.num_tokens) metadata = runner.build_afd_metadata(None, real_tokens) metadata.tokens_lens = [padded_tokens] @@ -113,7 +113,6 @@ def run_fullgraph( runner._afd_suppress_metadata_send = True 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, @@ -121,7 +120,7 @@ def run_fullgraph( result = original_run_fullgraph(desc) finally: runner._afd_is_graph_replaying = previous_is_graph_replaying - # ### PATCH END: publish one AFD pre-replay payload. + # ### PATCH END: publish AFD replay control metadata. return result try: @@ -472,16 +471,20 @@ def shutdown(self) -> None: finally: try: # ### PATCH START: AFD MRV2 DBO graph cleanup - if isinstance( - self.cudagraph_manager, - AFDModelAclGraphManagerV2, - ): - self.cudagraph_manager.clear_afd_graphs() + manager = self.cudagraph_manager + if isinstance(manager, AFDModelAclGraphManagerV2): + try: + manager.clear_afd_graphs() + finally: + del manager.ubatch_runner + del self.ubatch_runner # ### PATCH END: AFD MRV2 DBO graph cleanup - super().shutdown() finally: - self._afd_pending_metadata = None - self.connector.close() + try: + super().shutdown() + finally: + self._afd_pending_metadata = None + self.connector.close() # ### PATCH END: guarantee profiler/native/connector cleanup. diff --git a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py index 96cc213f..a5287395 100644 --- a/afd_plugin/v1/worker/npu/ubatch_runner_v2.py +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -269,10 +269,20 @@ def begin_capturable_run( forward_contexts, self.ready_barrier, ) + cancellation_event = contexts[0].cancellation_event # ### PATCH END: Ascend microbatch contexts outputs: dict[int, AscendModelOutput] = {} errors: dict[int, BaseException] = {} + # ### PATCH START: Failed worker release + def cancel_execution() -> None: + cancellation_event.set() + # Stages can block on different ring events, so cancellation must + # wake every event before joining any worker. + for context in contexts: + context.cpu_wait_event.set() + self.ready_barrier.abort() + @torch.inference_mode() def run_stage( context: AscendUBatchContext, @@ -285,10 +295,11 @@ def run_stage( with context: outputs[context.id] = model(**inputs) except BaseException as error: # noqa: BLE001 - errors[context.id] = error - # ### PATCH START: Failed worker release - self.ready_barrier.abort() - # ### PATCH END: Failed worker release + if not cancellation_event.is_set(): + errors[context.id] = error + cancel_execution() + + # ### PATCH END: Failed worker release stack = ExitStack() stack.enter_context(override_forward_context(None)) @@ -324,7 +335,7 @@ def raise_stage_error() -> None: threads.append(thread) self.ready_barrier.wait() except BaseException: # noqa: BLE001 - self.ready_barrier.abort() + cancel_execution() close_execution() raise_stage_error() raise diff --git a/afd_plugin/v1/worker/npu/ubatching.py b/afd_plugin/v1/worker/npu/ubatching.py index 13fed783..8756de1a 100644 --- a/afd_plugin/v1/worker/npu/ubatching.py +++ b/afd_plugin/v1/worker/npu/ubatching.py @@ -25,6 +25,7 @@ def __init__( ready_barrier: threading.Barrier, cpu_wait_event: threading.Event, cpu_signal_event: threading.Event, + cancellation_event: threading.Event, ): self.id = id self.compute_stream = compute_stream @@ -32,6 +33,7 @@ def __init__( self.ready_barrier = ready_barrier self.cpu_wait_event = cpu_wait_event self.cpu_signal_event = cpu_signal_event + self.cancellation_event = cancellation_event self.current_stream = compute_stream def __enter__(self): @@ -40,8 +42,7 @@ def __enter__(self): # ### PATCH START: Context failure cleanup try: self.ready_barrier.wait() - self.cpu_wait_event.wait() - self.cpu_wait_event.clear() + self._wait_for_turn() self._restore_context() self.update_stream(self.compute_stream) except BaseException: # noqa: BLE001 @@ -71,16 +72,23 @@ def update_stream(self, stream: torch.npu.Stream): if dbo_current_stream() != stream: dbo_set_stream(stream) + # ### PATCH START: Shared execution cancellation + def _wait_for_turn(self) -> None: + self.cpu_wait_event.wait() + self.cpu_wait_event.clear() + if self.cancellation_event.is_set(): + raise RuntimeError("AFD NPU microbatch execution cancelled") + def _cpu_yield(self): assert forward_context._forward_context == self.forward_context assert dbo_current_stream() == self.current_stream assert not self.cpu_wait_event.is_set() self.cpu_signal_event.set() - self.cpu_wait_event.wait() - self.cpu_wait_event.clear() + self._wait_for_turn() self._restore_context() self.update_stream(self.current_stream) + # ### PATCH END: Shared execution cancellation def yield_(self): self.current_stream = dbo_current_stream() @@ -135,6 +143,9 @@ def make_ubatch_contexts( if len(_CURRENT_CONTEXTS) < num_micro_batches: _CURRENT_CONTEXTS.extend([None] * (num_micro_batches - len(_CURRENT_CONTEXTS))) + # ### PATCH START: Shared execution cancellation + cancellation_event = threading.Event() + # ### PATCH END: Shared execution cancellation cpu_events = [threading.Event() for _ in range(num_micro_batches)] return [ AscendUBatchContext( @@ -144,6 +155,7 @@ def make_ubatch_contexts( ready_barrier=ready_barrier, cpu_wait_event=cpu_events[i], cpu_signal_event=cpu_events[(i + 1) % num_micro_batches], + cancellation_event=cancellation_event, ) for i in range(num_micro_batches) ] diff --git a/afd_plugin/validation.py b/afd_plugin/validation.py index deade7de..a7a42509 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -154,13 +154,15 @@ def validate_npu_model_runner_v2_config( raise RuntimeError( "AFD NPU ModelRunnerV2 DBO requires DP > 1 and exactly two ubatches", ) - if ( - not vllm_config.model_config.enforce_eager - and cudagraph_mode_name(vllm_config) != "FULL_DECODE_ONLY" - ): - raise RuntimeError( - "AFD NPU ModelRunnerV2 DBO ACL graph requires FULL_DECODE_ONLY", - ) + if not vllm_config.model_config.enforce_eager: + if cudagraph_mode_name(vllm_config) != "FULL_DECODE_ONLY": + raise RuntimeError( + "AFD NPU ModelRunnerV2 DBO ACL graph requires FULL_DECODE_ONLY", + ) + if not vllm_config.model_config.use_mla: + raise RuntimeError( + "AFD NPU ModelRunnerV2 DBO ACL graph requires MLA", + ) if ( vllm_config.speculative_config is not None or vllm_config.lora_config is not None diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index 035131d1..4b407a35 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -1,6 +1,12 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project + from __future__ import annotations import importlib.metadata +import os +import subprocess +import sys from pathlib import Path from types import SimpleNamespace @@ -15,6 +21,53 @@ def test_package_import_is_cpu_safe(): assert afd_plugin.AFDConfig().connector == "P2pNcclAFDConnector" +def test_wheel_contains_importable_backport_package(tmp_path): + root = Path(__file__).resolve().parents[3] + env = os.environ.copy() + env["AFD_BUILD_ASCEND_OPS"] = "0" + subprocess.run( + [ + sys.executable, + "-m", + "pip", + "wheel", + "--no-build-isolation", + "--no-deps", + "--wheel-dir", + str(tmp_path), + str(root), + ], + check=True, + capture_output=True, + env=env, + text=True, + ) + wheel = next(tmp_path.glob("vllm_afd_plugin-*.whl")) + subprocess.run( + [ + sys.executable, + "-I", + "-c", + """ +import importlib.util +import sys + +sys.path.insert(0, sys.argv[1]) +import afd_plugin.compat.backports as backports + +assert backports.__file__.startswith(sys.argv[1]) +assert importlib.util.find_spec( + "afd_plugin.compat.backports.vllm_v026_mrv2_dbo" +) is not None +""", + str(wheel), + ], + check=True, + capture_output=True, + text=True, + ) + + def test_register_afd_is_idempotent(): afd_plugin.register_afd() afd_plugin.register_afd() diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index c2f05747..19227a52 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -111,6 +111,7 @@ def _v2_config( tensor_parallel_size: int = 1, prefill_context_parallel_size: int = 1, enforce_eager: bool = True, + use_mla: bool = True, cudagraph_mode: CUDAGraphMode = CUDAGraphMode.NONE, cudagraph_capture_sizes: list[int] | None = None, max_cudagraph_capture_size: int = 0, @@ -127,6 +128,7 @@ def _v2_config( is_encoder_decoder=False, is_multimodal_model=False, enforce_eager=enforce_eager, + use_mla=use_mla, enable_prompt_embeds=False, enable_return_routed_experts=False, quantization=None, @@ -545,6 +547,7 @@ def test_npu_v2_validator_allows_eager_dbo_dp2(): num_attention_ranks=2, num_ffn_ranks=2, data_parallel_size=2, + use_mla=False, ) config.additional_config["afd"]["connector"] = "CAMP2pAFDConnector" config.parallel_config.enable_dbo = True @@ -580,6 +583,31 @@ def test_npu_v2_validator_allows_full_decode_only_dbo_dp2(): ) +def test_npu_v2_validator_rejects_non_mla_dbo_acl_graph(): + config = _v2_config( + architecture="Qwen3MoeForCausalLM", + num_attention_ranks=2, + num_ffn_ranks=2, + data_parallel_size=2, + enforce_eager=False, + use_mla=False, + cudagraph_mode=CUDAGraphMode.FULL_DECODE_ONLY, + cudagraph_capture_sizes=[TEST_CUDAGRAPH_CAPTURE_SIZE], + max_cudagraph_capture_size=TEST_CUDAGRAPH_CAPTURE_SIZE, + ) + config.additional_config["afd"]["connector"] = "CAMP2pAFDConnector" + config.parallel_config.enable_dbo = True + config.parallel_config.use_ubatching = True + config.parallel_config.num_ubatches = 2 + + with pytest.raises(RuntimeError, match="requires MLA"): + validate_npu_model_runner_v2_config( + config, + expected_role="attention", + device_type="npu", + ) + + @pytest.mark.parametrize( ("mutation", "message"), [ diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index 263bc92e..ec261d1c 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -33,6 +33,8 @@ AFDTransferState, ) +STAGE_FAILURE_JOIN_TIMEOUT_SECONDS = 5 + @contextmanager def _temporarily_reimport_module(module_name: str) -> Iterator[ModuleType]: @@ -95,6 +97,7 @@ class _RecordingConnector: def __init__(self): self.dp_metadata_updates = [] self.sent_dp_metadata_lists = [] + self.last_sent_payload: AFDControlPayload | None = None # The runners reach the control plane through connector.control_plane; # the fake serves as both. self.control_plane: _RecordingConnector | None = self @@ -111,6 +114,7 @@ def update_state_from_dp_metadata(self, payload): def send_dp_metadata_list(self, payload): assert isinstance(payload, AFDControlPayload) + self.last_sent_payload = payload self.sent_dp_metadata_lists.append( ( payload.dp_metadata_list, @@ -294,6 +298,178 @@ def _require_npu_runtime(): pytest.importorskip("torch_npu", reason="NPU runtime tests require torch-npu") +def test_npu_v2_dbo_fullgraph_replay_reaches_ffn_graph(): + _require_npu_runtime() + from vllm.config import CUDAGraphMode + + from afd_plugin.compat.backports.vllm_v026_mrv2_dbo import ( + AFDBatchExecutionDescriptor, + ) + from afd_plugin.v1.worker.npu import attention_model_runner_v2 + + runner = object.__new__(attention_model_runner_v2.AFDNPUAttentionModelRunnerV2) + runner.vllm_config = _vllm_config() + runner.connector = _RecordingConnector() + runner._is_warmup = False + runner._afd_is_graph_capturing = False + runner._afd_is_graph_replaying = False + runner._afd_is_profile = False + + ffn_runner = _new_ffn_runner() + ffn_runner.vllm_config = _vllm_config(role="ffn") + ffn_runner.connector = _FakeFFNConnector() + ffn_runner.max_num_tokens = 8 + ffn_runner.use_aclgraph = True + graph = _FakeGraph() + + class Manager: + def run_fullgraph(self, desc): + runner.send_dp_metadata(_FakeDPMetadata([desc.num_tokens]), None) + payload = runner.connector.last_sent_payload + ffn_runner._acl_graphs = { + ffn_runner._make_graph_key(payload.dp_metadata_list): {"graph": graph}, + } + ffn_runner.execute_model( + dp_metadata_list=payload.dp_metadata_list, + is_graph_replaying=payload.is_graph_replaying, + ) + return desc.num_tokens + + runner.cudagraph_manager = Manager() + descriptor = AFDBatchExecutionDescriptor( + cg_mode=CUDAGraphMode.FULL, + num_tokens=8, + num_reqs=4, + num_ubatches=2, + ) + + with attention_model_runner_v2._use_afd_fullgraph_replay_hook(runner, 6): + assert runner.cudagraph_manager.run_fullgraph(descriptor) == 8 + + payload = runner.connector.last_sent_payload + assert payload is not None + assert payload.is_graph_replaying is True + assert graph.replay_count == 1 + assert runner._afd_is_graph_replaying is False + + +def test_npu_v2_shutdown_releases_ubatch_before_native(monkeypatch): + _require_npu_runtime() + from afd_plugin.v1.worker.npu import attention_model_runner_v2 + from afd_plugin.v1.worker.npu.aclgraph_manager_v2 import ( + AFDModelAclGraphManagerV2, + ) + + events = [] + runner = object.__new__(attention_model_runner_v2.AFDNPUAttentionModelRunnerV2) + manager = object.__new__(AFDModelAclGraphManagerV2) + ubatch_runner = SimpleNamespace(model_state=object()) + manager.ubatch_runner = ubatch_runner + manager.clear_afd_graphs = MethodType( + lambda _self: events.append("clear_graphs"), + manager, + ) + runner.cudagraph_manager = manager + runner.ubatch_runner = ubatch_runner + runner.prof = object() + runner._afd_pending_metadata = object() + runner.connector = SimpleNamespace(close=lambda: events.append("close")) + + monkeypatch.setattr( + attention_model_runner_v2, + "stop_afd_npu_profiler", + lambda _prof: events.append("stop_profiler"), + ) + + def native_shutdown(_self): + assert not hasattr(runner, "ubatch_runner") + assert not hasattr(manager, "ubatch_runner") + events.append("native_shutdown") + + monkeypatch.setattr( + attention_model_runner_v2.NPUModelRunnerV2, + "shutdown", + native_shutdown, + ) + + runner.shutdown() + + assert events == ["stop_profiler", "clear_graphs", "native_shutdown", "close"] + assert runner._afd_pending_metadata is None + + +def test_npu_v2_ubatch_failure_cancels_sibling_event_wait(monkeypatch): + _require_npu_runtime() + import torch + from vllm.v1.worker.ubatch_utils import UBatchSlice + + from afd_plugin.v1.worker.npu import ubatch_runner_v2 + from afd_plugin.v1.worker.npu.ubatch_runner_v2 import ( + AFDAscendUBatchRunnerV2, + AFDAscendUBatchState, + ) + from afd_plugin.v1.worker.npu.ubatching import dbo_yield + + stream = object() + monkeypatch.setattr(torch.npu, "current_stream", lambda: stream) + monkeypatch.setattr(torch.npu, "set_device", lambda _device: None) + monkeypatch.setattr(torch.npu, "set_stream", lambda _stream: None) + + thread_class = threading.Thread + + def daemon_thread(*args, **kwargs): + kwargs["daemon"] = True + return thread_class(*args, **kwargs) + + monkeypatch.setattr( + ubatch_runner_v2, + "threading", + SimpleNamespace(Thread=daemon_thread), + ) + + runner = object.__new__(AFDAscendUBatchRunnerV2) + runner.num_ubatches = 2 + runner.device = torch.device("npu") + runner.ready_barrier = threading.Barrier(3) + state = AFDAscendUBatchState( + slices=[ + UBatchSlice(slice(0, 1), slice(0, 1)), + UBatchSlice(slice(1, 2), slice(1, 2)), + ], + attn_metadata=[], + forward_contexts=[SimpleNamespace(), SimpleNamespace()], + real_token_counts=[1, 1], + ) + model_inputs = { + "input_ids": torch.tensor([0, 1]), + "positions": torch.tensor([0, 1]), + } + + def model(input_ids, **_kwargs): + if input_ids[0].item() == 1: + raise ValueError("stage failure") + dbo_yield() + dbo_yield() + return input_ids + + execution_errors: list[BaseException] = [] + + def execute() -> None: + try: + runner.run(model, model_inputs, state) + except BaseException as error: # noqa: BLE001 + execution_errors.append(error) + + execution_thread = thread_class(target=execute, daemon=True) + execution_thread.start() + execution_thread.join(STAGE_FAILURE_JOIN_TIMEOUT_SECONDS) + + assert not execution_thread.is_alive() + assert len(execution_errors) == 1 + assert str(execution_errors[0]) == "AFD NPU microbatch 1 failed" + assert isinstance(execution_errors[0].__cause__, ValueError) + + def test_npu_v1_runner_signatures_match_pinned_ascend(): _require_npu_runtime() from vllm_ascend.worker.model_runner_v1 import NPUModelRunner