AFD Plugin Explained

The FFN runtime: a connector-driven server loop

AdvancedAFD conceptsabout 6 min read + 2 min deep dives

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

In one sentence

The FFN role runs a stand-alone runner, not a vLLM scheduler. A daemon loop reads the control plane once per step and drives the model.

The idea#

On the FFN role, AFD drops the vLLM serving loop. There is no real scheduler (only a no-op stand-in), no KV cache and no sampler. Instead, a daemon thread blocks on the control plane. When a payload arrives, it runs one step: for each decoder layer and stage, it receives hidden states, runs the FFN or experts, and sends the result back.

The runner is GPUFFNModelRunner, a stand-alone class, not a subclass of the vLLM runner (afd_plugin/v1/worker/ffn_model_runner.py:52). It keeps only what the FFN role needs: the model, the connector and an optional CUDA graph cache.

An everyday comparison. Think of a short-order kitchen with no front counter. Orders arrive on a ticket rail (the control plane). The cook reads one ticket, makes the dish, and sends it out. The comparison stops here: the "ticket" carries only token counts, and the dish is a tensor that goes back to the exact attention rank that asked.

The picture#

Figure: the daemon loop runs one step per control payload. Inside, _ffn_forward loops over layers and stages: receive, compute, send.
Open full size ↗

Read panel 1 as the outer loop in the worker. Panel 2 is the inner forward in the runner. The blue box is a receive from the attention side; the emerald boxes are FFN compute and the send back.

In the code#

Step 1: Construction keeps only the FFN parts#

The runner parses the AFD config, validates the V2 contract, resolves the graph policy and builds the connector (afd_plugin/v1/worker/ffn_model_runner.py:63-99). Like the attention runner, it requires a control-plane-driven connector and asserts it at construction (afd_plugin/v1/worker/ffn_model_runner.py:87-89).

afd_plugin/v1/worker/ffn_model_runner.py:80-96GitHub ↗
        self.connector = AFDConnectorFactory.create_connector(
            rank,
            local_rank,
            vllm_config,
            self.afd_config,
        )
        # TODO: Async GPU connector will be supported in the future
        assert self.connector.control_plane is not None, (
            "GPU model runner only supports control-plane-driven connectors"
        )

        self.model: Any | None = None
        self.model_memory_usage = 0
        self.num_layers = int(self.model_config.hf_text_config.num_hidden_layers)
        self.use_cuda_graph = bool(
            self.afd_cudagraph_policy.enable_ffn_graph_cache,
        )
Build the connector, assert a control plane, and read the layer count.

Step 2: load_model loads the FFN model (and only its weights)#

load_model calls the native model loader (afd_plugin/v1/worker/ffn_model_runner.py:111-126). The model it loads is the AFD model wrapper for the FFN role, which builds only the FFN modules and keeps only the FFN weights. See Model wrappers: the role-aware decoder layer for the weight filter.

afd_plugin/v1/worker/ffn_model_runner.py:111-126GitHub ↗
    def load_model(self, *, load_dummy_weights: bool = False, **kwargs: Any) -> None:
        """Load the vLLM model."""

        model_loader = get_model_loader(self.load_config)
        with DeviceMemoryProfiler() as profiler:
            if self.model is None:
                self.model = model_loader.load_model(
                    vllm_config=self.vllm_config,
                    model_config=self.model_config,
                )
            else:
                model_loader.load_weights(
                    self.model,
                    model_config=self.model_config,
                )
        self.model_memory_usage = profiler.consumed_memory
load_model: the native loader builds the FFN-role model.

Step 3: The server loop runs one step per control payload#

The worker starts the loop on a daemon thread (afd_plugin/v1/worker/ffn_worker.py:133-158). The thread reads one DP-metadata payload per step, then dispatches to graph capture or to execute (afd_plugin/v1/worker/ffn_worker.py:168-199).

afd_plugin/v1/worker/ffn_worker.py:168-199GitHub ↗
        while not event.is_set():
            if self.model_runner.connector.control_plane is None:
                raise NotImplementedError(
                    "GPU FFN only supports control-plane-driven connectors; "
                    "connectors without a control plane (control_plane is None) "
                    "are not supported.",
                )

            payload = self.model_runner.connector.control_plane.recv_dp_metadata_list()
            dp_metadata_list = payload.dp_metadata_list
            is_attn_graph_capturing = payload.is_graph_capturing
            is_warmup = payload.is_warmup
            is_graph_replaying = payload.is_graph_replaying

            if self.model_runner.use_cuda_graph and (
                is_warmup or is_attn_graph_capturing
            ):
                self.model_runner.capture_model(
                    dp_metadata_list=dp_metadata_list,
                    is_warmup=is_warmup,
                    is_attn_graph_capturing=is_attn_graph_capturing,
                )
            else:
                self.model_runner.execute_model(
                    dp_metadata_list=dp_metadata_list,
                    is_graph_capturing=is_attn_graph_capturing,
                    is_warmup=is_warmup,
                    is_graph_replaying=is_graph_replaying,
                )

            if self.device.type == "cuda":
                torch.cuda.synchronize()
One recv per step, then dispatch; a connector with no control plane is rejected.
  1. Lines 169-174. A connector with control_plane None raises; the GPU FFN needs a control plane.
  2. Line 176. One recv_dp_metadata_list per step.
  3. Lines 182-189. Graph warm-up or capture calls capture_model.
  4. Lines 190-196. Otherwise execute_model runs the step.

Step 4: _ffn_forward loops over layers and stages#

execute_model makes the graph key, then either replays a captured graph or calls _ffn_forward for the eager path (afd_plugin/v1/worker/ffn_model_runner.py:153-184). _ffn_forward loops over the layer indices, and over the stages inside each layer (afd_plugin/v1/worker/ffn_model_runner.py:222-269).

afd_plugin/v1/worker/ffn_model_runner.py:236-268GitHub ↗
                    payload = self.connector.recv_attn_output(
                        ubatch_idx=stage_idx,
                        **recv_kwargs,
                    )
                    hidden_states = payload.hidden_states
                    context = payload.context
                    metadata = context.metadata
                    metadata.layer_idx = layer_idx
                    metadata.stage_idx = stage_idx
                    if forward_context is not None:
                        forward_context.dp_metadata = ffn_dp_metadata_list[
                            metadata.stage_idx
                        ]
                        forward_context.additional_kwargs["afd_metadata"] = metadata
                        _set_moe_layer_index(forward_context, layer_idx)
                    if (
                        uses_remote_experts
                        and self.afd_config.compute_gate_on_attention
                    ):
                        router_logits = payload.router_logits
                        assert router_logits is not None
                        rank_ffn_output = model.compute_experts_output(
                            hidden_states,
                            layer_idx,
                            router_logits,
                        )
                    else:
                        rank_ffn_output = self._execute_eager_mode(
                            hidden_states,
                            layer_idx,
                            input_ids=payload.input_ids,
                        )
                    self.connector.send_ffn_output(rank_ffn_output, context)
For each layer and stage: receive, compute, send back.
  1. Line 236. recv_attn_output returns the hidden states and the transfer context.
  2. Lines 251-261. An expert layer with the gate on the attention side calls compute_experts_output with the router logits.
  3. Lines 262-267. Any other layer calls _execute_eager_mode, which calls compute_ffn_output.
  4. Line 268. send_ffn_output sends the result back to the attention rank.

Step 5: The FFN role never samples#

sample_tokens raises; the FFN role produces no tokens (afd_plugin/v1/worker/ffn_model_runner.py:405-406). The scheduler-driven execute_model on the worker also raises, because FFN work is connector-driven, not scheduler-driven (afd_plugin/v1/worker/ffn_worker.py:122-131).

Deep dive#

Deep dive

Eager vs graph path

execute_model resolves a run mode from the warm-up, capture and replay flags and whether a graph already exists (afd_plugin/v1/worker/ffn_model_runner.py:166-174).

ModeWhat runsLine
REPLAYthe captured CUDA graph replays; no recomputeafd_plugin/v1/worker/ffn_model_runner.py:175-177
eager / warm-up_ffn_forward runs the layers and stagesafd_plugin/v1/worker/ffn_model_runner.py:179-183
capturecapture_model records the graph into its own poolafd_plugin/v1/worker/ffn_model_runner.py:358-403

The graph cache is keyed by the per-stage token counts (afd_plugin/v1/worker/ffn_model_runner.py:166). AFD on GPU allows only FULL_DECODE_ONLY graphs. The CUDA graphs in AFD: FULL_DECODE_ONLY on both roles page reads capture and replay.

Deep dive

Worked example: one step on 2A2F with DBO

Take a 2A2F topology: two attention ranks (A0, A1) and two FFN ranks (F0, F1). The mapping sends A0 to F0 and A1 to F1 (afd_plugin/distributed/topology.py:135-136). Run with DBO, so there are two stages per layer.

Use DeepSeek-V2-Lite as the model. This model has 27 decoder layers (layer 0 dense, the rest MoE). The 27 and the dense/MoE split are inferred from the DeepSeek-V2-Lite checkpoint config; they are not stored in either repository. With compute_gate_on_attention=false, the FFN runs every layer.

  1. The control plane delivers one payload for the step, with one DP-metadata entry per stage (afd_plugin/v1/worker/ffn_worker.py:176).
  2. _ffn_forward sets layer_indices = range(27) because the gate is on the FFN side (afd_plugin/v1/worker/ffn_model_runner.py:211-215).
  3. For each of the 27 layers, the runner loops over the two stages (afd_plugin/v1/worker/ffn_model_runner.py:223-230).
  4. That is 27 x 2 = 54 receive/compute/send cycles on F0 (and 54 on F1) for the step. F0 serves only A0's tokens, because A0 maps to F0.

The data plane runs 54 times per FFN rank in this step; the control plane ran once. This is why AFD counts the control plane per step and the data plane per layer and stage.

Deep dive

Shutdown closes the connector

The worker stops the loop, then shuts the runner down (afd_plugin/v1/worker/ffn_worker.py:207-219). The runner's shutdown releases the graph state and native-equivalent GPU state, then closes the connector in a finally (afd_plugin/v1/worker/ffn_model_runner.py:442-456).

afd_plugin/v1/worker/ffn_model_runner.py:444-456GitHub ↗
        stop_afd_gpu_profiler(self.prof)
        try:
            for graph_info in self._cuda_graphs.values():
                graph_info["graph"].reset()
            self._cuda_graphs.clear()
            self._graph_memory_pool = None
            self.vllm_config.compilation_config.static_forward_context.clear()
            self.model = None
            _ROPE_DICT.clear()
            reset_workspace_manager()
        finally:
            self.connector.close()
        # ### PATCH END: release native-equivalent and AFD-owned GPU state.
shutdown: release GPU state, then connector.close().

What can go wrong#

Why it matters. The FFN side is passive: it does nothing until a control payload arrives. If the attention side does not publish, the loop blocks; if a connector has no control plane, the loop cannot start.

Check yourself#

  1. 1
    Does the FFN role have a KV cache?
    easy
    Show answer

    No. get_kv_cache_spec returns an empty dict and initialize_kv_cache is a no-op (afd_plugin/v1/worker/ffn_model_runner.py:147-151).

  2. 2
    What drives one FFN step?
    easy
    Show answer

    One recv_dp_metadata_list from the control plane, read by the daemon loop (afd_plugin/v1/worker/ffn_worker.py:176).

  3. 3
    Inside _ffn_forward, what is the order of the three connector and compute calls per layer and stage?
    medium
    Show answer

    recv_attn_output, then the compute (compute_experts_output or compute_ffn_output), then send_ffn_output (afd_plugin/v1/worker/ffn_model_runner.py:236-268).

  4. 4
    What happens when a captured graph exists for the step's token counts?
    medium
    Show answer

    The run mode is REPLAY: the captured CUDA graph replays and the runner returns without recompute (afd_plugin/v1/worker/ffn_model_runner.py:175-177).

  5. 5
    On 2A2F with DBO and a 27-layer model, how many data-plane cycles does one FFN rank run per step, and how many control-plane reads?
    hard
    Show answer

    27 layers x 2 stages = 54 receive/compute/send cycles on the data plane, and one control-plane read (afd_plugin/v1/worker/ffn_model_runner.py:223-230, afd_plugin/v1/worker/ffn_worker.py:176). The 27 is inferred from the DeepSeek-V2-Lite checkpoint.

  • The attention runtime: a seam over Model Runner V2The attention runner is a thin subclass of the native vLLM GPU Model Runner V2. It overrides a few seams for the connector and the control plane, and delegates the model forward, sampling and output to the native class.
  • The control plane: DP metadata once per stepBefore each layer loop, the attention side sends per-stage token counts to the FFN side. This page shows what the control plane carries, both publish paths, the wire format and the receive side.
  • The MoE handoff: where the gate runsThe compute_gate_on_attention flag decides whether the MoE gate runs on the attention side or the FFN side. The choice changes which classes build, what crosses the connector, and which layers the FFN runs.
  • v1/worker: the FFN worker and runnerThree files run the FFN role: the worker that starts the daemon loop, the stand-alone connector-driven runner, and the CPU-safe token-count helpers that map attention counts to FFN ranks.