From b3c24c2fb3485305c9d62669f91405e58fdbc42c Mon Sep 17 00:00:00 2001 From: Mohammad Date: Mon, 9 Nov 2026 14:39:33 +0200 Subject: [PATCH] refactor(moe): share remote MoE runners for GPU P2P On CUDA with P2pNcclAFDConnector, the Attention role now builds the MoE layers of the DeepSeek V2 family and Qwen3.5/3.6 through vLLM's own FusedMoEFactory, with an AFD runner (runner_cls) and parameter-free routed experts (routed_experts_cls), instead of the hand-made AFDAttentionFusedMoE stand-in. This is phase 1 of RFC #225 for GPU P2P. It ports #406 to vLLM v0.32.0 (5df63a72) and to the ffn/ layout of #562. ffn/remote.py: - AFDRemoteMoEMethod allocates no expert weights and runs no kernel; AFDRemoteRoutedExperts selects it. - AFDRemoteMoERunner sends the hidden states through the existing RemoteFFNProxy transport (send, DBO yield, receive). AFDExternalRoutingMoERunner also applies the gate it holds, as native MoERunner does, and sends the router logits. Both reject input_ids and shared_experts_input. - build_attention_moe_runner() accepts only CUDA with P2pNcclAFDConnector and takes the model's gate exactly when compute_gate_on_attention is set. It rejects Attention-local EPLB, redundant experts, routed-expert capture (aux_output_config in vLLM 0.32) and a non-finite routed_scaling_factor, and fixes the arguments that describe local experts. - AFDAttentionFusedMoE is deleted. Model adapters: - AFDDeepseekV2RemoteExpertsMoE takes vllm_config, keeps its gate (assigned before the experts, so the mlp.gate.* checkpoint names hold) and passes it to the factory with the native 0.32 arguments. Its forward override is gone; the shell runs vLLM's native forward. - The decoder-layer EPLB check moves into the factory. The DeepEP M2N shell keeps a parameter-free RemoteFFNProxy transport and rejects EPLB itself. - AFDQwen3_5RemoteExpertsMoE matches the native constructor signature and builds the gate-on-FFN runner. Outputs are unchanged: a new CUDA test checks that the Attention MoE sends the payload its gate placement requires and returns the output of the unsplit native MoE bit for bit (rtol=0, atol=0), for both gate placements and for 1 and 7 tokens. Co-authored-by: lirx-pd <616517220@qq.com> Co-authored-by: Claude Opus 5.5 Signed-off-by: Mohammad --- .../references/modules/model-integration.md | 12 +- afd_plugin/model_executor/ffn/remote.py | 252 +++++++++++-- .../model_executor/models/deepseek_v2.py | 125 ++++--- afd_plugin/model_executor/models/qwen3_5.py | 41 ++- docs/design/module/model_integration.md | 49 ++- docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md | 3 +- tests/gpu/test_attention_routing.py | 287 ++++++++++++++- tests/unit/fakes/vllm_config.py | 26 ++ tests/unit/model_executor/ffn/test_remote.py | 344 ++++++++++++++++++ .../models/test_deepseek_v2_construction.py | 323 ++++++++-------- .../models/test_deepseek_v2_proxy.py | 42 +-- .../models/test_deepseek_v2_weight_policy.py | 17 +- .../models/test_qwen3_5_contract.py | 74 ++++ 13 files changed, 1267 insertions(+), 328 deletions(-) create mode 100644 tests/unit/model_executor/ffn/test_remote.py diff --git a/.agents/skills/review-pr/references/modules/model-integration.md b/.agents/skills/review-pr/references/modules/model-integration.md index 450684a50..02f4e9682 100644 --- a/.agents/skills/review-pr/references/modules/model-integration.md +++ b/.agents/skills/review-pr/references/modules/model-integration.md @@ -6,7 +6,8 @@ head. Use for `afd_plugin/model_executor/`: the registration maps in `afd_plugin/__init__.py` (`_DEEPSEEK_MODEL_REGISTRATIONS`, `_QWEN_MODEL_REGISTRATIONS`, `_QWEN3_5_MODEL_REGISTRATIONS`), role-split model -wrappers (`RemoteFFNProxy`, `AFDAttentionFusedMoE`, `GateOnlyRemoteMoE`, the +wrappers (`RemoteFFNProxy`, the Attention-role MoE runners from +`build_attention_moe_runner()`, `GateOnlyRemoteMoE`, the `AFDDeepseekV2/V3/V4` and `AFDQwen3Moe`/`AFDQwen3_5` families, plus `AFDGlmMoeDsaForCausalLM` and `AFDDSparkDeepseekV4ForCausalLM`), checkpoint weight-role filtering, the NPU-specific forward/gate modules under @@ -27,6 +28,15 @@ registration mechanics. - Keep remote proxies faithful: `RemoteFFNProxy` and friends preserve upstream signatures, hidden-state layouts, and auxiliary outputs (residuals, gate values) across the role boundary. +- On CUDA with `P2pNcclAFDConnector`, build the experts of an Attention-role + native MoE shell through `build_attention_moe_runner()` (vLLM + `FusedMoEFactory` with an AFD `runner_cls` and `routed_experts_cls`): the + runner holds the model's gate only with `compute_gate_on_attention`, + allocates no expert weights, and runs no local top-k or expert kernel. A + change that does, or that passes an argument the factory fixes, is a + finding. DeepEP M2N (`AFDDeepseekV2RoutedExpertsMoE`) and the NPU gate + paths (`GateOnlyRemoteMoE`) route on Attention by design and keep their own + `RemoteFFNProxy` transports. - Access upstream model config attributes directly per AGENTS.md; a `getattr`/`hasattr` guard around vLLM model fields hides the exact incompatibility the pinned-version contract exists to surface. diff --git a/afd_plugin/model_executor/ffn/remote.py b/afd_plugin/model_executor/ffn/remote.py index 4c0575193..a4f1786e8 100644 --- a/afd_plugin/model_executor/ffn/remote.py +++ b/afd_plugin/model_executor/ffn/remote.py @@ -1,20 +1,42 @@ # 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.""" +"""Connector-independent remote dense/MoE proxies and Attention-side MoE +runners; no concrete model imports.""" from __future__ import annotations +import math +from typing import Any + import torch from torch import nn +from vllm.config import VllmConfig 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 import fused_moe +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + FusedMoEQuantConfig, +) +from vllm.model_executor.layers.fused_moe.fused_moe_method_base import ( + FusedMoEMethodBase, +) +from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts +from vllm.model_executor.layers.fused_moe.router.fused_moe_router import ( + FusedMoERouter, +) from vllm.model_executor.layers.fused_moe.router.fused_topk_router import fused_topk +from vllm.model_executor.layers.fused_moe.runner.moe_runner import MoERunner from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.platforms import current_platform +from afd_plugin.config import parse_afd_config from afd_plugin.connectors import ( AFDForwardContextMetadata, AFDTransferContext, @@ -24,6 +46,10 @@ from afd_plugin.v1.worker.dbo import maybe_apply_dbo_yield UNROUTED_EXPERT_ID = -1 +# The only placement validated for native-factory Attention MoE runners. +# DeepEP M2N routes on Attention through its own model shell instead. +ATTENTION_MOE_RUNNER_DEVICE = "cuda" +ATTENTION_MOE_RUNNER_CONNECTOR = "P2pNcclAFDConnector" def _unroute_tp_padding(topk_ids: torch.Tensor, num_tokens: int) -> torch.Tensor: @@ -144,41 +170,221 @@ def _afd_metadata() -> AFDForwardContextMetadata: return afd_metadata -class AFDAttentionFusedMoE(RemoteFFNProxy): - """Parameter-free native-MoE experts proxy for the Attention runtime.""" +class AFDRemoteMoEMethod(FusedMoEMethodBase): + """Native MoE method lifecycle without local expert weights or kernels. + + ``apply`` and ``apply_monolithic`` keep the base ``NotImplementedError``, + so computing experts on the Attention role fails loudly. + """ + + def create_weights( + self, + layer: RoutedExperts, + num_experts: int, + hidden_size: int, + intermediate_size_per_partition: int, + params_dtype: torch.dtype, + **extra_weight_attrs, + ): + """Allocate nothing: the FFN role owns the expert weights.""" + + def get_fused_moe_quant_config( + self, layer: RoutedExperts + ) -> FusedMoEQuantConfig | None: + return None + + def maybe_roundup_sizes( + self, + hidden_size: int, + intermediate_size_per_partition: int, + act_dtype: torch.dtype, + moe_parallel_config: FusedMoEParallelConfig, + ) -> tuple[int, int]: + # No local kernel or all2all buffer needs aligned sizes, so the + # runner's FusedMoEConfig keeps the model's dimensions. + return hidden_size, intermediate_size_per_partition + + +class AFDRemoteRoutedExperts(RoutedExperts): + """Native routed-experts descriptor that allocates no expert parameters.""" + + def _get_quant_method( + self, + prefix: str, + quant_config: QuantizationConfig | None, + moe_config: FusedMoEConfig, + ) -> FusedMoEMethodBase: + return AFDRemoteMoEMethod(moe_config) + + +class AFDRemoteMoERunner(MoERunner): + """Native MoE runner whose gate and experts run on the FFN role. + + The runner keeps the native layer registration but runs no local routing + or expert kernel: it sends this layer's hidden states through the + ``RemoteFFNProxy`` transport (send, DBO yield, receive) and returns the + FFN output. + """ def __init__( self, - *, - layer_idx: int, - external_routing: bool, + layer_name: str, + moe_config: FusedMoEConfig, + router: FusedMoERouter, + routed_experts: RoutedExperts, + enable_dbo: bool = False, + gate: nn.Module | None = None, + shared_experts: nn.Module | None = None, + shared_expert_gate: nn.Module | None = None, + routed_input_transform: nn.Module | None = None, + routed_output_transform: nn.Module | None = None, + routed_scaling_factor: float = 1.0, ) -> 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] + super().__init__( + layer_name=layer_name, + moe_config=moe_config, + router=router, + routed_experts=routed_experts, + enable_dbo=enable_dbo, + gate=gate, + shared_experts=shared_experts, + shared_expert_gate=shared_expert_gate, + routed_input_transform=routed_input_transform, + routed_output_transform=routed_output_transform, + routed_scaling_factor=routed_scaling_factor, + ) + # The native layer name carries the decoder layer index. + self.ffn_transport = RemoteFFNProxy(layer_idx=self.layer_id) + + def forward( self, hidden_states: torch.Tensor, router_logits: torch.Tensor, input_ids: torch.Tensor | None = None, + shared_experts_input: 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) + if shared_experts_input is not None: + raise NotImplementedError( + "Attention remote MoE has no local shared experts", + ) + return self.ffn_transport._send_and_receive( + hidden_states, + **self._routing_payload(hidden_states), + ) + + def _routing_payload( + self, + hidden_states: torch.Tensor, + ) -> dict[str, torch.Tensor]: + """Routing tensors sent with the hidden states; none when FFN gates.""" + return {} + + +class AFDExternalRoutingMoERunner(AFDRemoteMoERunner): + """Native MoE runner that applies the gate held on the Attention role. + + As native ``MoERunner._forward_impl`` does when it holds a gate, the + runner ignores its ``router_logits`` argument and computes the logits + from the hidden states. It sends them with the hidden states, and the FFN + role runs its external-router experts. + """ + + def _routing_payload( + self, + hidden_states: torch.Tensor, + ) -> dict[str, torch.Tensor]: + assert self.gate is not None + router_logits, _ = self.gate(hidden_states) + return {"router_logits": router_logits} + - def update_expert_map(self) -> None: - """Satisfy the native MoERunner EPLB interface without local experts.""" +def build_attention_moe_runner( + vllm_config: VllmConfig, + *, + gate: nn.Module | None, + **moe_kwargs: Any, +) -> MoERunner: + """Build an Attention-role MoE runner through native ``FusedMoEFactory``. + + ``moe_kwargs`` are the native model's ``FusedMoEFactory`` arguments. Pass + the model's gate exactly when ``compute_gate_on_attention`` is set; the + runner then applies it and sends the router logits. This function fixes + the arguments that describe local experts: the runner registers like a + native MoE layer but allocates no expert weights and runs no expert + kernel. + """ + afd_config = parse_afd_config(vllm_config, validate=False) + placement = (current_platform.device_type, afd_config.connector) + if placement != (ATTENTION_MOE_RUNNER_DEVICE, ATTENTION_MOE_RUNNER_CONNECTOR): + raise ValueError( + f"unsupported Attention remote MoE configuration: {placement!r}", + ) + if (gate is not None) != afd_config.compute_gate_on_attention: + raise ValueError( + "Attention remote MoE takes the model's gate if and only if " + "compute_gate_on_attention is set", + ) + parallel_config = vllm_config.parallel_config + if parallel_config.enable_eplb: + raise RuntimeError( + "Remote MoE does not support Attention-local EPLB (enable_eplb)", + ) + if parallel_config.eplb_config.num_redundant_experts != 0: + raise RuntimeError( + "Remote MoE does not support Attention-local redundant experts", + ) + if vllm_config.aux_output_config.enable_return_routed_experts: + raise RuntimeError( + "Remote MoE does not support Attention-local routed_experts capture", + ) + if "routed_scaling_factor" in moe_kwargs and not math.isfinite( + moe_kwargs["routed_scaling_factor"], + ): + raise ValueError("routed_scaling_factor must be finite") + # A parameter-free container: no local quantized weights, shared experts, + # transforms or EPLB, and unit TP/DP/PCP sizes, so no expert sharding. + # The FFN role applies routed_scaling_factor; the runner only records it. + factory_kwargs: dict[str, Any] = { + "gate": gate, + "quant_config": None, + "shared_experts": None, + "shared_expert_gate": None, + "routed_input_transform": None, + "routed_output_transform": None, + "n_shared_experts": None, + "fuse_shared_experts": False, + "apply_routed_scale_to_output": True, + "enable_eplb": False, + "num_redundant_experts": 0, + "is_sequence_parallel": False, + "tp_size": 1, + "dp_size": 1, + "pcp_size": 1, + "runner_cls": ( + AFDExternalRoutingMoERunner if gate is not None else AFDRemoteMoERunner + ), + "runner_args": None, + "routed_experts_cls": AFDRemoteRoutedExperts, + "routed_experts_args": None, + } + reserved = moe_kwargs.keys() & factory_kwargs.keys() + if reserved: + raise ValueError( + "Attention remote MoE factory reserves these arguments: " + + ", ".join(sorted(reserved)), + ) + return fused_moe.FusedMoEFactory(**moe_kwargs, **factory_kwargs) -__all__ = ["RemoteFFNProxy", "AFDAttentionFusedMoE"] +__all__ = [ + "AFDExternalRoutingMoERunner", + "AFDRemoteMoEMethod", + "AFDRemoteMoERunner", + "AFDRemoteRoutedExperts", + "RemoteFFNProxy", + "build_attention_moe_runner", +] diff --git a/afd_plugin/model_executor/models/deepseek_v2.py b/afd_plugin/model_executor/models/deepseek_v2.py index 078bd5a67..edfba302b 100644 --- a/afd_plugin/model_executor/models/deepseek_v2.py +++ b/afd_plugin/model_executor/models/deepseek_v2.py @@ -16,7 +16,7 @@ import torch import torch.nn as nn from transformers import DeepseekV2Config, DeepseekV3Config, GlmMoeDsaConfig -from vllm.config import ParallelConfig, VllmConfig +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers import fused_moe from vllm.model_executor.layers.fused_moe.router.router_factory import ( @@ -31,9 +31,9 @@ ) from afd_plugin.connectors import AFDExpertRoutingSpec from afd_plugin.model_executor.ffn.remote import ( - AFDAttentionFusedMoE, RemoteFFNProxy, _unroute_tp_padding, + build_attention_moe_runner, ) from afd_plugin.model_executor.models._weight_utils import ( ATTENTION_ROLE, @@ -181,34 +181,38 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: class AFDDeepseekV2RemoteExpertsMoE(native.DeepseekV2MoE): - """Native DeepSeek MoE forward with parameter-free remote experts.""" + """Native DeepSeek MoE forward with remote experts and no expert weights.""" # Patch reason: native DeepseekV2MoE constructs local routed/shared experts # through FusedMoEFactory. - # Patch functionality: preserve the native MoE forward contract while - # constructing only the gate owned by Attention and a parameter-free proxy. - # Signature: AFD-owned; adds layer_idx and compute_gate_on_attention and omits - # quant_config because no local expert kernel is constructed. + # Patch functionality: keep the native MoE forward. Construct only the + # gate owned by Attention, and build the experts through FusedMoEFactory + # with an AFD runner that runs them on the FFN role and allocates no + # expert weights. As native, the runner holds the gate and applies it. + # Signature: AFD-owned; takes vllm_config for the factory, adds layer_idx + # and compute_gate_on_attention, and omits quant_config because no local + # expert kernel is constructed. # Upstream: vLLM v0.32.0, vllm/model_executor/models/deepseek_v2.py # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be - # In v0.32.0 the native forward always delegates routing to the experts - # runner, so forward() below computes the Attention-side gate first when - # this shell owns one. def __init__( self, *, config: _DeepseekAdapterConfig, - parallel_config: ParallelConfig, + vllm_config: VllmConfig, layer_idx: int, prefix: str, compute_gate_on_attention: bool, ) -> None: # ### PATCH START: construct a remote-experts native MoE shell. nn.Module.__init__(self) + parallel_config = vllm_config.parallel_config self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe router_dtype = native._get_moe_router_dtype(config) self.router_dtype = router_dtype + # The experts runner holds this same gate, as native does. Assign it + # before the experts so its parameters keep their checkpoint names + # (mlp.gate.*). if compute_gate_on_attention: self.gate = native.GateLinear( config.hidden_size, @@ -235,47 +239,43 @@ def __init__( # No local shared-expert MLP is constructed, so fused shared experts # can never be enabled for this shell. self.is_fused_shared_expert_enabled = False - self.experts = AFDAttentionFusedMoE( + self.experts = self._build_remote_experts( + config=config, + vllm_config=vllm_config, layer_idx=layer_idx, - external_routing=compute_gate_on_attention, + prefix=prefix, ) # ### PATCH END: construct a remote-experts native MoE shell. - # Patch reason: v0.32.0's native forward always delegates routing to the - # experts runner (router_logits=hidden_states), so a shell that owns a - # gate would never compute it and Attention would ship hidden states as - # router logits. - # Patch functionality: compute the Attention-side gate before transport; - # gate-less shells keep the native delegation unchanged. - # Signature: matches upstream; no added parameters. - # Upstream: vLLM v0.32.0, vllm/model_executor/models/deepseek_v2.py - # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be - def forward( + def _build_remote_experts( self, - hidden_states: torch.Tensor, - already_sequence_parallel: bool = False, - ) -> torch.Tensor: - if self.gate is None: - return super().forward(hidden_states, already_sequence_parallel) - - # ### PATCH START: compute the Attention-side gate before transport. - num_tokens, hidden_dim = hidden_states.shape - hidden_states = hidden_states.view(-1, hidden_dim) - if self.is_sequence_parallel and not already_sequence_parallel: - hidden_states = native.sequence_parallel_chunk(hidden_states) - router_logits, _ = self.gate(hidden_states) - final_hidden_states = self.experts( - hidden_states=hidden_states, - router_logits=router_logits, + *, + config: _DeepseekAdapterConfig, + vllm_config: VllmConfig, + layer_idx: int, + prefix: str, + ) -> nn.Module: + """Native DeepseekV2MoE factory arguments, minus local-expert ones.""" + return build_attention_moe_runner( + vllm_config, + gate=self.gate, + num_experts=config.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + use_grouped_topk=True, + num_expert_group=getattr(config, "n_group", 1), + topk_group=getattr(config, "topk_group", 1), + prefix=f"{prefix}.experts", + scoring_func=getattr(config, "scoring_func", "softmax"), + routed_scaling_factor=getattr(config, "routed_scaling_factor", 1.0), + e_score_correction_bias=( + None if self.gate is None else self.gate.e_score_correction_bias + ), + shared_expert_prefix=f"{prefix}.shared_experts", + router_logits_dtype=self.router_dtype, ) - if self.is_sequence_parallel and not already_sequence_parallel: - final_hidden_states = native.tensor_model_parallel_all_gather( - final_hidden_states, - 0, - ) - final_hidden_states = final_hidden_states[:num_tokens] - return final_hidden_states.view(num_tokens, hidden_dim) - # ### PATCH END: compute the Attention-side gate before transport. class AFDDeepseekV2RoutedExpertsMoE(AFDDeepseekV2RemoteExpertsMoE): @@ -298,7 +298,9 @@ class AFDDeepseekV2RoutedExpertsMoE(AFDDeepseekV2RemoteExpertsMoE): # Patch reason: native DeepseekV2MoE runs the router and the shared experts # inside its local experts runner, which Attention does not build. # Patch functionality: build the native router and shared MLP next to the - # remote-experts shell; route, dispatch, run the shared MLP, combine. + # remote-experts shell; route, dispatch, run the shared MLP, combine. The + # connector carries the routes itself, so the experts are a parameter-free + # RemoteFFNProxy instead of the native factory runner. # Signature: AFD-owned; takes vllm_config and layer_idx. # Upstream: vLLM v0.32.0, vllm/model_executor/models/deepseek_v2.py # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be @@ -310,9 +312,13 @@ def __init__( layer_idx: int, prefix: str, ) -> None: + if vllm_config.parallel_config.enable_eplb: + raise RuntimeError( + "CUDA remote experts do not support EPLB on the Attention role", + ) super().__init__( config=config, - parallel_config=vllm_config.parallel_config, + vllm_config=vllm_config, layer_idx=layer_idx, prefix=prefix, compute_gate_on_attention=True, @@ -349,6 +355,17 @@ def __init__( else None ) + def _build_remote_experts( + self, + *, + config: _DeepseekAdapterConfig, + vllm_config: VllmConfig, + layer_idx: int, + prefix: str, + ) -> RemoteFFNProxy: + """Parameter-free transport; forward sends the routes it selects.""" + return RemoteFFNProxy(layer_idx=layer_idx) + def forward( self, hidden_states: torch.Tensor, @@ -456,15 +473,9 @@ def __init__( afd_config.compute_gate_on_attention, ) device_type = native.current_platform.device_type + # Attention-side EPLB is rejected by the remote experts: the native + # factory runner (P2P) and the DeepEP M2N shell. self.uses_remote_experts = device_type == "cuda" and self.is_moe_layer - if ( - afd_role == "attention" - and self.uses_remote_experts - and parallel_config.enable_eplb - ): - raise RuntimeError( - "CUDA remote experts do not support EPLB on the Attention role", - ) if self.compute_gate_on_attention and device_type not in ("cuda", "npu"): raise RuntimeError( "DeepSeekV2 compute_gate_on_attention requires CUDA or NPU", @@ -517,7 +528,7 @@ def __init__( elif self.uses_remote_experts: self.mlp = AFDDeepseekV2RemoteExpertsMoE( config=config, - parallel_config=parallel_config, + vllm_config=vllm_config, layer_idx=layer_idx, prefix=f"{prefix}.mlp", compute_gate_on_attention=self.compute_gate_on_attention, diff --git a/afd_plugin/model_executor/models/qwen3_5.py b/afd_plugin/model_executor/models/qwen3_5.py index 1e31262c3..31e1fbc26 100644 --- a/afd_plugin/model_executor/models/qwen3_5.py +++ b/afd_plugin/model_executor/models/qwen3_5.py @@ -12,7 +12,7 @@ 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.ffn.remote import AFDAttentionFusedMoE +from afd_plugin.model_executor.ffn.remote import build_attention_moe_runner from afd_plugin.model_executor.models._weight_utils import ( ATTENTION_ROLE, FFN_ROLE, @@ -65,21 +65,21 @@ def _checkpoint_weight_roles( class AFDQwen3_5RemoteExpertsMoE( # noqa: N801 native.Qwen3NextSparseMoeBlock, ): - """Native Qwen MoE shell with a parameter-free remote experts runner.""" - - # Patch reason: native Qwen3NextSparseMoeBlock allocates routed and shared - # experts on every rank. - # Patch functionality: preserve its native forward while keeping a - # parameter-free experts proxy with a local FFN router. - # Signature: AFD-owned; layer_idx is required for correlation metadata. + """Native Qwen MoE shell whose gate and experts run on the FFN role.""" + + # Patch reason: native Qwen3NextSparseMoeBlock allocates the gate and the + # routed and shared experts on every rank. + # Patch functionality: keep its native forward and build the experts + # through FusedMoEFactory with an AFD runner that sends hidden states to + # the FFN role, which owns the gate and all expert weights. + # Signature: matches upstream; no added parameters. # Upstream: vLLM v0.32.0, vllm/model_executor/models/qwen3_next.py # Commit: 5df63a7203cdb9328d824a27b682e77c7b9928be def __init__( self, - *, vllm_config: VllmConfig, - layer_idx: int, - prefix: str, + prefix: str = "", + reduce_results: bool = True, ) -> None: # ### PATCH START: construct a remote-experts native MoE shell. nn.Module.__init__(self) @@ -87,8 +87,6 @@ def __init__( parallel_config = vllm_config.parallel_config if parallel_config.use_sequence_parallel_moe: raise RuntimeError("AFD Qwen3.5/3.6 does not support SP MoE") - if parallel_config.enable_eplb: - raise RuntimeError("AFD Qwen3.5/3.6 does not support EPLB") self.tp_size = next_native.get_tensor_model_parallel_world_size() self.ep_group = next_native.get_ep_group().device_group @@ -106,9 +104,19 @@ def __init__( self.gate = None self.shared_expert = None self.shared_expert_gate = None - self.experts = AFDAttentionFusedMoE( - layer_idx=layer_idx, - external_routing=False, + # Native Qwen3NextSparseMoeBlock factory arguments, minus local-expert + # ones; EPLB and the other local-expert features are rejected there. + self.experts = build_attention_moe_runner( + vllm_config, + gate=None, + num_experts=self.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=getattr(config, "norm_topk_prob", True), + prefix=f"{prefix}.experts", + reduce_results=reduce_results, + shared_expert_prefix=f"{prefix}.shared_expert", ) # ### PATCH END: construct a remote-experts native MoE shell. @@ -195,7 +203,6 @@ def __init__( raise ValueError(f"Invalid layer_type {self.layer_type}") self.mlp = AFDQwen3_5RemoteExpertsMoE( vllm_config=vllm_config, - layer_idx=self.layer_idx, prefix=f"{prefix}.mlp", ) self.input_layernorm = native.Qwen3_5RMSNorm( diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index d38814c42..92917bd5d 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -27,6 +27,7 @@ upstream_refs: - "NPU target: vLLM 0.32.0 5df63a7203cdb9328d824a27b682e77c7b9928be" - "vLLM vllm.model_executor.models.deepseek_v2" - "vLLM vllm.model_executor.models.qwen3_moe" + - "vLLM vllm.model_executor.layers.fused_moe (FusedMoEFactory, MoERunner, RoutedExperts)" - "vLLM vllm.forward_context.ForwardContext" - "vLLM vllm.model_executor.model_loader" verified_platform_refs: @@ -40,6 +41,7 @@ related_issues: - "#88" - "#105" - "#129" + - "#225" last_reviewed: 2026-11-09 --- @@ -61,12 +63,14 @@ make a backend-specific worker class the shared model API. `model_executor/models/` owns model structure and checkpoint policy. `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 +Ascend models, and the Attention-role MoE runners built through vLLM's +`FusedMoEFactory`, which its factory currently enables for CUDA P2P only. +`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 | +| `remote.py` | Connector-independent dense/MoE proxies, routed TP padding, and the Attention-role MoE runners built through vLLM's `FusedMoEFactory` | Connector contracts, the plugin AFD config, public vLLM helpers, the native fused-MoE factory, runner and routed-experts classes, and package forward context; no CAM, concrete model or compute imports | | `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 | @@ -122,7 +126,7 @@ SFA/MLA/DSA-specific handling stays explicit inside that worker-side module. | 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 | [`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) | +| CUDA remote-experts boundary | [`ffn/remote.py`](../../../afd_plugin/model_executor/ffn/remote.py), [`gpu/p2p.py`](../../../afd_plugin/connectors/gpu/p2p.py) | [`test_remote.py`](../../../tests/unit/model_executor/ffn/test_remote.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), bit-exact native-kernel parity in [`test_attention_routing.py`](../../../tests/gpu/test_attention_routing.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/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) | @@ -169,8 +173,8 @@ needed by the split execution. | Layer/component | Attention role | FFN role | | --- | --- | --- | | Attention module and KV-facing computation | Constructed and executed. | Not constructed. | -| MoE with `compute_gate_on_attention=false` | CUDA constructs the native MoE shell with a parameter-free internal-router experts proxy; NPU sends after post-Attention normalization. | Native gate and experts are constructed and executed from connector input. | -| MoE with `compute_gate_on_attention=true`, non-routed connectors | CUDA keeps the native gate and uses an external-router experts proxy; NPU uses its Attention-side gate helper. | Expert MLP is constructed and consumes transferred router logits or routed payloads without rerunning the gate. | +| MoE with `compute_gate_on_attention=false` | CUDA constructs the native MoE shell without a gate; its experts are a `FusedMoEFactory` runner without expert weights that sends the hidden states. NPU sends after post-Attention normalization. | Native gate and experts are constructed and executed from connector input. | +| MoE with `compute_gate_on_attention=true`, non-routed connectors | CUDA keeps the native gate on the shell and passes it to the `FusedMoEFactory` runner, which applies it and sends the router logits with the hidden states; NPU uses its Attention-side gate helper. | Expert MLP is constructed and consumes transferred router logits or routed payloads without rerunning the gate. | | MoE with `compute_gate_on_attention=true`, DeepEP M2N (CUDA) | `AFDDeepseekV2RoutedExpertsMoE` constructs the native gate, router, and shared-expert MLP. It selects top-k, dispatches routed rows, runs the shared MLP while dispatch and the FFN experts are in flight, and adds it to the combined output. With TP > 1, each TP rank routes its `sequence_parallel_chunk` slice and runs a replicated (`is_sequence_parallel`, no all-reduce) shared MLP on it; the TP group then all-gathers the output. | Native `DeepseekV2MoE` without shared experts; runs only the routed experts on the dispatched top-k IDs and weights. | | Dense MLP, normal mode | Not constructed; output is sent after post-Attention normalization. | Constructed and executed from connector input. | | Dense MLP with `compute_gate_on_attention=true` | Constructed and executed locally because there is no routed MoE handoff. | Not constructed and a dense-layer FFN compute request is rejected. | @@ -178,15 +182,28 @@ needed by the split execution. CUDA MoE splits at the remote-experts boundary. With gate-on-FFN, and with gate-on-Attention on non-routed connectors, Attention keeps native -`DeepseekV2MoE.forward`. With gate-on-FFN, the proxy asks FFN to run its native -internal-router MoE. With gate-on-Attention, Attention runs the native gate and -FFN executes its external-router experts path. Routed-experts connectors -(DeepEP M2N) use a third mode with an AFD-owned forward: Attention runs the -gate, top-k selection, and shared experts; FFN runs only the routed experts, -and combine returns each token's weighted top-k sum. CUDA Attention-side remote -experts currently reject EPLB. The NPU gate helper supports unquantized and -Ascend W8A8 MoE expert computation; unsupported devices or quantization fail -explicitly. +`DeepseekV2MoE.forward`. Its experts come from vLLM's own `FusedMoEFactory` +through `build_attention_moe_runner()` in `ffn/remote.py`, with an AFD runner +(`runner_cls`) and parameter-free routed experts (`routed_experts_cls`, +`AFDRemoteRoutedExperts`, whose `AFDRemoteMoEMethod` allocates no weights). +The runner registers like a native MoE layer but runs no local routing or +expert kernel; it sends through the `RemoteFFNProxy` transport (send, DBO +yield, receive). With gate-on-FFN, `AFDRemoteMoERunner` sends the hidden states +and FFN runs its native internal-router MoE. With gate-on-Attention, the +runner holds the shell's native gate, as native `MoERunner` does; +`AFDExternalRoutingMoERunner` applies it and sends the router logits, and FFN +executes its external-router experts path. The factory accepts only CUDA with +`P2pNcclAFDConnector` and fixes the arguments that describe local experts +(quantization, shared experts, transforms, EPLB and TP/DP/PCP sharding). It +rejects Attention-local EPLB, redundant experts, routed-expert capture, a +non-finite `routed_scaling_factor`, and caller-supplied values for the +arguments it fixes. Routed-experts connectors (DeepEP M2N) use a third mode +with an AFD-owned forward: Attention runs the gate, top-k selection, and shared +experts; FFN runs only the routed experts, and combine returns each token's +weighted top-k sum. That shell sends through a parameter-free `RemoteFFNProxy` +instead of the factory runner and also rejects EPLB. The NPU gate helper +supports unquantized and Ascend W8A8 MoE expert computation; unsupported +devices or quantization fail explicitly. The full AFD model remains decorated with vLLM's compile support. Backend-only helpers are imported inside the NPU path so CUDA model import does not require @@ -278,7 +295,9 @@ parallelism, speculative decoding, and LoRA during construction. `AFDQwen3_5MoeForConditionalGeneration` is a distinct wrapper around vLLM's native hybrid Qwen3.5/3.6 architecture. Attention owns embeddings, norms, and -linear/full-attention state, then sends hidden states only. FFN owns the native +linear/full-attention state, then sends hidden states only: its MoE keeps the +native `Qwen3NextSparseMoeBlock.forward`, with experts built by the same +`FusedMoEFactory` runner as DeepSeek's gate-on-FFN path. FFN owns the native gate, routed experts, shared expert, and shared-expert gate. This path is CUDA-only, text-only, and requires `--language-model-only`, diff --git a/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md b/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md index d87287845..01baacfd4 100644 --- a/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md +++ b/docs/gpu/NCCL_P2P_CONNECTOR_USER_GUIDE.md @@ -7,7 +7,8 @@ P2pNcclAFDConnector (implemented with vLLM's PyNcclCommunicator) is a GPU-backed Use this connector for CUDA deployments that disaggregate Attention and FFN workers and exchange hidden states synchronously through NCCL point-to-point communication. On the GPU target vLLM 0.32, DeepSeek MoE layers keep the native MoE forward -contract and replace only the local experts with an AFD remote-experts proxy. +contract and build their experts through vLLM's `FusedMoEFactory` with an AFD +runner that sends them to the FFN role and allocates no expert weights. The MoE gate may run on Attention or FFN. It supports both prefill and decode which all support eager mode. CUDA graph support is currently limited to `FULL_DECODE_ONLY`, which is mainly used in decode instance. The checked-in DeepSeek V2 Lite recipes cover colocated and prefill/decode-disaggregated deployments. diff --git a/tests/gpu/test_attention_routing.py b/tests/gpu/test_attention_routing.py index 5986a8774..3814b054f 100644 --- a/tests/gpu/test_attention_routing.py +++ b/tests/gpu/test_attention_routing.py @@ -1,14 +1,53 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project -"""Compare AFD Attention-side CUDA routing with native and mathematical oracles.""" +"""Compare AFD Attention-side CUDA routing and MoE output with native and +mathematical oracles.""" + +import subprocess +import sys +from collections.abc import Callable +from functools import partial +from pathlib import Path +from typing import TYPE_CHECKING import pytest -torch = pytest.importorskip("torch") +if TYPE_CHECKING: + import torch +else: + torch = pytest.importorskip("torch") pytest.importorskip("vllm") +from torch import nn # noqa: E402 +from transformers import DeepseekV2Config # noqa: E402 +from vllm.config import VllmConfig, set_current_vllm_config # noqa: E402 +from vllm.distributed import ( # noqa: E402 + destroy_distributed_environment, + destroy_model_parallel, +) +from vllm.engine.arg_utils import EngineArgs # noqa: E402 +from vllm.forward_context import ( # noqa: E402 + ForwardContext, + override_forward_context, + set_forward_context, +) from vllm.model_executor.layers.fused_moe.router import fused_topk_router # noqa: E402 +from vllm.model_executor.model_loader.utils import ( # noqa: E402 + process_weights_after_loading, +) +from vllm.model_executor.models import deepseek_v2 as native # noqa: E402 +from vllm.utils.network_utils import get_open_port # noqa: E402 +from vllm.utils.torch_utils import set_default_torch_dtype # noqa: E402 +from vllm.v1.worker.gpu_worker import ( # noqa: E402 + init_worker_distributed_environment, +) +from vllm.v1.worker.workspace import ( # noqa: E402 + init_workspace_manager, + reset_workspace_manager, +) +from afd_plugin.connectors import AFDForwardContextMetadata # noqa: E402 from afd_plugin.model_executor.models import deepseek_v2 as adapter # noqa: E402 +from tests.conftest import REPO_ROOT # noqa: E402 from tests.unit.fakes.connector import install_proxy_context # noqa: E402 from tests.unit.fakes.model_layers import GateLinear # noqa: E402 @@ -18,6 +57,18 @@ ROUTING_NUM_EXPERTS = 8 ROUTING_TOP_K = 2 +# The parity child process imports vLLM, initializes its distributed +# environment and compiles the Triton MoE kernels on first use. +PARITY_TIMEOUT_S = 600 +PARITY_TOKEN_COUNTS = (1, 7) +PARITY_HIDDEN_SIZE = 128 +PARITY_MODEL_LEN = 128 +# The single rank of the child process's vLLM world. +PARITY_RANK = 0 +# Deterministic values ``(i % period - period // 2) / scale``, exact in BF16. +WEIGHT_PATTERN = (31, 128) +INPUT_PATTERN = (43, 32) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="fused top-k needs CUDA") @pytest.mark.parametrize("renormalize", [True, False]) @@ -68,3 +119,235 @@ def test_routing_proxy_matches_native_router(monkeypatch, renormalize): rtol=1e-3, atol=1e-3, ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="native MoE needs CUDA") +@pytest.mark.parametrize( + "routing", + [ + pytest.param( + { + "topk_method": "greedy", + "scoring_func": "softmax", + "n_group": 1, + "topk_group": 1, + "norm_topk_prob": False, + }, + id="greedy-softmax", + ), + pytest.param( + { + "topk_method": "noaux_tc", + "scoring_func": "sigmoid", + "n_group": 2, + "topk_group": 1, + "norm_topk_prob": True, + }, + id="noaux-tc-sigmoid", + ), + ], +) +def test_attention_moe_runner_matches_native_moe_bit_for_bit(tmp_path, routing): + """On CUDA, the Attention MoE built by vLLM's FusedMoEFactory sends the + payload the gate placement requires and returns, bit for bit, the output + of the unsplit native MoE, for both gate placements and for 1 and 7 + tokens. + + A loopback connector runs the FFN role's native MoE in-process on what + Attention sends: the complete MoE on the hidden states when FFN owns the + gate, or its external-router experts on the sent router logits when + Attention does. + + Why: splitting the MoE between the roles must not change any output. + Both sides run the same kernels on the same inputs, so the comparison + uses rtol=0 and atol=0. + """ + program = ( + "import runpy\n" + f"module = runpy.run_path({str(Path(__file__).resolve())!r})\n" + f"module['_run_attention_moe_parity']({str(tmp_path)!r}, {routing!r})\n" + ) + # Keep vLLM's process groups and MoE workspace out of this test process. + subprocess.run( + [sys.executable, "-c", program], + check=True, + cwd=REPO_ROOT, + timeout=PARITY_TIMEOUT_S, + ) + + +class _LoopbackFFN: + """Runs the FFN role's native MoE in-process on what Attention sends.""" + + def __init__( + self, + vllm_config: VllmConfig, + internal_router_moe: nn.Module, + external_router_moe: nn.Module, + ) -> None: + self.vllm_config = vllm_config + self.internal_router_moe = internal_router_moe + self.external_router_moe = external_router_moe + self.expected_payload: set[str] = set() + self.output: torch.Tensor | None = None + + def send_attn_output(self, hidden_states, context, **payload) -> None: + assert set(payload) == self.expected_payload + num_tokens = hidden_states.shape[0] + with set_forward_context(None, self.vllm_config, num_tokens=num_tokens): + if payload: + # FFN compute_experts_output: native runner, external routing. + self.output = self.external_router_moe.experts( + hidden_states=hidden_states, + router_logits=payload["router_logits"], + ).clone() + else: + # FFN compute_ffn_output: the complete native MoE. + self.output = self.internal_router_moe(hidden_states).clone() + + def recv_ffn_output(self, *, ref_tensor, ubatch_idx) -> torch.Tensor: + assert self.output is not None + return self.output + + +def _pattern(count: int, pattern: tuple[int, int]) -> torch.Tensor: + period, scale = pattern + values = torch.arange(count, device="cuda").remainder(period) - period // 2 + return values / scale + + +def _built_moe( + module_factory: Callable[[], nn.Module], + vllm_config: VllmConfig, +) -> nn.Module: + """Construct a MoE as the model loader does, with deterministic weights.""" + with set_default_torch_dtype(torch.bfloat16), torch.device("cuda"): + module = module_factory() + with torch.no_grad(): + for parameter in module.parameters(): + values = _pattern(parameter.numel(), WEIGHT_PATTERN) + parameter.copy_(values.view_as(parameter)) + process_weights_after_loading( + module, vllm_config.model_config, torch.device("cuda") + ) + return module + + +def _check_attention_moe_parity( + vllm_config: VllmConfig, + hf_config: DeepseekV2Config, +) -> None: + def native_moe(layer_idx: int) -> nn.Module: + return _built_moe( + partial( + native.DeepseekV2MoE, + config=hf_config, + parallel_config=vllm_config.parallel_config, + prefix=f"model.layers.{layer_idx}.mlp", + # As the FFN role builds it on CUDA. + apply_routed_scale_to_output=True, + ), + vllm_config, + ) + + internal_router_moe = native_moe(0) + external_router_moe = native_moe(1) + # As the FFN role builds it when Attention computes the gate. + external_router_moe.experts.gate = None + loopback = _LoopbackFFN(vllm_config, internal_router_moe, external_router_moe) + attention_context = ForwardContext( + no_compile_layers=vllm_config.compilation_config.static_forward_context, + attn_metadata={}, + slot_mapping={}, + additional_kwargs={ + "afd_metadata": AFDForwardContextMetadata( + tokens_start_loc=[0], + requests_start_loc=[0], + stage_idx=0, + connector=loopback, + tokens_lens=[0], + num_stages=1, + ), + }, + ) + for layer_idx, gate_on_attention in ((2, False), (3, True)): + vllm_config.additional_config["afd"] = { + "role": "attention", + "connector": "P2pNcclAFDConnector", + "compute_gate_on_attention": gate_on_attention, + } + # Equal shapes get equal values, so the gate matches the native one. + attention_moe = _built_moe( + partial( + adapter.AFDDeepseekV2RemoteExpertsMoE, + config=hf_config, + vllm_config=vllm_config, + layer_idx=layer_idx, + prefix=f"model.layers.{layer_idx}.mlp", + compute_gate_on_attention=gate_on_attention, + ), + vllm_config, + ) + loopback.expected_payload = {"router_logits"} if gate_on_attention else set() + for num_tokens in PARITY_TOKEN_COUNTS: + hidden_states = _pattern(num_tokens * PARITY_HIDDEN_SIZE, INPUT_PATTERN) + hidden_states = hidden_states.view(num_tokens, -1).to(torch.bfloat16) + with torch.inference_mode(): + with override_forward_context(attention_context): + actual = attention_moe(hidden_states.clone()).clone() + with set_forward_context(None, vllm_config, num_tokens=num_tokens): + native_output = internal_router_moe(hidden_states.clone()).clone() + torch.testing.assert_close(actual, native_output, rtol=0, atol=0) + + +def _run_attention_moe_parity( + model_dir: str, + routing: dict[str, str | int | bool], +) -> None: + """Child-process body of the parity test: one CUDA rank of vLLM.""" + hf_config = DeepseekV2Config( + architectures=["DeepseekV2ForCausalLM"], + vocab_size=128, + hidden_size=PARITY_HIDDEN_SIZE, + intermediate_size=256, + moe_intermediate_size=128, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + kv_lora_rank=64, + qk_nope_head_dim=32, + qk_rope_head_dim=32, + v_head_dim=32, + max_position_embeddings=PARITY_MODEL_LEN, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + first_k_dense_replace=0, + routed_scaling_factor=2.5, + **routing, + ) + hf_config.save_pretrained(model_dir) + vllm_config = EngineArgs( + model=model_dir, + skip_tokenizer_init=True, + dtype="bfloat16", + enforce_eager=True, + max_model_len=PARITY_MODEL_LEN, + max_num_batched_tokens=PARITY_MODEL_LEN, + ).create_engine_config() + device = torch.device("cuda", PARITY_RANK) + torch.cuda.set_device(device) + try: + with set_current_vllm_config(vllm_config): + init_worker_distributed_environment( + vllm_config, + rank=PARITY_RANK, + distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", + local_rank=PARITY_RANK, + ) + init_workspace_manager(device) + _check_attention_moe_parity(vllm_config, hf_config) + finally: + reset_workspace_manager() + destroy_model_parallel() + destroy_distributed_environment() diff --git a/tests/unit/fakes/vllm_config.py b/tests/unit/fakes/vllm_config.py index e67a400da..5f2fbdc57 100644 --- a/tests/unit/fakes/vllm_config.py +++ b/tests/unit/fakes/vllm_config.py @@ -253,6 +253,32 @@ def m2n_vllm_config( ) +def attention_moe_vllm_config( + *, + connector: str = "P2pNcclAFDConnector", + compute_gate_on_attention: bool = False, + **model_fields: Any, +): + """A real CPU ``VllmConfig`` for building Attention-role MoE runners. + + vLLM's ``FusedMoEFactory`` reads the full config, so this is not a + ``SimpleNamespace``. ``model_config`` holds only what the factory and the + AFD model shells read; ``model_fields`` adds model-specific fields. + """ + # Deferred: this fakes module must import on CPU-only hosts without torch. + import torch + from vllm.config import DeviceConfig, VllmConfig + + config = VllmConfig(device_config=DeviceConfig("cpu")) + config.additional_config["afd"] = { + "role": "attention", + "connector": connector, + "compute_gate_on_attention": compute_gate_on_attention, + } + config.model_config = SimpleNamespace(dtype=torch.float32, **model_fields) + return config + + def npu_dspark_config(*, role="attention", draft_eager=True): hf = SimpleNamespace( architectures=["DSparkDraftModel"], diff --git a/tests/unit/model_executor/ffn/test_remote.py b/tests/unit/model_executor/ffn/test_remote.py new file mode 100644 index 000000000..0a1f322df --- /dev/null +++ b/tests/unit/model_executor/ffn/test_remote.py @@ -0,0 +1,344 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""Attention-role MoE runners built through vLLM's native FusedMoEFactory.""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +torch = pytest.importorskip("torch") +pytest.importorskip("vllm") + +from vllm.config import set_current_vllm_config # noqa: E402 +from vllm.model_executor.layers.fused_moe import layer as fused_moe_layer # noqa: E402 +from vllm.model_executor.model_loader.utils import ( # noqa: E402 + process_weights_after_loading, +) + +from afd_plugin.model_executor.ffn import remote # noqa: E402 +from tests.unit.fakes.connector import install_proxy_context # noqa: E402 +from tests.unit.fakes.model_layers import GateLinear # noqa: E402 +from tests.unit.fakes.vllm_config import attention_moe_vllm_config # noqa: E402 + +LAYER_IDX = 3 +PREFIX = f"model.layers.{LAYER_IDX}.mlp.experts" +HIDDEN_SIZE = 7 +NUM_EXPERTS = 4 +TOP_K = 2 +INTERMEDIATE_SIZE = 11 +# ``ProxyConnector.recv_ffn_output`` returns the reference rows times this. +FFN_OUTPUT_SCALE = 0.25 +STAGE_IDX = 1 +# FusedMoEFactory arguments the AFD factory fixes for a weightless runner. +RESERVED_FACTORY_ARGUMENTS = ( + "quant_config", + "shared_experts", + "shared_expert_gate", + "routed_input_transform", + "routed_output_transform", + "n_shared_experts", + "fuse_shared_experts", + "apply_routed_scale_to_output", + "enable_eplb", + "num_redundant_experts", + "is_sequence_parallel", + "tp_size", + "dp_size", + "pcp_size", + "runner_cls", + "runner_args", + "routed_experts_cls", + "routed_experts_args", +) + + +def _build_runner(monkeypatch, config, *, device_type="cuda", gate=None, **kwargs): + monkeypatch.setattr( + remote, "current_platform", SimpleNamespace(device_type=device_type) + ) + with set_current_vllm_config(config): + return remote.build_attention_moe_runner( + config, + gate=gate, + num_experts=NUM_EXPERTS, + top_k=TOP_K, + hidden_size=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + prefix=PREFIX, + **kwargs, + ) + + +def _gate(): + return GateLinear(HIDDEN_SIZE, NUM_EXPERTS, out_dtype=torch.float32) + + +def _hidden_states(num_tokens: int): + return torch.arange(num_tokens * HIDDEN_SIZE, dtype=torch.float32).view( + num_tokens, HIDDEN_SIZE + ) + + +@pytest.mark.parametrize( + "compute_gate_on_attention", [False, True], ids=["ffn-gate", "attention-gate"] +) +def test_factory_registers_native_runner_without_expert_weights( + monkeypatch, compute_gate_on_attention +): + """The runner registers like a native MoE layer and holds no expert weights. + + Why: the Attention role must allocate no expert memory; the only + parameters its MoE may hold are those of the model's own gate. + """ + config = attention_moe_vllm_config( + compute_gate_on_attention=compute_gate_on_attention + ) + gate = _gate() if compute_gate_on_attention else None + + runner = _build_runner(monkeypatch, config, gate=gate) + + expected_type = ( + remote.AFDExternalRoutingMoERunner + if compute_gate_on_attention + else remote.AFDRemoteMoERunner + ) + assert type(runner) is expected_type + assert config.compilation_config.static_forward_context == {PREFIX: runner} + assert config.compilation_config.static_all_moe_layers == [PREFIX] + gate_parameters = [] if gate is None else list(gate.parameters()) + assert [id(p) for p in runner.parameters()] == [id(p) for p in gate_parameters] + assert list(runner.buffers()) == [] + + +def test_runner_stays_weightless_and_fails_on_expert_checkpoints(monkeypatch): + """Post-load processing keeps the runner weightless, and an expert weight + that reaches it fails the load. + + Why: Attention role filtering must keep FFN expert weights away; a leaked + weight must fail loudly instead of being dropped. + """ + runner = _build_runner(monkeypatch, attention_moe_vllm_config()) + model_config = SimpleNamespace( + quantization=None, word_embeddings_untied_by_checkpoint=False + ) + + process_weights_after_loading(runner, model_config, torch.device("cpu")) + + assert list(runner.parameters()) == [] + expert_weight = ("0.gate_proj.weight", torch.ones(INTERMEDIATE_SIZE, HIDDEN_SIZE)) + with pytest.raises(AttributeError, match="w13_weight"): + list(runner.load_weights([expert_weight])) + + +@pytest.mark.parametrize( + ("device_type", "connector"), + [ + ("cuda", "CAMAsyncAFDConnector"), + ("cuda", "DeepEPM2NAFDConnector"), + ("npu", "CAMP2pAFDConnector"), + ("cpu", "P2pNcclAFDConnector"), + ], +) +def test_factory_rejects_unsupported_placement_before_registration( + monkeypatch, device_type, connector +): + config = attention_moe_vllm_config(connector=connector) + + with pytest.raises(ValueError, match="unsupported Attention remote MoE"): + _build_runner(monkeypatch, config, device_type=device_type) + + assert config.compilation_config.static_forward_context == {} + + +@pytest.mark.parametrize( + "compute_gate_on_attention", [False, True], ids=["ffn-gate", "attention-gate"] +) +def test_factory_rejects_gate_that_contradicts_gate_placement( + monkeypatch, compute_gate_on_attention +): + """The model passes its gate exactly when Attention computes the gate. + + Why: the gate decides whether the FFN receives router logits; a mismatch + would send a payload the FFN role does not expect. + """ + config = attention_moe_vllm_config( + compute_gate_on_attention=compute_gate_on_attention + ) + gate = None if compute_gate_on_attention else _gate() + + with pytest.raises(ValueError, match="gate if and only if"): + _build_runner(monkeypatch, config, gate=gate) + + assert config.compilation_config.static_forward_context == {} + + +@pytest.mark.parametrize( + ("setting", "value", "message"), + [ + ("parallel_config.enable_eplb", True, "Attention-local EPLB"), + ( + "parallel_config.eplb_config.num_redundant_experts", + 1, + "Attention-local redundant experts", + ), + ( + "aux_output_config.enable_return_routed_experts", + True, + "Attention-local routed_experts capture", + ), + ], +) +def test_factory_rejects_attention_local_expert_features( + monkeypatch, setting, value, message +): + """EPLB, redundant experts and routed-expert capture need local experts. + + Why: the Attention runner owns no experts to balance, replicate or + capture; vLLM would otherwise wire these features to a runner that never + routes. + """ + config = attention_moe_vllm_config() + *parents, field = setting.split(".") + target = config + for parent in parents: + target = getattr(target, parent) + setattr(target, field, value) + + with pytest.raises(RuntimeError, match=message): + _build_runner(monkeypatch, config) + + assert config.compilation_config.static_forward_context == {} + + +@pytest.mark.parametrize("factor", [float("nan"), float("inf"), -float("inf")]) +def test_factory_rejects_nonfinite_routed_scaling_factor(monkeypatch, factor): + config = attention_moe_vllm_config() + + with pytest.raises(ValueError, match="routed_scaling_factor must be finite"): + _build_runner(monkeypatch, config, routed_scaling_factor=factor) + + assert config.compilation_config.static_forward_context == {} + + +@pytest.mark.parametrize("argument", RESERVED_FACTORY_ARGUMENTS) +def test_factory_rejects_reserved_native_arguments(monkeypatch, argument): + """The factory alone sets the arguments that describe local experts. + + Why: a caller-provided shared expert, quantization or sharding argument + would allocate or compute local experts on the Attention role. + """ + config = attention_moe_vllm_config() + + with pytest.raises(ValueError, match=f"reserves these arguments: {argument}"): + _build_runner(monkeypatch, config, **{argument: None}) + + assert config.compilation_config.static_forward_context == {} + + +def test_runner_descriptor_ignores_attention_parallel_layout(monkeypatch): + """The runner describes unsharded experts with the model's dimensions. + + Why: Attention TP and DP must not shard or map experts that live on the + FFN role. + """ + attention_tp_size = 4 + monkeypatch.setattr( + fused_moe_layer, + "get_tensor_model_parallel_world_size", + lambda: attention_tp_size, + ) + + runner = _build_runner(monkeypatch, attention_moe_vllm_config()) + + parallel = runner.moe_config.moe_parallel_config + assert (parallel.tp_size, parallel.dp_size, parallel.pcp_size) == (1, 1, 1) + assert (parallel.ep_size, parallel.use_ep) == (1, False) + assert runner.moe_config.hidden_dim == HIDDEN_SIZE + assert runner.moe_config.intermediate_size_per_partition == INTERMEDIATE_SIZE + assert runner.expert_map is None + + +@pytest.mark.parametrize("num_tokens", [1, 7]) +def test_runner_without_gate_sends_hidden_states_only(monkeypatch, num_tokens): + """With the gate on FFN, the runner sends only the hidden states, tagged + with its layer and stage, and returns the FFN output. + + Why: the FFN role runs its native internal-router MoE on the rows it + receives; the send, DBO yield and receive order is the stage contract. + """ + events: list[tuple] = [] + install_proxy_context(remote, monkeypatch, events, stage_idx=STAGE_IDX) + runner = _build_runner(monkeypatch, attention_moe_vllm_config()) + hidden_states = _hidden_states(num_tokens) + + output = runner(hidden_states=hidden_states, router_logits=hidden_states) + + assert [event[0] for event in events] == ["send", "yield", "recv"] + _, sent_states, context, send_kwargs = events[0] + assert sent_states is hidden_states + assert send_kwargs == {} + assert context.metadata.layer_idx == LAYER_IDX + assert context.metadata.stage_idx == STAGE_IDX + assert context.metadata.seq_lens == [num_tokens] + assert events[2][2] == STAGE_IDX + assert torch.equal(output, hidden_states * FFN_OUTPUT_SCALE) + + +def test_runner_with_gate_sends_its_router_logits(monkeypatch): + """With the gate on Attention, the runner applies its gate and sends the + logits with the hidden states, ignoring its router_logits argument. + + Why: native DeepseekV2MoE passes the hidden states as router_logits and + lets the runner's gate route, as native MoERunner does; the FFN role + needs the real logits for its external-router experts. + """ + events: list[tuple] = [] + install_proxy_context(remote, monkeypatch, events, stage_idx=STAGE_IDX) + gate = _gate() + config = attention_moe_vllm_config(compute_gate_on_attention=True) + runner = _build_runner(monkeypatch, config, gate=gate) + hidden_states = _hidden_states(3) + + output = runner(hidden_states=hidden_states, router_logits=hidden_states) + + assert [event[0] for event in events] == ["send", "yield", "recv"] + _, sent_states, context, send_kwargs = events[0] + assert sent_states is hidden_states + assert context.states is None + assert set(send_kwargs) == {"router_logits"} + assert torch.equal(send_kwargs["router_logits"], hidden_states @ gate.weight.t()) + assert torch.equal(output, hidden_states * FFN_OUTPUT_SCALE) + + +@pytest.mark.parametrize( + ("local_input", "message"), + [ + ("input_ids", "input_ids transport"), + ("shared_experts_input", "no local shared experts"), + ], +) +def test_runner_rejects_local_inputs_before_sending(monkeypatch, local_input, message): + """Inputs for local routing or shared experts fail before any transfer. + + Why: the FFN protocol carries no token ids for this boundary, and the + Attention role builds no shared experts to feed. + """ + events: list[tuple] = [] + install_proxy_context(remote, monkeypatch, events) + runner = _build_runner(monkeypatch, attention_moe_vllm_config()) + hidden_states = _hidden_states(2) + local_inputs = { + "input_ids": torch.tensor([1, 2]), + "shared_experts_input": hidden_states, + } + + with pytest.raises(NotImplementedError, match=message): + runner( + hidden_states=hidden_states, + router_logits=hidden_states, + **{local_input: local_inputs[local_input]}, + ) + + assert events == [] diff --git a/tests/unit/model_executor/models/test_deepseek_v2_construction.py b/tests/unit/model_executor/models/test_deepseek_v2_construction.py index ab846f7d9..195647bd5 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_construction.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_construction.py @@ -17,14 +17,16 @@ pytest.importorskip("vllm") -from vllm.config import CompilationMode # noqa: E402 +from vllm.config import CompilationMode, set_current_vllm_config # noqa: E402 from afd_plugin.config import DEEPEP_M2N_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 install_proxy_context # noqa: E402 from tests.unit.fakes.model_layers import GateLinear from tests.unit.fakes.model_layers import Stage as _FakeStage # noqa: E402 from tests.unit.fakes.model_layers import stage_type as _stage_type +from tests.unit.fakes.vllm_config import attention_moe_vllm_config # noqa: E402 @pytest.fixture @@ -135,11 +137,14 @@ def _make_layer( connector=connector, compute_gate_on_attention=attention_gate, ) - monkeypatch.setattr( - adapter, - "parse_afd_config", - lambda *_args, **_kwargs: afd_config, - ) + # The model shell and the native-factory runner read the same placement. + for module in (adapter, remote): + monkeypatch.setattr( + module, + "parse_afd_config", + lambda *_args, **_kwargs: afd_config, + ) + monkeypatch.setattr(remote, "current_platform", adapter.native.current_platform) if vllm_config is None: vllm_config = _vllm_config() return adapter.AFDDeepseekV2DecoderLayer( @@ -148,6 +153,27 @@ def _make_layer( ) +def _single_ep_rank(): + return SimpleNamespace( + device_group=SimpleNamespace(size=lambda: 1), rank_in_group=0 + ) + + +@pytest.fixture +def cuda_construction_env(monkeypatch, construction_env): + """A real vLLM config and CUDA placement for native-factory MoE runners.""" + config = attention_moe_vllm_config( + hf_config=_vllm_config().model_config.hf_config, + use_mla=False, + ) + monkeypatch.setattr( + adapter.native, "current_platform", SimpleNamespace(device_type="cuda") + ) + monkeypatch.setattr(adapter.native, "get_ep_group", _single_ep_rank) + with set_current_vllm_config(config): + yield config + + def _parameter_names(module: nn.Module) -> set[str]: return {name for name, _ in module.named_parameters()} @@ -235,84 +261,82 @@ def test_npu_gate_only_uses_native_gate_precision_and_checkpoint_name( def test_cuda_attention_gate_uses_native_gate_contract( monkeypatch, - construction_env, + cuda_construction_env, ): + """With the gate on Attention, the MoE builds the gate-applying native + factory runner, and its only parameters are the native gate's, at their + checkpoint names. - vllm_config = _vllm_config() - vllm_config.model_config.hf_config.moe_router_dtype = "float32" - monkeypatch.setattr( - adapter.native, - "current_platform", - SimpleNamespace(device_type="cuda"), - ) + Why: the gate's weights load from ``mlp.gate.*``; the experts must + allocate no local weights. + """ + config = cuda_construction_env + config.model_config.hf_config.moe_router_dtype = "float32" monkeypatch.setattr(adapter.native, "GateLinear", GateLinear) - monkeypatch.setattr( - adapter.native, - "get_tensor_model_parallel_world_size", - lambda: 1, - ) - monkeypatch.setattr( - adapter.native, - "get_tensor_model_parallel_rank", - lambda: 0, - ) - monkeypatch.setattr( - adapter.native, - "get_ep_group", - lambda: SimpleNamespace( - device_group=SimpleNamespace(size=lambda: 1), - rank_in_group=0, - ), - ) moe = _make_layer( monkeypatch, role="attention", layer_idx=1, attention_gate=True, - vllm_config=vllm_config, + vllm_config=config, ) assert isinstance(moe.mlp, adapter.AFDDeepseekV2RemoteExpertsMoE) - assert isinstance(moe.mlp.experts, adapter.AFDAttentionFusedMoE) - assert list(moe.mlp.experts.parameters()) == [] - assert list(moe.mlp.experts.buffers()) == [] - + assert type(moe.mlp.experts) is remote.AFDExternalRoutingMoERunner assert moe.mlp.gate.weight.dtype == torch.float32 - assert "mlp.gate.weight" in _parameter_names(moe) + assert {name for name in _parameter_names(moe) if name.startswith("mlp.")} == { + "mlp.gate.weight", + "mlp.gate.e_score_correction_bias", + } + assert list(moe.mlp.experts.buffers()) == [] +@pytest.mark.parametrize( + ("connector", "attention_gate", "message"), + [ + ("P2pNcclAFDConnector", False, "Attention-local EPLB"), + ("P2pNcclAFDConnector", True, "Attention-local EPLB"), + (DEEPEP_M2N_CONNECTOR, True, "do not support EPLB on the Attention role"), + ], + ids=["p2p-ffn-gate", "p2p-attention-gate", "m2n"], +) def test_cuda_remote_experts_reject_eplb_on_attention( monkeypatch, - construction_env, + cuda_construction_env, + connector, + attention_gate, + message, ): - vllm_config = _vllm_config() - vllm_config.parallel_config.enable_eplb = True - monkeypatch.setattr( - adapter.native, - "current_platform", - SimpleNamespace(device_type="cuda"), - ) + """Attention-side remote experts reject EPLB, with either transport. - with pytest.raises(RuntimeError, match="do not support EPLB"): + Why: the Attention role owns no expert weights to balance. + """ + config = cuda_construction_env + config.parallel_config.enable_eplb = True + + with pytest.raises(RuntimeError, match=message): _make_layer( monkeypatch, role="attention", layer_idx=1, - attention_gate=True, - vllm_config=vllm_config, + attention_gate=attention_gate, + vllm_config=config, + connector=connector, ) + assert config.compilation_config.static_forward_context == {} + -def test_cuda_ffn_gate_uses_parameter_free_external_router_shell( +def test_cuda_ffn_gate_shell_owns_no_mlp_parameters( monkeypatch, - construction_env, -): # long-test: native router shell ABI plus role-owned parameters - monkeypatch.setattr( - adapter.native, - "current_platform", - SimpleNamespace(device_type="cuda"), - ) + cuda_construction_env, +): + """With the gate on FFN, the Attention MoE builds no gate and no weights. + + Why: the FFN role owns the gate and all experts; the native factory + runner only sends hidden states. + """ monkeypatch.setattr( adapter.native, "GateLinear", @@ -320,45 +344,20 @@ def test_cuda_ffn_gate_uses_parameter_free_external_router_shell( "FFN-side gate must not construct an Attention gate", ), ) - monkeypatch.setattr( - adapter.native, - "get_tensor_model_parallel_world_size", - lambda: 1, - ) - monkeypatch.setattr( - adapter.native, - "get_tensor_model_parallel_rank", - lambda: 0, - ) - monkeypatch.setattr( - adapter.native, - "get_ep_group", - lambda: SimpleNamespace( - device_group=SimpleNamespace(size=lambda: 1), - rank_in_group=0, - ), - ) moe = _make_layer( monkeypatch, role="attention", layer_idx=1, attention_gate=False, + vllm_config=cuda_construction_env, ) assert isinstance(moe.mlp, adapter.AFDDeepseekV2RemoteExpertsMoE) - assert isinstance(moe.mlp.experts, adapter.AFDAttentionFusedMoE) + assert type(moe.mlp.experts) is remote.AFDRemoteMoERunner assert moe.mlp.gate is None - # Native contract: the gate-less shell is externally routed, so its proxy - # ships hidden states only and the FFN-side runner computes the routing. - assert moe.mlp.experts.external_routing is False - # 0.30's native MoE forward no longer computes an external gate, so - # the shell overrides forward and delegates to it while its gate is - # None (behavior covered by - # test_attention_shell_without_gate_keeps_native_delegation). - assert list(moe.mlp.experts.parameters()) == [] - assert list(moe.mlp.experts.buffers()) == [] assert not any(name.startswith("mlp.") for name in _parameter_names(moe)) + assert list(moe.mlp.buffers()) == [] def _use_cuda_remote_experts(monkeypatch) -> None: @@ -367,14 +366,7 @@ def _use_cuda_remote_experts(monkeypatch) -> None: "current_platform", SimpleNamespace(device_type="cuda"), ) - monkeypatch.setattr( - adapter.native, - "get_ep_group", - lambda: SimpleNamespace( - device_group=SimpleNamespace(size=lambda: 1), - rank_in_group=0, - ), - ) + monkeypatch.setattr(adapter.native, "get_ep_group", _single_ep_rank) monkeypatch.setattr(adapter.native, "_get_moe_router_dtype", lambda _c: None) monkeypatch.setattr( adapter, @@ -414,6 +406,33 @@ def test_deepep_m2n_attention_owns_gate_and_shared_experts( } +def test_deepep_m2n_shell_builds_parameter_free_proxy_transport( + monkeypatch, + construction_env, +): + """The DeepEP M2N shell builds with a parameter-free ``RemoteFFNProxy`` as + its transport, not the native factory runner. + + Why: the shell dispatches the top-k routes it selects itself, and the + native factory runner accepts only the P2P connector. + """ + _use_cuda_remote_experts(monkeypatch) + vllm_config = _vllm_config() + vllm_config.parallel_config.tensor_parallel_size = 1 + + layer = _make_layer( + monkeypatch, + role="attention", + layer_idx=1, + attention_gate=True, + vllm_config=vllm_config, + connector=DEEPEP_M2N_CONNECTOR, + ) + + assert type(layer.mlp.experts) is adapter.RemoteFFNProxy + assert list(layer.mlp.experts.parameters()) == [] + + @pytest.mark.parametrize("tp_size", [1, 2]) def test_deepep_m2n_shared_experts_replicated_under_attention_tp( monkeypatch, @@ -466,96 +485,64 @@ def test_ffn_constructs_no_real_attention( assert not any(name.startswith("self_attn.") for name in _parameter_names(moe)) -def test_attention_shell_computes_gate_before_transport( - monkeypatch, - construction_env, -): # long-test: native gate output must be present in the peer transport payload +def _remote_experts_shell(monkeypatch, *, compute_gate_on_attention: bool): + """A CUDA P2P Attention MoE shell with its native factory runner.""" + monkeypatch.setattr(remote, "current_platform", SimpleNamespace(device_type="cuda")) + monkeypatch.setattr(adapter.native, "get_ep_group", _single_ep_rank) + monkeypatch.setattr(adapter.native, "GateLinear", GateLinear) + config = attention_moe_vllm_config( + compute_gate_on_attention=compute_gate_on_attention + ) + with set_current_vllm_config(config): + return adapter.AFDDeepseekV2RemoteExpertsMoE( + config=_vllm_config().model_config.hf_config, + vllm_config=config, + layer_idx=1, + prefix="model.layers.1.mlp", + compute_gate_on_attention=compute_gate_on_attention, + ) + + +def test_attention_shell_computes_gate_before_transport(monkeypatch): """CUDA compute_gate_on_attention=True ships real router logits. - The pinned native MoE forward passes router_logits=hidden_states and - expects the experts runner to hold the gate. The Attention-side shell - owns the gate, so its forward must compute the logits before transport; - shipping hidden states breaks the P2P width contract (hidden_size vs - n_routed_experts) on the FFN side. + The native MoE forward passes router_logits=hidden_states and lets the + experts runner's gate route. The Attention runner holds the shell's gate, + so it computes the logits before transport; shipping hidden states would + break the P2P width contract (hidden_size vs n_routed_experts) on the FFN. """ - monkeypatch.setattr( - adapter.native, - "current_platform", - SimpleNamespace(device_type="cuda"), - ) - monkeypatch.setattr( - adapter.native, - "get_ep_group", - lambda: SimpleNamespace( - device_group=SimpleNamespace(size=lambda: 1), - rank_in_group=0, - ), - ) - monkeypatch.setattr(adapter.native, "_get_moe_router_dtype", lambda _c: None) - - monkeypatch.setattr(adapter.native, "GateLinear", GateLinear) events: list[tuple] = [] install_proxy_context(adapter, monkeypatch, events, stage_idx=0) - - vllm_config = _vllm_config() - shell = adapter.AFDDeepseekV2RemoteExpertsMoE( - config=vllm_config.model_config.hf_config, - parallel_config=vllm_config.parallel_config, - layer_idx=1, - prefix="model.layers.1.mlp", - compute_gate_on_attention=True, - ) - + shell = _remote_experts_shell(monkeypatch, compute_gate_on_attention=True) shell.gate.weight.data.fill_(1) hidden_states = torch.ones((2, 8), dtype=torch.bfloat16) - output = shell.forward(hidden_states) - assert [event[0] for event in events] == ["send", "yield", "recv"] - torch.testing.assert_close( - events[0][3]["router_logits"], torch.full((2, 4), 8.0, dtype=torch.bfloat16) - ) - torch.testing.assert_close(output, hidden_states * 0.25) + output = shell(hidden_states) -def test_attention_shell_without_gate_keeps_native_delegation( - monkeypatch, - construction_env, -): - """With the gate on the FFN side, the shell ships hidden states only.""" - monkeypatch.setattr( - adapter.native, - "get_ep_group", - lambda: SimpleNamespace( - device_group=SimpleNamespace(size=lambda: 1), - rank_in_group=0, - ), + assert [event[0] for event in events] == ["send", "yield", "recv"] + assert torch.equal( + events[0][3]["router_logits"], torch.full((2, 4), 8.0, dtype=torch.bfloat16) ) - monkeypatch.setattr(adapter.native, "_get_moe_router_dtype", lambda _c: None) - - sent = [] + assert torch.equal(output, hidden_states * 0.25) - class _RecordingProxy(adapter.AFDAttentionFusedMoE): - def _send_and_receive(self, hidden_states, **send_kwargs): - sent.append(send_kwargs) - return hidden_states - monkeypatch.setattr(adapter, "AFDAttentionFusedMoE", _RecordingProxy) +def test_attention_shell_without_gate_keeps_native_delegation(monkeypatch): + """With the gate on the FFN side, the native forward ships hidden states + only. - vllm_config = _vllm_config() - shell = adapter.AFDDeepseekV2RemoteExpertsMoE( - config=vllm_config.model_config.hf_config, - parallel_config=vllm_config.parallel_config, - layer_idx=1, - prefix="model.layers.1.mlp", - compute_gate_on_attention=False, - ) - assert shell.gate is None - assert shell.experts.external_routing is False + Why: the FFN role runs its own native gate on the received rows. + """ + events: list[tuple] = [] + install_proxy_context(adapter, monkeypatch, events) + shell = _remote_experts_shell(monkeypatch, compute_gate_on_attention=False) + hidden_states = torch.ones((2, 8)) - hidden_states = torch.zeros((2, 8)) - output = shell.forward(hidden_states) + output = shell(hidden_states) - assert sent == [{}] - assert torch.equal(output, hidden_states) + assert [event[0] for event in events] == ["send", "yield", "recv"] + assert torch.equal(events[0][1], hidden_states) + assert events[0][3] == {} + assert torch.equal(output, hidden_states * 0.25) def test_npu_ffn_refreshes_native_fused_moe_factory( 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 5e3567da2..15428d6fd 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_proxy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_proxy.py @@ -140,46 +140,6 @@ def compute_gate_topk(**kwargs): assert torch.equal(output, hidden_states * 0.25) -def test_remote_experts_proxy_sends_router_logits(monkeypatch): - events: list[tuple] = [] - install_proxy_context(adapter, monkeypatch, events, stage_idx=1) - proxy = adapter.AFDAttentionFusedMoE( - layer_idx=3, - external_routing=True, - ) - hidden_states = torch.ones(1, 4) - router_logits = torch.ones(1, 8) - - output = proxy(hidden_states, router_logits) - - assert [event[0] for event in events] == ["send", "yield", "recv"] - context = events[0][2] - assert context.metadata.layer_idx == 3 - assert context.metadata.stage_idx == 1 - assert context.states is None - assert events[0][3]["router_logits"] is router_logits - assert torch.equal(output, hidden_states * 0.25) - - -def test_remote_experts_proxy_without_external_routing_ships_hidden_states( - monkeypatch, -): - events: list[tuple] = [] - install_proxy_context(adapter, monkeypatch, events, stage_idx=1) - proxy = adapter.AFDAttentionFusedMoE( - layer_idx=2, - external_routing=False, - ) - hidden_states = torch.ones(1, 4) - - output = proxy(hidden_states, hidden_states) - - assert [event[0] for event in events] == ["send", "yield", "recv"] - assert events[0][3] == {} - assert torch.equal(output, hidden_states * 0.25) - assert list(proxy.parameters()) == [] - - COMBINED_VALUE = 1.0 ROUTE_WEIGHT = 0.5 @@ -223,7 +183,7 @@ def _routed_experts_moe(*, shared_experts, routed_scaling_factor, tp_size=1): ) moe.routed_scaling_factor = routed_scaling_factor moe.tp_size = tp_size - moe.experts = adapter.AFDAttentionFusedMoE(layer_idx=3, external_routing=True) + moe.experts = adapter.RemoteFFNProxy(layer_idx=3) moe.shared_experts = shared_experts return moe diff --git a/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py b/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py index fe20b7645..ebd79dc63 100644 --- a/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py +++ b/tests/unit/model_executor/models/test_deepseek_v2_weight_policy.py @@ -104,14 +104,27 @@ def test_weight_role_policy( ) +@pytest.mark.parametrize("connector", ["P2pNcclAFDConnector", "CAMP2pAFDConnector"]) def test_load_weights_passes_one_shot_generator_to_native_loader( monkeypatch: pytest.MonkeyPatch, + connector: str, ) -> None: + """Attention forwards only its own weights, once, to the native loader, + for the P2P connector as for CAMP2P. + + Why: with P2P on CUDA the Attention MoE is a native runner without expert + parameters, so an FFN-owned gate, expert or shared-expert weight that + reached the native loader would fail the load. The P2P case fails if P2P + is ever classified as a connector whose Attention owns shared experts. + """ names = [ "model.embed_tokens.weight", "model.layers.0.self_attn.q_proj.weight", "model.layers.0.mlp.gate_proj.weight", "model.layers.3.mlp.experts.0.down_proj.weight", + "model.layers.3.mlp.experts.0.down_proj.weight_scale_inv", + "model.layers.3.mlp.gate.weight", + "model.layers.3.mlp.shared_experts.gate_proj.weight", ] weights = _OneShotWeights(names) seen: list[str] = [] @@ -132,9 +145,7 @@ def fake_native_loader(self, filtered_weights): object.__setattr__( model, "afd_config", - SimpleNamespace( - compute_gate_on_attention=False, connector="CAMP2pAFDConnector" - ), + SimpleNamespace(compute_gate_on_attention=False, connector=connector), ) object.__setattr__(model, "config", _config()) diff --git a/tests/unit/model_executor/models/test_qwen3_5_contract.py b/tests/unit/model_executor/models/test_qwen3_5_contract.py index ca48798cd..858bace3b 100644 --- a/tests/unit/model_executor/models/test_qwen3_5_contract.py +++ b/tests/unit/model_executor/models/test_qwen3_5_contract.py @@ -9,11 +9,85 @@ torch = pytest.importorskip("torch") pytest.importorskip("vllm") from torch import nn # noqa: E402 +from vllm.config import set_current_vllm_config # noqa: E402 from vllm.config.multimodal import MultiModalConfig # noqa: E402 from vllm.model_executor.models import qwen3_5 as native # noqa: E402 from vllm.model_executor.models.utils import StageMissingLayer # noqa: E402 +from afd_plugin.model_executor.ffn import remote # noqa: E402 from afd_plugin.model_executor.models import qwen3_5 as adapter # noqa: E402 +from tests.unit.fakes.connector import install_proxy_context # noqa: E402 +from tests.unit.fakes.vllm_config import attention_moe_vllm_config # noqa: E402 + +QWEN_LAYER_IDX = 7 +QWEN_LAYER_PREFIX = f"model.layers.{QWEN_LAYER_IDX}" +QWEN_HIDDEN_SIZE = 8 +# ``ProxyConnector.recv_ffn_output`` returns the reference rows times this. +FFN_OUTPUT_SCALE = 0.25 + + +@pytest.fixture +def attention_qwen_layer(monkeypatch): + """A CUDA P2P Attention-role Qwen3.5 decoder layer and its vLLM config.""" + text_config = SimpleNamespace( + model_type="qwen3_5_moe_text", + rms_norm_eps=1e-6, + num_experts=4, + num_experts_per_tok=2, + hidden_size=QWEN_HIDDEN_SIZE, + moe_intermediate_size=16, + norm_topk_prob=True, + ) + config = attention_moe_vllm_config( + hf_text_config=text_config, + multimodal_config=None, + ) + monkeypatch.setattr(remote, "current_platform", SimpleNamespace(device_type="cuda")) + monkeypatch.setattr( + adapter.next_native, "get_tensor_model_parallel_world_size", lambda: 1 + ) + monkeypatch.setattr( + adapter.next_native, + "get_ep_group", + lambda: SimpleNamespace(device_group=SimpleNamespace(size=lambda: 1)), + ) + monkeypatch.setattr( + adapter.native, "Qwen3NextAttention", lambda *_args, **_kwargs: nn.Identity() + ) + with set_current_vllm_config(config): + layer = adapter.AFDQwen3_5DecoderLayer( + config, layer_type="full_attention", prefix=QWEN_LAYER_PREFIX + ) + return SimpleNamespace(layer=layer, config=config) + + +def test_attention_moe_sends_hidden_states_through_weightless_runner( + monkeypatch, attention_qwen_layer +): + """The Attention MoE is a native factory runner without gate or expert + weights; the native forward sends only the hidden states, tagged with the + decoder layer. + + Why: the FFN role owns the gate, routed experts and shared expert, and + runs its native internal-router MoE on the rows it receives. + """ + events: list[tuple] = [] + install_proxy_context(remote, monkeypatch, events) + moe = attention_qwen_layer.layer.mlp + hidden_states = torch.ones(2, QWEN_HIDDEN_SIZE) + + output = moe(hidden_states) + + assert type(moe.experts) is remote.AFDRemoteMoERunner + assert list(moe.parameters()) == [] + assert attention_qwen_layer.config.compilation_config.static_all_moe_layers == [ + f"{QWEN_LAYER_PREFIX}.mlp.experts" + ] + assert [event[0] for event in events] == ["send", "yield", "recv"] + assert torch.equal(events[0][1], hidden_states) + assert events[0][2].metadata.layer_idx == QWEN_LAYER_IDX + assert events[0][3] == {} + assert torch.equal(output, hidden_states * FFN_OUTPUT_SCALE) def test_qwen_remote_experts_proxy_preserves_completed_ffn_output_under_tp(