diff --git a/afd_plugin/compat/npu/feature_validation.py b/afd_plugin/compat/npu/feature_validation.py index a8278e93..a17ad4bd 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, @@ -52,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): @@ -68,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: @@ -130,8 +134,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/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index 8dcec3bf..b862e965 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -224,6 +224,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 +513,17 @@ 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 + 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. 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 +535,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 +569,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,20 +628,30 @@ 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 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 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)) + compute_gate_mode = 1 if recv_input_ids else 0 batch_size = _num_tokens_for_ffn_rank( self.dp_metadata_list, ubatch_idx, @@ -582,14 +694,28 @@ def recv_attn_output( self.world_rank, group_ep, custom_states.aiv_num, - 0, + 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. + 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( @@ -850,6 +976,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 +995,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 +1025,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 +1100,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/afd_plugin/model_executor/models/npu/deepseek_v4.py b/afd_plugin/model_executor/models/npu/deepseek_v4.py index 32a74a6b..81e1771c 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,15 @@ 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 + # 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,)) return frozenset((_FFN_ROLE,)) # HC parameters and any future shared layer parameters are required by # both role-local model instances. @@ -97,12 +115,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 +287,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() @@ -303,6 +356,11 @@ def compute_ffn_output( # topk_weights, which CAM applies during combine-recv. routed_scale_applied_in_topk=True, ) + # The native MoE runs the gate internally when the gate is not on + # 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) @@ -533,6 +591,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 +631,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/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/afd_plugin/v1/worker/npu/ffn_model_runner.py b/afd_plugin/v1/worker/npu/ffn_model_runner.py index 2f42c6bc..5d60dd71 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 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, + 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 = payload.input_ids 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) 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/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/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) | 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 df513da6..99b8c008 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -149,6 +149,47 @@ 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. 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") + 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.input_ids.tolist() == [0, 2, 4, 6, 8] + assert calls[1][-1] == 0 + assert without_ids.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..1ecfae17 --- /dev/null +++ b/tests/unit/connectors/test_camp2p_token_ids.py @@ -0,0 +1,362 @@ +# 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. + +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 + +import contextlib +import logging +import sys +import types +from collections.abc import Iterator +from types import SimpleNamespace +from typing import Any + +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.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) + + 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) + + +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/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..46d5f88d --- /dev/null +++ b/tests/unit/model_executor/test_deepseek_v4_hash_ids.py @@ -0,0 +1,310 @@ +# 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 ast +import contextlib +import logging +import sys +import types +from collections.abc import Iterator +from pathlib import Path +from typing import Any + +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: 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( # 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 # 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) + 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, + ) + + +_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 new file mode 100644 index 00000000..7130916d --- /dev/null +++ b/tests/unit/model_executor/test_deepseek_v4_npu_weight_roles.py @@ -0,0 +1,209 @@ +# 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. 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 +---------------------------- +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")), + } + + 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_follows_the_gate_ownership(name: str) -> None: + """The Hash table belongs to every role that built a router. + + 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. + """ + + 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( + "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"}) 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(