Repository: vllm-project/vllm
File: vllm/v1/worker/gpu/model_runner.py (Python)

Code before:
        assert self.execute_model_state is not None
        hidden_states, input_batch, kv_connector_output = self.execute_model_state
        self.execute_model_state = None  # type: ignore

        # Non-last PP rank: receive tokens and update state.
        # Don't need to sample tokens here - just receive from last rank.
        if hidden_states is None:
            received = self.pp_handler.maybe_receive_sampled_tokens(
                input_batch.num_reqs,
                self.device,
                max_sample_len=self.num_speculative_steps + 1,
            )
            if received is not None:
                sampled, num_sampled, num_rejected = received
                self.postprocess(input_batch, sampled, num_sampled, num_rejected)
            return None

        # Last rank: sample tokens
        sampler_output, num_sampled, num_rejected = self.sample(
            hidden_states, input_batch, grammar_output
        )

        # Broadcast to non-last ranks (handles spec decode multi-token)
        self.pp_handler.maybe_broadcast_sampled_tokens(
            sampler_output, num_sampled, num_rejected
        )

        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
            self.req_states.prefill_token_ids.gpu,
            self.req_states.num_computed_tokens.gpu,
            self.req_states.prompt_len,
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )

        # Prepare the model runner output.
        model_runner_output = ModelRunnerOutput(
            req_ids=input_batch.req_ids,
            # NOTE(woosuk): req_id_to_index is unused in this model runner.
            # Only for compatibility with the existing model runner and scheduler.
            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
            sampled_token_ids=None,  # type: ignore
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]
            kv_connector_output=kv_connector_output,
        )
        async_output = AsyncOutput(
            model_runner_output=model_runner_output,
            sampler_output=sampler_output,

Review comment (suggestion):
I think it'd be nice if we can do something like
```suggestion
        if self.use_pp:
            # Broadcast to non-last ranks (handles spec decode multi-token)
            self.pp_handler.maybe_broadcast_sampled_tokens(
                sampler_output, num_sampled, num_rejected
            )
```

Code after:
    ) -> AsyncOutput | ModelRunnerOutput | None:
        assert self.execute_model_state is not None
        hidden_states, input_batch, kv_connector_output = self.execute_model_state
        self.execute_model_state = None  # type: ignore

        # Non-last PP rank: hidden_states is None because this rank produced
        # IntermediateTensors instead of final hidden states. Receive the
        # sampled tokens broadcast by the last rank and update local state.
        if self.use_pp and not get_pp_group().is_last_rank:
            assert self.pp_handler is not None
            received = self.pp_handler.maybe_receive_sampled_tokens(
                input_batch.num_reqs,
                self.device,
                max_sample_len=self.num_speculative_steps + 1,
            )
            if received is not None:
                sampled, num_sampled, num_rejected = received
                self.postprocess(input_batch, sampled, num_sampled, num_rejected)
            return None

        # Last rank: sample tokens
        sampler_output, num_sampled, num_rejected = self.sample(
            hidden_states, input_batch, grammar_output
        )

        # Broadcast to non-last PP ranks (handles spec decode multi-token).
        if self.use_pp:
            assert self.pp_handler is not None
            self.pp_handler.maybe_broadcast_sampled_tokens(
                sampler_output, num_sampled, num_rejected
            )

        prompt_logprobs_dict = self.prompt_logprobs_worker.compute_prompt_logprobs(
            self.model.compute_logits,
            hidden_states,
            input_batch,
            self.req_states.prefill_token_ids.gpu,
            self.req_states.num_computed_tokens.gpu,
            self.req_states.prompt_len,
            self.req_states.prefill_len.np,
            self.req_states.num_computed_prefill_tokens,
        )

        # Prepare the model runner output.
        model_runner_output = ModelRunnerOutput(
            req_ids=input_batch.req_ids,
            # NOTE(woosuk): req_id_to_index is unused in this model runner.
            # Only for compatibility with the existing model runner and scheduler.
            req_id_to_index={req_id: i for i, req_id in enumerate(input_batch.req_ids)},
            sampled_token_ids=None,  # type: ignore
            prompt_logprobs_dict=prompt_logprobs_dict,  # type: ignore[arg-type]