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/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..c229d11b --- /dev/null +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/__init__.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Temporary vLLM 0.26 ModelRunnerV2 DBO backport.""" + +from .runtime import ( + AFDBatchExecutionDescriptor, + create_ubatch_slices, + dispatch_afd_dbo_and_sync_dp, + prepare_attn_for_ubatch, + slice_input_batch, + slice_model_inputs, + use_two_metadata_builders, +) + +__all__ = [ + "AFDBatchExecutionDescriptor", + "create_ubatch_slices", + "dispatch_afd_dbo_and_sync_dp", + "prepare_attn_for_ubatch", + "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..d8da7424 --- /dev/null +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/execute.py @@ -0,0 +1,251 @@ +# 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. + +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 + +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.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 + + +# 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, + 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 or FULL 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, + ) + # ### 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, + 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, + 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) + + num_ubatches = ( + batch_desc.num_ubatches + 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) + 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: + # ### 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: + block_tables = None + slot_mappings = None + + 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( + slot_mappings, + runner.kv_cache_config, + ) + attn_metadata = runner.model_state.prepare_attn( + input_batch, + batch_desc.cg_mode, + 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, + ) + + # ### 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( + 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, + 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, + ) + # ### 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) + model_output = runner.cudagraph_manager.run_fullgraph(batch_desc) + 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=batch_desc.cg_mode, + 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..60a3b08f --- /dev/null +++ b/afd_plugin/compat/backports/vllm_v026_mrv2_dbo/runtime.py @@ -0,0 +1,456 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""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 +``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, replace +from typing import TYPE_CHECKING, 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, + is_last_ubatch_empty, + maybe_create_ubatch_slices, +) +from vllm.v1.worker.utils import AttentionGroup + +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.""" + + num_ubatches: int = 1 + + +# ### 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, + num_tokens: int, + uniform_token_count: int | None, + dp_size: int, + dp_rank: int, + parallel_config: ParallelConfig, + decode_query_len: int, + allow_ubatching: bool, + cudagraph_manager: AFDModelAclGraphManagerV2 | None = None, + need_eager: bool = False, +) -> tuple[BatchExecutionDescriptor, torch.Tensor | None]: + """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. + """ + + # ### PATCH START: AFD v0.26 graph dispatch adapter + def dispatch( + tokens: int, + ubatches: int, + uniform_tokens: int | None = uniform_token_count, + force_eager: bool = False, + ) -> 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, + ) + 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, + ) + + assert cudagraph_manager is not None + 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] + 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()) + # ### 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 + or not torch.all(tensor[1] == synced_uniform_token_count).item() + ): + 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) + and check_ubatch_thresholds( + parallel_config, + int(num_tokens_across_dp.min().item()), + uniform_decode=uniform_decode, + ) + ) + if not should_ubatch: + 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 + + # ### 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( + 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, + 1, + synced_uniform_token_count, + 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 + + +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 + }, + ) + # ### PATCH START: Ascend auxiliary hidden-state output + 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 + # ### PATCH END: Ascend auxiliary hidden-state output + 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 + + # 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( + self, + vllm_config, + device, + kernel_block_size: int | None = None, + num_metadata_builders: int = 1, + ): + # ### PATCH START: AFD two metadata builders + del num_metadata_builders + original( + 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 + yield + finally: + AttentionGroup.create_metadata_builders = original + + +# 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, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + 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, + cg_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=for_capture, + ) + + # ### PATCH START: AFD per-microbatch metadata builder + 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, + cg_mode, + block_tables, + slot_mappings, + attn_groups, + kv_cache_config, + for_capture=for_capture, + ) + finally: + for group in reversed(swapped): + group.metadata_builders[0], group.metadata_builders[ubatch_index] = ( + 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 78ec51c0..9e8e375e 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): + # ### 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 backend 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/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..4a707857 --- /dev/null +++ b/afd_plugin/compat/patches/npu/model_runner_v2_dbo.py @@ -0,0 +1,94 @@ +# 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. + +Upstream source: ``vllm_ascend/worker/v2/model_runner.py`` at commit +``d543ccee0``, function ``graph_manager_wrapper``. +""" + +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.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(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, + ): + # ### PATCH START: AFD MRV2 DBO graph manager + ubatch_runner = AFDAscendUBatchRunnerV2( + vllm_config, + device, + model_runner.model_state, + model_runner.attn_groups, + model_runner.kv_cache_config, + model_runner.max_num_reqs, + ) + model_runner.ubatch_runner = ubatch_runner + return AFDModelAclGraphManagerV2( + vllm_config, + device, + cudagraph_mode, + decode_query_len, + model_runner, + ubatch_runner, + lora_capture_cases=lora_capture_cases, + ) + # ### PATCH END: AFD MRV2 DBO graph manager + + 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/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/aclgraph_manager_v2.py b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py new file mode 100644 index 00000000..32762892 --- /dev/null +++ b/afd_plugin/v1/worker/npu/aclgraph_manager_v2.py @@ -0,0 +1,429 @@ +# 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. + +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 TYPE_CHECKING + +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 ( + GraphParams, + 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, +) + +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: torch.npu.NPUGraph + output: AscendModelOutput + graph_params: tuple[GraphParams, GraphParams] + workspace: torch.Tensor + + +@dataclass +class _PendingReplay: + descriptor: AFDBatchExecutionDescriptor + state: AFDAscendUBatchState + + +class AFDModelAclGraphManagerV2(ModelAclGraphManager): + """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: AFDNPUAttentionModelRunnerV2, + 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, + ) + # ### PATCH START: AFD DBO graph registry + 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 + # ### PATCH END: AFD DBO graph registry + + 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_ubatches( + self, + base: BatchExecutionDescriptor, + num_ubatches: int, + ) -> BatchExecutionDescriptor: + 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=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, + ) + + 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() + + # 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, + 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, + ) + # ### PATCH START: AFD DBO twin graph capture + 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 + # ### PATCH END: AFD DBO twin graph capture + + 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[GraphParams, GraphParams] | 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, + ) + + # 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: + 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 + 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( + 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, + ) + # ### PATCH END: AFD DBO FULL graph replay + 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 656f4f12..0da5dc03 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 @@ -22,12 +23,22 @@ 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, + 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 ( create_afd_npu_profiler, 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, @@ -39,6 +50,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" @@ -56,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") @@ -81,13 +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: 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] @@ -95,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, @@ -103,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: @@ -169,6 +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: @@ -211,6 +231,30 @@ 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.""" + + if not self.vllm_config.parallel_config.use_ubatching: + super().initialize_kv_cache(kv_cache_config) + return + + # ### 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) + # ### 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 # that lifecycle to AFD's control plane. @@ -381,6 +425,17 @@ 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: + 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, + ) + # ### PATCH END: AFD MRV2 DBO execute backport return super().execute_model( scheduler_output, intermediate_tensors, @@ -415,10 +470,21 @@ def shutdown(self) -> None: stop_afd_npu_profiler(self.prof) finally: try: - super().shutdown() + # ### PATCH START: AFD MRV2 DBO graph cleanup + 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 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/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 2b6e9281..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): @@ -402,12 +411,21 @@ 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 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 + # ### PATCH END: AFD MRV2 repeated capture replay logger.debug("AFD NPU FFN capturing ACL graph for key=%s", graph_key) graph = torch.npu.NPUGraph() @@ -486,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 ) @@ -557,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)): @@ -596,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/afd_plugin/v1/worker/npu/forward_context.py b/afd_plugin/v1/worker/npu/forward_context.py index e7ca5d2e..8da7d157 100644 --- a/afd_plugin/v1/worker/npu/forward_context.py +++ b/afd_plugin/v1/worker/npu/forward_context.py @@ -26,6 +26,15 @@ from vllm_ascend.compilation.acl_graph import GraphParams +# 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, @@ -41,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( @@ -75,77 +84,126 @@ 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 - ) - 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 - ) - 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 - ) - 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 + # ### 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: + 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, ) - new_forward_context.padded_length = padded_length - new_forward_context.pad_size = padded_length - num_tokens + 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 - new_forward_context.max_tokens_across_dp = max_tokens_across_dp + 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 + ) + 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 + ) - 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, + 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 ) - mc2_mask[:num_tokens] = True - mc2_mask[num_tokens:] = False - new_forward_context.mc2_mask = 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,), + 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 new file mode 100644 index 00000000..a5287395 --- /dev/null +++ b/afd_plugin/v1/worker/npu/ubatch_runner_v2.py @@ -0,0 +1,374 @@ +# 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. + +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 + +import threading +from collections.abc import Callable +from contextlib import ExitStack +from dataclasses import dataclass, replace +from typing import Any, cast + +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.model_states.interface import ModelState +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 ( + 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 ( + AscendModelOutput, + _all_gather_ubatch_output, +) +from afd_plugin.v1.worker.npu.ubatching import ( + AscendUBatchContext, + make_ubatch_contexts, +) + +AFD_NPU_MRV2_NUM_UBATCHES = 2 + + +@dataclass +class AFDAscendUBatchState: + """Prepared inputs for one eager or FULL graph DBO step.""" + + slices: UBatchSlices + attn_metadata: list[dict[str, Any]] + forward_contexts: list[ForwardContext] | None + real_token_counts: list[int] + + +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, + 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 = 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) + ] + self.seq_lens_buffers = [ + 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, + block_tables: tuple[torch.Tensor, ...], + slot_mappings: torch.Tensor, + slices: UBatchSlices, + parent_context: ForwardContext | None, + *, + cg_mode: CUDAGraphMode = CUDAGraphMode.NONE, + context_cg_mode: CUDAGraphMode | None = None, + for_capture: bool = False, + 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 = 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( + 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) + # ### 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, + stage_block_tables, + stage_slot_mappings, + self.attn_groups, + self.kv_cache_config, + stage_index, + 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: + continue + 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, + ) + # ### PATCH START: Ascend microbatch forward context + context = create_ascend_forward_context( + parent_context, + attn_metadata, + self.vllm_config, + slices, + ubatch_num=stage_index, + dp_metadata=dp_metadata, + 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, + self.kv_cache_config, + ) + # ### PATCH END: Ascend microbatch forward context + forward_contexts.append(context) + + return AFDAscendUBatchState( + slices=slices, + attn_metadata=attn_metadata_list, + forward_contexts=forward_contexts or None, + real_token_counts=real_token_counts, + ) + + @staticmethod + def _with_ascend_fields( + parent_batch: AscendInputBatch, + child_batch: AscendInputBatch, + stage: UBatchSlice, + ) -> AscendInputBatch: + 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, + ) + + # 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], + 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], + ubatch_state: AFDAscendUBatchState, + for_capture: bool = False, + ) -> Callable[[], AscendModelOutput]: + """Start both stages outside capture and return a one-shot finisher.""" + + 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() + ) + contexts = make_ubatch_contexts( + self.num_ubatches, + compute_stream, + 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, + 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 + 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)) + threads: list[threading.Thread] = [] + + # ### PATCH START: Failed execution cleanup + def close_execution() -> None: + try: + 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 + cancel_execution() + 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"]: + ordered_outputs = [ + _all_gather_ubatch_output( + output, + context.additional_kwargs["pad_size"], + ) + for output, context in zip( + ordered_outputs, + forward_contexts, + strict=True, + ) + ] + # ### PATCH END: Ascend FlashComm output gathering + return merge_ubatch_outputs(ordered_outputs) + + return finish + + +__all__ = [ + "AFD_NPU_MRV2_NUM_UBATCHES", + "AFDAscendUBatchRunnerV2", + "AFDAscendUBatchState", +] diff --git a/afd_plugin/v1/worker/npu/ubatching.py b/afd_plugin/v1/worker/npu/ubatching.py index fad0de28..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,24 +33,36 @@ 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): _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._wait_for_turn() + 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 @@ -59,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() @@ -123,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( @@ -132,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 9d5f6642..a7a42509 100644 --- a/afd_plugin/validation.py +++ b/afd_plugin/validation.py @@ -147,8 +147,33 @@ 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") + # ### 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: + raise RuntimeError( + "AFD NPU ModelRunnerV2 DBO requires DP > 1 and exactly two ubatches", + ) + 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 + 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", + ) + # ### 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 new file mode 100644 index 00000000..972ccbec --- /dev/null +++ b/tests/unit/compat/backports/test_vllm_v026_mrv2_dbo.py @@ -0,0 +1,414 @@ +# 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 + +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.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 + 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_full_graph_dispatch_signature(): + parameters = inspect.signature(dispatch_afd_dbo_and_sync_dp).parameters + + 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" + 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, **_kwargs): + 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_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(("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): + 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.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=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[: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(): + 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..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 @@ -58,6 +63,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 +99,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 @@ -105,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__( @@ -184,6 +191,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 +411,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", @@ -516,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/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 eed48913..19227a52 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)) @@ -109,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, @@ -125,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, @@ -383,6 +387,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", @@ -530,6 +542,107 @@ 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, + use_mla=False, + ) + 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", + ) + + +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", + ) + + +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"), + [ + (_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), + "FULL_DECODE_ONLY", + ), + ( + 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"), [ @@ -926,7 +1039,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): @@ -1059,7 +1172,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={}, @@ -1547,7 +1660,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 07f88561..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 @@ -557,6 +562,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, @@ -925,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"), @@ -952,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=[], @@ -964,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 3c4ecc4a..ec261d1c 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 @@ -31,6 +33,8 @@ AFDTransferState, ) +STAGE_FAILURE_JOIN_TIMEOUT_SECONDS = 5 + @contextmanager def _temporarily_reimport_module(module_name: str) -> Iterator[ModuleType]: @@ -93,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 @@ -109,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, @@ -292,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 @@ -1337,6 +1515,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 +1533,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() @@ -1670,6 +1906,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)