vllm.model_executor.utils

Utils for model executor.

Functions:

• copy_weight_ –

  dst.copy_(src) for weight loaders.

• get_live_param –

  getattr(layer, name, None), treating retired or empty params as absent.

• get_moe_expert_mapping –

  Get the expert mapping from a model.

• kernel_parameters –

  The layer's own parameters as kernels see them: live parameters only, each

• move_parameter –

  Register the src_name parameter object under dst_name as well.

• refill_derived –

  Set obj.name (a plain attribute, never a parameter or buffer) to a

• replace_parameter –

  Replace a parameter of a layer while maintaining the ability to reload the

• retire_parameter –

  Mark a loaded parameter that no kernel reads after post-processing.

• set_kernel_parameter –

  replace_parameter, for a tensor kernels read under name: the parameter

• set_weight_attrs –

  Set attributes on a weight tensor.

• storage_view –

  A dtype tensor on storage at element offset. set_ does

• weights_already_processed –

  Mark that weights are already in post-processed (runtime) format, so

_extent(t)

Number of elements spanned by t in its storage.

Source code in vllm/model_executor/utils.py

def _extent(t: torch.Tensor) -> int:
    """Number of elements spanned by ``t`` in its storage."""
    if t.numel() == 0:
        return 0
    return 1 + sum((n - 1) * s for n, s in zip(t.shape, t.stride()))

_in_own_storage(old, new, whole_storage=False)

new copied into the bytes old spans if it fits, so the loaded storage is kept (with its data_ptr); else new itself. An equal-size result fits too: a repack to the same byte count (e.g. FP8 → Marlin int32) then keeps the address captured CUDA graphs read.

With whole_storage (direct reload), new may use the rest of old's storage, so a weight that post-load processing grew keeps its address on a reload too, as do a padded view with gaps (ROCm FP8 padding) and an emptied weight (the checkpoint bytes stay for the next reload).

Source code in vllm/model_executor/utils.py

def _in_own_storage(
    old: torch.Tensor, new: torch.Tensor, whole_storage: bool = False
) -> torch.Tensor:
    """``new`` copied into the bytes ``old`` spans if it fits, so the loaded
    storage is kept (with its ``data_ptr``); else ``new`` itself. An equal-size
    result fits too: a repack to the same byte count (e.g. FP8 → Marlin int32)
    then keeps the address captured CUDA graphs read.

    With ``whole_storage`` (direct reload), ``new`` may use the rest of
    ``old``'s storage, so a weight that post-load processing grew keeps its
    address on a reload too, as do a padded view with gaps (ROCm FP8 padding)
    and an emptied weight (the checkpoint bytes stay for the next reload)."""
    storage = old.untyped_storage()
    offset, size = old.storage_offset() * old.element_size(), new.element_size()
    span = _extent(new)
    if whole_storage:
        fits = (
            span >= new.numel()  # not overlapping
            and offset + span * size <= storage.nbytes()
        )
    else:
        fits = (
            new.numel() > 0
            and span == new.numel()  # neither overlapping nor with gaps
            # its gaps may hold another tensor's data
            and _extent(old) == old.numel()
            and new.numel() * size <= old.numel() * old.element_size()
        )
    if (
        not fits
        or old.device != new.device
        or new.untyped_storage().data_ptr() == storage.data_ptr()
        or offset % size
    ):
        return new
    view = storage_view(storage, offset // size, new.shape, new.stride(), new.dtype)
    return view.copy_(new)

copy_weight_(dst, src)

dst.copy_(src) for weight loaders.

Model-specific weight loaders should call this instead of .copy_() when copying a checkpoint tensor to the device. Off integrated GPUs it is a plain copy_.

On integrated GPUs, a host-to-device copy straight from a safetensors mmap takes its page faults inside the driver copy, which is several times slower than faulting them on the CPU, so touch one byte per page of src first. PyTorch already stages non-contiguous and dtype-converting copies through a CPU temporary. Sources with at most one element are skipped: they count as contiguous with any stride, which view(torch.uint8) rejects.

Source code in vllm/model_executor/utils.py

def copy_weight_(dst: torch.Tensor, src: torch.Tensor) -> torch.Tensor:
    """`dst.copy_(src)` for weight loaders.

    Model-specific weight loaders should call this instead of `.copy_()` when
    copying a checkpoint tensor to the device. Off integrated GPUs it is a
    plain `copy_`.

    On integrated GPUs, a host-to-device copy straight from a safetensors mmap
    takes its page faults inside the driver copy, which is several times slower
    than faulting them on the CPU, so touch one byte per page of `src` first.
    PyTorch already stages non-contiguous and dtype-converting copies through
    a CPU temporary. Sources with at most one element are skipped: they count
    as contiguous with any stride, which `view(torch.uint8)` rejects.
    """
    if (
        src.device.type == "cpu"
        and src.dtype == dst.dtype
        and src.numel() > 1
        and src.is_contiguous()
        and _is_integrated_gpu(dst.device)
    ):
        flat = src.reshape(-1).view(torch.uint8)
        flat[:: mmap.PAGESIZE].sum()
        flat[-1:].sum()
    return dst.copy_(src)

get_live_param(layer, name)

getattr(layer, name, None), treating retired or empty params as absent.

Source code in vllm/model_executor/utils.py

def get_live_param(layer: torch.nn.Module, name: str) -> torch.Tensor | None:
    """``getattr(layer, name, None)``, treating retired or empty params as absent."""
    tensor = getattr(layer, name, None)
    if isinstance(tensor, torch.nn.Parameter) and not _is_live(tensor):
        return None
    return tensor

get_moe_expert_mapping(model)

Get the expert mapping from a model.

It will be retrieved from the first module that has a get_expert_mapping method. If the model manually implements get_expert_mapping, it will be used. Otherwise, it will use the first RoutedExperts layer.

Source code in vllm/model_executor/utils.py

def get_moe_expert_mapping(
    model: torch.nn.Module,
) -> list[tuple[str, str, int, str]]:
    """Get the expert mapping from a model.

    It will be retrieved from the first module that has a `get_expert_mapping` method.
    If the model manually implements `get_expert_mapping`, it will be used.
    Otherwise, it will use the first RoutedExperts layer."""
    for _, module in model.named_modules():
        get_mapping = getattr(module, "get_expert_mapping", None)
        if get_mapping is not None:
            return get_mapping()
    raise ValueError("No module in the model has a `get_expert_mapping` method.")

kernel_parameters(layer)

The layer's own parameters as kernels see them: live parameters only, each object once, under its runtime name (see move_parameter).

Source code in vllm/model_executor/utils.py

def kernel_parameters(layer: torch.nn.Module) -> dict[str, torch.nn.Parameter]:
    """The layer's own parameters as kernels see them: live parameters only, each
    object once, under its runtime name (see `move_parameter`)."""
    names: dict[int, str] = {}
    params: dict[str, torch.nn.Parameter] = {}
    for name, param in layer._parameters.items():
        if not _is_live(param):
            continue
        first = names.setdefault(id(param), name)
        if first == name:
            params[name] = param
        elif getattr(param, "runtime_name", None) == name:
            del params[first]
            params[name] = param
            names[id(param)] = name
    return params

move_parameter(layer, src_name, dst_name)

Register the src_name parameter object under dst_name as well.

Both names then refer to one object and so to one storage: the loader writes through the checkpoint name, kernels read the runtime name, and a later replace_parameter on either name updates both. dst_name becomes the name kernel_parameters reports for the object.

Source code in vllm/model_executor/utils.py

def move_parameter(layer: torch.nn.Module, src_name: str, dst_name: str) -> None:
    """Register the ``src_name`` parameter object under ``dst_name`` as well.

    Both names then refer to one object and so to one storage: the loader writes
    through the checkpoint name, kernels read the runtime name, and a later
    `replace_parameter` on either name updates both. ``dst_name`` becomes the
    name `kernel_parameters` reports for the object."""
    param = getattr(layer, src_name)
    existing = layer._parameters.get(dst_name)
    if existing is not None and existing is not param:
        raise ValueError(f"{dst_name!r} is already registered as another parameter")
    layer.register_parameter(dst_name, param)
    param.runtime_name = dst_name

refill_derived(obj, name, new)

refill_derived(obj: object, name: str, new: None) -> None

refill_derived(
    obj: object, name: str, new: torch.Tensor
) -> torch.Tensor

Set obj.name (a plain attribute, never a parameter or buffer) to a tensor derived from the weights. A held tensor of the same shape, dtype and device is refilled in place, so captured CUDA graphs keep its address; it must own its storage. Returns the tensor now held under name.

Source code in vllm/model_executor/utils.py

def refill_derived(
    obj: object, name: str, new: torch.Tensor | None
) -> torch.Tensor | None:
    """Set ``obj.name`` (a plain attribute, never a parameter or buffer) to a
    tensor derived from the weights. A held tensor of the same shape, dtype and
    device is refilled in place, so captured CUDA graphs keep its address; it
    must own its storage. Returns the tensor now held under ``name``."""
    old = getattr(obj, name, None)
    if (
        isinstance(old, torch.Tensor)
        and new is not None
        and (old.shape, old.dtype, old.device) == (new.shape, new.dtype, new.device)
    ):
        return old.copy_(new)
    setattr(obj, name, new)
    return new

replace_parameter(layer, param_name, new_tensor)

Replace a parameter of a layer while maintaining the ability to reload the weight. Called within implementations of the process_weights_after_loading method.

The existing parameter object is kept (and so its class and loader attributes for weight reloading): new_tensor is copied into it when shape, stride, dtype and device match, which preserves data_ptr for captured CUDA graphs; otherwise a smaller new_tensor is copied into the parameter's own storage (so the loaded storage never shrinks), else .data is rebound to it. A new parameter is registered only if none exists under param_name, or if the existing one is an online requantization source (released after conversion).

Custom attributes set on new_tensor (e.g. kernel dispatch flags such as is_shuffled) are carried over to the parameter, except weight_loader, which is always taken from the existing parameter.

Under direct reload, a new_tensor larger than the parameter at cold start may be written into the parameter's existing storage on a reload, so do not read a local alias of the old data (w = layer.weight.data) after this call.

This function should not be called on weights which are tied/shared

Parameters:

• layer

  (Module) –

  Layer containing parameter to replace

• param_name

  (str) –

  Name of parameter to replace

• new_tensor

  (Tensor | None) –

  New data of the new parameter, or None to set the parameter to None

Source code in vllm/model_executor/utils.py

def replace_parameter(
    layer: torch.nn.Module,
    param_name: str,
    new_tensor: torch.Tensor | None,
):
    """Replace a parameter of a layer while maintaining the ability to reload the
    weight. Called within implementations of the `process_weights_after_loading` method.

    The existing parameter object is kept (and so its class and loader attributes
    for weight reloading): ``new_tensor`` is copied into it when shape, stride,
    dtype and device match, which preserves ``data_ptr`` for captured CUDA graphs;
    otherwise a smaller ``new_tensor`` is copied into the parameter's own storage
    (so the loaded storage never shrinks), else ``.data`` is rebound to it. A new
    parameter is registered only if none exists under ``param_name``, or if the
    existing one is an online requantization source (released after conversion).

    Custom attributes set on ``new_tensor`` (e.g. kernel dispatch flags such as
    ``is_shuffled``) are carried over to the parameter, except ``weight_loader``,
    which is always taken from the existing parameter.

    Under direct reload, a ``new_tensor`` larger than the parameter at cold start
    may be written into the parameter's existing storage on a reload, so do not
    read a local alias of the old data (``w = layer.weight.data``) after this call.

    This function should not be called on weights which are tied/shared

    Args:
        layer: Layer containing parameter to replace
        param_name: Name of parameter to replace
        new_tensor: New data of the new parameter, or None to set the parameter
            to None

    """
    # should not be used on a tied/shared param

    # If new_tensor is None, set the parameter to None
    if new_tensor is None:
        setattr(layer, param_name, None)
        return

    old_param: torch.nn.Parameter | None = getattr(layer, param_name, None)

    new_tensor_attrs = dict(new_tensor.__dict__)

    # `weight_loader` is the only attribute not ported over from new_tensor,
    # old_param.weight_loader only is supported. `_weight_loader` is
    # `BasevLLMParameter`'s backing field for its `weight_loader` property, so
    # a loader riding along under that name has to be dropped as well.
    new_tensor_attrs.pop("weight_loader", None)
    new_tensor_attrs.pop("_weight_loader", None)
    # A checkpoint record belongs to the parameter it was recorded on (direct
    # reload); one on new_tensor (another checkpoint parameter) must not overwrite it.
    new_tensor_attrs.pop("_ckpt", None)

    if isinstance(new_tensor, torch.nn.Parameter):
        new_tensor = new_tensor.data

    # Keep the object, except online requantization sources (released later)
    # and meta params (no storage; `.data` cannot be rebound to a real tensor).
    if (
        isinstance(old_param, torch.nn.Parameter)
        and not getattr(old_param, "requantization_source", False)
        and not (old_param.is_meta and not new_tensor.is_meta)
    ):
        if (
            old_param.shape == new_tensor.shape
            and old_param.stride() == new_tensor.stride()
            and old_param.dtype == new_tensor.dtype
            and old_param.device == new_tensor.device
        ):
            old_param.data.copy_(new_tensor)
        elif getattr(getattr(layer, "quant_method", None), "uses_meta_device", False):
            old_param.data = new_tensor  # online quantization frees the source
        else:
            from vllm.model_executor.model_loader.reload.direct import (
                direct_reload_enabled,
            )

            old_param.data = _in_own_storage(
                old_param, new_tensor, whole_storage=direct_reload_enabled()
            )
        for attr_name, attr in new_tensor_attrs.items():
            setattr(old_param, attr_name, attr)
        return

    new_param = torch.nn.Parameter(new_tensor, requires_grad=False)

    for attr_name, attr in new_tensor_attrs.items():
        setattr(new_param, attr_name, attr)

    if old_param is not None and hasattr(old_param, "weight_loader"):
        weight_loader = old_param.weight_loader
        set_weight_attrs(new_param, {"weight_loader": weight_loader})

    setattr(layer, param_name, new_param)

retire_parameter(layer, name)

Mark a loaded parameter that no kernel reads after post-processing.

It stays registered with its loaded data (so it can be loaded again); kernel_parameters and get_live_param skip it.

Source code in vllm/model_executor/utils.py

def retire_parameter(layer: torch.nn.Module, name: str) -> None:
    """Mark a loaded parameter that no kernel reads after post-processing.

    It stays registered with its loaded data (so it can be loaded again);
    `kernel_parameters` and `get_live_param` skip it."""
    getattr(layer, name).retired = True

set_kernel_parameter(layer, name, tensor)

replace_parameter, for a tensor kernels read under name: the parameter is live even if it was retired before, and name becomes its runtime name if it is registered under several names (see move_parameter).

Source code in vllm/model_executor/utils.py

def set_kernel_parameter(
    layer: torch.nn.Module, name: str, tensor: torch.Tensor
) -> None:
    """`replace_parameter`, for a tensor kernels read under ``name``: the parameter
    is live even if it was retired before, and ``name`` becomes its runtime name
    if it is registered under several names (see `move_parameter`)."""
    replace_parameter(layer, name, tensor)
    param = getattr(layer, name)
    param.__dict__.pop("retired", None)
    if hasattr(param, "runtime_name"):
        param.runtime_name = name

set_weight_attrs(weight, weight_attrs)

Set attributes on a weight tensor.

This method is used to set attributes on a weight tensor. This method will not overwrite existing attributes.

Parameters:

• weight

  (Tensor) –

  The weight tensor.

• weight_attrs

  (dict[str, Any] | None) –

  A dictionary of attributes to set on the weight tensor.

Source code in vllm/model_executor/utils.py

def set_weight_attrs(
    weight: torch.Tensor,
    weight_attrs: dict[str, Any] | None,
):
    """Set attributes on a weight tensor.

    This method is used to set attributes on a weight tensor. This method
    will not overwrite existing attributes.

    Args:
        weight: The weight tensor.
        weight_attrs: A dictionary of attributes to set on the weight tensor.

    """
    if weight_attrs is None:
        return
    for key, value in weight_attrs.items():
        assert not hasattr(weight, key), f"Overwriting existing tensor attribute: {key}"

        # NOTE(woosuk): During weight loading, we often do something like:
        # narrowed_tensor = param.data.narrow(0, offset, len)
        # narrowed_tensor.copy_(real_weight)
        # expecting narrowed_tensor and param.data to share the same storage.
        # However, on TPUs, narrowed_tensor will lazily propagate to the base
        # tensor, which is param.data, leading to the redundant memory usage.
        # This sometimes causes OOM errors during model loading. To avoid this,
        # we sync the param tensor after its weight loader is called.
        # TODO(woosuk): Remove this hack once we have a better solution.
        from vllm.platforms import current_platform

        if current_platform.use_sync_weight_loader() and key == "weight_loader":
            value = current_platform.make_synced_weight_loader(value)
        setattr(weight, key, value)

storage_view(storage, offset, shape, stride, dtype)

A dtype tensor on storage at element offset. set_ does not bounds-check: callers check that the view fits.

Source code in vllm/model_executor/utils.py

def storage_view(
    storage: torch.UntypedStorage,
    offset: int,
    shape: tuple[int, ...],
    stride: tuple[int, ...],
    dtype: torch.dtype,
) -> torch.Tensor:
    """A ``dtype`` tensor on ``storage`` at element ``offset``. ``set_`` does
    not bounds-check: callers check that the view fits."""
    view = torch.empty(0, dtype=dtype, device=storage.device)
    return view.set_(storage, offset, shape, stride)

weights_already_processed()

Mark that weights are already in post-processed (runtime) format, so process_weights_after_loading skips tensor transforms (used by the weight cache IPC loader).

Source code in vllm/model_executor/utils.py

@contextmanager
def weights_already_processed():
    """Mark that weights are already in post-processed (runtime) format, so
    ``process_weights_after_loading`` skips tensor transforms (used by the
    weight cache IPC loader).
    """
    token = _weights_pre_processed.set(True)
    try:
        yield
    finally:
        _weights_pre_processed.reset(token)