| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import UnionType |
| 4 | from typing import Annotated, Any, Union, get_args, get_origin, get_type_hints |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.logger import init_logger |
| 9 | |
| 10 | logger = init_logger(__name__) |
| 11 | |
| 12 | |
| 13 | class TensorShape: |
| 14 | def __init__( |
| 15 | self, |
| 16 | *dims: int | str, |
| 17 | dynamic_dims: set[str] | None = None, |
| 18 | ) -> None: |
| 19 | super().__init__() |
| 20 | |
| 21 | self.dims = dims |
| 22 | self.dynamic_dims = dynamic_dims if dynamic_dims else set() |
| 23 | |
| 24 | def resolve(self, **bindings: int) -> tuple[int | str, ...]: |
| 25 | resolved = list[int | str]() |
| 26 | for dim in self.dims: |
| 27 | if isinstance(dim, str) and dim in bindings: |
| 28 | resolved.append(bindings[dim]) |
| 29 | else: |
| 30 | resolved.append(dim) |
| 31 | return tuple(resolved) |
| 32 | |
| 33 | def __str__(self) -> str: |
| 34 | """Return a string representation of the tensor shape.""" |
| 35 | dim_strs = [] |
| 36 | for dim in self.dims: |
| 37 | if isinstance(dim, str): |
| 38 | if dim in self.dynamic_dims: |
| 39 | dim_strs.append(f"{dim}*") # Mark dynamic dimensions with * |
| 40 | else: |
| 41 | dim_strs.append(dim) |
| 42 | else: |
| 43 | dim_strs.append(str(dim)) |
| 44 | return f"({', '.join(dim_strs)})" |
| 45 | |
| 46 | |
| 47 | class TensorSchema: |
| 48 | def __init__( |
| 49 | self, |
| 50 | *, |
| 51 | validate: bool = True, |
| 52 | resolve_bindings: dict[str, int] | None = None, |
| 53 | **kwargs: Any, |
| 54 | ) -> None: |
| 55 | super().__init__() |
| 56 | |
| 57 | self._resolve_bindings = resolve_bindings if resolve_bindings else {} |
| 58 | |
| 59 | for key, value in kwargs.items(): |
| 60 | setattr(self, key, value) |
| 61 | |
| 62 | if validate: |
| 63 | self.validate() |
| 64 | |
| 65 | def __getitem__(self, key: str) -> Any: |
| 66 | return getattr(self, key) |
| 67 | |
| 68 | def get(self, key: str, default: Any = None) -> Any: |
| 69 | return getattr(self, key, default) |
| 70 | |
| 71 | def _match_shape_with_dynamic( |
| 72 | self, |
| 73 | actual: tuple[int, ...], |
| 74 | reference: tuple[int, ...], |
| 75 | expected_shape: tuple[int | str, ...], |
| 76 | dynamic_dims: set[str], |
| 77 | ) -> bool: |
| 78 | if len(actual) != len(reference) or len(actual) > len(expected_shape): |
| 79 | return False |
| 80 | |
| 81 | for i, (a, r) in enumerate(zip(actual, reference)): |
| 82 | # When validating list inputs, we match shape suffixes only |
| 83 | # (e.g. "p", 3, "h", "w"), assuming the list length corresponds |
| 84 | # to the leading symbolic dim (e.g. "bn"). This allows comparing |
| 85 | # only the trailing dimensions of each element in the list. |
| 86 | dim = expected_shape[-len(actual) + i] |
| 87 | # Skip this dimension if it's marked dynamic |
| 88 | if dim in dynamic_dims: |
| 89 | continue |
| 90 | if a != r: |
| 91 | return False |
| 92 | return True |
| 93 | |
| 94 | def _fmt_indexer(self, idxs: tuple[int, ...]) -> str: |
| 95 | if not idxs: |
| 96 | return "" |
| 97 | |
| 98 | return str(list(idxs)) |
| 99 | |
| 100 | def _validate_field( |
| 101 | self, |
| 102 | value: object, |
| 103 | field_name: str, |
| 104 | expected_shape: tuple[int | str, ...], |
| 105 | dynamic_dims: set[str], |
| 106 | leading_idxs: tuple[int, ...] = (), |
| 107 | ) -> tuple[int, ...]: |
| 108 | """Validate a field and return the actual shape.""" |
| 109 | if isinstance(value, (int, float)): |
| 110 | return () # Scalar |
| 111 | if isinstance(value, torch.Tensor): |
| 112 | return value.shape |
| 113 | |
| 114 | if not isinstance(value, (list, tuple)): |
| 115 | raise TypeError( |
| 116 | f"{field_name}{self._fmt_indexer(leading_idxs)} is not " |
| 117 | f"one of the expected types: int, float, Tensor, list, tuple. " |
| 118 | f"Got: {type(value)}" |
| 119 | ) |
| 120 | |
| 121 | if len(value) == 0: |
| 122 | raise ValueError( |
| 123 | f"{field_name}{self._fmt_indexer(leading_idxs)} is an empty sequence" |
| 124 | ) |
| 125 | |
| 126 | # Ensure all tensors in the list have the same |
| 127 | # shape, besides dynamic dimensions |
| 128 | for i, v in enumerate(value): |
| 129 | shape = self._validate_field( |
| 130 | v, |
| 131 | field_name, |
| 132 | expected_shape[1:], |
| 133 | dynamic_dims, |
| 134 | leading_idxs=leading_idxs + (i,), |
| 135 | ) |
| 136 | |
| 137 | if i == 0: |
| 138 | first_shape = shape |
| 139 | elif not self._match_shape_with_dynamic( |
| 140 | shape, |
| 141 | first_shape, |
| 142 | expected_shape, |
| 143 | dynamic_dims, |
| 144 | ): |
| 145 | raise ValueError( |
| 146 | f"{field_name}{self._fmt_indexer(leading_idxs)} " |
| 147 | f"contains inconsistent shapes: {first_shape} " |
| 148 | f"(index 0) vs {shape} (index {i})" |
| 149 | ) |
| 150 | |
| 151 | # Treat the list as a stacked tensor: |
| 152 | # shape = (len(list), *tensor.shape) |
| 153 | return (len(value),) + first_shape |
| 154 | |
| 155 | def _validate_tensor_shape_expected( |
| 156 | self, |
| 157 | actual_shape: tuple[int, ...], |
| 158 | expected_shape: tuple[int | str, ...], |
| 159 | field_name: str, |
| 160 | shape_env: dict[str, int], |
| 161 | dynamic_dims: set[str], |
| 162 | ) -> None: |
| 163 | """Validate that the actual tensor shape matches the expected shape.""" |
| 164 | if len(actual_shape) != len(expected_shape): |
| 165 | raise ValueError( |
| 166 | f"{field_name} has rank {len(actual_shape)} " |
| 167 | f"but expected {len(expected_shape)}. " |
| 168 | f"Expected shape: {expected_shape}, " |
| 169 | f"but got {actual_shape}" |
| 170 | ) |
| 171 | |
| 172 | for i, dim in enumerate(expected_shape): |
| 173 | if dim in dynamic_dims: |
| 174 | continue |
| 175 | elif isinstance(dim, int): |
| 176 | if actual_shape[i] != dim: |
| 177 | raise ValueError( |
| 178 | f"{field_name} dim[{i}] expected " |
| 179 | f"{dim}, got {actual_shape[i]}. " |
| 180 | f"Expected shape: {expected_shape}, " |
| 181 | f"but got {actual_shape}" |
| 182 | ) |
| 183 | elif isinstance(dim, str): |
| 184 | if dim in shape_env: |
| 185 | if actual_shape[i] != shape_env[dim]: |
| 186 | raise ValueError( |
| 187 | f"{field_name} dim[{i}] expected " |
| 188 | f"'{dim}'={shape_env[dim]}, got " |
| 189 | f"{actual_shape[i]}" |
| 190 | ) |
| 191 | else: |
| 192 | shape_env[dim] = actual_shape[i] |
| 193 | else: |
| 194 | raise TypeError( |
| 195 | f"{field_name} dim[{i}] has unsupported type: {type(dim)}" |
| 196 | ) |
| 197 | |
| 198 | def validate(self) -> None: |
| 199 | type_hints = get_type_hints(self.__class__, include_extras=True) |
| 200 | shape_env = dict[str, int]() |
| 201 | |
| 202 | for field_name, field_type in type_hints.items(): |
| 203 | # Check if field is missing |
| 204 | if not hasattr(self, field_name) or getattr(self, field_name) is None: |
| 205 | # Check if field is marked as optional |
| 206 | actual_type = field_type |
| 207 | if get_origin(field_type) is Annotated: |
| 208 | args = get_args(field_type) |
| 209 | actual_type = args[0] |
| 210 | |
| 211 | # Check arg was provided as Union |
| 212 | if get_origin(actual_type) in {Union, UnionType}: |
| 213 | # Union for Union[X, Y] and UnionType for X | Y |
| 214 | args = get_args(actual_type) |
| 215 | # Skip validation when Union contains None |
| 216 | if type(None) in args: |
| 217 | continue |
| 218 | # Otherwise field is required, raise error |
| 219 | raise ValueError(f"Required field '{field_name}' is missing") |
| 220 | |
| 221 | # Field exists, proceed with validation |
| 222 | value = getattr(self, field_name) |
| 223 | if get_origin(field_type) is not None: |
| 224 | args = get_args(field_type) |
| 225 | |
| 226 | for arg in args: |
| 227 | if isinstance(arg, TensorShape): |
| 228 | expected_shape = arg.resolve(**self._resolve_bindings) |
| 229 | actual_shape = self._validate_field( |
| 230 | value, |
| 231 | field_name, |
| 232 | expected_shape, |
| 233 | arg.dynamic_dims, |
| 234 | ) |
| 235 | |
| 236 | self._validate_tensor_shape_expected( |
| 237 | actual_shape, |
| 238 | expected_shape, |
| 239 | field_name, |
| 240 | shape_env, |
| 241 | arg.dynamic_dims, |
| 242 | ) |
| 243 | |
| 244 | def print_shapes(self) -> None: |
| 245 | """Print TensorShape annotations for debugging.""" |
| 246 | logger.debug("Shapes in %s:", self.__class__.__name__) |
| 247 | type_hints = get_type_hints(self.__class__, include_extras=True) |
| 248 | |
| 249 | for field_name, field_type in type_hints.items(): |
| 250 | if get_origin(field_type) is not None: |
| 251 | args = get_args(field_type) |
| 252 | for arg in args: |
| 253 | if isinstance(arg, TensorShape): |
| 254 | logger.debug(" %s: %s", field_name, str(arg)) |