v1/worker/gpu_worker.py: the GPU Worker

            visible_device_index = (
                current_platform.logical_device_id_to_visible_device_id(self.local_rank)
            )
            self.device = torch.device(f"cuda:{visible_device_index}")
            torch.accelerator.set_device_index(self.device)

            current_platform.check_if_supports_dtype(self.model_config.dtype)

            # Initialize the distributed environment BEFORE taking
            # memory snapshot
            # This ensures NCCL buffers are allocated before we measure
            # available memory
            init_worker_distributed_environment(
                self.vllm_config,
                self.rank,
                self.distributed_init_method,
                self.local_rank,
                current_platform.dist_backend,
            )

            if self.use_v2_model_runner:
                logger.info_once("Using V2 Model Runner")
            warn_if_flashinfer_kernels_missing()

            # Set random seed.
            set_random_seed(self.model_config.seed)

            # cuBLAS(Lt) asks CUDA for its handle memory directly. When torch's own
            # cudaMalloc fails, torch frees its unused cache and retries, but a cuBLAS
            # request never triggers that because torch doesn't see it. Profiling can
            # grow torch's cache until the GPU is full, and the first cuBLAS call would
            # then fail even though torch holds idle memory it could have given back.
            # Create the handles now, before the snapshot, so they count as baseline
            # like the NCCL buffers. mm uses cuBLAS; linear with a bias uses cuBLASLt.
            x = torch.ones(16, 16, dtype=self.model_config.dtype, device=self.device)
            torch.mm(x, x)
            torch.nn.functional.linear(x, x, x[0])
            del x

            # Now take memory snapshot after NCCL is initialized
            gc.collect()
            torch.accelerator.empty_cache()

            # Weights from external model loader process
            external_weight_memory = get_model_loader(
                self.load_config
            ).get_external_weight_memory(self.vllm_config)

            # take current memory snapshot after external weights are ready
            self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
            self.requested_memory = request_memory(
                init_snapshot, self.cache_config, external_weight_memory
            )
            logger.debug("worker init memory snapshot: %r", self.init_snapshot)
            logger.debug(
                "worker requested memory: %sGiB", format_gib(self.requested_memory)
            )
        else:
            raise RuntimeError(f"Unsupported device type: {self.device_config.device}")

        self._init_workspace_and_model_runner()

    def _make_model_runner(self) -> "GPUModelRunner":
        """Construct the platform's model runner; subclasses override this."""
        if self.use_v2_model_runner:
            if self.vllm_config.is_mm_encoder_only:
                from vllm.v1.worker.mm_encoder_model_runner import (
                    MMEncoderModelRunner as GPUModelRunnerV2,
                )
            else:
                from vllm.v1.worker.gpu.model_runner import (  # type: ignore[assignment]
                    GPUModelRunner as GPUModelRunnerV2,
                )

            # HACK(woosuk): This is a temporary fix to avoid type errors.
            return GPUModelRunnerV2(self.vllm_config, self.device)  # type: ignore
        from vllm.v1.worker.gpu_model_runner import GPUModelRunner as GPUModelRunnerV1

        return GPUModelRunnerV1(self.vllm_config, self.device)

        # Execute a forward pass with dummy inputs to profile the memory usage
        # of the model.
        with (
            memory_profiling(
                self.init_snapshot,
                weights_memory=int(self.model_runner.model_memory_usage),
            ) as profile_result,
            # Workspaces (e.g. the MoE workspace) grow in steps during this pass,
            # freeing each smaller buffer. Without a split limit, a later small
            # allocation can be carved out of a freed multi-GiB block and pin the
            # whole segment past the empty_cache() in memory_profiling, so it is
            # counted as consumed and taken from the KV cache. Blocks above the
            # limit are never split, so they stay releasable. Exits before
            # memory_profiling measures, restoring the original limit.
            self._scoped_allocator_max_split(max_split_size_mb=20),
            set_current_vllm_config(self.vllm_config),
        ):
            self.model_runner.profile_run(randomize_inputs=self.randomize_dummy_inputs)

        # Profile CUDA graph memory if graphs will be captured.
        # ROCm is included: #44825 moved the profiler to
        # torch.accelerator.get_memory_info (reliable on ROCm, as used by
        # the AMD-CI mem tests), and graph_pool_handle resolves to the same
        # torch.cuda handle the live capture path already uses on ROCm.
        cudagraph_memory_estimate = 0
        if (
            current_platform.is_cuda_alike() or current_platform.is_xpu()
        ) and self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
            cudagraph_memory_estimate = self.model_runner.profile_cudagraph_memory()

        # Respect the opt-in flag as originally designed.
        cudagraph_memory_estimate_applied = (
            cudagraph_memory_estimate
            if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS
            else 0
        )

        self.total_consumed = profile_result.total_consumed
        self.peak_activation_memory = (
            profile_result.transient_peak_headroom + cudagraph_memory_estimate_applied
        )
        self.cudagraph_memory_estimate = cudagraph_memory_estimate

        self.available_kv_cache_memory_bytes = (
            self.requested_memory
            - profile_result.non_kv_cache_memory
            - cudagraph_memory_estimate_applied
        )

    # Measure total consumption via mem_get_info() instead of
    # memory_reserved(), which goes negative when pluggable allocators
    # (e.g. cumem) bypass PyTorch's tracking.
    result.total_consumed = (
        result.before_create.free_memory - result.after_profile.free_memory
    )

    # total_consumed already covers persistent torch allocations; add only the
    # transient peak headroom to avoid double-counting.
    result.transient_peak_headroom = (
        result.after_profile.torch_peak - result.after_profile.torch_allocated
    )
    result.non_kv_cache_memory = result.total_consumed + result.transient_peak_headroom

    def initialize_from_config(self, kv_cache_config: KVCacheConfig) -> None:
        """Allocate GPU KV cache with the specified kv_cache_config."""
        # Update local config with adjusted num blocks after profiling,
        # so that it's available to the warmup stage.
        self.cache_config.num_gpu_blocks = kv_cache_config.num_blocks

        # Adopt the engine core's layout and prefix-cache granularity; workers
        # spawned after resolution (e.g. elastic EP scale-up) only see them
        # through the config.
        if kv_cache_config.kv_cache_layout is not None:
            record_kv_cache_layout(self.cache_config, kv_cache_config.kv_cache_layout)
        self.cache_config.hash_block_size = kv_cache_config.hash_block_size
        self.cache_config.cache_hit_alignment_tokens = (
            kv_cache_config.cache_hit_alignment_tokens
        )

        # Init kv cache connector here, because it requires
        # `kv_cache_config`.
        # NOTE(Kuntai): This need to be done before `initialize_kv_cache`,
        # because `initialize_kv_cache` will inject kv cache groups not
        # related to kv cache connector (e.g. kv cache sharing layers).
        ensure_kv_transfer_initialized(self.vllm_config, kv_cache_config)

        # If the connector provides a custom memory pool (e.g. Mooncake
        # NVLink/BAREX), use it for KV cache allocation; otherwise fall
        # back to the standard CuMem pool.
        mem_pool_context = (
            get_kv_transfer_group().get_mem_pool_context()
            if has_kv_transfer_group()
            else None
        )
        if mem_pool_context is None:
            mem_pool_context = self._maybe_get_memory_pool_context(tag="kv_cache")

        # Offload KV-init state, except on XPU, whose outermost pool would win.
        runtime_pool = (
            nullcontext()
            if current_platform.is_xpu()
            else self._maybe_get_memory_pool_context(tag="runtime")
        )
        with runtime_pool:
            self.model_runner.initialize_kv_cache(
                kv_cache_config,
                kv_cache_allocation_context=mem_pool_context,
            )

        # We skip EPLB here since we don't want to record dummy metrics
        for size in sorted(warmup_sizes, reverse=True):
            logger.info("Compile and warming up model for size %d", size)
            self.model_runner._dummy_run(
                size,
                skip_eplb=True,
                remove_lora=False,
                randomize_inputs=self.randomize_dummy_inputs,
            )
        self.model_runner.maybe_remove_all_loras(self.model_runner.lora_config)

        # Warmup and tune the kernels used during model execution before
        # cuda graph capture.
        kernel_warmup(self)

        if self.use_v2_model_runner:  # noqa: SIM102
            if self.vllm_config.kernel_config.enable_jit_warmup:
                # A workspace resize after capture frees what the graphs point at.
                warmup_kernels(
                    self.model_runner, self.execute_model, self.sample_tokens
                )

        cuda_graph_memory_bytes = 0
        if not self.model_config.enforce_eager:
            with self._get_cudagraph_capture_context():
                cuda_graph_memory_bytes = self.model_runner.capture_model()

    @torch.inference_mode()
    @with_gpu_sync_check
    def sample_tokens(
        self, grammar_output: "GrammarOutput | None"
    ) -> ModelRunnerOutput | AsyncModelRunnerOutput:
        return self.model_runner.sample_tokens(grammar_output)

    @torch.inference_mode()
    @with_gpu_sync_check
    def execute_model(
        self, scheduler_output: "SchedulerOutput"
    ) -> ModelRunnerOutput | AsyncModelRunnerOutput | None:
        # Wait for the previous step's sends so this forward pass cannot
        # overwrite buffers they are still reading.
        if self._pp_send_work:
            for handle in self._pp_send_work:
                handle.wait()
            self._pp_send_work = []

        intermediate_tensors = None
        forward_pass = scheduler_output.total_num_scheduled_tokens > 0
        if forward_pass and not get_pp_group().is_first_rank:
            tensor_dict, comm_handles, comm_postprocess = (
                get_pp_group().irecv_tensor_dict(
                    all_gather_group=get_tp_group(),
                )
            )
            assert tensor_dict is not None
            intermediate_tensors = AsyncIntermediateTensors(
                tensor_dict,
                comm_handles=comm_handles,
                comm_postprocess=comm_postprocess,
            )

        with self.annotate_profile(scheduler_output):
            output = self.model_runner.execute_model(
                scheduler_output, intermediate_tensors
            )
            if (
                self.use_v2_model_runner
                and self.model_runner.is_pooling_model
                and output is None
            ):
                output = self.model_runner.pool()  # type: ignore
            if isinstance(
                output, ModelRunnerOutput | AsyncModelRunnerOutput | NoneType
            ):
                return output

    def sleep(self, level: int = 1) -> None:
        torch.accelerator.synchronize()
        free_bytes_before_sleep = torch.accelerator.get_memory_info()[0]

        # Save the buffers before level 2 sleep
        if level == 2:
            model = self.model_runner.model
            self._save_sleep_parameters(model)
            self._sleep_saved_buffers = {
                name: buffer.cpu().clone() for name, buffer in model.named_buffers()
            }
            draft = self.get_draft_model()
            if draft is not None:
                self._sleep_saved_draft_buffers = {
                    name: buffer.cpu().clone() for name, buffer in draft.named_buffers()
                }

        if has_kv_transfer_group():
            get_kv_transfer_group().unregister_kv_caches()
        self.sleep_mode_backend.suspend(level)
        if self.vllm_config.model_config.enable_nccl_comm_suspend:
            suspend_device_comms()