Tour: Life of a request

Attention and FFN may be started in either order. Send requests only to the
Attention API server. FFN workers are connector-driven; scheduler-driven FFN
`execute_model()` calls fail fast.

    def init_device(self):
        """Initialize the native GPU worker with the selected AFD runner."""

        # ### PATCH START: select and directly construct the AFD runner.
        assert_compatible_afd_stack(
            self.vllm_config,
            caller="AFDAttentionWorker.init_device",
            expected_role="attention",
        )
        validate_gpu_model_runner_v2_config(
            self.vllm_config,
            expected_role="attention",
            device_type=self.device_config.device_type,
        )

        afd_model_config = get_afd_model_config(
            self.vllm_config.model_config,
            device_type="cuda",
        )
        self.vllm_config.model_config = afd_model_config
        self.model_config = afd_model_config

        native_runner_cls = gpu_model_runner_v2.GPUModelRunner
        gpu_model_runner_v2.GPUModelRunner = AFDAttentionModelRunnerV2
        try:
            super().init_device()
        finally:
            gpu_model_runner_v2.GPUModelRunner = native_runner_cls
        # ### PATCH END: select and directly construct the AFD runner.

        # ### 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,
            )

    try:
        with (
            use_gpu_mrv2_dbo_dispatch(),
            replay_scope,
            _use_afd_ubatch_preparation(runner),
            use_afd_metadata_provider(
                runner.install_afd_metadata_on_forward_context,
            ),
        ):
            yield

        payload = AFDControlPayload(
            dp_metadata_list=dp_metadata_list,
            is_graph_capturing=is_graph_capturing,
            is_warmup=is_warmup,
            is_graph_replaying=is_graph_replaying,
            is_profile=self._afd_is_profile,
        )
        self.connector.control_plane.update_state_from_dp_metadata(payload)
        self.connector.control_plane.send_dp_metadata_list(payload)

class RemoteFFNProxy(nn.Module):
    """Parameter-free FFN stage executed through the AFD connector."""

    def __init__(self, *, layer_idx: int) -> None:
        super().__init__()
        self.layer_idx = layer_idx

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        return self._send_and_receive(hidden_states)

        metadata = AFDTransferMetadata.create_attention_metadata(
            layer_idx=self.layer_idx,
            stage_idx=stage_idx,
            seq_len=int(hidden_states.shape[0]),
        )
        context = AFDTransferContext(metadata=metadata)
        afd_metadata.connector.send_attn_output(
            hidden_states,
            context,
            **send_kwargs,
        )
        hidden_states = maybe_apply_dbo_yield(
            hidden_states,
            role="attention",
        )
        return afd_metadata.connector.recv_ffn_output(
            ref_tensor=hidden_states,
            ubatch_idx=stage_idx,
        )

        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,
                )

                for stage_idx in stage_ids:
                    recv_kwargs: dict[str, Any] = {}
                    if routing_spec is not None:
                        recv_kwargs["routing_spec"] = routing_spec
                    if recv_input_ids:
                        recv_kwargs["recv_input_ids"] = True
                    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)
        return rank_ffn_output

    def sample_tokens(self, grammar_output: Any = None) -> Any:
        raise RuntimeError("FFN runners do not sample tokens")