chore: import upstream snapshot with attribution

python/sglang/srt/layers/utils/common.py

@@ -0,0 +1,127 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import logging
+import re
+
+import torch
+from torch.nn.parameter import Parameter
+
+logger = logging.getLogger(__name__)
+
+
+def get_layer_id(weight_name):
+    # example weight name: model.layers.10.self_attn.qkv_proj.weight
+    match = re.search(r"layers\.(\d+)\.", weight_name)
+    if match:
+        return int(match.group(1))
+    return None
+
+
+def pad_or_narrow_weight(
+    loaded_weight: torch.Tensor, input_dim: int, start_idx: int, shard_size: int
+) -> torch.Tensor:
+    # Padding with zeros for special case such as qwen2_5_VL's mlp which is not 8-aligned
+    valid_size = max(loaded_weight.shape[input_dim] - start_idx, 0)
+
+    if valid_size > 0:
+        loaded_slice = loaded_weight.narrow(input_dim, start_idx, valid_size)
+        pad_shape = list(loaded_weight.shape)
+        pad_shape[input_dim] = shard_size - valid_size
+        pad = torch.zeros(
+            pad_shape, dtype=loaded_weight.dtype, device=loaded_weight.device
+        )
+        return torch.cat([loaded_slice, pad], dim=input_dim)
+
+    # All padding
+    pad_shape = list(loaded_weight.shape)
+    pad_shape[input_dim] = shard_size
+    return torch.zeros(
+        pad_shape, dtype=loaded_weight.dtype, device=loaded_weight.device
+    )
+
+
+def is_strict_contiguous(x: torch.Tensor) -> bool:
+    expected_stride = 1
+    for size, stride in zip(reversed(x.shape), reversed(x.stride())):
+        if stride != expected_stride:
+            return False
+        expected_stride *= size
+    return True
+
+
+def strict_contiguous(x: torch.Tensor) -> torch.Tensor:
+    if is_strict_contiguous(x):
+        return x
+    return x.clone(memory_format=torch.contiguous_format)
+
+
+def copy_or_rebind_param(
+    module: torch.nn.Module, name: str, new_value: torch.Tensor
+) -> None:
+    """Keep parameter identities stable for CUDA graph reuse and hot reload."""
+    new_value = new_value.detach()
+    param = getattr(module, name, None)
+    if isinstance(param, Parameter):
+        if param.data.shape == new_value.shape and param.data.dtype == new_value.dtype:
+            param.data.copy_(new_value)
+        else:
+            param.data = new_value
+        param.requires_grad_(False)
+    else:
+        setattr(module, name, Parameter(new_value, requires_grad=False))
+
+
+def alias_or_bind_derived_param(
+    module: torch.nn.Module,
+    source_name: str,
+    derived_name: str,
+    derived_value: torch.Tensor,
+) -> None:
+    """Bind a post-processed (derived) tensor to a derived attribute name.
+
+    When `derived_value` is broadcastable to the source Parameter's shape (and
+    dtype matches), write it broadcast-filled into the source's storage in
+    place and register `derived_name` as an alias of the source Parameter. The
+    two attribute names then share one underlying buffer, so:
+      - apply() can read via `derived_name`
+      - update_weights_from_disk can keep refilling `source_name` (the loader
+        re-runs process_weights_after_loading which re-derives in place)
+      - peak GPU memory is the source size, not source + derived.
+
+    When the shapes are not broadcast-compatible, fall back to allocating a
+    separate Parameter under `derived_name` via copy_or_rebind_param.
+    """
+    derived_value = derived_value.detach()
+    source = getattr(module, source_name, None)
+    if isinstance(source, Parameter) and source.data.dtype == derived_value.dtype:
+        try:
+            broadcast = torch.broadcast_to(derived_value, source.data.shape)
+        except RuntimeError:
+            broadcast = None
+        if broadcast is not None:
+            source.data.copy_(broadcast)
+            source.requires_grad_(False)
+            setattr(module, derived_name, source)
+            return
+    copy_or_rebind_param(module, derived_name, derived_value)
+
+
+class PPMissingLayer(torch.nn.Identity):
+    # Adapted from
+    # https://github.com/vllm-project/vllm/blob/18ed3132d2bfe1df9a74729457b69243955221e8/vllm/model_executor/models/utils.py#L468C1-L486C1
+    """
+    A placeholder layer for missing layers in a pipeline parallel model.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        self.return_tuple = kwargs.get("return_tuple", False)
+
+    def forward(self, *args, **kwargs):
+        """
+        Return the first arg from args or the first value from kwargs.
+
+        Wraps the input in a tuple if `self.return_tuple` is True.
+        """
+        input = args[0] if args else next(iter(kwargs.values()))
+        return (input,) if self.return_tuple else input