The attention runtime: a seam over Model Runner V2

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

The idea#

On the attention role, AFD does not rewrite the vLLM model runner. It subclasses it. AFDAttentionModelRunnerV2 inherits the native GPU Model Runner V2 and an AFD metadata mixin (afd_plugin/v1/worker/attention_model_runner_v2.py:358). The native class keeps request state, input preparation, attention, the KV cache, sampling and the output. The AFD subclass adds only the connector lifecycle and the control-plane publish.

The runner's own docstring states this split: native V2 keeps execution, dummy runs and the graph lifecycle; the AFD methods are connector and context plumbing (afd_plugin/v1/worker/attention_model_runner_v2.py:358-365).

An everyday comparison. Think of a seat cover on a car seat. It adds a pocket and a strap, but the seat, the frame and the belt stay the same. The comparison stops here: the runner's "pocket" must open and close at the exact moment vLLM builds each forward context.

The picture#

Read panel 1 as a thin layer over the native class. Every override calls super() for the real work. Panel 2 shows the one control-plane publish per step, before the layer loop; the data plane then runs once per layer and stage.

In the code#

Step 1: Construction makes the connector and checks for a control plane#

The constructor validates the V2 contract, calls super().__init__, then builds the connector through the factory (afd_plugin/v1/worker/attention_model_runner_v2.py:369-411). It requires a control-plane-driven connector and raises otherwise (afd_plugin/v1/worker/attention_model_runner_v2.py:394-397).

            self.afd_config = self.parse_config(self.vllm_config)
            rank, local_rank = _resolve_world_ranks()
            connector = AFDConnectorFactory.create_connector(
                rank,
                local_rank,
                self.vllm_config,
                self.afd_config,
            )
            self.connector = connector
            # The connector rendezvous is deferred to the end of ``load_model()``
            # so Attention and FFN weight loading overlap, matching the V1
            # Attention runner lifecycle.
            if connector.control_plane is None:
                raise RuntimeError(
                    "AFD ModelRunnerV2 requires a control-plane-driven connector",
                )

Step 2: load_model defers the connector rendezvous to the end#

load_model calls native loading first, then initializes the connector. The comment explains the order: the rendezvous waits until weights are loaded, so attention and FFN weight loading overlap (afd_plugin/v1/worker/attention_model_runner_v2.py:454-462).

    def load_model(
        self,
        load_dummy_weights: bool = False,
        *args,
        **kwargs,
    ) -> None:
        super().load_model(load_dummy_weights, *args, **kwargs)
        if not self.connector.is_initialized:
            self.connector.init_afd_connector()

Step 3: execute_model wraps native V2 in the AFD context#

execute_model steps the profiler, opens _use_afd_execution_context, then delegates the whole forward to super().execute_model (afd_plugin/v1/worker/attention_model_runner_v2.py:569-595).

        # ### PATCH START: scope AFD metadata provider/replay and profiler step.
        step_afd_gpu_profiler(self.prof)
        with _use_afd_execution_context(
            self,
            int(scheduler_output.total_num_scheduled_tokens),
        ):
            return super().execute_model(
                scheduler_output,
                intermediate_tensors,
                dummy_run=dummy_run,
                skip_attn_for_dummy_run=skip_attn_for_dummy_run,
                is_profile=is_profile,
                context_len=context_len,
                valid_dummy_state_slots=valid_dummy_state_slots,
            )
        # ### PATCH END: scope AFD metadata provider/replay and profiler step.

The forward is native, so logits and sampling are native too. The delegation at line 586 runs the model forward (afd_plugin/v1/worker/attention_model_runner_v2.py:586). Native V2 does the sampling in a separate method, sample_tokens, after the forward (vllm:vllm/v1/worker/gpu/model_runner.py:1997, vllm:vllm/v1/worker/gpu/model_runner.py:2048). The attention runner does not override sample_tokens, so it samples natively. The FFN runner overrides it to raise (afd_plugin/v1/worker/ffn_model_runner.py:405-406). Sampling runs once per step, after the last layer, on the attention side.

Step 4: The AFD metadata provider installs the sidecar#

Model wrappers read AFD metadata from the vLLM forward context. The provider installs it as vLLM makes each context. The plugin cannot extend set_forward_context, so it wraps create_forward_context for the scope and mutates additional_kwargs (afd_plugin/model_executor/models/forward_context.py:37-67).

    original_create = forward_context_module.create_forward_context

    @wraps(original_create)
    def create_forward_context_with_afd(*args: Any, **kwargs: Any) -> ForwardContext:
        forward_context = original_create(*args, **kwargs)
        installer(forward_context)
        return forward_context

    forward_context_module.create_forward_context = create_forward_context_with_afd
    try:
        yield
    finally:
        forward_context_module.create_forward_context = original_create

The installer and the publish logic live in the shared mixin, AFDMetadataProviderMixin (afd_plugin/v1/worker/attention_metadata.py:21). Both the V2 runner and the NPU runners reuse it.

Step 5: The control plane publishes once per step#

The DP metadata crosses once per step, before the layer loop. One of two seams sends it. With DBO, the ubatch prepare seam sends it after it stages the microbatches (afd_plugin/v1/worker/attention_model_runner_v2.py:153-158).

        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)
        return prepared

Without DBO, the forward-context provider sends it once, unless a graph lifecycle already pre-sent it (afd_plugin/v1/worker/attention_metadata.py:248-257). Both paths end in send_dp_metadata, which updates the connector state and sends the payload (afd_plugin/v1/worker/attention_metadata.py:115-116).

Deep dive#

What can go wrong#

• A connector with no control plane. The runner raises at construction; the GPU attention runner needs a control plane (afd_plugin/v1/worker/attention_model_runner_v2.py:394-397).
• A KV-only context with no batch descriptor. The runner skips the AFD install for such a context, so it does not start a spurious transaction (afd_plugin/v1/worker/attention_model_runner_v2.py:434-443).
• Construction failure. If any step after super().__init__ raises, the runner closes the connector and shuts down native state, then re-raises (afd_plugin/v1/worker/attention_model_runner_v2.py:405-411).
• Missing graph manager in FULL_DECODE_ONLY. The replay hook raises if the graph manager is not initialized (afd_plugin/v1/worker/attention_model_runner_v2.py:244-248).

Why it matters. The runner must publish the DP metadata before the layers start, and attach the sidecar on every context. If a seam misses, the FFN side waits for a payload that never arrives, or a layer runs with no connector in scope.