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

Code before:
            self.speculator.run_model(
                self.max_num_tokens,
                attn_metadata=None,
                slot_mappings=None,
                num_tokens_across_dp=num_tokens_across_dp,
            )
        torch.cuda.synchronize()
        del hidden_states, sample_hidden_states
        gc.collect()

    def reset_mm_cache(self) -> None:
        if self.supports_mm_inputs:
            self.encoder_runner.reset_mm_cache()

    def reset_encoder_cache(self) -> None:
        if self.supports_mm_inputs:
            self.encoder_runner.reset_encoder_cache()

    def _get_num_input_tokens(self, num_scheduled_tokens: int) -> int:
        # SP is not supported yet.
        return num_scheduled_tokens

    def _get_cudagraph_runtime_mode(
        self,
        scheduler_output: SchedulerOutput,
    ) -> tuple[CUDAGraphMode, int | None]:
        # Determine if this is a uniform decode batch
        num_reqs = len(scheduler_output.num_scheduled_tokens)
        max_num_scheduled_tokens = max(scheduler_output.num_scheduled_tokens.values())
        is_uniform_decode = (
            max_num_scheduled_tokens == self.uniform_decode_query_len
        ) and (
            scheduler_output.total_num_scheduled_tokens
            == max_num_scheduled_tokens * num_reqs
        )

        # Get local_cudagraph_size using the local token count
        local_cudagraph_size = self.cudagraph_manager.get_cudagraph_size(
            scheduler_output.total_num_scheduled_tokens,
            uniform_decode=is_uniform_decode,
        )

        cudagraph_mode = self.cudagraph_manager.cudagraph_mode
        if local_cudagraph_size is None:
            local_cudagraph_runtime_mode = CUDAGraphMode.NONE
        elif is_uniform_decode:
            local_cudagraph_runtime_mode = cudagraph_mode.decode_mode()
        else:
            local_cudagraph_runtime_mode = cudagraph_mode.mixed_mode()

        return local_cudagraph_runtime_mode, local_cudagraph_size

Review comment (refactor):
Could you move this method into `cudagraph_utils.py` / `CudaGraphManager`.

Also I don't think it needs to take scheduler_output, we can pass some of these things directly, such as num_reqs. With a recent PR merged, we have `max_query_len` here too (same as max_num_scheduled_tokens)

Code after:
        self.execute_model(
            dummy_scheduler_output,
            intermediate_tensors=intermediate_tensors,
            dummy_run=True,
            skip_attn_for_dummy_run=skip_attn,
        )
        self.kv_connector.set_disabled(False)

        # Non-last PP ranks don't produce output for sampling.
        if not self.is_last_pp_rank:
            return None, None

        assert self.execute_model_state is not None
        hidden_states, input_batch, _ = self.execute_model_state
        assert hidden_states is not None  # Last PP rank always has hidden_states
        sample_hidden_states = hidden_states[input_batch.logits_indices]
        return hidden_states, sample_hidden_states

    @torch.inference_mode()
    def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
        num_reqs = hidden_states.shape[0]
        logits = self.model.compute_logits(hidden_states)
        idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=self.device)
        idx_mapping_np = np.arange(num_reqs, dtype=np.int32)
        pos = torch.zeros(num_reqs, dtype=torch.int64, device=self.device)
        dummy_input_ids = torch.zeros(num_reqs, dtype=torch.int32, device=self.device)
        expanded_local_pos = torch.zeros(
            num_reqs, dtype=torch.int32, device=self.device
        )
        # NOTE(woosuk): During the initial memory profiling, the sampler may skip
        # top_k, top_p, and logprobs, using less GPU memory than what is possible
        # during actual execution.
        self.sampler(
            logits,
            idx_mapping,
            idx_mapping_np,
            idx_mapping_np,
            pos,
            dummy_input_ids,
            expanded_local_pos,
        )

    @torch.inference_mode()
    def profile_run(self) -> None:
        hidden_states, sample_hidden_states = self._dummy_run(
            self.max_num_tokens, skip_attn=True
        )

        # Only run sampler on last PP rank (non-last ranks return None).
        if self.is_last_pp_rank:
            assert sample_hidden_states is not None