From 836af54c45f6fa94d847f7be08c7f4362050250a Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Sun, 8 Nov 2026 19:29:00 -0800 Subject: [PATCH 1/4] refactor: decouple shared FFN components from model adapters Signed-off-by: lirx-pd <616517220@qq.com> --- .../compat/patches/gpu_deepseek_v4_dbo.py | 2 +- .../ffn_model/deepseek_v4/npu/v4_weights.py | 12 +- afd_plugin/ffn_faas/platforms/npu/model.py | 10 +- .../models => }/forward_context.py | 6 +- afd_plugin/model_executor/ffn/__init__.py | 3 + .../async_cam_layout.py => ffn/async_cam.py} | 46 ++- afd_plugin/model_executor/ffn/compute.py | 231 ++++++++++++++ .../model_executor/ffn/expert_weights.py | 205 ++++++++++++ afd_plugin/model_executor/ffn/remote.py | 184 +++++++++++ afd_plugin/model_executor/models/__init__.py | 22 +- .../model_executor/models/_weight_utils.py | 32 ++ .../model_executor/models/deepseek_v2.py | 226 ++----------- .../model_executor/models/deepseek_v4.py | 41 +-- .../models/npu/async_cam_dbo.py | 49 --- .../npu/deepseek_v2_async_cam_forward.py | 8 +- .../models/npu/deepseek_v2_attention_gate.py | 2 +- .../model_executor/models/npu/deepseek_v4.py | 82 ++--- .../models/npu/deepseek_v4_dspark.py | 8 +- afd_plugin/model_executor/models/qwen3_5.py | 42 +-- afd_plugin/model_executor/models/qwen3_moe.py | 47 +-- .../model_executor/npu/async_cam_layered.py | 298 ------------------ .../model_executor/npu/async_cam_mega.py | 92 ------ .../v1/worker/attention_model_runner_v2.py | 2 +- .../v1/worker/npu/async_ffn_model_runner.py | 54 ++-- .../worker/npu/attention_metadata.py} | 18 +- .../worker/npu/attention_model_runner_v2.py | 2 +- afd_plugin/v1/worker/npu/ubatch_runner.py | 8 +- docs/design/module/attention_runtime.md | 9 +- .../module/compatibility_and_patches.md | 2 +- docs/design/module/connector_contracts.md | 2 +- docs/design/module/execution_platforms.md | 2 +- docs/design/module/ffn_runtime.md | 21 +- docs/design/module/model_integration.md | 72 ++++- .../deepseek_v4_flash/dspark_evidence.py | 16 +- .../dspark_ffn_graph_evidence.py | 6 +- .../deepseek_v4_flash/ffn_graph_evidence.py | 26 +- tests/npu/README.md | 6 +- tests/npu/async_cam_runtime.py | 8 +- tests/npu/test_async_cam_layered_w4a8.py | 8 +- tests/npu/test_async_cam_layered_w8a8.py | 8 +- tests/npu/test_mrv2_dbo.py | 2 +- tests/npu/test_mrv2_graph_inputs.py | 8 +- .../npu/test_async_cam_startup_wiring.py | 125 +++++++- tests/unit/fakes/async_ffn.py | 3 +- tests/unit/fakes/connector.py | 21 +- tests/unit/ffn_faas/test_model.py | 2 +- .../unit/model_executor/ffn}/__init__.py | 2 +- .../test_async_cam.py} | 67 +++- .../test_compute.py} | 56 ++-- .../npu/test_deepseek_v2_attention_gate.py | 6 +- .../models/npu/test_deepseek_v4.py | 6 +- ... => test_deepseek_v2_async_cam_forward.py} | 26 +- .../models/test_deepseek_v2_proxy.py | 37 +-- .../test_deepseek_v2_sequence_parallel.py | 6 +- .../models/test_deepseek_v4_cam_layout.py | 10 +- .../models/test_weight_utils.py | 44 +++ .../model_executor/test_async_cam_layout.py | 59 ---- tests/unit/test_forward_context.py | 32 ++ .../worker/npu/test_attention_metadata.py} | 20 +- .../unit/v1/worker/npu/test_ubatch_runner.py | 4 +- .../v1/worker/test_attention_model_runner.py | 4 +- .../v1/worker/test_npu_model_runner_v2.py | 4 +- 62 files changed, 1348 insertions(+), 1114 deletions(-) rename afd_plugin/{model_executor/models => }/forward_context.py (92%) create mode 100644 afd_plugin/model_executor/ffn/__init__.py rename afd_plugin/model_executor/{models/npu/async_cam_layout.py => ffn/async_cam.py} (82%) create mode 100644 afd_plugin/model_executor/ffn/compute.py create mode 100644 afd_plugin/model_executor/ffn/expert_weights.py create mode 100644 afd_plugin/model_executor/ffn/remote.py create mode 100644 afd_plugin/model_executor/models/_weight_utils.py delete mode 100644 afd_plugin/model_executor/models/npu/async_cam_dbo.py delete mode 100644 afd_plugin/model_executor/npu/async_cam_layered.py delete mode 100644 afd_plugin/model_executor/npu/async_cam_mega.py rename afd_plugin/{model_executor/models/npu/deepseek_attention_metadata.py => v1/worker/npu/attention_metadata.py} (93%) rename {afd_plugin/model_executor/npu => tests/unit/model_executor/ffn}/__init__.py (59%) rename tests/unit/model_executor/{test_async_cam_dbo.py => ffn/test_async_cam.py} (56%) rename tests/unit/model_executor/{test_async_cam_layered.py => ffn/test_compute.py} (93%) rename tests/unit/model_executor/models/{test_forward_context.py => test_deepseek_v2_async_cam_forward.py} (79%) create mode 100644 tests/unit/model_executor/models/test_weight_utils.py delete mode 100644 tests/unit/model_executor/test_async_cam_layout.py create mode 100644 tests/unit/test_forward_context.py rename tests/unit/{model_executor/test_deepseek_attention_metadata.py => v1/worker/npu/test_attention_metadata.py} (93%) diff --git a/afd_plugin/compat/patches/gpu_deepseek_v4_dbo.py b/afd_plugin/compat/patches/gpu_deepseek_v4_dbo.py index 35e3cbc0e..d6473b192 100644 --- a/afd_plugin/compat/patches/gpu_deepseek_v4_dbo.py +++ b/afd_plugin/compat/patches/gpu_deepseek_v4_dbo.py @@ -42,7 +42,7 @@ from vllm.v1.worker.utils import AttentionGroup from vllm.v1.worker.workspace import WorkspaceManager, current_workspace_manager -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.forward_context import get_afd_metadata_from_forward_context if TYPE_CHECKING: from afd_plugin.v1.worker.attention_model_runner_v2 import AFDAttentionModelRunnerV2 diff --git a/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py index 45325c4e1..67490d8f1 100644 --- a/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py +++ b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py @@ -24,9 +24,9 @@ V4CheckpointDescription, ) from afd_plugin.ffn_faas.platforms.npu.settings import validate_static_experts -from afd_plugin.model_executor.npu.async_cam_layered import ( - AsyncCAMLayeredExecutor, - LayeredFFNWeights, +from afd_plugin.model_executor.ffn.expert_weights import ( + QuantizedExpertWeights, + validate_expert_weights, ) @@ -158,7 +158,7 @@ def routed_experts(self, layer_index: int) -> RoutedExperts: """Keep the native V4 module path inside the model adapter.""" return self.model.layers[str(layer_index)].mlp.experts.routed_experts - def weight_views(self) -> tuple[LayeredFFNWeights, ...]: + def weight_views(self) -> tuple[QuantizedExpertWeights, ...]: from vllm_ascend.quantization.quant_type import QuantType weights = [] @@ -174,7 +174,7 @@ def weight_views(self) -> tuple[LayeredFFNWeights, ...]: scheme = routed.quant_method.quant_method if scheme.quant_type != QuantType.W4A8: raise ValueError("V4 FFN requires native W4A8 routed parameters") - view = LayeredFFNWeights( + view = QuantizedExpertWeights( layer_idx=layer.layer_idx, w13=routed.w13_weight, w2=routed.w2_weight, @@ -187,7 +187,7 @@ def weight_views(self) -> tuple[LayeredFFNWeights, ...]: routed_scaling_factor=1.0, quantization="W4A8", ) - AsyncCAMLayeredExecutor._validate_layer(view) + validate_expert_weights([view]) weights.append(view) if not weights or any(w.signature() != weights[0].signature() for w in weights): raise ValueError("V4 FFN requires homogeneous native W4A8 layers") diff --git a/afd_plugin/ffn_faas/platforms/npu/model.py b/afd_plugin/ffn_faas/platforms/npu/model.py index 8209002d6..303f0339d 100644 --- a/afd_plugin/ffn_faas/platforms/npu/model.py +++ b/afd_plugin/ffn_faas/platforms/npu/model.py @@ -22,7 +22,7 @@ from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts - from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights + from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights class RoutedWeightOwner(Protocol): @@ -32,7 +32,7 @@ def create_layer(self, layer_index: int) -> nn.Module: ... def routed_experts(self, layer_index: int) -> RoutedExperts: ... def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: ... def finish_layer(self, layer_index: int) -> None: ... - def weight_views(self) -> tuple[LayeredFFNWeights, ...]: ... + def weight_views(self) -> tuple[QuantizedExpertWeights, ...]: ... def parameters(self) -> Iterator[nn.Parameter]: ... @@ -65,7 +65,7 @@ def __init__( self.adapter, self.ffn_ranks, self.ffn_rank = adapter, ffn_ranks, ffn_rank self.local_experts = adapter.num_experts // ffn_ranks self._owner: RoutedWeightOwner | None = None - self._views: tuple[LayeredFFNWeights, ...] | None = None + self._views: tuple[QuantizedExpertWeights, ...] | None = None self._loading = False self._loaded = False self._parameter_bytes = adapter.estimated_parameter_bytes(self.local_experts) @@ -79,7 +79,7 @@ def parameter_bytes(self) -> int: return self._parameter_bytes @property - def native_view(self) -> tuple[LayeredFFNWeights, ...]: + def native_view(self) -> tuple[QuantizedExpertWeights, ...]: if not self._loaded or self._views is None: raise RuntimeError("FFN weights have not been loaded") return self._views @@ -177,7 +177,7 @@ def _load_native( def get_experts_layer_indices(self) -> tuple[int, ...]: return self.adapter.layer_indices - def get_async_cam_layered_weights(self) -> tuple[list[LayeredFFNWeights], str]: + def get_quantized_expert_weights(self) -> tuple[list[QuantizedExpertWeights], str]: return list(self.native_view), "" def compute_ffn_output( diff --git a/afd_plugin/model_executor/models/forward_context.py b/afd_plugin/forward_context.py similarity index 92% rename from afd_plugin/model_executor/models/forward_context.py rename to afd_plugin/forward_context.py index 0cdd3a0aa..8cc72f86a 100644 --- a/afd_plugin/model_executor/models/forward_context.py +++ b/afd_plugin/forward_context.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Helpers for plugin-owned model wrappers to read AFD forward metadata.""" +"""AFD forward metadata installation and access shared by runners and models.""" from __future__ import annotations @@ -20,8 +20,8 @@ def get_afd_metadata_from_forward_context( ) -> AFDForwardContextMetadata | None: """Return AFD metadata from vLLM ``ForwardContext.additional_kwargs``. - Model wrappers use this helper so AFD metadata stays outside vLLM's - ``ForwardContext`` schema. + Models and remote FFN proxies read the metadata installed by runners + without extending vLLM's ``ForwardContext`` schema. """ if forward_context is None: diff --git a/afd_plugin/model_executor/ffn/__init__.py b/afd_plugin/model_executor/ffn/__init__.py new file mode 100644 index 000000000..d71e61b11 --- /dev/null +++ b/afd_plugin/model_executor/ffn/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Shared remote calls, expert weight contracts and FFN compute implementations.""" diff --git a/afd_plugin/model_executor/models/npu/async_cam_layout.py b/afd_plugin/model_executor/ffn/async_cam.py similarity index 82% rename from afd_plugin/model_executor/models/npu/async_cam_layout.py rename to afd_plugin/model_executor/ffn/async_cam.py index f9b373507..a9d4a4191 100644 --- a/afd_plugin/model_executor/models/npu/async_cam_layout.py +++ b/afd_plugin/model_executor/ffn/async_cam.py @@ -1,10 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Attention TP token-layout conversion at the Async CAM boundary.""" +"""Model-side Async CAM tensor layout and native DBO receive handoff.""" from __future__ import annotations import os +from collections.abc import Callable from dataclasses import dataclass from typing import Final @@ -12,6 +13,13 @@ from vllm.distributed import tensor_model_parallel_all_gather from vllm.distributed.parallel_state import get_tp_group from vllm.logger import init_logger +from vllm.v1.worker.ubatching import ( + dbo_enabled, + dbo_maybe_run_recv_hook, + dbo_register_recv_hook, +) + +from afd_plugin.connectors import AFDConnectorBase, AFDTransferContext ASYNC_MOE_LAYOUT_LOG_ENV: Final[str] = "AFD_ASYNC_MOE_LAYOUT_LOG" _TRUE_ENV_VALUES: Final[frozenset[str]] = frozenset({"1", "true", "yes", "on"}) @@ -223,10 +231,46 @@ def _pad_first_dim( ) +def send_async_cam_remote_call( + connector: AFDConnectorBase, + hidden_states: torch.Tensor, + context: AFDTransferContext, + **send_kwargs: torch.Tensor | None, +) -> Callable[[], torch.Tensor]: + """Send once and let the peer submit raw receive before its next send. + + The returned closure belongs to this call. Layout restoration, TP + collectives and shared/HC/scaling remain with the resumed owner. + """ + overlapping = dbo_enabled() + if overlapping: + dbo_maybe_run_recv_hook() + connector.send_attn_output(hidden_states, context, **send_kwargs) + owner_stage = context.metadata.stage_idx + receive_started = False + output: torch.Tensor | None = None + + def receive() -> torch.Tensor: + nonlocal receive_started, output + if not receive_started: + receive_started = True + output = connector.recv_ffn_output( + ref_tensor=hidden_states, ubatch_idx=owner_stage, context=context + ) + if output is None: + raise RuntimeError("Async CAM receive failed; restart the A/F worker group") + return output + + if overlapping: + dbo_register_recv_hook(receive) + return receive + + __all__ = [ "ASYNC_MOE_LAYOUT_LOG_ENV", "CAMDispatchLayout", "CAMDispatchPayload", "prepare_cam_dispatch_payload", "restore_cam_dispatch_output", + "send_async_cam_remote_call", ] diff --git a/afd_plugin/model_executor/ffn/compute.py b/afd_plugin/model_executor/ffn/compute.py new file mode 100644 index 000000000..256c3f214 --- /dev/null +++ b/afd_plugin/model_executor/ffn/compute.py @@ -0,0 +1,231 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""NPU Async CAM compute: independent Layered and Mega native calls.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +import torch + +from afd_plugin.model_executor.ffn.expert_weights import ( + QuantizedExpertWeights, + expert_geometry, + validate_expert_weights, +) + +if TYPE_CHECKING: + from afd_plugin.connectors.npu.async_cam import AFDAsyncTopology + +CAM_LAYER_INDEX = 2 + + +class AsyncCAMLayeredExecutor: + """Device-controlled Layered FFN; retains only compute tensors and mappings.""" + + def __init__(self, layers: list[QuantizedExpertWeights]) -> None: + layers = validate_expert_weights(layers) + self.layer_ids = tuple(layer.layer_idx for layer in layers) + first = layers[0] + self.w13 = [layer.w13 for layer in layers] + self.w2 = [layer.w2 for layer in layers] + self.w13_scale = [layer.w13_scale for layer in layers] + self.w2_scale = [layer.w2_scale for layer in layers] + self.quantization = first.quantization + self.w13_bias: list[torch.Tensor] + self.w2_bias: list[torch.Tensor] + if first.quantization == "W8A8": + self.w13_bias = [] + # The existing W2 binding requires a bias for each layer. Symmetric + # W8A8 has none: share one zero INT32 tensor across the entire table. + # Adding zero preserves the integer accumulator without changing + # the native binding or allocating compensation per model layer. + zero_bias = torch.zeros_like(first.w2_scale, dtype=torch.int32) + self.w2_bias = [zero_bias] * len(layers) + else: + self.w13_bias = [cast(torch.Tensor, layer.w13_bias) for layer in layers] + self.w2_bias = [cast(torch.Tensor, layer.w2_bias) for layer in layers] + self.per_channel = first.per_channel + # W4A8 applies the model limit; W8A8 keeps its unclipped computation. + self.swiglu_limit = first.swiglu_limit + self.applied_swiglu_limit = ( + first.swiglu_limit if first.quantization == "W4A8" else 0.0 + ) + self.output_dtype = first.output_dtype + self.routed_scaling_factor = first.routed_scaling_factor + self.layer_id_to_slot = None + if self.layer_ids != tuple(range(len(layers))): + mapping = [-1] * (self.layer_ids[-1] + 1) + for slot, layer_idx in enumerate(self.layer_ids): + mapping[layer_idx] = slot + self.layer_id_to_slot = torch.tensor( + mapping, + dtype=torch.int64, + device=first.w13.device, + ) + # Resolve both operators once, before dispatch can receive work. + self.w13_op = torch.ops.afd_ascend.gmm_swiglu_quant_v2_layered + self.w2_op = torch.ops.afd_ascend.grouped_matmul_layered + + def __call__( + self, + hidden_states: torch.Tensor, + dynamic_scales: torch.Tensor, + group_list: torch.Tensor, + batch_info: torch.Tensor, + *, + retained_tensors: list[torch.Tensor] | None = None, + ) -> torch.Tensor: + layer_index = batch_info[CAM_LAYER_INDEX : CAM_LAYER_INDEX + 1] + if self.layer_id_to_slot is not None: + layer_index = self.layer_id_to_slot.index_select(0, layer_index) + else: + # The CANN tiling path checks the storage shape. A slice still + # carries batch_info's larger storage, so materialize one element. + layer_index = layer_index.clone() + activation, scale = self.w13_op( + hidden_states, + self.w13, + self.w13_scale, + self.w13_bias, + dynamic_scales, + group_list, + layer_index, + dequant_mode=0 if self.per_channel else 1, + quant_mode=0, + group_list_type=1, + swiglu_limit=self.applied_swiglu_limit, + ) + output = self.w2_op( + [activation], + self.w2, + self.w2_bias, + self.w2_scale, + layer_index, + group_list, + per_token_scale=scale, + group_list_type=1, + split_item=3, + output_dtype=self.output_dtype, + )[0] + if retained_tensors is not None: + # A graph session owns these capture-time addresses through its + # last replay. Eager calls retain only the usual stream lifetimes. + retained_tensors.extend((layer_index, activation, scale, output)) + if self.routed_scaling_factor != 1.0: + output.mul_(self.routed_scaling_factor) + return output + + +# Mirror Pull::MetadataWords in afd_pull_layout.h (all fields are int64_t): +# [total, source, layer, round_begin, round_end], tp source prefixes, +# tp * experts original counts, then [session, sequence]. +METADATA_WORD_BYTES = torch.int64.itemsize +METADATA_FIXED_WORDS = 7 +METADATA_PREFIX_WORDS_PER_SOURCE = 1 +# NativeMegaWorkspace in afd_native_mega_workspace.h aligns scratch regions to +# 512 bytes. Match its counts offset, including padding after batch metadata. +WORKSPACE_ALIGNMENT_BYTES = 512 + + +class AsyncCAMMegaExecutor: + """Pull transport with a compiled W4A8 or W8A8 execution plan.""" + + def __init__( + self, + layers: list[QuantizedExpertWeights], + *, + topology: AFDAsyncTopology, + max_num_batched_tokens: int, + topk: int, + tp_size: int, + capacity: int, + ) -> None: + layers = validate_expert_weights(layers) + first = layers[0] + if not first.per_channel: + raise ValueError("native Mega requires per-channel W4A8/W8A8 weights") + self.w13 = [layer.w13 for layer in layers] + self.w2 = [layer.w2 for layer in layers] + self.w13_scale = [layer.w13_scale for layer in layers] + self.w2_scale = [layer.w2_scale for layer in layers] + if first.quantization == "W4A8": + self.w13_bias = [cast(torch.Tensor, layer.w13_bias) for layer in layers] + self.w2_bias = [cast(torch.Tensor, layer.w2_bias) for layer in layers] + else: + self.w13_bias = [] + self.w2_bias = [] + self.layer_ids = tuple(layer.layer_idx for layer in layers) + self.quantization = first.quantization + self.routed_scaling_factor = first.routed_scaling_factor + self.applied_swiglu_limit = ( + first.swiglu_limit if first.quantization == "W4A8" else 0.0 + ) + self.capacity = capacity + self.comm_args: torch.Tensor | None = None + self.group_name = "" + experts, hidden, intermediate = expert_geometry(first) + self.config = [ + max_num_batched_tokens, + hidden, + intermediate, + topk, + topology.ffn_size, + topology.attn_size, + experts, + topology.world_rank, + topology.world_size, + tp_size, + self.capacity, + 1 if first.quantization == "W4A8" else 2, + 1, + len(self.layer_ids), + ] + device = first.w13.device + size = torch.ops.afd_ascend.afd_mega_moe_pull_workspace_size( + self.capacity, hidden, intermediate, experts, tp_size + ) + self.scratch = torch.empty(size, dtype=torch.uint8, device=device) + metadata_words = METADATA_FIXED_WORDS + tp_size * ( + METADATA_PREFIX_WORDS_PER_SOURCE + experts + ) + metadata_bytes = METADATA_WORD_BYTES * metadata_words + count_offset = ( + (metadata_bytes + WORKSPACE_ALIGNMENT_BYTES - 1) + // WORKSPACE_ALIGNMENT_BYTES + * WORKSPACE_ALIGNMENT_BYTES + ) + count_bytes = METADATA_WORD_BYTES * experts + self.batch_info = self.scratch[:metadata_bytes].view(torch.int64) + self.group_list = self.scratch[count_offset : count_offset + count_bytes].view( + torch.int64 + ) + self.anchor = torch.empty(1, dtype=torch.bfloat16, device=device) + mapping = [-1] * (self.layer_ids[-1] + 1) + for slot, layer_id in enumerate(self.layer_ids): + mapping[layer_id] = slot + self.layer_map = torch.tensor(mapping, dtype=torch.int64, device=device) + self.op = torch.ops.afd_ascend.afd_mega_moe_pull_native + + def bind_communication(self, comm_args: torch.Tensor, group_name: str) -> None: + self.comm_args, self.group_name = comm_args, group_name + + def __call__(self) -> torch.Tensor: + if self.comm_args is None: + raise RuntimeError("Mega communication must be bound before execution") + return self.op( + self.anchor, + self.comm_args, + self.layer_map, + self.scratch, + self.w13, + self.w2, + self.w13_scale, + self.w2_scale, + self.w13_bias, + self.w2_bias, + self.config, + self.group_name, + self.routed_scaling_factor, + self.applied_swiglu_limit, + ) diff --git a/afd_plugin/model_executor/ffn/expert_weights.py b/afd_plugin/model_executor/ffn/expert_weights.py new file mode 100644 index 000000000..d41386e12 --- /dev/null +++ b/afd_plugin/model_executor/ffn/expert_weights.py @@ -0,0 +1,205 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Loaded expert tensors and shared quantized layout contracts; no native ops.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from typing import Literal, cast + +import torch + +INT4_PER_INT32 = 8 + + +@dataclass(frozen=True) +class QuantizedExpertWeights: + layer_idx: int + w13: torch.Tensor + w2: torch.Tensor + w13_scale: torch.Tensor + w2_scale: torch.Tensor + w13_bias: torch.Tensor | None + w2_bias: torch.Tensor | None + per_channel: bool + swiglu_limit: float + routed_scaling_factor: float + # Async CAM quantizes to INT8; the existing W4A8 MLP returns BF16. + output_dtype: torch.dtype = torch.bfloat16 + quantization: Literal["W4A8", "W8A8"] = "W4A8" + + def signature(self) -> tuple: + return ( + tuple( + (tuple(t.shape), t.dtype, t.stride(), t.device) + if t is not None + else None + for t in ( + self.w13, + self.w2, + self.w13_scale, + self.w2_scale, + self.w13_bias, + self.w2_bias, + ) + ), + self.per_channel, + self.swiglu_limit, + self.routed_scaling_factor, + self.output_dtype, + self.quantization, + ) + + +def validate_expert_weights( + layers: list[QuantizedExpertWeights], +) -> list[QuantizedExpertWeights]: + layers = sorted(layers, key=lambda layer: layer.layer_idx) + if not layers: + raise ValueError("layered GMM requires at least one quantized MoE layer") + layer_ids = tuple(layer.layer_idx for layer in layers) + if layer_ids[0] < 0 or len(set(layer_ids)) != len(layers): + raise ValueError("layered GMM requires unique nonnegative model layer IDs") + if any(layer.signature() != layers[0].signature() for layer in layers): + raise ValueError("layered GMM requires homogeneous quantized layers") + for layer in layers: + _validate_layer(layer) + return layers + + +def expert_geometry(layer: QuantizedExpertWeights) -> tuple[int, int, int]: + """Return (experts, hidden, intermediate) from the loaded weight layout. + + INT4 weights pack eight nibbles per INT32, so the stored last dimension + is a quarter of the gate/up width; INT8 stores that width directly. + Callers that need only the logical geometry must not re-derive it. + """ + experts, hidden, last = layer.w13.shape + if layer.quantization == "W8A8": + if last % 2: + raise ValueError("layered GMM requires an even W13 gate/up width") + return experts, hidden, last // 2 + return experts, hidden, last * INT4_PER_INT32 // 2 + + +def expert_weight_signature(layers: list[QuantizedExpertWeights], capacity: int) -> str: + """Preserve the negotiated layout, including Layered's W8A8 zero bias.""" + layers = validate_expert_weights(layers) + first = layers[0] + tensors = [ + first.w13, + first.w2, + first.w13_scale, + first.w2_scale, + first.w13_bias, + first.w2_bias, + ] + if first.quantization == "W8A8": + # Meta describes zeros_like's layout without allocating a device buffer. + tensors[-1] = torch.empty_like(first.w2_scale, dtype=torch.int32, device="meta") + return json.dumps( + { + "capacity": capacity, + "layers": tuple(layer.layer_idx for layer in layers), + "weights": [ + (tuple(t.shape), str(t.dtype), t.stride()) if t is not None else None + for t in tensors + ], + "per_channel": first.per_channel, + "quantization": first.quantization, + "limit": first.swiglu_limit, + "routed_scaling": first.routed_scaling_factor, + "output_dtype": str(first.output_dtype), + }, + sort_keys=True, + ) + + +def _validate_layer(layer: QuantizedExpertWeights) -> None: + prefix = f"layered GMM layer {layer.layer_idx}: " + if layer.quantization == "W8A8": + _validate_w8a8_layer(layer) + return + if layer.quantization != "W4A8": + raise ValueError(prefix + "requires W4A8 or W8A8 dynamic quantization") + if layer.w13.dim() != 3 or layer.w2.dim() != 3: + raise ValueError(prefix + "weights must be packed [E, K, N/8] tensors") + experts, hidden, packed_intermediate = layer.w13.shape + intermediate = packed_intermediate * INT4_PER_INT32 // 2 + if layer.w2.shape != (experts, intermediate, hidden // INT4_PER_INT32): + raise ValueError(prefix + "W13/W2 hidden and intermediate dimensions disagree") + for name, tensor in (("w13", layer.w13), ("w2", layer.w2)): + if tensor.dtype != torch.int32 or not tensor.is_contiguous(): + raise ValueError( + prefix + name + " requires contiguous INT32-packed INT4 storage" + ) + for name, scale, width, k in ( + ("w13_scale", layer.w13_scale, 2 * intermediate, hidden), + ("w2_scale", layer.w2_scale, hidden, intermediate), + ): + expected_rank = 2 if name == "w13_scale" and layer.per_channel else 3 + if ( + scale.dtype != torch.int64 + or scale.dim() != expected_rank + or scale.shape[0] != experts + or scale.shape[-1] != width + ): + raise ValueError( + prefix + name + " has incompatible shape or encoded INT64 dtype" + ) + if scale.dim() == 3: + groups = scale.shape[1] + if groups <= 0 or k % groups or (layer.per_channel and groups != 1): + raise ValueError( + prefix + name + " has incompatible quantization groups" + ) + for name, bias, width in ( + ("w13_scale_bias", layer.w13_bias, 2 * intermediate), + ("w2_scale_bias", layer.w2_bias, hidden), + ): + if ( + bias is None + or bias.dtype != torch.float32 + or bias.shape != (experts, width) + ): + raise ValueError(prefix + name + " requires [E, N] FP32 compensation") + tensors = ( + layer.w2, + layer.w13_scale, + layer.w2_scale, + cast(torch.Tensor, layer.w13_bias), + cast(torch.Tensor, layer.w2_bias), + ) + if any(t.device != layer.w13.device for t in tensors): + raise ValueError( + prefix + "all weights, scales and compensation must share a device" + ) + + +def _validate_w8a8_layer(layer: QuantizedExpertWeights) -> None: + prefix = f"layered GMM layer {layer.layer_idx}: " + if layer.w13.dim() != 3 or layer.w2.dim() != 3: + raise ValueError(prefix + "W8A8 weights require [E, K, N] tensors") + experts, hidden, gate_up = layer.w13.shape + if gate_up % 2 or layer.w2.shape != (experts, gate_up // 2, hidden): + raise ValueError(prefix + "W13/W2 hidden and intermediate dimensions disagree") + if any( + t.dtype != torch.int8 or not t.is_contiguous() for t in (layer.w13, layer.w2) + ): + raise ValueError(prefix + "W8A8 requires contiguous INT8 weights") + if not layer.per_channel or layer.w13_scale.dtype != torch.float32: + raise ValueError(prefix + "W8A8 requires FP32 per-channel W13 scales") + if layer.w13_scale.shape != (experts, gate_up): + raise ValueError(prefix + "W8A8 w13_scale requires [E, 2I]") + if layer.w2_scale.shape != (experts, hidden): + raise ValueError(prefix + "W8A8 w2_scale requires [E, H] scales") + if layer.w2_scale.dtype != torch.bfloat16 or layer.output_dtype != torch.bfloat16: + raise ValueError(prefix + "W8A8 requires BF16 W2 scales and output") + if layer.w13_bias is not None or layer.w2_bias is not None: + raise ValueError(prefix + "W8A8 does not use W4A8 compensation") + if any( + t.device != layer.w13.device + for t in (layer.w2, layer.w13_scale, layer.w2_scale) + ): + raise ValueError(prefix + "all weights and scales must share a device") diff --git a/afd_plugin/model_executor/ffn/remote.py b/afd_plugin/model_executor/ffn/remote.py new file mode 100644 index 000000000..4c0575193 --- /dev/null +++ b/afd_plugin/model_executor/ffn/remote.py @@ -0,0 +1,184 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Connector-independent remote dense/MoE proxies; no concrete model imports.""" + +from __future__ import annotations + +import torch +from torch import nn +from vllm.distributed import ( + get_tensor_model_parallel_rank, + tensor_model_parallel_all_gather, +) +from vllm.forward_context import get_forward_context +from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.models.utils import sequence_parallel_chunk + +from afd_plugin.connectors import ( + AFDForwardContextMetadata, + AFDTransferContext, + AFDTransferMetadata, +) +from afd_plugin.forward_context import get_afd_metadata_from_forward_context +from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield + +UNROUTED_EXPERT_ID = -1 + + +def _unroute_tp_padding(topk_ids: torch.Tensor, num_tokens: int) -> torch.Tensor: + """Route the padding rows of this TP rank's ``sequence_parallel_chunk`` + slice of ``num_tokens`` tokens to ``UNROUTED_EXPERT_ID``.""" + local_tokens = topk_ids.shape[0] + token_index = get_tensor_model_parallel_rank() * local_tokens + token_index = token_index + torch.arange(local_tokens, device=topk_ids.device) + return topk_ids.masked_fill( + (token_index >= num_tokens).unsqueeze(1), UNROUTED_EXPERT_ID + ) + + +class RemoteFFNProxy(nn.Module): + """FFN stage executed through the AFD connector. + + By default the proxy is parameter-free: it sends the hidden states and the + FFN role runs the complete MLP or MoE block. A model whose connector routes + on Attention (an M2N routed-experts connector such as DeepEP M2N) sets + ``gate``, ``top_k``, ``renormalize`` and ``tp_size`` after construction, + from its model and parallel config. The proxy then computes the gate and + vLLM's default fused top-k (``FusedTopKRouter``; EPLB is not supported) + and sends ``topk_ids`` and ``topk_weights``; combine returns the top-k + weighted sum, which replaces the native MoE block's output. + + With Attention TP > 1 every TP rank holds the same tokens after the + attention all-reduce; dispatching them from every rank would make the FFN + run each token TP times. Each rank instead pads the tokens to a multiple + of ``tp_size`` (``sequence_parallel_chunk``), routes and dispatches only + its slice, and the TP group all-gathers the combined slices, as vLLM's + sequence-parallel MoE does. Padding rows get ``UNROUTED_EXPERT_ID`` and + reach no FFN rank; they are dropped after the all-gather. + """ + + def __init__(self, *, layer_idx: int) -> None: + super().__init__() + self.layer_idx = layer_idx + self.gate: ReplicatedLinear | None = None + self.top_k = 0 + self.renormalize = False + self.tp_size = 1 + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + if self.gate is None: + return self._send_and_receive(hidden_states) + num_tokens = hidden_states.shape[0] + if self.tp_size > 1: + hidden_states = sequence_parallel_chunk(hidden_states) + router_logits, _ = self.gate(hidden_states) + topk_weights, topk_ids, _ = fused_topk( + hidden_states=hidden_states, + gating_output=router_logits, + topk=self.top_k, + renormalize=self.renormalize, + ) + if self.tp_size > 1: + topk_ids = _unroute_tp_padding(topk_ids, num_tokens) + output = self._send_and_receive( + hidden_states, + topk_ids=topk_ids, + topk_weights=topk_weights, + ) + if self.tp_size > 1: + output = tensor_model_parallel_all_gather(output, 0) + output = output[:num_tokens] + return output + + def _send_and_receive( + self, + hidden_states: torch.Tensor, + **send_kwargs: torch.Tensor, + ) -> torch.Tensor: + stage_idx = self._send(hidden_states, **send_kwargs) + hidden_states = maybe_apply_dbo_yield( + hidden_states, + role="attention", + ) + return self._receive(hidden_states, stage_idx) + + def _send( + self, + hidden_states: torch.Tensor, + **send_kwargs: torch.Tensor, + ) -> int: + """Send this layer's Attention output; return the stage to receive.""" + afd_metadata = self._afd_metadata() + forward_context = get_forward_context() + stage_idx = int( + getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), + ) + afd_metadata.stage_idx = stage_idx + metadata = AFDTransferMetadata.create_attention_metadata( + layer_idx=self.layer_idx, + stage_idx=stage_idx, + seq_len=int(hidden_states.shape[0]), + ) + context = AFDTransferContext(metadata=metadata) + afd_metadata.connector.send_attn_output( + hidden_states, + context, + **send_kwargs, + ) + return stage_idx + + def _receive(self, hidden_states: torch.Tensor, stage_idx: int) -> torch.Tensor: + """Receive the FFN output; a connector may write it into + ``hidden_states``.""" + return self._afd_metadata().connector.recv_ffn_output( + ref_tensor=hidden_states, + ubatch_idx=stage_idx, + ) + + @staticmethod + def _afd_metadata() -> AFDForwardContextMetadata: + afd_metadata = get_afd_metadata_from_forward_context() + if afd_metadata is None: + raise RuntimeError("RemoteFFNProxy requires AFD forward metadata") + return afd_metadata + + +class AFDAttentionFusedMoE(RemoteFFNProxy): + """Parameter-free native-MoE experts proxy for the Attention runtime.""" + + def __init__( + self, + *, + layer_idx: int, + external_routing: bool, + ) -> None: + super().__init__(layer_idx=layer_idx) + # vLLM 5df63a72 fused_moe/runner/moe_runner.py::MoERunner.forward + # enters _forward_impl, which consumes transported router_logits only + # when its gate is None (external routing); a held gate routes internally + # and ignores the forwarded logits argument. + self.external_routing = external_routing + + # The native MoERunner boundary passes router_logits and input_ids; only + # the proxy's transport is reused, not its kernel execution contract. + def forward( # type: ignore[override] + self, + hidden_states: torch.Tensor, + router_logits: torch.Tensor, + input_ids: torch.Tensor | None = None, + ) -> torch.Tensor: + if input_ids is not None: + raise NotImplementedError( + "experts-boundary input_ids transport is not implemented", + ) + send_kwargs: dict[str, torch.Tensor] = {} + if self.external_routing: + send_kwargs["router_logits"] = router_logits + return self._send_and_receive(hidden_states, **send_kwargs) + + def update_expert_map(self) -> None: + """Satisfy the native MoERunner EPLB interface without local experts.""" + + +__all__ = ["RemoteFFNProxy", "AFDAttentionFusedMoE"] diff --git a/afd_plugin/model_executor/models/__init__.py b/afd_plugin/model_executor/models/__init__.py index 622fdaeea..639fb9257 100644 --- a/afd_plugin/model_executor/models/__init__.py +++ b/afd_plugin/model_executor/models/__init__.py @@ -1,11 +1,19 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""AFD model wrapper namespace.""" +"""AFD model wrappers; shared metadata access remains a lazy compatibility export.""" -from afd_plugin.model_executor.models.forward_context import ( - get_afd_metadata_from_forward_context, -) +from typing import TYPE_CHECKING -__all__ = [ - "get_afd_metadata_from_forward_context", -] +if TYPE_CHECKING: + from afd_plugin.forward_context import get_afd_metadata_from_forward_context + + +def __getattr__(name: str): + if name == "get_afd_metadata_from_forward_context": + from afd_plugin.forward_context import get_afd_metadata_from_forward_context + + return get_afd_metadata_from_forward_context + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + +__all__ = ["get_afd_metadata_from_forward_context"] diff --git a/afd_plugin/model_executor/models/_weight_utils.py b/afd_plugin/model_executor/models/_weight_utils.py new file mode 100644 index 000000000..236e33cfb --- /dev/null +++ b/afd_plugin/model_executor/models/_weight_utils.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Checkpoint path parsing and one-pass filtering; models own role policies.""" + +from collections.abc import Callable, Iterable, Iterator +from typing import TypeVar + +_Weight = TypeVar("_Weight") + + +def layer_weight_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: + parts = name.split(".") + for index, part in enumerate(parts[:-2]): + if part != "layers": + continue + try: + layer_idx = int(parts[index + 1]) + except ValueError: + continue + return layer_idx, parts[index + 2], tuple(parts[index + 3 :]) + return None + + +def filter_role_weights( + weights: Iterable[tuple[str, _Weight]], + *, + role: str | None, + classify: Callable[[str], frozenset[str]], +) -> Iterator[tuple[str, _Weight]]: + for name, loaded_weight in weights: + if role is None or role in classify(name): + yield name, loaded_weight diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 54f115f34..110b8a073 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -8,22 +8,20 @@ hidden states between the Attention and FFN roles through the AFD connector. """ -from collections.abc import Iterable, Iterator +from collections.abc import Iterable from copy import copy +from functools import partial from typing import TYPE_CHECKING, Any, TypeAlias import torch import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config, GlmMoeDsaConfig from vllm.config import ParallelConfig, VllmConfig -from vllm.forward_context import get_forward_context from vllm.logger import init_logger from vllm.model_executor.layers import fused_moe -from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk from vllm.model_executor.layers.fused_moe.router.router_factory import ( create_fused_moe_router, ) -from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.models import deepseek_v2 as native from afd_plugin.config import ( @@ -31,13 +29,16 @@ DEEPEP_M2N_CONNECTORS, parse_afd_config, ) -from afd_plugin.connectors import ( - AFDExpertRoutingSpec, - AFDForwardContextMetadata, - AFDTransferContext, - AFDTransferMetadata, +from afd_plugin.connectors import AFDExpertRoutingSpec +from afd_plugin.model_executor.ffn.remote import ( + AFDAttentionFusedMoE, + RemoteFFNProxy, + _unroute_tp_padding, +) +from afd_plugin.model_executor.models._weight_utils import ( + filter_role_weights, + layer_weight_path, ) -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield if TYPE_CHECKING: @@ -68,22 +69,6 @@ def _is_moe_layer(config: _DeepseekAdapterConfig, layer_idx: int) -> bool: AFD_ASYNC_CONNECTOR, *DEEPEP_M2N_CONNECTORS, ) -# Expert id of a padding row: a routed-experts connector sends it nowhere. -UNROUTED_EXPERT_ID = -1 - - -def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: - """Return ``(layer index, stage, remainder)`` for a decoder weight.""" - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) - return None def _checkpoint_weight_roles( @@ -94,7 +79,7 @@ def _checkpoint_weight_roles( attention_shared_experts: bool = False, ) -> frozenset[str]: """Classify one native checkpoint path by its AFD execution owner.""" - layer_path = _weight_layer_path(name) + layer_path = layer_weight_path(name) if layer_path is None: return _BOTH_ROLES @@ -115,180 +100,6 @@ def _checkpoint_weight_roles( return _FFN_ROLE -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str, - config: _DeepseekAdapterConfig, - compute_gate_on_attention: bool, - attention_shared_experts: bool = False, -) -> Iterator[tuple[str, torch.Tensor]]: - """Consume a checkpoint iterator once and retain this role's paths.""" - for name, loaded_weight in weights: - if role in _checkpoint_weight_roles( - name, - config, - compute_gate_on_attention=compute_gate_on_attention, - attention_shared_experts=attention_shared_experts, - ): - yield name, loaded_weight - - -def _unroute_tp_padding(topk_ids: torch.Tensor, num_tokens: int) -> torch.Tensor: - """Route the padding rows of this TP rank's ``sequence_parallel_chunk`` - slice of ``num_tokens`` tokens to ``UNROUTED_EXPERT_ID``.""" - local_tokens = topk_ids.shape[0] - token_index = native.get_tensor_model_parallel_rank() * local_tokens - token_index = token_index + torch.arange(local_tokens, device=topk_ids.device) - return topk_ids.masked_fill( - (token_index >= num_tokens).unsqueeze(1), UNROUTED_EXPERT_ID - ) - - -class RemoteFFNProxy(nn.Module): - """FFN stage executed through the AFD connector. - - By default the proxy is parameter-free: it sends the hidden states and the - FFN role runs the complete MLP or MoE block. A model whose connector routes - on Attention (an M2N routed-experts connector such as DeepEP M2N) sets - ``gate``, ``top_k``, ``renormalize`` and ``tp_size`` after construction, - from its model and parallel config. The proxy then computes the gate and - vLLM's default fused top-k (``FusedTopKRouter``; EPLB is not supported) - and sends ``topk_ids`` and ``topk_weights``; combine returns the top-k - weighted sum, which replaces the native MoE block's output. - - With Attention TP > 1 every TP rank holds the same tokens after the - attention all-reduce; dispatching them from every rank would make the FFN - run each token TP times. Each rank instead pads the tokens to a multiple - of ``tp_size`` (``sequence_parallel_chunk``), routes and dispatches only - its slice, and the TP group all-gathers the combined slices, as vLLM's - sequence-parallel MoE does. Padding rows get ``UNROUTED_EXPERT_ID`` and - reach no FFN rank; they are dropped after the all-gather. - """ - - def __init__(self, *, layer_idx: int) -> None: - super().__init__() - self.layer_idx = layer_idx - self.gate: ReplicatedLinear | None = None - self.top_k = 0 - self.renormalize = False - self.tp_size = 1 - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - if self.gate is None: - return self._send_and_receive(hidden_states) - num_tokens = hidden_states.shape[0] - if self.tp_size > 1: - hidden_states = native.sequence_parallel_chunk(hidden_states) - router_logits, _ = self.gate(hidden_states) - topk_weights, topk_ids, _ = fused_topk( - hidden_states=hidden_states, - gating_output=router_logits, - topk=self.top_k, - renormalize=self.renormalize, - ) - if self.tp_size > 1: - topk_ids = _unroute_tp_padding(topk_ids, num_tokens) - output = self._send_and_receive( - hidden_states, - topk_ids=topk_ids, - topk_weights=topk_weights, - ) - if self.tp_size > 1: - output = native.tensor_model_parallel_all_gather(output, 0) - output = output[:num_tokens] - return output - - def _send_and_receive( - self, - hidden_states: torch.Tensor, - **send_kwargs: torch.Tensor, - ) -> torch.Tensor: - stage_idx = self._send(hidden_states, **send_kwargs) - hidden_states = maybe_apply_dbo_yield( - hidden_states, - role="attention", - ) - return self._receive(hidden_states, stage_idx) - - def _send( - self, - hidden_states: torch.Tensor, - **send_kwargs: torch.Tensor, - ) -> int: - """Send this layer's Attention output; return the stage to receive.""" - afd_metadata = self._afd_metadata() - forward_context = get_forward_context() - stage_idx = int( - getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), - ) - afd_metadata.stage_idx = stage_idx - metadata = AFDTransferMetadata.create_attention_metadata( - layer_idx=self.layer_idx, - stage_idx=stage_idx, - seq_len=int(hidden_states.shape[0]), - ) - context = AFDTransferContext(metadata=metadata) - afd_metadata.connector.send_attn_output( - hidden_states, - context, - **send_kwargs, - ) - return stage_idx - - def _receive(self, hidden_states: torch.Tensor, stage_idx: int) -> torch.Tensor: - """Receive the FFN output; a connector may write it into - ``hidden_states``.""" - return self._afd_metadata().connector.recv_ffn_output( - ref_tensor=hidden_states, - ubatch_idx=stage_idx, - ) - - @staticmethod - def _afd_metadata() -> AFDForwardContextMetadata: - afd_metadata = get_afd_metadata_from_forward_context() - if afd_metadata is None: - raise RuntimeError("RemoteFFNProxy requires AFD forward metadata") - return afd_metadata - - -class AFDAttentionFusedMoE(RemoteFFNProxy): - """Parameter-free native-MoE experts proxy for the Attention runtime.""" - - def __init__( - self, - *, - layer_idx: int, - external_routing: bool, - ) -> None: - super().__init__(layer_idx=layer_idx) - # vLLM 5df63a72 fused_moe/runner/moe_runner.py::MoERunner.forward - # enters _forward_impl, which consumes transported router_logits only - # when its gate is None (external routing); a held gate routes internally - # and ignores the forwarded logits argument. - self.external_routing = external_routing - - # The native MoERunner boundary passes router_logits and input_ids; only - # the proxy's transport is reused, not its kernel execution contract. - def forward( # type: ignore[override] - self, - hidden_states: torch.Tensor, - router_logits: torch.Tensor, - input_ids: torch.Tensor | None = None, - ) -> torch.Tensor: - if input_ids is not None: - raise NotImplementedError( - "experts-boundary input_ids transport is not implemented", - ) - send_kwargs: dict[str, torch.Tensor] = {} - if self.external_routing: - send_kwargs["router_logits"] = router_logits - return self._send_and_receive(hidden_states, **send_kwargs) - - def update_expert_map(self) -> None: - """Satisfy the native MoERunner EPLB interface without local experts.""" - - class GateOnlyRemoteMoE(RemoteFFNProxy): """Attention-side MoE gate with experts delegated to the FFN role.""" @@ -1159,13 +970,16 @@ def get_experts_routing_spec( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return super().load_weights( - _iter_role_weights( + filter_role_weights( weights, role=self.afd_role, - config=self.config, - compute_gate_on_attention=self.afd_config.compute_gate_on_attention, - attention_shared_experts=self.afd_config.connector - in _ATTENTION_SHARED_EXPERTS_CONNECTORS, + classify=partial( + _checkpoint_weight_roles, + config=self.config, + compute_gate_on_attention=self.afd_config.compute_gate_on_attention, + attention_shared_experts=self.afd_config.connector + in _ATTENTION_SHARED_EXPERTS_CONNECTORS, + ), ) ) diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index 148788ce0..291d371fb 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -8,7 +8,7 @@ FFN activation is the sole FFN-to-Attention tensor. """ -from collections.abc import Iterable, Iterator +from collections.abc import Iterable from typing import Any import torch @@ -20,7 +20,11 @@ from afd_plugin.compat.patches.gpu_deepseek_v4_dbo import forward_v4_attention from afd_plugin.config import parse_afd_config from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.forward_context import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.models._weight_utils import ( + filter_role_weights, + layer_weight_path, +) from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield _ATTENTION_ROLE = frozenset(("attention",)) @@ -28,20 +32,6 @@ _BOTH_ROLES = frozenset(("attention", "ffn")) -def _weight_layer_path(name: str) -> tuple[int, str] | None: - """Extract the decoder layer index and first layer-local path component.""" - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2] - return None - - def _checkpoint_weight_roles(name: str) -> frozenset[str]: """Classify a native DeepSeek-V4 checkpoint path by execution owner.""" if name in { @@ -53,26 +43,15 @@ def _checkpoint_weight_roles(name: str) -> frozenset[str]: "model.hc_head_scale", }: return _ATTENTION_ROLE - layer_path = _weight_layer_path(name) + layer_path = layer_weight_path(name) if layer_path is None: return _BOTH_ROLES - _, stage = layer_path + _, stage, _ = layer_path if stage == "ffn": return _FFN_ROLE return _ATTENTION_ROLE -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str, -) -> Iterator[tuple[str, torch.Tensor]]: - """Consume a checkpoint iterator once and retain only role-owned paths.""" - for name, loaded_weight in weights: - if role in _checkpoint_weight_roles(name): - yield name, loaded_weight - - class RemoteDeepseekV4FFN(nn.Module): """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" @@ -615,7 +594,9 @@ def get_experts_layer_indices(self) -> tuple[int, ...]: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return super().load_weights( - _iter_role_weights(weights, role=self.afd_role), + filter_role_weights( + weights, role=self.afd_role, classify=_checkpoint_weight_roles + ), ) diff --git a/afd_plugin/model_executor/models/npu/async_cam_dbo.py b/afd_plugin/model_executor/models/npu/async_cam_dbo.py deleted file mode 100644 index 938548247..000000000 --- a/afd_plugin/model_executor/models/npu/async_cam_dbo.py +++ /dev/null @@ -1,49 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Single-window CAM handoff through the native DBO receive hook.""" - -from collections.abc import Callable - -import torch -from vllm.v1.worker.ubatching import ( - dbo_enabled, - dbo_maybe_run_recv_hook, - dbo_register_recv_hook, -) - -from afd_plugin.connectors import AFDConnectorBase, AFDTransferContext - - -def send_async_cam_remote_call( - connector: AFDConnectorBase, - hidden_states: torch.Tensor, - context: AFDTransferContext, - **send_kwargs: torch.Tensor | None, -) -> Callable[[], torch.Tensor]: - """Send once and let the peer submit raw receive before its next send. - - The returned closure belongs to this call. Layout restoration, TP - collectives and shared/HC/scaling remain with the resumed owner. - """ - overlapping = dbo_enabled() - if overlapping: - dbo_maybe_run_recv_hook() - connector.send_attn_output(hidden_states, context, **send_kwargs) - owner_stage = context.metadata.stage_idx - receive_started = False - output: torch.Tensor | None = None - - def receive() -> torch.Tensor: - nonlocal receive_started, output - if not receive_started: - receive_started = True - output = connector.recv_ffn_output( - ref_tensor=hidden_states, ubatch_idx=owner_stage, context=context - ) - if output is None: - raise RuntimeError("Async CAM receive failed; restart the A/F worker group") - return output - - if overlapping: - dbo_register_recv_hook(receive) - return receive diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index c43837594..a3ddff99d 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py @@ -23,14 +23,12 @@ AFDTransferContext, AFDTransferMetadata, ) -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context -from afd_plugin.model_executor.models.npu.async_cam_dbo import ( - send_async_cam_remote_call, -) -from afd_plugin.model_executor.models.npu.async_cam_layout import ( +from afd_plugin.forward_context import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.ffn.async_cam import ( CAMDispatchLayout, prepare_cam_dispatch_payload, restore_cam_dispatch_output, + send_async_cam_remote_call, ) from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py b/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py index 5e230e2c7..7b0e59693 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py @@ -10,7 +10,7 @@ import torch from afd_plugin.envs import force_balanced_topk_ids_enabled -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context +from afd_plugin.forward_context import get_afd_metadata_from_forward_context if TYPE_CHECKING: from vllm.config import VllmConfig diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index a7d45536d..f5a7965c9 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -11,8 +11,9 @@ ``compute_ffn_output`` hook. """ -from collections.abc import Iterable, Iterator +from collections.abc import Iterable from copy import copy +from functools import partial from typing import TYPE_CHECKING, Any, cast import torch @@ -34,14 +35,16 @@ AFDTransferContext, AFDTransferMetadata, ) -from afd_plugin.model_executor.models import get_afd_metadata_from_forward_context -from afd_plugin.model_executor.models.deepseek_v2 import RemoteFFNProxy -from afd_plugin.model_executor.models.npu.async_cam_dbo import ( - send_async_cam_remote_call, -) -from afd_plugin.model_executor.models.npu.async_cam_layout import ( +from afd_plugin.forward_context import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.ffn.async_cam import ( prepare_cam_dispatch_payload, restore_cam_dispatch_output, + send_async_cam_remote_call, +) +from afd_plugin.model_executor.ffn.remote import RemoteFFNProxy +from afd_plugin.model_executor.models._weight_utils import ( + filter_role_weights, + layer_weight_path, ) from afd_plugin.model_executor.models.npu.deepseek_v4_shared_experts import ( AFDDeepseekV4SharedExperts, @@ -65,7 +68,7 @@ AscendW4A8DynamicFusedMoEMethod, ) - from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights + from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights _ATTENTION_ROLE = "attention" @@ -82,19 +85,6 @@ def _refresh_ascend_fused_moe() -> None: native.FusedMoEFactory = fused_moe.FusedMoEFactory -def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) - return None - - def _checkpoint_weight_roles( name: str, *, @@ -114,7 +104,7 @@ def _checkpoint_weight_roles( skipping. """ - layer_path = _weight_layer_path(name) + layer_path = layer_weight_path(name) if layer_path is None: return _BOTH_ROLES @@ -140,22 +130,6 @@ def _checkpoint_weight_roles( return _BOTH_ROLES -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str, - attn_owns_gate: bool = True, - attention_shared_experts: bool = False, -) -> Iterator[tuple[str, torch.Tensor]]: - for name, loaded_weight in weights: - if role in _checkpoint_weight_roles( - name, - attn_owns_gate=attn_owns_gate, - attention_shared_experts=attention_shared_experts, - ): - yield name, loaded_weight - - class AFDDeepseekV4RemoteMoE(RemoteFFNProxy): """DSV4 gate-on-FFN shell that sends Hash ids alongside the activations. @@ -887,8 +861,10 @@ def process_weights_after_loading(self) -> None: prepare_attention_gate_w8a8_scales(self.model.layers) # ### PATCH END: retain FFN gate-on-A W8A8 scales after quant finalize - def get_async_cam_layered_weights(self) -> "tuple[list[LayeredFFNWeights], str]": - return _extract_async_cam_layered_weights(self.model.layers) + def get_quantized_expert_weights( + self, + ) -> "tuple[list[QuantizedExpertWeights], str]": + return _extract_quantized_expert_weights(self.model.layers) def compute_ffn_output( self, @@ -922,27 +898,27 @@ def compute_experts_output( def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # Only the gate-on-Attention configuration gives Attention a router; with # the gate on FFN its MoE slot is a parameter-free transfer shell. - # The shared experts follow the router, so their weights load on Attention - # for either connector; keying this on the async connector alone would - # construct the modules on the sync path without ever loading them. - role_weights = _iter_role_weights( + # Shared experts follow the router on both Async CAM and sync paths. + role_weights = filter_role_weights( weights, role=self.afd_role, - attn_owns_gate=bool(self.afd_config.compute_gate_on_attention), - attention_shared_experts=bool(self.afd_config.compute_gate_on_attention), + classify=partial( + _checkpoint_weight_roles, + attn_owns_gate=bool(self.afd_config.compute_gate_on_attention), + attention_shared_experts=bool(self.afd_config.compute_gate_on_attention), + ), ) - loaded = super().load_weights(role_weights) - return loaded + return super().load_weights(role_weights) -def _extract_async_cam_layered_weights( +def _extract_quantized_expert_weights( layers: Iterable[AFDDeepseekV4DecoderLayer], -) -> "tuple[list[LayeredFFNWeights], str]": - """Expose loaded DeepSeek V4 routed-expert weights to the layered executor.""" +) -> "tuple[list[QuantizedExpertWeights], str]": + """Expose loaded DeepSeek V4 routed-expert weights to Async CAM compute.""" from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm_ascend.quantization.quant_type import QuantType - from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights + from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights layers = [layer for layer in layers if layer.is_moe_layer] # Only the FFN role exposes loaded experts; Ascend's factory returns its @@ -985,7 +961,7 @@ def _extract_async_cam_layered_weights( f"layered GMM layer {layer.layer_idx}: missing loaded {name}" ) weights.append( - LayeredFFNWeights( + QuantizedExpertWeights( layer_idx=layer.layer_idx, # The presence checks above establish these dynamically # registered parameters as loaded tensors on the routed owner. diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py b/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py index 93fe1fc49..0c7aeb5ed 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py @@ -33,11 +33,11 @@ from afd_plugin.config import parse_afd_config from afd_plugin.model_executor.models.npu.deepseek_v4 import ( AFDDeepseekV4DecoderLayer, - _extract_async_cam_layered_weights, + _extract_quantized_expert_weights, ) if TYPE_CHECKING: - from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights + from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights def _iter_role_weights( @@ -253,8 +253,8 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: def get_experts_layer_indices(self) -> tuple[int, ...]: return tuple(int(index) for index in self.model.layers) - def get_async_cam_layered_weights(self) -> tuple[list[LayeredFFNWeights], str]: - return _extract_async_cam_layered_weights(self.model.layers.values()) + def get_quantized_expert_weights(self) -> tuple[list[QuantizedExpertWeights], str]: + return _extract_quantized_expert_weights(self.model.layers.values()) def compute_ffn_output( self, diff --git a/afd_plugin/model_executor/models/qwen3_5.py b/afd_plugin/model_executor/models/qwen3_5.py index fbfb9ae0a..05648e6e4 100644 --- a/afd_plugin/model_executor/models/qwen3_5.py +++ b/afd_plugin/model_executor/models/qwen3_5.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project """Qwen3.5/3.6 MoE AFD wrapper for the native vLLM lifecycle.""" -from collections.abc import Iterable, Iterator +from collections.abc import Iterable from typing import Any import torch @@ -12,7 +12,11 @@ from vllm.model_executor.models import qwen3_next as next_native from afd_plugin.config import parse_optional_afd_config -from afd_plugin.model_executor.models.deepseek_v2 import AFDAttentionFusedMoE +from afd_plugin.model_executor.ffn.remote import AFDAttentionFusedMoE +from afd_plugin.model_executor.models._weight_utils import ( + filter_role_weights, + layer_weight_path, +) _ATTENTION_ROLE = frozenset(("attention",)) _FFN_ROLE = frozenset(("ffn",)) @@ -29,20 +33,6 @@ def _validate_qwen_text_only(model_config: ModelConfig) -> None: ) -def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: - """Return ``(layer index, stage, remainder)`` for a decoder weight.""" - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) - return None - - def _checkpoint_weight_roles( name: str, ) -> frozenset[str]: @@ -51,7 +41,7 @@ def _checkpoint_weight_roles( if "visual" in parts or "mtp" in parts: return _NO_ROLES - layer_path = _weight_layer_path(name) + layer_path = layer_weight_path(name) if layer_path is None: # Embeddings, final norm, and lm_head are required only by Attention. return _ATTENTION_ROLE @@ -72,19 +62,6 @@ def _checkpoint_weight_roles( raise RuntimeError(f"unclassified Qwen MoE checkpoint weight: {name}") -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str | None, -) -> Iterator[tuple[str, torch.Tensor]]: - """Consume a checkpoint iterator once and retain only this role's paths.""" - for name, loaded_weight in weights: - if role is None or role in _checkpoint_weight_roles( - name, - ): - yield name, loaded_weight - - class AFDQwen3_5RemoteExpertsMoE( # noqa: N801 native.Qwen3NextSparseMoeBlock, ): @@ -551,9 +528,8 @@ def get_experts_layer_indices(self) -> tuple[int, ...]: def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # ### PATCH START: load each checkpoint path only on its owning role. loaded = super().load_weights( - _iter_role_weights( - weights, - role=self.afd_role, + filter_role_weights( + weights, role=self.afd_role, classify=_checkpoint_weight_roles ), ) # ### PATCH END diff --git a/afd_plugin/model_executor/models/qwen3_moe.py b/afd_plugin/model_executor/models/qwen3_moe.py index 54dc0f406..fd7cbedcd 100644 --- a/afd_plugin/model_executor/models/qwen3_moe.py +++ b/afd_plugin/model_executor/models/qwen3_moe.py @@ -10,7 +10,8 @@ the MoE gate and top-k, and FFN runs only the routed experts. """ -from collections.abc import Iterable, Iterator +from collections.abc import Iterable +from functools import partial from typing import TYPE_CHECKING import torch @@ -25,7 +26,11 @@ AFDConfig, parse_afd_config, ) -from afd_plugin.model_executor.models.deepseek_v2 import RemoteFFNProxy +from afd_plugin.model_executor.ffn.remote import RemoteFFNProxy +from afd_plugin.model_executor.models._weight_utils import ( + filter_role_weights, + layer_weight_path, +) if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts @@ -49,20 +54,6 @@ def _routes_on_attention(afd_config: AFDConfig) -> bool: return afd_config.connector in DEEPEP_M2N_CONNECTORS -def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: - """Return ``(layer index, stage, remainder)`` for a native decoder path.""" - parts = name.split(".") - for marker_idx, part in enumerate(parts[:-2]): - if part != "layers": - continue - try: - layer_idx = int(parts[marker_idx + 1]) - except ValueError: - continue - return layer_idx, parts[marker_idx + 2], tuple(parts[marker_idx + 3 :]) - return None - - def _checkpoint_weight_roles( name: str, *, @@ -74,7 +65,7 @@ def _checkpoint_weight_roles( Attention routes with it, and FFN keeps the native gate parameter that vLLM's strict loader requires but never runs it. """ - layer_path = _weight_layer_path(name) + layer_path = layer_weight_path(name) if layer_path is None: return _BOTH_ROLES @@ -89,21 +80,6 @@ def _checkpoint_weight_roles( return _BOTH_ROLES -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], - *, - role: str, - routes_on_attention: bool = False, -) -> Iterator[tuple[str, torch.Tensor]]: - """Consume a checkpoint iterator once and retain this role's paths.""" - for name, loaded_weight in weights: - if role in _checkpoint_weight_roles( - name, - routes_on_attention=routes_on_attention, - ): - yield name, loaded_weight - - def _validate_supported_config( vllm_config: VllmConfig, afd_config: AFDConfig, @@ -402,10 +378,13 @@ def get_routed_experts(self, layer_idx: int) -> "RoutedExperts": # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: # ### PATCH START: filter the checkpoint stream by AFD execution role. - role_weights = _iter_role_weights( + role_weights = filter_role_weights( weights, role=self.afd_role, - routes_on_attention=_routes_on_attention(self.afd_config), + classify=partial( + _checkpoint_weight_roles, + routes_on_attention=_routes_on_attention(self.afd_config), + ), ) # ### PATCH END: filter the checkpoint stream by AFD execution role. return super().load_weights(role_weights) diff --git a/afd_plugin/model_executor/npu/async_cam_layered.py b/afd_plugin/model_executor/npu/async_cam_layered.py deleted file mode 100644 index 06962e7a4..000000000 --- a/afd_plugin/model_executor/npu/async_cam_layered.py +++ /dev/null @@ -1,298 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Static W4A8/W8A8 weights and device-controlled async CAM FFN execution.""" - -from __future__ import annotations - -import json -from dataclasses import dataclass -from typing import Literal, cast - -import torch - -INT4_PER_INT32 = 8 -CAM_LAYER_INDEX = 2 - - -@dataclass(frozen=True) -class LayeredFFNWeights: - layer_idx: int - w13: torch.Tensor - w2: torch.Tensor - w13_scale: torch.Tensor - w2_scale: torch.Tensor - w13_bias: torch.Tensor | None - w2_bias: torch.Tensor | None - per_channel: bool - swiglu_limit: float - routed_scaling_factor: float - # Async CAM quantizes to INT8; the existing W4A8 MLP returns BF16. - output_dtype: torch.dtype = torch.bfloat16 - quantization: Literal["W4A8", "W8A8"] = "W4A8" - - def signature(self) -> tuple: - return ( - tuple( - (tuple(t.shape), t.dtype, t.stride(), t.device) - if t is not None - else None - for t in ( - self.w13, - self.w2, - self.w13_scale, - self.w2_scale, - self.w13_bias, - self.w2_bias, - ) - ), - self.per_channel, - self.swiglu_limit, - self.routed_scaling_factor, - self.output_dtype, - self.quantization, - ) - - -class AsyncCAMLayeredExecutor: - """Reuse layer tensors; all per-work-item control data stays on device.""" - - def __init__(self, layers: list[LayeredFFNWeights]) -> None: - layers = sorted(layers, key=lambda layer: layer.layer_idx) - if not layers: - raise ValueError("layered GMM requires at least one quantized MoE layer") - self.layer_ids = tuple(layer.layer_idx for layer in layers) - if self.layer_ids[0] < 0 or len(set(self.layer_ids)) != len(layers): - raise ValueError("layered GMM requires unique nonnegative model layer IDs") - first = layers[0] - if any(layer.signature() != first.signature() for layer in layers): - raise ValueError("layered GMM requires homogeneous quantized layers") - for layer in layers: - self._validate_layer(layer) - self.w13 = [layer.w13 for layer in layers] - self.w2 = [layer.w2 for layer in layers] - self.w13_scale = [layer.w13_scale for layer in layers] - self.w2_scale = [layer.w2_scale for layer in layers] - self.quantization = first.quantization - self.w13_bias: list[torch.Tensor] - self.w2_bias: list[torch.Tensor] - if first.quantization == "W8A8": - self.w13_bias = [] - # The existing W2 binding requires a bias for each layer. Symmetric - # W8A8 has none: share one zero INT32 tensor across the entire table. - # Adding zero preserves the integer accumulator without changing - # the native binding or allocating compensation per model layer. - zero_bias = torch.zeros_like(first.w2_scale, dtype=torch.int32) - self.w2_bias = [zero_bias] * len(layers) - else: - self.w13_bias = [cast(torch.Tensor, layer.w13_bias) for layer in layers] - self.w2_bias = [cast(torch.Tensor, layer.w2_bias) for layer in layers] - self.per_channel = first.per_channel - # W4A8 applies the model limit; W8A8 keeps its unclipped computation. - self.swiglu_limit = first.swiglu_limit - self.applied_swiglu_limit = ( - first.swiglu_limit if first.quantization == "W4A8" else 0.0 - ) - self.output_dtype = first.output_dtype - self.routed_scaling_factor = first.routed_scaling_factor - self.layer_id_to_slot = None - if self.layer_ids != tuple(range(len(layers))): - mapping = [-1] * (self.layer_ids[-1] + 1) - for slot, layer_idx in enumerate(self.layer_ids): - mapping[layer_idx] = slot - self.layer_id_to_slot = torch.tensor( - mapping, - dtype=torch.int64, - device=first.w13.device, - ) - # Resolve both operators once, before dispatch can receive work. - self.w13_op = torch.ops.afd_ascend.gmm_swiglu_quant_v2_layered - self.w2_op = torch.ops.afd_ascend.grouped_matmul_layered - - def geometry(self) -> tuple[int, int, int]: - """Return (experts, hidden, intermediate) from the loaded weight layout. - - INT4 weights pack eight nibbles per INT32, so the stored last dimension - is a quarter of the gate/up width; INT8 stores that width directly. - Callers that need only the logical geometry must not re-derive it. - """ - experts, hidden, last = self.w13[0].shape - if self.quantization == "W8A8": - if last % 2: - raise ValueError("layered GMM requires an even W13 gate/up width") - return experts, hidden, last // 2 - return experts, hidden, last * INT4_PER_INT32 // 2 - - def graph_signature(self, capacity: int) -> str: - """Compare static layouts across ranks without local device addresses.""" - tensors = ( - self.w13, - self.w2, - self.w13_scale, - self.w2_scale, - self.w13_bias, - self.w2_bias, - ) - return json.dumps( - { - "capacity": capacity, - "layers": self.layer_ids, - "weights": [ - (tuple(group[0].shape), str(group[0].dtype), group[0].stride()) - if group - else None - for group in tensors - ], - "per_channel": self.per_channel, - "quantization": self.quantization, - "limit": self.swiglu_limit, - "routed_scaling": self.routed_scaling_factor, - "output_dtype": str(self.output_dtype), - }, - sort_keys=True, - ) - - @staticmethod - def _validate_layer(layer: LayeredFFNWeights) -> None: - prefix = f"layered GMM layer {layer.layer_idx}: " - if layer.quantization == "W8A8": - AsyncCAMLayeredExecutor._validate_w8a8_layer(layer) - return - if layer.quantization != "W4A8": - raise ValueError(prefix + "requires W4A8 or W8A8 dynamic quantization") - if layer.w13.dim() != 3 or layer.w2.dim() != 3: - raise ValueError(prefix + "weights must be packed [E, K, N/8] tensors") - experts, hidden, packed_intermediate = layer.w13.shape - intermediate = packed_intermediate * INT4_PER_INT32 // 2 - if layer.w2.shape != (experts, intermediate, hidden // INT4_PER_INT32): - raise ValueError( - prefix + "W13/W2 hidden and intermediate dimensions disagree" - ) - for name, tensor in (("w13", layer.w13), ("w2", layer.w2)): - if tensor.dtype != torch.int32 or not tensor.is_contiguous(): - raise ValueError( - prefix + name + " requires contiguous INT32-packed INT4 storage" - ) - for name, scale, width, k in ( - ("w13_scale", layer.w13_scale, 2 * intermediate, hidden), - ("w2_scale", layer.w2_scale, hidden, intermediate), - ): - expected_rank = 2 if name == "w13_scale" and layer.per_channel else 3 - if ( - scale.dtype != torch.int64 - or scale.dim() != expected_rank - or scale.shape[0] != experts - or scale.shape[-1] != width - ): - raise ValueError( - prefix + name + " has incompatible shape or encoded INT64 dtype" - ) - if scale.dim() == 3: - groups = scale.shape[1] - if groups <= 0 or k % groups or (layer.per_channel and groups != 1): - raise ValueError( - prefix + name + " has incompatible quantization groups" - ) - for name, bias, width in ( - ("w13_scale_bias", layer.w13_bias, 2 * intermediate), - ("w2_scale_bias", layer.w2_bias, hidden), - ): - if ( - bias is None - or bias.dtype != torch.float32 - or bias.shape != (experts, width) - ): - raise ValueError(prefix + name + " requires [E, N] FP32 compensation") - tensors = ( - layer.w2, - layer.w13_scale, - layer.w2_scale, - cast(torch.Tensor, layer.w13_bias), - cast(torch.Tensor, layer.w2_bias), - ) - if any(t.device != layer.w13.device for t in tensors): - raise ValueError( - prefix + "all weights, scales and compensation must share a device" - ) - - @staticmethod - def _validate_w8a8_layer(layer: LayeredFFNWeights) -> None: - prefix = f"layered GMM layer {layer.layer_idx}: " - if layer.w13.dim() != 3 or layer.w2.dim() != 3: - raise ValueError(prefix + "W8A8 weights require [E, K, N] tensors") - experts, hidden, gate_up = layer.w13.shape - if gate_up % 2 or layer.w2.shape != (experts, gate_up // 2, hidden): - raise ValueError( - prefix + "W13/W2 hidden and intermediate dimensions disagree" - ) - if any( - t.dtype != torch.int8 or not t.is_contiguous() - for t in (layer.w13, layer.w2) - ): - raise ValueError(prefix + "W8A8 requires contiguous INT8 weights") - if not layer.per_channel or layer.w13_scale.dtype != torch.float32: - raise ValueError(prefix + "W8A8 requires FP32 per-channel W13 scales") - if layer.w13_scale.shape != (experts, gate_up): - raise ValueError(prefix + "W8A8 w13_scale requires [E, 2I]") - if layer.w2_scale.shape != (experts, hidden): - raise ValueError(prefix + "W8A8 w2_scale requires [E, H] scales") - if ( - layer.w2_scale.dtype != torch.bfloat16 - or layer.output_dtype != torch.bfloat16 - ): - raise ValueError(prefix + "W8A8 requires BF16 W2 scales and output") - if layer.w13_bias is not None or layer.w2_bias is not None: - raise ValueError(prefix + "W8A8 does not use W4A8 compensation") - if any( - t.device != layer.w13.device - for t in (layer.w2, layer.w13_scale, layer.w2_scale) - ): - raise ValueError(prefix + "all weights and scales must share a device") - - def __call__( - self, - hidden_states: torch.Tensor, - dynamic_scales: torch.Tensor, - group_list: torch.Tensor, - batch_info: torch.Tensor, - *, - retained_tensors: list[torch.Tensor] | None = None, - ) -> torch.Tensor: - layer_index = batch_info[CAM_LAYER_INDEX : CAM_LAYER_INDEX + 1] - if self.layer_id_to_slot is not None: - layer_index = self.layer_id_to_slot.index_select(0, layer_index) - else: - # The CANN tiling path checks the storage shape. A slice still - # carries batch_info's larger storage, so materialize one element. - layer_index = layer_index.clone() - activation, scale = self.w13_op( - hidden_states, - self.w13, - self.w13_scale, - self.w13_bias, - dynamic_scales, - group_list, - layer_index, - dequant_mode=0 if self.per_channel else 1, - quant_mode=0, - group_list_type=1, - swiglu_limit=self.applied_swiglu_limit, - ) - output = self.w2_op( - [activation], - self.w2, - self.w2_bias, - self.w2_scale, - layer_index, - group_list, - per_token_scale=scale, - group_list_type=1, - split_item=3, - output_dtype=self.output_dtype, - )[0] - if retained_tensors is not None: - # A graph session owns these capture-time addresses through its - # last replay. Eager calls retain only the usual stream lifetimes. - retained_tensors.extend((layer_index, activation, scale, output)) - if self.routed_scaling_factor != 1.0: - output.mul_(self.routed_scaling_factor) - return output diff --git a/afd_plugin/model_executor/npu/async_cam_mega.py b/afd_plugin/model_executor/npu/async_cam_mega.py deleted file mode 100644 index 1370346de..000000000 --- a/afd_plugin/model_executor/npu/async_cam_mega.py +++ /dev/null @@ -1,92 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""One native launch drains an exclusive quantized transaction per layer.""" - -from __future__ import annotations - -import torch - -from afd_plugin.model_executor.npu.async_cam_layered import AsyncCAMLayeredExecutor -from afd_plugin.validation import cam_receive_capacity - -# Mirror Pull::MetadataWords in afd_pull_layout.h (all fields are int64_t): -# [total, source, layer, round_begin, round_end], tp source prefixes, -# tp * experts original counts, then [session, sequence]. -METADATA_WORD_BYTES = torch.int64.itemsize -METADATA_FIXED_WORDS = 7 -METADATA_PREFIX_WORDS_PER_SOURCE = 1 -# NativeMegaWorkspace in afd_native_mega_workspace.h aligns scratch regions to -# 512 bytes. Match its counts offset, including padding after batch metadata. -WORKSPACE_ALIGNMENT_BYTES = 512 - - -class AsyncCAMMegaExecutor: - """Pull transport with a compiled W4A8 or W8A8 execution plan.""" - - def __init__(self, weights: AsyncCAMLayeredExecutor, connector) -> None: - if not weights.per_channel: - raise ValueError("native Mega requires per-channel W4A8/W8A8 weights") - self.weights = weights - self.connector = connector - self.capacity = cam_receive_capacity() - experts, hidden, intermediate = weights.geometry() - self.config = [ - connector.max_num_batched_tokens, - hidden, - intermediate, - connector.topk, - connector.ffn_size, - connector.attn_size, - experts, - connector.world_rank, - connector.topology.world_size, - connector.tp_size, - self.capacity, - 1 if weights.quantization == "W4A8" else 2, - 1, - len(weights.layer_ids), - ] - device = weights.w13[0].device - size = torch.ops.afd_ascend.afd_mega_moe_pull_workspace_size( - self.capacity, hidden, intermediate, experts, connector.tp_size - ) - self.scratch = torch.empty(size, dtype=torch.uint8, device=device) - metadata_words = METADATA_FIXED_WORDS + connector.tp_size * ( - METADATA_PREFIX_WORDS_PER_SOURCE + experts - ) - metadata_bytes = METADATA_WORD_BYTES * metadata_words - count_offset = ( - (metadata_bytes + WORKSPACE_ALIGNMENT_BYTES - 1) - // WORKSPACE_ALIGNMENT_BYTES - * WORKSPACE_ALIGNMENT_BYTES - ) - count_bytes = METADATA_WORD_BYTES * experts - self.batch_info = self.scratch[:metadata_bytes].view(torch.int64) - self.group_list = self.scratch[count_offset : count_offset + count_bytes].view( - torch.int64 - ) - self.anchor = torch.empty(1, dtype=torch.bfloat16, device=device) - mapping = [-1] * (weights.layer_ids[-1] + 1) - for slot, layer_id in enumerate(weights.layer_ids): - mapping[layer_id] = slot - self.layer_map = torch.tensor(mapping, dtype=torch.int64, device=device) - self.op = torch.ops.afd_ascend.afd_mega_moe_pull_native - - def __call__(self) -> torch.Tensor: - c, w = self.connector, self.weights - return self.op( - self.anchor, - c.comm_args, - self.layer_map, - self.scratch, - w.w13, - w.w2, - w.w13_scale, - w.w2_scale, - w.w13_bias if w.quantization == "W4A8" else [], - w.w2_bias if w.quantization == "W4A8" else [], - self.config, - c.group_name, - w.routed_scaling_factor, - w.applied_swiglu_limit, - ) diff --git a/afd_plugin/v1/worker/attention_model_runner_v2.py b/afd_plugin/v1/worker/attention_model_runner_v2.py index 3ff495685..e6423bfeb 100644 --- a/afd_plugin/v1/worker/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/attention_model_runner_v2.py @@ -41,7 +41,7 @@ AFDForwardContextMetadata, AFDRoutedExpertsConnector, ) -from afd_plugin.model_executor.models.forward_context import use_afd_metadata_provider +from afd_plugin.forward_context import use_afd_metadata_provider from afd_plugin.v1.worker.attention_metadata import ( AFDMetadataProviderMixin, _resolve_world_ranks, diff --git a/afd_plugin/v1/worker/npu/async_ffn_model_runner.py b/afd_plugin/v1/worker/npu/async_ffn_model_runner.py index 8ce1e175a..0f020228d 100644 --- a/afd_plugin/v1/worker/npu/async_ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/async_ffn_model_runner.py @@ -30,11 +30,14 @@ POLL_SECONDS, FFNGraphPlan, ) -from afd_plugin.model_executor.npu.async_cam_layered import ( +from afd_plugin.model_executor.ffn.compute import ( AsyncCAMLayeredExecutor, - LayeredFFNWeights, + AsyncCAMMegaExecutor, +) +from afd_plugin.model_executor.ffn.expert_weights import ( + QuantizedExpertWeights, + expert_weight_signature, ) -from afd_plugin.model_executor.npu.async_cam_mega import AsyncCAMMegaExecutor from afd_plugin.validation import ( cam_receive_capacity, requested_ffn_mode, @@ -52,7 +55,9 @@ class FFNModel(Protocol): def get_experts_layer_indices(self) -> tuple[int, ...]: ... - def get_async_cam_layered_weights(self) -> tuple[list[LayeredFFNWeights], str]: ... + def get_quantized_expert_weights( + self, + ) -> tuple[list[QuantizedExpertWeights], str]: ... class AsyncNPUFFNModelRunner: @@ -169,7 +174,6 @@ def initialize_afd_connector(self) -> None: return assert self.model is not None executor_error = signature = "" - config = self.connector.extra_info try: models = (self.model,) + ( (self.draft_model,) if self.draft_model is not None else () @@ -180,17 +184,9 @@ def initialize_afd_connector(self) -> None: if len(layer_ids) != len(set(layer_ids)): raise ValueError("Target and draft remote MoE layer IDs overlap") self.layer_ids = tuple(sorted(layer_ids)) - self._initialize_layered_executor() - assert self._layered_executor is not None - signature = self._layered_executor.graph_signature(cam_receive_capacity()) - if config.ffn_mega_moe: - self._mega_executor = AsyncCAMMegaExecutor( - self._layered_executor, self.connector - ) - signature = "native_mc2:" + signature + signature = self._initialize_ffn_executor() except Exception as exc: executor_error = str(exc) or type(exc).__name__ - executor = self._layered_executor self.connector.graph_startup.plan = FFNGraphPlan( requested=self.requested_mode, executor_error=executor_error, @@ -199,9 +195,13 @@ def initialize_afd_connector(self) -> None: ) try: self.connector.init_afd_connector() + if self._mega_executor is not None: + assert self.connector.comm_args is not None + self._mega_executor.bind_communication( + self.connector.comm_args, self.connector.group_name + ) self.effective_mode = self.connector.graph_startup.mode if self.effective_mode == "full": - assert executor is not None chunks = self.connector.graph_startup_chunks() * ( self.connector.attn_size // self.connector.tp_size ) @@ -253,7 +253,7 @@ def _capture_graph(self) -> None: self._check_graph_live() if self._graph is not None: raise RuntimeError("Async FFN may only capture once") - assert self._layered_executor is not None + assert self._layered_executor is not None or self._mega_executor is not None self._graph = torch.npu.NPUGraph() try: with torch.npu.graph( @@ -263,6 +263,7 @@ def _capture_graph(self) -> None: self._graph_payload = None self._graph_output = self._mega_executor() else: + assert self._layered_executor is not None self._graph_payload = self.connector.recv_attn_output() state = cast( AFDAsyncTransferState, self._graph_payload.context.states @@ -337,13 +338,13 @@ def _close_graph(self) -> None: del self._graph_pool, self._graph_stream self._graph_closed = True - def _initialize_layered_executor(self) -> None: + def _initialize_ffn_executor(self) -> str: assert self.model is not None - layers, reason = self.model.get_async_cam_layered_weights() + layers, reason = self.model.get_quantized_expert_weights() if not layers: raise ValueError(f"layered GMM target weights are unsupported: {reason}") if self.draft_model is not None: - draft_layers, reason = self.draft_model.get_async_cam_layered_weights() + draft_layers, reason = self.draft_model.get_quantized_expert_weights() if not draft_layers: raise ValueError(f"layered GMM draft weights are unsupported: {reason}") layers = [*layers, *draft_layers] @@ -356,6 +357,18 @@ def _initialize_layered_executor(self) -> None: "layered GMM weights do not cover the remote MoE layer IDs" ) ensure_cam_async_ops_available() + if self.connector.extra_info.ffn_mega_moe: + self._mega_executor = AsyncCAMMegaExecutor( + layers, + topology=self.connector.topology, + max_num_batched_tokens=self.connector.max_num_batched_tokens, + topk=self.connector.topk, + tp_size=self.connector.tp_size, + capacity=cam_receive_capacity(), + ) + return "native_mc2:" + expert_weight_signature( + layers, cam_receive_capacity() + ) self._layered_executor = AsyncCAMLayeredExecutor(layers) logger.info( "Async CAM layered GMM quant=%s layers=%d mode=%s layer_mapping=%s " @@ -367,6 +380,7 @@ def _initialize_layered_executor(self) -> None: layers[0].swiglu_limit, layers[0].quantization == "W8A8" and bool(layers[0].swiglu_limit), ) + return expert_weight_signature(layers, cam_receive_capacity()) @torch.inference_mode() def execute_connector_driven_step(self, shutdown_requested: bool = False) -> bool: @@ -379,7 +393,7 @@ def execute_connector_driven_step(self, shutdown_requested: bool = False) -> boo self._graph_stream.wait_stream(torch.npu.current_stream(self.device)) self._first_step = False if self._graph is None: - # Eager and graph modes use the same layered executor. + # Eager and graph modes use the same selected compute implementation. self._ffn_forward_connector_driven() self.session.after_step() return True diff --git a/afd_plugin/model_executor/models/npu/deepseek_attention_metadata.py b/afd_plugin/v1/worker/npu/attention_metadata.py similarity index 93% rename from afd_plugin/model_executor/models/npu/deepseek_attention_metadata.py rename to afd_plugin/v1/worker/npu/attention_metadata.py index b2dacdc86..ebb0fde54 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_attention_metadata.py +++ b/afd_plugin/v1/worker/npu/attention_metadata.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""DeepSeek attention metadata ownership for staged NPU execution.""" +"""Ascend attention metadata and RoPE buffer ownership for staged execution.""" from __future__ import annotations @@ -21,11 +21,11 @@ from vllm_ascend.ops.rope_dsv4 import RopeDataProxy, get_cos_and_sin_dsa -def isolate_deepseek_attention_builder_inputs( +def isolate_attention_builder_inputs( builder: AttentionMetadataBuilder, metadata: AscendCommonAttentionMetadata, ) -> None: - """Detach common metadata that a DeepSeek builder mutates in place. + """Detach common metadata that an Ascend attention builder mutates in place. Async CAM builds the full batch before both execution stages. Stage builders therefore must not write through views into the full-batch @@ -49,7 +49,7 @@ def isolate_deepseek_attention_builder_inputs( metadata.block_table_tensor = metadata.block_table_tensor.clone() -def materialize_deepseek_attention_metadata( +def materialize_attention_metadata( metadata: AttentionMetadata, input_positions: torch.Tensor, num_input_tokens: int, @@ -153,7 +153,7 @@ def _own_rope( return buffers[field][: value.shape[0]].copy_(value) -def materialize_deepseek_attention_metadata_by_layer( +def materialize_attention_metadata_by_layer( metadata_by_layer: Mapping[str, AttentionMetadata], input_positions: torch.Tensor, num_input_tokens: int, @@ -168,7 +168,7 @@ def materialize_deepseek_attention_metadata_by_layer( metadata_id = id(metadata) if metadata_id in materialized_ids: continue - materialize_deepseek_attention_metadata( + materialize_attention_metadata( metadata, input_positions, num_input_tokens, @@ -181,7 +181,7 @@ def materialize_deepseek_attention_metadata_by_layer( __all__ = [ - "isolate_deepseek_attention_builder_inputs", - "materialize_deepseek_attention_metadata", - "materialize_deepseek_attention_metadata_by_layer", + "isolate_attention_builder_inputs", + "materialize_attention_metadata", + "materialize_attention_metadata_by_layer", ] 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 8f791cca7..d3a2db741 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner_v2.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner_v2.py @@ -42,7 +42,7 @@ AFDConnectorFactory, AFDForwardContextMetadata, ) -from afd_plugin.model_executor.models.forward_context import use_afd_metadata_provider +from afd_plugin.forward_context import use_afd_metadata_provider from afd_plugin.v1.worker.attention_metadata import ( AFDMetadataProviderMixin, _resolve_world_ranks, diff --git a/afd_plugin/v1/worker/npu/ubatch_runner.py b/afd_plugin/v1/worker/npu/ubatch_runner.py index 0c3fdcc5b..255466562 100644 --- a/afd_plugin/v1/worker/npu/ubatch_runner.py +++ b/afd_plugin/v1/worker/npu/ubatch_runner.py @@ -54,8 +54,8 @@ use_npu_ubatch_contexts, use_npu_ubatch_output_merge, ) -from afd_plugin.model_executor.models.npu.deepseek_attention_metadata import ( - materialize_deepseek_attention_metadata_by_layer, +from afd_plugin.v1.worker.npu.attention_metadata import ( + materialize_attention_metadata_by_layer, ) from afd_plugin.v1.worker.npu.mla_graph import ( merge_mla_graph_params, @@ -208,7 +208,7 @@ def _prepare_attn( full_graph_mode=cudagraph_mode == CUDAGraphMode.FULL, ) if cudagraph_mode == CUDAGraphMode.FULL: - materialize_deepseek_attention_metadata_by_layer( + materialize_attention_metadata_by_layer( metadata, input_batch.positions, num_input_tokens, @@ -216,7 +216,7 @@ def _prepare_attn( rope_capacity=self.vllm_config.scheduler_config.max_num_batched_tokens, ) else: - materialize_deepseek_attention_metadata_by_layer( + materialize_attention_metadata_by_layer( metadata, input_batch.positions, num_input_tokens ) return metadata diff --git a/docs/design/module/attention_runtime.md b/docs/design/module/attention_runtime.md index 57188eeda..2412771fc 100644 --- a/docs/design/module/attention_runtime.md +++ b/docs/design/module/attention_runtime.md @@ -14,6 +14,7 @@ primary_code_paths: - "afd_plugin/v1/worker/npu/attention_worker.py" - "afd_plugin/v1/worker/npu/dspark_speculator.py" - "afd_plugin/v1/worker/npu/dspark_graph.py" + - "afd_plugin/v1/worker/npu/attention_metadata.py" related_code_paths: - "afd_plugin/connectors/**" - "afd_plugin/model_executor/**" @@ -32,7 +33,8 @@ validation_paths: - "tests/unit/v1/worker/test_npu_model_runner_v2.py" - "tests/unit/v1/worker/npu/test_ubatch_runner.py" - "tests/unit/v1/worker/npu/test_attention_model_runner_v2.py" - - "tests/unit/model_executor/test_async_cam_dbo.py" + - "tests/unit/v1/worker/npu/test_attention_metadata.py" + - "tests/unit/model_executor/ffn/test_async_cam.py" - "tests/npu/test_mrv2_dbo.py" - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v4_flash/test_gpu.py" @@ -317,9 +319,8 @@ Enable vLLM native DBO with `--enable-dbo` on Attention. It: Native DBO owns request/token boundaries. Enable DBO on Attention only; Async FFN rejects native DBO flags and continues its receive loop. -`models/npu/async_cam_layout.py` retains only communication layout -conversion, and `models/npu/async_cam_dbo.py` supplies the single-window receive -hook. FlashComm1/SP is Attention-local and +`model_executor/ffn/async_cam.py` owns CAM communication layout conversion and +the single-window receive hook, shared by the model adapters. FlashComm1/SP is Attention-local and FFN topology validation is independent. See the CAM async user guide for the supported topology and environment matrix. diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index 86a936d58..e9d72f3bf 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/docs/design/module/compatibility_and_patches.md @@ -137,7 +137,7 @@ not the package dependency policy. | [`npu/ascend_platform.py`](../../../afd_plugin/compat/patches/npu/ascend_platform.py): `NPUPlatform.check_and_update_config` | Snapshots AFD DBO state, runs upstream normalization, and restores configured `enable_dbo`, `ubatch_size`, and `all2all_backend` in `finally`; non-AFD behavior is unchanged. | Installed through the config facade during AFD NPU normalization and verified again by the worker runtime facade; no version guard. Saves the original on the class and uses a class sentinel. The runtime facade caches success only after the wrapper is installed, so an early missing vLLM-Ascend import remains retryable. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) and [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py). | Remove when vLLM-Ascend recognizes plugin-owned DBO workers or no longer clears these fields. | | [`npu/ascend_config.py`](../../../afd_plugin/compat/patches/npu/ascend_config.py): `vllm_ascend.ascend_config.init_ascend_config`, loaded factory aliases, and the AFD DP child entry point | Excludes the AFD namespace from native constructor arguments; preserves config identity, validation, singleton publication, and cache refresh. DP children install the factory before native Ascend initialization. | Installed by the NPU config facade; rebinds already-loaded aliases. The DP entry point is selected only for AFD synchronous Attention or FFN. | [`test_ascend_config_lifecycle.py`](../../../tests/unit/compat/npu/test_ascend_config_lifecycle.py), [`test_ascend_engine_core_binding.py`](../../../tests/unit/compat/npu/test_ascend_engine_core_binding.py). | Remove when Ascend accepts plugin namespaces and loads general plugins before config validation in every DP child. | | [`npu/mla_graph.py`](../../../afd_plugin/compat/patches/npu/mla_graph.py): `vllm_ascend.attention.mla_v1.get_graph_params` | Resolves an AFD ubatch-owned `GraphParams` registry from the active forward context and otherwise delegates to the saved upstream process-global resolver. | Called through `apply_afd_ascend_patches_if_needed` during NPU worker startup; no version guard. Saves the original on the upstream module and uses a module sentinel. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) covers AFD-context resolution, upstream fallback, and idempotence; [`test_npu_mla_graph.py`](../../../tests/unit/v1/worker/test_npu_mla_graph.py) covers the owning graph lifecycle. | Remove when vLLM-Ascend accepts a forward-context-local MLA graph registry or exposes an equivalent resolver hook. | -| [`npu/mrv2_ubatching.py`](../../../afd_plugin/compat/patches/npu/mrv2_ubatching.py): native ubatch factory, dispatch, child DP metadata, input slicing, context and stream aliases | Builds `NPUUBatchRunner`, enables Async CAM DBO at DP1, omits Async child DP coordination, applies native truncation to Ascend CPU sequence lengths, preserves model-specific PP slices/nested auxiliary outputs, and provides NPU stream/failure handling during eager and capture. Native MRV2 owns execution and graph lifecycle. | Scoped to AFD initialization, preparation or execution; every replaced alias is restored in `finally`. The scope temporarily changes module symbols, so it is not an independent concurrent-runner API. | [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py), [`test_async_cam_dbo.py`](../../../tests/unit/model_executor/test_async_cam_dbo.py), [`test_mrv2_dbo.py`](../../../tests/npu/test_mrv2_dbo.py). | Remove when native vLLM/Ascend supports NPU child builders and streams, optional child DP metadata, DP1 DBO and peer cancellation. | +| [`npu/mrv2_ubatching.py`](../../../afd_plugin/compat/patches/npu/mrv2_ubatching.py): native ubatch factory, dispatch, child DP metadata, input slicing, context and stream aliases | Builds `NPUUBatchRunner`, enables Async CAM DBO at DP1, omits Async child DP coordination, applies native truncation to Ascend CPU sequence lengths, preserves model-specific PP slices/nested auxiliary outputs, and provides NPU stream/failure handling during eager and capture. Native MRV2 owns execution and graph lifecycle. | Scoped to AFD initialization, preparation or execution; every replaced alias is restored in `finally`. The scope temporarily changes module symbols, so it is not an independent concurrent-runner API. | [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py), [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/test_async_cam.py), [`test_mrv2_dbo.py`](../../../tests/npu/test_mrv2_dbo.py). | Remove when native vLLM/Ascend supports NPU child builders and streams, optional child DP metadata, DP1 DBO and peer cancellation. | ## Temporary CUDA DeepSeek-V4 DBO fixes diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 8403e52b5..f2cd20450 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -64,7 +64,7 @@ depend on role worker implementations. | Area | Source | Focused validation | | --- | --- | --- | | Base surface and lazy factory | [`base.py`](../../../afd_plugin/connectors/base.py), [`factory.py`](../../../afd_plugin/connectors/factory.py) | [P2P factory and construction tests](../../../tests/unit/connectors/test_p2p_connector.py), [metadata contracts](../../../tests/unit/connectors/test_metadata.py) | -| Payloads and control-plane codec | [`metadata.py`](../../../afd_plugin/connectors/metadata.py) | Connector unit tests and [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | +| Payloads and control-plane codec | [`metadata.py`](../../../afd_plugin/connectors/metadata.py) | Connector unit tests and [`test_forward_context.py`](../../../tests/unit/test_forward_context.py) | | CUDA P2P | [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py), [`topology.py`](../../../afd_plugin/distributed/topology.py) | [`test_p2p_connector.py`](../../../tests/unit/connectors/test_p2p_connector.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | | CUDA DeepEP M2N | [`gpu/deepep_m2n/`](../../../afd_plugin/connectors/gpu/deepep_m2n/), `AFDRoutedExpertsConnector` in [`base.py`](../../../afd_plugin/connectors/base.py) | [`test_deepep_m2n_connector.py`](../../../tests/unit/connectors/test_deepep_m2n_connector.py), [GPU device tests](../../../tests/gpu/README.md), [Qwen3 MoE E2E](../../../tests/e2e/models/qwen3_moe/test_qwen3_moe.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) (`test_deepseek_v2_lite_deepep_m2n`) | | Ascend CAMP2P | [`npu/camp2p/connector.py`](../../../afd_plugin/connectors/npu/camp2p/connector.py) | [`test_camp2p_connector.py`](../../../tests/unit/connectors/test_camp2p_connector.py), [`test_camp2p_token_ids.py`](../../../tests/unit/connectors/test_camp2p_token_ids.py), [`test_external_routing.py`](../../../tests/unit/connectors/test_external_routing.py), [DeepSeek-V2-Lite E2E](../../../tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py) | diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md index 1992482ca..dff666576 100644 --- a/docs/design/module/execution_platforms.md +++ b/docs/design/module/execution_platforms.md @@ -35,7 +35,7 @@ validation_paths: - "tests/unit/v1/worker/test_npu_mla_graph.py" - "tests/unit/v1/worker/test_npu_runtime.py" - "tests/unit/v1/worker/test_npu_model_runner_v2.py" - - "tests/unit/model_executor/test_async_cam_dbo.py" + - "tests/unit/model_executor/ffn/test_async_cam.py" - "tests/npu/**" - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v4_flash/test_gpu.py" diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index 4b35c9074..5179cca6e 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -28,7 +28,7 @@ depends_on: - "compatibility_and_patches.md" validation_paths: - "tests/unit/ffn_faas/test_*.py" - - "tests/unit/model_executor/test_async_cam_layered.py" + - "tests/unit/model_executor/ffn/test_compute.py" - "tests/npu/test_async_cam_layered_w8a8.py" - "tests/unit/v1/worker/test_ffn_metadata.py" - "tests/unit/v1/worker/test_ffn_model_runner.py" @@ -142,7 +142,7 @@ and CAM runtime composition live under `ffn_faas/platforms/npu/`. communication, computation and Graph execution to the Async CAM runtime. A shared cross-platform, cross-connector Step Loop and a pure `compute(item)` backend are not implemented. The runtime consumes the existing -`LayeredFFNWeights` interface with W4A8 weights, without per-step model selection; +`QuantizedExpertWeights` contract with W4A8 weights, without per-step model selection; there is no separate V4 execution runtime. @@ -322,6 +322,23 @@ synchronization remains after a work-item group. The `swiglu_limit` behavior and validation scope are documented in the [CAM async guide](../../npu/CAM_ASYNC_CONNECTOR_USER_GUIDE.md#w4a8-and-w8a8-layered-gmm). +## Async compute ownership + +The Async runner loads target/draft `QuantizedExpertWeights` through the model +export contract, verifies layer coverage and selects Layered or Mega directly. +`model_executor/ffn/expert_weights.py` owns shared shape validation, logical +geometry and graph signatures. `compute.py` owns native operator handles, +workspace, zero bias and layer maps. There is no execution layer beneath the +ModelRunner and no access to private model parameters from compute. + +Mega consumes the exported tensors independently of Layered. It prepares its +workspace before startup negotiation and binds live `comm_args` and the group +name after connector initialization, before warmup/capture. Binding failures +are published through Startup and retain resources. Negotiation preserves the +existing signature format, including `native_mc2:` and the W8A8 zero-bias layout; +W8A8 Mega does not allocate Layered's unused zero bias. Eager/FULL transaction, +graph and shutdown orchestration remain runner-owned. + ## Control-plane-driven forward The current runner contract for one control payload is: diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 7af50ca4a..1bc31dd25 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -7,6 +7,7 @@ owners: - "@jiangkuaixue123" primary_code_paths: - "afd_plugin/model_executor/**/*.py" + - "afd_plugin/forward_context.py" related_code_paths: - "afd_plugin/connectors/metadata.py" - "afd_plugin/v1/worker/dbo.py" @@ -17,6 +18,7 @@ depends_on: - "execution_platforms.md" validation_paths: - "tests/unit/model_executor/**" + - "tests/unit/test_forward_context.py" - "tests/gpu/test_attention_routing.py" - "tests/e2e/models/**" - "tests/e2e/accuracy/**" @@ -55,22 +57,73 @@ Model integration consumes plugin configuration, connector payload contracts, and platform mechanisms. It must not reach into concrete worker instances or make a backend-specific worker class the shared model API. +## Shared FFN components + +`model_executor/models/` owns model structure and checkpoint policy. +`model_executor/ffn/` contains four modules with separate responsibilities: + +| Component | Responsibility | Dependencies | +| --- | --- | --- | +| `remote.py` | Connector-independent dense/MoE proxies and routed TP padding | Connector contracts, public vLLM helpers and package forward context; no CAM, concrete model or compute imports | +| `async_cam.py` | CAM model-side TP/SP layout conversion, padding, output restoration and native DBO receive handoff | Connector contracts and native DBO hooks; no concrete model or compute imports | +| `expert_weights.py` | Loaded `QuantizedExpertWeights`, homogeneous layout validation, geometry and graph signatures | Torch tensors; no native operator lookup or workspace allocation | +| `compute.py` | Independent NPU Layered and Mega native implementations and their buffers | Expert weights and explicit configuration/communication arguments; no model internals | + +```mermaid +flowchart TD + Worker[Worker] --> Runner[ModelRunner] + Runner --> Model[Concrete model adapter] + Model --> Remote[ffn/remote.py] + Remote --> Connector[Connector / Session] + Model --> CAM[ffn/async_cam.py] + CAM --> Connector + Model --> Weights[ffn/expert_weights.py] + Runner --> Compute[ffn/compute.py: Layered OR Mega] + Compute --> Weights + Compute --> Native[Native NPU operators] +``` + +DeepSeek V2, Qwen3, Qwen3.5 and Ascend V4 use the shared remote components. +The CAM helpers stay in `ffn/async_cam.py` because they adapt model tensors +to remote FFN calls across model families. Connector/Session owns transport +windows, pending work and Push/Pull protocol state. The generic proxies do +not import the CAM module, and neither module imports compute. +Gate, Hash routing, shared experts, dense-layer behavior and DSpark target/draft +identities remain model-owned. The Async runner receives loaded weights through +`get_quantized_expert_weights()` and selects one compute implementation. Mega +neither constructs Layered nor keeps the connector; its live communication +arguments are bound after connector initialization. + +Five model adapters share checkpoint path parsing and streaming role filtering +in `models/_weight_utils.py`; each keeps its classifier. DSpark retains its +specialized draft-prefix policy. Role-aware constructor copies remain where +upstream provides no hook for omitting one role's modules. Existing `model_cls` +and decoder-class injection are used where supported; no global model factory, +registry or new monkey patch is added. + +Shared metadata installation and access live in `afd_plugin/forward_context.py`, +used by GPU/NPU runners, models and remote proxies. Ascend attention metadata +and RoPE buffer ownership live in `v1/worker/npu/attention_metadata.py`: the +NPU ubatch runner owns stage isolation and stable FULL-graph buffer addresses. +SFA/MLA/DSA-specific handling stays explicit inside that worker-side module. + ## Implementation evidence | Area | Source | Focused validation | | --- | --- | --- | | Registration map | [`afd_plugin/__init__.py`](../../../afd_plugin/__init__.py) | [`test_package.py`](../../../tests/unit/package/test_package.py) | -| Role-aware model and weight loading | [`deepseek_v2.py`](../../../afd_plugin/model_executor/models/deepseek_v2.py) | [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py), model and accuracy E2E suites | +| Role-aware model and weight loading | [`deepseek_v2.py`](../../../afd_plugin/model_executor/models/deepseek_v2.py), [`_weight_utils.py`](../../../afd_plugin/model_executor/models/_weight_utils.py) | [`test_deepseek_v2_construction.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_construction.py), [`test_weight_utils.py`](../../../tests/unit/model_executor/models/test_weight_utils.py), model and accuracy E2E suites | | DeepSeek V4 CUDA role boundary | [`deepseek_v4.py`](../../../afd_plugin/model_executor/models/deepseek_v4.py) | [`test_deepseek_v4_construction.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_construction.py), [`test_deepseek_v4_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_proxy.py), [`test_deepseek_v4_weight_policy.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_weight_policy.py) | | DeepSeek V4 Ascend Hash-id boundary | [`npu/deepseek_v4.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4.py), [`npu/deepseek_v4_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py) | [`test_deepseek_v4_hash_ids.py`](../../../tests/unit/model_executor/test_deepseek_v4_hash_ids.py), [NPU model role tests](../../../tests/unit/model_executor/models/npu/test_deepseek_v4.py), ids-mode cases in [`test_camp2p_token_ids.py`](../../../tests/unit/connectors/test_camp2p_token_ids.py) | | Qwen3 MoE role-aware model and weight loading | [`qwen3_moe.py`](../../../afd_plugin/model_executor/models/qwen3_moe.py) | [`test_qwen3_moe_construction.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_construction.py), [`test_qwen3_moe_weight_policy.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py) | -| CUDA remote-experts boundary | [`deepseek_v2.py`](../../../afd_plugin/model_executor/models/deepseek_v2.py), [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py) | [`test_p2p_experts_contract.py`](../../../tests/unit/connectors/test_p2p_experts_contract.py), [`test_deepseek_v2_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_proxy.py) | -| Forward-context adapter | [`forward_context.py`](../../../afd_plugin/model_executor/models/forward_context.py) | [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | -| NPU Async CAM native DBO handoff | [`npu/async_cam_dbo.py`](../../../afd_plugin/model_executor/models/npu/async_cam_dbo.py), [`npu/ubatch_runner.py`](../../../afd_plugin/v1/worker/npu/ubatch_runner.py) | [`test_async_cam_dbo.py`](../../../tests/unit/model_executor/test_async_cam_dbo.py), [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py) | -| Ascend Attention-side gate | [`npu/deepseek_v2_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py) | Attention-gate unit cases in [`test_forward_context.py`](../../../tests/unit/model_executor/models/test_forward_context.py) | +| CUDA remote-experts boundary | [`ffn/remote.py`](../../../afd_plugin/model_executor/ffn/remote.py), [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py) | [`test_p2p_experts_contract.py`](../../../tests/unit/connectors/test_p2p_experts_contract.py), [`test_deepseek_v2_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_proxy.py) | +| Forward-context adapter | [`forward_context.py`](../../../afd_plugin/forward_context.py) | [`test_forward_context.py`](../../../tests/unit/test_forward_context.py) | +| NPU Async CAM native DBO handoff | [`ffn/async_cam.py`](../../../afd_plugin/model_executor/ffn/async_cam.py), [`npu/ubatch_runner.py`](../../../afd_plugin/v1/worker/npu/ubatch_runner.py) | [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/test_async_cam.py), [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py) | +| Ascend Attention-side gate | [`npu/deepseek_v2_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py) | [`test_deepseek_v2_attention_gate.py`](../../../tests/unit/model_executor/models/npu/test_deepseek_v2_attention_gate.py), profile I/O in [`test_deepseek_v2_async_cam_forward.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py) | +| Ascend CAM dispatch layout | [`ffn/async_cam.py`](../../../afd_plugin/model_executor/ffn/async_cam.py) | [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/test_async_cam.py) | +| Ascend stage metadata and RoPE ownership | [`worker/npu/attention_metadata.py`](../../../afd_plugin/v1/worker/npu/attention_metadata.py) | [`test_attention_metadata.py`](../../../tests/unit/v1/worker/npu/test_attention_metadata.py) | | Ascend external-routing MoE execution | [`npu/deepseek_v4_external_routing_moe.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4_external_routing_moe.py) | [`test_deepseek_v4_external_routing_moe_cpu.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_external_routing_moe_cpu.py) | -| Ascend CAM dispatch layout | [`npu/async_cam_layout.py`](../../../afd_plugin/model_executor/models/npu/async_cam_layout.py) | [`test_async_cam_layout.py`](../../../tests/unit/model_executor/test_async_cam_layout.py) | -| DeepSeek Ascend stage metadata ownership | [`npu/deepseek_attention_metadata.py`](../../../afd_plugin/model_executor/models/npu/deepseek_attention_metadata.py) | [`test_deepseek_attention_metadata.py`](../../../tests/unit/model_executor/test_deepseek_attention_metadata.py) | + | Ascend CAM orchestration | [`npu/deepseek_v4.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4.py) | [`test_async_cam_npu.py`](../../../tests/e2e/models/deepseek_v4_flash/test_async_cam_npu.py) | ## Model registration @@ -248,9 +301,8 @@ global provider. Async CAM reuses native DBO child slices, attention metadata and execution. The generic adapter installs `afd_metadata` on the native child context; CAM transfer metadata and -TP/SP layout conversion remain necessary for the communication boundary in -`models/npu/async_cam_layout.py`. The single-window receive hook lives in -`models/npu/async_cam_dbo.py`. Metadata ownership and the live connector reference +TP/SP layout conversion and the single-window receive hook share the +model-side communication boundary in `model_executor/ffn/async_cam.py`. Metadata ownership and the live connector reference remain **draft** under [#88](https://github.com/JiusiServe/afd-plugin/issues/88) and [#105](https://github.com/JiusiServe/afd-plugin/issues/105). diff --git a/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py b/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py index 80998d5cc..e9fe381ab 100644 --- a/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector -from afd_plugin.model_executor.npu.async_cam_layered import CAM_LAYER_INDEX +from afd_plugin.model_executor.ffn.compute import CAM_LAYER_INDEX from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker @@ -28,6 +28,12 @@ def snapshot(runner) -> dict: is_ffn = isinstance(runner, AsyncNPUFFNModelRunner) + executor = None + execution = "uninitialized" + if is_ffn: + executor = runner._mega_executor or runner._layered_executor + if executor is not None: + execution = "mega" if runner._mega_executor is not None else "layered" draft = runner.draft_model if is_ffn else runner.speculator.model if draft is None: return { @@ -59,11 +65,9 @@ def snapshot(runner) -> dict: "draft_kv_groups": kv_groups, **( { - "ffn_execution": "layered" - if runner._layered_executor is not None - else "uninitialized", - "executor_layer_ids": list(runner._layered_executor.layer_ids) - if runner._layered_executor is not None + "ffn_execution": execution, + "executor_layer_ids": list(executor.layer_ids) + if executor is not None else [], "ffn_mode": runner.effective_mode, "graph_capture_count": runner._graph_capture_count, diff --git a/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py b/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py index f758a481b..9795bdf13 100644 --- a/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py @@ -17,7 +17,7 @@ import torch -from afd_plugin.model_executor.npu.async_cam_layered import ( +from afd_plugin.model_executor.ffn.compute import ( CAM_LAYER_INDEX, AsyncCAMLayeredExecutor, ) @@ -32,8 +32,8 @@ class DSparkFFNGraphEvidence(DSparkEagerEvidence): class GraphComparison: def __init__(self, runner): self.runner = runner - target, _ = runner.model.get_async_cam_layered_weights() - draft, _ = runner.draft_model.get_async_cam_layered_weights() + target, _ = runner.model.get_quantized_expert_weights() + draft, _ = runner.draft_model.get_quantized_expert_weights() self.target = AsyncCAMLayeredExecutor(target) self.draft = AsyncCAMLayeredExecutor(draft) self.records = {} diff --git a/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py b/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py index ac1c31e15..1cb4bffc8 100644 --- a/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py @@ -10,11 +10,12 @@ from functools import wraps from pathlib import Path from types import SimpleNamespace +from weakref import WeakKeyDictionary import torch from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector -from afd_plugin.model_executor.npu.async_cam_mega import AsyncCAMMegaExecutor +from afd_plugin.model_executor.ffn.compute import AsyncCAMMegaExecutor from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( AFDNPUAttentionModelRunnerV2, @@ -31,7 +32,13 @@ def snapshot(runner) -> dict: "transport_live": connector.is_initialized, } if isinstance(runner, AsyncNPUFFNModelRunner): + execution = "uninitialized" + if runner._mega_executor is not None: + execution = "mega" + elif runner._layered_executor is not None: + execution = "layered" result.update( + execution=execution, mode=runner.effective_mode, requested=runner.requested_mode, capture_count=runner._graph_capture_count @@ -72,6 +79,10 @@ def afd_ffn_graph_snapshot(self) -> dict: def install_observers() -> None: + initialize = AsyncNPUFFNModelRunner.initialize_afd_connector + owners: WeakKeyDictionary[AsyncCAMMegaExecutor, CAMAsyncAFDConnector] = ( + WeakKeyDictionary() + ) complete = AsyncNPUFFNModelRunner._complete_graph send_eager = CAMAsyncAFDConnector.send_ffn_output execute_mega = AsyncCAMMegaExecutor.__call__ @@ -121,13 +132,21 @@ def completed(self): observe(self.connector, context) return result + @wraps(initialize) + def initialized(self): + result = initialize(self) + if self._mega_executor is not None: + owners[self._mega_executor] = self.connector + return result + @wraps(execute_mega) def mega_executed(self): result = execute_mega(self) - if self.connector.graph_startup.mode == "eager": + connector = owners.get(self) + if connector is not None and connector.graph_startup.mode == "eager": # Never read device headers while FULL is warming up or capturing. torch.npu.current_stream().synchronize() - observe(self.connector, mega_context(self)) + observe(connector, mega_context(self)) return result @wraps(send_eager) @@ -193,6 +212,7 @@ def ffn_stop_requested(self): record(self.model_runner, "stop_rpc_complete") return result + AsyncNPUFFNModelRunner.initialize_afd_connector = initialized AsyncNPUFFNModelRunner._complete_graph = completed CAMAsyncAFDConnector.send_ffn_output = sent AsyncCAMMegaExecutor.__call__ = mega_executed diff --git a/tests/npu/README.md b/tests/npu/README.md index e68210f4e..05b4221e8 100644 --- a/tests/npu/README.md +++ b/tests/npu/README.md @@ -74,8 +74,8 @@ Run the focused contract checks in the matching runtime environment: python3 -m pytest -q -ra \ tests/unit/compat/npu/test_runtime_config.py \ tests/unit/config/test_validation.py \ - tests/unit/model_executor/test_async_cam_dbo.py \ - tests/unit/model_executor/test_async_cam_layout.py \ + tests/unit/model_executor/ffn/test_async_cam.py \ + tests/unit/v1/worker/npu/test_attention_metadata.py \ tests/unit/v1/worker/test_npu_model_runner_v2.py \ tests/unit/v1/worker/test_npu_ubatch_outputs.py \ tests/unit/v1/worker/test_npu_mla_graph.py @@ -278,7 +278,7 @@ device and no collective. Run the CPU contracts first, then rebuild the matching extension/vendor library and run the device checks: ```bash -pytest -q tests/unit/test_envs.py tests/unit/model_executor/test_async_cam_layered.py +pytest -q tests/unit/test_envs.py tests/unit/model_executor/ffn/test_compute.py SOC_VERSION=910c pytest -q tests/npu/test_gmm_layered.py \ tests/npu/test_gmm_swiglu_quant_v2_layered.py ``` diff --git a/tests/npu/async_cam_runtime.py b/tests/npu/async_cam_runtime.py index 74c7e3c59..7d1a85c62 100644 --- a/tests/npu/async_cam_runtime.py +++ b/tests/npu/async_cam_runtime.py @@ -29,7 +29,7 @@ from afd_plugin.config import parse_afd_config from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector -from afd_plugin.model_executor.npu.async_cam_layered import LayeredFFNWeights +from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker @@ -81,7 +81,7 @@ def __init__( experts, ) self.layers.append( - LayeredFFNWeights( + QuantizedExpertWeights( layer_id, w13, w2, @@ -152,7 +152,7 @@ def make_weight( def get_experts_layer_indices(self): return LAYERS - def get_async_cam_layered_weights(self): + def get_quantized_expert_weights(self): return self.layers, "" @@ -312,7 +312,7 @@ def main(): ) print(json.dumps({"rank": rank, "stage": "weights_ready"}), flush=True) runner.initialize_afd_connector() - assert runner._layered_executor is not None + assert runner.layer_ids def work_state(): if args.mega: diff --git a/tests/npu/test_async_cam_layered_w4a8.py b/tests/npu/test_async_cam_layered_w4a8.py index 5df33587a..f8f20e90c 100644 --- a/tests/npu/test_async_cam_layered_w4a8.py +++ b/tests/npu/test_async_cam_layered_w4a8.py @@ -27,10 +27,8 @@ def test_layered_w4a8_valid_rows_and_capacity_tail( torch = pytest.importorskip("torch") torch_npu = pytest.importorskip("torch_npu") from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available - from afd_plugin.model_executor.npu.async_cam_layered import ( - AsyncCAMLayeredExecutor, - LayeredFFNWeights, - ) + from afd_plugin.model_executor.ffn.compute import AsyncCAMLayeredExecutor + from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights ensure_cam_async_ops_available() torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) @@ -86,7 +84,7 @@ def make_weight(k, n, factor, squeeze): w13, s13, b13, ref13 = make_weight(hidden, 2 * intermediate, factor, True) w2, s2, b2, ref2 = make_weight(intermediate, hidden, factor, False) layers.append( - LayeredFFNWeights( + QuantizedExpertWeights( idx, w13, w2, s13, s2, b13, b2, per_channel, swiglu_limit, 1.0 ) ) diff --git a/tests/npu/test_async_cam_layered_w8a8.py b/tests/npu/test_async_cam_layered_w8a8.py index a5cc30805..b04dc80f0 100644 --- a/tests/npu/test_async_cam_layered_w8a8.py +++ b/tests/npu/test_async_cam_layered_w8a8.py @@ -27,10 +27,8 @@ def w8a8_runtime(request): torch = pytest.importorskip("torch") torch_npu = pytest.importorskip("torch_npu") from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available - from afd_plugin.model_executor.npu.async_cam_layered import ( - AsyncCAMLayeredExecutor, - LayeredFFNWeights, - ) + from afd_plugin.model_executor.ffn.compute import AsyncCAMLayeredExecutor + from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights ensure_cam_async_ops_available() torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) @@ -54,7 +52,7 @@ def w8a8_runtime(request): generator=generator, ) layers.append( - LayeredFFNWeights( + QuantizedExpertWeights( layer_idx=layer_id, w13=torch_npu.npu_format_cast(w13.npu(), 29), w2=torch_npu.npu_format_cast(w2.npu(), 29), diff --git a/tests/npu/test_mrv2_dbo.py b/tests/npu/test_mrv2_dbo.py index 7a3fed6e2..a621785bb 100644 --- a/tests/npu/test_mrv2_dbo.py +++ b/tests/npu/test_mrv2_dbo.py @@ -42,7 +42,7 @@ from vllm_ascend.worker.v2.utils import torch_cuda_wrapper # noqa: E402 from afd_plugin.connectors import AFDTransferContext, AFDTransferMetadata # noqa: E402 -from afd_plugin.model_executor.models.npu.async_cam_dbo import ( # noqa: E402 +from afd_plugin.model_executor.ffn.async_cam import ( # noqa: E402 send_async_cam_remote_call, ) from afd_plugin.v1.worker.npu.ubatch_runner import NPUUBatchRunner # noqa: E402 diff --git a/tests/npu/test_mrv2_graph_inputs.py b/tests/npu/test_mrv2_graph_inputs.py index 0abf47aac..2485e2078 100644 --- a/tests/npu/test_mrv2_graph_inputs.py +++ b/tests/npu/test_mrv2_graph_inputs.py @@ -28,10 +28,10 @@ from vllm_ascend.ops import rope_dsv4 # noqa: E402 from vllm_ascend.ops.rope_dsv4 import RopeDataProxy # noqa: E402 -from afd_plugin.model_executor.models.npu import ( # noqa: E402 - deepseek_attention_metadata, +from afd_plugin.v1.worker.npu import ( # noqa: E402 + attention_metadata, + ubatch_runner, ) -from afd_plugin.v1.worker.npu import ubatch_runner # noqa: E402 @pytest.mark.parametrize("backend", ["mla", "sfa", "dsa"]) @@ -85,7 +85,7 @@ def build(**kwargs): monkeypatch.setattr(ubatch_runner, "build_attn_metadata", build) monkeypatch.setattr( - deepseek_attention_metadata, + attention_metadata, "get_cos_and_sin_dsa", lambda positions, use_cache: dsa_rope( positions.float().view(2, 1) * 3, diff --git a/tests/unit/connectors/npu/test_async_cam_startup_wiring.py b/tests/unit/connectors/npu/test_async_cam_startup_wiring.py index 452364201..68b2147cc 100644 --- a/tests/unit/connectors/npu/test_async_cam_startup_wiring.py +++ b/tests/unit/connectors/npu/test_async_cam_startup_wiring.py @@ -112,6 +112,11 @@ def executor(layers): ) monkeypatch.setattr(runtime, "AsyncCAMLayeredExecutor", executor) + monkeypatch.setattr( + runtime, + "expert_weight_signature", + lambda layers, capacity: repr(tuple(layer.layer_idx for layer in layers)), + ) target = ( [] if invalid == "target" @@ -127,11 +132,11 @@ def executor(layers): runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) runner.model = SimpleNamespace( get_experts_layer_indices=lambda: (2,), - get_async_cam_layered_weights=lambda: (target, "unsupported target"), + get_quantized_expert_weights=lambda: (target, "unsupported target"), ) runner.draft_model = SimpleNamespace( get_experts_layer_indices=lambda: (5,), - get_async_cam_layered_weights=lambda: (draft, "unsupported draft"), + get_quantized_expert_weights=lambda: (draft, "unsupported draft"), ) connector = SimpleNamespace( is_initialized=False, @@ -143,7 +148,7 @@ def executor(layers): [connector.graph_startup.plan] ) runner.connector = connector - runner._failure = runner._layered_executor = None + runner._failure = runner._layered_executor = runner._mega_executor = None runner.requested_mode = mode monkeypatch.setattr( runtime, "get_ascend_device_type", lambda: runtime.AscendDeviceType.A3 @@ -181,26 +186,32 @@ def test_ffn_startup_validates_target_and_draft_weights( def test_mega_executor_failure_is_published_during_startup(modules, monkeypatch, mode): _, _, runtime = modules runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) - runner.model = SimpleNamespace(get_experts_layer_indices=lambda: (2,)) + runner.model = SimpleNamespace( + get_experts_layer_indices=lambda: (2,), + get_quantized_expert_weights=lambda: ([SimpleNamespace(layer_idx=2)], ""), + ) runner.draft_model = None runner._failure = runner._layered_executor = runner._mega_executor = None runner.requested_mode = mode connector = SimpleNamespace( is_initialized=False, extra_info=SimpleNamespace(ffn_mega_moe=True), + dynamic_quant=1, + topology=None, + max_num_batched_tokens=8, + topk=2, + tp_size=1, graph_startup=AsyncCAMGraphStartup(1, 1, 2), ) connector.init_afd_connector = lambda: resolve_ffn_mode( [connector.graph_startup.plan] ) runner.connector = connector - runner._initialize_layered_executor = lambda: setattr( - runner, - "_layered_executor", - SimpleNamespace(graph_signature=lambda _: "layered"), + monkeypatch.setattr( + runtime, "get_ascend_device_type", lambda: runtime.AscendDeviceType.A3 ) - def fail_executor(*_args): + def fail_executor(*_args, **_kwargs): raise ValueError("unsupported Mega weights") monkeypatch.setattr(runtime, "AsyncCAMMegaExecutor", fail_executor) @@ -209,6 +220,98 @@ def fail_executor(*_args): assert connector.graph_startup.plan.executor_error == "unsupported Mega weights" +@pytest.fixture +def independent_mega_runner(modules, monkeypatch): + _, _, runtime = modules + peers = startup_world(attn=1, ffn=1) + model = SimpleNamespace( + get_experts_layer_indices=lambda: (2,), + get_quantized_expert_weights=lambda: ([SimpleNamespace(layer_idx=2)], ""), + ) + connector = SimpleNamespace( + is_initialized=False, + dynamic_quant=1, + extra_info=SimpleNamespace(ffn_mega_moe=True), + graph_startup=peers[1], + topology=None, + max_num_batched_tokens=8, + topk=2, + tp_size=1, + comm_args=object(), + group_name="ready", + ) + + def initialize(): + connector.graph_startup.mode = resolve_ffn_mode([connector.graph_startup.plan]) + connector.is_initialized = True + + connector.init_afd_connector = initialize + fail_binding = [False] + mega = MagicMock() + + def bind(*args): + if not connector.is_initialized: + raise RuntimeError("communication was bound before initialization") + if fail_binding[0]: + raise RuntimeError("communication binding failed") + + mega.bind_communication = bind + + def forbidden_layered(*args): + raise AssertionError("Mega must not construct Layered") + + monkeypatch.setattr(runtime, "AsyncCAMLayeredExecutor", forbidden_layered) + monkeypatch.setattr(runtime, "AsyncCAMMegaExecutor", lambda *args, **kw: mega) + monkeypatch.setattr( + runtime, "expert_weight_signature", lambda layers, capacity: "weights" + ) + monkeypatch.setattr( + runtime, "get_ascend_device_type", lambda: runtime.AscendDeviceType.A3 + ) + runner = runtime.AsyncNPUFFNModelRunner.__new__(runtime.AsyncNPUFFNModelRunner) + runner.__dict__.update( + model=model, + draft_model=None, + connector=connector, + requested_mode="eager", + _failure=None, + _layered_executor=None, + _mega_executor=None, + _first_step=False, + _graph=None, + _is_shutdown=False, + prof=None, + session=sys.modules["afd_plugin.connectors.base"].FFNExecutionSession(), + ) + return SimpleNamespace( + runner=runner, connector=connector, peer=peers[0], fail_binding=fail_binding + ) + + +def test_mega_serves_without_constructing_layered(independent_mega_runner): + h = independent_mega_runner + h.runner.initialize_afd_connector() + assert h.connector.graph_startup.plan.signature == "native_mc2:weights" + assert h.runner.execute_connector_driven_step() + h.peer.check_failure() + + +@pytest.mark.parametrize("mode", ["eager", "full"]) +def test_mega_binding_failure_reaches_peer_and_retains_resources( + independent_mega_runner, mode +): + h = independent_mega_runner + h.runner.requested_mode = mode + h.fail_binding[0] = True + with pytest.raises(RuntimeError, match="communication binding failed"): + h.runner.initialize_afd_connector() + with pytest.raises(RuntimeError, match="communication binding failed"): + h.peer.check_failure() + with pytest.raises(RuntimeError, match="startup retains resources"): + h.runner.shutdown() + assert h.connector.is_initialized + + @pytest.mark.parametrize( "field", [ @@ -409,8 +512,8 @@ def capture(): device="npu:0", model=SimpleNamespace(get_experts_layer_indices=lambda: (5, 2)), draft_model=None, - _initialize_layered_executor=lambda: None, - _layered_executor=SimpleNamespace(graph_signature=lambda capacity: "weights"), + _initialize_ffn_executor=lambda: "weights", + _layered_executor=None, _mega_executor=None, connector=SimpleNamespace( is_initialized=False, diff --git a/tests/unit/fakes/async_ffn.py b/tests/unit/fakes/async_ffn.py index b87b880e5..7f918eddf 100644 --- a/tests/unit/fakes/async_ffn.py +++ b/tests/unit/fakes/async_ffn.py @@ -42,7 +42,8 @@ def stub(name, **attributes): "afd_plugin.compat.npu.ops", "afd_plugin.compat.npu.forward_context", "afd_plugin.compat.npu.profiler", - "afd_plugin.model_executor.npu.async_cam_layered", + "afd_plugin.model_executor.ffn.compute", + "afd_plugin.model_executor.ffn.expert_weights", "afd_plugin.v1.worker.ffn_metadata", ): stub(name) diff --git a/tests/unit/fakes/connector.py b/tests/unit/fakes/connector.py index 7222e2270..8c361f01a 100644 --- a/tests/unit/fakes/connector.py +++ b/tests/unit/fakes/connector.py @@ -131,20 +131,19 @@ def recv_ffn_output(self, *, ref_tensor, ubatch_idx): def install_proxy_context(adapter, monkeypatch, events, *, stage_idx=2): connector = ProxyConnector(events) afd_metadata = SimpleNamespace(connector=connector, stage_idx=9) - monkeypatch.setattr( - adapter, - "get_afd_metadata_from_forward_context", - lambda: afd_metadata, - ) - monkeypatch.setattr( - adapter, - "get_forward_context", - lambda: SimpleNamespace(ubatch_idx=stage_idx), - ) + from afd_plugin.model_executor.ffn import remote def record_yield(hidden_states, *, role): events.append(("yield", hidden_states, role)) return hidden_states - monkeypatch.setattr(adapter, "maybe_apply_dbo_yield", record_yield) + monkeypatch.setattr( + remote, "get_afd_metadata_from_forward_context", lambda: afd_metadata + ) + monkeypatch.setattr( + remote, "get_forward_context", lambda: SimpleNamespace(ubatch_idx=stage_idx) + ) + # Model-specific forwards and shared proxies both yield to DBO. + for module in (adapter, remote): + monkeypatch.setattr(module, "maybe_apply_dbo_yield", record_yield) return afd_metadata diff --git a/tests/unit/ffn_faas/test_model.py b/tests/unit/ffn_faas/test_model.py index 3621867c0..c104ebafd 100644 --- a/tests/unit/ffn_faas/test_model.py +++ b/tests/unit/ffn_faas/test_model.py @@ -16,7 +16,7 @@ def test_resource_uses_adapter_routed_layers_and_loading_estimate(tmp_path): assert model.get_experts_layer_indices() == (2, 7) assert model.parameter_bytes == 2048 with pytest.raises(RuntimeError, match="not been loaded"): - model.get_async_cam_layered_weights() + model.get_quantized_expert_weights() with pytest.raises(RuntimeError, match="not finished loading"): _ = model.native_owner model.close() diff --git a/afd_plugin/model_executor/npu/__init__.py b/tests/unit/model_executor/ffn/__init__.py similarity index 59% rename from afd_plugin/model_executor/npu/__init__.py rename to tests/unit/model_executor/ffn/__init__.py index 32bc539c1..2c50c051d 100644 --- a/afd_plugin/model_executor/npu/__init__.py +++ b/tests/unit/model_executor/ffn/__init__.py @@ -1,3 +1,3 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""NPU-specific model execution helpers.""" +"""Shared FFN component tests, isolated from connector test module names.""" diff --git a/tests/unit/model_executor/test_async_cam_dbo.py b/tests/unit/model_executor/ffn/test_async_cam.py similarity index 56% rename from tests/unit/model_executor/test_async_cam_dbo.py rename to tests/unit/model_executor/ffn/test_async_cam.py index e19baf39d..592fbe31a 100644 --- a/tests/unit/model_executor/test_async_cam_dbo.py +++ b/tests/unit/model_executor/ffn/test_async_cam.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""CAM's physical window and owner result survive native DBO handoffs.""" +"""Async CAM model-side layout and native DBO handoff contracts.""" import threading from types import SimpleNamespace @@ -13,7 +13,54 @@ from vllm.v1.worker import ubatching # noqa: E402 from afd_plugin.connectors import AFDTransferContext, AFDTransferMetadata # noqa: E402 -from afd_plugin.model_executor.models.npu import async_cam_dbo # noqa: E402 +from afd_plugin.model_executor.ffn import async_cam # noqa: E402 + + +def test_plain_tp_cam_boundary_shards_and_restores_replicated_tokens(monkeypatch): + tp_group = SimpleNamespace(world_size=2, rank_in_group=0) + monkeypatch.setattr(async_cam, "get_tp_group", lambda: tp_group) + hidden_states = torch.arange(10, dtype=torch.float32).reshape(5, 2) + topk_weights = torch.arange(10, dtype=torch.float32).reshape(5, 2) + topk_ids = torch.arange(10, dtype=torch.int32).reshape(5, 2) + router_logits = torch.arange(20, dtype=torch.float32).reshape(5, 4) + padded_output = torch.arange(12, dtype=torch.float32).reshape(6, 2) + 100 + + monkeypatch.setattr( + async_cam, + "tensor_model_parallel_all_gather", + lambda tensor, token_dim: padded_output, + ) + + expected_hidden_rows = ( + hidden_states[:3], + torch.cat((hidden_states[3:], hidden_states.new_zeros((1, 2)))), + ) + for tp_rank in range(2): + tp_group.rank_in_group = tp_rank + payload = async_cam.prepare_cam_dispatch_payload( + hidden_states, + topk_weights, + topk_ids, + router_logits, + use_sequence_parallel=False, + ) + + assert torch.equal(payload.hidden_states, expected_hidden_rows[tp_rank]) + assert payload.hidden_states.shape[0] == 3 + assert payload.topk_weights.shape[0] == 3 + assert payload.topk_ids.shape[0] == 3 + assert payload.router_logits is not None + assert payload.router_logits.shape[0] == 3 + assert payload.layout.parent_tokens == 5 + assert payload.layout.padded_tokens == 6 + assert payload.layout.requires_tp_all_gather is True + + local_output = padded_output[tp_rank * 3 : (tp_rank + 1) * 3] + restored = async_cam.restore_cam_dispatch_output( + local_output, + payload.layout, + ) + assert torch.equal(restored, padded_output[:5]) def _context(stage): @@ -47,9 +94,9 @@ def receive(*, ref_tensor, ubatch_idx, context): connector = SimpleNamespace(send_attn_output=send, recv_ffn_output=receive) hidden = [torch.full((2, 1), stage) for stage in range(2)] - first = async_cam_dbo.send_async_cam_remote_call(connector, hidden[0], _context(0)) + first = async_cam.send_async_cam_remote_call(connector, hidden[0], _context(0)) registry[identity] = 1 - second = async_cam_dbo.send_async_cam_remote_call(connector, hidden[1], _context(1)) + second = async_cam.send_async_cam_remote_call(connector, hidden[1], _context(1)) assert events == [("send", 0, 0), ("recv", 0, 1), ("send", 1, 1)] registry[identity] = 0 @@ -64,12 +111,12 @@ def receive(*, ref_tensor, ubatch_idx, context): def test_without_dbo_receives_directly_without_inactive_hooks(monkeypatch): - monkeypatch.setattr(async_cam_dbo, "dbo_enabled", lambda: False) + monkeypatch.setattr(async_cam, "dbo_enabled", lambda: False) monkeypatch.setattr( - async_cam_dbo, "dbo_maybe_run_recv_hook", lambda: pytest.fail("inactive hook") + async_cam, "dbo_maybe_run_recv_hook", lambda: pytest.fail("inactive hook") ) monkeypatch.setattr( - async_cam_dbo, "dbo_register_recv_hook", lambda _: pytest.fail("inactive hook") + async_cam, "dbo_register_recv_hook", lambda _: pytest.fail("inactive hook") ) calls = [] value = torch.ones(2, 1) @@ -82,14 +129,14 @@ def recv(**kwargs): send_attn_output=lambda *args, **kwargs: calls.append("send"), recv_ffn_output=recv, ) - receive = async_cam_dbo.send_async_cam_remote_call(connector, value, _context(0)) + receive = async_cam.send_async_cam_remote_call(connector, value, _context(0)) assert calls == ["send"] assert receive() is value and receive() is value assert calls == ["send", "recv"] def test_failed_receive_is_never_resubmitted(monkeypatch): - monkeypatch.setattr(async_cam_dbo, "dbo_enabled", lambda: False) + monkeypatch.setattr(async_cam, "dbo_enabled", lambda: False) calls = [] def fail(**kwargs): @@ -99,7 +146,7 @@ def fail(**kwargs): connector = SimpleNamespace( send_attn_output=lambda *a, **k: None, recv_ffn_output=fail ) - receive = async_cam_dbo.send_async_cam_remote_call( + receive = async_cam.send_async_cam_remote_call( connector, torch.ones(2, 1), _context(0) ) with pytest.raises(ValueError, match="wire failure"): diff --git a/tests/unit/model_executor/test_async_cam_layered.py b/tests/unit/model_executor/ffn/test_compute.py similarity index 93% rename from tests/unit/model_executor/test_async_cam_layered.py rename to tests/unit/model_executor/ffn/test_compute.py index 6b6e90b1a..d68ffdd5d 100644 --- a/tests/unit/model_executor/test_async_cam_layered.py +++ b/tests/unit/model_executor/ffn/test_compute.py @@ -15,19 +15,20 @@ torch = pytest.importorskip("torch") -from afd_plugin.model_executor.npu.async_cam_layered import ( # noqa: E402 +from afd_plugin.model_executor.ffn.compute import ( # noqa: E402 AsyncCAMLayeredExecutor, - LayeredFFNWeights, -) -from afd_plugin.model_executor.npu.async_cam_mega import ( # noqa: E402 AsyncCAMMegaExecutor, ) +from afd_plugin.model_executor.ffn.expert_weights import ( # noqa: E402 + QuantizedExpertWeights, + expert_weight_signature, +) def weights(layer_idx=0, per_channel=True, quantization="W4A8"): experts, hidden, intermediate = 2, 32, 16 if quantization == "W8A8": - return LayeredFFNWeights( + return QuantizedExpertWeights( layer_idx=layer_idx, w13=torch.empty(experts, hidden, 2 * intermediate, dtype=torch.int8), w2=torch.empty(experts, intermediate, hidden, dtype=torch.int8), @@ -40,7 +41,7 @@ def weights(layer_idx=0, per_channel=True, quantization="W4A8"): routed_scaling_factor=2.0, quantization="W8A8", ) - return LayeredFFNWeights( + return QuantizedExpertWeights( layer_idx=layer_idx, w13=torch.empty(experts, hidden, intermediate // 4, dtype=torch.int32), w2=torch.empty(experts, intermediate, hidden // 8, dtype=torch.int32), @@ -62,9 +63,15 @@ def weights(layer_idx=0, per_channel=True, quantization="W4A8"): def test_native_rejects_grouped_scale_before_receiving(operators): - executor = AsyncCAMLayeredExecutor([weights(per_channel=False)]) with pytest.raises(ValueError, match="per-channel"): - AsyncCAMMegaExecutor(executor, None) + AsyncCAMMegaExecutor( + [weights(per_channel=False)], + topology=None, + max_num_batched_tokens=8, + topk=2, + tp_size=1, + capacity=16, + ) @pytest.mark.parametrize("quantization", ["W4A8", "W8A8"]) @@ -86,21 +93,14 @@ def execute(*args): torch.ops.afd_ascend, "afd_mega_moe_pull_native", execute, raising=False ) executor = AsyncCAMMegaExecutor( - AsyncCAMLayeredExecutor( - [replace(weights(quantization=quantization), swiglu_limit=10.0)] - ), - SimpleNamespace( - max_num_batched_tokens=8, - topk=2, - ffn_size=2, - attn_size=1, - world_rank=1, - topology=SimpleNamespace(world_size=3), - tp_size=1, - comm_args=torch.empty(512, dtype=torch.float16), - group_name="test-pull", - ), + [replace(weights(quantization=quantization), swiglu_limit=10.0)], + topology=SimpleNamespace(ffn_size=2, attn_size=1, world_rank=1, world_size=3), + max_num_batched_tokens=8, + topk=2, + tp_size=1, + capacity=16, ) + executor.bind_communication(torch.empty(512, dtype=torch.float16), "test-pull") executor() assert calls[0][-1] == (10.0 if quantization == "W4A8" else 0.0) assert calls[0][10][11] == (1 if quantization == "W4A8" else 2) @@ -236,10 +236,10 @@ def test_w8a8_rejects_incompatible_layout_before_receive(operators, change, mess assert operators == [] -def test_graph_signature_distinguishes_quantization(operators): - w4 = AsyncCAMLayeredExecutor([weights()]) - w8 = AsyncCAMLayeredExecutor([weights(quantization="W8A8")]) - assert w4.graph_signature(32) != w8.graph_signature(32) +def test_graph_signature_distinguishes_quantization(): + assert expert_weight_signature([weights()], 32) != expert_weight_signature( + [weights(quantization="W8A8")], 32 + ) with pytest.raises(ValueError, match="homogeneous"): AsyncCAMLayeredExecutor([weights(0), weights(1, quantization="W8A8")]) @@ -340,7 +340,7 @@ def model_weights_method(monkeypatch): node for node in ast.parse(source).body if isinstance(node, ast.FunctionDef) - and node.name == "_extract_async_cam_layered_weights" + and node.name == "_extract_quantized_expert_weights" ) monkeypatch.setitem( sys.modules, @@ -365,7 +365,7 @@ def model_weights_method(monkeypatch): exec( compile(ast.fix_missing_locations(module), "", "exec"), namespace ) - return namespace["_extract_async_cam_layered_weights"] + return namespace["_extract_quantized_expert_weights"] def model_layer(idx, quantization="W4A8"): diff --git a/tests/unit/model_executor/models/npu/test_deepseek_v2_attention_gate.py b/tests/unit/model_executor/models/npu/test_deepseek_v2_attention_gate.py index 561f40bca..52a968d75 100644 --- a/tests/unit/model_executor/models/npu/test_deepseek_v2_attention_gate.py +++ b/tests/unit/model_executor/models/npu/test_deepseek_v2_attention_gate.py @@ -14,9 +14,9 @@ @pytest.fixture def preparation(monkeypatch): - models = ModuleType("afd_plugin.model_executor.models") - models.__dict__["get_afd_metadata_from_forward_context"] = lambda: None - monkeypatch.setitem(sys.modules, models.__name__, models) + context = ModuleType("afd_plugin.forward_context") + context.__dict__["get_afd_metadata_from_forward_context"] = lambda: None + monkeypatch.setitem(sys.modules, context.__name__, context) quantization = ModuleType("vllm_ascend.quantization.quant_type") quant_type = SimpleNamespace(W8A8="w8a8", W4A8="w4a8", NONE="none") quantization.__dict__["QuantType"] = quant_type diff --git a/tests/unit/model_executor/models/npu/test_deepseek_v4.py b/tests/unit/model_executor/models/npu/test_deepseek_v4.py index c41d9444f..e47f65130 100644 --- a/tests/unit/model_executor/models/npu/test_deepseek_v4.py +++ b/tests/unit/model_executor/models/npu/test_deepseek_v4.py @@ -249,7 +249,7 @@ def checkpoint_weight_roles() -> Callable[..., frozenset[str]]: for node in ast.parse(path.read_text()).body if ( isinstance(node, ast.FunctionDef) - and node.name in ("_weight_layer_path", "_checkpoint_weight_roles") + and node.name == "_checkpoint_weight_roles" ) or ( isinstance(node, ast.Assign) @@ -257,7 +257,9 @@ def checkpoint_weight_roles() -> Callable[..., frozenset[str]]: and node.targets[0].id in ("_ATTENTION_ROLE", "_FFN_ROLE", "_BOTH_ROLES") ) ] - namespace: dict[str, object] = {} + from afd_plugin.model_executor.models._weight_utils import layer_weight_path + + namespace: dict[str, object] = {"layer_weight_path": layer_weight_path} module = ast.Module(body=helpers, type_ignores=[]) exec(compile(module, str(path), "exec"), namespace) return cast(Callable[..., frozenset[str]], namespace["_checkpoint_weight_roles"]) diff --git a/tests/unit/model_executor/models/test_forward_context.py b/tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py similarity index 79% rename from tests/unit/model_executor/models/test_forward_context.py rename to tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py index 1a6894638..1f5892ec6 100644 --- a/tests/unit/model_executor/models/test_forward_context.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py @@ -10,28 +10,6 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") -from afd_plugin.model_executor.models import ( # noqa: E402 - get_afd_metadata_from_forward_context, -) - - -def test_get_afd_metadata_from_additional_kwargs(): - forward_context = SimpleNamespace( - additional_kwargs={"afd_metadata": {"stage": 0}}, - afd_metadata={"stage": 1}, - ) - - assert get_afd_metadata_from_forward_context(forward_context) == {"stage": 0} - - -def test_get_afd_metadata_ignores_forward_context_attribute(): - forward_context = SimpleNamespace( - additional_kwargs={}, - afd_metadata={"stage": 0}, - ) - - assert get_afd_metadata_from_forward_context(forward_context) is None - def test_async_cam_profile_forward_runs_matched_connector_io( monkeypatch, @@ -53,10 +31,10 @@ def test_async_cam_profile_forward_runs_matched_connector_io( ) connector_calls: list[str] = [] - from afd_plugin.model_executor.models.npu import async_cam_layout + from afd_plugin.model_executor.ffn import async_cam monkeypatch.setattr( - async_cam_layout, + async_cam, "get_tp_group", lambda: SimpleNamespace(world_size=1, rank_in_group=0), ) diff --git a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py index 2bac7facd..5e3567da2 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -12,6 +12,7 @@ from torch import nn # noqa: E402 from afd_plugin.config import AFD_ASYNC_CONNECTOR, AFDConfig # noqa: E402 +from afd_plugin.model_executor.ffn import remote # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 from tests.unit.fakes.connector import ProxyConnector as _FakeConnector # noqa: E402 from tests.unit.fakes.connector import install_proxy_context # noqa: E402 @@ -270,7 +271,7 @@ def test_m2n_moe_sends_topk_and_adds_shared_output_of_dispatched_input( def test_remote_proxy_requires_forward_metadata(monkeypatch): monkeypatch.setattr( - adapter, + remote, "get_afd_metadata_from_forward_context", lambda: None, ) @@ -283,7 +284,7 @@ def test_remote_proxy_exchanges_cam_during_profile(monkeypatch): events: list[tuple] = [] install_proxy_context(adapter, monkeypatch, events) monkeypatch.setattr( - adapter, + remote, "get_forward_context", lambda: SimpleNamespace(in_profile_run=True, ubatch_idx=0), ) @@ -450,7 +451,7 @@ def test_routing_proxy_routes_on_attention_and_sends_topk(monkeypatch): def fake_fused_topk(**kwargs): return topk_weights, topk_ids, torch.zeros_like(topk_ids) - monkeypatch.setattr(adapter, "fused_topk", fake_fused_topk) + monkeypatch.setattr(remote, "fused_topk", fake_fused_topk) proxy = _routing_proxy( GateLinear(ROUTING_HIDDEN_SIZE, ROUTING_NUM_EXPERTS, out_dtype=torch.float32), renormalize=False, @@ -521,7 +522,7 @@ def _tokens(num_tokens: int): def _tp_routing_proxy(monkeypatch, tp_size: int): - monkeypatch.setattr(adapter, "fused_topk", _torch_topk) + monkeypatch.setattr(remote, "fused_topk", _torch_topk) proxy = _routing_proxy(_TokenGate()) proxy.tp_size = tp_size return proxy @@ -538,7 +539,7 @@ def no_split(x): events: list[tuple] = [] install_proxy_context(adapter, monkeypatch, events) - monkeypatch.setattr(adapter.native, "sequence_parallel_chunk", no_split) + monkeypatch.setattr(remote, "sequence_parallel_chunk", no_split) hidden_states = _tokens(3) output = _tp_routing_proxy(monkeypatch, tp_size=1)(hidden_states) @@ -567,18 +568,16 @@ def test_routing_proxy_tp2_rank_dispatches_its_slice_and_gathers( events: list[tuple] = [] install_proxy_context(adapter, monkeypatch, events) monkeypatch.setattr( - adapter.native, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) - ) - monkeypatch.setattr( - adapter.native, "get_tensor_model_parallel_rank", lambda: tp_rank + remote, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) ) + monkeypatch.setattr(remote, "get_tensor_model_parallel_rank", lambda: tp_rank) gathered_inputs = [] def all_gather(local_output, dim): gathered_inputs.append(local_output.clone()) return torch.cat([local_output, local_output + PEER_OFFSET], dim=dim) - monkeypatch.setattr(adapter.native, "tensor_model_parallel_all_gather", all_gather) + monkeypatch.setattr(remote, "tensor_model_parallel_all_gather", all_gather) output = _tp_routing_proxy(monkeypatch, tp_size)(_tokens(3)) @@ -586,9 +585,9 @@ def all_gather(local_output, dim): sent_states, sent_kwargs = send[1], send[3] assert sent_states[:, 0].tolist() == sent_tokens topk_ids = sent_kwargs["topk_ids"] - unrouted = (topk_ids == adapter.UNROUTED_EXPERT_ID).all(dim=1) + unrouted = (topk_ids == remote.UNROUTED_EXPERT_ID).all(dim=1) assert unrouted.nonzero().flatten().tolist() == padding_rows - assert not (topk_ids[~unrouted] == adapter.UNROUTED_EXPERT_ID).any() + assert not (topk_ids[~unrouted] == remote.UNROUTED_EXPERT_ID).any() (local_output,) = gathered_inputs assert torch.equal(local_output, sent_states * COMBINE_SCALE) expected = torch.cat([local_output, local_output + PEER_OFFSET])[:3] @@ -606,16 +605,16 @@ def test_routing_proxy_tp4_with_one_token_routes_only_rank0(monkeypatch): events: list[tuple] = [] install_proxy_context(adapter, monkeypatch, events) monkeypatch.setattr( - adapter.native, + remote, "tensor_model_parallel_all_gather", lambda local, dim: torch.cat([local] * tp_size, dim=dim), ) for tp_rank in range(tp_size): monkeypatch.setattr( - adapter.native, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) + remote, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) ) monkeypatch.setattr( - adapter.native, + remote, "get_tensor_model_parallel_rank", lambda rank=tp_rank: rank, ) @@ -624,7 +623,7 @@ def test_routing_proxy_tp4_with_one_token_routes_only_rank0(monkeypatch): assert output.shape == (1, ROUTING_HIDDEN_SIZE) unrouted = [ - bool((event[3]["topk_ids"] == adapter.UNROUTED_EXPERT_ID).all()) + bool((event[3]["topk_ids"] == remote.UNROUTED_EXPERT_ID).all()) for event in events if event[0] == "send" ] @@ -653,9 +652,7 @@ def test_m2n_moe_tp2_rank_dispatches_its_slice_and_gathers_shared_sum( monkeypatch.setattr( adapter.native, "sequence_parallel_chunk", _chunk_for_rank(tp_rank, tp_size) ) - monkeypatch.setattr( - adapter.native, "get_tensor_model_parallel_rank", lambda: tp_rank - ) + monkeypatch.setattr(remote, "get_tensor_model_parallel_rank", lambda: tp_rank) gathered_inputs = [] def all_gather(local_output, dim): @@ -674,7 +671,7 @@ def all_gather(local_output, dim): assert [event[0] for event in events] == ["send", "yield", "recv"] sent_states, sent_kwargs = events[0][1], events[0][3] assert sent_states[:, 0].tolist() == sent_tokens - unrouted = (sent_kwargs["topk_ids"] == adapter.UNROUTED_EXPERT_ID).all(dim=1) + unrouted = (sent_kwargs["topk_ids"] == remote.UNROUTED_EXPERT_ID).all(dim=1) assert unrouted.nonzero().flatten().tolist() == padding_rows (local_output,) = gathered_inputs expected_local = COMBINED_VALUE + shared_factor * torch.tensor( diff --git a/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py b/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py index 4299eb12e..59c50e8ca 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py @@ -13,10 +13,8 @@ pytest.importorskip("vllm") +from afd_plugin.model_executor.ffn import async_cam # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 -from afd_plugin.model_executor.models.npu import ( # noqa: E402 - async_cam_layout, -) from afd_plugin.model_executor.models.npu import ( # noqa: E402 deepseek_v2_async_cam_forward as schedule, ) @@ -140,7 +138,7 @@ def gather(states, dim): return torch.cat((states, torch.zeros_like(states))) monkeypatch.setattr( - async_cam_layout, + async_cam, "get_tp_group", lambda: SimpleNamespace(world_size=2, rank_in_group=0), ) diff --git a/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py b/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py index 31dc1f740..ad6402d9b 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py @@ -11,11 +11,11 @@ pytest.importorskip("vllm") pytest.importorskip("vllm_ascend.models.deepseek_v4.model") +from afd_plugin.model_executor.ffn import async_cam # noqa: E402 +from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 from afd_plugin.model_executor.models.npu import ( # noqa: E402 - async_cam_layout, # noqa: E402 deepseek_v4_attention_gate, ) -from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 def _make_proxy(monkeypatch, *, use_sequence_parallel: bool, in_profile_run: bool): @@ -55,11 +55,11 @@ def test_dsv4_profile_gate_shards_plain_tp8_and_restores(monkeypatch): in_profile_run=True, ) tp_group = SimpleNamespace(world_size=8, rank_in_group=3) - monkeypatch.setattr(async_cam_layout, "get_tp_group", lambda: tp_group) + monkeypatch.setattr(async_cam, "get_tp_group", lambda: tp_group) hidden_states = torch.arange(64, dtype=torch.float32).reshape(16, 4) gathered_output = hidden_states + 100 monkeypatch.setattr( - async_cam_layout, + async_cam, "tensor_model_parallel_all_gather", lambda _tensor, token_dim: gathered_output, ) @@ -89,7 +89,7 @@ def test_dsv4_gate_keeps_sequence_parallel_tokens_rank_local(monkeypatch): in_profile_run=False, ) monkeypatch.setattr( - async_cam_layout, + async_cam, "get_tp_group", lambda: SimpleNamespace(world_size=8, rank_in_group=3), ) diff --git a/tests/unit/model_executor/models/test_weight_utils.py b/tests/unit/model_executor/models/test_weight_utils.py new file mode 100644 index 000000000..7947835ea --- /dev/null +++ b/tests/unit/model_executor/models/test_weight_utils.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Shared checkpoint mechanics remain independent of torch and model classes.""" + +import pytest + +from afd_plugin.model_executor.models._weight_utils import ( + filter_role_weights, + layer_weight_path, +) + + +@pytest.mark.parametrize( + "name,expected", + [ + ("model.layers.3.mlp.experts.weight", (3, "mlp", ("experts", "weight"))), + ( + "model.language_model.layers.5.self_attn.weight", + (5, "self_attn", ("weight",)), + ), + ("layers.invalid.layers.2.mlp.weight", (2, "mlp", ("weight",))), + ("model.layers.2", None), + ("model.embed_tokens.weight", None), + ], +) +def test_checkpoint_path_keeps_layer_stage_and_remainder(name, expected): + assert layer_weight_path(name) == expected + + +@pytest.mark.parametrize("role", [None, "attention", "ffn"]) +def test_filter_streams_once_and_preserves_weights(role): + first, second = object(), object() + source = iter([("attention", first), ("ffn", second)]) + + def classify(name): + assert role is not None + return frozenset((name,)) + + actual = list(filter_role_weights(source, role=role, classify=classify)) + expected = [("attention", first), ("ffn", second)] + if role is not None: + expected = [item for item in expected if item[0] == role] + assert actual == expected + assert list(source) == [] diff --git a/tests/unit/model_executor/test_async_cam_layout.py b/tests/unit/model_executor/test_async_cam_layout.py deleted file mode 100644 index 60022ac8b..000000000 --- a/tests/unit/model_executor/test_async_cam_layout.py +++ /dev/null @@ -1,59 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project - -from types import SimpleNamespace - -import pytest - -torch = pytest.importorskip("torch") - -pytest.importorskip("vllm") - -from afd_plugin.model_executor.models.npu import async_cam_layout # noqa: E402 - - -def test_plain_tp_cam_boundary_shards_and_restores_replicated_tokens(monkeypatch): - tp_group = SimpleNamespace(world_size=2, rank_in_group=0) - monkeypatch.setattr(async_cam_layout, "get_tp_group", lambda: tp_group) - hidden_states = torch.arange(10, dtype=torch.float32).reshape(5, 2) - topk_weights = torch.arange(10, dtype=torch.float32).reshape(5, 2) - topk_ids = torch.arange(10, dtype=torch.int32).reshape(5, 2) - router_logits = torch.arange(20, dtype=torch.float32).reshape(5, 4) - padded_output = torch.arange(12, dtype=torch.float32).reshape(6, 2) + 100 - - monkeypatch.setattr( - async_cam_layout, - "tensor_model_parallel_all_gather", - lambda tensor, token_dim: padded_output, - ) - - expected_hidden_rows = ( - hidden_states[:3], - torch.cat((hidden_states[3:], hidden_states.new_zeros((1, 2)))), - ) - for tp_rank in range(2): - tp_group.rank_in_group = tp_rank - payload = async_cam_layout.prepare_cam_dispatch_payload( - hidden_states, - topk_weights, - topk_ids, - router_logits, - use_sequence_parallel=False, - ) - - assert torch.equal(payload.hidden_states, expected_hidden_rows[tp_rank]) - assert payload.hidden_states.shape[0] == 3 - assert payload.topk_weights.shape[0] == 3 - assert payload.topk_ids.shape[0] == 3 - assert payload.router_logits is not None - assert payload.router_logits.shape[0] == 3 - assert payload.layout.parent_tokens == 5 - assert payload.layout.padded_tokens == 6 - assert payload.layout.requires_tp_all_gather is True - - local_output = padded_output[tp_rank * 3 : (tp_rank + 1) * 3] - restored = async_cam_layout.restore_cam_dispatch_output( - local_output, - payload.layout, - ) - assert torch.equal(restored, padded_output[:5]) diff --git a/tests/unit/test_forward_context.py b/tests/unit/test_forward_context.py new file mode 100644 index 000000000..4669f4003 --- /dev/null +++ b/tests/unit/test_forward_context.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Package-level AFD forward metadata reads.""" + +from types import SimpleNamespace + +import pytest + +pytest.importorskip("torch") +pytest.importorskip("vllm") + +from afd_plugin.forward_context import ( + get_afd_metadata_from_forward_context, # noqa: E402 +) + + +def test_get_afd_metadata_from_additional_kwargs(): + forward_context = SimpleNamespace( + additional_kwargs={"afd_metadata": {"stage": 0}}, + afd_metadata={"stage": 1}, + ) + + assert get_afd_metadata_from_forward_context(forward_context) == {"stage": 0} + + +def test_get_afd_metadata_ignores_forward_context_attribute(): + forward_context = SimpleNamespace( + additional_kwargs={}, + afd_metadata={"stage": 0}, + ) + + assert get_afd_metadata_from_forward_context(forward_context) is None diff --git a/tests/unit/model_executor/test_deepseek_attention_metadata.py b/tests/unit/v1/worker/npu/test_attention_metadata.py similarity index 93% rename from tests/unit/model_executor/test_deepseek_attention_metadata.py rename to tests/unit/v1/worker/npu/test_attention_metadata.py index f480f6fe6..375c851ba 100644 --- a/tests/unit/model_executor/test_deepseek_attention_metadata.py +++ b/tests/unit/v1/worker/npu/test_attention_metadata.py @@ -23,9 +23,7 @@ from vllm_ascend.ops import rope_dsv4 # noqa: E402 from vllm_ascend.ops.rope_dsv4 import RopeDataProxy # noqa: E402 -from afd_plugin.model_executor.models.npu import ( # noqa: E402 - deepseek_attention_metadata, -) +from afd_plugin.v1.worker.npu import attention_metadata # noqa: E402 def test_materializes_sfa_mla_and_dsa_rope_metadata( @@ -36,7 +34,7 @@ def test_materializes_sfa_mla_and_dsa_rope_metadata( sfa_metadata = object.__new__(AscendSFAMetadata) sfa_metadata.cos = sfa_cos_source[:4] sfa_metadata.sin = sfa_sin_source[:4] - deepseek_attention_metadata.materialize_deepseek_attention_metadata( + attention_metadata.materialize_attention_metadata( sfa_metadata, torch.arange(4), 4, @@ -53,7 +51,7 @@ def test_materializes_sfa_mla_and_dsa_rope_metadata( cos=mla_cos_source[5:], sin=mla_sin_source[5:], ) - deepseek_attention_metadata.materialize_deepseek_attention_metadata( + attention_metadata.materialize_attention_metadata( mla_metadata, torch.arange(8), 8, @@ -82,13 +80,13 @@ def get_dsa_rope(positions, use_cache=False): state.layer_info["layer"] = ("test", ["default"]) monkeypatch.setattr(rope_dsv4, "_ROPE_STATE", state) monkeypatch.setattr( - deepseek_attention_metadata, + attention_metadata, "get_cos_and_sin_dsa", get_dsa_rope, ) dsa_metadata = object.__new__(AscendDSAMetadata) dsa_metadata.req_metadata = SimpleNamespace(cos=object(), sin=object()) - deepseek_attention_metadata.materialize_deepseek_attention_metadata( + attention_metadata.materialize_attention_metadata( dsa_metadata, torch.arange(6), 4, @@ -110,7 +108,7 @@ def test_isolates_mutable_sfa_and_dsa_builder_inputs(): group_key_idx=group_key_idx, group_key_cache_idx=group_key_cache_idx, ) - deepseek_attention_metadata.isolate_deepseek_attention_builder_inputs( + attention_metadata.isolate_attention_builder_inputs( object.__new__(AscendSFAMetadataBuilder), sfa_common_metadata, ) @@ -123,7 +121,7 @@ def test_isolates_mutable_sfa_and_dsa_builder_inputs(): block_table = torch.arange(6, dtype=torch.int32).reshape(2, 3) dsa_common_metadata = SimpleNamespace(block_table_tensor=block_table) - deepseek_attention_metadata.isolate_deepseek_attention_builder_inputs( + attention_metadata.isolate_attention_builder_inputs( object.__new__(AscendDSAMetadataBuilder), dsa_common_metadata, ) @@ -157,13 +155,13 @@ def get_rope(positions, use_cache=False): returned.append(result) return result - monkeypatch.setattr(deepseek_attention_metadata, "get_cos_and_sin_dsa", get_rope) + monkeypatch.setattr(attention_metadata, "get_cos_and_sin_dsa", get_rope) def materialize(positions, num_tokens, buffers=None): metadata = object.__new__(AscendDSAMetadata) metadata.req_metadata = SimpleNamespace(cos=None, sin=None) # Both layers share one request metadata object: one getter call is enough. - deepseek_attention_metadata.materialize_deepseek_attention_metadata_by_layer( + attention_metadata.materialize_attention_metadata_by_layer( {"layer-a": metadata, "layer-b": metadata}, positions, num_tokens, diff --git a/tests/unit/v1/worker/npu/test_ubatch_runner.py b/tests/unit/v1/worker/npu/test_ubatch_runner.py index fee4b746e..4585be59d 100644 --- a/tests/unit/v1/worker/npu/test_ubatch_runner.py +++ b/tests/unit/v1/worker/npu/test_ubatch_runner.py @@ -198,9 +198,7 @@ def stub(name, **attributes): stub( "vllm_ascend.ops.rope_dsv4", RopeDataProxy=_RopeProxy, get_cos_and_sin_dsa=_rope ) - load_host_module( - monkeypatch, "afd_plugin.model_executor.models.npu.deepseek_attention_metadata" - ) + load_host_module(monkeypatch, "afd_plugin.v1.worker.npu.attention_metadata") module = load_host_module(monkeypatch, "afd_plugin.v1.worker.npu.ubatch_runner") return module diff --git a/tests/unit/v1/worker/test_attention_model_runner.py b/tests/unit/v1/worker/test_attention_model_runner.py index 80583945c..ac0812a66 100644 --- a/tests/unit/v1/worker/test_attention_model_runner.py +++ b/tests/unit/v1/worker/test_attention_model_runner.py @@ -13,10 +13,10 @@ from vllm.config import CUDAGraphMode from vllm.forward_context import BatchDescriptor, ForwardContext -import afd_plugin.model_executor.models.forward_context as afd_forward_context +import afd_plugin.forward_context as afd_forward_context from afd_plugin.config import AFDConfig from afd_plugin.connectors import AFDControlPayload -from afd_plugin.model_executor.models.forward_context import ( +from afd_plugin.forward_context import ( get_afd_metadata_from_forward_context, use_afd_metadata_provider, ) diff --git a/tests/unit/v1/worker/test_npu_model_runner_v2.py b/tests/unit/v1/worker/test_npu_model_runner_v2.py index 7cb0283c8..b792eeaea 100644 --- a/tests/unit/v1/worker/test_npu_model_runner_v2.py +++ b/tests/unit/v1/worker/test_npu_model_runner_v2.py @@ -586,7 +586,7 @@ def unexpected_state(*args): monkeypatch.setattr(npu_ubatch, "build_attn_state", unexpected_state) monkeypatch.setattr( npu_ubatch, - "materialize_deepseek_attention_metadata_by_layer", + "materialize_attention_metadata_by_layer", lambda *a, **kw: None, ) with patches.use_npu_input_batch_slicing(): @@ -648,7 +648,7 @@ def materialize(metadata, positions, count): monkeypatch.setattr(npu_ubatch, "build_attn_metadata", build) monkeypatch.setattr(npu_ubatch, "build_attn_state", lambda *args: "child-state") monkeypatch.setattr( - npu_ubatch, "materialize_deepseek_attention_metadata_by_layer", materialize + npu_ubatch, "materialize_attention_metadata_by_layer", materialize ) table = torch.ones((1, 2)) metadata = [ From aaf876b36d8b30a0127a3d6e69c1270df67437c3 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Sun, 8 Nov 2026 22:06:40 -0800 Subject: [PATCH 2/4] refactor: simplify V4 proxy and CAM environment handling Signed-off-by: lirx-pd <616517220@qq.com> --- afd_plugin/envs.py | 7 +++ afd_plugin/model_executor/ffn/async_cam.py | 9 +-- .../model_executor/models/deepseek_v4.py | 60 ++++--------------- docs/design/module/model_integration.md | 8 ++- .../models/test_deepseek_v4_construction.py | 43 +++++++++++++ .../models/test_deepseek_v4_proxy.py | 37 ++++++++---- tests/unit/test_envs.py | 22 +++++-- 7 files changed, 115 insertions(+), 71 deletions(-) diff --git a/afd_plugin/envs.py b/afd_plugin/envs.py index 1f7e10c8e..666d760fb 100644 --- a/afd_plugin/envs.py +++ b/afd_plugin/envs.py @@ -7,6 +7,7 @@ import os AFD_FORCE_BALANCED_TOPK_IDS = "AFD_FORCE_BALANCED_TOPK_IDS" +AFD_ASYNC_MOE_LAYOUT_LOG = "AFD_ASYNC_MOE_LAYOUT_LOG" ENV_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) @@ -14,7 +15,13 @@ def force_balanced_topk_ids_enabled() -> bool: return os.environ.get(AFD_FORCE_BALANCED_TOPK_IDS, "").lower() in ENV_TRUE_VALUES +def async_moe_layout_log_enabled() -> bool: + return os.environ.get(AFD_ASYNC_MOE_LAYOUT_LOG, "").lower() in ENV_TRUE_VALUES + + __all__ = [ + "AFD_ASYNC_MOE_LAYOUT_LOG", "AFD_FORCE_BALANCED_TOPK_IDS", + "async_moe_layout_log_enabled", "force_balanced_topk_ids_enabled", ] diff --git a/afd_plugin/model_executor/ffn/async_cam.py b/afd_plugin/model_executor/ffn/async_cam.py index a9d4a4191..3ebf1ab45 100644 --- a/afd_plugin/model_executor/ffn/async_cam.py +++ b/afd_plugin/model_executor/ffn/async_cam.py @@ -4,10 +4,8 @@ from __future__ import annotations -import os from collections.abc import Callable from dataclasses import dataclass -from typing import Final import torch from vllm.distributed import tensor_model_parallel_all_gather @@ -20,9 +18,8 @@ ) from afd_plugin.connectors import AFDConnectorBase, AFDTransferContext - -ASYNC_MOE_LAYOUT_LOG_ENV: Final[str] = "AFD_ASYNC_MOE_LAYOUT_LOG" -_TRUE_ENV_VALUES: Final[frozenset[str]] = frozenset({"1", "true", "yes", "on"}) +from afd_plugin.envs import AFD_ASYNC_MOE_LAYOUT_LOG as ASYNC_MOE_LAYOUT_LOG_ENV +from afd_plugin.envs import async_moe_layout_log_enabled logger = init_logger(__name__) @@ -192,7 +189,7 @@ def _log_cam_layout( phase: str, layout: CAMDispatchLayout, ) -> None: - if os.environ.get(ASYNC_MOE_LAYOUT_LOG_ENV, "").lower() not in _TRUE_ENV_VALUES: + if not async_moe_layout_log_enabled(): return logger.warning( "AFD Async CAM layout; phase=%s sequence_parallel=%s " diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index 291d371fb..36d6b817e 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -14,18 +14,17 @@ import torch import torch.nn as nn from vllm.config import VllmConfig -from vllm.forward_context import get_forward_context from vllm.models.deepseek_v4.nvidia import model as native from afd_plugin.compat.patches.gpu_deepseek_v4_dbo import forward_v4_attention from afd_plugin.config import parse_afd_config -from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata +from afd_plugin.connectors.metadata import AFDForwardContextMetadata from afd_plugin.forward_context import get_afd_metadata_from_forward_context +from afd_plugin.model_executor.ffn.remote import RemoteFFNProxy from afd_plugin.model_executor.models._weight_utils import ( filter_role_weights, layer_weight_path, ) -from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield _ATTENTION_ROLE = frozenset(("attention",)) _FFN_ROLE = frozenset(("ffn",)) @@ -52,14 +51,10 @@ def _checkpoint_weight_roles(name: str) -> frozenset[str]: return _ATTENTION_ROLE -class RemoteDeepseekV4FFN(nn.Module): +class RemoteDeepseekV4FFN(RemoteFFNProxy): """Parameter-free FFN proxy carrying V4 hash-router token identifiers.""" - def __init__(self, *, layer_idx: int) -> None: - super().__init__() - self.layer_idx = layer_idx - - def forward( + def forward( # type: ignore[override] self, hidden_states: torch.Tensor, input_ids: torch.Tensor | None, @@ -71,33 +66,15 @@ def forward( "DeepSeek-V4 input_ids must be one-dimensional and token-aligned", ) + return self._send_and_receive(hidden_states, input_ids=input_ids) + + @staticmethod + def _afd_metadata() -> AFDForwardContextMetadata: + # Keep V4's missing-context error while sharing the transport. afd_metadata = get_afd_metadata_from_forward_context() if afd_metadata is None: raise RuntimeError("RemoteDeepseekV4FFN requires AFD forward metadata") - forward_context = get_forward_context() - stage_idx = int( - getattr(forward_context, "ubatch_idx", afd_metadata.stage_idx), - ) - afd_metadata.stage_idx = stage_idx - metadata = AFDTransferMetadata.create_attention_metadata( - layer_idx=self.layer_idx, - stage_idx=stage_idx, - seq_len=int(hidden_states.shape[0]), - ) - context = AFDTransferContext(metadata=metadata) - afd_metadata.connector.send_attn_output( - hidden_states, - context, - input_ids=input_ids, - ) - hidden_states = maybe_apply_dbo_yield( - hidden_states, - role="attention", - ) - return afd_metadata.connector.recv_ffn_output( - ref_tensor=hidden_states, - ubatch_idx=stage_idx, - ) + return afd_metadata class AFDDeepseekV4DecoderLayer(native.DeepseekV4DecoderLayer): @@ -546,22 +523,7 @@ def finalize_mhc_broadcast_weights(self) -> None: """Finalize only the Attention-owned first-layer broadcast matrix.""" if self.afd_config.role == "ffn": return - if ( - not native.get_pp_group().is_first_rank - or self.start_layer >= self.end_layer - ): - return - layer = self.layers[self.start_layer] - if isinstance(layer, AFDDeepseekV4DecoderLayer): - broadcast = ( - layer.hc_attn_fn.detach() - .view(-1, layer.hc_mult, layer.hidden_size) - .sum(dim=1) - ) - if layer.hc_attn_fn_broadcast is None: - layer.hc_attn_fn_broadcast = broadcast - else: - layer.hc_attn_fn_broadcast.copy_(broadcast) + super().finalize_mhc_broadcast_weights() class AFDDeepseekV4ForCausalLM(native.DeepseekV4ForCausalLM): diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 1bc31dd25..60e5a54f0 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -83,7 +83,7 @@ flowchart TD Compute --> Native[Native NPU operators] ``` -DeepSeek V2, Qwen3, Qwen3.5 and Ascend V4 use the shared remote components. +DeepSeek V2, Qwen3, Qwen3.5 and CUDA/Ascend V4 use the shared remote components. The CAM helpers stay in `ffn/async_cam.py` because they adapt model tensors to remote FFN calls across model families. Connector/Session owns transport windows, pending work and Push/Pull protocol state. The generic proxies do @@ -196,6 +196,12 @@ two-dimensional FFN activation plus token-aligned `input_ids` through a parameter-free `RemoteDeepseekV4FFN`. FFN owns the complete native `DeepseekV4MoE`, including its hash router, and returns the FFN activation. +`RemoteDeepseekV4FFN` inherits the shared `RemoteFFNProxy` transport and keeps +V4's token-ID validation and missing-context error. The common implementation +owns metadata creation, stage selection and send/yield/receive ordering. mHC +broadcast-weight finalization delegates to the native model on Attention; +the FFN role skips it because it owns no mHC parameters. + Role-aware weight filtering keeps layer `.ffn` parameters on FFN and all other layer-local and mHC head parameters on Attention; common non-layer paths remain available to both roles for the native loader lifecycle. The adapter is diff --git a/tests/unit/model_executor/models/test_deepseek_v4_construction.py b/tests/unit/model_executor/models/test_deepseek_v4_construction.py index fa78586b9..0806908c5 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_construction.py @@ -193,6 +193,49 @@ def test_v4_attention_model_owns_head_mhc_parameters( ) +@pytest.mark.parametrize( + "role,first_rank,layer_count", + [ + ("attention", True, 1), + ("attention", False, 1), + ("attention", True, 0), + ("ffn", True, 1), + ], + ids=["attention", "non-first-rank", "empty-model", "ffn"], +) +def test_v4_finalizes_only_attention_owned_mhc_broadcast( + monkeypatch, construction_env, role, first_rank, layer_count +): + model = _make_model( + monkeypatch, role=role, vllm_config=_vllm_config(layer_count=layer_count) + ) + monkeypatch.setattr( + adapter.native, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=first_rank), + ) + if role == "attention" and layer_count: + layer = model.layers[0] + layer.hc_attn_fn.data.fill_(2) + + model.finalize_mhc_broadcast_weights() + + if role != "attention" or not layer_count: + return + if not first_rank: + assert layer.hc_attn_fn_broadcast is None + return + expected = torch.full( + (layer.hc_attn_fn.shape[0], layer.hidden_size), 2.0 * layer.hc_mult + ) + torch.testing.assert_close(layer.hc_attn_fn_broadcast, expected) + broadcast = layer.hc_attn_fn_broadcast + layer.hc_attn_fn.data.fill_(3) + model.finalize_mhc_broadcast_weights() + assert layer.hc_attn_fn_broadcast is broadcast + torch.testing.assert_close(broadcast, expected * 1.5) + + def test_v4_attention_model_allocates_mtp_buffer_for_spec_decode( monkeypatch, construction_env, diff --git a/tests/unit/model_executor/models/test_deepseek_v4_proxy.py b/tests/unit/model_executor/models/test_deepseek_v4_proxy.py index 4818915b8..d4122a494 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_proxy.py @@ -12,11 +12,13 @@ from torch import nn # noqa: E402 +from afd_plugin.model_executor.ffn import remote # noqa: E402 from afd_plugin.model_executor.models import deepseek_v4 as adapter # noqa: E402 from tests.unit.fakes.connector import ProxyConnector as _FakeConnector # noqa: E402 -def test_remote_v4_ffn_sends_token_ids(monkeypatch): +@pytest.mark.parametrize("stage_idx", [None, 2], ids=["metadata-stage", "ubatch-stage"]) +def test_remote_v4_ffn_sends_token_ids(monkeypatch, stage_idx): events: list[tuple] = [] connector = _FakeConnector(events) afd_metadata = SimpleNamespace( @@ -28,17 +30,22 @@ def test_remote_v4_ffn_sends_token_ids(monkeypatch): "get_afd_metadata_from_forward_context", lambda: afd_metadata, ) + forward_context = SimpleNamespace() + if stage_idx is not None: + forward_context.ubatch_idx = stage_idx + expected_stage = 9 if stage_idx is None else stage_idx monkeypatch.setattr( - adapter, + remote, "get_forward_context", - lambda: SimpleNamespace(ubatch_idx=2, slot_mapping={}), + lambda: forward_context, ) def record_yield(hidden_states, *, role): events.append(("yield", hidden_states, role)) + forward_context.ubatch_idx = 23 return hidden_states - monkeypatch.setattr(adapter, "maybe_apply_dbo_yield", record_yield) + monkeypatch.setattr(remote, "maybe_apply_dbo_yield", record_yield) proxy = adapter.RemoteDeepseekV4FFN(layer_idx=3) hidden_states = torch.full((2, 4), 8.0, dtype=torch.float16) input_ids = torch.tensor([11, 13], dtype=torch.int32) @@ -49,15 +56,15 @@ def record_yield(hidden_states, *, role): assert [event[0] for event in events] == ["send", "yield", "recv"] sent_context = events[0][2] assert sent_context.metadata.layer_idx == 3 - assert sent_context.metadata.stage_idx == 2 + assert sent_context.metadata.stage_idx == expected_stage assert sent_context.metadata.seq_lens == [2] assert sent_context.states is None assert set(events[0][3]) == {"input_ids"} assert events[0][3]["input_ids"] is input_ids assert events[0][3]["input_ids"].dtype is torch.int32 assert events[1][2] == "attention" - assert events[2][2] == 2 - assert afd_metadata.stage_idx == 2 + assert events[2][2] == expected_stage + assert afd_metadata.stage_idx == expected_stage assert torch.equal(output, hidden_states * 0.25) @@ -71,7 +78,7 @@ def test_remote_v4_ffn_preserves_ids_in_padding_slots(monkeypatch): lambda: afd_metadata, ) monkeypatch.setattr( - adapter, + remote, "get_forward_context", lambda: SimpleNamespace( ubatch_idx=0, @@ -79,7 +86,7 @@ def test_remote_v4_ffn_preserves_ids_in_padding_slots(monkeypatch): ), ) monkeypatch.setattr( - adapter, + remote, "maybe_apply_dbo_yield", lambda hidden_states, *, role: hidden_states, ) @@ -112,7 +119,7 @@ def test_remote_v4_ffn_rejects_bad_ids_before_sending( lambda: SimpleNamespace(connector=_FakeConnector(events), stage_idx=0), ) monkeypatch.setattr( - adapter, + remote, "get_forward_context", lambda: SimpleNamespace(ubatch_idx=0, slot_mapping={}), ) @@ -129,6 +136,16 @@ def test_remote_v4_ffn_rejects_bad_ids_before_sending( assert events == [] +def test_remote_v4_ffn_preserves_missing_metadata_error(monkeypatch): + monkeypatch.setattr(adapter, "get_afd_metadata_from_forward_context", lambda: None) + proxy = adapter.RemoteDeepseekV4FFN(layer_idx=0) + + with pytest.raises( + RuntimeError, match="^RemoteDeepseekV4FFN requires AFD forward metadata$" + ): + proxy(torch.ones((2, 4)), torch.tensor([11, 13], dtype=torch.int32)) + + def test_v4_decoder_forward_rejects_ffn_role(monkeypatch): class FakeMissingLayer(nn.Module): pass diff --git a/tests/unit/test_envs.py b/tests/unit/test_envs.py index 3bb40ccc2..d22cf8d06 100644 --- a/tests/unit/test_envs.py +++ b/tests/unit/test_envs.py @@ -4,9 +4,21 @@ import pytest -from afd_plugin.envs import AFD_FORCE_BALANCED_TOPK_IDS, force_balanced_topk_ids_enabled +from afd_plugin.envs import ( + AFD_ASYNC_MOE_LAYOUT_LOG, + AFD_FORCE_BALANCED_TOPK_IDS, + async_moe_layout_log_enabled, + force_balanced_topk_ids_enabled, +) +@pytest.mark.parametrize( + "name,enabled", + [ + (AFD_FORCE_BALANCED_TOPK_IDS, force_balanced_topk_ids_enabled), + (AFD_ASYNC_MOE_LAYOUT_LOG, async_moe_layout_log_enabled), + ], +) @pytest.mark.parametrize( "value,expected", [ @@ -20,9 +32,9 @@ ("", False), ], ) -def test_force_balanced_topk_ids_environment(monkeypatch, value, expected): +def test_diagnostic_environment_flags(monkeypatch, name, enabled, value, expected): if value is None: - monkeypatch.delenv(AFD_FORCE_BALANCED_TOPK_IDS, raising=False) + monkeypatch.delenv(name, raising=False) else: - monkeypatch.setenv(AFD_FORCE_BALANCED_TOPK_IDS, value) - assert force_balanced_topk_ids_enabled() is expected + monkeypatch.setenv(name, value) + assert enabled() is expected From ef5ea66b2f9dd8e8e2f2242bbf13cb638ad3c825 Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Mon, 9 Nov 2026 19:50:54 +0800 Subject: [PATCH 3/4] refactor: isolate NPU FFN components Signed-off-by: lirx-pd <616517220@qq.com> --- .../ffn_model/deepseek_v4/npu/v4_weights.py | 2 +- afd_plugin/ffn_faas/platforms/npu/model.py | 2 +- afd_plugin/model_executor/ffn/__init__.py | 2 +- afd_plugin/model_executor/ffn/npu/__init__.py | 3 ++ .../model_executor/ffn/{ => npu}/async_cam.py | 0 .../model_executor/ffn/{ => npu}/compute.py | 2 +- .../ffn/{ => npu}/expert_weights.py | 0 .../npu/deepseek_v2_async_cam_forward.py | 2 +- .../model_executor/models/npu/deepseek_v4.py | 6 ++-- .../models/npu/deepseek_v4_dspark.py | 2 +- .../v1/worker/npu/async_ffn_model_runner.py | 4 +-- docs/design/module/attention_runtime.md | 4 +-- .../module/compatibility_and_patches.md | 2 +- docs/design/module/execution_platforms.md | 2 +- docs/design/module/ffn_runtime.md | 7 ++-- docs/design/module/model_integration.md | 34 +++++++++++-------- .../deepseek_v4_flash/dspark_evidence.py | 2 +- .../dspark_ffn_graph_evidence.py | 2 +- .../deepseek_v4_flash/ffn_graph_evidence.py | 2 +- tests/npu/README.md | 4 +-- tests/npu/async_cam_runtime.py | 2 +- tests/npu/test_async_cam_layered_w4a8.py | 4 +-- tests/npu/test_async_cam_layered_w8a8.py | 4 +-- tests/npu/test_mrv2_dbo.py | 2 +- tests/unit/fakes/async_ffn.py | 4 +-- tests/unit/model_executor/ffn/npu/__init__.py | 3 ++ .../ffn/{ => npu}/test_async_cam.py | 2 +- .../ffn/{ => npu}/test_compute.py | 4 +-- .../test_deepseek_v2_async_cam_forward.py | 2 +- .../test_deepseek_v2_sequence_parallel.py | 2 +- .../models/test_deepseek_v4_cam_layout.py | 2 +- tests/unit/v1/worker/test_model_runner_v2.py | 4 +-- 32 files changed, 65 insertions(+), 54 deletions(-) create mode 100644 afd_plugin/model_executor/ffn/npu/__init__.py rename afd_plugin/model_executor/ffn/{ => npu}/async_cam.py (100%) rename afd_plugin/model_executor/ffn/{ => npu}/compute.py (99%) rename afd_plugin/model_executor/ffn/{ => npu}/expert_weights.py (100%) create mode 100644 tests/unit/model_executor/ffn/npu/__init__.py rename tests/unit/model_executor/ffn/{ => npu}/test_async_cam.py (98%) rename tests/unit/model_executor/ffn/{ => npu}/test_compute.py (99%) diff --git a/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py index 67490d8f1..87ff424d8 100644 --- a/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py +++ b/afd_plugin/ffn_faas/platforms/ffn_model/deepseek_v4/npu/v4_weights.py @@ -24,7 +24,7 @@ V4CheckpointDescription, ) from afd_plugin.ffn_faas.platforms.npu.settings import validate_static_experts -from afd_plugin.model_executor.ffn.expert_weights import ( +from afd_plugin.model_executor.ffn.npu.expert_weights import ( QuantizedExpertWeights, validate_expert_weights, ) diff --git a/afd_plugin/ffn_faas/platforms/npu/model.py b/afd_plugin/ffn_faas/platforms/npu/model.py index 303f0339d..5dd265476 100644 --- a/afd_plugin/ffn_faas/platforms/npu/model.py +++ b/afd_plugin/ffn_faas/platforms/npu/model.py @@ -22,7 +22,7 @@ from vllm.config import VllmConfig from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts - from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights + from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights class RoutedWeightOwner(Protocol): diff --git a/afd_plugin/model_executor/ffn/__init__.py b/afd_plugin/model_executor/ffn/__init__.py index d71e61b11..9c6385422 100644 --- a/afd_plugin/model_executor/ffn/__init__.py +++ b/afd_plugin/model_executor/ffn/__init__.py @@ -1,3 +1,3 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Shared remote calls, expert weight contracts and FFN compute implementations.""" +"""Shared remote FFN calls and backend-specific FFN implementations.""" diff --git a/afd_plugin/model_executor/ffn/npu/__init__.py b/afd_plugin/model_executor/ffn/npu/__init__.py new file mode 100644 index 000000000..b04405dc7 --- /dev/null +++ b/afd_plugin/model_executor/ffn/npu/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""NPU Async CAM layouts, quantized weights and native FFN compute.""" diff --git a/afd_plugin/model_executor/ffn/async_cam.py b/afd_plugin/model_executor/ffn/npu/async_cam.py similarity index 100% rename from afd_plugin/model_executor/ffn/async_cam.py rename to afd_plugin/model_executor/ffn/npu/async_cam.py diff --git a/afd_plugin/model_executor/ffn/compute.py b/afd_plugin/model_executor/ffn/npu/compute.py similarity index 99% rename from afd_plugin/model_executor/ffn/compute.py rename to afd_plugin/model_executor/ffn/npu/compute.py index 256c3f214..6d7538e0c 100644 --- a/afd_plugin/model_executor/ffn/compute.py +++ b/afd_plugin/model_executor/ffn/npu/compute.py @@ -8,7 +8,7 @@ import torch -from afd_plugin.model_executor.ffn.expert_weights import ( +from afd_plugin.model_executor.ffn.npu.expert_weights import ( QuantizedExpertWeights, expert_geometry, validate_expert_weights, diff --git a/afd_plugin/model_executor/ffn/expert_weights.py b/afd_plugin/model_executor/ffn/npu/expert_weights.py similarity index 100% rename from afd_plugin/model_executor/ffn/expert_weights.py rename to afd_plugin/model_executor/ffn/npu/expert_weights.py diff --git a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py index a3ddff99d..e910416e6 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v2_async_cam_forward.py @@ -24,7 +24,7 @@ AFDTransferMetadata, ) from afd_plugin.forward_context import get_afd_metadata_from_forward_context -from afd_plugin.model_executor.ffn.async_cam import ( +from afd_plugin.model_executor.ffn.npu.async_cam import ( CAMDispatchLayout, prepare_cam_dispatch_payload, restore_cam_dispatch_output, diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index f5a7965c9..d70d08f9d 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -36,7 +36,7 @@ AFDTransferMetadata, ) from afd_plugin.forward_context import get_afd_metadata_from_forward_context -from afd_plugin.model_executor.ffn.async_cam import ( +from afd_plugin.model_executor.ffn.npu.async_cam import ( prepare_cam_dispatch_payload, restore_cam_dispatch_output, send_async_cam_remote_call, @@ -68,7 +68,7 @@ AscendW4A8DynamicFusedMoEMethod, ) - from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights + from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights _ATTENTION_ROLE = "attention" @@ -918,7 +918,7 @@ def _extract_quantized_expert_weights( from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm_ascend.quantization.quant_type import QuantType - from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights + from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights layers = [layer for layer in layers if layer.is_moe_layer] # Only the FFN role exposes loaded experts; Ascend's factory returns its diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py b/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py index 0c7aeb5ed..dc6245712 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py @@ -37,7 +37,7 @@ ) if TYPE_CHECKING: - from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights + from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights def _iter_role_weights( diff --git a/afd_plugin/v1/worker/npu/async_ffn_model_runner.py b/afd_plugin/v1/worker/npu/async_ffn_model_runner.py index 0f020228d..70fa62788 100644 --- a/afd_plugin/v1/worker/npu/async_ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/async_ffn_model_runner.py @@ -30,11 +30,11 @@ POLL_SECONDS, FFNGraphPlan, ) -from afd_plugin.model_executor.ffn.compute import ( +from afd_plugin.model_executor.ffn.npu.compute import ( AsyncCAMLayeredExecutor, AsyncCAMMegaExecutor, ) -from afd_plugin.model_executor.ffn.expert_weights import ( +from afd_plugin.model_executor.ffn.npu.expert_weights import ( QuantizedExpertWeights, expert_weight_signature, ) diff --git a/docs/design/module/attention_runtime.md b/docs/design/module/attention_runtime.md index 2412771fc..1e49bed39 100644 --- a/docs/design/module/attention_runtime.md +++ b/docs/design/module/attention_runtime.md @@ -34,7 +34,7 @@ validation_paths: - "tests/unit/v1/worker/npu/test_ubatch_runner.py" - "tests/unit/v1/worker/npu/test_attention_model_runner_v2.py" - "tests/unit/v1/worker/npu/test_attention_metadata.py" - - "tests/unit/model_executor/ffn/test_async_cam.py" + - "tests/unit/model_executor/ffn/npu/test_async_cam.py" - "tests/npu/test_mrv2_dbo.py" - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v4_flash/test_gpu.py" @@ -319,7 +319,7 @@ Enable vLLM native DBO with `--enable-dbo` on Attention. It: Native DBO owns request/token boundaries. Enable DBO on Attention only; Async FFN rejects native DBO flags and continues its receive loop. -`model_executor/ffn/async_cam.py` owns CAM communication layout conversion and +`model_executor/ffn/npu/async_cam.py` owns CAM communication layout conversion and the single-window receive hook, shared by the model adapters. FlashComm1/SP is Attention-local and FFN topology validation is independent. See the CAM async user guide for the supported topology and environment matrix. diff --git a/docs/design/module/compatibility_and_patches.md b/docs/design/module/compatibility_and_patches.md index e9d72f3bf..980562755 100644 --- a/docs/design/module/compatibility_and_patches.md +++ b/docs/design/module/compatibility_and_patches.md @@ -137,7 +137,7 @@ not the package dependency policy. | [`npu/ascend_platform.py`](../../../afd_plugin/compat/patches/npu/ascend_platform.py): `NPUPlatform.check_and_update_config` | Snapshots AFD DBO state, runs upstream normalization, and restores configured `enable_dbo`, `ubatch_size`, and `all2all_backend` in `finally`; non-AFD behavior is unchanged. | Installed through the config facade during AFD NPU normalization and verified again by the worker runtime facade; no version guard. Saves the original on the class and uses a class sentinel. The runtime facade caches success only after the wrapper is installed, so an early missing vLLM-Ascend import remains retryable. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) and [`test_npu_runtime.py`](../../../tests/unit/v1/worker/test_npu_runtime.py). | Remove when vLLM-Ascend recognizes plugin-owned DBO workers or no longer clears these fields. | | [`npu/ascend_config.py`](../../../afd_plugin/compat/patches/npu/ascend_config.py): `vllm_ascend.ascend_config.init_ascend_config`, loaded factory aliases, and the AFD DP child entry point | Excludes the AFD namespace from native constructor arguments; preserves config identity, validation, singleton publication, and cache refresh. DP children install the factory before native Ascend initialization. | Installed by the NPU config facade; rebinds already-loaded aliases. The DP entry point is selected only for AFD synchronous Attention or FFN. | [`test_ascend_config_lifecycle.py`](../../../tests/unit/compat/npu/test_ascend_config_lifecycle.py), [`test_ascend_engine_core_binding.py`](../../../tests/unit/compat/npu/test_ascend_engine_core_binding.py). | Remove when Ascend accepts plugin namespaces and loads general plugins before config validation in every DP child. | | [`npu/mla_graph.py`](../../../afd_plugin/compat/patches/npu/mla_graph.py): `vllm_ascend.attention.mla_v1.get_graph_params` | Resolves an AFD ubatch-owned `GraphParams` registry from the active forward context and otherwise delegates to the saved upstream process-global resolver. | Called through `apply_afd_ascend_patches_if_needed` during NPU worker startup; no version guard. Saves the original on the upstream module and uses a module sentinel. | [`test_runtime.py`](../../../tests/unit/compat/test_runtime.py) covers AFD-context resolution, upstream fallback, and idempotence; [`test_npu_mla_graph.py`](../../../tests/unit/v1/worker/test_npu_mla_graph.py) covers the owning graph lifecycle. | Remove when vLLM-Ascend accepts a forward-context-local MLA graph registry or exposes an equivalent resolver hook. | -| [`npu/mrv2_ubatching.py`](../../../afd_plugin/compat/patches/npu/mrv2_ubatching.py): native ubatch factory, dispatch, child DP metadata, input slicing, context and stream aliases | Builds `NPUUBatchRunner`, enables Async CAM DBO at DP1, omits Async child DP coordination, applies native truncation to Ascend CPU sequence lengths, preserves model-specific PP slices/nested auxiliary outputs, and provides NPU stream/failure handling during eager and capture. Native MRV2 owns execution and graph lifecycle. | Scoped to AFD initialization, preparation or execution; every replaced alias is restored in `finally`. The scope temporarily changes module symbols, so it is not an independent concurrent-runner API. | [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py), [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/test_async_cam.py), [`test_mrv2_dbo.py`](../../../tests/npu/test_mrv2_dbo.py). | Remove when native vLLM/Ascend supports NPU child builders and streams, optional child DP metadata, DP1 DBO and peer cancellation. | +| [`npu/mrv2_ubatching.py`](../../../afd_plugin/compat/patches/npu/mrv2_ubatching.py): native ubatch factory, dispatch, child DP metadata, input slicing, context and stream aliases | Builds `NPUUBatchRunner`, enables Async CAM DBO at DP1, omits Async child DP coordination, applies native truncation to Ascend CPU sequence lengths, preserves model-specific PP slices/nested auxiliary outputs, and provides NPU stream/failure handling during eager and capture. Native MRV2 owns execution and graph lifecycle. | Scoped to AFD initialization, preparation or execution; every replaced alias is restored in `finally`. The scope temporarily changes module symbols, so it is not an independent concurrent-runner API. | [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py), [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/npu/test_async_cam.py), [`test_mrv2_dbo.py`](../../../tests/npu/test_mrv2_dbo.py). | Remove when native vLLM/Ascend supports NPU child builders and streams, optional child DP metadata, DP1 DBO and peer cancellation. | ## Temporary CUDA DeepSeek-V4 DBO fixes diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md index dff666576..1b8e06351 100644 --- a/docs/design/module/execution_platforms.md +++ b/docs/design/module/execution_platforms.md @@ -35,7 +35,7 @@ validation_paths: - "tests/unit/v1/worker/test_npu_mla_graph.py" - "tests/unit/v1/worker/test_npu_runtime.py" - "tests/unit/v1/worker/test_npu_model_runner_v2.py" - - "tests/unit/model_executor/ffn/test_async_cam.py" + - "tests/unit/model_executor/ffn/npu/test_async_cam.py" - "tests/npu/**" - "tests/e2e/models/deepseek_v2_lite/test_deepseek_v2_lite.py" - "tests/e2e/models/deepseek_v4_flash/test_gpu.py" diff --git a/docs/design/module/ffn_runtime.md b/docs/design/module/ffn_runtime.md index 5179cca6e..729ee0b14 100644 --- a/docs/design/module/ffn_runtime.md +++ b/docs/design/module/ffn_runtime.md @@ -28,7 +28,7 @@ depends_on: - "compatibility_and_patches.md" validation_paths: - "tests/unit/ffn_faas/test_*.py" - - "tests/unit/model_executor/ffn/test_compute.py" + - "tests/unit/model_executor/ffn/npu/test_compute.py" - "tests/npu/test_async_cam_layered_w8a8.py" - "tests/unit/v1/worker/test_ffn_metadata.py" - "tests/unit/v1/worker/test_ffn_model_runner.py" @@ -326,8 +326,9 @@ and validation scope are documented in the The Async runner loads target/draft `QuantizedExpertWeights` through the model export contract, verifies layer coverage and selects Layered or Mega directly. -`model_executor/ffn/expert_weights.py` owns shared shape validation, logical -geometry and graph signatures. `compute.py` owns native operator handles, +`model_executor/ffn/npu/expert_weights.py` owns the Ascend quantized weight +layout validation, logical geometry and graph signatures shared by Layered +and Mega. `model_executor/ffn/npu/compute.py` owns native operator handles, workspace, zero bias and layer maps. There is no execution layer beneath the ModelRunner and no access to private model parameters from compute. diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 60e5a54f0..604dae4cc 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -60,14 +60,16 @@ make a backend-specific worker class the shared model API. ## Shared FFN components `model_executor/models/` owns model structure and checkpoint policy. -`model_executor/ffn/` contains four modules with separate responsibilities: +`model_executor/ffn/remote.py` contains the remote proxies shared by CUDA and +Ascend models. `model_executor/ffn/npu/` owns the NPU-specific Async CAM +layout, quantized weight contracts and native compute implementations: | Component | Responsibility | Dependencies | | --- | --- | --- | | `remote.py` | Connector-independent dense/MoE proxies and routed TP padding | Connector contracts, public vLLM helpers and package forward context; no CAM, concrete model or compute imports | -| `async_cam.py` | CAM model-side TP/SP layout conversion, padding, output restoration and native DBO receive handoff | Connector contracts and native DBO hooks; no concrete model or compute imports | -| `expert_weights.py` | Loaded `QuantizedExpertWeights`, homogeneous layout validation, geometry and graph signatures | Torch tensors; no native operator lookup or workspace allocation | -| `compute.py` | Independent NPU Layered and Mega native implementations and their buffers | Expert weights and explicit configuration/communication arguments; no model internals | +| `npu/async_cam.py` | CAM model-side TP/SP layout conversion, padding, output restoration and native DBO receive handoff | Connector contracts and native DBO hooks; no concrete model or compute imports | +| `npu/expert_weights.py` | Loaded `QuantizedExpertWeights`, Ascend quantized layout validation, geometry and graph signatures | Torch tensors with the Layered/Mega W4A8/W8A8 storage and scale contracts; no native operator lookup or workspace allocation | +| `npu/compute.py` | Independent NPU Layered and Mega native implementations and their buffers | Expert weights and explicit configuration/communication arguments; no model internals | ```mermaid flowchart TD @@ -75,19 +77,23 @@ flowchart TD Runner --> Model[Concrete model adapter] Model --> Remote[ffn/remote.py] Remote --> Connector[Connector / Session] - Model --> CAM[ffn/async_cam.py] + Model --> CAM[ffn/npu/async_cam.py] CAM --> Connector - Model --> Weights[ffn/expert_weights.py] - Runner --> Compute[ffn/compute.py: Layered OR Mega] + Model --> Weights[ffn/npu/expert_weights.py] + Runner --> Compute[ffn/npu/compute.py: Layered OR Mega] Compute --> Weights Compute --> Native[Native NPU operators] ``` DeepSeek V2, Qwen3, Qwen3.5 and CUDA/Ascend V4 use the shared remote components. -The CAM helpers stay in `ffn/async_cam.py` because they adapt model tensors -to remote FFN calls across model families. Connector/Session owns transport -windows, pending work and Push/Pull protocol state. The generic proxies do -not import the CAM module, and neither module imports compute. +The NPU CAM helpers adapt model tensors to remote FFN calls across Ascend +model families. Their placement under `ffn/npu/` reflects the backend-specific +CAM contract. The quantized weight contract is shared by Layered and Mega; +its packed weights, encoded scales and graph signatures describe the Ascend +layout rather than a cross-platform quantization interface. Connector/Session +owns transport windows, pending work and Push/Pull protocol state. The shared +remote module does not import `ffn/npu/`, and neither remote module imports +compute. Gate, Hash routing, shared experts, dense-layer behavior and DSpark target/draft identities remain model-owned. The Async runner receives loaded weights through `get_quantized_expert_weights()` and selects one compute implementation. Mega @@ -118,9 +124,9 @@ SFA/MLA/DSA-specific handling stays explicit inside that worker-side module. | Qwen3 MoE role-aware model and weight loading | [`qwen3_moe.py`](../../../afd_plugin/model_executor/models/qwen3_moe.py) | [`test_qwen3_moe_construction.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_construction.py), [`test_qwen3_moe_weight_policy.py`](../../../tests/unit/model_executor/models/test_qwen3_moe_weight_policy.py) | | CUDA remote-experts boundary | [`ffn/remote.py`](../../../afd_plugin/model_executor/ffn/remote.py), [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py) | [`test_p2p_experts_contract.py`](../../../tests/unit/connectors/test_p2p_experts_contract.py), [`test_deepseek_v2_proxy.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_proxy.py) | | Forward-context adapter | [`forward_context.py`](../../../afd_plugin/forward_context.py) | [`test_forward_context.py`](../../../tests/unit/test_forward_context.py) | -| NPU Async CAM native DBO handoff | [`ffn/async_cam.py`](../../../afd_plugin/model_executor/ffn/async_cam.py), [`npu/ubatch_runner.py`](../../../afd_plugin/v1/worker/npu/ubatch_runner.py) | [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/test_async_cam.py), [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py) | +| NPU Async CAM native DBO handoff | [`ffn/npu/async_cam.py`](../../../afd_plugin/model_executor/ffn/npu/async_cam.py), [`npu/ubatch_runner.py`](../../../afd_plugin/v1/worker/npu/ubatch_runner.py) | [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/npu/test_async_cam.py), [`test_npu_model_runner_v2.py`](../../../tests/unit/v1/worker/test_npu_model_runner_v2.py) | | Ascend Attention-side gate | [`npu/deepseek_v2_attention_gate.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v2_attention_gate.py) | [`test_deepseek_v2_attention_gate.py`](../../../tests/unit/model_executor/models/npu/test_deepseek_v2_attention_gate.py), profile I/O in [`test_deepseek_v2_async_cam_forward.py`](../../../tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py) | -| Ascend CAM dispatch layout | [`ffn/async_cam.py`](../../../afd_plugin/model_executor/ffn/async_cam.py) | [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/test_async_cam.py) | +| Ascend CAM dispatch layout | [`ffn/npu/async_cam.py`](../../../afd_plugin/model_executor/ffn/npu/async_cam.py) | [`test_async_cam.py`](../../../tests/unit/model_executor/ffn/npu/test_async_cam.py) | | Ascend stage metadata and RoPE ownership | [`worker/npu/attention_metadata.py`](../../../afd_plugin/v1/worker/npu/attention_metadata.py) | [`test_attention_metadata.py`](../../../tests/unit/v1/worker/npu/test_attention_metadata.py) | | Ascend external-routing MoE execution | [`npu/deepseek_v4_external_routing_moe.py`](../../../afd_plugin/model_executor/models/npu/deepseek_v4_external_routing_moe.py) | [`test_deepseek_v4_external_routing_moe_cpu.py`](../../../tests/unit/model_executor/models/test_deepseek_v4_external_routing_moe_cpu.py) | @@ -308,7 +314,7 @@ global provider. Async CAM reuses native DBO child slices, attention metadata and execution. The generic adapter installs `afd_metadata` on the native child context; CAM transfer metadata and TP/SP layout conversion and the single-window receive hook share the -model-side communication boundary in `model_executor/ffn/async_cam.py`. Metadata ownership and the live connector reference +model-side communication boundary in `model_executor/ffn/npu/async_cam.py`. Metadata ownership and the live connector reference remain **draft** under [#88](https://github.com/JiusiServe/afd-plugin/issues/88) and [#105](https://github.com/JiusiServe/afd-plugin/issues/105). diff --git a/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py b/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py index e9fe381ab..02000481f 100644 --- a/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/dspark_evidence.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector -from afd_plugin.model_executor.ffn.compute import CAM_LAYER_INDEX +from afd_plugin.model_executor.ffn.npu.compute import CAM_LAYER_INDEX from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker diff --git a/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py b/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py index 9795bdf13..36d13ab0a 100644 --- a/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/dspark_ffn_graph_evidence.py @@ -17,7 +17,7 @@ import torch -from afd_plugin.model_executor.ffn.compute import ( +from afd_plugin.model_executor.ffn.npu.compute import ( CAM_LAYER_INDEX, AsyncCAMLayeredExecutor, ) diff --git a/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py b/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py index 1cb4bffc8..66537c6e4 100644 --- a/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py +++ b/tests/e2e/models/deepseek_v4_flash/ffn_graph_evidence.py @@ -15,7 +15,7 @@ import torch from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector -from afd_plugin.model_executor.ffn.compute import AsyncCAMMegaExecutor +from afd_plugin.model_executor.ffn.npu.compute import AsyncCAMMegaExecutor from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.attention_model_runner_v2 import ( AFDNPUAttentionModelRunnerV2, diff --git a/tests/npu/README.md b/tests/npu/README.md index 05b4221e8..3f8fb3346 100644 --- a/tests/npu/README.md +++ b/tests/npu/README.md @@ -74,7 +74,7 @@ Run the focused contract checks in the matching runtime environment: python3 -m pytest -q -ra \ tests/unit/compat/npu/test_runtime_config.py \ tests/unit/config/test_validation.py \ - tests/unit/model_executor/ffn/test_async_cam.py \ + tests/unit/model_executor/ffn/npu/test_async_cam.py \ tests/unit/v1/worker/npu/test_attention_metadata.py \ tests/unit/v1/worker/test_npu_model_runner_v2.py \ tests/unit/v1/worker/test_npu_ubatch_outputs.py \ @@ -278,7 +278,7 @@ device and no collective. Run the CPU contracts first, then rebuild the matching extension/vendor library and run the device checks: ```bash -pytest -q tests/unit/test_envs.py tests/unit/model_executor/ffn/test_compute.py +pytest -q tests/unit/test_envs.py tests/unit/model_executor/ffn/npu/test_compute.py SOC_VERSION=910c pytest -q tests/npu/test_gmm_layered.py \ tests/npu/test_gmm_swiglu_quant_v2_layered.py ``` diff --git a/tests/npu/async_cam_runtime.py b/tests/npu/async_cam_runtime.py index 7d1a85c62..edc7cbee6 100644 --- a/tests/npu/async_cam_runtime.py +++ b/tests/npu/async_cam_runtime.py @@ -29,7 +29,7 @@ from afd_plugin.config import parse_afd_config from afd_plugin.connectors.metadata import AFDTransferContext, AFDTransferMetadata from afd_plugin.connectors.npu.async_cam import CAMAsyncAFDConnector -from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights +from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights from afd_plugin.v1.worker.npu.async_ffn_model_runner import AsyncNPUFFNModelRunner from afd_plugin.v1.worker.npu.ffn_worker import AFDNPUFFNWorker diff --git a/tests/npu/test_async_cam_layered_w4a8.py b/tests/npu/test_async_cam_layered_w4a8.py index f8f20e90c..f4387717e 100644 --- a/tests/npu/test_async_cam_layered_w4a8.py +++ b/tests/npu/test_async_cam_layered_w4a8.py @@ -27,8 +27,8 @@ def test_layered_w4a8_valid_rows_and_capacity_tail( torch = pytest.importorskip("torch") torch_npu = pytest.importorskip("torch_npu") from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available - from afd_plugin.model_executor.ffn.compute import AsyncCAMLayeredExecutor - from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights + from afd_plugin.model_executor.ffn.npu.compute import AsyncCAMLayeredExecutor + from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights ensure_cam_async_ops_available() torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) diff --git a/tests/npu/test_async_cam_layered_w8a8.py b/tests/npu/test_async_cam_layered_w8a8.py index b04dc80f0..16414c615 100644 --- a/tests/npu/test_async_cam_layered_w8a8.py +++ b/tests/npu/test_async_cam_layered_w8a8.py @@ -27,8 +27,8 @@ def w8a8_runtime(request): torch = pytest.importorskip("torch") torch_npu = pytest.importorskip("torch_npu") from afd_plugin.compat.npu.ops import ensure_cam_async_ops_available - from afd_plugin.model_executor.ffn.compute import AsyncCAMLayeredExecutor - from afd_plugin.model_executor.ffn.expert_weights import QuantizedExpertWeights + from afd_plugin.model_executor.ffn.npu.compute import AsyncCAMLayeredExecutor + from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights ensure_cam_async_ops_available() torch.npu.set_device(int(os.environ.get("LOCAL_RANK", "0"))) diff --git a/tests/npu/test_mrv2_dbo.py b/tests/npu/test_mrv2_dbo.py index a621785bb..8d6d03a42 100644 --- a/tests/npu/test_mrv2_dbo.py +++ b/tests/npu/test_mrv2_dbo.py @@ -42,7 +42,7 @@ from vllm_ascend.worker.v2.utils import torch_cuda_wrapper # noqa: E402 from afd_plugin.connectors import AFDTransferContext, AFDTransferMetadata # noqa: E402 -from afd_plugin.model_executor.ffn.async_cam import ( # noqa: E402 +from afd_plugin.model_executor.ffn.npu.async_cam import ( # noqa: E402 send_async_cam_remote_call, ) from afd_plugin.v1.worker.npu.ubatch_runner import NPUUBatchRunner # noqa: E402 diff --git a/tests/unit/fakes/async_ffn.py b/tests/unit/fakes/async_ffn.py index 7f918eddf..51f3fc473 100644 --- a/tests/unit/fakes/async_ffn.py +++ b/tests/unit/fakes/async_ffn.py @@ -42,8 +42,8 @@ def stub(name, **attributes): "afd_plugin.compat.npu.ops", "afd_plugin.compat.npu.forward_context", "afd_plugin.compat.npu.profiler", - "afd_plugin.model_executor.ffn.compute", - "afd_plugin.model_executor.ffn.expert_weights", + "afd_plugin.model_executor.ffn.npu.compute", + "afd_plugin.model_executor.ffn.npu.expert_weights", "afd_plugin.v1.worker.ffn_metadata", ): stub(name) diff --git a/tests/unit/model_executor/ffn/npu/__init__.py b/tests/unit/model_executor/ffn/npu/__init__.py new file mode 100644 index 000000000..6d1f3d085 --- /dev/null +++ b/tests/unit/model_executor/ffn/npu/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""NPU FFN contracts, isolated from connector test module names.""" diff --git a/tests/unit/model_executor/ffn/test_async_cam.py b/tests/unit/model_executor/ffn/npu/test_async_cam.py similarity index 98% rename from tests/unit/model_executor/ffn/test_async_cam.py rename to tests/unit/model_executor/ffn/npu/test_async_cam.py index 592fbe31a..9c21fe569 100644 --- a/tests/unit/model_executor/ffn/test_async_cam.py +++ b/tests/unit/model_executor/ffn/npu/test_async_cam.py @@ -13,7 +13,7 @@ from vllm.v1.worker import ubatching # noqa: E402 from afd_plugin.connectors import AFDTransferContext, AFDTransferMetadata # noqa: E402 -from afd_plugin.model_executor.ffn import async_cam # noqa: E402 +from afd_plugin.model_executor.ffn.npu import async_cam # noqa: E402 def test_plain_tp_cam_boundary_shards_and_restores_replicated_tokens(monkeypatch): diff --git a/tests/unit/model_executor/ffn/test_compute.py b/tests/unit/model_executor/ffn/npu/test_compute.py similarity index 99% rename from tests/unit/model_executor/ffn/test_compute.py rename to tests/unit/model_executor/ffn/npu/test_compute.py index d68ffdd5d..b0c2d4388 100644 --- a/tests/unit/model_executor/ffn/test_compute.py +++ b/tests/unit/model_executor/ffn/npu/test_compute.py @@ -15,11 +15,11 @@ torch = pytest.importorskip("torch") -from afd_plugin.model_executor.ffn.compute import ( # noqa: E402 +from afd_plugin.model_executor.ffn.npu.compute import ( # noqa: E402 AsyncCAMLayeredExecutor, AsyncCAMMegaExecutor, ) -from afd_plugin.model_executor.ffn.expert_weights import ( # noqa: E402 +from afd_plugin.model_executor.ffn.npu.expert_weights import ( # noqa: E402 QuantizedExpertWeights, expert_weight_signature, ) diff --git a/tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py b/tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py index 1f5892ec6..79140893d 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_async_cam_forward.py @@ -31,7 +31,7 @@ def test_async_cam_profile_forward_runs_matched_connector_io( ) connector_calls: list[str] = [] - from afd_plugin.model_executor.ffn import async_cam + from afd_plugin.model_executor.ffn.npu import async_cam monkeypatch.setattr( async_cam, diff --git a/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py b/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py index 59c50e8ca..1754082ac 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_sequence_parallel.py @@ -13,7 +13,7 @@ pytest.importorskip("vllm") -from afd_plugin.model_executor.ffn import async_cam # noqa: E402 +from afd_plugin.model_executor.ffn.npu import async_cam # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 from afd_plugin.model_executor.models.npu import ( # noqa: E402 deepseek_v2_async_cam_forward as schedule, diff --git a/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py b/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py index ad6402d9b..36e8ee692 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_cam_layout.py @@ -11,7 +11,7 @@ pytest.importorskip("vllm") pytest.importorskip("vllm_ascend.models.deepseek_v4.model") -from afd_plugin.model_executor.ffn import async_cam # noqa: E402 +from afd_plugin.model_executor.ffn.npu import async_cam # noqa: E402 from afd_plugin.model_executor.models.npu import deepseek_v4 as adapter # noqa: E402 from afd_plugin.model_executor.models.npu import ( # noqa: E402 deepseek_v4_attention_gate, diff --git a/tests/unit/v1/worker/test_model_runner_v2.py b/tests/unit/v1/worker/test_model_runner_v2.py index 8c7e35c58..5db35f5a7 100644 --- a/tests/unit/v1/worker/test_model_runner_v2.py +++ b/tests/unit/v1/worker/test_model_runner_v2.py @@ -25,9 +25,7 @@ from vllm.v1.worker.gpu.ubatch_utils import UBatchState # noqa: E402 from vllm.v1.worker.ubatch_utils import UBatchSlice # noqa: E402 -from afd_plugin.model_executor.models import ( # noqa: E402 - forward_context as afd_forward_context, -) +from afd_plugin import forward_context as afd_forward_context # noqa: E402 from afd_plugin.v1.worker import ( # noqa: E402 attention_model_runner_v2 as v2_runner_module, ) From afff5cbf8e9586553ff6398788b99d0e1bade5cb Mon Sep 17 00:00:00 2001 From: lirx-pd <616517220@qq.com> Date: Tue, 10 Nov 2026 19:55:54 -0800 Subject: [PATCH 4/4] refactor: consolidate model weight ownership Signed-off-by: lirx-pd <616517220@qq.com> --- .../model_executor/models/_weight_utils.py | 6 ++- .../model_executor/models/deepseek_v2.py | 20 ++++---- .../model_executor/models/deepseek_v4.py | 15 +++--- .../model_executor/models/npu/deepseek_v4.py | 36 ++++++------- .../models/npu/deepseek_v4_dspark.py | 51 ++++++++++++------- afd_plugin/model_executor/models/qwen3_5.py | 14 ++--- afd_plugin/model_executor/models/qwen3_moe.py | 17 +++---- .../models/npu/test_deepseek_v4.py | 30 +++++------ .../models/npu/test_deepseek_v4_dspark.py | 33 ++++++++---- tests/unit/package/test_package.py | 11 ++++ 10 files changed, 139 insertions(+), 94 deletions(-) diff --git a/afd_plugin/model_executor/models/_weight_utils.py b/afd_plugin/model_executor/models/_weight_utils.py index 236e33cfb..4a77658c3 100644 --- a/afd_plugin/model_executor/models/_weight_utils.py +++ b/afd_plugin/model_executor/models/_weight_utils.py @@ -3,10 +3,14 @@ """Checkpoint path parsing and one-pass filtering; models own role policies.""" from collections.abc import Callable, Iterable, Iterator -from typing import TypeVar +from typing import Final, TypeVar _Weight = TypeVar("_Weight") +ATTENTION_ROLE: Final[frozenset[str]] = frozenset(("attention",)) +FFN_ROLE: Final[frozenset[str]] = frozenset(("ffn",)) +BOTH_ROLES: Final[frozenset[str]] = ATTENTION_ROLE | FFN_ROLE + def layer_weight_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: parts = name.split(".") diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 110b8a073..078bd5a67 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -36,6 +36,9 @@ _unroute_tp_padding, ) from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + BOTH_ROLES, + FFN_ROLE, filter_role_weights, layer_weight_path, ) @@ -60,9 +63,6 @@ def _is_moe_layer(config: _DeepseekAdapterConfig, layer_idx: int) -> bool: ) -_ATTENTION_ROLE = frozenset(("attention",)) -_FFN_ROLE = frozenset(("ffn",)) -_BOTH_ROLES = frozenset(("attention", "ffn")) # Connectors whose Attention runs the shared experts, so their FFN builds # none: Async CAM, and DeepEP M2N whose FFN runs routed experts only. _ATTENTION_SHARED_EXPERTS_CONNECTORS = ( @@ -81,23 +81,23 @@ def _checkpoint_weight_roles( """Classify one native checkpoint path by its AFD execution owner.""" layer_path = layer_weight_path(name) if layer_path is None: - return _BOTH_ROLES + return BOTH_ROLES layer_idx, stage, remainder = layer_path if stage == "self_attn": - return _ATTENTION_ROLE + return ATTENTION_ROLE if stage != "mlp": - return _BOTH_ROLES + return BOTH_ROLES if not _is_moe_layer(config, layer_idx): - return _ATTENTION_ROLE if compute_gate_on_attention else _FFN_ROLE + return ATTENTION_ROLE if compute_gate_on_attention else FFN_ROLE if remainder and remainder[0] == "shared_experts" and attention_shared_experts: - return _ATTENTION_ROLE + return ATTENTION_ROLE is_moe_gate = bool(remainder) and remainder[0] == "gate" if is_moe_gate and compute_gate_on_attention: - return _BOTH_ROLES - return _FFN_ROLE + return BOTH_ROLES + return FFN_ROLE class GateOnlyRemoteMoE(RemoteFFNProxy): diff --git a/afd_plugin/model_executor/models/deepseek_v4.py b/afd_plugin/model_executor/models/deepseek_v4.py index 36d6b817e..b38b3e7cb 100644 --- a/afd_plugin/model_executor/models/deepseek_v4.py +++ b/afd_plugin/model_executor/models/deepseek_v4.py @@ -22,14 +22,13 @@ from afd_plugin.forward_context import get_afd_metadata_from_forward_context from afd_plugin.model_executor.ffn.remote import RemoteFFNProxy from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + BOTH_ROLES, + FFN_ROLE, filter_role_weights, layer_weight_path, ) -_ATTENTION_ROLE = frozenset(("attention",)) -_FFN_ROLE = frozenset(("ffn",)) -_BOTH_ROLES = frozenset(("attention", "ffn")) - def _checkpoint_weight_roles(name: str) -> frozenset[str]: """Classify a native DeepSeek-V4 checkpoint path by execution owner.""" @@ -41,14 +40,14 @@ def _checkpoint_weight_roles(name: str) -> frozenset[str]: "model.hc_head_base", "model.hc_head_scale", }: - return _ATTENTION_ROLE + return ATTENTION_ROLE layer_path = layer_weight_path(name) if layer_path is None: - return _BOTH_ROLES + return BOTH_ROLES _, stage, _ = layer_path if stage == "ffn": - return _FFN_ROLE - return _ATTENTION_ROLE + return FFN_ROLE + return ATTENTION_ROLE class RemoteDeepseekV4FFN(RemoteFFNProxy): diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index d70d08f9d..8a09e3341 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -43,6 +43,9 @@ ) from afd_plugin.model_executor.ffn.remote import RemoteFFNProxy from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + BOTH_ROLES, + FFN_ROLE, filter_role_weights, layer_weight_path, ) @@ -71,11 +74,6 @@ from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights -_ATTENTION_ROLE = "attention" -_FFN_ROLE = "ffn" -_BOTH_ROLES = frozenset((_ATTENTION_ROLE, _FFN_ROLE)) - - def _refresh_ascend_fused_moe() -> None: """Bind native DSV4 MoE construction to the Ascend implementation.""" # vLLM-Ascend applies this replacement during platform initialization, @@ -106,14 +104,14 @@ def _checkpoint_weight_roles( layer_path = layer_weight_path(name) if layer_path is None: - return _BOTH_ROLES + return BOTH_ROLES _, stage, remainder = layer_path if stage in ("attn", "self_attn"): - return frozenset((_ATTENTION_ROLE,)) + return ATTENTION_ROLE if stage in ("ffn", "mlp"): if attention_shared_experts and remainder and remainder[0] == "shared_experts": - return frozenset((_ATTENTION_ROLE,)) + return ATTENTION_ROLE if remainder and remainder[0] == "gate": # Every router parameter, including the Hash id table, belongs to # each role that built a router: with the gate on Attention, the @@ -122,12 +120,12 @@ def _checkpoint_weight_roles( # With the gate on FFN the Attention MoE slot is parameter-free, so # Attention must not receive any of them. if attn_owns_gate: - return _BOTH_ROLES - return frozenset((_FFN_ROLE,)) - return frozenset((_FFN_ROLE,)) + return BOTH_ROLES + return FFN_ROLE + return FFN_ROLE # HC parameters and any future shared layer parameters are required by # both role-local model instances. - return _BOTH_ROLES + return BOTH_ROLES class AFDDeepseekV4RemoteMoE(RemoteFFNProxy): @@ -435,7 +433,7 @@ def __init__( | AFDDeepseekV4RemoteMoE | native.DeepseekV4MoE ) - if afd_config.role == _ATTENTION_ROLE: + if afd_config.role == "attention": self.self_attn = native.DeepseekV4Attention( vllm_config=vllm_config, config=config, @@ -464,7 +462,7 @@ def __init__( ) else: self.mlp = AFDDeepseekV4RemoteMoE(layer_idx=layer_idx) - elif afd_config.role == _FFN_ROLE: + elif afd_config.role == "ffn": self.self_attn = native.PPMissingLayer() _refresh_ascend_fused_moe() moe_config = config @@ -495,7 +493,7 @@ def __init__( # FFN path, even though it also computes the gate on the Attention side. if ( self.uses_external_routing - and afd_config.role == _FFN_ROLE + and afd_config.role == "ffn" and isinstance(self.mlp, native.DeepseekV4MoE) ): from afd_plugin.model_executor.models.npu import ( @@ -654,7 +652,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.vocab_size = config.vocab_size self.is_v32 = hasattr(config, "index_topk") # ### PATCH START: Attention-owned index buffer - if self.is_v32 and afd_config.role == _ATTENTION_ROLE: + if self.is_v32 and afd_config.role == "attention": topk_tokens = config.index_topk topk_indices_buffer = torch.empty( vllm_config.scheduler_config.max_num_batched_tokens, @@ -696,7 +694,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): for layer in self.layers: if isinstance(layer, native.DeepseekV4DecoderLayer): # ### PATCH START: FFN has no Attention implementation - if layer.afd_role != _ATTENTION_ROLE: + if layer.afd_role != "attention": continue # ### PATCH END: FFN has no Attention implementation impl = layer.self_attn.dsa_attn.dsa_attn.impl @@ -905,7 +903,9 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: classify=partial( _checkpoint_weight_roles, attn_owns_gate=bool(self.afd_config.compute_gate_on_attention), - attention_shared_experts=bool(self.afd_config.compute_gate_on_attention), + attention_shared_experts=bool( + self.afd_config.compute_gate_on_attention + ), ), ) return super().load_weights(role_weights) diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py b/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py index dc6245712..592f50120 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_dspark.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections.abc import Iterable, Iterator +from collections.abc import Iterable from typing import TYPE_CHECKING, cast import torch @@ -31,6 +31,12 @@ from vllm_ascend.utils import get_rotation_path from afd_plugin.config import parse_afd_config +from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + BOTH_ROLES, + FFN_ROLE, + filter_role_weights, +) from afd_plugin.model_executor.models.npu.deepseek_v4 import ( AFDDeepseekV4DecoderLayer, _extract_quantized_expert_weights, @@ -40,21 +46,24 @@ from afd_plugin.model_executor.ffn.npu.expert_weights import QuantizedExpertWeights -def _iter_role_weights( - weights: Iterable[tuple[str, torch.Tensor]], *, role: str -) -> Iterator[tuple[str, torch.Tensor]]: - """Filter once before the native loader resolves its strict parameter names.""" - for name, weight in weights: - parts = name.split(".") - routed = ( - len(parts) > 4 - and parts[0] == "mtp" - and parts[2] in ("ffn", "mlp") - and parts[3] in ("experts", "gate") - ) - expert = routed and parts[3] == "experts" - if (role == "ffn" and routed) or (role == "attention" and not expert): - yield name, weight +def _dspark_weight_roles(name: str) -> frozenset[str]: + """Classify one DSpark MTP checkpoint path by its AFD execution owner. + + The draft stack has no ``layers.`` segment, so this cannot reuse + ``layer_weight_path``. Expert weights are FFN-only; the router (including + its Hash id table) belongs to every role that builds a router, matching the + native MoE ownership rule; everything else stays with Attention. + """ + parts = name.split(".") + routed = ( + len(parts) > 4 + and parts[0] == "mtp" + and parts[2] in ("ffn", "mlp") + and parts[3] in ("experts", "gate") + ) + if not routed: + return ATTENTION_ROLE + return FFN_ROLE if parts[3] == "experts" else BOTH_ROLES class AFDDeepseekV4DSparkModel(native.DeepseekV4DSparkModel): @@ -241,7 +250,15 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: # ### PATCH END: construct only this role's draft modules def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - loaded = super().load_weights(_iter_role_weights(weights, role=self.afd_role)) + # None must not bypass filtering for a role-specific DSpark model. + role: str = self.afd_role if self.afd_role in ("attention", "ffn") else "" + loaded = super().load_weights( + filter_role_weights( + weights, + role=role, + classify=_dspark_weight_roles, + ) + ) native.logger.info( "AFD DSpark role=%s remote MoE layer IDs=%s loaded parameters=%d", self.afd_role, diff --git a/afd_plugin/model_executor/models/qwen3_5.py b/afd_plugin/model_executor/models/qwen3_5.py index 05648e6e4..1e31262c3 100644 --- a/afd_plugin/model_executor/models/qwen3_5.py +++ b/afd_plugin/model_executor/models/qwen3_5.py @@ -14,12 +14,12 @@ from afd_plugin.config import parse_optional_afd_config from afd_plugin.model_executor.ffn.remote import AFDAttentionFusedMoE from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + FFN_ROLE, filter_role_weights, layer_weight_path, ) -_ATTENTION_ROLE = frozenset(("attention",)) -_FFN_ROLE = frozenset(("ffn",)) _NO_ROLES: frozenset[str] = frozenset() @@ -44,21 +44,21 @@ def _checkpoint_weight_roles( layer_path = layer_weight_path(name) if layer_path is None: # Embeddings, final norm, and lm_head are required only by Attention. - return _ATTENTION_ROLE + return ATTENTION_ROLE _layer_idx, stage, remainder = layer_path if stage in ("linear_attn", "self_attn"): - return _ATTENTION_ROLE + return ATTENTION_ROLE if stage != "mlp": - return _ATTENTION_ROLE + return ATTENTION_ROLE if remainder and remainder[0] == "gate": - return _FFN_ROLE + return FFN_ROLE if remainder and remainder[0] in ( "experts", "shared_expert", "shared_expert_gate", ): - return _FFN_ROLE + return FFN_ROLE raise RuntimeError(f"unclassified Qwen MoE checkpoint weight: {name}") diff --git a/afd_plugin/model_executor/models/qwen3_moe.py b/afd_plugin/model_executor/models/qwen3_moe.py index fd7cbedcd..5ba89a437 100644 --- a/afd_plugin/model_executor/models/qwen3_moe.py +++ b/afd_plugin/model_executor/models/qwen3_moe.py @@ -28,6 +28,9 @@ ) from afd_plugin.model_executor.ffn.remote import RemoteFFNProxy from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + BOTH_ROLES, + FFN_ROLE, filter_role_weights, layer_weight_path, ) @@ -35,10 +38,6 @@ if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts -_ATTENTION_ROLE = frozenset(("attention",)) -_FFN_ROLE = frozenset(("ffn",)) -_BOTH_ROLES = frozenset(("attention", "ffn")) - def _is_moe_layer(config: Qwen3MoeConfig, layer_idx: int) -> bool: """Return whether the native Qwen3 schedule selects MoE for one layer.""" @@ -67,17 +66,17 @@ def _checkpoint_weight_roles( """ layer_path = layer_weight_path(name) if layer_path is None: - return _BOTH_ROLES + return BOTH_ROLES _layer_idx, stage, remainder = layer_path if stage == "mlp": is_moe_gate = bool(remainder) and remainder[0] == "gate" if is_moe_gate and routes_on_attention: - return _BOTH_ROLES - return _FFN_ROLE + return BOTH_ROLES + return FFN_ROLE if stage in {"self_attn", "input_layernorm", "post_attention_layernorm"}: - return _ATTENTION_ROLE - return _BOTH_ROLES + return ATTENTION_ROLE + return BOTH_ROLES def _validate_supported_config( diff --git a/tests/unit/model_executor/models/npu/test_deepseek_v4.py b/tests/unit/model_executor/models/npu/test_deepseek_v4.py index e47f65130..790197270 100644 --- a/tests/unit/model_executor/models/npu/test_deepseek_v4.py +++ b/tests/unit/model_executor/models/npu/test_deepseek_v4.py @@ -138,8 +138,6 @@ def make_layers(count, factory, *, prefix): "cast": cast, "nn": nn, "native": native, - "_ATTENTION_ROLE": "attention", - "_FFN_ROLE": "ffn", "AFD_ASYNC_CONNECTOR": "CAMAsyncAFDConnector", "parse_afd_config": lambda config, **kwargs: config.afd, "_refresh_ascend_fused_moe": lambda: None, @@ -238,8 +236,8 @@ def test_mtp_target_hidden_states_exist_before_execution( @pytest.fixture(scope="module") def checkpoint_weight_roles() -> Callable[..., frozenset[str]]: - # Execute the pure helpers and their real role constants without importing - # torch or the native model class graph. + # Execute the pure helper against the shared role constants without + # importing torch or the native model class graph. path = ( Path(__file__).resolve().parents[5] / "afd_plugin/model_executor/models/npu/deepseek_v4.py" @@ -247,19 +245,21 @@ def checkpoint_weight_roles() -> Callable[..., frozenset[str]]: helpers: list[ast.stmt] = [ node for node in ast.parse(path.read_text()).body - if ( - isinstance(node, ast.FunctionDef) - and node.name == "_checkpoint_weight_roles" - ) - or ( - isinstance(node, ast.Assign) - and isinstance(node.targets[0], ast.Name) - and node.targets[0].id in ("_ATTENTION_ROLE", "_FFN_ROLE", "_BOTH_ROLES") - ) + if isinstance(node, ast.FunctionDef) and node.name == "_checkpoint_weight_roles" ] - from afd_plugin.model_executor.models._weight_utils import layer_weight_path + from afd_plugin.model_executor.models._weight_utils import ( + ATTENTION_ROLE, + BOTH_ROLES, + FFN_ROLE, + layer_weight_path, + ) - namespace: dict[str, object] = {"layer_weight_path": layer_weight_path} + namespace: dict[str, object] = { + "ATTENTION_ROLE": ATTENTION_ROLE, + "BOTH_ROLES": BOTH_ROLES, + "FFN_ROLE": FFN_ROLE, + "layer_weight_path": layer_weight_path, + } module = ast.Module(body=helpers, type_ignores=[]) exec(compile(module, str(path), "exec"), namespace) return cast(Callable[..., frozenset[str]], namespace["_checkpoint_weight_roles"]) diff --git a/tests/unit/model_executor/models/npu/test_deepseek_v4_dspark.py b/tests/unit/model_executor/models/npu/test_deepseek_v4_dspark.py index d1f4687c5..75e1bf562 100644 --- a/tests/unit/model_executor/models/npu/test_deepseek_v4_dspark.py +++ b/tests/unit/model_executor/models/npu/test_deepseek_v4_dspark.py @@ -12,9 +12,6 @@ from vllm.config.compilation import CompilationMode # noqa: E402 from afd_plugin.model_executor.models.npu import deepseek_v4_dspark # noqa: E402 -from afd_plugin.model_executor.models.npu.deepseek_v4_dspark import ( # noqa: E402 - _iter_role_weights, -) from tests.unit.fakes.vllm_config import npu_dspark_config # noqa: E402 @@ -61,10 +58,28 @@ def test_role_constructor_preserves_native_eager_call(eager_role_model): ("head.weight", {"attention"}), ], ) -def test_role_receives_only_owned_checkpoint_weights(name, roles): +def test_dspark_loader_streams_only_role_owned_weights( + name, roles, eager_role_model, monkeypatch +): weight = torch.ones(1) - for role in ("attention", "ffn"): - actual = list(_iter_role_weights(iter([(name, weight)]), role=role)) - assert [key for key, _ in actual] == ([name] if role in roles else []) - if actual: - assert actual[0][1] is weight + loaded: list[tuple[str, object]] = [] + + def load_weights(_model, weights): + loaded[:] = weights + return {key for key, _ in loaded} + + monkeypatch.setattr( + deepseek_v4_dspark.native.DSparkDeepseekV4ForCausalLM, + "load_weights", + load_weights, + ) + roles_to_check: tuple[object, ...] = ("attention", "ffn", None, "invalid", []) + for role in roles_to_check: + eager_role_model.afd_role = role + should_load = role in ("attention", "ffn") and role in roles + assert eager_role_model.load_weights(iter([(name, weight)])) == ( + {name} if should_load else set() + ) + assert [key for key, _ in loaded] == ([name] if should_load else []) + if loaded: + assert loaded[0][1] is weight diff --git a/tests/unit/package/test_package.py b/tests/unit/package/test_package.py index d025f73c7..97fe0589f 100644 --- a/tests/unit/package/test_package.py +++ b/tests/unit/package/test_package.py @@ -16,6 +16,17 @@ import afd_plugin +def test_models_metadata_compat_export(monkeypatch): + from afd_plugin.model_executor import models + + context = ModuleType("afd_plugin.forward_context") + metadata_reader = object() + context.__dict__["get_afd_metadata_from_forward_context"] = metadata_reader + monkeypatch.setitem(sys.modules, context.__name__, context) + + assert models.get_afd_metadata_from_forward_context is metadata_reader + + @pytest.fixture(params=["cuda", "npu"]) def registration(monkeypatch, request): native = {