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

Code before:
                self.prepare_dummy_attn_metadata(input_batch)
            # FIXME(woosuk): Fix warmup for LoRA.

        # Run model.
        if use_cudagraph:
            # Run CUDA graph.
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
            self.kv_connector.pre_forward(scheduler_output)
            hidden_states = self.cudagraph_manager.run(
                input_batch.num_tokens_after_padding
            )
        else:
            # Run PyTorch model in eager mode.
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
                positions = input_batch.mrope_positions

            # PP input preparation: handler centralizes all PP input logic.
            model_inputs = self.pp_handler.prepare_model_inputs(
                input_batch.input_ids,
                positions,
                input_batch.inputs_embeds,
                intermediate_tensors,
            )

            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)
                hidden_states = self.model(**model_inputs)

        # PP output handling: use handler so all PP output logic lives in one place.
        kv_connector_output = self.kv_connector.post_forward(scheduler_output)
        output = self.pp_handler.prepare_output(hidden_states, kv_connector_output)

        if isinstance(output, IntermediateTensors):
            # Non-last rank: return intermediate tensors to worker for sending
            self.execute_model_state = (None, input_batch, kv_connector_output)
            return output

        # Last rank: store hidden states for sampling
        self.execute_model_state = (output, input_batch, kv_connector_output)
        return None

Review comment (refactor):
Can we keep the model inputs explicit? I’m not a fan of encapsulating them into a separate object. For example, when supporting CUDA graphs, it’s critical to reason about all inputs and ensure they use consistent memory addresses. This abstraction makes that harder to see.

Code after:
            # NOTE(woosuk): Here, we don't need to pass the input tensors,
            # because they are already copied to the CUDA graph input buffers.
            self.kv_connector.pre_forward(scheduler_output)
            hidden_states = self.cudagraph_manager.run(
                input_batch.num_tokens_after_padding
            )
        else:
            # Run PyTorch model in eager mode.
            positions = input_batch.positions
            if self.uses_mrope:
                assert input_batch.mrope_positions is not None
                positions = input_batch.mrope_positions

            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)