| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """Fused Normalisation on the Device. |
| 5 | |
| 6 | Equivalent to:: |
| 7 | |
| 8 | output = (input * rescale_factor - image_mean) / image_std |
| 9 | |
| 10 | This is implemented as a single per-channel affine transform:: |
| 11 | |
| 12 | output = input * weight[c] + bias[c] |
| 13 | |
| 14 | where:: |
| 15 | |
| 16 | weight = rescale_factor / image_std |
| 17 | bias = -image_mean / image_std |
| 18 | """ |
| 19 | |
| 20 | import math |
| 21 | from typing import NamedTuple |
| 22 | |
| 23 | import torch |
| 24 | from torch import nn |
| 25 | |
| 26 | from vllm.config import ModelConfig |
| 27 | from vllm.model_executor.custom_op import CustomOp |
| 28 | from vllm.transformers_utils.processor import get_processor_config |
| 29 | from vllm.triton_utils import tl, triton |
| 30 | |
| 31 | |
| 32 | @triton.jit |
| 33 | def _fused_mm_input_norm_kernel( |
| 34 | x_ptr, |
| 35 | y_ptr, |
| 36 | w_ptr, |
| 37 | b_ptr, |
| 38 | numel, |
| 39 | L, |
| 40 | C: tl.constexpr, |
| 41 | BLOCK: tl.constexpr, |
| 42 | ): |
| 43 | # 1D grid over the flattened (N, C, L) tensor. Each program processes |
| 44 | # BLOCK contiguous elements; the channel index is recovered from the |
| 45 | # flat offset via `(offs // L) % C`. int32 offsets: int64 division is |
| 46 | # ~2x slower here (measured on GB200); the wrapper asserts |
| 47 | # numel < 2**31 so the flat index never overflows. |
| 48 | pid = tl.program_id(0) |
| 49 | offs = pid * BLOCK + tl.arange(0, BLOCK) |
| 50 | mask = offs < numel |
| 51 | c = (offs // L) % C |
| 52 | |
| 53 | x = tl.load( |
| 54 | x_ptr + offs, |
| 55 | mask=mask, |
| 56 | other=0, |
| 57 | eviction_policy="evict_first", |
| 58 | ).to(tl.float32) |
| 59 | w = tl.load(w_ptr + c, mask=mask, other=0).to(tl.float32) |
| 60 | b = tl.load(b_ptr + c, mask=mask, other=0).to(tl.float32) |
| 61 | tl.store( |
| 62 | y_ptr + offs, |
| 63 | x * w + b, |
| 64 | mask=mask, |
| 65 | eviction_policy="evict_first", |
| 66 | ) |
| 67 | |
| 68 | |
| 69 | def fused_mm_input_norm_triton( |
| 70 | inputs: torch.Tensor, |
| 71 | outputs: torch.Tensor, |
| 72 | weight: torch.Tensor, |
| 73 | bias: torch.Tensor, |
| 74 | ) -> torch.Tensor: |
| 75 | """Fused per-channel affine transform for normalisation. |
| 76 | |
| 77 | Args: |
| 78 | inputs: Input tensor, shape ``(N, C, L)``. A contiguous copy is |
| 79 | materialized internally if needed. |
| 80 | outputs: Output tensor, contiguous and shaped exactly ``(N, C, L)``. |
| 81 | weight: Per-channel scale, shape ``(C,)``. |
| 82 | bias: Per-channel shift, shape ``(C,)``. |
| 83 | |
| 84 | Returns: |
| 85 | ``outputs``, for chaining. |
| 86 | |
| 87 | """ |
| 88 | N, C, L = inputs.shape |
| 89 | assert outputs.shape == (N, C, L), ( |
| 90 | f"outputs shape {tuple(outputs.shape)} != inputs shape {(N, C, L)}" |
| 91 | ) |
| 92 | assert weight.numel() == C and bias.numel() == C, ( |
| 93 | f"weight/bias must have {C} elements, got {weight.numel()} / {bias.numel()}" |
| 94 | ) |
| 95 | # The flat 1D kernel writes outputs by raw offset; a non-contiguous |
| 96 | # buffer would be written at the wrong locations and cannot be |
| 97 | # auto-materialized without breaking the caller's aliasing. |
| 98 | assert outputs.is_contiguous(), "outputs must be contiguous" |
| 99 | |
| 100 | # Tile size along the flattened element axis. 1024 with 4 warps gives |
| 101 | # 8 elements per lane and benchmarks fastest across realistic image |
| 102 | # sizes; larger tiles starve the SMs on small inputs. |
| 103 | BLOCK_SIZE = 1024 |
| 104 | |
| 105 | numel = N * C * L |
| 106 | # The kernel indexes with int32 offsets (int64 division costs ~2x |
| 107 | # throughput); guard the flat index range here. |
| 108 | assert numel < 2**31, f"numel={numel} exceeds the int32 index range" |
| 109 | |
| 110 | grid = (triton.cdiv(numel, BLOCK_SIZE),) |
| 111 | |
| 112 | _fused_mm_input_norm_kernel[grid]( |
| 113 | inputs.contiguous(), |
| 114 | outputs, |
| 115 | weight.contiguous(), |
| 116 | bias.contiguous(), |
| 117 | numel, |
| 118 | L, |
| 119 | C=C, |
| 120 | BLOCK=BLOCK_SIZE, |
| 121 | num_warps=4, # 8 elements per lane at BLOCK_SIZE=1024 |
| 122 | ) |
| 123 | return outputs |
| 124 | |
| 125 | |
| 126 | class NormParams(NamedTuple): |
| 127 | """Resolved per-channel affine parameters (flags already folded in).""" |
| 128 | |
| 129 | image_mean: list[float] |
| 130 | image_std: list[float] |
| 131 | rescale_factor: float |
| 132 | |
| 133 | @property |
| 134 | def is_identity(self) -> bool: |
| 135 | """Whether ``weight = rescale/std`` and ``bias = -mean/std`` are |
| 136 | numerically the identity transform (mirrors ``torch.allclose``'s |
| 137 | default fp32 tolerances).""" |
| 138 | return all( |
| 139 | math.isclose(self.rescale_factor / s, 1.0, rel_tol=1e-5, abs_tol=1e-8) |
| 140 | and math.isclose(m / s, 0.0, abs_tol=1e-8) |
| 141 | for m, s in zip(self.image_mean, self.image_std) |
| 142 | ) |
| 143 | |
| 144 | |
| 145 | def _load_norm_params(model_config: ModelConfig) -> NormParams: |
| 146 | """Resolve the per-channel affine parameters ``(image_mean, image_std, |
| 147 | rescale_factor)`` from the processor config.""" |
| 148 | config = get_processor_config(model_config.model, revision=model_config.revision) |
| 149 | config = config.get("image_processor", config) |
| 150 | if not any( |
| 151 | key in config |
| 152 | for key in ("do_normalize", "do_rescale", "image_mean", "image_std") |
| 153 | ): |
| 154 | config = model_config.hf_image_processor_config |
| 155 | |
| 156 | has_norm_params = "image_mean" in config and "image_std" in config |
| 157 | do_normalize = bool(config.get("do_normalize", has_norm_params)) |
| 158 | do_rescale = bool(config.get("do_rescale", has_norm_params)) |
| 159 | |
| 160 | # Parameters whose flag is off are unused; default them to no-ops |
| 161 | # without resolving them. |
| 162 | rescale_factor = config.get("rescale_factor", 1 / 255) if do_rescale else 1.0 |
| 163 | image_mean = config.get("image_mean") if do_normalize else [0.0] * 3 |
| 164 | image_std = config.get("image_std") if do_normalize else [1.0] * 3 |
| 165 | |
| 166 | assert rescale_factor is not None, "rescale_factor is still None after resolution." |
| 167 | assert image_mean is not None, "image_mean is still None after resolution." |
| 168 | assert image_std is not None, "image_std is still None after resolution." |
| 169 | assert len(image_mean) == len(image_std), ( |
| 170 | f"image_mean and image_std have different lengths: " |
| 171 | f"{len(image_mean)} vs {len(image_std)}" |
| 172 | ) |
| 173 | |
| 174 | return NormParams( |
| 175 | image_mean=[float(v) for v in image_mean], |
| 176 | image_std=[float(v) for v in image_std], |
| 177 | rescale_factor=float(rescale_factor), |
| 178 | ) |
| 179 | |
| 180 | |
| 181 | class IdentityInputNorm(nn.Module): |
| 182 | """Stand-in used when the processor requires no rescale/normalise. |
| 183 | |
| 184 | Not a no-op: with ``mm_device_do_normalize`` enabled, raw ``uint8`` |
| 185 | pixels travel to the device unprocessed and must be cast to |
| 186 | ``visual_dtype`` here. |
| 187 | """ |
| 188 | |
| 189 | @property |
| 190 | def input_dtype(self) -> torch.dtype | None: |
| 191 | return None |
| 192 | |
| 193 | def forward( |
| 194 | self, pixel_values: torch.Tensor, visual_dtype: torch.dtype |
| 195 | ) -> torch.Tensor: |
| 196 | return pixel_values.to(visual_dtype, copy=False) |
| 197 | |
| 198 | |
| 199 | @CustomOp.register("fused_mm_input_norm") |
| 200 | class FusedMMInputNorm(CustomOp): |
| 201 | """Module that applies rescaling and normalisation to input images. |
| 202 | Equivalent to: output = (input * rescale_factor - mean) / std |
| 203 | |
| 204 | Dtype semantics: |
| 205 | |
| 206 | * Input dtype — the dtype of the ``pixel_values`` argument to ``forward_*``. |
| 207 | It is ``uint8`` when ``mm_device_do_normalize`` is enabled (raw bytes |
| 208 | travel to the device unprocessed) and equals ``visual_dtype`` |
| 209 | otherwise. |
| 210 | * Output dtype — the ``visual_dtype`` argument of ``forward_*``. The |
| 211 | computation itself is always fp32, independent of the output dtype |
| 212 | (e.g. compute fp32, emit bf16). |
| 213 | |
| 214 | Platform dispatch: |
| 215 | |
| 216 | * ``forward_native`` — pure PyTorch eager path, used as the semantic |
| 217 | reference and default fallback on all platforms. |
| 218 | * ``forward_cuda`` — Triton kernel path for CUDA devices. |
| 219 | * ``forward_xpu`` — custom XPU kernel path. |
| 220 | * ``forward_oot`` — out-of-tree platform override entry point; falls back |
| 221 | to ``forward_native`` unless a plugin overrides it. |
| 222 | """ |
| 223 | |
| 224 | def __init__( |
| 225 | self, |
| 226 | image_mean: list[float], |
| 227 | image_std: list[float], |
| 228 | rescale_factor: float, |
| 229 | channel: int = 3, |
| 230 | ): |
| 231 | super().__init__() |
| 232 | |
| 233 | assert len(image_mean) == len(image_std) == channel, ( |
| 234 | f"image_mean/image_std must have {channel} entries, " |
| 235 | f"got {len(image_mean)} / {len(image_std)}" |
| 236 | ) |
| 237 | assert rescale_factor != 0.0, "rescale_factor must be non-zero" |
| 238 | |
| 239 | self.channel = channel |
| 240 | |
| 241 | # Model construction can set the accelerator as PyTorch's default |
| 242 | # device; build the buffers on CPU first, then move them over. |
| 243 | mean = torch.tensor(image_mean, dtype=torch.float32, device="cpu") |
| 244 | std = torch.tensor(image_std, dtype=torch.float32, device="cpu") |
| 245 | device = torch.get_default_device() |
| 246 | self.register_buffer("weight", (rescale_factor / std).to(device)) |
| 247 | self.register_buffer("bias", (-mean / std).to(device)) |
| 248 | |
| 249 | @classmethod |
| 250 | def enabled(cls) -> bool: |
| 251 | return True |
| 252 | |
| 253 | @property |
| 254 | def input_dtype(self) -> torch.dtype | None: |
| 255 | return torch.uint8 |
| 256 | |
| 257 | # ------------------------------------------------------------------ |
| 258 | # Internal helpers shared by the platform-specific forward_* methods |
| 259 | # ------------------------------------------------------------------ |
| 260 | |
| 261 | @staticmethod |
| 262 | def _unpack_2d(pixel_values: torch.Tensor) -> tuple[int, int]: |
| 263 | assert pixel_values.ndim == 2, ( |
| 264 | f"pixel_values must be 2D (patches, size), got {pixel_values.dim()}D " |
| 265 | f"with shape {tuple(pixel_values.shape)}" |
| 266 | ) |
| 267 | patches, size = pixel_values.shape |
| 268 | return patches, size |
| 269 | |
| 270 | def _patch_size(self, size: int) -> int: |
| 271 | assert size % self.channel == 0, ( |
| 272 | f"size={size} is not divisible by channel={self.channel}" |
| 273 | ) |
| 274 | return size // self.channel |
| 275 | |
| 276 | # ------------------------------------------------------------------ |
| 277 | # Platform-specific implementations |
| 278 | # ------------------------------------------------------------------ |
| 279 | |
| 280 | def forward_native( |
| 281 | self, pixel_values: torch.Tensor, visual_dtype: torch.dtype |
| 282 | ) -> torch.Tensor: |
| 283 | """Pure PyTorch eager implementation. |
| 284 | |
| 285 | This is the semantic reference implementation and the fallback used |
| 286 | on any platform without a specialised kernel. |
| 287 | """ |
| 288 | patches, size = self._unpack_2d(pixel_values) |
| 289 | patch_size = self._patch_size(size) |
| 290 | |
| 291 | # weight/bias are fp32, so type promotion makes the arithmetic fp32 |
| 292 | x = pixel_values.reshape(patches, self.channel, patch_size) |
| 293 | x = x * self.weight.view(1, self.channel, 1) + self.bias.view( |
| 294 | 1, self.channel, 1 |
| 295 | ) |
| 296 | return x.view(patches, size).to(visual_dtype) |
| 297 | |
| 298 | def forward_cuda( |
| 299 | self, pixel_values: torch.Tensor, visual_dtype: torch.dtype |
| 300 | ) -> torch.Tensor: |
| 301 | """Triton kernel path for CUDA devices.""" |
| 302 | patches, size = self._unpack_2d(pixel_values) |
| 303 | patch_size = self._patch_size(size) |
| 304 | |
| 305 | x3 = pixel_values.reshape(patches, self.channel, patch_size) |
| 306 | y = torch.empty((patches, size), dtype=visual_dtype, device=pixel_values.device) |
| 307 | y3 = y.view(patches, self.channel, patch_size) |
| 308 | |
| 309 | fused_mm_input_norm_triton(x3, y3, self.weight, self.bias) |
| 310 | return y |
| 311 | |
| 312 | def forward_xpu( |
| 313 | self, pixel_values: torch.Tensor, visual_dtype: torch.dtype |
| 314 | ) -> torch.Tensor: |
| 315 | """XPU fused custom kernel path. |
| 316 | |
| 317 | On XPU, fuse the whole rescale + normalise into a single custom |
| 318 | kernel. The eager path materializes an fp32 intermediate and then |
| 319 | casts back, which adds device-side compute that cancels the |
| 320 | bandwidth saving of transferring uint8 pixel_values. The fused |
| 321 | kernel reads uint8 directly and writes ``visual_dtype`` in one pass. |
| 322 | """ |
| 323 | # The out-of-tree XPU kernel only supports the uint8 input that |
| 324 | # device-side normalisation guarantees; fail loudly otherwise. |
| 325 | assert pixel_values.dtype == torch.uint8, ( |
| 326 | f"xpu_fused_input_norm requires uint8 input, got {pixel_values.dtype}" |
| 327 | ) |
| 328 | return torch.ops.vllm.xpu_fused_input_norm( |
| 329 | pixel_values, self.weight, self.bias, visual_dtype |
| 330 | ) |
| 331 | |
| 332 | def forward_oot( |
| 333 | self, pixel_values: torch.Tensor, visual_dtype: torch.dtype |
| 334 | ) -> torch.Tensor: |
| 335 | """Out-of-tree platform override entrypoint.""" |
| 336 | return self.forward_native(pixel_values, visual_dtype) |
| 337 | |
| 338 | |
| 339 | def build_mm_input_norm(model_config: ModelConfig) -> nn.Module: |
| 340 | """Build the input normalisation module for a model. |
| 341 | |
| 342 | Returns an ``IdentityInputNorm`` when device-side normalisation is |
| 343 | disabled or the processor's rescale/normalise is numerically the identity |
| 344 | transform; otherwise a ``FusedMMInputNorm`` built from the processor's |
| 345 | parameters. |
| 346 | """ |
| 347 | mm_config = getattr(model_config, "multimodal_config", None) |
| 348 | if not getattr(mm_config, "mm_device_do_normalize", False): |
| 349 | return IdentityInputNorm() |
| 350 | |
| 351 | params = _load_norm_params(model_config) |
| 352 | |
| 353 | # Flags off are already defaulted to identity values by |
| 354 | # ``_load_norm_params``; this also catches configs that request |
| 355 | # rescale/normalise but are numerically the identity transform. |
| 356 | if params.is_identity: |
| 357 | return IdentityInputNorm() |
| 358 | |
| 359 | return FusedMMInputNorm( |
| 360 | image_mean=params.image_mean, |
| 361 | image_std=params.image_std, |
| 362 | rescale_factor=params.rescale_factor, |
| 363 | channel=len(params.image_mean), |
| 364 | ) |