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

Code before:
        )

        # 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,
            num_sampled_tokens=num_sampled,
            copy_stream=self.output_copy_stream,
            copy_event=self.output_copy_event,
        )

        # Postprocess results and update request states.
        # NOTE: This is intentionally done after creating the AsyncOutput,
        # ensuring that `copy_event` is recorded before calling postprocess.
        # This sequencing may slightly reduce latency as async D2H copy does not
        # need to wait for the postprocess to finish.
        self.postprocess(
            input_batch, sampler_output.sampled_token_ids, num_sampled, num_rejected
        )
        if self.do_spec_decode:
            draft_tokens = self.propose_draft(
                input_batch,
                hidden_states,
                None,  # aux_hidden_states
                num_sampled,
                num_rejected,
            )
            self.req_states.draft_tokens[input_batch.idx_mapping] = draft_tokens
            self.draft_tokens_handler.set_draft_tokens(input_batch, draft_tokens)

        if self.use_async_scheduling:
            return async_output
        return async_output.get_output()

    def take_draft_token_ids(self) -> DraftTokenIds | None:
        return self.draft_tokens_handler.get_draft_tokens()


Review comment (performance):
suggested:
```suggestion
                if get_pp_group().is_first_rank:
                    input_ids = input_batch.input_ids
                    inputs_embeds = input_batch.inputs_embeds
                else:
                    # Non-first PP rank: forward with intermediate tensors.
                    input_ids, inputs_embeds = None, None
                    assert intermediate_tensors is not None

                hidden_states = self.model(
                    input_ids=input_ids,
                    positions=positions,
                    inputs_embeds=inputs_embeds,
                    intermediate_tensors=intermediate_tensors,
                )
```

Code after:

            with set_forward_context(
                input_batch.attn_metadata,
                self.vllm_config,
                num_tokens=input_batch.num_tokens_after_padding,
                # TODO(woosuk): Support piecewise CUDA graph.
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
                num_tokens_across_dp=num_tokens_across_dp,
                slot_mapping=input_batch.slot_mappings,
            ):
                self.kv_connector.pre_forward(scheduler_output)
                if self.use_pp and not get_pp_group().is_first_rank:
                    # Non-first PP rank: forward with intermediate tensors.
                    assert intermediate_tensors is not None
                    hidden_states = self.model(
                        input_ids=None,
                        positions=positions,
                        inputs_embeds=None,
                        intermediate_tensors=intermediate_tensors,
                    )
                else:
                    hidden_states = self.model(
                        input_ids=input_batch.input_ids,
                        positions=positions,
                        inputs_embeds=input_batch.inputs_embeds,
                    )

        kv_connector_output = self.kv_connector.post_forward(scheduler_output)

        if self.use_pp and not get_pp_group().is_last_rank:
            # Non-last PP rank: return IntermediateTensors for sending.
            assert isinstance(hidden_states, IntermediateTensors)
            hidden_states.kv_connector_output = kv_connector_output
            self.execute_model_state = (None, input_batch, kv_connector_output)
            return hidden_states

        assert isinstance(hidden_states, torch.Tensor)
        # Last rank (or no PP): hidden_states is a tensor for sampling.
        self.execute_model_state = (hidden_states, input_batch, kv_connector_output)
        return None

    @torch.inference_mode()
    def sample_tokens(
        self, grammar_output: GrammarOutput | None
    ) -> 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