The control plane: DP metadata once per step

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

The idea#

The FFN side must know the token count of each stage before it receives any hidden states. The sizes set the shapes of the wire tensors and the receive buffers. The attention side sends these counts once per step, before the layer loop, over the control plane. This is separate from the data plane, which moves the hidden states once per layer and stage.

The payload is small: the per-DP-rank token counts, plus a few flags for graph capture and warm-up. It travels as JSON over a dedicated process group.

An everyday comparison. Think of a manifest sent ahead of a convoy of trucks. The manifest lists how many boxes each truck carries, so the warehouse can set up the right number of bays before the trucks arrive. The comparison stops here: one connector, the async NPU one, carries no manifest at all; its receive loop sizes each transfer as it comes.

The picture#

Read panel 1 as the per-step order: control plane first, then the data-plane layer loop, then sampling. Read panel 2 as the path of one payload from the attention side to the FFN side.

In the code#

Step 1: Once per step, before the layers#

The control plane runs once per step, before the layer loop. On the attention side the publish happens during input preparation, before the model forward (afd_plugin/v1/worker/attention_metadata.py:115-116). On the FFN side the loop receives one payload per step, before it runs any layer (afd_plugin/v1/worker/ffn_worker.py:176). The data plane then runs once per layer and stage inside the forward.

Step 2: What the payload carries#

The payload is an AFDControlPayload: a map from stage index to DP metadata, plus the graph and warm-up flags (afd_plugin/connectors/metadata.py:107-150).

    dp_metadata_list: dict[int, AFDDPMetadata]
    is_graph_capturing: bool
    is_warmup: bool
    is_graph_replaying: bool = False
    is_profile: bool = False

Each stage's DP metadata is an AFDDPMetadata. Its one required field is the per-DP-rank token count, stored as an int32 CPU tensor (afd_plugin/connectors/metadata.py:23-44). This plugin-owned type keeps the wire format decoupled from vLLM's own DP metadata object (afd_plugin/connectors/metadata.py:320-328).

Step 3: The two attention publish paths#

The attention runner publishes the DP metadata on one of two paths, both once per step before the layer loop. For a DBO step, a prepare seam publishes one complete payload before the microbatch threads start (afd_plugin/v1/worker/attention_model_runner_v2.py:153-157).

        runner.install_mrv2_ubatch_metadata(
            prepared.slices, prepared.forward_contexts, input_batch.num_tokens
        )
        if runner.connector.control_plane is not None:
            runner.send_dp_metadata(None, prepared.slices)

For a non-DBO step, the metadata is attached to the forward context and sent there (afd_plugin/v1/worker/attention_metadata.py:252-257). Both paths end in the same call, which updates the local state and then sends (afd_plugin/v1/worker/attention_metadata.py:115-116).

        self.connector.control_plane.update_state_from_dp_metadata(payload)
        self.connector.control_plane.send_dp_metadata_list(payload)

Step 4: The wire format and the separate group#

The payload is encoded as compact JSON: only the per-stage token counts and the flags (afd_plugin/connectors/metadata.py:320-346). The send is two messages: a long size tensor, then a uint8 byte tensor (afd_plugin/connectors/metadata.py:376-401).

The GPU connector carries this over a dedicated NCCL process group, the "p2p" group, not the per-subgroup data communicators. The connector owns the group; the control plane uses it (afd_plugin/connectors/gpu/p2p.py:724-733, afd_plugin/connectors/gpu/p2p.py:853-883). Only the designated sender ranks transmit; the rest no-op (afd_plugin/connectors/gpu/p2p.py:870-875).

Step 5: The FFN side receives and derives the shapes#

The FFN loop blocks on recv_dp_metadata_list until its mapped attention sender transmits (afd_plugin/connectors/gpu/p2p.py:885-908). The FFN runner then calls update_state_from_dp_metadata once per step, before the layer loop (afd_plugin/v1/worker/ffn_model_runner.py:194-203).

        if update_connector_state:
            cast(
                AFDControlPlane, self.connector.control_plane
            ).update_state_from_dp_metadata(
                _make_dp_metadata_payload(
                    dp_metadata_list,
                    is_graph_capturing=is_graph_capturing,
                    is_warmup=is_warmup,
                ),
            )

update_state_from_dp_metadata derives the expected wire-tensor shape for each stage and, in graph mode, preallocates the receive buffers so capture and replay reuse stable addresses (afd_plugin/connectors/gpu/p2p.py:744-851).

Step 6: The async NPU connector has no control plane#

One connector sets control_plane = None: the async NPU connector (afd_plugin/connectors/npu/async_cam.py:283). It has no control plane; its receive loop drives the FFN side instead. See Ascend NPU overview: the same contract, different transports.

Deep dive#

What can go wrong#

• A GPU connector with no control plane. The GPU FFN loop needs a control-plane-driven connector and raises for one that sets control_plane = None (afd_plugin/v1/worker/ffn_worker.py:169-174).
• Fragments from different ranks. The receive raises if the size message and the byte message come from different ranks (afd_plugin/connectors/metadata.py:418-419).
• A stall before the first transfer. The FFN loop blocks on the control plane each step (afd_plugin/connectors/gpu/p2p.py:885-908). If the attention side never publishes, the FFN loop waits.

Why it matters. The DP metadata sets the wire shapes. Without it, the FFN side cannot size its receive buffers, and the per-layer transfers cannot start.