Topology and role ranks: the AFD world

    role_rank = (dp_rank * pcp_size + pcp_rank) * tp_size + tp_rank
    if config.role == "attention":
        role_size = config.num_attention_ranks
    elif config.role == "ffn":
        role_size = config.num_ffn_ranks
    else:
        raise ValueError(f"unknown AFD role {config.role!r}")
    if not 0 <= role_rank < role_size:
        raise ValueError(
            "AFD role rank derived from distributed ranks is out of range: "
            f"role={config.role!r}, dp_rank={dp_rank}, pcp_rank={pcp_rank}, "
            f"tp_rank={tp_rank}, role_size={role_size}",
        )

    if config.role == "attention":
        if role_rank >= attention_size:
            raise ValueError(
                "Attention role rank must be within attention size "
                f"(rank={role_rank}, size={attention_size})",
            )
        world_rank = ffn_size + role_rank
        subgroup_index = role_rank * ffn_size // attention_size
    elif config.role == "ffn":
        if role_rank >= ffn_size:
            raise ValueError(
                "FFN role rank must be within FFN size "
                f"(rank={role_rank}, size={ffn_size})",
            )
        world_rank = role_rank
        subgroup_index = role_rank

    min_size = min(ffn_size, attention_size)
    subgroup_attention_ranks = [
        ffn_size + attention_rank
        for attention_rank in subgroup_attention_block(
            subgroup_index, attention_size, ffn_size
        )
    ]
    subgroup_ranks = tuple([subgroup_index] + subgroup_attention_ranks)
    rank_in_subgroup = subgroup_ranks.index(world_rank)

def validate_p2p_topology(config: AFDConfig) -> None:
    attention_size, ffn_size = topology_from_config(config)
    if attention_size < ffn_size:
        raise ValueError(
            "P2pNcclAFDConnector currently requires num_attention_ranks >= "
            f"num_ffn_ranks, got {attention_size} < {ffn_size}",
        )

    destinations: list[int] = []
    if ffn_size <= world_rank < ffn_size + min_size:
        local_attention_rank = world_rank - ffn_size
        destination = local_attention_rank
        while destination < ffn_size:
            destinations.append(destination)
            destination += min_size

        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),
        )

        with DefaultProcessGroupSwitcher(_get_default_group(), afd_pg):
            # The subgroup only has to hand rank 0's ncclUniqueId to its
            # members, so it reuses the store the AFD world already
            # rendezvoused on. StatelessProcessGroup.create would stand up a
            # second store bound to ``host``, which is only correct while
            # every subgroup's rank 0 sits on that host.
            subgroup_store = PrefixStore(
                f"afd_subgroup_{self.mapping.subgroup_index}",
                afd_pg.get_group_store(),
            )
            subgroup_store.set_timeout(timedelta(seconds=300))
            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)

        if self.mapping.participates_in_dp_metadata_group:
            self.p2p_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.min_size,
                rank=self.p2p_rank,
                group_name="p2p",
                timeout=timedelta(minutes=30),
            )