The connector contract: the Attention/FFN transport

class AFDConnectorBase(ABC):
    """Base class for plugin-owned AFD Attention/FFN connectors.

    An AFD connector owns the communication contract between the Attention
    runtime and the FFN runtime. Implementations are responsible for three
    kinds of work:

    1. Initializing and releasing backend communication resources.
    2. Moving hidden states between Attention and FFN ranks.
    3. Applying DP metadata control-plane payloads so later tensor transfers
       can use backend-specific shapes, buffers, or connector data.

    Implementations may use very different transport mechanisms. For example,
    the GPU P2P connector uses explicit point-to-point tensor transfers, while
    the NPU CAMP2P connector carries additional backend state through its
    ``AFDTransferState`` subclass on ``AFDTransferContext.states``. This base
    class documents the common runtime contract shared by those implementations.
    """

    @abstractmethod
    def close(self) -> None:
        """Release connector-owned communication resources.

        Implementations should destroy process groups, communicators, cached
        handles, and backend-specific state owned by the connector. The method
        should be safe to call during worker shutdown and should leave
        ``is_initialized`` false after successful cleanup.
        """
        raise NotImplementedError

    @abstractmethod
    def init_afd_connector(self) -> None:
        """Initialize backend communication resources.

        This method is called once the owning runtime is ready to create AFD
        communication state. Implementations typically create process groups,
        register custom ops, initialize backend communicators, and precompute
        topology-derived rank lists.
        """
        raise NotImplementedError

    @property
    @abstractmethod
    def is_initialized(self) -> bool:
        """Return whether backend communication resources are initialized."""
        raise NotImplementedError

    @abstractmethod
    def send_attn_output(
        self,
        hidden_states: torch.Tensor,
        context: AFDTransferContext,
        **kwargs: Any,
    ) -> None:
        """Send Attention hidden states to the FFN runtime.

        This method is called on Attention-side ranks after an Attention block
        produces hidden states for one layer/stage.

        Args:
            hidden_states: Tensor to send to the FFN side. The leading
                dimension should match ``context.metadata.total_tokens`` unless
                the backend explicitly supports a different compiled/captured
                calling convention.
            context: Per-transfer context describing layer, stage, and token
                layout. Backends may also consume ``context.states``.
            **kwargs: Optional backend-specific arguments.

        Raises:
            ValueError: If tensor shape does not match metadata.
            RuntimeError: If the connector is not initialized or required
                backend-specific metadata is missing.
        """
        raise NotImplementedError

    @abstractmethod
    def recv_ffn_output(
        self,
        ref_tensor: torch.Tensor,
        ubatch_idx: int = 0,
        **kwargs: Any,
    ) -> torch.Tensor:

    @abstractmethod
    def recv_attn_output(
        self,
        ubatch_idx: int = 0,
        **kwargs: Any,
    ) -> AFDA2FTransferPayload:
        """Receive Attention hidden states on the FFN side.

        This method is called by FFN-side execution before running the FFN
        compute for a layer/stage. It returns both the hidden states and the
        metadata needed by later FFN-side calls.

        Args:
            ubatch_idx: Microbatch/stage index to receive. Defaults to ``0``.
            **kwargs: Optional backend-specific receive arguments.

        Returns:
            ``AFDA2FTransferPayload`` containing received hidden states and the
            transfer context (metadata + backend-produced transfer state)
            for the receive operation.
        """
        raise NotImplementedError

    @abstractmethod
    def send_ffn_output(
        self,
        ffn_output: torch.Tensor,
        context: AFDTransferContext,
        **kwargs: Any,
    ) -> None:

    # An optional DP metadata control plane. When set, FFN steps are driven by
    # DP metadata received over the control plane. When ``None``, the connector
    # has no control plane and FFN steps are driven directly by the connector
    # receive loop.
    control_plane: AFDControlPlane | None = None

class AFDControlPlane(ABC):
    """DP metadata control plane"""

    @abstractmethod
    def update_state_from_dp_metadata(
        self,
        payload: AFDControlPayload,
    ) -> None:
        """Apply a DP metadata payload to local connector state.

    @classmethod
    @abstractmethod
    def parse_extra_config(
        cls,
        raw: Mapping[str, Any] | None,
    ) -> ConnectorExtraInfo:
        """Parse connector-owned config using this connector's schema."""

        raise NotImplementedError

    @classmethod
    def create_connector(
        cls,
        rank: int,
        local_rank: int,
        vllm_config: VllmConfig,
        afd_config: AFDConfig | None = None,
    ) -> AFDConnectorBase:
        config = afd_config or parse_afd_config(vllm_config)
        if config.connector not in cls._registry:
            raise ValueError(f"unsupported AFD connector type: {config.connector}")
        connector_cls = cls._registry[config.connector]()
        role_rank = resolve_role_rank(vllm_config, config)
        return connector_cls(
            rank,
            local_rank,
            vllm_config,
            config,
            role_rank,
        )

AFDConnectorFactory.register_connector(
    "P2pNcclAFDConnector",
    "afd_plugin.connectors.gpu.p2p",
    "P2pNcclAFDConnector",
)
AFDConnectorFactory.register_connector(
    "CAMP2pAFDConnector",
    "afd_plugin.connectors.npu.camp2p",
    "CAMP2pAFDConnector",
)
AFDConnectorFactory.register_connector(
    "CAMAsyncAFDConnector",
    "afd_plugin.connectors.npu.async_cam",
    "CAMAsyncAFDConnector",
)