Tour: One step on many GPUs

            bool is_any_point = (decode_step != 0) && (decode_step % prefill_any_attempt_interval == 0);
            bool is_all_point = (decode_step % prefill_attempt_interval == 0);
            if (!is_any_point && !is_all_point) {
                return _schedule_decode(acquire);
            }
            bool want = (get_gpu_cache_usage() <= GPU_CACHE_PREFILL_WATERMARK) && _has_prefill_waiting();
            bool do_prefill;
            if (is_any_point) {
                do_prefill = (_dp_allreduce(want ? 0 : 1, c10d::ReduceOp::MIN) == 0); // ANY：任一 rank 想做即触发
            } else {
                do_prefill = (_dp_allreduce(want ? 1 : 0, c10d::ReduceOp::MIN) == 1); // ALL：全员想做才触发
            }

    def _execute_model_async_step(
        self, seq_metadata_list: list[SequenceMetadata], scheduler_outputs: SchedulerOutputs
    ) -> np.ndarray[np.uint64]:
        with cpu_timer.record("run_workers"):
            request_outputs = self._run_workers(
                "execute_model_step",
                driver_kwargs=dict(
                    seq_metadata_list=seq_metadata_list,
                    scheduler_outputs=scheduler_outputs,
                ),
                is_prompt=scheduler_outputs.prompt_run,
                is_placeholder=scheduler_outputs.is_empty(),
            )
        self.deferred_release.collect_request_outputs(request_outputs)
        return request_outputs

    def run_workers(
        self,
        method: str,
        driver_kwargs: dict[str, Any] | None = None,
        **kwargs: Any,
    ) -> Any:
        with cpu_timer.record("broadcast_method", enable=False):
            assert self.is_driver_worker
            cp_group = PG.cpu_cp_group
            if cp_group and cp_group.size() > 1:
                cp_broadcast_object((method, kwargs))

        return getattr(self, method)(**(driver_kwargs or {}), **kwargs)

        if PG.context_parallel_group().size() > 1:
            with cpu_timer.record("broadcast_input_metadata"):
                input_metadata_dict = {
                    "is_prompt": is_prompt,
                    "type": type(input_metadata),
                    "run_cuda_graph": run_cuda_graph,
                    "total_q": total_q,
                    "num_reqs": num_reqs,
                }
                if not run_cuda_graph:
                    assert hasattr(input_metadata, "is_prompt")
                    input_metadata_dict.update(input_metadata.to_flatten_dict())
                broadcast_tensor_dict(input_metadata_dict, src=0)

        return input_metadata, total_q, run_cuda_graph, cuda_graph_used, graph_size

    def _prepare_input_follower(self) -> tuple[InputMetadata, int, bool, CudaGraphData | None, int]:
        input_metadata_dict = broadcast_tensor_dict(src=0)
        # 只 pop 非 InputMetadata 字段的 key；is_prompt/total_q/num_reqs 本身也是字段，留给 from_flatten_dict。
        input_metadata_type: type[InputMetadata] = input_metadata_dict.pop("type")
        run_cuda_graph: bool = input_metadata_dict.pop("run_cuda_graph")
        total_q = input_metadata_dict["total_q"]
        num_reqs = input_metadata_dict["num_reqs"]
        if run_cuda_graph:
            graph_size, cuda_graph_used = self.cuda_graph_runner.select(total_q)
            input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)
        else:
            input_metadata = input_metadata_type.from_flatten_dict(input_metadata_dict)
            cuda_graph_used = None
            graph_size = 0
        return input_metadata, total_q, run_cuda_graph, cuda_graph_used, graph_size

        if hidden_states is None:
            hidden_states = self.embed_tokens(input_metadata.input_ids)

        residual = None
        # 跨层持有的运行时值 (shared_hidden / pre_rope_kv / topk): 逐层传同一引用，由各层原地写入。
        cross_layer_runtime = CrossLayerRuntime()

        if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
            input_metadata.prepare_for_mirror_cp()
            hidden_states = gather(hidden_states, input_metadata.mirror_cp_index)

        use_cp = input_metadata.is_prompt and PG.context_parallel_group().size() > 1
        shared_hidden_bf16 = self._resolve_shared_hidden(hidden_states_bf16, shared_hidden_bf16)
        projected_kv, index_weights = self._project_current_hidden_inputs(hidden_states_bf16)
        yoco_full_kv_stored = cross_layer_runtime.consumer_full_metadata is not None
        if yoco_full_kv_stored:
            # YOCO compact prefill: the query has been trimmed, but the KV still has to cover the full context.
            # Here the full shared hidden is used to generate and write the main/index KV, and the later main path skips the KV update.
            # With CP > 1 an inner cp session is opened here (before the main session below, with no nesting).
            self._store_yoco_full_kv(cross_layer_runtime, kv_cache)
        elif self.shared_hidden_kv_proj is not None:
            # The full KV update path above was not taken and this layer needs to generate KV from the
            # shared hidden (decode / non-compact prefill). Only the KV projection is done here; afterwards
            # _compute_kv_ca_index calls _store_main_kv to do the compression and the write.
            assert projected_kv is None and shared_hidden_bf16 is not None
            projected_kv = self.shared_hidden_kv_proj(shared_hidden_bf16)

        with userbuffer.cp_session(enabled=use_cp):
            # CP: launch the all_gather of the compressor input first, overlapping it with the Q projection below.
            # Under the yoco_full_kv_stored scenario the main branch does not update kv, so this communication step is not needed
            wait_projected_kv = _make_cp_gather_getter(projected_kv, use_cp) if not yoco_full_kv_stored else None

            q_a, kv_swa = self._project_q_and_swa(
                hidden_states,
                shared_hidden_bf16,
            )
            del hidden_states, hidden_states_bf16, shared_hidden_bf16
            kv_swa, wait_kv_swa = self._prepare_swa_kv(
                kv_swa,
                kv_cache.swa_cache,
                input_metadata,
                use_cp,
            )

    def forward(
        self,
        hidden_states: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
        hidden_states_bf16: torch.Tensor,
        input_metadata: InputMetadata,
    ) -> torch.Tensor:
        if self.use_mega_moe:
            return self.mega_moe_forward(hidden_states, hidden_states_bf16, input_metadata)

        topk_idx, topk_weights = self.gate_forward(hidden_states_bf16, input_metadata)

        try:
            # Dispatch
            event_overlap = self.dispatch(
                hidden_states,
                topk_idx,
                topk_weights,
                input_metadata,
                None,
            )
            moe_in_hidden, _dispatched_topk_idx, dispatched_topk_weights, handle = event_overlap.current_stream_wait()

            # MoE compute
            with cuda_graph_timer.record(
                "routed_experts",
                enabled=(self.layer_idx == 1 and not input_metadata.is_prompt),
                scale=eplb.current().num_moe_layers,
            ):
                moe_out_hidden = self.routed_experts_forward(
                    moe_in_hidden,
                    dispatched_topk_weights,
                    handle.psum_num_recv_tokens_per_expert,
                    expected_m=None,
                )

            # Combine
            combine_pre_evt = self.ep_buffer.capture()
            with limit_num_sms(envs.DEEPINFER_USERBUFFER_COMM_SMS, enabled=PLATFORM.is_gpu):
                shared_out = self.shared_forward(hidden_states)
            event_overlap = self.combine(
                moe_out_hidden,
                shared_out,
                handle,
                combine_pre_evt,
            )
            moe_combined_hidden, _combined_topk_weights = event_overlap.current_stream_wait()

        expert_lb = eplb.current()
        if expert_lb is not None:  # dense 模型没有 EPLB
            expert_lb.step(input_metadata.is_prompt)

            def _moe_lb_post_func():
                sampler_output.moe_load_balance_ratio = expert_lb.take_moe_load_balance_ratio()

            sampler_output.post_funcs.append(_moe_lb_post_func)

    def step(self, is_prompt: bool) -> None:
        """由 DeepseekForCausalLM.sample 每 step 调用，更新重排与负载统计。"""
        if self._ep_group.size() <= 1:
            return

        self._update_rearrangement(is_prompt)
        self._update_load_balance_stats()
        self.expert_count.zero_()

        if input_metadata.is_prompt and PG.context_parallel_group().size() > 1:
            with userbuffer.cp_session():
                gathered_hidden_states = userbuffer.cp.all_gather(hidden_states)()
                hidden_states = gather(gathered_hidden_states, input_metadata.mirror_cp_reverse_index)

        logits = get_logits(hidden_states, embedding, self.vocab_size)
        del hidden_states

        if driver_inputs is None:
            if need_prompt_logprobs:
                # 结果只在 driver 上有效；follower 只参与 embedding gather 集合通信。
                compute_prompt_logprobs(
                    prompt_hidden_states,
                    embedding,
                    None,
                    vocab_size=self.vocab_size,
                    force_topk=self.force_topk,
                )
            if self.max_MTP_value > 0 and not is_prompt:
                mtp_module.decode_forward(
                    main_input_metadata=main_input_metadata,
                    nextn_kv_caches=nextn_kv_caches,
                    decode_input=MTPDecodeInput(first_token_positions=decode_inputs.mtp_first_token_positions, driver=None),
                )
            return SamplerOutput(post_funcs=[])

        with cpu_timer.record("worker.post_process.sync"):
            step_state.sample_end_evt.synchronize()
        with cpu_timer.record("worker.post_process.graph_timer_flush"):
            # flush the time of completed graph, no effect on the running graph.
            cuda_graph_timer.flush_out()
        with cpu_timer.record("worker.post_process.do_process"):
            self.model_runner.post_process(step_state.seq_metadata_list, step_state.sampler_output, step_state.is_prompt)
            request_outputs = np.empty(0, dtype=np.uint64)
            if self.is_driver_worker:
                request_outputs = self.engine_callbacks.post_process_fn(step_state)
        return request_outputs