Data on the wire: tensors per layer and stage

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

The idea#

For each layer and stage, the attention rank sends its hidden states to its FFN rank. The FFN rank may own more than one attention peer, so it concatenates the inputs along the token dimension, runs the feed-forward, then splits the output back and returns each slice to the rank that produced it.

Each transfer carries a small metadata record: the layer index, the stage index and the per-peer token lengths. The token lengths set the expected tensor shape and the split sizes. The transfers go through custom ops, so they stay usable under torch.compile and CUDA graph capture.

An everyday comparison. Think of parcels with shipping labels. The label says how many boxes and for which address, so the depot can combine shipments and later separate them again. The comparison stops here: the "parcels" also carry optional extras next to the hidden states, the router logits and the token IDs, which only some models send.

The picture#

Read panel 1 as one exchange for one layer and stage. Read panel 2 as the fan-in when one FFN rank owns two attention peers.

In the code#

Step 1: What the attention rank sends#

send_attn_output sends the hidden states first, then the optional router logits, then the optional token IDs (afd_plugin/connectors/gpu/p2p.py:323-390). The hidden states are a CUDA tensor of shape (num_tokens, hidden_size); the hidden size comes from the model config (afd_plugin/connectors/gpu/p2p.py:195-197).

        metadata = context.metadata
        if not torch.compiler.is_compiling() and not metadata.validate_tensor_shape(
            tuple(hidden_states.shape),
        ):
            raise ValueError(
                f"hidden_states shape {hidden_states.shape!r} does not match "
                f"AFD metadata token count {metadata.total_tokens}",
            )
        router_logits: torch.Tensor | None = kwargs.get("router_logits")
        input_ids: torch.Tensor | None = kwargs.get("input_ids")
        if (
            router_logits is not None
            and router_logits.shape[0] != hidden_states.shape[0]
        ):
            raise ValueError(
                "router_logits and hidden_states must have equal token counts",
            )
        if input_ids is not None:
            if input_ids.ndim != 1 or input_ids.shape[0] != hidden_states.shape[0]:
                raise ValueError(
                    "input_ids must be one-dimensional and token-aligned with "
                    "hidden_states",
                )
            if input_ids.dtype != _INPUT_IDS_DTYPE:
                raise ValueError("P2P AFD input_ids must use torch.int32")

The leading dimension must equal total_tokens, unless torch.compile is tracing (afd_plugin/connectors/gpu/p2p.py:347-353). The token IDs must be a one-dimensional int32 tensor (afd_plugin/connectors/gpu/p2p.py:363-370).

Step 2: The transfer metadata#

Each transfer carries an AFDTransferMetadata: the layer index, the stage index and the per-peer sequence lengths (afd_plugin/connectors/metadata.py:188-190). The sum of the lengths is the expected leading dimension (afd_plugin/connectors/metadata.py:198-200).

    @property
    def total_tokens(self) -> int:
        return sum(self.seq_lens)

The shape check compares shape[0] with total_tokens (afd_plugin/connectors/metadata.py:230-231).

Step 3: The FFN rank fans in its peers#

recv_attn_output receives one tensor from every attention peer in the subgroup, concatenates them along the token dimension, and records each peer's length in the returned metadata (afd_plugin/connectors/gpu/p2p.py:428-552).

        hidden_states = (
            torch.cat(hidden_states_list, dim=0)
            if len(hidden_states_list) > 1
            else hidden_states_list[0]
        )
        router_logits = None
        if routing_spec is not None:
            router_logits = (
                torch.cat(router_logits_list, dim=0)
                if len(router_logits_list) > 1
                else router_logits_list[0]
            )
        input_ids = None
        if input_ids_list:
            input_ids = (
                torch.cat(input_ids_list, dim=0)
                if len(input_ids_list) > 1
                else input_ids_list[0]
            )

        metadata = AFDTransferMetadata.create_ffn_metadata(
            layer_idx=0,
            stage_idx=ubatch_idx,
            seq_lens=[tensor.shape[0] for tensor in hidden_states_list],
        )

With one peer the single tensor is the whole input (afd_plugin/connectors/gpu/p2p.py:520-524). All of a step's receives run in one NCCL group, so the peers' transfers proceed together (afd_plugin/connectors/gpu/p2p.py:514-519).

Step 4: The FFN rank splits the output back#

send_ffn_output splits the output along the token dimension by the recorded sequence lengths, then sends each slice to the attention rank that produced it (afd_plugin/connectors/gpu/p2p.py:554-613).

        output_slices: list[torch.Tensor] = []
        dsts: list[int] = []
        start = 0
        for dst, token_count in zip(
            range(1, self.group_size),
            split_sizes,
            strict=False,
        ):
            end = start + token_count
            output_slices.append(ffn_output[start:end])
            dsts.append(dst)
            start = end
        self._send_tensors(output_slices, dsts, self.e2a_group, self.e2a_comm_id)

When the recorded lengths do not cover every peer, the split falls back to an even split, which raises if the tokens do not divide evenly (afd_plugin/connectors/gpu/p2p.py:590-599).

Step 5: The transfers are custom ops#

The sends and receives are wrapped in four custom ops: afd_p2p_send, afd_p2p_recv, and the _group variants (afd_plugin/connectors/gpu/p2p.py:979-1122). The wrapping keeps the transfers traceable by torch.compile and capturable by CUDA graphs (afd_plugin/connectors/gpu/p2p.py:983-986). Each op runs on the current CUDA stream (afd_plugin/connectors/gpu/p2p.py:1002-1006). The connectors/gpu/p2p.py: the GPU P2P connector page reads the custom ops in full.

Deep dive#

What can go wrong#

• Shape mismatch. A hidden-state tensor whose leading dimension is not total_tokens raises, unless torch.compile is tracing (afd_plugin/connectors/gpu/p2p.py:347-353).
• An uneven split. When the recorded lengths do not cover every peer and the tokens do not divide evenly, the split raises (afd_plugin/connectors/gpu/p2p.py:590-599).
• A CPU tensor. The hidden states must be on the GPU; a CPU tensor is rejected (afd_plugin/connectors/gpu/p2p.py:634-635).
• Wrong token-ID dtype. The token IDs must be int32 and one-dimensional (afd_plugin/connectors/gpu/p2p.py:363-370).

Why it matters. The two roles must agree on the shapes. The recorded sequence lengths are the only record of which tokens belong to which peer, so the FFN rank uses them to split the output back.