CUDA graphs in AFD: FULL_DECODE_ONLY on both roles

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

The idea#

A CUDA graph records a fixed sequence of GPU work once, then replays it with almost no CPU cost. AFD splits one model across two processes, so each role records its own graph. The attention side and the FFN side capture and replay separately. They share nothing but the tensors that cross the connector.

AFD allows one graph mode: FULL_DECODE_ONLY. In this mode the graph covers pure-decode steps, where every request adds one token. Prefill and mixed batches run eager (no graph). This is set in one small policy module (afd_plugin/v1/worker/cuda_graph.py:20-21).

An everyday comparison. Think of a stage play with two actors in separate rooms, linked by a hatch. Each actor rehearses their own lines until they are automatic (capture), then performs them fast (replay). The comparison stops here: the two actors must hand the exact same size of parcel through the hatch every time, or the fast performance breaks.

The picture#

Read the top panel as the attention role and the bottom panel as the FFN role. The attention side does not own the graph, native vLLM does; AFD only publishes control around it. The FFN side owns its graphs in a plain dictionary, one per token shape.

In the code#

Step 1: One supported mode#

The policy module lists the supported modes. The set holds one name (afd_plugin/v1/worker/cuda_graph.py:20-21). validate_cuda_graph_mode returns a disabled policy when enforce_eager is set, and raises for any mode that is not FULL_DECODE_ONLY (afd_plugin/v1/worker/cuda_graph.py:42-85).

def validate_cuda_graph_mode(
    vllm_config: VllmConfig,
    *,
    role: str | None = None,
) -> AFDCUDAGraphPolicy:
    """Return the CUDA graph policy or raise for unsupported AFD modes."""

    enforce_eager = bool(getattr(vllm_config.model_config, "enforce_eager", False))
    mode_name = cudagraph_mode_name(vllm_config)
    graph_enabled = not enforce_eager

    if not graph_enabled:
        return AFDCUDAGraphPolicy(
            enabled=False,
            mode_name=mode_name,
            allow_attention_full_decode_only=False,
            enable_ffn_graph_cache=False,
        )

    if mode_name not in _SUPPORTED_GRAPH_MODES:
        role_suffix = f" for {role}" if role else ""
        raise RuntimeError(
            "AFD only supports CUDA graph mode "
            f"{FULL_DECODE_ONLY}{role_suffix}; got {mode_name!r}.",
        )

On Ascend NPU, the counterpart is an ACL graph; the NPU validator accepts the FULL and FULL_DECODE_ONLY modes (afd_plugin/validation.py:233-239). The NPU runtime is covered on Ascend NPU overview: the same contract, different transports.

Step 2: The attention side wraps native capture#

The AFD attention runner does not capture graphs itself. It wraps the native capture_model and publishes one warm-up payload and one capture payload for each FULL descriptor (afd_plugin/v1/worker/attention_model_runner_v2.py:509-553). An event tracker checks that native capture prepares inputs exactly twice per descriptor, in order (afd_plugin/v1/worker/attention_model_runner_v2.py:55-97).

        try:
            with (
                _use_afd_capture_input_preparation(self, event_tracker),
                _use_afd_ubatch_preparation(self, event_tracker),
            ):
                with use_afd_metadata_provider(
                    self.install_afd_metadata_on_forward_context,
                ):
                    result = super().capture_model(profile_only=profile_only)
                event_tracker.assert_complete()
            return result

Step 3: The attention side publishes control before each replay#

Native FULL replay does not build a forward context, so the AFD metadata provider cannot run there. The runner wraps the graph manager's run_fullgraph and sends one padded control payload just before each native replay (afd_plugin/v1/worker/attention_model_runner_v2.py:270-303). This hook is used only in FULL_DECODE_ONLY mode, and only when a graph manager exists (afd_plugin/v1/worker/attention_model_runner_v2.py:323-327).

        try:
            padded_tokens = int(desc.num_tokens)
            metadata = runner.build_afd_metadata(None, real_tokens)
            metadata.tokens_lens = [padded_tokens]
            runner._afd_pending_metadata = metadata
            runner._afd_suppress_metadata_send = True
            runner._is_warmup = False
            runner._afd_is_graph_capturing = False
            runner._afd_is_graph_replaying = True
            runner.send_dp_metadata(
                runner.build_capture_dp_metadata(padded_tokens),
                None,
            )
            result = original_run_fullgraph(desc)
        finally:
            runner._afd_is_graph_replaying = previous_is_graph_replaying
        # ### PATCH END: publish one AFD pre-replay payload.
        return result

Step 4: The FFN side keys one graph per shape#

The FFN runner builds a hashable FFN graph key from the per-stage DP metadata, then looks up a stored graph (afd_plugin/v1/worker/ffn_model_runner.py:166-174). On a hit in replay mode it calls graph.replay() and returns (afd_plugin/v1/worker/ffn_model_runner.py:175-177).

        graph_key = make_ffn_graph_key(dp_metadata_list)
        cuda_graph_info = self._cuda_graphs.get(graph_key)
        run_mode = graph_run_mode(
            is_warmup=is_warmup,
            is_graph_capturing=is_graph_capturing,
            is_graph_replaying=is_graph_replaying,
            graph_enabled=bool(self.use_cuda_graph),
            graph_exists=cuda_graph_info is not None,
        )
        if run_mode is AFDGraphRunMode.REPLAY:
            self._cuda_graphs[graph_key]["graph"].replay()
            return None

        self._ffn_forward(
            dp_metadata_list=dp_metadata_list,
            is_graph_capturing=is_graph_capturing,
            is_warmup=is_warmup,
        )
        return None

The run mode is chosen by a small helper. It returns REPLAY only when the graph is enabled and a graph for this key exists (afd_plugin/v1/worker/cuda_graph.py:136-150).

Step 5: The FFN side captures its own graph#

The FFN server loop reads the control payload, then calls capture_model for warm-up and capture steps, or execute_model otherwise (afd_plugin/v1/worker/ffn_worker.py:182-196). Capture runs _dummy_run, which records one torch.cuda.CUDAGraph into the runner's own graph pool and stores it under the key (afd_plugin/v1/worker/ffn_model_runner.py:316-351).

        if mode_name == "FULL":
            if self._graph_memory_pool is None:
                self._graph_memory_pool = torch.cuda.graph_pool_handle()
            graph_key = make_ffn_graph_key(dp_metadata_list)
            cudagraph = torch.cuda.CUDAGraph()
            # DP metadata receive/update is a control-plane side effect and must
            # complete before CUDA graph capture starts.
            cast(
                AFDControlPlane, self.connector.control_plane
            ).update_state_from_dp_metadata(
                _make_dp_metadata_payload(
                    dp_metadata_list,
                    is_graph_capturing=is_attn_graph_capturing,
                ),
            )
            with torch.cuda.graph(cudagraph, pool=self._graph_memory_pool):
                output = self._ffn_forward(
                    dp_metadata_list=dp_metadata_list,
                    is_graph_capturing=is_attn_graph_capturing,
                    update_connector_state=False,
                )
            self._cuda_graphs[graph_key] = {
                "graph": cudagraph,
                "input_hidden_states": output,
                "output": output,
            }

Step 6: Custom ops keep the transfers capturable#

The data path uses vLLM's PyNcclCommunicator send and receive, wrapped in four custom ops (afd_plugin/connectors/gpu/p2p.py:979-1122). The wrapping keeps the transfers traceable by torch.compile and capturable by a CUDA graph (afd_plugin/connectors/gpu/p2p.py:36-42).

Deep dive#

The blocks below give the capture handshake, the FFN graph cache, why the two sides must agree on shapes, and a worked example.

What can go wrong#

• An unsupported graph mode. Any mode other than FULL_DECODE_ONLY raises at config validation (afd_plugin/v1/worker/cuda_graph.py:61-66).
• Ubatching with the wrong count. A graph with ubatching needs exactly two ubatches, or the policy raises (afd_plugin/v1/worker/cuda_graph.py:72-77).
• A shape that no captured key covers. The FFN side does not crash; it runs that step eager, because the run mode falls back to EAGER (afd_plugin/v1/worker/cuda_graph.py:148-150).
• Capture order drift. If native capture prepares inputs the wrong number of times, the event tracker raises, so a silent mismatch cannot happen (afd_plugin/v1/worker/attention_model_runner_v2.py:73-88).

Why it matters. A CUDA graph replays fixed addresses. The two roles must agree on the token shapes, or the FFN side cannot reuse a graph. AFD keeps the shapes stable by limiting graphs to pure-decode steps and keying each graph on the DP metadata.