Tour: One decode step

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

        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

        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

        if desc.num_ubatches > 1:
            # Native execute_model already prepared live stage inputs and our
            # prepare seam sent their complete control payload.
            assert runner._afd_pending_metadata is not None
            assert runner._afd_pending_metadata.num_stages == desc.num_ubatches
            return original_run_fullgraph(desc)

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

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

        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

                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)