vLLM/vllm/v1/worker/gpu_worker.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""A GPU worker class."""

import gc
import os
import time
from collections.abc import Callable
from contextlib import AbstractContextManager, contextmanager, nullcontext
from datetime import timedelta
from fnmatch import filter as fnmatch_filter
from types import NoneType
from typing import TYPE_CHECKING, Any, cast

import torch
import torch.nn as nn

import vllm.envs as envs
from vllm.config import CUDAGraphMode, VllmConfig, set_current_vllm_config
from vllm.config.compilation import CompilationMode
from vllm.config.profiler import validate_profile_prefix
from vllm.device_allocator import get_mem_allocator_instance
from vllm.distributed import (
    ensure_model_parallel_initialized,
    init_distributed_environment,
    set_custom_all_reduce,
)
from vllm.distributed.ec_transfer import (
    ensure_ec_transfer_initialized,
    ensure_ec_transfer_shutdown,
    get_ec_transfer,
    has_ec_transfer,
)
from vllm.distributed.eplb.eplb_utils import override_envs_for_eplb
from vllm.distributed.kv_transfer import (
    ensure_kv_transfer_initialized,
    ensure_kv_transfer_shutdown,
    get_kv_transfer_group,
    has_kv_transfer_group,
)
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
    KVConnectorHandshakeMetadata,
)
from vllm.distributed.parallel_state import (
    Handle,
    checkpoint_prepare_distributed_state,
    checkpoint_restore_distributed_state,
    get_pcp_group,
    get_pp_group,
    get_tp_group,
    resume_device_comms,
    suspend_device_comms,
)
from vllm.distributed.weight_transfer import (
    WeightTransferEngine,
    WeightTransferEngineFactory,
)
from vllm.logger import init_logger
from vllm.lora.request import LoRARequest
from vllm.model_executor.model_loader import get_model_loader
from vllm.model_executor.model_loader.weight_checksum import (
    compute_tensor_digests,
    zero_weights,
)
from vllm.model_executor.warmup.kernel_warmup import kernel_warmup
from vllm.multimodal.gpu_ipc_memory import reserve_mm_ipc_gpu_memory
from vllm.platforms import current_platform
from vllm.profiler.wrapper import (
    create_graph_capture_profiler,
    create_worker_profiler,
    validate_worker_profiler_config,
)
from vllm.sequence import IntermediateTensors
from vllm.tasks import SupportedTask
from vllm.tracing import instrument
from vllm.utils.flashinfer import warn_if_flashinfer_kernels_missing
from vllm.utils.gc_utils import freeze_gc_heap, maybe_attach_gc_debug_callback
from vllm.utils.gpu_sync_debug import enable_gpu_sync_check, with_gpu_sync_check
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.mem_utils import (
    MemoryProfilingResult,
    MemorySnapshot,
    format_gib,
    memory_profiling,
)
from vllm.utils.torch_utils import set_random_seed, set_torch_threads_for_runtime
from vllm.v1.attention.backends.utils import record_kv_cache_layout
from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput
from vllm.v1.kv_cache_interface import KVCacheConfig, KVCacheSpec
from vllm.v1.outputs import (
    AsyncModelRunnerOutput,
    DraftTokenIds,
    ModelRunnerOutput,
)
from vllm.v1.utils import compute_iteration_details, report_usage_stats
from vllm.v1.worker.sentinel.gpu_worker_sentinel import WorkerSentinel
from vllm.v1.worker.startup_plan import (
    maybe_apply_startup_plan,
    maybe_save_startup_plan,
)
from vllm.v1.worker.worker_base import CompilationTimes, WorkerBase
from vllm.v1.worker.workspace import init_workspace_manager

from ...model_executor.model_loader import TensorizerLoader
from .gpu.cudagraph_utils import has_compiled_submodule
from .gpu.warmup import warmup_kernels
from .utils import request_memory

logger = init_logger(__name__)


def _num_workspace_lanes(vllm_config: VllmConfig, use_v2_model_runner: bool) -> int:
    spec_config = vllm_config.speculative_config
    return (
        2
        if use_v2_model_runner and spec_config is not None and spec_config.use_dspark()
        else 1
    )


def maybe_rocm_profiling_fallback(profile_result: MemoryProfilingResult) -> int | None:
    """Memory in bytes to size the KV cache from when profiling measured a
    release on ROCm, or None to keep what profiling measured.

    The measurement is the drop in free device memory across startup, so it comes
    out negative when anything else on the device released memory meanwhile. Torch
    reserved is process local, so it stands in as a lower bound instead of refusing
    to start. Kept to ROCm, where the AMD CI groups hit this, following #45490.
    """
    if profile_result.total_consumed >= 0 or not current_platform.is_rocm():
        return None

    baseline = profile_result.before_create
    after_profile = profile_result.after_profile
    torch_reserved = max(after_profile.torch_memory - baseline.torch_memory, 0)
    logger.warning(
        "Free memory grew by %s GiB during profiling (initial %s GiB, current "
        "%s GiB). Sizing the KV cache from the %s GiB torch reserved instead of "
        "the profiling difference.",
        format_gib(-profile_result.total_consumed),
        format_gib(baseline.free_memory),
        format_gib(after_profile.free_memory),
        format_gib(torch_reserved),
    )
    return torch_reserved


if TYPE_CHECKING:
    from vllm.device_allocator.sleep_mode_backend import SleepModeBackend
    from vllm.model_executor.model_loader.tensorizer import TensorizerConfig
    from vllm.v1.worker.gpu.model_runner import GPUModelRunner as GPUModelRunnerV2
    from vllm.v1.worker.gpu_model_runner import GPUModelRunner


class AsyncIntermediateTensors(IntermediateTensors):
    """IntermediateTensors with lazy comm synchronization"""

    def __init__(
        self,
        tensors: dict[str, torch.Tensor],
        comm_handles: list[Handle] | None = None,
        comm_postprocess: list[Callable[[], None]] | None = None,
    ) -> None:
        super().__init__(tensors)
        self._comm_handles = comm_handles
        self._comm_postprocess = comm_postprocess
        self._comm_waited = False

    def wait_for_comm(self) -> None:
        if self._comm_waited:
            return
        if self._comm_handles:
            for handle in self._comm_handles:
                handle.wait()
        if self._comm_postprocess:
            for fn in self._comm_postprocess:
                fn()
        self._comm_waited = True

    def __getattribute__(self, name: str):
        # ensure `.tensors` is ready before use
        if name == "tensors" and not object.__getattribute__(self, "_comm_waited"):
            object.__getattribute__(self, "wait_for_comm")()
        return object.__getattribute__(self, name)


class Worker(WorkerBase):
    def __init__(
        self,
        vllm_config: VllmConfig,
        local_rank: int,
        rank: int,
        distributed_init_method: str,
        is_driver_worker: bool = False,
    ):
        super().__init__(
            vllm_config=vllm_config,
            local_rank=local_rank,
            rank=rank,
            distributed_init_method=distributed_init_method,
            is_driver_worker=is_driver_worker,
        )

        # configure float32 matmul precision according to vLLM env.
        precision = envs.VLLM_FLOAT32_MATMUL_PRECISION
        torch.set_float32_matmul_precision(precision)

        from vllm.distributed.elastic_ep.elastic_execute import ElasticEPScalingExecutor

        self.elastic_ep_executor = ElasticEPScalingExecutor(self)
        self.worker_sentinel: WorkerSentinel | None = None
        if self.parallel_config.enable_fault_tolerance:
            self.worker_sentinel = WorkerSentinel(worker=self)
        # Buffers saved before sleep
        self._sleep_saved_buffers: dict[str, torch.Tensor] = {}
        self._sleep_saved_parameters: dict[str, torch.Tensor] = {}
        self._sleep_saved_draft_buffers: dict[str, torch.Tensor] = {}

        # Weight transfer engine is created in `load_model` once the model
        # is available, since the engine needs a reference to the model.
        self.weight_transfer_engine: WeightTransferEngine | None = None
        self._weight_update_active = False
        self._weight_update_is_draft = False

        # Worker profiler. Enabled and configured through profiler_config.
        # Profiler wrapper is created lazily in profile() when start is called,
        # so we have all the information needed for proper trace naming.
        self.profiler: Any | None = None
        self.profiler_config = vllm_config.profiler_config
        validate_worker_profiler_config(self.profiler_config)

        self.use_v2_model_runner = vllm_config.use_v2_model_runner

        # Device handles of the previous step's PP intermediate-tensor send.
        self._pp_send_work: list[Handle] = []

        # Resolved lazily on first sleep/wake; persists worker-process state.
        self._sleep_mode_backend: SleepModeBackend | None = None

    @property
    def sleep_mode_backend(self) -> "SleepModeBackend":
        if self._sleep_mode_backend is None:
            from vllm.device_allocator.sleep_mode_backend import (
                SleepModeBackendFactory,
            )

            self._sleep_mode_backend = SleepModeBackendFactory.create_backend(
                self.vllm_config.model_config
            )
        return self._sleep_mode_backend

    def _save_sleep_parameters(self, model: nn.Module) -> None:
        patterns = self.model_config.sleep_preserve_parameter_names
        if not patterns:
            self._sleep_saved_parameters = {}
            return
        parameters = dict(model.named_parameters())
        names: set[str] = set()
        for pattern in patterns:
            matches = fnmatch_filter(parameters, pattern)
            if not matches:
                raise ValueError(f"No parameter matches sleep retention: {pattern}")
            names.update(matches)
        self._sleep_saved_parameters = {
            name: param.detach().to("cpu")
            for name, param in parameters.items()
            if name in names and not param.is_cpu
        }

    @torch.no_grad()
    def _restore_sleep_parameters(self, model: nn.Module) -> None:
        for name, value in self._sleep_saved_parameters.items():
            model.get_parameter(name).copy_(value)
        self._sleep_saved_parameters.clear()

    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()
                }

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

        torch.accelerator.synchronize()
        deadline = time.monotonic() + (5.0 if current_platform.is_rocm() else 0)
        while True:
            free_bytes_after_sleep, total = torch.accelerator.get_memory_info()
            freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
            if freed_bytes >= 0 or time.monotonic() >= deadline:
                break
            time.sleep(0.1)

        used_bytes = total - free_bytes_after_sleep
        assert freed_bytes >= 0, "Memory usage increased after sleeping."
        logger.info(
            "Sleep mode freed %s GiB memory, %s GiB memory is still in use.",
            format_gib(freed_bytes),
            format_gib(used_bytes),
        )

    def wake_up(self, tags: list[str] | None = None) -> None:
        self.sleep_mode_backend.resume(tags)
        if self.vllm_config.model_config.enable_nccl_comm_suspend:
            resume_device_comms()

        # Restore the buffers after level 2 sleep
        wake_weights = tags is None or "weights" in tags
        if wake_weights and self._sleep_saved_parameters:
            self._restore_sleep_parameters(self.model_runner.model)
        if wake_weights and len(self._sleep_saved_buffers):
            model = self.model_runner.model
            for name, buffer in model.named_buffers():
                if name in self._sleep_saved_buffers:
                    buffer.data.copy_(self._sleep_saved_buffers[name].data)
            self._sleep_saved_buffers = {}

        if wake_weights and len(self._sleep_saved_draft_buffers):
            draft = self.get_draft_model()
            if draft is not None:
                for name, buffer in draft.named_buffers():
                    if name in self._sleep_saved_draft_buffers:
                        buffer.data.copy_(self._sleep_saved_draft_buffers[name].data)
            self._sleep_saved_draft_buffers = {}

        self.synchronize_device()

    def discard(self, tags: tuple[str, ...]) -> None:
        self.sleep_mode_backend.discard(tags)

    def checkpoint_prepare(self) -> None:
        checkpoint_prepare_distributed_state()

    def checkpoint_restore(self) -> None:
        checkpoint_restore_distributed_state()

    def _weight_checksum_prefix(self) -> str:
        # data_parallel_rank is 0 in dense workers; only the index is unique.
        dp = self.parallel_config.data_parallel_index
        pp = get_pp_group().rank_in_group
        pcp = get_pcp_group().rank_in_group
        tp = get_tp_group().rank_in_group
        return f"dp{dp}:pp{pp}:pcp{pcp}:tp{tp}:"

    def compute_weight_checksums(self) -> dict[str, str]:
        prefix = self._weight_checksum_prefix()
        digests = compute_tensor_digests(self.model_runner.get_model())
        return {prefix + name: digest for name, digest in digests.items()}

    def reset_weights(self) -> None:
        zero_weights(self.model_runner.get_model())

    def _maybe_get_memory_pool_context(self, tag: str) -> AbstractContextManager:
        if (
            current_platform.is_cuda_alike()
            and not self.vllm_config.model_config.enable_cumem_allocator
        ):
            return nullcontext()

        if (
            current_platform.is_xpu()
            and not self.vllm_config.model_config.enable_sleep_mode
        ):
            return nullcontext()

        if current_platform.is_cpu():
            return nullcontext()

        allocator = get_mem_allocator_instance()
        if tag == "weights":
            assert allocator.get_current_usage() == 0, (
                "CuMem allocator can only be used for one instance per process."
            )
        return allocator.use_memory_pool(tag=tag)

    @contextmanager
    def _scoped_allocator_max_split(self, max_split_size_mb: int):
        """Limit CUDA/ROCm allocator splitting, restoring settings on exit."""
        if (
            not current_platform.is_cuda_alike()
            or torch.cuda.memory.get_allocator_backend() != "native"
        ):
            yield
            return

        get_settings = getattr(torch._C, "_accelerator_getAllocatorSettings", None)
        if get_settings is not None:
            original_settings = get_settings()
        else:
            original_settings = torch.cuda.memory._snapshot()["allocator_settings"][
                "PYTORCH_CUDA_ALLOC_CONF"
            ]

        # The setter resets other allocator options unless they are included.
        prefix = original_settings.strip().rstrip(",")
        settings = f"max_split_size_mb:{max_split_size_mb}"
        if prefix:
            settings = f"{prefix},{settings}"
        try:
            torch._C._accelerator_setAllocatorSettings(settings)
            yield
        finally:
            torch._C._accelerator_setAllocatorSettings(original_settings)

    @instrument(span_name="Init device")
    def init_device(self):
        if self.device_config.device_type == "cuda":
            # This env var set by Ray causes exceptions with graph building.
            os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
            parallel_config = self.parallel_config
            if (
                parallel_config.distributed_executor_backend
                not in ("ray", "external_launcher")
                and parallel_config.data_parallel_backend != "ray"
                and parallel_config.nnodes_within_dp == 1
            ):
                # Use local DP rank if available, otherwise use global DP rank.
                dp_local_rank = self.parallel_config.data_parallel_rank_local
                if dp_local_rank is None:
                    dp_local_rank = self.parallel_config.data_parallel_index

                tp_pcp_pp_world_size = (
                    self.parallel_config.pipeline_parallel_size
                    * self.parallel_config.prefill_context_parallel_size
                    * self.parallel_config.tensor_parallel_size
                )

                # DP_LOCAL_RANK * TP_PCP_PP_WORLD_SIZE + TP_LOCAL_RANK
                self.local_rank += dp_local_rank * tp_pcp_pp_world_size

            # Publish the logical-to-physical mapping for topology queries
            # such as NIC affinity and P2P checks.
            assigned_physical_gpu_ids = parallel_config.assigned_physical_gpu_ids
            if assigned_physical_gpu_ids is not None:
                from vllm.platforms.interface import set_assigned_physical_gpu_ids

                set_assigned_physical_gpu_ids(assigned_physical_gpu_ids)
                assert self.local_rank < len(assigned_physical_gpu_ids), (
                    f"local_rank {self.local_rank} is out of bounds for "
                    f"assigned_physical_gpu_ids {assigned_physical_gpu_ids}"
                )
                # NOTE(patch pr45026): local_world_size is derived from
                # parallel_config.nnodes, which is only set for the "mp"
                # multi-node backend. With the "ray"/"external_launcher"
                # backends nnodes stays 1, so local_world_size collapses to
                # the full world_size and this check wrongly fires on
                # cross-node deployments. assigned_physical_gpu_ids is already
                # per-node and the local_rank bound above fully validates the
                # mapping for these backends, so skip the check for them.
                if parallel_config.distributed_executor_backend not in (
                    "ray",
                    "external_launcher",
                ):
                    assert self.parallel_config.local_world_size <= len(
                        assigned_physical_gpu_ids
                    ), (
                        f"local_world_size ({self.parallel_config.local_world_size})"
                        " exceeds assigned_physical_gpu_ids count "
                        f"({len(assigned_physical_gpu_ids)})"
                    )
            else:
                assert self.local_rank < torch.accelerator.device_count(), (
                    f"DP adjusted local rank {self.local_rank} is out of "
                    f"bounds for {torch.accelerator.device_count()} devices."
                )

            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)

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

            # take current memory snapshot
            self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
            # Weights from external model loader process
            external_weight_memory = get_model_loader(
                self.load_config
            ).get_external_weight_memory(self.vllm_config)
            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 _init_workspace_and_model_runner(self) -> None:
        """Set up the workspace manager, build the model runner, report usage."""
        # DSpark target and draft CUDA graphs retain workspace views concurrently.
        num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1
        init_workspace_manager(
            self.device,
            num_ubatches,
            _num_workspace_lanes(self.vllm_config, self.use_v2_model_runner),
            # Scratch holds no state across steps: discard it on sleep.
            alloc_context=lambda: self._maybe_get_memory_pool_context("workspace"),
        )
        self.model_runner: GPUModelRunner = self._make_model_runner()
        if self.rank == 0:
            # If usage stat is enabled, collect relevant info.
            report_usage_stats(self.vllm_config)

    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)

    def handle_ft_command(self, ft_request):
        assert self.worker_sentinel is not None
        return self.worker_sentinel.handle_command(ft_request)

    # FIXME(youkaichao & ywang96): Use TorchDispatchMode instead of memory pool
    # to hijack tensor allocation.
    def load_model(self, *, load_dummy_weights: bool = False) -> None:
        with (
            self._maybe_get_memory_pool_context(tag="weights"),
            set_current_vllm_config(self.vllm_config),
            # 20 MiB is the minimum PyTorch allows for max_split_size_mb.
            self._scoped_allocator_max_split(max_split_size_mb=20),
        ):
            self.model_runner.load_model(load_dummy_weights=load_dummy_weights)

        if has_ec_transfer():
            get_ec_transfer().start_worker_services()

        if self.vllm_config.weight_transfer_config is not None:
            self.weight_transfer_engine = WeightTransferEngineFactory.create_engine(
                self.vllm_config.weight_transfer_config,
                self.vllm_config,
                self.device,
                self.model_runner.get_model(),
            )

        # Preserve parallel weight loading, then use serving's thread count
        # for profiling and compilation so Dynamo's global-state guards remain
        # valid when requests arrive.
        set_torch_threads_for_runtime()

    def update_config(self, overrides: dict[str, Any]) -> None:
        self.model_runner.update_config(overrides)

    def reload_weights(self, *args, **kwargs) -> None:
        with set_current_vllm_config(self.vllm_config):
            self.model_runner.reload_weights(*args, **kwargs)

    @torch.inference_mode()
    def determine_available_memory(self) -> int:
        """Profiles the peak memory usage of the model to determine how much
        memory can be used for KV cache without OOMs.

        The engine will first conduct a profiling of the existing memory usage.
        Then, it calculates the free memory that can be used for KV cache in
        bytes.

        Tip:
            You may limit the usage of GPU memory
            by adjusting the `gpu_memory_utilization` parameter.

        """
        maybe_apply_startup_plan(self)

        if kv_cache_memory_bytes := self.cache_config.kv_cache_memory_bytes:
            # still need a profile run which compiles the model for
            # max_num_batched_tokens
            with set_current_vllm_config(self.vllm_config):
                self.model_runner.profile_run(
                    randomize_inputs=self.randomize_dummy_inputs
                )

            msg = (
                f"Initial free memory {format_gib(self.init_snapshot.free_memory)} "
                f"GiB, reserved {format_gib(kv_cache_memory_bytes)} GiB memory for "
                "KV Cache as specified by kv_cache_memory_bytes config and "
                "skipped memory profiling. This does not respect the "
                "gpu_memory_utilization config. Only use kv_cache_memory_bytes "
                "config when you want manual control of KV cache memory "
                "size. If OOM'ed, check the difference of initial free "
                "memory between the current run and the previous run "
                "where kv_cache_memory_bytes is suggested and update it "
                "correspondingly."
            )
            logger.info(msg)
            return reserve_mm_ipc_gpu_memory(
                kv_cache_memory_bytes,
                self.model_config.multimodal_config,
                getattr(self.parallel_config, "_api_process_count", 1),
            )

        # 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
        )

        init_free_memory = self.init_snapshot.free_memory
        free_gpu_memory = profile_result.after_profile.free_memory
        rocm_fallback = maybe_rocm_profiling_fallback(profile_result)
        if rocm_fallback is None:
            # NOTE(woosuk): Here we assume that the other processes using the same
            # GPU did not change their memory usage during the profiling.
            assert init_free_memory >= free_gpu_memory, (
                "Error in memory profiling. "
                f"Initial free memory {format_gib(init_free_memory)} GiB, "
                f"current free memory {format_gib(free_gpu_memory)} GiB. "
                "This happens when other processes sharing the same container "
                "release GPU memory while vLLM is profiling during initialization. "
                "To fix this, ensure consistent GPU memory allocation or "
                "isolate vLLM in its own container."
            )
        else:
            profile_result.total_consumed = rocm_fallback
            profile_result.non_kv_cache_memory = (
                profile_result.total_consumed + profile_result.transient_peak_headroom
            )

        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
        )

        unrequested_memory = self.init_snapshot.free_memory - self.requested_memory
        logger.debug(
            "Initial free memory: %s GiB; Requested memory: %f (util), %s GiB",
            format_gib(self.init_snapshot.free_memory),
            self.cache_config.gpu_memory_utilization,
            format_gib(self.requested_memory),
        )
        logger.debug(
            "Free memory after profiling: %s GiB (total), %s GiB (within requested)",
            format_gib(free_gpu_memory),
            format_gib(free_gpu_memory - unrequested_memory),
        )
        logger.debug(profile_result)
        logger.info_once(
            "Available KV cache memory: %s GiB",
            format_gib(self.available_kv_cache_memory_bytes),
        )

        if cudagraph_memory_estimate > 0:
            total_mem = self.init_snapshot.total_memory
            current_util = self.cache_config.gpu_memory_utilization
            cg_util_delta = cudagraph_memory_estimate / total_mem
            if envs.VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS:
                equiv_util = round(current_util - cg_util_delta, 4)
                suggested_util = min(
                    round(current_util + cg_util_delta, 4),
                    1.0,
                )
                logger.info(
                    "CUDA graph memory profiling is enabled (default since "
                    "v0.21.0). The current --gpu-memory-utilization=%.4f is "
                    "equivalent to --gpu-memory-utilization=%.4f without "
                    "CUDA graph memory profiling. To maintain the same "
                    "effective KV cache size as before, increase "
                    "--gpu-memory-utilization to %.4f. To disable, set "
                    "VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0.",
                    current_util,
                    equiv_util,
                    suggested_util,
                )
            else:
                suggested_util = min(
                    round(current_util + cg_util_delta, 4),
                    1.0,
                )
                logger.warning(
                    "CUDA graph memory profiling is disabled "
                    "(VLLM_MEMORY_PROFILER_ESTIMATE_CUDAGRAPHS=0). "
                    "Without it, CUDA graph memory is not accounted for "
                    "during KV cache allocation, which may require lowering "
                    "--gpu-memory-utilization to avoid OOM. Consider "
                    "re-enabling it (the default as of v0.21.0) and increasing "
                    "--gpu-memory-utilization from %.4f to %.4f.",
                    current_util,
                    suggested_util,
                )

        return reserve_mm_ipc_gpu_memory(
            int(self.available_kv_cache_memory_bytes),
            self.model_config.multimodal_config,
            getattr(self.parallel_config, "_api_process_count", 1),
        )

    def get_kv_connector_handshake_metadata(
        self,
    ) -> dict[tuple[int, int], KVConnectorHandshakeMetadata] | None:
        """Get KV connector metadata from this worker if available.

        Returned dict is keyed by `(pp_rank, tp_rank)`.
        """
        if not has_kv_transfer_group():
            return None

        connector = get_kv_transfer_group()
        # Return None for connectors that don't need to exchange handshake
        # metadata across workers.
        if (metadata := connector.get_handshake_metadata()) is None:
            return None

        pp_rank = get_pp_group().rank_in_group
        tp_rank = get_tp_group().rank_in_group
        parallel_config = self.vllm_config.parallel_config
        if (
            parallel_config.prefill_context_parallel_size > 1
            and parallel_config.decode_context_parallel_size > 1
        ):
            tp_rank += (
                get_pcp_group().rank_in_group * parallel_config.tensor_parallel_size
            )
        return {(pp_rank, tp_rank): metadata}

    def get_kv_cache_spec(self) -> dict[str, KVCacheSpec]:
        return self.model_runner.get_kv_cache_spec()

    def update_max_model_len(self, max_model_len: int) -> None:
        """Update max_model_len after auto-fit to GPU memory.
        This is called when max_model_len=-1 is used and the engine
        automatically determines the maximum context length that fits
        in GPU memory. Workers need to update their cached max_model_len
        to match the engine's decision.
        """
        self.model_config.max_model_len = max_model_len
        if self.model_runner is not None:
            self.model_runner.update_max_model_len(max_model_len)
        logger.debug("Updated max_model_len to %d", max_model_len)

    @instrument(span_name="Allocate KV cache")
    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,
            )

        # Build KV-zero metadata outside the CuMem pool so the bookkeeping
        # GPU tensors (seg_addrs, block-id buffers) use the standard PyTorch
        # allocator and are not discarded during sleep/wake cycles.
        if kv_cache_config.needs_kv_cache_zeroing and hasattr(
            self.model_runner, "_init_kv_zero_meta"
        ):
            self.model_runner._init_kv_zero_meta()

    @instrument(span_name="Warmup (GPU)")
    def compile_or_warm_up_model(self) -> CompilationTimes:
        warmup_sizes: list[int] = []

        if (
            self.vllm_config.compilation_config.mode == CompilationMode.VLLM_COMPILE
            and has_compiled_submodule(self.model_runner.get_model())
        ):
            # warm up sizes that are not in cudagraph capture sizes,
            # but users still want to compile for better performance,
            # e.g. for the max-num-batched token size in chunked prefill.
            compile_sizes = self.vllm_config.compilation_config.compile_sizes
            warmup_sizes = compile_sizes.copy() if compile_sizes is not None else []  # type: ignore[assignment]
            cg_capture_sizes: list[int] = []

            if self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE:
                cg_sizes = self.vllm_config.compilation_config.cudagraph_capture_sizes
                cg_capture_sizes = [] if cg_sizes is None else cg_sizes
                warmup_sizes = [x for x in warmup_sizes if x not in cg_capture_sizes]

            compile_ranges = self.vllm_config.compilation_config.get_compile_ranges()
            # For each compile_range, if none of the batch sizes
            # in warmup_sizes or cudagraph_capture_sizes are in the range,
            # add the end of the range to ensure compilation/warmup.
            all_sizes = set(cg_capture_sizes)
            all_sizes.update([x for x in warmup_sizes if isinstance(x, int)])
            for compile_range in compile_ranges:
                if not any(x in compile_range for x in all_sizes):
                    warmup_sizes.append(compile_range.end)

        # 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()

        # Compare actual vs estimated CUDA graph memory (if we did profiling)
        if (
            hasattr(self, "cudagraph_memory_estimate")
            and self.cudagraph_memory_estimate > 0
        ):
            GiB = lambda b: round(b / GiB_bytes, 2)
            diff = abs(cuda_graph_memory_bytes - self.cudagraph_memory_estimate)
            logger.info(
                "CUDA graph pool memory: %s GiB (actual), %s GiB (estimated), "
                "difference: %s GiB (%.1f%%).",
                GiB(cuda_graph_memory_bytes),
                GiB(self.cudagraph_memory_estimate),
                GiB(diff),
                100 * diff / max(cuda_graph_memory_bytes, 1),
            )

        if self.cache_config.kv_cache_memory_bytes is None and hasattr(
            self, "peak_activation_memory"
        ):
            # Suggests optimal kv cache memory size if we rely on
            # memory_profiling to guess the kv cache memory size which
            # provides peak_activation_memory and a few other memory
            # consumption. `memory_profiling` does not consider
            # CUDAGraph memory size and may not utilize all gpu memory.
            # Users may want fine-grained control to specify kv cache
            # memory size.

            # empirically observed that the memory profiling may
            # slightly underestimate the memory consumption.
            # So leave a small buffer (=150MiB) to avoid OOM.
            redundancy_buffer_memory = 150 * (1 << 20)

            non_kv_cache_memory = (
                self.total_consumed
                + self.peak_activation_memory
                + cuda_graph_memory_bytes
            )
            kv_cache_memory_bytes_to_gpu_limit = (
                self.init_snapshot.free_memory
                - non_kv_cache_memory
                - redundancy_buffer_memory
            )
            kv_cache_memory_bytes_to_requested_limit = (
                int(self.requested_memory)
                - non_kv_cache_memory
                - redundancy_buffer_memory
            )

            msg = (
                f"Free memory on device "
                f"({format_gib(self.init_snapshot.free_memory)}/"
                f"{format_gib(self.init_snapshot.total_memory)} GiB) on startup. "
                f"Desired GPU memory utilization is "
                f"({self.cache_config.gpu_memory_utilization}, "
                f"{format_gib(self.requested_memory)} GiB). "
                f"Actual usage is {format_gib(self.total_consumed)} "
                f"GiB for consumed memory (weights + non-torch), "
                f"{format_gib(self.peak_activation_memory)} GiB "
                f"for peak activation, and {format_gib(cuda_graph_memory_bytes)} "
                f"GiB for CUDAGraph memory. Replace gpu_memory_utilization "
                f"config with `--kv-cache-memory="
                f"{kv_cache_memory_bytes_to_requested_limit}` "
                f"({format_gib(kv_cache_memory_bytes_to_requested_limit)} GiB) to fit "
                f"into requested memory, or `--kv-cache-memory="
                f"{kv_cache_memory_bytes_to_gpu_limit}` "
                f"({format_gib(kv_cache_memory_bytes_to_gpu_limit)} GiB) to fully "
                f"utilize gpu memory. Current kv cache memory in use is "
                f"{format_gib(self.available_kv_cache_memory_bytes)} GiB."
            )

            logger.info(msg)

            maybe_save_startup_plan(self, kv_cache_memory_bytes_to_requested_limit)

        if not self.use_v2_model_runner and get_pp_group().is_last_rank:
            # V1: Warm up sampler and preallocate memory buffer for logits and other
            # sampling related tensors of max possible shape to avoid memory
            # fragmentation issue.
            # NOTE: This is called after `capture_model` on purpose to prevent
            # memory buffers from being cleared by `torch.accelerator.empty_cache`.
            max_num_reqs = min(
                self.scheduler_config.max_num_seqs,
                self.scheduler_config.max_num_batched_tokens,
            )

            # We skip EPLB here since we don't want to record dummy metrics
            hidden_states, last_hidden_states = self.model_runner._dummy_run(
                num_tokens=max_num_reqs,
                skip_eplb=True,
                cudagraph_runtime_mode=CUDAGraphMode.NONE,
            )
            if self.model_runner.is_pooling_model:
                self.model_runner._dummy_pooler_run(hidden_states)
            else:
                self.model_runner._dummy_sampler_run(hidden_states=last_hidden_states)

        # Reset the seed to ensure that the random state is not affected by
        # the model initialization and profiling.
        set_random_seed(
            self.model_config.seed, self.parallel_config.data_parallel_index
        )

        # Eagerly trigger inductor's once-per-process lazy inits during
        # warmup (rather than on a later compile cache-miss at runtime).
        c_config = self.compilation_config
        if c_config.mode != CompilationMode.NONE and c_config.backend == "inductor":
            from vllm.compilation.compiler_interface import (
                trigger_inductor_lazy_init,
            )

            trigger_inductor_lazy_init(self.device)

        # All warmup is done — start monitoring for unexpected JIT
        # compilations that would cause latency spikes during inference.
        self._maybe_activate_jit_monitor()

        # Freeze the worker heap so the GC won't scan static objects
        # (model weights, KV caches, CUDA graphs) during inference.
        freeze_gc_heap()
        maybe_attach_gc_debug_callback()

        # Warmup / first-compile is done — activate the `VLLM_GPU_SYNC_CHECK`
        # gate so subsequent `execute_model` / `sample_tokens` calls enforce it.
        enable_gpu_sync_check()

        return CompilationTimes(
            language_model=self.compilation_config.compilation_time,
            encoder=self.compilation_config.encoder_compilation_time,
        )

    def _maybe_activate_jit_monitor(self) -> None:
        # When JIT warmup is disabled (e.g. enforce_eager), runtime JIT
        # compilation is expected, so monitoring would only produce noise
        # (or spurious errors in "error" mode).
        if not self.vllm_config.kernel_config.enable_jit_warmup:
            return

        from vllm.utils.jit_monitor import activate as activate_jit_monitor

        activate_jit_monitor(
            mode=self.observability_config.jit_monitor_mode,
            verbose=self.observability_config.jit_monitor_verbose,
        )

    def _get_cudagraph_capture_context(self) -> AbstractContextManager[None]:
        """Let the configured profiler observe CUDA graph capture."""
        if not self.use_v2_model_runner:
            return nullcontext()
        if self.profiler is None:
            model_runner = cast("GPUModelRunnerV2", self.model_runner)
            if not model_runner.needs_cudagraph_capture():
                return nullcontext()
            self.profiler = create_graph_capture_profiler(
                self.profiler_config, global_rank=self.rank
            )
            if self.profiler is None:
                return nullcontext()
        return self.profiler.capture_cuda_graphs()

    def reset_mm_cache(self) -> None:
        self.model_runner.reset_mm_cache()

    def reset_encoder_cache(self) -> None:
        self.model_runner.reset_encoder_cache()

    def get_model(self) -> nn.Module:
        return self.model_runner.get_model()

    def get_draft_model(self) -> nn.Module | None:
        return self.model_runner.get_draft_model()

    def supports_draft_weight_updates(self) -> bool:
        engine = self.weight_transfer_engine
        speculative_config = self.speculative_config
        get_draft_model = getattr(self.model_runner, "get_draft_model", None)
        return (
            engine is not None
            and engine.supports_draft_weight_update
            and callable(get_draft_model)
            and get_draft_model() is not None
            and speculative_config is not None
            and speculative_config.draft_model_config is not None
        )

    def _set_draft_weight_update_target(self) -> None:
        assert self.weight_transfer_engine is not None

        draft_model = self.get_draft_model()
        if draft_model is None:
            raise RuntimeError(
                "Draft model weight update requested, but no draft model is configured."
            )

        speculative_config = self.speculative_config
        if speculative_config is None or speculative_config.draft_model_config is None:
            raise RuntimeError(
                "Draft model weight update requested, but no draft model "
                "config is configured."
            )

        self.weight_transfer_engine.set_weight_update_target(
            draft_model, speculative_config.draft_model_config
        )

    def get_supported_tasks(self) -> tuple[SupportedTask, ...]:
        return self.model_runner.get_supported_tasks()

    def get_compilation_match_table(self) -> dict[str, int]:
        from vllm.compilation.passes.vllm_inductor_pass import get_match_table

        return get_match_table()

    def get_encoder_timing_stats(self) -> dict[str, dict[str, float | int]]:
        """Get encoder timing stats from model runner."""
        return self.model_runner.get_encoder_timing_stats()

    def annotate_profile(self, scheduler_output):
        # add trace annotation so that we can easily distinguish
        # context/generation request numbers in each iteration.
        # A context request is a request that has not yet generated any tokens
        if not self.profiler:
            return nullcontext()

        self.profiler.step()
        if not self.profiler.should_annotate:
            return nullcontext()

        iteration_details = compute_iteration_details(scheduler_output)

        if self.vllm_config.profiler_config.detailed_trace_annotation:
            # Compute roofline-model metrics per request, split by phase
            # (context vs generation). These help estimate compute and
            # memory intensity from the trace.
            #
            # Per-request quantities:
            #   query_len = number of scheduled (new) tokens for this request
            #   seq_len   = total sequence length (computed + scheduled tokens)
            #
            # Aggregated across requests in each phase
            # (ctx_=context, gen_=generation):
            #   seq_len_sum = sum of seq_len   (total KV length)
            #   qq_compute  = sum of query_len*query_len
            #                 (proxy for QK^T compute cost)
            #   qk_compute  = sum of query_len*seq_len
            #                 (proxy for QK^T compute cost for decode and
            #                  chunked prefill)
            #   total_scheduled_tokens = scheduled tokens across all requests
            ctx_seq_len_sum = 0
            ctx_qq_compute = 0
            ctx_qk_compute = 0
            gen_seq_len_sum = 0
            gen_qq_compute = 0
            gen_qk_compute = 0
            total_scheduled_tokens = 0

            # Build a map of req_id -> num_computed_tokens for all requests
            new_req_ids = {
                new_req.req_id for new_req in scheduler_output.scheduled_new_reqs
            }
            num_computed_tokens_ids = {
                new_req.req_id: new_req.num_computed_tokens
                for new_req in scheduler_output.scheduled_new_reqs
            }
            for req_id, num_computed_tokens in zip(
                scheduler_output.scheduled_cached_reqs.req_ids,
                scheduler_output.scheduled_cached_reqs.num_computed_tokens,
            ):
                num_computed_tokens_ids[req_id] = num_computed_tokens

            # Accumulate per-phase metrics
            for req_id, num_tokens in scheduler_output.num_scheduled_tokens.items():
                query_len = num_tokens
                total_scheduled_tokens += query_len
                seq_len = num_computed_tokens_ids.get(req_id, 0) + query_len
                if (
                    scheduler_output.scheduled_cached_reqs.is_context_phase(req_id)
                    or req_id in new_req_ids
                ):
                    ctx_seq_len_sum += seq_len
                    ctx_qq_compute += query_len * query_len
                    ctx_qk_compute += query_len * seq_len
                else:
                    gen_seq_len_sum += seq_len
                    gen_qq_compute += query_len * query_len
                    gen_qk_compute += query_len * seq_len
            annotation = "".join(
                [
                    "execute_",
                    str(total_scheduled_tokens),
                    "_context_",
                    str(iteration_details.num_ctx_requests),
                    "(sq",
                    str(iteration_details.num_ctx_tokens),
                    "sk",
                    str(ctx_seq_len_sum),
                    "sqsq",
                    str(ctx_qq_compute),
                    "sqsk",
                    str(ctx_qk_compute),
                    ")_generation_",
                    str(iteration_details.num_generation_requests),
                    "(sq",
                    str(iteration_details.num_generation_tokens),
                    "sk",
                    str(gen_seq_len_sum),
                    "sqsq",
                    str(gen_qq_compute),
                    "sqsk",
                    str(gen_qk_compute),
                    ")",
                ]
            )
        else:
            annotation = "".join(
                [
                    "execute_context_",
                    str(iteration_details.num_ctx_requests),
                    "(",
                    str(iteration_details.num_ctx_tokens),
                    ")",
                    "_generation_",
                    str(iteration_details.num_generation_requests),
                    "(",
                    str(iteration_details.num_generation_tokens),
                    ")",
                ]
            )
        return self.profiler.annotate_context_manager(annotation)

    @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

        assert isinstance(output, IntermediateTensors)
        parallel_config = self.vllm_config.parallel_config
        assert (
            parallel_config.distributed_executor_backend != "external_launcher"
            and not get_pp_group().is_last_rank
        )

        # Non-blocking send of the intermediate tensors. The metadata handle
        # is reaped lazily by the GroupCoordinator; the device handles are
        # waited at the top of the next step.
        handles = get_pp_group().isend_tensor_dict(
            output.tensors,
            all_gather_group=get_tp_group(),
        )
        self._pp_send_work = handles[1:]

        if self.use_v2_model_runner and self.model_runner.is_pooling_model:
            return self.model_runner.pool()  # type: ignore
        return None

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

    def profile(
        self,
        is_start: bool = True,
        profile_prefix: str | None = None,
        *,
        delay_iterations: int | None = None,
        max_iterations: int | None = None,
    ):
        # Check if profiling is enabled
        if self.profiler_config is None or self.profiler_config.profiler is None:
            raise RuntimeError(
                "Profiling is not enabled. Please set --profiler-config to enable "
                "profiling. Example: "
                "'--profiler-config.profiler=torch --profiler-config.torch_profiler_dir"
                "=YOUR_DIR_PATH_TO_DUMP_TRACE'"
            )

        if is_start:
            validate_profile_prefix(profile_prefix)

            # Generate the trace name by combining prefix with comprehensive rank suffix
            from vllm.distributed.utils import get_worker_rank_suffix

            rank_suffix = get_worker_rank_suffix(global_rank=self.rank)

            # Build the full trace name
            if profile_prefix:
                trace_name = f"{profile_prefix}_{rank_suffix}"
            else:
                trace_name = rank_suffix

            # Create the profiler wrapper only on the first start call
            if self.profiler is None:
                self.profiler = create_worker_profiler(
                    self.profiler_config,
                    worker_name=trace_name,
                    local_rank=self.local_rank,
                )

            self.profiler.set_output_name(trace_name)
            self.profiler.start(
                delay_iterations=delay_iterations,
                max_iterations=max_iterations,
            )
        else:
            if self.profiler is None:
                logger.warning("Profiler was not started, nothing to stop.")
                return
            try:
                self.profiler.stop()
            finally:
                if self.profiler_config.profiler == "proton" and not (
                    self.profiler.has_cuda_graph_session
                ):
                    # Proton output names are fixed when the wrapper is constructed.
                    # Recreate it so the next profile_prefix is honored.
                    self.profiler = None

    @property
    def randomize_dummy_inputs(self) -> bool:
        # Not cached: elastic EP rewrites parallel_config in place on reconfigure.
        return (
            envs.VLLM_RANDOMIZE_DP_DUMMY_INPUTS
            and self.parallel_config.data_parallel_size > 1
        )

    def execute_dummy_batch(self) -> None:
        num_tokens = getattr(self.model_runner, "uniform_decode_query_len", 1)
        self.model_runner._dummy_run(
            num_tokens,
            uniform_decode=True,
            randomize_inputs=self.randomize_dummy_inputs,
        )

    def add_lora(self, lora_request: LoRARequest) -> bool:
        return self.model_runner.add_lora(lora_request)

    def remove_lora(self, lora_id: int) -> bool:
        return self.model_runner.remove_lora(lora_id)

    def list_loras(self) -> set[int]:
        return self.model_runner.list_loras()

    def pin_lora(self, lora_id: int) -> bool:
        return self.model_runner.pin_lora(lora_id)

    def check_health(self) -> None:
        # worker will always be healthy as long as it's running.
        return

    def save_sharded_state(
        self,
        path: str,
        pattern: str | None = None,
        max_size: int | None = None,
    ) -> None:
        from vllm.model_executor.model_loader import ShardedStateLoader

        ShardedStateLoader.save_model(
            self.model_runner.model,
            path,
            pattern=pattern,
            max_size=max_size,
        )

    def save_tensorized_model(self, tensorizer_config: "TensorizerConfig") -> None:
        TensorizerLoader.save_model(
            self.get_model(),
            tensorizer_config=tensorizer_config,
            model_config=self.model_config,
        )

    def _check_weight_transfer_engine(self) -> None:
        if self.weight_transfer_engine is None:
            raise RuntimeError(
                "Weight transfer not configured. "
                "Please set weight_transfer_config to enable weight transfer."
            )

    def init_weight_transfer_engine(self, init_info: dict) -> None:
        """Initialize weight transfer mechanism.
        For NCCL backend, this creates a process group with the trainer.

        Args:
            init_info: Dictionary containing backend-specific initialization info

        """
        self._check_weight_transfer_engine()
        assert self.weight_transfer_engine is not None
        # Parse dict into backend-specific typed dataclass
        typed_init_info = self.weight_transfer_engine.parse_init_info(init_info)
        self.weight_transfer_engine.init_transfer_engine(typed_init_info)

    def start_weight_update(self) -> None:
        """Start a new weight update session.

        Delegates engine-specific preparation (e.g. layerwise reload setup) to
        the configured weight transfer engine. The worker only tracks that a
        session is active.
        """
        with set_current_vllm_config(self.vllm_config):
            self._start_weight_update()

    def start_draft_weight_update(self) -> None:
        """Like start_weight_update, but retargets the engine at the speculative
        draft model for this session.
        """
        with set_current_vllm_config(self.vllm_config):
            self._start_weight_update(is_draft=True)

    def _start_weight_update(self, is_draft: bool = False) -> None:
        self._check_weight_transfer_engine()
        assert self.weight_transfer_engine is not None

        if is_draft and not self.weight_transfer_engine.supports_draft_weight_update:
            raise RuntimeError(
                f"{type(self.weight_transfer_engine).__name__} does not support "
                "draft model weight updates."
            )

        if self._weight_update_active:
            raise RuntimeError(
                "start_weight_update called while a weight update is already "
                "active. Call finish_weight_update first."
            )

        try:
            if is_draft:
                self._set_draft_weight_update_target()
            self.weight_transfer_engine.start_weight_update()
        except BaseException:
            self.weight_transfer_engine.reset_weight_update_target()
            raise
        self._weight_update_active = True
        self._weight_update_is_draft = is_draft

    def update_weights(self, update_info: dict | list[dict]) -> None:
        """Receive one weight update chunk from the trainer.

        start_weight_update must be called before update_weights and
        finish_weight_update must be called after all chunks have been sent.
        Every chunk loads into whichever model the session's start_weight_update
        / start_draft_weight_update call selected.

        Args:
            update_info: Backend-specific update info, or a list indexed by
                global worker rank across data parallel replicas.

        """
        self._check_weight_transfer_engine()
        assert self.weight_transfer_engine is not None

        if not self._weight_update_active:
            raise RuntimeError(
                "start_weight_update must be called before update_weights."
            )

        with set_current_vllm_config(self.vllm_config):
            try:
                if isinstance(update_info, list):
                    parallel_config = self.vllm_config.parallel_config
                    local_update_info = update_info[
                        parallel_config.data_parallel_index * parallel_config.world_size
                        + self.rank
                    ]
                else:
                    local_update_info = update_info
                self.weight_transfer_engine.update_weights(local_update_info)
            except BaseException:
                self._weight_update_active = False
                self.weight_transfer_engine.reset_weight_update_target()
                raise

    def finish_weight_update(self) -> None:
        """Finish the current weight update session."""
        self._check_weight_transfer_engine()
        assert self.weight_transfer_engine is not None

        if not self._weight_update_active:
            raise RuntimeError(
                "finish_weight_update called without a matching start_weight_update."
            )

        with set_current_vllm_config(self.vllm_config):
            self.weight_transfer_engine.finish_weight_update()
            self.weight_transfer_engine.reset_weight_update_target()
            self._weight_update_active = False

        # Weight transfer bypasses GPUModelRunner.reload_weights().
        if not self._weight_update_is_draft:
            self.model_runner.reset_lora_state()

    def shutdown(self) -> None:
        gc.unfreeze()

        # has_kv_transfer_group can be None during interpreter shutdown.
        if ensure_kv_transfer_shutdown is not None:
            ensure_kv_transfer_shutdown()
        if ensure_ec_transfer_shutdown is not None:
            ensure_ec_transfer_shutdown()
        if self.profiler is not None:
            self.profiler.shutdown()

        if weight_transfer_engine := getattr(self, "weight_transfer_engine", None):
            weight_transfer_engine.shutdown()

        self.elastic_ep_executor.shutdown()

        # Release GPU resources held by the model runner so that memory
        # can be reclaimed when running in-process
        if model_runner := getattr(self, "model_runner", None):
            model_runner.shutdown()

        # Release kept-alive cumem pools while the pluggable allocator wrappers
        # and callbacks are still alive, so MemPool teardown is not deferred to
        # interpreter finalization (pytorch/pytorch#145168).
        if current_platform.is_cuda_alike():
            from vllm.device_allocator.cumem import CuMemAllocator

            if CuMemAllocator.instance is not None:
                CuMemAllocator.instance.release_pools()

    def elastic_ep_execute(self, execute_method: str, *args, **kwargs):
        return self.elastic_ep_executor.execute(execute_method, *args, **kwargs)


def init_worker_distributed_environment(
    vllm_config: VllmConfig,
    rank: int,
    distributed_init_method: str | None = None,
    local_rank: int = -1,
    backend: str = "nccl",
) -> None:
    """Initialize the distributed environment."""
    parallel_config = vllm_config.parallel_config
    from vllm.model_executor.determinism.batch_invariant import init_batch_invariance

    init_batch_invariance()
    override_envs_for_eplb(
        parallel_config,
        moe_backend=getattr(vllm_config.kernel_config, "moe_backend", None),
    )
    set_custom_all_reduce(not parallel_config.disable_custom_all_reduce)

    init_method = distributed_init_method or "env://"

    timeout = None
    if parallel_config.distributed_timeout_seconds is not None:
        timeout = timedelta(seconds=parallel_config.distributed_timeout_seconds)

    init_distributed_environment(
        parallel_config.world_size,
        rank,
        init_method,
        local_rank,
        backend,
        timeout,
    )

    ensure_model_parallel_initialized(
        parallel_config.tensor_parallel_size,
        parallel_config.pipeline_parallel_size,
        parallel_config.prefill_context_parallel_size,
        parallel_config.decode_context_parallel_size,
    )

    # Init ec connector here before KV caches init
    # NOTE: We do not init KV caches for Encoder-only instance in EPD disagg mode
    ensure_ec_transfer_initialized(vllm_config)