connectors/gpu/p2p.py: the GPU P2P connector

Code version: afd-plugin @ 7452d52 · 2026-10-27 · vLLM @ ced6857

Purpose#

This file implements the connector contract for CUDA. It builds the AFD world, the per-subgroup communicators and the control-plane group, and moves tensors through four custom ops. The contract itself is on The connector contract: the Attention/FFN transport; this page reads the implementation.

Where it sits#

Who calls it#

The factory loads this module by name and builds the connector for a runner (afd_plugin/connectors/factory.py:88-92). The runner then calls init_afd_connector once (afd_plugin/v1/worker/ffn_model_runner.py:108-109), then the four data calls per layer and stage, and close at shutdown (afd_plugin/v1/worker/ffn_model_runner.py:455).

What it calls#

The data transfers go through vLLM's PyNcclCommunicator.send and recv (vllm:vllm/distributed/device_communicators/pynccl.py:91). The process groups come from the plugin's own helper (afd_plugin/distributed/afd_process_group.py:62-120) and the rank mapping from the topology module (afd_plugin/distributed/topology.py:118).

Key classes and functions#

Walk through the main methods#

init_afd_connector: the collective setup#

This is a collective call: every attention and FFN rank must run it with matching host, port and rank counts (afd_plugin/connectors/gpu/p2p.py:245-262). It does three things. First, it registers the custom ops and joins the AFD world group, rendezvoused at tcp://host:port (afd_plugin/connectors/gpu/p2p.py:267-276).

        afd_pg = init_afd_process_group(
            backend="nccl",
            init_method=f"tcp://{self.afd_config.host}:{self.afd_config.port}",
            world_size=self.ffn_size + self.attn_size,
            rank=self.world_rank,
            group_name="afd",
            timeout=timedelta(seconds=self.afd_config.afd_process_group_timeout_s),
        )

Second, it builds this rank's subgroup communicators over a prefix of the world's rendezvous store (afd_plugin/connectors/gpu/p2p.py:278-304). It creates two PyNcclCommunicator instances, attention-to-FFN and FFN-to-attention, and registers each for the custom ops.

            self.a2e_group = StatelessProcessGroup(
                rank=self.mapping.rank_in_subgroup,
                world_size=len(self.mapping.subgroup_ranks),
                store=subgroup_store,
            )
            self.e2a_group = self.a2e_group
            self.a2e_pynccl = PyNcclCommunicator(
                group=self.a2e_group,
                device=self.local_rank,
            )
            self.a2e_comm_id = _register_comm(self.a2e_pynccl)
            self.e2a_pynccl = PyNcclCommunicator(
                group=self.e2a_group,
                device=self.local_rank,
            )
            self.e2a_comm_id = _register_comm(self.e2a_pynccl)

Third, on ranks that join the control plane, it creates the "p2p" group (afd_plugin/connectors/gpu/p2p.py:306-314). The Topology and role ranks: the AFD world page reads the world order and the group sizes.

The four data calls#

The data path is four methods. The attention side calls two, the FFN side calls two.

• send_attn_output sends hidden states, then optional router logits and token IDs, to subgroup rank 0 (afd_plugin/connectors/gpu/p2p.py:371-390).
• recv_ffn_output receives this rank's output slice into a preallocated ref_tensor (afd_plugin/connectors/gpu/p2p.py:415-421).
• recv_attn_output receives every peer's tensor in one NCCL group, concatenates them, and records each peer's length (afd_plugin/connectors/gpu/p2p.py:514-544).
• send_ffn_output splits the output by those lengths and sends each slice back (afd_plugin/connectors/gpu/p2p.py:590-613).

The Data on the wire: tensors per layer and stage page reads the shapes, the fan-in and the split in full.

P2pNcclAFDControlPlane: the DP-metadata path#

The connector owns one control-plane instance and exposes it as control_plane (afd_plugin/connectors/gpu/p2p.py:222, afd_plugin/connectors/gpu/p2p.py:724-733). It has three methods: apply a payload locally, send one, receive one.

    def recv_dp_metadata_list(self) -> AFDControlPayload:
        """Receive a DP metadata payload on an FFN rank.

        Blocks on the ``p2p`` process group until the mapped Attention sender
        rank transmits the next payload.

        Returns:
            The DP metadata payload sent by the Attention side.

        Raises:
            RuntimeError: If the DP metadata process group is not initialized
                on this rank.
        """
        connector = self.connector
        if connector.p2p_pg is None:
            raise RuntimeError("P2P DP metadata process group is not initialized")

        src = connector.p2p_rank % connector.min_size + connector.ffn_size

update_state_from_dp_metadata derives the wire-tensor shapes and preallocates the graph-mode buffers (afd_plugin/connectors/gpu/p2p.py:744-851); send_dp_metadata_list transmits only on the designated sender ranks (afd_plugin/connectors/gpu/p2p.py:870-883). The The control plane: DP metadata once per step page reads these.

_register_p2p_custom_ops: four ops, once per process#

The module registers four custom ops: afd_p2p_send, afd_p2p_recv, and the _group variants (afd_plugin/connectors/gpu/p2p.py:979-1122). Each op wraps the NCCL send or receive on the current stream and has a fake implementation for tracing.

        if hasattr(torch.ops.vllm, op_name):
            return
        direct_register_custom_op(
            op_name=op_name,
            op_func=op_func,
            mutates_args=mutates_args,
            fake_impl=fake_impl,
        )

1. Line 1089-1090. If the op is already on torch.ops.vllm, return; torch rejects a re-definition.
2. Line 1091-1096. Otherwise register the op, its mutated args and its fake implementation.
3. Line 1122. Set the module flag so the registration runs once per process.

Deep dive#

Related concepts#

• The connector contract: the Attention/FFN transport — the contract this file implements.
• Topology and role ranks: the AFD world — the world order and the group sizes.
• The control plane: DP metadata once per step — the DP-metadata path.
• Data on the wire: tensors per layer and stage — the data-plane shapes and the fan-in.