v1/worker: the FFN worker and runner

AdvancedCode mapabout 4 min read + 1 min deep dives

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

In one sentence

The FFN role is a native vLLM worker that starts a daemon loop, a stand-alone runner driven by the control plane, and small token-count helpers.

Purpose#

These three files run the FFN side with no real scheduler (a no-op stand-in) and no KV cache. The concept is on The FFN runtime: a connector-driven server loop; this page maps the files.

Item	Value
Files (lines, origin)	`ffn_worker.py` (226), `ffn_model_runner.py` (483), `ffn_metadata.py` (93); all AFD-owned
Main classes	`AFDFFNWorker` (`afd_plugin/v1/worker/ffn_worker.py:32`), `GPUFFNModelRunner` (`afd_plugin/v1/worker/ffn_model_runner.py:52`)
Runs in	each FFN worker process; the loop runs on a daemon thread; compute on the current CUDA stream
Called by	vLLM, through the rewritten `worker_cls` (`afd_plugin/validation.py:290-291`); then the control plane drives the loop
Calls	native `Worker` and the model loader, the connector, the model wrapper's FFN compute
Not in this file	the attention runtime (v1/worker: the attention worker and runner), the connector transport (connectors/gpu/p2p.py: the GPU P2P connector), the model split (model_executor/models/deepseek_v2.py: the DeepSeek wrapper)

Where it sits#

Who calls it#

The AFD config patch rewrites worker_cls to AFDFFNWorker when the role is FFN (afd_plugin/validation.py:280-296). The FFN EngineCore returns after the executor is built, so no scheduler drives this worker; the loop is driven by the control plane instead (afd_plugin/v1/worker/ffn_worker.py:133-158).

What it calls#

The worker swaps the runner class and delegates to native Worker.init_device (afd_plugin/v1/worker/ffn_worker.py:73-100). The runner builds the connector (afd_plugin/connectors/factory.py:49-68) and calls the model wrapper's compute_ffn_output or compute_experts_output.

Key classes and functions#

Name	Line	Role
`AFDFFNWorker.init_device`	`afd_plugin/v1/worker/ffn_worker.py:73-100`	swap in the FFN runner for the native window
`initialize_from_config`	`afd_plugin/v1/worker/ffn_worker.py:107-113`	skip KV cache, init connector, start the loop
`_run_ffn_server_loop`	`afd_plugin/v1/worker/ffn_worker.py:160-199`	one recv per step, then dispatch
`GPUFFNModelRunner.execute_model`	`afd_plugin/v1/worker/ffn_model_runner.py:153-184`	replay a graph, or run `_ffn_forward`
`_ffn_forward`	`afd_plugin/v1/worker/ffn_model_runner.py:186-269`	loop layers x stages: receive, compute, send
`sample_tokens`	`afd_plugin/v1/worker/ffn_model_runner.py:405-406`	raises; the FFN role never samples
`aggregate_ffn_token_counts`	`afd_plugin/v1/worker/ffn_metadata.py:15-60`	sum attention counts per FFN rank

Walk through the main methods#

_run_ffn_server_loop: one step per control payload#

The loop runs until the shutdown event is set. Each pass reads one DP-metadata payload, then dispatches to capture or execute (afd_plugin/v1/worker/ffn_worker.py:160-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()

The loop: reject a connector with no control plane, recv once, then dispatch.

Lines 169-174. A connector with control_plane None raises; the GPU FFN needs one.
Line 176. One recv_dp_metadata_list per step.
Lines 182-189. Warm-up or graph capture calls capture_model.
Lines 190-196. Otherwise execute_model runs the step; then the loop synchronizes.

_ffn_forward: layers x stages#

The forward sets the layer set from the gate flag, then loops layers and stages (afd_plugin/v1/worker/ffn_model_runner.py:207-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)

recv_attn_output, then compute, then send_ffn_output.

Line 236. recv_attn_output returns the hidden states and the context.
Lines 251-261. An attention-side-gate expert layer calls compute_experts_output.
Lines 262-267. Any other layer calls compute_ffn_output through _execute_eager_mode.
Line 268. send_ffn_output sends the result back.

aggregate_ffn_token_counts: attention counts to FFN ranks#

Each FFN rank sums the contiguous block of attention ranks mapped to it (afd_plugin/v1/worker/ffn_metadata.py:15-60). Every attention peer contributes at least one row, so a zero-token peer still adds a placeholder.

afd_plugin/v1/worker/ffn_metadata.py:49-60GitHub ↗

    # Each FFN rank sums the block that ``build_rank_mapping`` assigns to it.
    return tuple(
        sum(
            max(1, int(expanded_counts[attention_rank]))
            if attention_rank < len(expanded_counts)
            else fallback_count
            for attention_rank in subgroup_attention_block(
                ffn_rank, attention_size, ffn_size
            )
        )
        for ffn_rank in range(ffn_size)
    )

Sum the block that build_rank_mapping assigns to each FFN rank.

Deep dive#

Deep dive

Why the worker and runner both reject a no-control-plane connector

The GPU FFN path supports only control-plane-driven connectors. The runner asserts it at construction (afd_plugin/v1/worker/ffn_model_runner.py:87-89), and the loop re-checks it each pass (afd_plugin/v1/worker/ffn_worker.py:169-174). The async NPU connector sets control_plane = None and drives the FFN side through its own receive loop; that path is on Ascend NPU overview: the same contract, different transports.