Skip to content

vllm.model_executor.utils

Utils for model executor.

Functions:

_extent(t)

Number of elements spanned by t in its storage.

Source code in vllm/model_executor/utils.py 
71
72
73
74
75
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 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
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)

get_live_param(layer, name)

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

Source code in vllm/model_executor/utils.py 
269
270
271
272
273
274
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 
318
319
320
321
322
323
324
325
326
327
328
329
330
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 
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
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 
227
228
229
230
231
232
233
234
235
236
237
238
239
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

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 
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
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 
242
243
244
245
246
247
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 
250
251
252
253
254
255
256
257
258
259
260
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 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
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 
78
79
80
81
82
83
84
85
86
87
88
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 
19
20
21
22
23
24
25
26
27
28
29
@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)