diff --git a/afd_plugin/connectors/npu/camp2p.py b/afd_plugin/connectors/npu/camp2p.py index 32f25424..45e5dde1 100644 --- a/afd_plugin/connectors/npu/camp2p.py +++ b/afd_plugin/connectors/npu/camp2p.py @@ -290,6 +290,21 @@ def is_initialized(self) -> bool: """Return ``True`` after all CAMP2p connections have been created.""" return self._initialized + def _get_afd_pg(self, ubatch_idx: int) -> ProcessGroup: + """Return the AFD HCCL process group for a ubatch (A5 p2p path).""" + if not self.afd_pg_list: + raise RuntimeError("CAMP2P connector has no AFD process groups") + if ubatch_idx < 0: + raise RuntimeError( + f"CAMP2P ubatch index must be non-negative: {ubatch_idx}", + ) + if ubatch_idx >= len(self.afd_pg_list): + raise RuntimeError( + f"CAMP2P ubatch {ubatch_idx} requires " + f"{ubatch_idx + 1} AFD process groups", + ) + return self.afd_pg_list[ubatch_idx] + def init_afd_connector(self) -> None: """Connect this process to the other Attention and FFN processes. @@ -309,9 +324,13 @@ def init_afd_connector(self) -> None: return import torch_npu # noqa: F401 - ensure_cam_p2p_ops_available() - - _register_camp2p_custom_ops() + if is_a5(): + # The a2e/e2a custom ops are 910C-only (not registered for + # ascend950); Route B2 on A5 uses plain HCCL p2p instead. + logger.info("CAMP2P on A5: using HCCL p2p route (B2), skipping custom ops") + else: + ensure_cam_p2p_ops_available() + _register_camp2p_custom_ops() num_ubatches = max(1, self.vllm_config.parallel_config.num_ubatches) self.afd_pg_list = [] @@ -425,6 +444,20 @@ def send_attn_output( f"hidden_states shape {hidden_states.shape!r} does not match " f"CAMP2P metadata token count {metadata.total_tokens}", ) + if is_a5(): + # Route B2: plain HCCL p2p send to the mapped FFN rank. The + # a2e/e2a custom ops (and native MC2 ops) are unusable on A5. + if torch.compiler.is_compiling(): + return None + ubatch_idx = metadata.stage_idx + get_forward_context().ubatch_idx = ubatch_idx + dst_ffn = dst_ffn_for_attention( + self.world_rank - self.ffn_size, + self.attn_size, + self.ffn_size, + ) + p2p_send(self._get_afd_pg(ubatch_idx), hidden_states, dst_ffn) + return None transfer_state = CAMP2PTransferState( aiv_num=self.aiv_num, batch_size=metadata.total_tokens, @@ -475,6 +508,23 @@ def recv_ffn_output( """ if not self._initialized: raise RuntimeError("CAMP2P connector is not initialized") + if is_a5(): + # Route B2: recv the FFN result over HCCL p2p from the mapped FFN + # rank, sized by the ref tensor (same tokens as this rank sent). + if torch.compiler.is_compiling(): + return ref_tensor + src_ffn = dst_ffn_for_attention( + self.world_rank - self.ffn_size, + self.attn_size, + self.ffn_size, + ) + return p2p_recv( + self._get_afd_pg(ubatch_idx), + tuple(ref_tensor.shape), + ref_tensor.dtype, + ref_tensor.device, + src_ffn, + ) transfer_state = getattr(get_forward_context(), "cam_afdtransfer_state", None) if transfer_state is None: raise RuntimeError("CAMP2P Attention side is missing connector data") @@ -527,6 +577,57 @@ def recv_attn_output( ffn_size=self.ffn_size, fallback=max_num_tokens, ) + if is_a5(): + # Route B2: recv each mapped Attention peer's token block over + # HCCL p2p and concatenate. Per-peer counts come from the DP + # metadata control plane (no equal-ratio split like e2a). + peers = attention_peers_for_ffn( + self.role_rank, + self.attn_size, + self.ffn_size, + ) + counts = attention_token_counts( + self.dp_metadata_list, + ubatch_idx, + self.attn_size, + ) + if counts is not None: + seq_lens = [max(1, counts[peer]) for peer in peers] + else: + # Metadata missing (e.g. warmup): even split of the total. + seq_lens = [max(1, batch_size // len(peers))] * len(peers) + dtype = getattr( + self.vllm_config.model_config, + "dtype", + torch.bfloat16, + ) + pg = self._get_afd_pg(ubatch_idx) + blocks = [ + p2p_recv( + pg, + (seq_lens[i], self.hidden_size), + dtype, + torch.device("npu"), + self.ffn_size + peer, + ) + for i, peer in enumerate(peers) + ] + hidden_states = torch.cat(blocks, dim=0) + a5_metadata = AFDTransferMetadata.create_ffn_metadata( + layer_idx=layer_idx, + stage_idx=ubatch_idx, + seq_lens=seq_lens, + ) + a5_states = CAMP2PTransferState( + aiv_num=self.aiv_num, + batch_size=int(sum(seq_lens)), + h=self.hidden_size, + k=self.num_experts_per_tok, + ) + return AFDA2FTransferPayload( + hidden_states=hidden_states, + context=AFDTransferContext(metadata=a5_metadata, states=a5_states), + ) metadata = AFDTransferMetadata.create_ffn_metadata( layer_idx=layer_idx, stage_idx=ubatch_idx, @@ -593,9 +694,31 @@ def send_ffn_output( if not self._initialized: raise RuntimeError("CAMP2P connector is not initialized") states = context.states + ubatch_idx = int(kwargs.get("ubatch_idx", context.metadata.stage_idx)) + if is_a5(): + # Route B2: send each Attention peer its slice of the result back + # over HCCL p2p, split by the per-peer counts captured at recv. + if torch.compiler.is_compiling(): + return None + peers = attention_peers_for_ffn( + self.role_rank, + self.attn_size, + self.ffn_size, + ) + split_sizes = list(context.metadata.seq_lens) + if sum(split_sizes) != ffn_output.shape[0]: + # Inconsistent per-peer counts: fall back to an even split. + split_sizes = [ffn_output.shape[0] // len(peers)] * len(peers) + for i in range(ffn_output.shape[0] % len(peers)): + split_sizes[i] += 1 + pg = self._get_afd_pg(ubatch_idx) + offset = 0 + for peer, n in zip(peers, split_sizes, strict=True): + p2p_send(pg, ffn_output[offset : offset + n], self.ffn_size + peer) + offset += n + return None if states.atten_batch_size is None: raise RuntimeError("CAMP2P FFN side is missing A2E atten_batch_size") - ubatch_idx = int(kwargs.get("ubatch_idx", context.metadata.stage_idx)) group_ep = _get_group_ep( ubatch_idx, self.hccl_comm_name, @@ -741,6 +864,31 @@ def build_camp2p_topology( ) +def attention_token_counts( + dp_metadata_list: Mapping[int, DPMetadata | AFDDPMetadata], + stage_idx: int, + attention_size: int, +) -> list[int] | None: + """Per-Attention-rank token counts for a stage (DP expanded to AFD ranks). + + Mirrors the DP -> AFD expansion used by ``_num_tokens_for_ffn_rank``: when + TP creates several Attention workers per DP rank, the DP token count is + replicated ``tp_size`` times. Returns ``None`` when the metadata is missing + or cannot be expanded (the caller falls back to an even split). + """ + dp_metadata = dp_metadata_list.get(stage_idx) + if dp_metadata is None: + return None + token_counts = dp_metadata.num_tokens_across_dp_cpu + counts = token_counts.flatten().tolist() + if len(counts) < attention_size and attention_size % len(counts) == 0: + tp_size = attention_size // len(counts) + counts = [counts[i // tp_size] for i in range(attention_size)] + if len(counts) < attention_size: + return None + return counts + + def _num_tokens_for_ffn_rank( dp_metadata_list: dict[int, DPMetadata | AFDDPMetadata], stage_idx: int, @@ -767,28 +915,94 @@ def _num_tokens_for_ffn_rank( Returns: The number of tokens this FFN rank should receive, always at least one. """ - dp_metadata = dp_metadata_list[stage_idx] - if dp_metadata is None: - return max(1, fallback) - token_counts = dp_metadata.num_tokens_across_dp_cpu - counts = token_counts.flatten().tolist() - # Expand DP-level counts to AFD-level when TP > 1. - # num_tokens_across_dp_cpu has dp_size entries, but attention_size - # = num_attention_ranks includes TP workers. Each DP rank's count - # is replicated tp_size times. - if len(counts) < attention_size and attention_size % len(counts) == 0: - tp_size = attention_size // len(counts) - counts = [counts[i // tp_size] for i in range(attention_size)] - if len(counts) < attention_size: + counts = attention_token_counts(dp_metadata_list, stage_idx, attention_size) + if counts is None: return max(1, fallback) if attention_size >= ffn_size and attention_size % ffn_size == 0: group_size = attention_size // ffn_size start_idx = ffn_rank * group_size - end_idx = start_idx + group_size - return max(1, sum(counts[start_idx:end_idx])) + return max(1, sum(counts[start_idx : start_idx + group_size])) return max(1, fallback) +# --------------------------------------------------------------------------- +# A5 (Ascend950) HCCL p2p transport (Route B2). +# +# On Atlas A5 the a2e/e2a custom ops and the torch_npu native MoE +# dispatch/combine ops both fail on the AFD mixed group with 507035 (MTE +# out-of-range; the remote HCCL window entries are not kernel-addressable on +# A5). Route B2 therefore moves hidden states between Attention and FFN ranks +# with plain HCCL ``dist.send``/``recv`` over the per-ubatch ``afd`` groups; +# the FFN side runs its MoE through the standard vLLM-Ascend EP path, where +# the native ops are proven. Only ``hidden_states`` crosses the wire (gate +# stays on FFN, ``compute_gate_on_attention=false``). +# +# The Attention<->FFN rank mapping matches ``_num_tokens_for_ffn_rank``: with +# ``group_size = attention_size // ffn_size``, Attention local rank ``i`` maps +# to FFN rank ``i // group_size`` and FFN rank ``j`` receives from the +# consecutive Attention local ranks ``[j*group_size, (j+1)*group_size)``. +# --------------------------------------------------------------------------- + +_is_a5: bool | None = None + + +def is_a5() -> bool: + """Return whether this process runs on Atlas A5 (Ascend950).""" + global _is_a5 + if _is_a5 is None: + try: + from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type + + _is_a5 = get_ascend_device_type() == AscendDeviceType.A5 + except Exception: + # Not on an Ascend platform or vllm_ascend unavailable. + _is_a5 = False + return _is_a5 + + +def attention_group_size(attention_size: int, ffn_size: int) -> int: + """Number of Attention ranks mapped to each FFN rank.""" + return attention_size // ffn_size + + +def dst_ffn_for_attention( + attn_local_rank: int, + attention_size: int, + ffn_size: int, +) -> int: + """FFN rank (0-based) that an Attention rank sends its tokens to.""" + return attn_local_rank // attention_group_size(attention_size, ffn_size) + + +def attention_peers_for_ffn( + ffn_rank: int, + attention_size: int, + ffn_size: int, +) -> list[int]: + """Attention local ranks that an FFN rank receives from (ascending).""" + group_size = attention_group_size(attention_size, ffn_size) + start = ffn_rank * group_size + return list(range(start, start + group_size)) + + +def p2p_send(pg, tensor: torch.Tensor, dst_rank: int) -> None: + """Blocking HCCL send of ``tensor`` to ``dst_rank`` over the AFD group.""" + dist.send(tensor.contiguous(), dst=dst_rank, group=pg) + + +def p2p_recv( + pg, + shape: tuple[int, ...], + dtype: torch.dtype, + device: torch.device, + src_rank: int, +) -> torch.Tensor: + """Blocking HCCL recv of a tensor with the given shape from ``src_rank``.""" + buf = torch.empty(shape, dtype=dtype, device=device) + dist.recv(buf, src=src_rank, group=pg) + return buf + + def _get_group_ep( ubatch_idx: int, hccl_comm_name1: str, diff --git a/tests/unit/connectors/test_camp2p_connector.py b/tests/unit/connectors/test_camp2p_connector.py index 1abc3acb..88e83954 100644 --- a/tests/unit/connectors/test_camp2p_connector.py +++ b/tests/unit/connectors/test_camp2p_connector.py @@ -21,7 +21,11 @@ CAMP2pAFDConnector, CAMP2PExtraInfo, CAMP2PTransferState, + attention_group_size, + attention_peers_for_ffn, + attention_token_counts, build_camp2p_topology, + dst_ffn_for_attention, ) @@ -146,6 +150,46 @@ def test_camp2p_recv_attn_output_uses_original_contiguous_af_grouping(monkeypatc assert context0.states.k == 2 +def test_camp2p_a5_attention_to_ffn_mapping_is_contiguous_groups(): + # attn_size=4, ffn_size=2 -> group_size=2. Attention local ranks 0,1 map + # to FFN 0 and 2,3 to FFN 1 - the same grouping as _num_tokens_for_ffn_rank + # and the CAMP2P recv test (seq_lens [5] for FFN0, [12] for FFN1). + assert attention_group_size(4, 2) == 2 + assert dst_ffn_for_attention(0, 4, 2) == 0 + assert dst_ffn_for_attention(1, 4, 2) == 0 + assert dst_ffn_for_attention(2, 4, 2) == 1 + assert dst_ffn_for_attention(3, 4, 2) == 1 + assert attention_peers_for_ffn(0, 4, 2) == [0, 1] + assert attention_peers_for_ffn(1, 4, 2) == [2, 3] + + +def test_camp2p_a5_mapping_round_trip_covers_all_attention_ranks(): + for ffn_size, attn_size in [(1, 1), (1, 2), (2, 4), (2, 6), (4, 4)]: + groups = [ + attention_peers_for_ffn(j, attn_size, ffn_size) for j in range(ffn_size) + ] + assert sorted(p for group in groups for p in group) == list( + range(attn_size), + ) + for i in range(attn_size): + assert i in attention_peers_for_ffn( + dst_ffn_for_attention(i, attn_size, ffn_size), + attn_size, + ffn_size, + ) + + +def test_camp2p_a5_attention_token_counts_expands_dp_to_tp(): + # dp_size=2, attention_size=4 -> each DP count is replicated 2x. + counts = attention_token_counts({0: _FakeDPMetadata([2, 3])}, 0, 4) + assert counts == [2, 2, 3, 3] + + +def test_camp2p_a5_attention_token_counts_missing_metadata_returns_none(): + assert attention_token_counts({}, 0, 4) is None + assert attention_token_counts({0: None}, 0, 4) 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})