From 6805d58cd18345b8294167abd9f6d26078fd2f9a Mon Sep 17 00:00:00 2001 From: ksiyuan Date: Mon, 14 Sep 2026 09:54:05 +0800 Subject: [PATCH 1/5] feat(npu): wire the CAMP2P token-id channel for token-keyed routers Models whose router is keyed by token identity route on input_ids, which only the Attention role holds. The a2e operator already provides the whole transport for this: its expert_ids input slot is int32, and it returns a (base_batch_size, topk) int32 ids tensor plus a same-shape float32 scales tensor. The plugin left that path idle by sending None/None with a literal compute_gate=0, and the FFN side consumed only expand_x. Wire it up: * send_attn_output accepts an optional token-aligned input_ids. When present it packs ids and inert zero scales, and raises compute_gate to 1. * the FFN side exposes the received ids on CAMP2PTransferState, requested explicitly through recv_attn_output(recv_input_ids=True), and passes that mode on to the operator instead of leaving it at the default. compute_gate selects the operator's layout on each side independently, so both roles have to agree on it before anything is sent. The receiving rank declares it through recv_input_ids. The ids output is only written in the ids mode: reading it in the other mode yields an untouched at::empty allocation, and a token-keyed router turns that into an out-of-range table read. That failure surfaces as a device abort inside the routing operator rather than at the actual defect, which is why the mode is passed through explicitly rather than inferred from the returned tensor. Add prepare_token_id_transfer and received_token_ids with CPU tests. The received_token_ids token-count check is the alignment invariant for this transport: if the ids that arrive do not describe exactly the tokens the FFN rank computes on, a token-keyed router would silently select experts for the wrong tokens, so it fails loudly instead. Tests for the ids mode itself live in tests/unit/connectors/test_camp2p_connector.py, which requires torch_npu and therefore does not run on a CPU-only host; the helper tests above cover the transport contract everywhere. Signed-off-by: ksiyuan --- afd_plugin/connectors/npu/camp2p.py | 151 +++++++++++++++- .../unit/connectors/test_camp2p_connector.py | 39 +++++ .../unit/connectors/test_camp2p_token_ids.py | 163 ++++++++++++++++++ 3 files changed, 346 insertions(+), 7 deletions(-) create mode 100644 tests/unit/connectors/test_camp2p_token_ids.py diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index 8dcec3bf..b5b09bce 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -178,6 +178,13 @@ class CAMP2PTransferState(AFDTransferState): A2E-returned Attention token count that the FFN-to-Attention send requires. ``x_active_mask`` and ``cam_p2p_ep_name`` are the A2E-returned active-token mask and HCCL endpoint name captured on the receive path. + + ``input_ids`` holds the token-aligned ids that Attention sent alongside the + hidden states, as received by the FFN rank. It is populated only when the + receiving rank declared ``recv_input_ids``, which ``compute_gate_mode`` + records. ``compute_gate_mode`` is the operator's ids mode for this transfer + and has to equal the mode the sending rank selected, because the operator + only writes the ids slot in that mode. """ aiv_num: int = 8 @@ -187,6 +194,8 @@ class CAMP2PTransferState(AFDTransferState): atten_batch_size: torch.Tensor | None = None x_active_mask: torch.Tensor | None = None cam_p2p_ep_name: str | None = None + input_ids: torch.Tensor | None = None + compute_gate_mode: int = 0 @dataclass(frozen=True, slots=True) @@ -224,6 +233,90 @@ def is_attn_top_min_size_rank(self) -> bool: return self.ffn_size <= self.world_rank < self.ffn_size + self.min_size +def prepare_token_id_transfer( + input_ids: torch.Tensor, + *, + topk: int, + expected_tokens: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pack token ids and inert scales for the A2E ids channel. + + The operator transports an ``int32`` ids tensor and a ``float32`` scales + tensor of shape ``(batch, topk)``. The id repeats across the columns, which + is what the receiving side collapses back. The scales carry no routing + weight: this channel moves token identity, not router output. + + Args: + input_ids: Token-aligned ids for the local Attention tokens. + topk: Number of routed experts per token; the operator's column count. + expected_tokens: Token count the transfer metadata declares. + + Returns: + The ``(ids, scales)`` pair to hand to the operator. + + Raises: + ValueError: If the ids do not describe exactly ``expected_tokens`` + tokens, which is the alignment invariant for this channel. + """ + + num_tokens = int(input_ids.numel()) + if num_tokens != expected_tokens: + raise ValueError( + f"input_ids token count {num_tokens} does not match the AFD " + f"transfer token count {expected_tokens}", + ) + + ids = ( + input_ids.reshape(-1) + .to(dtype=torch.int32) + .unsqueeze(1) + .expand(-1, topk) + .contiguous() + ) + scales = torch.zeros( + (num_tokens, topk), + dtype=torch.float32, + device=input_ids.device, + ) + return ids, scales + + +def received_token_ids( + sim_expert_ids: torch.Tensor, + *, + expected_tokens: int, +) -> torch.Tensor: + """Collapse the A2E ids channel back to a token-aligned id vector. + + The operator returns ``(N, topk)`` with every column of a row repeating that + token's id, and sizes ``N`` to the capacity it was given, so this trims to + the tokens the FFN rank actually computes on. + + Args: + sim_expert_ids: The operator's ids output. + expected_tokens: Token count derived from the FFN rank's DP metadata, + which is what the FFN compute will actually run on. + + Returns: + A one-dimensional ``int32`` tensor of length ``expected_tokens``. + + Raises: + ValueError: If fewer ids arrived than the FFN rank computes on. That is + the misalignment this channel has to catch: a token-keyed router + would otherwise select experts for the wrong tokens. + """ + + received_tokens = int(sim_expert_ids.shape[0]) + if received_tokens < expected_tokens: + raise ValueError( + f"received {received_tokens} token ids but the FFN rank computes on " + f"{expected_tokens} tokens; the ids channel is not aligned with the " + "FFN token layout", + ) + # The columns are replicas of the same id, so the first one carries it. + return sim_expert_ids[:expected_tokens, 0].contiguous() + + class CAMP2pAFDConnector(AFDConnectorBase): """Move model data between Attention and FFN workers on Ascend NPU. @@ -429,12 +522,16 @@ def send_attn_output( hidden_states: Model data with shape ``(tokens, hidden_size)``. context: Transfer context whose ``metadata`` supplies the layer number, ubatch number, and token count for this transfer. - **kwargs: Extra arguments accepted for interface compatibility. + **kwargs: An optional token-aligned ``input_ids`` tensor. When it is + supplied, the transfer runs with ``compute_gate=1`` so the ids + reach the FFN rank through the operator's ids channel, and the + local FFN-side receive exposes them on ``CAMP2PTransferState``. Raises: RuntimeError: If the communication groups are not ready. ValueError: If the number of tokens in ``hidden_states`` does not - match ``context.metadata`` outside a ``torch.compile`` trace. + match ``context.metadata`` outside a ``torch.compile`` trace, or + if a supplied ``input_ids`` tensor is malformed. """ if not self._initialized: raise RuntimeError("CAMP2P connector is not initialized") @@ -446,6 +543,17 @@ def send_attn_output( f"hidden_states shape {hidden_states.shape!r} does not match " f"CAMP2P metadata token count {metadata.total_tokens}", ) + input_ids = cast(torch.Tensor | None, kwargs.get("input_ids")) + expert_ids: torch.Tensor | None = None + expert_scales: torch.Tensor | None = None + compute_gate = 0 + if input_ids is not None: + expert_ids, expert_scales = prepare_token_id_transfer( + input_ids, + topk=self.num_experts_per_tok, + expected_tokens=metadata.total_tokens, + ) + compute_gate = 1 transfer_state = CAMP2PTransferState( aiv_num=self.aiv_num, batch_size=metadata.total_tokens, @@ -469,7 +577,9 @@ def send_attn_output( self.attn_size, self.world_rank, transfer_state.aiv_num, - 0, + compute_gate, + expert_ids, + expert_scales, ) return None @@ -526,7 +636,12 @@ def recv_attn_output( Args: ubatch_idx: Ubatch number, starting from ``0``. **kwargs: May provide existing transfer information or the layer - number needed to create it. + number needed to create it. ``recv_input_ids`` states that this + FFN rank expects token ids on the transfer, which selects the + operator's ids mode and makes the connector validate and expose + the operator's ids slot. The sending rank has to select the same + mode, so only request ids for a run whose Attention role + transports them. Returns: The received hidden states and the information FFN needs to process @@ -535,11 +650,14 @@ def recv_attn_output( Raises: RuntimeError: If communication is not ready, transfer information is missing, or the requested ubatch group does not exist. + ValueError: If ids were requested but do not align with the FFN + rank's token layout. """ if not self._initialized: raise RuntimeError("CAMP2P connector is not initialized") layer_idx: int = kwargs.get("layer_idx", 0) max_num_tokens: int = kwargs.get("max_num_tokens", 0) + recv_input_ids: bool = bool(kwargs.get("recv_input_ids", False)) batch_size = _num_tokens_for_ffn_rank( self.dp_metadata_list, ubatch_idx, @@ -558,6 +676,7 @@ def recv_attn_output( batch_size=batch_size, h=self.hidden_size, k=self.num_experts_per_tok, + compute_gate_mode=1 if recv_input_ids else 0, ) context = AFDTransferContext( metadata=metadata, @@ -582,11 +701,23 @@ def recv_attn_output( self.world_rank, group_ep, custom_states.aiv_num, - 0, + custom_states.compute_gate_mode, ) custom_states.atten_batch_size = outputs[3] custom_states.x_active_mask = outputs[4] custom_states.cam_p2p_ep_name = self.hccl_comm_name1 + # The ids slot is only written in the operator's ids mode, so the mode has + # to match the sending rank's ``compute_gate``. It is declared here by the + # receiving FFN rank through ``recv_input_ids`` rather than inferred from + # the returned tensor: a genuine single-token layer would otherwise be + # indistinguishable from the operator's placeholder. Reading the slot in + # the other mode would hand the model uninitialised device memory as token + # ids, which a token-keyed router turns into an out-of-range table read. + if custom_states.compute_gate_mode == 1: + custom_states.input_ids = received_token_ids( + outputs[1], + expected_tokens=batch_size, + ) return AFDA2FTransferPayload( hidden_states=outputs[0], context=context, @@ -850,6 +981,8 @@ def send_attn_output_impl( world_rank: int, aiv_num: int, compute_gate: int, + expert_ids: torch.Tensor | None, + expert_scales: torch.Tensor | None, ) -> torch.Tensor: transfer_state = getattr(get_forward_context(), "cam_afdtransfer_state", None) if transfer_state is None: @@ -867,8 +1000,8 @@ def send_attn_output_impl( outputs = torch.ops.afd_ascend.a2e( hidden_states, - None, - None, + expert_ids, + expert_scales, transfer_state.batch_size, transfer_state.h, transfer_state.k, @@ -897,6 +1030,8 @@ def send_attn_output_fake_impl( world_rank: int, aiv_num: int, compute_gate: int, + expert_ids: torch.Tensor | None, + expert_scales: torch.Tensor | None, ) -> torch.Tensor: """Return the input unchanged while PyTorch inspects the send operation.""" return hidden_states @@ -970,6 +1105,8 @@ def recv_ffn_output_fake_impl( "world_rank": int, "aiv_num": int, "compute_gate": int, + "expert_ids": torch.Tensor | None, + "expert_scales": torch.Tensor | None, "return": torch.Tensor, } recv_annotations = { diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index df513da6..86e5239c 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -149,6 +149,45 @@ def test_camp2p_recv_attn_output_uses_original_contiguous_af_grouping(monkeypatc assert context0.states.k == 2 +def test_camp2p_recv_attn_output_drives_the_operator_ids_mode(monkeypatch): + """Requesting ids must reach the operator, not stay a connector-local flag. + + The ``a2e`` operator only writes its ids slot in the ids mode, and the + sending rank selects that mode independently. A receiving rank that reads the + slot in the other mode would install uninitialised device memory as token + ids, which a token-keyed router turns into an out-of-range table read. + """ + + torch = pytest.importorskip("torch") + calls: list[tuple] = [] + + def fake_a2e(*args): + calls.append(args) + tokens, topk = args[3], args[5] + ids = torch.arange(tokens * topk, dtype=torch.int32).reshape(tokens, topk) + return ("hidden", ids, None, "atten-batch", "active-mask") + + monkeypatch.setattr(torch.ops.afd_ascend, "a2e", fake_a2e, raising=False) + connector = _init_ffn_connector(0, _vllm_config()) + connector.dp_metadata_list = {0: _FakeDPMetadata([2, 3, 5, 7])} + + with_ids = connector.recv_attn_output( + ubatch_idx=0, + layer_idx=0, + recv_input_ids=True, + ) + without_ids = connector.recv_attn_output( + ubatch_idx=0, + layer_idx=0, + recv_input_ids=False, + ) + + assert calls[0][-1] == 1 + assert with_ids.context.states.input_ids.tolist() == [0, 2, 4, 6, 8] + assert calls[1][-1] == 0 + assert without_ids.context.states.input_ids is None + + def test_camp2p_extra_info_rejects_unknown_mix_placement(): with pytest.raises(ValueError, match="unknown CAMP2P connector_extra_config"): CAMP2PExtraInfo.from_mapping({"mix_placement": True}) diff --git a/tests/unit/connectors/test_camp2p_token_ids.py b/tests/unit/connectors/test_camp2p_token_ids.py new file mode 100644 index 00000000..845b991b --- /dev/null +++ b/tests/unit/connectors/test_camp2p_token_ids.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for CAMP2P token-id transport helpers. + +These helpers are pure tensor logic, so they are testable without an Ascend +device. The surrounding connector needs ``torch_npu`` and is covered by +``test_camp2p_connector.py`` instead. + +The module-level ``vllm`` stub mirrors the pattern used by the compat patch +tests: importing any connector module pulls in ``vllm`` through +``afd_plugin.connectors``, so the import is satisfied with the minimal surface +the helpers' module needs. +""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import types +from collections.abc import Iterator + +import pytest + +torch = pytest.importorskip("torch") + +_STUB_MODULES = ( + "vllm", + "vllm.forward_context", + "vllm.logger", + "vllm.utils", + "vllm.utils.torch_utils", + "vllm.distributed", + "vllm.distributed.parallel_state", +) + + +@contextlib.contextmanager +def _vllm_stub() -> Iterator[None]: + """Expose a minimal ``vllm`` surface for the duration of one import. + + Importing any connector module pulls in ``vllm`` through + ``afd_plugin.connectors``. This stub supplies only the names the helpers' + module needs at import time. + + The stub is removed again immediately afterwards. Leaving a partial + ``vllm`` in ``sys.modules`` would mask the real "vllm is not installed" + failure for every other test module in the same pytest session, turning a + clear error into a confusing one. + """ + + missing = [name for name in _STUB_MODULES if name not in sys.modules] + if not missing: + yield + return + + saved = {name: sys.modules.get(name) for name in missing} + + def make(name: str, **attributes: object) -> None: + module = types.ModuleType(name) + module.__path__ = [] # type: ignore[attr-defined] + for attribute, value in attributes.items(): + setattr(module, attribute, value) + sys.modules[name] = module + + make("vllm") + make( + "vllm.forward_context", + DPMetadata=type("DPMetadata", (), {}), + get_forward_context=lambda: None, + ) + make("vllm.logger", init_logger=lambda *args, **kwargs: logging.getLogger("test")) + make("vllm.utils") + make( + "vllm.utils.torch_utils", + direct_register_custom_op=lambda **kwargs: None, + is_torch_equal_or_newer=lambda *args, **kwargs: True, + ) + make("vllm.distributed") + make( + "vllm.distributed.parallel_state", + get_pcp_group=None, + get_tensor_model_parallel_rank=None, + ) + try: + yield + finally: + for name in missing: + sys.modules.pop(name, None) + for name, module in saved.items(): + if module is not None: + sys.modules[name] = module + + +with _vllm_stub(): + from afd_plugin.connectors.npu.camp2p import ( # noqa: E402 + prepare_token_id_transfer, + received_token_ids, + ) + + +def test_prepare_token_id_transfer_replicates_ids_across_columns(): + input_ids = torch.tensor([7, 11, 13], dtype=torch.int64) + + ids, scales = prepare_token_id_transfer( + input_ids, + topk=2, + expected_tokens=3, + ) + + assert ids.dtype == torch.int32 + assert tuple(ids.shape) == (3, 2) + assert ids[:, 0].tolist() == [7, 11, 13] + assert ids[:, 1].tolist() == [7, 11, 13] + assert scales.dtype == torch.float32 + assert tuple(scales.shape) == (3, 2) + # Scales accompany token identity, not routing weights, so they are inert. + assert torch.count_nonzero(scales) == 0 + + +def test_prepare_token_id_transfer_rejects_token_count_mismatch(): + input_ids = torch.tensor([7, 11], dtype=torch.int32) + + with pytest.raises(ValueError, match="does not match the AFD transfer"): + prepare_token_id_transfer(input_ids, topk=2, expected_tokens=3) + + +def test_received_token_ids_collapses_replicated_columns(): + sent_ids, _ = prepare_token_id_transfer( + torch.tensor([7, 11, 13], dtype=torch.int32), + topk=2, + expected_tokens=3, + ) + + received = received_token_ids(sent_ids, expected_tokens=3) + + assert received.dtype == torch.int32 + assert received.tolist() == [7, 11, 13] + + +def test_received_token_ids_trims_operator_padded_capacity(): + # The operator works on a padded capacity, so extra trailing rows are normal. + sent_ids, _ = prepare_token_id_transfer( + torch.tensor([7, 11, 13, 99, 99], dtype=torch.int32), + topk=2, + expected_tokens=5, + ) + + received = received_token_ids(sent_ids, expected_tokens=3) + + assert received.tolist() == [7, 11, 13] + + +def test_received_token_ids_rejects_alignment_shortfall(): + """Misaligned ids must fail loudly rather than route the wrong tokens.""" + sent_ids, _ = prepare_token_id_transfer( + torch.tensor([7, 11], dtype=torch.int32), + topk=2, + expected_tokens=2, + ) + + with pytest.raises(ValueError, match="not aligned with the FFN token layout"): + received_token_ids(sent_ids, expected_tokens=5) From 500b5fec9f5bd144a59a67c656b9b279ab169b96 Mon Sep 17 00:00:00 2001 From: ksiyuan Date: Mon, 14 Sep 2026 10:34:07 +0800 Subject: [PATCH 2/5] feat(npu): run DeepSeek V4 over the synchronous CAMP2P boundary DSV4 was restricted to CAMAsyncAFDConnector, which needs CAM/UMDK operator packages that ship for 910C today. That made DSV4 unrunnable through AFD on Ascend 950 (A5), whose only AFD transport is the plugin's own a2e/e2a operator pair. The blocker is narrow: a DSV4 Hash layer routes by token identity, so whichever role owns the gate needs the tokens' ids, and only Attention holds input_ids. The FFN side consumes them through the forward context. vLLM-Ascend's fused-expert selector reads forward_context.input_ids for Hash routing, so installing the transported ids there lets the native MoE pick them up without changes to its internals. Attention: * AFDDeepseekV4RemoteMoE replaces the plain RemoteFFNProxy for gate-on-FFN layers and sends the rank-local ids alongside the activations. Connectors that do not transport ids ignore the extra argument. * local_hash_input_ids lifts the id selection out of the Hash routing path so the send side reuses one implementation. The send-side slice and the local-routing slice have to agree, and a second copy would eventually drift into routing the wrong tokens. * hash_input_ids_from_context reuses that helper for the send side and raises when the forward context carries no ids. Whether ids cross the boundary is a run-level decision the two roles share through the model's afd_requires_input_ids declaration, not a per-layer one: the FFN role waits for the operator's ids channel, so a quiet activations-only fallback would leave it reading a slot the operator never wrote. FFN: * ascend_forward_context gains input_ids and installs it. * The runner receives the transfer before building the forward context rather than inside it, because the ids arrive with that transfer; the previous order could not have installed them. It resolves recv_input_ids once per forward from the model's afd_requires_input_ids, outside the layer loop, and compute_ffn_output forwards the ids to the MoE. Role-aware weight filtering. The upstream Ascend loader indexes its parameter dict by name without a membership check, so a checkpoint path handed to a role that never registered it raises KeyError instead of being skipped. Two paths were misassigned for the gate-on-FFN configuration and are fixed here: * gate.tid2eid is registered only where the Hash MoE is built, which is FFN, so Attention must not receive it. * With the gate on FFN, Attention builds no router at all, so no gate.* parameter belongs to it. Gate ownership now follows the configured placement rather than assuming a shared router. Validation admits CAMP2P for DSV4. Per-connector checks still apply, so CAMP2P keeps requiring compute_gate_on_attention=false and quant_mode=0, while CAM async keeps requiring gate-on-Attention. A CAMP2P_CONNECTOR constant replaces the remaining duplicate connector-name literal. Verified on Ascend 950 (A5): DSV4 completes a full request through the synchronous CAMP2P boundary and produces sensible output. Accuracy has not been compared against a native A5 run, so that remains open. Signed-off-by: ksiyuan --- afd_plugin/compat/npu/feature_validation.py | 15 +- afd_plugin/compat/npu/forward_context.py | 13 +- afd_plugin/config.py | 4 +- .../model_executor/models/npu/deepseek_v4.py | 89 ++++++- .../models/npu/deepseek_v4_attention_gate.py | 125 +++++++--- afd_plugin/v1/worker/npu/ffn_model_runner.py | 31 ++- .../models/test_deepseek_v4_attention_gate.py | 25 +- .../test_deepseek_v4_hash_ids.py | 222 ++++++++++++++++++ .../test_deepseek_v4_npu_weight_roles.py | 217 +++++++++++++++++ 9 files changed, 682 insertions(+), 59 deletions(-) create mode 100644 tests/unit/model_executor/test_deepseek_v4_hash_ids.py create mode 100644 tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index a8278e93..ec6f3214 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -8,6 +8,7 @@ from afd_plugin.config import ( AFD_ASYNC_CONNECTOR, + CAMP2P_CONNECTOR, AFDConfig, is_afd_async_dp, parse_afd_config, @@ -130,8 +131,18 @@ def _fail_if_unsupported_dsv4_async_features( def _fail_if_unsupported_dsv4_connector(afd_config: AFDConfig) -> None: - if afd_config.connector != AFD_ASYNC_CONNECTOR: - raise RuntimeError("DSV4 NPU AFD supports only CAMAsyncAFDConnector") + """Reject DSV4 on connectors that cannot carry token ids. + + A DSV4 Hash layer routes by token identity, so whichever role owns the gate + needs the tokens' ids. CAM async reads them from its dispatch metadata; + CAMP2P moves them over the A2E ids channel with the gate left on FFN. + """ + + if afd_config.connector not in (AFD_ASYNC_CONNECTOR, CAMP2P_CONNECTOR): + raise RuntimeError( + "DSV4 NPU AFD supports only CAMAsyncAFDConnector and " + f"CAMP2pAFDConnector; got {afd_config.connector!r}", + ) def _fail_if_unsupported_npu_afd_async_features( diff --git a/afd_plugin/compat/npu/forward_context.py b/afd_plugin/compat/npu/forward_context.py index df50f941..b5c71a1f 100644 --- a/afd_plugin/compat/npu/forward_context.py +++ b/afd_plugin/compat/npu/forward_context.py @@ -27,8 +27,15 @@ def ascend_forward_context( in_profile_run: bool = False, aclgraph_runtime_mode: CUDAGraphMode | None = None, skip_mc2_mask: bool = False, + input_ids: torch.Tensor | None = None, ) -> Iterator[ForwardContext]: - """Create the minimal forward context needed by connector-driven FFN steps.""" + """Create the minimal forward context needed by connector-driven FFN steps. + + ``input_ids`` carries the token ids of exactly the tokens this FFN rank + computes on. Models whose router is keyed by token identity read it from the + forward context, so a connector that transports ids must install them here + before running the FFN compute. + """ from vllm.config import CUDAGraphMode from vllm.forward_context import get_forward_context @@ -60,6 +67,8 @@ def ascend_forward_context( forward_context = get_forward_context() forward_context.additional_kwargs["afd_metadata"] = afd_metadata forward_context.additional_kwargs["model_instance"] = model_instance + if input_ids is not None: + forward_context.input_ids = input_ids yield forward_context return @@ -81,6 +90,8 @@ def ascend_forward_context( if forward_context.additional_kwargs is None: forward_context.additional_kwargs = {} forward_context.additional_kwargs["afd_metadata"] = afd_metadata + if input_ids is not None: + forward_context.input_ids = input_ids yield forward_context diff --git a/afd_plugin/config.py b/afd_plugin/config.py index 256590ab..8d792bb1 100644 --- a/afd_plugin/config.py +++ b/afd_plugin/config.py @@ -17,12 +17,13 @@ AFD_ADDITIONAL_CONFIG_KEY: Final[str] = "afd" AFD_ASYNC_CONNECTOR: Final[str] = "CAMAsyncAFDConnector" +CAMP2P_CONNECTOR: Final[str] = "CAMP2pAFDConnector" AFDRole = Literal["attention", "ffn"] SUPPORTED_AFD_ROLES: Final[tuple[str, ...]] = ("attention", "ffn") SUPPORTED_AFD_CONNECTORS: Final[tuple[str, ...]] = ( "P2pNcclAFDConnector", - "CAMP2pAFDConnector", + CAMP2P_CONNECTOR, AFD_ASYNC_CONNECTOR, ) @@ -335,6 +336,7 @@ def validate_afd_config( "afd_config_from_mapping", "AFD_ADDITIONAL_CONFIG_KEY", "AFDRole", + "CAMP2P_CONNECTOR", "SUPPORTED_AFD_CONNECTORS", "SUPPORTED_AFD_ROLES", "connector_extra_config_from_mapping", diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index 32a74a6b..42da9aa6 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -69,12 +69,22 @@ def _weight_layer_path(name: str) -> tuple[int, str, tuple[str, ...]] | None: return None -def _checkpoint_weight_roles(name: str) -> frozenset[str]: +def _checkpoint_weight_roles( + name: str, + *, + attn_owns_gate: bool = True, +) -> frozenset[str]: """Return the AFD owner for a DSV4 checkpoint path. DSV4 checkpoints use ``attn``/``ffn`` names while the Ascend runtime model exposes ``self_attn``/``mlp``. The native loader performs that name conversion later, so filtering must understand both spellings here. + + ``attn_owns_gate`` describes whether the Attention role built a router for + the current configuration. Handing a role a path it never registered is not + a harmless no-op: the upstream Ascend loader indexes its parameter dict by + name without a membership check, so it raises ``KeyError`` instead of + skipping. """ layer_path = _weight_layer_path(name) @@ -86,7 +96,18 @@ def _checkpoint_weight_roles(name: str) -> frozenset[str]: return frozenset((_ATTENTION_ROLE,)) if stage in ("ffn", "mlp"): if remainder and remainder[0] == "gate": - return _BOTH_ROLES + # The Hash id table is a parameter only where the Hash MoE is + # built, which is the FFN role, whatever the gate placement: with + # the gate on FFN the table lives on the FFN router, and with the + # gate on Attention the Hash path routes from the table instead of + # from a gate weight. + if "tid2eid" in remainder: + return frozenset((_FFN_ROLE,)) + # The remaining gate parameters belong to every role that built a + # router. With the gate on FFN, Attention has none. + if attn_owns_gate: + return _BOTH_ROLES + return frozenset((_FFN_ROLE,)) return frozenset((_FFN_ROLE,)) # HC parameters and any future shared layer parameters are required by # both role-local model instances. @@ -97,12 +118,47 @@ def _iter_role_weights( weights: Iterable[tuple[str, torch.Tensor]], *, role: str, + attn_owns_gate: bool = True, ) -> Iterator[tuple[str, torch.Tensor]]: for name, loaded_weight in weights: - if role in _checkpoint_weight_roles(name): + if role in _checkpoint_weight_roles(name, attn_owns_gate=attn_owns_gate): yield name, loaded_weight +class AFDDeepseekV4RemoteMoE(RemoteFFNProxy): + """DSV4 gate-on-FFN shell that sends Hash ids alongside the activations. + + The FFN role owns the gate for this configuration. Its Hash layers route by + token identity, and only Attention holds ``input_ids``, so Attention sends + the rank-local ids that the FFN rank's tokens correspond to. Connectors that + do not transport ids ignore the extra argument, which keeps this shell valid + for gate-on-FFN configurations in general. + + Whether ids cross the boundary is a run-level decision the two roles share + through the model's ``afd_requires_input_ids`` declaration, not a per-layer + one, because the FFN role cannot tell Hash layers from non-Hash ones. A + forward context with no ids is therefore an error here rather than a silent + activations-only send. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + from afd_plugin.model_executor.models.npu.deepseek_v4_attention_gate import ( + hash_input_ids_from_context, + ) + + # The FFN rank asks for the operator's ids channel on every layer, so + # this side must send ids rather than fall back to an activations-only + # transfer. ``hash_input_ids_from_context`` raises if the forward context + # cannot supply them. + return self._send_and_receive( + hidden_states, + input_ids=hash_input_ids_from_context( + forward_context=get_forward_context(), + router_tokens=int(hidden_states.shape[0]), + ), + ) + + class AFDDeepseekV4AttentionGateRemoteMoE(RemoteFFNProxy): """DSV4 gate shell that routes local Attention tokens through Async CAM.""" @@ -234,7 +290,7 @@ def __init__( prefix=f"{prefix}.mlp", ) else: - self.mlp = RemoteFFNProxy(layer_idx=layer_idx) + self.mlp = AFDDeepseekV4RemoteMoE(layer_idx=layer_idx) elif afd_config.role == _FFN_ROLE: self.self_attn = native.PPMissingLayer() _refresh_ascend_fused_moe() @@ -275,6 +331,7 @@ def compute_ffn_output( dynamic_scales_shared: torch.Tensor | None = None, topk_scales: torch.Tensor | None = None, group_list_type: int = 1, + input_ids: torch.Tensor | None = None, **_: Any, ) -> torch.Tensor | AFDF2ATransferPayload: if not isinstance(self.mlp, native.DeepseekV4MoE): @@ -303,7 +360,14 @@ def compute_ffn_output( # topk_weights, which CAM applies during combine-recv. routed_scale_applied_in_topk=True, ) - return self.mlp(hidden_states) + # ### PATCH START: FFN-side Hash routing needs the transported ids. + # The native MoE runs the gate internally when the gate is not on + # Attention, and its Hash layers route by token identity: vLLM-Ascend's + # FusedMoE reads `forward_context.input_ids`, which AFD installs from the + # transfer. Pass them on to the native MoE as well so the ids travel with + # the call rather than only through ambient context. + return self.mlp(hidden_states, input_ids=input_ids) + # ### PATCH END: FFN-side Hash routing needs the transported ids. @native.support_torch_compile @@ -533,6 +597,11 @@ class AFDDeepseekV4ForCausalLM(native.AscendDeepseekV4ForCausalLM): model_cls = AFDDeepseekV4Model + # DSV4 Hash layers route by token identity. The FFN role does not hold + # input_ids, so the connector must transport them and the FFN runner + # installs them in the forward context before the FFN compute. + afd_requires_input_ids = True + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: self.afd_config = parse_afd_config(vllm_config, validate=False) self.afd_role = self.afd_config.role @@ -568,7 +637,15 @@ def compute_experts_output( ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - return super().load_weights(_iter_role_weights(weights, role=self.afd_role)) + # Only the gate-on-Attention configuration gives Attention a router; with + # the gate on FFN its MoE slot is a parameter-free transfer shell. + role_weights = _iter_role_weights( + weights, + role=self.afd_role, + attn_owns_gate=bool(self.afd_config.compute_gate_on_attention), + ) + loaded = super().load_weights(role_weights) + return loaded __all__ = [ diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py b/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py index 49acdd3b..2ac96050 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py @@ -9,11 +9,96 @@ import torch if TYPE_CHECKING: + from vllm.forward_context import ForwardContext + from afd_plugin.model_executor.models.npu.deepseek_v4 import ( AFDDeepseekV4AttentionGateRemoteMoE, ) +def local_hash_input_ids( + *, + input_ids: torch.Tensor | None, + router_tokens: int, + flash_comm_v1_enabled: bool, + pad_size: int, +) -> torch.Tensor: + """Return the rank-local token ids that a Hash layer routes on. + + A DSV4 Hash layer routes by token identity rather than by router logits, so + the FFN rank executing that layer needs the ids of exactly the tokens it + computes on. On Attention the forward context carries the *global* ids while + FlashComm v1 shards router logits across TP ranks, so the global vector must + receive the same padding and contiguous TP split as the logits before it can + be sent. + + Both the local routing path and the AFD send path call this, so the ids that + cross the boundary describe the same tokens the Attention-side routing used. + + Args: + input_ids: Global ids from the forward context, or ``None``. + router_tokens: Token count of this rank's router logits. + flash_comm_v1_enabled: Whether FlashComm v1 is active for this forward. + pad_size: FlashComm v1 padding applied to the activation. + + Returns: + A one-dimensional ``int64`` tensor of ``router_tokens`` local ids. + + Raises: + RuntimeError: If ids are unavailable, or if the ids do not describe + exactly ``router_tokens`` tokens. Both would otherwise let a + token-keyed router select experts for the wrong tokens. + """ + + if input_ids is None: + raise RuntimeError( + "DSV4 Hash routing requires input_ids to send towards the FFN role, " + "but the forward context carries none. This path routes by token " + "identity and has no fallback, so the runner must install the " + "request's input_ids before the model forward.", + ) + ids = input_ids.reshape(-1).to(torch.int64) + if flash_comm_v1_enabled and ids.numel() != router_tokens: + from vllm.distributed import get_tp_group + from vllm_ascend.distributed.utils import split_tensor_along_first_dim + + if pad_size > 0: + ids = torch.nn.functional.pad(ids, (0, pad_size)) + group = get_tp_group() + ids = split_tensor_along_first_dim( + ids, + num_partitions=group.world_size, + contiguous_split_chunks=True, + )[group.rank_in_group] + if ids.numel() != router_tokens: + raise RuntimeError( + "DSV4 Hash routing cannot align the ids sent to FFN with the local " + f"tokens: ids={ids.numel()} router_tokens={router_tokens}", + ) + return ids + + +def hash_input_ids_from_context( + *, + forward_context: ForwardContext, + router_tokens: int, +) -> torch.Tensor: + """Return the ids to send for a Hash layer, raising if the context has none. + + The ids channel belongs to the transfer rather than to one layer: the FFN + role cannot tell a Hash layer from a non-Hash one, so it asks for ids on + every layer. Answering with activations alone would leave it reading an ids + slot the operator never wrote. + """ + + return local_hash_input_ids( + input_ids=forward_context.input_ids, + router_tokens=router_tokens, + flash_comm_v1_enabled=forward_context.flash_comm_v1_enabled, + pad_size=forward_context.pad_size, + ) + + def compute_attention_gate_topk( moe: AFDDeepseekV4AttentionGateRemoteMoE, hidden_states: torch.Tensor, @@ -54,40 +139,12 @@ def _compute_sqrtsoftplus_topk( from vllm.forward_context import get_forward_context forward_context = get_forward_context() - input_ids = getattr(forward_context, "input_ids", None) - if input_ids is None: - raise RuntimeError( - "DSV4 Hash routing requires local input_ids in the forward context", - ) - input_ids = input_ids.reshape(-1).to(torch.int64) - # FlashComm v1 shards router logits across TP ranks, but the forward - # context still carries global input IDs. Apply the same padding and - # contiguous TP split so Hash routing receives rank-local token IDs. - if ( - forward_context.flash_comm_v1_enabled - and input_ids.numel() != router_logits.shape[0] - ): - from vllm.distributed import get_tp_group - from vllm_ascend.distributed.utils import ( - split_tensor_along_first_dim, - ) - - if forward_context.pad_size > 0: - input_ids = torch.nn.functional.pad( - input_ids, - (0, forward_context.pad_size), - ) - tp_group = get_tp_group() - input_ids = split_tensor_along_first_dim( - input_ids, - num_partitions=tp_group.world_size, - contiguous_split_chunks=True, - )[tp_group.rank_in_group] - if input_ids.numel() != router_logits.shape[0]: - raise RuntimeError( - "DSV4 Hash routing input_ids/token count mismatch on Attention: " - f"input_ids={input_ids.numel()} router_tokens={router_logits.shape[0]}", - ) + input_ids = local_hash_input_ids( + input_ids=getattr(forward_context, "input_ids", None), + router_tokens=router_logits.shape[0], + flash_comm_v1_enabled=forward_context.flash_comm_v1_enabled, + pad_size=forward_context.pad_size, + ) input_ids = torch.where(input_ids == -1, 0, input_ids) tid2eid = tid2eid.to(torch.int32) correction_bias = moe.gate.e_score_correction_bias diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 2f42c6bc..9289b6ab 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -244,6 +244,10 @@ def _ffn_forward( ) stage_ids = sorted(int(stage_idx) for stage_idx in dp_metadata_list) or [0] rank_ffn_output = None + # A model whose router is keyed by token identity declares that the FFN + # role needs the tokens' ids. Resolve it once per forward, outside the + # layer loop, so a per-layer re-read cannot drift within a step. + recv_input_ids = getattr(self.model, "afd_requires_input_ids", False) for layer_idx in _ffn_layer_indices(self): for stage_idx in stage_ids: @@ -264,6 +268,22 @@ def _ffn_forward( num_tokens_across_dp, dp_size=int(self.vllm_config.parallel_config.data_parallel_size), ) + # A model whose router is keyed by token identity needs the ids + # of the tokens this rank computes on installed in the forward + # context before the FFN compute runs. They arrive with the + # transfer, so the receive must happen before the context is + # built rather than inside it. + payload = self.connector.recv_attn_output( + ubatch_idx=stage_idx, + layer_idx=layer_idx, + max_num_tokens=self.max_num_tokens, + recv_input_ids=recv_input_ids, + ) + context = payload.context + metadata = context.metadata + states = context.states + hidden_states = payload.hidden_states + received_input_ids = getattr(states, "input_ids", None) with ascend_forward_context( vllm_config=self.vllm_config, afd_metadata=afd_metadata, @@ -272,16 +292,8 @@ def _ffn_forward( num_tokens_across_dp=dp_num_tokens_across_dp, in_profile_run=is_profile, aclgraph_runtime_mode=aclgraph_runtime_mode, + input_ids=received_input_ids, ) as forward_context: - payload = self.connector.recv_attn_output( - ubatch_idx=stage_idx, - layer_idx=layer_idx, - max_num_tokens=self.max_num_tokens, - ) - context = payload.context - metadata = context.metadata - states = context.states - hidden_states = payload.hidden_states metadata.layer_idx = layer_idx metadata.stage_idx = stage_idx forward_context.dp_metadata = dp_metadata_list.get(stage_idx) @@ -292,6 +304,7 @@ def _ffn_forward( rank_ffn_output = self.model.compute_ffn_output( hidden_states=hidden_states, layer_idx=layer_idx, + input_ids=received_input_ids, ) _send_ffn_output( self.connector, diff --git a/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py b/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py index 2eb99cd3..89bd16e3 100644 --- a/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py +++ b/tests/unit/model_executor/models/test_deepseek_v4_attention_gate.py @@ -27,17 +27,30 @@ def test_dsv4_async_gate_bypasses_native_moe_communicator() -> None: def test_dsv4_async_gate_validates_local_hash_token_alignment() -> None: + """The Hash ids sent to FFN must be sliced like the router logits. + + The alignment logic now lives in ``local_hash_input_ids`` so that the + send-side selection and this local routing path cannot drift apart. These + assertions are structural: the helper's numerics are covered by + ``tests/unit/model_executor/test_deepseek_v4_hash_ids.py``. + """ + source = Path( "afd_plugin/model_executor/models/npu/deepseek_v4_attention_gate.py", ).read_text() - assert "DSV4 Hash routing input_ids/token count mismatch on Attention" in source - assert "input_ids = input_ids.reshape(-1).to(torch.int64)" in source - assert "forward_context.flash_comm_v1_enabled" in source - assert "and input_ids.numel() != router_logits.shape[0]" in source + assert "def local_hash_input_ids(" in source + assert "DSV4 Hash routing cannot align the ids sent to FFN" in source + assert "ids = input_ids.reshape(-1).to(torch.int64)" in source + assert "flash_comm_v1_enabled" in source + assert "ids.numel() != router_tokens" in source assert "split_tensor_along_first_dim(" in source - assert "num_partitions=tp_group.world_size" in source - assert ")[tp_group.rank_in_group]" in source + assert "num_partitions=group.world_size" in source + assert ")[group.rank_in_group]" in source + # The send-side selection must reuse this helper rather than re-deriving a + # slice; its numerics live in test_deepseek_v4_hash_ids.py. + assert "def hash_input_ids_from_context(" in source + assert "return local_hash_input_ids(" in source def test_dsv4_ffn_does_not_reapply_gate_routed_scale() -> None: diff --git a/tests/unit/model_executor/test_deepseek_v4_hash_ids.py b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py new file mode 100644 index 00000000..4c9c73f7 --- /dev/null +++ b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for the DSV4 Attention-side Hash id selection helper. + +``local_hash_input_ids`` decides which token ids Attention sends towards the FFN +role for a Hash layer. The helper is pure tensor logic, which matters here: an +ids/token misalignment would let a token-keyed router select experts for the +wrong tokens without raising anywhere, so the alignment behaviour is worth +testing without an Ascend device. +""" + +from __future__ import annotations + +import contextlib +import logging +import sys +import types +from collections.abc import Iterator + +import pytest + +torch = pytest.importorskip("torch") + +_STUB_MODULES = ( + "vllm", + "vllm.forward_context", + "vllm.logger", +) + + +@contextlib.contextmanager +def _vllm_stub() -> Iterator[None]: + """Expose a minimal ``vllm`` surface for the duration of one import. + + Importing the model package pulls in ``vllm.forward_context``. The stub is + removed again immediately afterwards so a partial ``vllm`` cannot mask the + real "vllm is not installed" failure for other test modules in the same + session. + """ + + missing = [name for name in _STUB_MODULES if name not in sys.modules] + if not missing: + yield + return + + saved = {name: sys.modules.get(name) for name in missing} + + def make(name: str, **attributes: object) -> None: + module = types.ModuleType(name) + module.__path__ = [] # type: ignore[attr-defined] + for attribute, value in attributes.items(): + setattr(module, attribute, value) + sys.modules[name] = module + + make("vllm") + make( + "vllm.forward_context", + DPMetadata=type("DPMetadata", (), {}), + ForwardContext=type("ForwardContext", (), {}), + get_forward_context=lambda: None, + ) + make( + "vllm.logger", + init_logger=lambda *args, **kwargs: logging.getLogger("test"), + ) + try: + yield + finally: + for name in missing: + sys.modules.pop(name, None) + for name, module in saved.items(): + if module is not None: + sys.modules[name] = module + + +with _vllm_stub(): + from afd_plugin.model_executor.models.npu.deepseek_v4_attention_gate import ( + hash_input_ids_from_context, + local_hash_input_ids, + ) + + +def _make_model(**attributes: object) -> types.SimpleNamespace: + return types.SimpleNamespace(**attributes) + + +def _forward_context(**overrides: object) -> types.SimpleNamespace: + """Build a forward context carrying the fields the helper reads. + + A real Ascend forward context always defines all three, so a fixture that + omitted one would exercise an ``AttributeError`` rather than a routing + condition. + """ + + fields: dict[str, object] = { + "input_ids": None, + "flash_comm_v1_enabled": False, + "pad_size": 0, + } + fields.update(overrides) + return types.SimpleNamespace(**fields) + + +def test_returns_global_ids_when_no_flash_comm_split_is_needed(): + ids = torch.tensor([5, 6, 7], dtype=torch.int32) + + result = local_hash_input_ids( + input_ids=ids, + router_tokens=3, + flash_comm_v1_enabled=False, + pad_size=0, + ) + + assert result.dtype == torch.int64 + assert result.tolist() == [5, 6, 7] + + +def test_flattens_multi_dimensional_ids(): + ids = torch.tensor([[5, 6], [7, 8]], dtype=torch.int64) + + result = local_hash_input_ids( + input_ids=ids, + router_tokens=4, + flash_comm_v1_enabled=False, + pad_size=0, + ) + + assert result.tolist() == [5, 6, 7, 8] + + +def test_rejects_missing_ids(): + with pytest.raises(RuntimeError, match="requires input_ids to send"): + local_hash_input_ids( + input_ids=None, + router_tokens=3, + flash_comm_v1_enabled=False, + pad_size=0, + ) + + +def test_rejects_unalignable_token_count(): + """A count mismatch must fail here, before any cross-role transfer.""" + ids = torch.tensor([5, 6], dtype=torch.int64) + + with pytest.raises(RuntimeError, match="cannot align the ids sent to FFN"): + local_hash_input_ids( + input_ids=ids, + router_tokens=3, + flash_comm_v1_enabled=False, + pad_size=0, + ) + + +def test_applies_flash_comm_padding_and_tp_split(monkeypatch): + """FlashComm v1 must slice ids exactly like the router logits it mirrors.""" + captured: dict[str, object] = {} + + def fake_split(tensor: torch.Tensor, *, num_partitions: int, **kwargs: object): + captured["num_partitions"] = num_partitions + captured["kwargs"] = kwargs + return list(torch.chunk(tensor, num_partitions)) + + distributed = types.ModuleType("vllm.distributed") + distributed.get_tp_group = lambda: _make_model(world_size=2, rank_in_group=1) + ascend_distributed = types.ModuleType("vllm_ascend.distributed") + ascend_utils = types.ModuleType("vllm_ascend.distributed.utils") + ascend_utils.split_tensor_along_first_dim = fake_split + ascend_distributed.utils = ascend_utils + + monkeypatch.setitem(sys.modules, "vllm.distributed", distributed) + monkeypatch.setitem(sys.modules, "vllm_ascend.distributed", ascend_distributed) + monkeypatch.setitem(sys.modules, "vllm_ascend.distributed.utils", ascend_utils) + + # 4 global ids plus one padded slot gives 5. torch.chunk(5, 2) splits that + # into [3, 2] rather than evenly, so rank 1 receives the trailing pair. + result = local_hash_input_ids( + input_ids=torch.tensor([5, 6, 7, 8], dtype=torch.int64), + router_tokens=2, + flash_comm_v1_enabled=True, + pad_size=1, + ) + + assert captured["num_partitions"] == 2 + assert captured["kwargs"] == {"contiguous_split_chunks": True} + assert result.tolist() == [8, 0] + + +def test_context_without_ids_is_an_error_not_a_quiet_fallback(): + """A context with no ids must fail here, not downgrade the transfer. + + The FFN role decides once per run that it expects ids and cannot tell Hash + layers from non-Hash ones, so "this layer does not need them" is not a case + that can be answered by sending activations alone. + """ + + with pytest.raises(RuntimeError, match="requires input_ids to send"): + hash_input_ids_from_context( + forward_context=_forward_context(input_ids=None), + router_tokens=3, + ) + + +def test_context_ids_are_returned_as_router_aligned_tokens(): + result = hash_input_ids_from_context( + forward_context=_forward_context( + input_ids=torch.tensor([5, 6, 7], dtype=torch.int32), + ), + router_tokens=3, + ) + + assert result.tolist() == [5, 6, 7] + + +def test_context_ids_are_still_validated_against_the_local_token_count(): + """Carrying ids opts into validation; a mismatch must not pass silently.""" + with pytest.raises(RuntimeError, match="cannot align the ids sent to FFN"): + hash_input_ids_from_context( + forward_context=_forward_context( + input_ids=torch.tensor([5, 6], dtype=torch.int32), + ), + router_tokens=3, + ) diff --git a/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py b/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py new file mode 100644 index 00000000..8a73e7a1 --- /dev/null +++ b/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py @@ -0,0 +1,217 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the AFD plugin project +"""CPU tests for the DSV4 NPU role-aware checkpoint filter. + +``_checkpoint_weight_roles`` decides which AFD role receives each checkpoint +path. The distinction is not cosmetic: the upstream Ascend loader indexes its +parameter dict by name without a membership check, so a path handed to a role +that never registered it raises ``KeyError`` during weight loading instead of +being skipped. + +The Hash id table is the case that matters. It is a parameter only where the +Hash MoE is built, which is the FFN role, so it must not be handed to Attention. + +How the functions are loaded +---------------------------- +The NPU DSV4 module defines these helpers next to its model classes, and +importing it pulls in the whole ``vllm``, ``vllm_ascend`` and ``transformers`` +class surface. Stubbing that surface would make this file break on every +upstream addition and would test the stubs as much as the code. + +Instead the two helpers are read from the module source and executed in an +isolated namespace. They are module-level functions that reference nothing but +each other and three role constants, so this exercises the real code from the +real file without importing it. +""" + +from __future__ import annotations + +import ast +from pathlib import Path +from types import ModuleType + +import pytest + +pytest.importorskip("torch") + +_MODULE_PATH = ( + Path(__file__).resolve().parents[3] + / "afd_plugin" + / "model_executor" + / "models" + / "npu" + / "deepseek_v4.py" +) + +_HELPER_NAMES = ( + "_weight_layer_path", + "_checkpoint_weight_roles", +) + + +def _load_helpers() -> ModuleType: + """Execute the role-filter helpers in an isolated namespace. + + Returns: + A module-like namespace holding ``_checkpoint_weight_roles``. + + Raises: + AssertionError: If the module no longer defines the helpers, which means + this test needs to be repointed rather than silently pass. + """ + + source = _MODULE_PATH.read_text(encoding="utf-8") + tree = ast.parse(source) + + namespace: dict[str, object] = { + "frozenset": frozenset, + "tuple": tuple, + "int": int, + "str": str, + "None": None, + "_ATTENTION_ROLE": "attention", + "_FFN_ROLE": "ffn", + "_BOTH_ROLES": frozenset(("attention", "ffn")), + } + + # Module-level constants the helpers read. Evaluating them from the source + # keeps the tests from drifting when a constant changes. + for node in tree.body: + if not isinstance(node, ast.Assign): + continue + for target in node.targets: + name = getattr(target, "id", None) + if name and name.startswith("_HASH_"): + namespace[name] = eval( # noqa: S307 - this repository's own source + compile(ast.Expression(body=node.value), "", "eval"), + namespace, + ) + + found: set[str] = set() + for node in tree.body: + if not isinstance(node, ast.FunctionDef) or node.name not in _HELPER_NAMES: + continue + # ``from __future__ import annotations`` makes the annotations lazy, so + # the definitions can be executed without their referenced types. + code = compile(ast.Module(body=[node], type_ignores=[]), "", "exec") + exec(code, namespace) # noqa: S102 - source is this repository's own file + found.add(node.name) + + assert found == set(_HELPER_NAMES), ( + f"{_MODULE_PATH} no longer defines {sorted(_HELPER_NAMES)}; " + f"found {sorted(found)}" + ) + + module = ModuleType("dsv4_role_helpers") + module.__dict__.update(namespace) + return module + + +_helpers = _load_helpers() +_checkpoint_weight_roles = _helpers._checkpoint_weight_roles # type: ignore[attr-defined] + + +def test_module_and_helpers_are_present() -> None: + assert _MODULE_PATH.is_file() + + +@pytest.mark.parametrize( + "name", + [ + "model.layers.0.mlp.gate.tid2eid", + "model.layers.7.ffn.gate.tid2eid", + ], +) +def test_hash_id_table_is_ffn_owned(name: str) -> None: + """Only the Hash MoE registers this parameter, so Attention must not see it. + + Handing it to Attention is exactly the mismatch that raises + ``KeyError: 'model.layers.0.mlp.gate.tid2eid'`` while loading that rank, + under either gate placement. + """ + + for attn_owns_gate in (True, False): + assert _checkpoint_weight_roles( + name, + attn_owns_gate=attn_owns_gate, + ) == frozenset({"ffn"}) + + +@pytest.mark.parametrize( + "name", + [ + "model.layers.0.mlp.gate.weight", + "model.layers.3.ffn.gate.weight", + "model.layers.3.mlp.gate.e_score_correction_bias", + ], +) +def test_gate_paths_skip_attention_when_the_gate_is_on_ffn(name: str) -> None: + """With the gate on FFN the Attention MoE slot is parameter-free. + + This is the supported CAMP2P configuration, and handing Attention these paths + raises ``KeyError: 'model.layers.0.mlp.gate.weight'`` because it registered + no gate at all. + """ + + assert _checkpoint_weight_roles( + name, + attn_owns_gate=False, + ) == frozenset({"ffn"}) + + +@pytest.mark.parametrize( + "name", + [ + "model.layers.0.mlp.gate.weight", + "model.layers.3.mlp.gate.e_score_correction_bias", + ], +) +def test_gate_paths_stay_shared_when_attention_owns_the_gate(name: str) -> None: + """Gate-on-Attention builds a router there, so both roles load it.""" + + assert _checkpoint_weight_roles( + name, + attn_owns_gate=True, + ) == frozenset({"attention", "ffn"}) + + +@pytest.mark.parametrize( + "name", + [ + "model.layers.0.attn.q_a_proj.weight", + "model.layers.0.self_attn.o_proj.weight", + ], +) +def test_attention_paths_are_attention_owned(name: str) -> None: + assert _checkpoint_weight_roles(name) == frozenset({"attention"}) + + +@pytest.mark.parametrize( + "name", + [ + "model.layers.0.mlp.experts.0.gate_proj.weight", + "model.layers.0.ffn.shared_experts.gate_proj.weight", + ], +) +def test_expert_paths_are_ffn_owned(name: str) -> None: + for attn_owns_gate in (True, False): + assert _checkpoint_weight_roles( + name, + attn_owns_gate=attn_owns_gate, + ) == frozenset({"ffn"}) + + +@pytest.mark.parametrize( + "name", + [ + "model.embed_tokens.weight", + "model.layers.0.hc_attn_fn", + "lm_head.weight", + ], +) +def test_shared_and_non_layer_paths_are_shared(name: str) -> None: + for attn_owns_gate in (True, False): + assert _checkpoint_weight_roles( + name, + attn_owns_gate=attn_owns_gate, + ) == frozenset({"attention", "ffn"}) From ca014f74d861aef782b5c74a902d3f842c9951f0 Mon Sep 17 00:00:00 2001 From: ksiyuan Date: Wed, 16 Sep 2026 10:05:51 +0800 Subject: [PATCH 3/5] fix(npu): keep the DSV4 Hash id table with its gate owner Review follow-up on the CAMP2P token-id channel. Four defects, all found after the head was pushed and none of them caught by CI, which has never run on this branch. Role filtering: _checkpoint_weight_roles pinned every gate.tid2eid path to FFN. That breaks the configuration that already worked: with the gate on Attention, AFDDeepseekV4AttentionGateRemoteMoE registers gate.tid2eid for Hash layers and _compute_sqrtsoftplus_topk routes from it, so the Attention rank has to load the table. Pinned to FFN it kept the zeros initialiser and every Hash token routed to expert 0 without raising, because the upstream loader has no completeness check. The Hash table now follows attn_owns_gate like the other gate paths. Dead ids argument: compute_ffn_output passed input_ids into the native MoE forward, which neither takes nor uses such an argument; vLLM-Ascend's fused-expert selector reads forward_context.input_ids and overrides its own parameter with it. Drop the argument and the patch marker so the ambient forward context stays the single channel. Payload contract: the received ids were stored on CAMP2PTransferState, but model-specific tensors are payload fields rather than transfer state (connector_contracts.md), and the CUDA P2P path already returns them on AFDA2FTransferPayload. recv_attn_output now returns them there, the FFN runner reads payload.input_ids, and both added state fields are gone. Stale tests: test_dsv4_rejects_camp2p_connector asserted the rejection this change removes, and the NPU FFN runner test asserted the kwarg the dead argument added. Replace the first with the connector that is still rejected, and drop the argument so the second holds again. New CPU coverage: send-side compute_gate selection, receive-side mode and payload contract, and AFDDeepseekV4RemoteMoE, which had no test reference at all. The ids-mode connector cases stay NPU-gated. Validation: - ruff 0.15.13 check and format: clean on all changed files. - mypy 1.11.1 with --python-version 3.10 (CI hook arguments): clean for both the library and the tests group. - pytest on the four affected CPU files: 35 passed. - The NPU-gated cases (test_camp2p_connector.py, test_npu_runtime.py) cannot run on this host; they were reviewed against the new signatures. Signed-off-by: ksiyuan --- afd_plugin/compat/npu/feature_validation.py | 2 +- afd_plugin/connectors/npu/camp2p.py | 33 ++- .../model_executor/models/npu/deepseek_v4.py | 28 +-- afd_plugin/v1/worker/npu/ffn_model_runner.py | 9 +- .../compat/npu/test_dsv4_async_validation.py | 17 +- .../unit/connectors/test_camp2p_connector.py | 8 +- .../unit/connectors/test_camp2p_token_ids.py | 201 +++++++++++++++++- .../test_deepseek_v4_hash_ids.py | 96 ++++++++- .../test_deepseek_v4_npu_weight_roles.py | 42 ++-- 9 files changed, 358 insertions(+), 78 deletions(-) diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index ec6f3214..3a58a3d4 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -53,7 +53,7 @@ def fail_if_unsupported_npu_afd_features( raise RuntimeError( "AFD NPU runtime does not support compute_gate_on_attention=true yet", ) - if afd_config.connector == "CAMP2pAFDConnector": + if afd_config.connector == CAMP2P_CONNECTOR: from afd_plugin.connectors.npu.camp2p import CAMP2PExtraInfo if not isinstance(extra_info, CAMP2PExtraInfo): diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index b5b09bce..b862e965 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -178,13 +178,6 @@ class CAMP2PTransferState(AFDTransferState): A2E-returned Attention token count that the FFN-to-Attention send requires. ``x_active_mask`` and ``cam_p2p_ep_name`` are the A2E-returned active-token mask and HCCL endpoint name captured on the receive path. - - ``input_ids`` holds the token-aligned ids that Attention sent alongside the - hidden states, as received by the FFN rank. It is populated only when the - receiving rank declared ``recv_input_ids``, which ``compute_gate_mode`` - records. ``compute_gate_mode`` is the operator's ids mode for this transfer - and has to equal the mode the sending rank selected, because the operator - only writes the ids slot in that mode. """ aiv_num: int = 8 @@ -194,8 +187,6 @@ class CAMP2PTransferState(AFDTransferState): atten_batch_size: torch.Tensor | None = None x_active_mask: torch.Tensor | None = None cam_p2p_ep_name: str | None = None - input_ids: torch.Tensor | None = None - compute_gate_mode: int = 0 @dataclass(frozen=True, slots=True) @@ -525,7 +516,8 @@ def send_attn_output( **kwargs: An optional token-aligned ``input_ids`` tensor. When it is supplied, the transfer runs with ``compute_gate=1`` so the ids reach the FFN rank through the operator's ids channel, and the - local FFN-side receive exposes them on ``CAMP2PTransferState``. + matching ``recv_attn_output(recv_input_ids=True)`` returns them + on the payload's ``input_ids`` field. Raises: RuntimeError: If the communication groups are not ready. @@ -639,13 +631,14 @@ def recv_attn_output( number needed to create it. ``recv_input_ids`` states that this FFN rank expects token ids on the transfer, which selects the operator's ids mode and makes the connector validate and expose - the operator's ids slot. The sending rank has to select the same - mode, so only request ids for a run whose Attention role - transports them. + the operator's ids slot on the payload. The sending rank has to + select the same mode, so only request ids for a run whose + Attention role transports them. Returns: - The received hidden states and the information FFN needs to process - them and send the result back. + The received hidden states, the information FFN needs to process them + and send the result back, and the transported ``input_ids`` when the + ids mode was selected. Raises: RuntimeError: If communication is not ready, transfer information @@ -658,6 +651,7 @@ def recv_attn_output( layer_idx: int = kwargs.get("layer_idx", 0) max_num_tokens: int = kwargs.get("max_num_tokens", 0) recv_input_ids: bool = bool(kwargs.get("recv_input_ids", False)) + compute_gate_mode = 1 if recv_input_ids else 0 batch_size = _num_tokens_for_ffn_rank( self.dp_metadata_list, ubatch_idx, @@ -676,7 +670,6 @@ def recv_attn_output( batch_size=batch_size, h=self.hidden_size, k=self.num_experts_per_tok, - compute_gate_mode=1 if recv_input_ids else 0, ) context = AFDTransferContext( metadata=metadata, @@ -701,7 +694,7 @@ def recv_attn_output( self.world_rank, group_ep, custom_states.aiv_num, - custom_states.compute_gate_mode, + compute_gate_mode, ) custom_states.atten_batch_size = outputs[3] custom_states.x_active_mask = outputs[4] @@ -713,14 +706,16 @@ def recv_attn_output( # indistinguishable from the operator's placeholder. Reading the slot in # the other mode would hand the model uninitialised device memory as token # ids, which a token-keyed router turns into an out-of-range table read. - if custom_states.compute_gate_mode == 1: - custom_states.input_ids = received_token_ids( + received_ids: torch.Tensor | None = None + if compute_gate_mode == 1: + received_ids = received_token_ids( outputs[1], expected_tokens=batch_size, ) return AFDA2FTransferPayload( hidden_states=outputs[0], context=context, + input_ids=received_ids, ) def send_ffn_output( diff --git a/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index 42da9aa6..81e1771c 100644 --- a/afd_plugin/model_executor/models/npu/deepseek_v4.py +++ b/afd_plugin/model_executor/models/npu/deepseek_v4.py @@ -96,15 +96,12 @@ def _checkpoint_weight_roles( return frozenset((_ATTENTION_ROLE,)) if stage in ("ffn", "mlp"): if remainder and remainder[0] == "gate": - # The Hash id table is a parameter only where the Hash MoE is - # built, which is the FFN role, whatever the gate placement: with - # the gate on FFN the table lives on the FFN router, and with the - # gate on Attention the Hash path routes from the table instead of - # from a gate weight. - if "tid2eid" in remainder: - return frozenset((_FFN_ROLE,)) - # The remaining gate parameters belong to every role that built a - # router. With the gate on FFN, Attention has none. + # Every router parameter, including the Hash id table, belongs to + # each role that built a router: with the gate on Attention, the + # Attention gate shell registers ``tid2eid`` for Hash layers and + # routes from it, and the native FFN MoE registers its own copy. + # 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,)) @@ -331,7 +328,6 @@ def compute_ffn_output( dynamic_scales_shared: torch.Tensor | None = None, topk_scales: torch.Tensor | None = None, group_list_type: int = 1, - input_ids: torch.Tensor | None = None, **_: Any, ) -> torch.Tensor | AFDF2ATransferPayload: if not isinstance(self.mlp, native.DeepseekV4MoE): @@ -360,14 +356,12 @@ def compute_ffn_output( # topk_weights, which CAM applies during combine-recv. routed_scale_applied_in_topk=True, ) - # ### PATCH START: FFN-side Hash routing needs the transported ids. # The native MoE runs the gate internally when the gate is not on - # Attention, and its Hash layers route by token identity: vLLM-Ascend's - # FusedMoE reads `forward_context.input_ids`, which AFD installs from the - # transfer. Pass them on to the native MoE as well so the ids travel with - # the call rather than only through ambient context. - return self.mlp(hidden_states, input_ids=input_ids) - # ### PATCH END: FFN-side Hash routing needs the transported ids. + # Attention. Its Hash layers route by token identity, and vLLM-Ascend's + # fused-expert selector reads ``forward_context.input_ids``, which the + # FFN runner installs from the transfer. The native forward takes no + # ``input_ids`` argument, so the ambient context is the whole channel. + return self.mlp(hidden_states) @native.support_torch_compile diff --git a/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 9289b6ab..5d60dd71 100644 --- a/afd_plugin/v1/worker/npu/ffn_model_runner.py +++ b/afd_plugin/v1/worker/npu/ffn_model_runner.py @@ -270,9 +270,9 @@ def _ffn_forward( ) # A model whose router is keyed by token identity needs the ids # of the tokens this rank computes on installed in the forward - # context before the FFN compute runs. They arrive with the - # transfer, so the receive must happen before the context is - # built rather than inside it. + # context before the FFN compute runs. They arrive on the + # transfer payload, so the receive must happen before the + # context is built rather than inside it. payload = self.connector.recv_attn_output( ubatch_idx=stage_idx, layer_idx=layer_idx, @@ -283,7 +283,7 @@ def _ffn_forward( metadata = context.metadata states = context.states hidden_states = payload.hidden_states - received_input_ids = getattr(states, "input_ids", None) + received_input_ids = payload.input_ids with ascend_forward_context( vllm_config=self.vllm_config, afd_metadata=afd_metadata, @@ -304,7 +304,6 @@ def _ffn_forward( rank_ffn_output = self.model.compute_ffn_output( hidden_states=hidden_states, layer_idx=layer_idx, - input_ids=received_input_ids, ) _send_ffn_output( self.connector, diff --git a/tests/unit/compat/npu/test_dsv4_async_validation.py b/tests/unit/compat/npu/test_dsv4_async_validation.py index 1d74027c..5e71770b 100644 --- a/tests/unit/compat/npu/test_dsv4_async_validation.py +++ b/tests/unit/compat/npu/test_dsv4_async_validation.py @@ -24,16 +24,27 @@ def _afd_config( ) -def test_dsv4_rejects_camp2p_connector() -> None: - with pytest.raises(RuntimeError, match="only CAMAsyncAFDConnector"): +def test_dsv4_rejects_connectors_without_a_token_id_channel() -> None: + with pytest.raises(RuntimeError, match="supports only CAMAsyncAFDConnector"): _fail_if_unsupported_dsv4_connector( _afd_config( compute_gate_on_attention=False, - connector="CAMP2pAFDConnector", + connector="P2pNcclAFDConnector", ), ) +def test_dsv4_accepts_the_camp2p_connector() -> None: + """CAMP2P carries the Hash ids over the A2E ids channel with the gate on FFN.""" + + _fail_if_unsupported_dsv4_connector( + _afd_config( + compute_gate_on_attention=False, + connector="CAMP2pAFDConnector", + ), + ) + + def test_dsv4_async_requires_attention_side_gate() -> None: with pytest.raises(RuntimeError, match="compute_gate_on_attention"): _fail_if_unsupported_dsv4_async_features( diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 86e5239c..99b8c008 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -155,7 +155,9 @@ def test_camp2p_recv_attn_output_drives_the_operator_ids_mode(monkeypatch): The ``a2e`` operator only writes its ids slot in the ids mode, and the sending rank selects that mode independently. A receiving rank that reads the slot in the other mode would install uninitialised device memory as token - ids, which a token-keyed router turns into an out-of-range table read. + ids, which a token-keyed router turns into an out-of-range table read. The + ids that arrive are model-specific tensors, so they travel on the payload + rather than in the backend transfer state. """ torch = pytest.importorskip("torch") @@ -183,9 +185,9 @@ def fake_a2e(*args): ) assert calls[0][-1] == 1 - assert with_ids.context.states.input_ids.tolist() == [0, 2, 4, 6, 8] + assert with_ids.input_ids.tolist() == [0, 2, 4, 6, 8] assert calls[1][-1] == 0 - assert without_ids.context.states.input_ids is None + assert without_ids.input_ids is None def test_camp2p_extra_info_rejects_unknown_mix_placement(): diff --git a/tests/unit/connectors/test_camp2p_token_ids.py b/tests/unit/connectors/test_camp2p_token_ids.py index 845b991b..1ecfae17 100644 --- a/tests/unit/connectors/test_camp2p_token_ids.py +++ b/tests/unit/connectors/test_camp2p_token_ids.py @@ -10,6 +10,12 @@ tests: importing any connector module pulls in ``vllm`` through ``afd_plugin.connectors``, so the import is satisfied with the minimal surface the helpers' module needs. + +The mode tests at the end drive the real ``send_attn_output`` and +``recv_attn_output`` against recorded operator calls. The operator's +``compute_gate`` mode is selected independently on each side and the operator +only writes its ids slot in that mode, so both sides deriving it from the same +run-level decision is the invariant worth pinning here. """ from __future__ import annotations @@ -19,6 +25,8 @@ import sys import types from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any import pytest @@ -93,12 +101,93 @@ def make(name: str, **attributes: object) -> None: with _vllm_stub(): - from afd_plugin.connectors.npu.camp2p import ( # noqa: E402 + from afd_plugin.config import AFDConfig + from afd_plugin.connectors.metadata import ( + AFDTransferContext, + AFDTransferMetadata, + ) + from afd_plugin.connectors.npu import camp2p as camp2p_module + from afd_plugin.connectors.npu.camp2p import ( + CAMP2pAFDConnector, prepare_token_id_transfer, received_token_ids, ) +class _CpuTorch: + """``torch`` with the NPU device mapped to CPU for host-side tests. + + The connector hands the operator empty device tensors, which only exist on an + Ascend host. Only ``tensor()`` is adapted; every other attribute, including + ``ops``, stays the real one so the recorded operator calls are the calls the + connector makes. + """ + + def __getattr__(self, name: str) -> Any: + return getattr(torch, name) + + def tensor(self, *args: Any, **kwargs: Any) -> Any: + if kwargs.get("device") == "npu": + kwargs["device"] = "cpu" + return torch.tensor(*args, **kwargs) + + +class _FakeDPMetadata: + def __init__(self, values: list[int]) -> None: + # The connector counts tokens with .flatten().tolist(), so this has to be + # a tensor like the real DP metadata rather than a plain list. + self.num_tokens_across_dp_cpu = torch.tensor(values, dtype=torch.int32) + + +def _vllm_config() -> SimpleNamespace: + return SimpleNamespace( + additional_config={"afd": {"connector_extra_config": {}}}, + parallel_config=SimpleNamespace( + data_parallel_size=1, + data_parallel_rank=0, + prefill_context_parallel_size=1, + tensor_parallel_size=1, + num_ubatches=1, + ), + scheduler_config=SimpleNamespace(max_num_seqs=8), + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + hidden_size=16, + num_experts_per_tok=2, + n_routed_experts=4, + n_shared_experts=0, + ), + ), + ) + + +def _afd_config(*, role: str) -> AFDConfig: + return AFDConfig( + connector="CAMP2pAFDConnector", + role=role, + num_attention_ranks=4, + num_ffn_ranks=2, + ) + + +def _connector(*, role: str, rank: int) -> CAMP2pAFDConnector: + """Build a connector whose communication groups are not needed yet.""" + + connector = CAMP2pAFDConnector( + rank, + rank, + _vllm_config(), + _afd_config(role=role), + rank, + ) + connector._initialized = True + connector.hccl_comm_name = "hccl0" + connector.hccl_comm_name2 = "hccl1" + connector.hccl_comm_name3 = "" + connector.hccl_comm_name1 = "moe" + return connector + + def test_prepare_token_id_transfer_replicates_ids_across_columns(): input_ids = torch.tensor([7, 11, 13], dtype=torch.int64) @@ -161,3 +250,113 @@ def test_received_token_ids_rejects_alignment_shortfall(): with pytest.raises(ValueError, match="not aligned with the FFN token layout"): received_token_ids(sent_ids, expected_tokens=5) + + +def test_send_attn_output_selects_the_operator_ids_mode(monkeypatch): + """Sending ids must raise ``compute_gate`` and fill the operator's id slot. + + The receiving rank selects the same mode from its own declaration, so the + sending rank has to derive it from the presence of ids alone: an + activations-only send must leave the slot untouched. + """ + + calls: list[tuple[Any, ...]] = [] + monkeypatch.setattr( + torch.ops.vllm, + "afd_camp2p_send_attn_output", + lambda *args: calls.append(args), + raising=False, + ) + forward_context = SimpleNamespace() + monkeypatch.setattr( + camp2p_module, + "get_forward_context", + lambda: forward_context, + ) + connector = _connector(role="attention", rank=0) + hidden_states = torch.zeros(3, connector.hidden_size) + context = AFDTransferContext( + metadata=AFDTransferMetadata.create_attention_metadata( + layer_idx=0, + stage_idx=0, + seq_len=3, + ), + ) + + connector.send_attn_output( + hidden_states, + context, + input_ids=torch.tensor([7, 11, 13], dtype=torch.int64), + ) + connector.send_attn_output(hidden_states, context) + + with_ids, without_ids = calls + # Trailing operator arguments: aiv_num, compute_gate, ids, scales. + assert with_ids[-3] == 1 + assert with_ids[-2].dtype == torch.int32 + assert with_ids[-2][:, 0].tolist() == [7, 11, 13] + assert with_ids[-1].dtype == torch.float32 + assert torch.count_nonzero(with_ids[-1]) == 0 + assert without_ids[-3] == 0 + assert without_ids[-2] is None + assert without_ids[-1] is None + + +def test_recv_attn_output_mode_and_ids_follow_the_receiver_declaration(monkeypatch): + """The receiving rank declares the mode and gets ids only in that mode. + + ``recv_input_ids`` is the FFN side's half of the run-level decision: it + selects the operator's mode and decides whether the id slot may be read. The + ids are model-specific tensors, so they travel on the payload rather than in + the backend transfer state. + """ + + calls: list[tuple[Any, ...]] = [] + + def fake_a2e(*args: Any) -> tuple[Any, ...]: + calls.append(args) + tokens, topk = int(args[3]), int(args[5]) + ids = ( + torch.arange(tokens, dtype=torch.int32) + .mul(10) + .unsqueeze(1) + .expand(tokens, topk) + .contiguous() + ) + return ("hidden", ids, None, "atten-batch", "active-mask") + + monkeypatch.setattr(torch.ops.afd_ascend, "a2e", fake_a2e, raising=False) + monkeypatch.setattr(camp2p_module, "torch", _CpuTorch()) + connector = _connector(role="ffn", rank=1) + # FFN rank 1 owns attention ranks 2 and 3, so it computes on 5 + 7 tokens. + connector.dp_metadata_list = {0: _FakeDPMetadata([2, 3, 5, 7])} + + with_ids = connector.recv_attn_output( + ubatch_idx=0, + layer_idx=0, + recv_input_ids=True, + ) + without_ids = connector.recv_attn_output( + ubatch_idx=0, + layer_idx=0, + recv_input_ids=False, + ) + + assert calls[0][-1] == 1 + assert with_ids.input_ids is not None + assert with_ids.input_ids.tolist() == [ + 0, + 10, + 20, + 30, + 40, + 50, + 60, + 70, + 80, + 90, + 100, + 110, + ] + assert calls[1][-1] == 0 + assert without_ids.input_ids is None diff --git a/tests/unit/model_executor/test_deepseek_v4_hash_ids.py b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py index 4c9c73f7..46d5f88d 100644 --- a/tests/unit/model_executor/test_deepseek_v4_hash_ids.py +++ b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py @@ -11,11 +11,14 @@ from __future__ import annotations +import ast import contextlib import logging import sys import types from collections.abc import Iterator +from pathlib import Path +from typing import Any import pytest @@ -155,17 +158,20 @@ def test_applies_flash_comm_padding_and_tp_split(monkeypatch): """FlashComm v1 must slice ids exactly like the router logits it mirrors.""" captured: dict[str, object] = {} - def fake_split(tensor: torch.Tensor, *, num_partitions: int, **kwargs: object): + def fake_split(tensor: Any, *, num_partitions: int, **kwargs: object): captured["num_partitions"] = num_partitions captured["kwargs"] = kwargs return list(torch.chunk(tensor, num_partitions)) distributed = types.ModuleType("vllm.distributed") - distributed.get_tp_group = lambda: _make_model(world_size=2, rank_in_group=1) + distributed.get_tp_group = lambda: _make_model( # type: ignore[attr-defined] + world_size=2, + rank_in_group=1, + ) ascend_distributed = types.ModuleType("vllm_ascend.distributed") ascend_utils = types.ModuleType("vllm_ascend.distributed.utils") - ascend_utils.split_tensor_along_first_dim = fake_split - ascend_distributed.utils = ascend_utils + ascend_utils.split_tensor_along_first_dim = fake_split # type: ignore[attr-defined] + ascend_distributed.utils = ascend_utils # type: ignore[attr-defined] monkeypatch.setitem(sys.modules, "vllm.distributed", distributed) monkeypatch.setitem(sys.modules, "vllm_ascend.distributed", ascend_distributed) @@ -220,3 +226,85 @@ def test_context_ids_are_still_validated_against_the_local_token_count(): ), router_tokens=3, ) + + +_DSV4_MODULE_PATH = ( + Path(__file__).resolve().parents[3] + / "afd_plugin" + / "model_executor" + / "models" + / "npu" + / "deepseek_v4.py" +) +_REMOTE_MOE_CLASS = "AFDDeepseekV4RemoteMoE" + + +class _RecordingProxy: + """Stand-in base class recording the transfer the shell requests.""" + + def __init__(self) -> None: + self.sent: list[dict[str, Any]] = [] + + def _send_and_receive(self, hidden_states: Any, **send_kwargs: Any) -> str: + self.sent.append({"hidden_states": hidden_states, **send_kwargs}) + return "ffn-output" + + +def _load_remote_moe_class(forward_context: Any) -> Any: + """Execute the real ``AFDDeepseekV4RemoteMoE`` from the module source. + + Importing the DSV4 NPU module pulls in the whole ``vllm``/``vllm_ascend`` + model surface, so the class body is executed against a parameter-free base + class instead. The body under test is read from the file unchanged. + + Raises: + AssertionError: If the module no longer defines the class, which means + this test needs to be repointed rather than silently pass. + """ + + tree = ast.parse(_DSV4_MODULE_PATH.read_text(encoding="utf-8")) + node = next( + ( + item + for item in tree.body + if isinstance(item, ast.ClassDef) and item.name == _REMOTE_MOE_CLASS + ), + None, + ) + assert node is not None, ( + f"{_DSV4_MODULE_PATH} no longer defines {_REMOTE_MOE_CLASS}" + ) + + namespace: dict[str, Any] = { + "RemoteFFNProxy": _RecordingProxy, + "torch": torch, + "get_forward_context": lambda: forward_context, + } + code = compile(ast.Module(body=[node], type_ignores=[]), "", "exec") + exec(code, namespace) # noqa: S102 - source is this repository's own file + return namespace[_REMOTE_MOE_CLASS] + + +def test_remote_moe_sends_the_context_ids_alongside_the_activations(): + """The gate-on-FFN shell must transport ids, since FFN cannot route without them.""" + + forward_context = _forward_context( + input_ids=torch.tensor([5, 6, 7], dtype=torch.int32), + ) + layer = _load_remote_moe_class(forward_context)() + + output = layer.forward(torch.zeros(3, 8)) + + assert output == "ffn-output" + assert layer.sent[0]["input_ids"].tolist() == [5, 6, 7] + + +def test_remote_moe_without_context_ids_raises_instead_of_sending_activations_only(): + """A quiet activations-only fallback would leave FFN reading an unwritten slot.""" + + layer = _load_remote_moe_class(_forward_context(input_ids=None))() + + with pytest.raises(RuntimeError, match="requires input_ids to send"): + layer.forward(torch.zeros(3, 8)) + + assert layer.sent == [] diff --git a/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py b/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py index 8a73e7a1..7130916d 100644 --- a/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py +++ b/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py @@ -8,8 +8,9 @@ that never registered it raises ``KeyError`` during weight loading instead of being skipped. -The Hash id table is the case that matters. It is a parameter only where the -Hash MoE is built, which is the FFN role, so it must not be handed to Attention. +The Hash id table is the case that matters. Both roles register their own copy +whenever they build a router, so the table follows the same ownership rule as +the other gate paths rather than being pinned to FFN. How the functions are loaded ---------------------------- @@ -74,19 +75,6 @@ def _load_helpers() -> ModuleType: "_BOTH_ROLES": frozenset(("attention", "ffn")), } - # Module-level constants the helpers read. Evaluating them from the source - # keeps the tests from drifting when a constant changes. - for node in tree.body: - if not isinstance(node, ast.Assign): - continue - for target in node.targets: - name = getattr(target, "id", None) - if name and name.startswith("_HASH_"): - namespace[name] = eval( # noqa: S307 - this repository's own source - compile(ast.Expression(body=node.value), "", "eval"), - namespace, - ) - found: set[str] = set() for node in tree.body: if not isinstance(node, ast.FunctionDef) or node.name not in _HELPER_NAMES: @@ -122,19 +110,23 @@ def test_module_and_helpers_are_present() -> None: "model.layers.7.ffn.gate.tid2eid", ], ) -def test_hash_id_table_is_ffn_owned(name: str) -> None: - """Only the Hash MoE registers this parameter, so Attention must not see it. +def test_hash_id_table_follows_the_gate_ownership(name: str) -> None: + """The Hash table belongs to every role that built a router. - Handing it to Attention is exactly the mismatch that raises - ``KeyError: 'model.layers.0.mlp.gate.tid2eid'`` while loading that rank, - under either gate placement. + With the gate on Attention, the Attention gate shell registers ``tid2eid`` + for its Hash layers and routes from it, so the table is shared. With the gate + on FFN the Attention MoE slot is parameter-free, so the path is FFN-only; + handing it to Attention raises ``KeyError`` from the upstream loader. """ - for attn_owns_gate in (True, False): - assert _checkpoint_weight_roles( - name, - attn_owns_gate=attn_owns_gate, - ) == frozenset({"ffn"}) + assert _checkpoint_weight_roles( + name, + attn_owns_gate=True, + ) == frozenset({"attention", "ffn"}) + assert _checkpoint_weight_roles( + name, + attn_owns_gate=False, + ) == frozenset({"ffn"}) @pytest.mark.parametrize( From 6cf0137dc448440cce251dcde9749b827695b3d6 Mon Sep 17 00:00:00 2001 From: ksiyuan Date: Wed, 16 Sep 2026 10:05:55 +0800 Subject: [PATCH 4/5] docs(npu): record the DSV4 Ascend Hash-id boundary Add the boundary to the model-integration evidence table with its focused validation, note that it has a manual Ascend 950 run rather than an accuracy comparison, and extend the connector contract to state how CAMP2P transports token ids and why both roles must select the same operator mode. Signed-off-by: ksiyuan --- docs/design/module/connector_contracts.md | 6 +++++- docs/design/module/model_integration.md | 2 ++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/docs/design/module/connector_contracts.md b/docs/design/module/connector_contracts.md index 0750d3f7..f9a93ad9 100644 --- a/docs/design/module/connector_contracts.md +++ b/docs/design/module/connector_contracts.md @@ -184,7 +184,11 @@ are resolved. Model-specific tensors remain explicit payload fields rather than transfer state. P2P can send optional router logits and one-dimensional, token-aligned -`torch.int32` input IDs after hidden states. FFN requests only the fields its +`torch.int32` input IDs after hidden states. CAMP2P sends the same ids through +the A2E operator's ids channel, which the receiving FFN rank selects with +`recv_input_ids`; both roles derive that mode from the model's +`afd_requires_input_ids` declaration, because the operator only writes the ids +slot in the mode the sender chose. FFN requests only the fields its model declares, concatenates them with the same per-peer token order as hidden states, and returns them in `AFDA2FTransferPayload`. The input-ID path supports DeepSeek V4's native hash router and has graph-stable receive buffers; it does diff --git a/docs/design/module/model_integration.md b/docs/design/module/model_integration.md index 26b3c77e..ac90e14c 100644 --- a/docs/design/module/model_integration.md +++ b/docs/design/module/model_integration.md @@ -28,6 +28,7 @@ verified_platform_refs: - "DeepSeek V2 Lite GPU and NPU model E2E paths" - "CAM async NPU model E2E path" - "DeepSeek V4 CUDA boundary has focused unit coverage only" + - "DeepSeek V4 Ascend Hash-id boundary has focused unit coverage and a manual Ascend 950 run" related_issues: - "#86" - "#88" @@ -57,6 +58,7 @@ make a backend-specific worker class the shared model API. | 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 | | 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), [`test_deepseek_v4_npu_weight_roles.py`](../../../tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.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) | From 614ae0e8eb8fd208ea029b285df171bf8f0f81cd Mon Sep 17 00:00:00 2001 From: ksiyuan Date: Wed, 16 Sep 2026 16:44:13 +0800 Subject: [PATCH 5/5] fix(npu): keep compressor models off the MLA DBO full graph path DeepSeek V4 declares both `index_topk` and `compress_ratios`, so the pinned NPUModelRunner selects the DSA backend: its `update_graph_params()` is a no-op and it never registers an FIA workspace. The ubatch wrapper still classified `use_mla and not use_sparse` as the MLA DBO full graph path, so a DBO run with a FULL graph aborted in `_new_mla_capture_params()` with "MLA DBO FULL graph requires the single-batch FIA workspace for 8 tokens" - a workspace that model cannot register, which made graph-mode DBO impossible for it. Mirror upstream's own graph-params condition, which excludes the compressor backend as well, and align the feature-validation predicate with the same backend selection. Compressor models now take the generic two-stage path; the plain MLA path is unchanged. Signed-off-by: ksiyuan --- afd_plugin/compat/npu/feature_validation.py | 13 +++-- .../v1/worker/npu/attention_model_runner.py | 15 ++++- docs/design/module/execution_platforms.md | 5 ++ tests/unit/v1/worker/test_npu_runtime.py | 57 +++++++++++++++++++ 4 files changed, 82 insertions(+), 8 deletions(-) diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index 3a58a3d4..a17ad4bd 100644 --- a/afd_plugin/compat/npu/feature_validation.py +++ b/afd_plugin/compat/npu/feature_validation.py @@ -69,16 +69,19 @@ def fail_if_unsupported_npu_afd_features( "AFD NPU runtime supports exactly two ubatches when DBO is enabled", ) model_config = vllm_config.model_config - # Match the pinned NPUModelRunner's sparse-attention backend selection. - uses_sparse_mla = hasattr( - model_config.hf_text_config, - "index_topk", - ) + # Mirror the pinned NPUModelRunner's attention backend selection: a sparse + # (SFA) or compressor (DSA) configuration selects a backend whose + # graph-params update is a no-op, so only plain MLA takes the MLA DBO full + # graph path that owns the merged registry. + hf_text_config = model_config.hf_text_config + uses_sparse_mla = hasattr(hf_text_config, "index_topk") + uses_mla_compressor = hasattr(hf_text_config, "compress_ratios") cudagraph_mode = vllm_config.compilation_config.cudagraph_mode uses_mla_dbo_full_graph = ( uses_ubatching and model_config.use_mla and not uses_sparse_mla + and not uses_mla_compressor and cudagraph_mode.has_full_cudagraphs() ) if uses_mla_dbo_full_graph and vllm_config.speculative_config is not None: diff --git a/afd_plugin/v1/worker/npu/attention_model_runner.py b/afd_plugin/v1/worker/npu/attention_model_runner.py index 15f53e9b..8f586d04 100644 --- a/afd_plugin/v1/worker/npu/attention_model_runner.py +++ b/afd_plugin/v1/worker/npu/attention_model_runner.py @@ -1580,14 +1580,23 @@ def _install_ascend_ubatch_wrapper(self) -> None: runtime_mode = CUDAGraphMode.FULL elif self.compilation_config.cudagraph_mode.has_full_cudagraphs(): runtime_mode = CUDAGraphMode.FULL + # Only the plain MLA backend owns the merged graph-params registry. + # The sparse (SFA) and compressor (DSA) backends define upstream's + # update_graph_params as a no-op and register no FIA workspace, which is + # why upstream itself skips its graph-params update for them. A + # compressor model such as DeepSeek V4 therefore takes the generic + # two-stage path instead of the MLA one. + mla_full_graph_enabled = ( + self.vllm_config.model_config.use_mla + and not self.use_sparse + and not self.use_compress + ) self.model = AscendUBatchWrapper( model, self.vllm_config, runtime_mode, self.device, - mla_full_graph_enabled=( - self.vllm_config.model_config.use_mla and not self.use_sparse - ), + mla_full_graph_enabled=mla_full_graph_enabled, full_graph_params_updater=self._update_full_graph_params_if_needed, enable_enpu=self.enable_enpu, ) diff --git a/docs/design/module/execution_platforms.md b/docs/design/module/execution_platforms.md index 605c6a9d..0348abb7 100644 --- a/docs/design/module/execution_platforms.md +++ b/docs/design/module/execution_platforms.md @@ -273,6 +273,11 @@ registry is exposed only through the active forward context under `afd_mla_graph_params`; the compatibility resolver falls back to upstream process-global state outside that scope. +This protocol belongs to the plain MLA backend alone. Upstream's sparse (SFA) +and compressor (DSA) backends implement `update_graph_params()` as a no-op and +register no FIA workspace, so those models, DeepSeek V4 among them, take the +generic two-stage path with no MLA registries. + The NPU V2 runner supports eager, `FULL`, and `FULL_DECODE_ONLY`. Like CUDA V2, it publishes descriptor-matched warmup/capture control outside formal graph capture and installs an instance-scoped pre-replay hook because native full diff --git a/tests/unit/v1/worker/test_npu_runtime.py b/tests/unit/v1/worker/test_npu_runtime.py index e5d4306c..5ca0645c 100644 --- a/tests/unit/v1/worker/test_npu_runtime.py +++ b/tests/unit/v1/worker/test_npu_runtime.py @@ -559,6 +559,7 @@ def __init__(self, *args: object, **kwargs: object): cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: True), ) runner.use_sparse = False + runner.use_compress = False runner.enable_enpu = False runner._install_ascend_ubatch_wrapper() @@ -572,6 +573,45 @@ def __init__(self, *args: object, **kwargs: object): assert updater.__self__ is runner +def test_npu_attention_runner_keeps_compressor_models_off_the_mla_path(monkeypatch): + """DeepSeek V4 selects the DSA backend, whose graph params are a no-op.""" + _require_npu_runtime() + from afd_plugin.v1.worker.npu import attention_model_runner + + captured_kwargs: list[dict[str, object]] = [] + + class RecordingUBatchWrapper: + def __init__(self, *args: object, **kwargs: object): + captured_kwargs.append(kwargs) + + monkeypatch.setattr( + attention_model_runner, + "AscendUBatchWrapper", + RecordingUBatchWrapper, + ) + runner = object.__new__( + attention_model_runner.AFDNPUAttentionModelRunner, + ) + runner.model = "model" + runner.device = "npu" + runner.vllm_config = SimpleNamespace( + model_config=SimpleNamespace(use_mla=True), + ) + runner.compilation_config = SimpleNamespace( + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: True), + ) + runner.use_sparse = False + runner.use_compress = True + runner.enable_enpu = False + + runner._install_ascend_ubatch_wrapper() + + # The MLA DBO full graph path requires a single-batch FIA workspace that a + # compressor model never registers, so those models stay on the generic + # two-stage path. + assert captured_kwargs[0]["mla_full_graph_enabled"] is False + + def test_npu_attention_runner_builds_and_sets_metadata(): torch = pytest.importorskip("torch") runner = _new_attention_runner() @@ -2327,6 +2367,23 @@ def test_npu_feature_validation_requires_decode_only_full_graph_for_mla_dbo( sparse_config.model_config.hf_text_config = SimpleNamespace(index_topk=8) fail_if_unsupported_npu_afd_features(sparse_config) + # A compressor model such as DeepSeek V4 selects the DSA backend, whose + # graph-params update is a no-op, so the MLA DBO full graph rules do not + # apply to it either. + compressor_config = _vllm_config( + use_mla=True, + cudagraph_mode="FULL", + enable_dbo=True, + use_ubatching=True, + num_ubatches=2, + ubatch_size=4, + ) + compressor_config.model_config.hf_text_config = SimpleNamespace( + index_topk=512, + compress_ratios=[0, 4, 128], + ) + fail_if_unsupported_npu_afd_features(compressor_config) + def test_npu_feature_validation_rejects_speculative_mla_dbo_full_graph(): config = _vllm_config(