@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 8a269919f4

[MM] Add Triton kernel for mm_input_normal. (#56798)

Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Signed-off-by: wang.yuqi <noooop@126.com> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Co-authored-by: Isotr0py <2037008807@qq.com> Co-authored-by: Isotr0py <Isotr0py@outlook.com>

wang.yuqi committed Sep 25, 2026 at 12:31 UTC 8a269919f4b96c2d84aaa70da6e125671ca95263

9 files changed +749 -304

docs/design/mm_processing.md

+23 -17

@@ -42,34 +42,39 @@ When new data is passed in, we first check which items are in the cache, and whi

43 ### Fused Normalisation on the Device

45 -To accelerate the multi‑modal data pipeline (decoding, resizing, normalisation, and rescaling), we offload the heavy numerical preprocessing from the CPU to the GPU and optimise data movement.

45 +To accelerate the multi‑modal data pipeline (decoding, resizing, normalisation, and rescaling), we offload the **normalisation and rescaling** steps from the CPU to the device and optimises data movement.

47 #### Fusing Normalisation and Rescaling on the GPU

49 -Traditionally, the CPU would divide pixel values by 255, then subtract the mean and divide by the standard deviation. We fuse these steps into one operation and run it entirely on the GPU.

49 +Traditionally, the CPU would divide pixel values by 255, then subtracts the mean and divides by the standard deviation. We fuse these steps into a single per-channel affine transform on the device.

51 -**How it works**: We use a dedicated `FusedInputNorm` module that bakes the rescale factor (1/255) directly into the layer's `weight` and `bias` parameters. Instead of performing three separate steps (scale, subtract, divide), the module does everything in a single affine transformation: `y = x * weight + bias`.

51 +`FusedMMInputNorm` dispatches to `fused_mm_input_norm_triton` on CUDA, which applies:

53 -The parameters are set as follows:

53 + y = x * weight[c] + bias[c]

54 + = (x * rescale_factor - image_mean[c]) / image_std[c]

55 -- `weight` controls both the standard deviation and the rescale factor

56 -- `bias` centers the data using the mean and the same rescale factor

56 +with:

58 -This means the module takes raw pixel values (0–255) and outputs properly normalised values without ever explicitly dividing by 255 as a separate step.

58 + weight[c] = rescale_factor / image_std[c]

59 + bias[c] = -image_mean[c] / image_std[c]

60 -#### Optimized Data Path for Fused Normalisation

61 +The kernel accepts `uint8`, `float16`, `bfloat16` and `float32` inputs. The `uint8` path is the primary fast path: raw bytes travel to the device unprocessed, and the rescale factor is folded into `weight` — no separate divide-by-255 step. Compute is always done in fp32 inside the kernel.

62 -Performing fused normalisation directly on the device allows us to keep the entire transfer path—from **Entrypoint** through **Engine Core** to **GPU memory**—in **`uint8`**. This halves PCIe bandwidth and reduces CPU memory footprint.

63 +#### Optimized Data Path

64 -Only after data reaches GPU memory do we cast to `fp32` for the `FusedInputNorm` layer (to ensure numerical accuracy), then cast to `bf16` for subsequent layers—all within the GPU, avoiding any host‑side conversions.

65 + Before: CPU decode → CPU resize → CPU rescale (÷255) → CPU normalize

66 + → cast to bf16 → PCIe (2 B/elem) → GPU

66 -Overall path: **`Entrypoint (uint8) → Engine Core (uint8) → GPU Memory (uint8)`** → GPU‑local `fp32` `FusedInputNorm` → `bf16` output.

68 + After: CPU decode → CPU resize → uint8 pixel_values

69 + → PCIe (1 B/elem) → GPU fused affine (fp32 compute) → bf16

70 +

71 +The transfer path **Entrypoint (API server / offline LLM) → Engine Core(scheduler + workers) → Device Memory** stays in `uint8`. On the device,`fused_mm_input_norm_triton` computes in fp32 internally and writes the requested `visual_dtype` (a per-call argument to `forward_*`, commonly `bf16`) directly, without materialising a global fp32 intermediate.

73 #### Toggle: `mm_device_do_normalize`

70 -This GPU‑side fusion is controlled by a config flag called **`mm_device_do_normalize`**.

75 +GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:

72 -- When `True`, normalisation and rescaling are done on the GPU using the `FusedInputNorm` layer; when `False`, we fall back to the old CPU‑side path.

77 +- When `True`, normalisation and rescaling are done on the GPU using the `FusedMMInputNorm` layer; when `False`, we fall back to the old CPU‑side path.

78 - The flag is **enabled by default** for all models that support it.

79 - Currently, it’s on by default for these architectures:

@@ -78,8 +83,9 @@ This GPU‑side fusion is controlled by a config flag called **`mm_device_do_nor

83 | `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |

84 | `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |

81 -#### What We Gain Overall

86 +#### Key Properties and Gains

83 -- **CPU offload**: The arithmetic for normalisation and rescaling is completely gone from the CPU.

84 -- **PCIe savings**: Sending `uint8` (1 byte) instead of `bf16` (2 bytes) slashes data transfer volume by **50%** .

85 -- **GPU overhead**: The fused kernel is very lightweight and can often be merged with subsequent CUDA operations, so it hardly adds any extra cost.

88 +- **CPU offload & 50% PCIe savings** — normalisation/rescaling leaves the CPU entirely; sending `uint8` (1 byte) instead of `bf16` (2 bytes) halves transfer volume.

89 +- **Single-pass fusion** — `y = x * weight[c] + bias[c]` in one kernel launch via `fused_mm_input_norm_triton`; input is read in its native dtype (`uint8`).

90 +- **float32 compute for free** — arithmetic runs in fp32 inside the kernel regardless of I/O dtypes. Accuracy matches fp32, with no extra bandwidth: intermediates stay in registers, and no global fp32 tensor is materialised.

91 +- **Preallocated output with batch-dim padding** — callers can pass a larger buffer; only the leading `N` rows are written, so buffers are reusable across calls.

tests/model_executor/layers/test_mm_input_norm.py new

+320

@@ -0,0 +1,320 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +"""Tests for FusedMMInputNorm and the fused affine transform kernel."""

5 +

6 +import pytest

7 +import torch

8 +

9 +from vllm.model_executor.layers.fusion.mm_input_norm import (

10 + FusedMMInputNorm,

11 + IdentityInputNorm,

12 + NormParams,

13 + fused_mm_input_norm_triton,

14 +)

15 +from vllm.platforms import current_platform

16 +from vllm.triton_utils import HAS_TRITON

17 +from vllm.utils.torch_utils import set_random_seed

18 +

19 +# Module-level tests run on any platform: the CPU platform exercises

20 +# forward_native, accelerators exercise the fused kernel. Only the raw

21 +# kernel tests require an accelerator.

22 +_DEVICE = torch.device(current_platform.device_type)

23 +

24 +requires_accelerator = pytest.mark.skipif(

25 + current_platform.is_cpu(),

26 + reason="fused Triton kernel requires a CUDA/XPU accelerator",

27 +)

28 +requires_triton = pytest.mark.skipif(not HAS_TRITON, reason="requires Triton")

29 +requires_vllm_config = pytest.mark.usefixtures("default_vllm_config")

30 +

31 +

32 +def _reference_input_norm(

33 + pixel_values: torch.Tensor,

34 + image_mean: list[float],

35 + image_std: list[float],

36 + rescale_factor: float,

37 + channel: int,

38 + out_dtype: torch.dtype,

39 +) -> torch.Tensor:

40 + """Straightforward per-channel affine: (x * rescale - mean) / std."""

41 + patches, size = pixel_values.shape

42 + patch_size = size // channel

43 + mean = torch.tensor(

44 + image_mean, dtype=torch.float32, device=pixel_values.device

45 + ).view(1, channel, 1)

46 + std = torch.tensor(image_std, dtype=torch.float32, device=pixel_values.device).view(

47 + 1, channel, 1

48 + )

49 + x = pixel_values.to(torch.float32).reshape(patches, channel, patch_size)

50 + x = (x * rescale_factor - mean) / std

51 + return x.view(patches, size).to(out_dtype)

52 +

53 +

54 +# CLIP-style RGB normalization constants reused across tests.

55 +_RGB_MEAN = [0.48145466, 0.4578275, 0.40821073]

56 +_RGB_STD = [0.26862954, 0.26130258, 0.27577711]

57 +_RGB_RESCALE = 1.0 / 255.0

58 +

59 +

60 +def _check_against_reference(

61 + pixel_values: torch.Tensor,

62 + out_dtype: torch.dtype,

63 + channel: int = 3,

64 + image_mean: list[float] | None = None,

65 + image_std: list[float] | None = None,

66 + rescale_factor: float = _RGB_RESCALE,

67 + **close_kwargs,

68 +):

69 + """Run ``FusedMMInputNorm`` on ``pixel_values`` and compare against the

70 + eager per-channel affine reference. Uses the RGB constants by default."""

71 + if image_mean is None:

72 + image_mean = list(_RGB_MEAN)

73 + if image_std is None:

74 + image_std = list(_RGB_STD)

75 + norm = FusedMMInputNorm(

76 + image_mean=image_mean,

77 + image_std=image_std,

78 + rescale_factor=rescale_factor,

79 + channel=channel,

80 + ).to(_DEVICE)

81 + expected = _reference_input_norm(

82 + pixel_values, image_mean, image_std, rescale_factor, channel, out_dtype

83 + )

84 + torch.testing.assert_close(

85 + norm(pixel_values, visual_dtype=out_dtype), expected, **close_kwargs

86 + )

87 +

88 +

89 +# ===========================================================================

90 +# Module-level behavior

91 +# ===========================================================================

92 +@requires_vllm_config

93 +class TestFusedMMInputNormModule:

94 + def test_matches_reference_above_cudnn_grid_limit(self):

95 + """num_patches above the old cuDNN batch-norm grid limit (~65535)

96 + used to raise CUDNN_STATUS_INTERNAL_ERROR."""

97 + set_random_seed(0)

98 + pixel_values = torch.randint(

99 + 0, 256, (70000, 3 * 196), dtype=torch.float32, device=_DEVICE

100 + )

101 + _check_against_reference(pixel_values, torch.float32)

102 +

103 + def test_identity_passthrough(self):

104 + norm = IdentityInputNorm()

105 +

106 + pixel_values = torch.randn(8, 3 * 196, dtype=torch.float32, device=_DEVICE)

107 + out = norm(pixel_values, visual_dtype=torch.bfloat16)

108 + torch.testing.assert_close(out, pixel_values.to(torch.bfloat16))

109 +

110 +

111 +# ===========================================================================

112 +# dtype coverage

113 +# ===========================================================================

114 +@requires_vllm_config

115 +class TestFusedMMInputNormDtypes:

116 + @pytest.mark.parametrize(

117 + "in_dtype,out_dtype",

118 + [

119 + (torch.float32, torch.float32),

120 + (torch.float32, torch.bfloat16),

121 + (torch.float32, torch.float16),

122 + (torch.bfloat16, torch.bfloat16),

123 + (torch.bfloat16, torch.float32),

124 + (torch.float16, torch.float16),

125 + (torch.uint8, torch.bfloat16),

126 + (torch.uint8, torch.float32),

127 + ],

128 + )

129 + def test_dtype_combinations(self, in_dtype: torch.dtype, out_dtype: torch.dtype):

130 + set_random_seed(0)

131 + if in_dtype == torch.uint8:

132 + pixel_values = torch.randint(

133 + 0, 256, (32, 3 * 16), dtype=torch.uint8, device=_DEVICE

134 + )

135 + else:

136 + pixel_values = torch.rand(

137 + (32, 3 * 16), dtype=torch.float32, device=_DEVICE

138 + ).to(in_dtype)

139 +

140 + _check_against_reference(

141 + pixel_values,

142 + out_dtype,

143 + image_mean=[0.5, 0.5, 0.5],

144 + image_std=[0.5, 0.5, 0.5],

145 + atol=1e-2,

146 + rtol=1e-2,

147 + )

148 +

149 +

150 +# ===========================================================================

151 +# Shape / channel coverage

152 +# ===========================================================================

153 +@requires_vllm_config

154 +class TestFusedMMInputNormShapes:

155 + @pytest.mark.parametrize("channel", [1, 3, 4])

156 + def test_channel_variants(self, channel: int):

157 + set_random_seed(0)

158 + pixel_values = torch.randint(

159 + 0, 256, (16, channel * 64), dtype=torch.float32, device=_DEVICE

160 + )

161 + _check_against_reference(

162 + pixel_values,

163 + torch.float32,

164 + channel=channel,

165 + image_mean=[0.5] * channel,

166 + image_std=[0.25] * channel,

167 + )

168 +

169 + @pytest.mark.parametrize("patch_size", [1, 2048])

170 + def test_block_boundaries(self, patch_size: int):

171 + """Masked-tail (12 elements) and exact-multiple (24 blocks) paths."""

172 + set_random_seed(0)

173 + pixel_values = torch.randint(

174 + 0, 256, (4, 3 * patch_size), dtype=torch.float32, device=_DEVICE

175 + )

176 + _check_against_reference(pixel_values, torch.float32)

177 +

178 +

179 +# ===========================================================================

180 +# Input handling: non-contiguous inputs

181 +# ===========================================================================

182 +@requires_vllm_config

183 +class TestFusedMMInputNormInputHandling:

184 + def test_non_contiguous_input_matches_reference(self):

185 + set_random_seed(0)

186 + base = torch.randint(

187 + 0, 256, (8, 3 * 32, 2), dtype=torch.float32, device=_DEVICE

188 + )

189 + non_contig = base[..., 0]

190 + assert not non_contig.is_contiguous()

191 +

192 + _check_against_reference(non_contig, torch.float32)

193 +

194 +

195 +# ===========================================================================

196 +# Raw kernel entry point

197 +# ===========================================================================

198 +@requires_accelerator

199 +@requires_triton

200 +class TestFusedMMInputNormKernel:

201 + @pytest.mark.parametrize(

202 + "N, C, L",

203 + [

204 + (1, 1, 1),

205 + (1, 3, 1),

206 + (3, 3, 7), # tile straddles channel boundaries

207 + (2, 8, 63),

208 + (4, 3, 1000),

209 + ],

210 + )

211 + def test_channel_boundary_crossing(self, N: int, C: int, L: int):

212 + """The 1D kernel recovers ``c = (offs // L) % C`` per lane, so a tile

213 + can straddle a channel boundary and gather weight/bias per lane.

214 + These shapes force that path by keeping ``L`` well below the block."""

215 + set_random_seed(0)

216 + x = torch.randn(N, C, L, dtype=torch.float32, device=_DEVICE)

217 + w = torch.randn(C, dtype=torch.float32, device=_DEVICE)

218 + b = torch.randn(C, dtype=torch.float32, device=_DEVICE)

219 +

220 + out = torch.empty_like(x)

221 + fused_mm_input_norm_triton(x, out, w, b)

222 +

223 + expected = x * w.view(1, C, 1) + b.view(1, C, 1)

224 + torch.testing.assert_close(out, expected)

225 +

226 + def test_non_contiguous_inputs_materialized(self):

227 + """Read-only tensors are made contiguous internally."""

228 + N, C, L = 3, 3, 100

229 + set_random_seed(0)

230 + base = torch.randn(N, C, L, 2, dtype=torch.float32, device=_DEVICE)

231 + x = base[..., 0]

232 + assert not x.is_contiguous()

233 + w = torch.randn(C, dtype=torch.float32, device=_DEVICE)

234 + b = torch.randn(C, dtype=torch.float32, device=_DEVICE)

235 +

236 + out = torch.empty(N, C, L, dtype=torch.float32, device=_DEVICE)

237 + fused_mm_input_norm_triton(x, out, w, b)

238 +

239 + expected = x * w.view(1, C, 1) + b.view(1, C, 1)

240 + torch.testing.assert_close(out, expected)

241 +

242 + def test_rejects_mismatched_output_shape(self):

243 + """``outputs`` must be shaped exactly like ``inputs``."""

244 + N, C, L = 3, 3, 100

245 + x = torch.randn(N, C, L, dtype=torch.float32, device=_DEVICE)

246 + w = torch.randn(C, dtype=torch.float32, device=_DEVICE)

247 + b = torch.randn(C, dtype=torch.float32, device=_DEVICE)

248 +

249 + for out_shape in ((N, C + 1, L), (N, C, L + 1), (N + 2, C, L)):

250 + with pytest.raises(AssertionError):

251 + fused_mm_input_norm_triton(

252 + x,

253 + torch.empty(*out_shape, dtype=torch.float32, device=_DEVICE),

254 + w,

255 + b,

256 + )

257 +

258 +

259 +# ===========================================================================

260 +# Construction / configuration

261 +# ===========================================================================

262 +@requires_vllm_config

263 +class TestFusedMMInputNormConstruction:

264 + """Weight/bias buffer semantics at init time."""

265 +

266 + def test_identity_config_buffers(self):

267 + """Numerically identity parameters still build plain weight/bias

268 + buffers; the identity shortcut lives in ``from_model_config``."""

269 + norm = FusedMMInputNorm(

270 + image_mean=[0.0, 0.0, 0.0],

271 + image_std=[1.0, 1.0, 1.0],

272 + rescale_factor=1.0,

273 + )

274 + torch.testing.assert_close(norm.weight, torch.ones(3))

275 + torch.testing.assert_close(norm.bias, torch.zeros(3))

276 +

277 + def test_non_identity_buffers(self):

278 + channel = 3

279 + norm = FusedMMInputNorm(

280 + image_mean=_RGB_MEAN,

281 + image_std=_RGB_STD,

282 + rescale_factor=_RGB_RESCALE,

283 + channel=channel,

284 + )

285 + assert norm.weight.shape == (channel,)

286 + assert norm.bias.shape == (channel,)

287 +

288 + mean = torch.tensor(_RGB_MEAN, dtype=torch.float32)

289 + std = torch.tensor(_RGB_STD, dtype=torch.float32)

290 + torch.testing.assert_close(norm.weight, _RGB_RESCALE / std)

291 + torch.testing.assert_close(norm.bias, -mean / std)

292 +

293 + def test_is_identity_params(self):

294 + """The numeric-identity check mirrors the old allclose detection."""

295 + assert NormParams([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0).is_identity

296 + assert not NormParams(_RGB_MEAN, _RGB_STD, _RGB_RESCALE).is_identity

297 +

298 + @pytest.mark.parametrize(

299 + ("image_mean", "image_std", "rescale_factor"),

300 + [

301 + ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0),

302 + ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25], 1 / 255),

303 + ],

304 + )

305 + def test_buffers_initialized_on_default_device(

306 + self,

307 + image_mean: list[float],

308 + image_std: list[float],

309 + rescale_factor: float,

310 + ):

311 + """Buffers are built on CPU first, then moved to the caller's

312 + default device, so init never touches the accelerator."""

313 + # The meta device gives the CPU-only test shard the same non-CPU

314 + # default-device semantics without requiring a CUDA build.

315 + default_device = "cuda" if torch.cuda.is_available() else "meta"

316 + with torch.device(default_device):

317 + norm = FusedMMInputNorm(image_mean, image_std, rescale_factor)

318 +

319 + assert norm.weight.device.type == default_device

320 + assert norm.bias.device.type == default_device

tests/models/multimodal/processing/test_qwen2_vl.py

+20 -48

@@ -8,10 +8,11 @@ import pytest

8 import torch

9 from PIL import Image

11 -from vllm.model_executor.models.vision import FusedInputNorm

11 +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm

12 from vllm.multimodal import MULTIMODAL_REGISTRY

13 from vllm.multimodal.cache import MultiModalProcessorOnlyCache

14 from vllm.multimodal.inputs import batched_tensors_equal

15 +from vllm.platforms import current_platform

17 from ....conftest import ImageTestAssets

18 from ...utils import build_model_context

@@ -65,45 +66,6 @@ def test_jina_vl_processing_order() -> None:

66 ]

68 -@pytest.mark.parametrize(

69 - ("image_mean", "image_std", "rescale_factor", "is_identity"),

70 - [

71 - ([0.0, 0.0, 0.0], [1.0, 1.0, 1.0], 1.0, True),

72 - ([0.5, 0.5, 0.5], [0.25, 0.25, 0.25], 1 / 255, False),

73 - ],

74 -)

75 -def test_fused_input_norm_initialization_on_device(

76 - monkeypatch: pytest.MonkeyPatch,

77 - image_mean: list[float],

78 - image_std: list[float],

79 - rescale_factor: float,

80 - is_identity: bool,

81 -):

82 - """Identity detection must not synchronize the default device."""

83 - original_allclose = torch.allclose

84 -

85 - def cpu_allclose(input: torch.Tensor, other: torch.Tensor, *args, **kwargs):

86 - assert input.device.type == "cpu"

87 - assert other.device.type == "cpu"

88 - return original_allclose(input, other, *args, **kwargs)

89 -

90 - monkeypatch.setattr(torch, "allclose", cpu_allclose)

91 - # Exercise the real accelerator when available. The meta device gives the

92 - # CPU-only test shard the same non-CPU default-device semantics without

93 - # requiring a CUDA-enabled PyTorch build.

94 - default_device = "cuda" if torch.cuda.is_available() else "meta"

95 - with torch.device(default_device):

96 - input_norm = FusedInputNorm(image_mean, image_std, rescale_factor)

97 -

98 - assert input_norm.is_identity is is_identity

99 - if is_identity:

100 - assert input_norm.weight is None

101 - assert input_norm.bias is None

102 - else:

103 - assert input_norm.weight.device.type == default_device

104 - assert input_norm.bias.device.type == default_device

105 -

106 -

69 @pytest.mark.parametrize("model_id", ["Qwen/Qwen2-VL-2B-Instruct"])

70 @pytest.mark.parametrize(

71 ("mm_processor_kwargs", "expected_toks_per_img", "expected_pixels_shape"),

@@ -240,16 +202,20 @@ def test_qwen2_5_vl_video_fps_to_second_per_grid_ts(model_id: str) -> None:

202 assert float(spg) == pytest.approx(temporal_patch_size / fps, rel=2e-2)

203

204

205 +@pytest.mark.usefixtures("default_vllm_config")

206 @pytest.mark.parametrize(

207 "model_id", ["Qwen/Qwen2-VL-2B-Instruct", "Qwen/Qwen2.5-VL-3B-Instruct"]

208 )

209 @pytest.mark.parametrize("num_imgs", [1, 2])

210 def test_mm_device_do_normalize(

248 - image_assets: ImageTestAssets,

249 - model_id: str,

250 - num_imgs: int,

251 -):

252 - """Ensure that enable mm_device_do_normalize yields the correct result."""

211 + image_assets: ImageTestAssets, model_id: str, num_imgs: int

212 +) -> None:

213 + """Device-side normalisation must reproduce the on-CPU processor result.

214 +

215 + Runs on any platform: the CPU platform exercises the ``forward_native``

216 + fallback, accelerators exercise the fused kernel.

217 + """

218 + device = current_platform.device_type

219 ctx = build_model_context(

220 model_id,

221 limit_mm_per_prompt={"image": num_imgs},

@@ -280,9 +246,15 @@ def test_mm_device_do_normalize(

246 ].get_data()["pixel_values"]

247

248 ctx.model_config.multimodal_config.mm_device_do_normalize = True

283 - input_norm = FusedInputNorm.from_model_config(ctx.model_config)

249 + input_norm = build_mm_input_norm(ctx.model_config).to(device)

250 +

251 + # With normalisation disabled, the processor emits raw uint8 pixels,

252 + # matching the production mm_device_do_normalize path.

253 + assert pixel_values_without_normalize.dtype == torch.uint8

254 pixel_values_do_input_norm = input_norm(

285 - pixel_values_without_normalize.to(dtype), dtype

255 + pixel_values_without_normalize.to(device), dtype

256 )

257

288 - torch.testing.assert_close(pixel_values_with_normalize, pixel_values_do_input_norm)

258 + torch.testing.assert_close(

259 + pixel_values_with_normalize.to(device), pixel_values_do_input_norm

260 + )

tests/models/test_vision.py

-60

@@ -13,7 +13,6 @@ from vllm.distributed.parallel_state import (

13 initialize_model_parallel,

14 )

15 from vllm.model_executor.models.vision import (

16 - FusedInputNorm,

16 get_load_balance_assignment,

17 resolve_visual_encoder_outputs,

18 run_dp_sharded_mrope_vision_model,

@@ -488,62 +487,3 @@ def test_simple_mrope_vision_model_spatial_merge(spatial_merge_size: int):

487

488 assert output.shape[0] == expected_output_patches

489 assert output.shape[1] == vision_model.out_hidden_size

491 -

492 -

493 -def _reference_input_norm(

494 - pixel_values: torch.Tensor,

495 - image_mean: list[float],

496 - image_std: list[float],

497 - rescale_factor: float,

498 - channel: int,

499 - out_dtype: torch.dtype,

500 -) -> torch.Tensor:

501 - """Straightforward per-channel affine: (x * rescale - mean) / std."""

502 - patches, size = pixel_values.shape

503 - patch_size = size // channel

504 - mean = torch.tensor(image_mean, dtype=torch.float32).view(1, channel, 1)

505 - std = torch.tensor(image_std, dtype=torch.float32).view(1, channel, 1)

506 - x = pixel_values.to(torch.float32).view(patches, channel, patch_size)

507 - x = (x * rescale_factor - mean) / std

508 - return x.view(patches, size).to(out_dtype)

509 -

510 -

511 -@pytest.mark.parametrize("num_patches", [1, 37, 70000])

512 -def test_fused_input_norm_matches_reference(num_patches: int):

513 - """FusedInputNorm must equal the plain affine, including for num_patches

514 - above the cuDNN batch-norm grid limit (~65535) that previously raised

515 - CUDNN_STATUS_INTERNAL_ERROR (issue #51717)."""

516 - channel = 3

517 - patch_size = 14 * 14

518 - image_mean = [0.48145466, 0.4578275, 0.40821073]

519 - image_std = [0.26862954, 0.26130258, 0.27577711]

520 - rescale_factor = 1.0 / 255.0

521 -

522 - set_random_seed(0)

523 - pixel_values = torch.randint(

524 - 0, 256, (num_patches, channel * patch_size), dtype=torch.float32

525 - )

526 -

527 - norm = FusedInputNorm(

528 - image_mean=image_mean,

529 - image_std=image_std,

530 - rescale_factor=rescale_factor,

531 - channel=channel,

532 - )

533 - assert not norm.is_identity

534 -

535 - out = norm(pixel_values, visual_dtype=torch.float32)

536 - expected = _reference_input_norm(

537 - pixel_values, image_mean, image_std, rescale_factor, channel, torch.float32

538 - )

539 - torch.testing.assert_close(out, expected)

540 -

541 -

542 -def test_fused_input_norm_identity_passthrough():

543 - """The identity configuration returns the input unchanged (cast only)."""

544 - norm = FusedInputNorm.identity()

545 - assert norm.is_identity

546 -

547 - pixel_values = torch.randn(8, 3 * 196, dtype=torch.float32)

548 - out = norm(pixel_values, visual_dtype=torch.bfloat16)

549 - torch.testing.assert_close(out, pixel_values.to(torch.bfloat16))

vllm/config/multimodal.py

+12 -4

@@ -534,11 +534,19 @@ class MultiModalConfig:

534 """Get the keyword arguments to pass to the multi-modal processor

535 according to the extra arguments passed during inference.

536 """

537 - kwargs = self.mm_processor_kwargs or {}

537 + merged = dict(self.mm_processor_kwargs or {}) | dict(inference_kwargs)

538 if self.mm_device_do_normalize:

539 - kwargs["do_normalize"] = False

540 - kwargs["do_rescale"] = False

541 - return kwargs | dict(inference_kwargs)

539 + # The model expects raw pixels and normalises on device;

540 + # overriding these per request would silently double-normalise.

541 + for key in ("do_normalize", "do_rescale"):

542 + if key in inference_kwargs:

543 + logger.warning_once(

544 + "mm_device_do_normalize is enabled; ignoring "

545 + "per-request mm_processor_kwargs[%r].",

546 + key,

547 + )

548 + merged[key] = False

549 + return merged

550

551 def use_gpu_video_backend(self) -> bool:

552 """Return whether the configured video loader or codec uses the GPU."""

vllm/model_executor/layers/fusion/mm_input_norm.py new

+361

@@ -0,0 +1,361 @@

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 Any, 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 cached_get_processor, 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 +):

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, falling back to the image

148 + processor object."""

149 + model = model_config.model

150 + revision = model_config.revision

151 +

152 + config = get_processor_config(model, revision=revision)

153 + # NOTE: do not use cached_image_processor_from_config here — it merges

154 + # mm_processor_kwargs, which mm_device_do_normalize poisons with

155 + # do_normalize=False.

156 + image_processor = cached_get_processor(

157 + model, revision=revision, trust_remote_code=model_config.trust_remote_code

158 + ).image_processor

159 +

160 + def resolve(key: str) -> Any:

161 + """Processor config value, falling back to the image_processor."""

162 + if (value := config.get(key)) is not None:

163 + return value

164 + return getattr(image_processor, key, None)

165 +

166 + do_rescale = bool(resolve("do_rescale"))

167 + do_normalize = bool(resolve("do_normalize"))

168 +

169 + # Parameters whose flag is off are unused; default them to no-ops

170 + # without resolving them.

171 + rescale_factor = resolve("rescale_factor") if do_rescale else 1.0

172 + image_mean = resolve("image_mean") if do_normalize else [0.0] * 3

173 + image_std = resolve("image_std") if do_normalize else [1.0] * 3

174 +

175 + assert rescale_factor is not None, "rescale_factor is still None after resolution."

176 + assert image_mean is not None, "image_mean is still None after resolution."

177 + assert image_std is not None, "image_std is still None after resolution."

178 + assert len(image_mean) == len(image_std), (

179 + f"image_mean and image_std have different lengths: "

180 + f"{len(image_mean)} vs {len(image_std)}"

181 + )

182 +

183 + return NormParams(

184 + image_mean=[float(v) for v in image_mean],

185 + image_std=[float(v) for v in image_std],

186 + rescale_factor=float(rescale_factor),

187 + )

188 +

189 +

190 +class IdentityInputNorm(nn.Module):

191 + """Stand-in used when the processor requires no rescale/normalise.

192 +

193 + Not a no-op: with ``mm_device_do_normalize`` enabled, raw ``uint8``

194 + pixels travel to the device unprocessed and must be cast to

195 + ``visual_dtype`` here.

196 + """

197 +

198 + def forward(

199 + self, pixel_values: torch.Tensor, visual_dtype: torch.dtype

200 + ) -> torch.Tensor:

201 + return pixel_values.to(visual_dtype, copy=False)

202 +

203 +

204 +@CustomOp.register("fused_mm_input_norm")

205 +class FusedMMInputNorm(CustomOp):

206 + """Module that applies rescaling and normalisation to input images.

207 + Equivalent to: output = (input * rescale_factor - mean) / std

208 +

209 + Dtype semantics:

210 +

211 + * Input dtype — the dtype of the ``pixel_values`` argument to ``forward_*``.

212 + It is ``uint8`` when ``mm_device_do_normalize`` is enabled (raw bytes

213 + travel to the device unprocessed) and equals ``visual_dtype``

214 + otherwise.

215 + * Output dtype — the ``visual_dtype`` argument of ``forward_*``. The

216 + computation itself is always fp32, independent of the output dtype

217 + (e.g. compute fp32, emit bf16).

218 +

219 + Platform dispatch:

220 +

221 + * ``forward_native`` — pure PyTorch eager path, used as the semantic

222 + reference and default fallback on all platforms.

223 + * ``forward_cuda`` — Triton kernel path for CUDA devices.

224 + * ``forward_xpu`` — custom XPU kernel path.

225 + * ``forward_oot`` — out-of-tree platform override entry point; falls back

226 + to ``forward_native`` unless a plugin overrides it.

227 + """

228 +

229 + def __init__(

230 + self,

231 + image_mean: list[float],

232 + image_std: list[float],

233 + rescale_factor: float,

234 + channel: int = 3,

235 + ):

236 + super().__init__()

237 +

238 + assert len(image_mean) == len(image_std) == channel, (

239 + f"image_mean/image_std must have {channel} entries, "

240 + f"got {len(image_mean)} / {len(image_std)}"

241 + )

242 + assert rescale_factor != 0.0, "rescale_factor must be non-zero"

243 +

244 + self.channel = channel

245 +

246 + # Model construction can set the accelerator as PyTorch's default

247 + # device; build the buffers on CPU first, then move them over.

248 + mean = torch.tensor(image_mean, dtype=torch.float32, device="cpu")

249 + std = torch.tensor(image_std, dtype=torch.float32, device="cpu")

250 + device = torch.get_default_device()

251 + self.register_buffer("weight", (rescale_factor / std).to(device))

252 + self.register_buffer("bias", (-mean / std).to(device))

253 +

254 + # ------------------------------------------------------------------

255 + # Internal helpers shared by the platform-specific forward_* methods

256 + # ------------------------------------------------------------------

257 +

258 + @staticmethod

259 + def _unpack_2d(pixel_values: torch.Tensor) -> tuple[int, int]:

260 + assert pixel_values.ndim == 2, (

261 + f"pixel_values must be 2D (patches, size), got {pixel_values.dim()}D "

262 + f"with shape {tuple(pixel_values.shape)}"

263 + )

264 + patches, size = pixel_values.shape

265 + return patches, size

266 +

267 + def _patch_size(self, size: int) -> int:

268 + assert size % self.channel == 0, (

269 + f"size={size} is not divisible by channel={self.channel}"

270 + )

271 + return size // self.channel

272 +

273 + # ------------------------------------------------------------------

274 + # Platform-specific implementations

275 + # ------------------------------------------------------------------

276 +

277 + def forward_native(

278 + self, pixel_values: torch.Tensor, visual_dtype: torch.dtype

279 + ) -> torch.Tensor:

280 + """Pure PyTorch eager implementation.

281 +

282 + This is the semantic reference implementation and the fallback used

283 + on any platform without a specialised kernel.

284 + """

285 + patches, size = self._unpack_2d(pixel_values)

286 + patch_size = self._patch_size(size)

287 +

288 + # weight/bias are fp32, so type promotion makes the arithmetic fp32

289 + x = pixel_values.reshape(patches, self.channel, patch_size)

290 + x = x * self.weight.view(1, self.channel, 1) + self.bias.view(

291 + 1, self.channel, 1

292 + )

293 + return x.view(patches, size).to(visual_dtype)

294 +

295 + def forward_cuda(

296 + self, pixel_values: torch.Tensor, visual_dtype: torch.dtype

297 + ) -> torch.Tensor:

298 + """Triton kernel path for CUDA devices."""

299 + patches, size = self._unpack_2d(pixel_values)

300 + patch_size = self._patch_size(size)

301 +

302 + x3 = pixel_values.reshape(patches, self.channel, patch_size)

303 + y = torch.empty((patches, size), dtype=visual_dtype, device=pixel_values.device)

304 + y3 = y.view(patches, self.channel, patch_size)

305 +

306 + fused_mm_input_norm_triton(x3, y3, self.weight, self.bias)

307 + return y

308 +

309 + def forward_xpu(

310 + self, pixel_values: torch.Tensor, visual_dtype: torch.dtype

311 + ) -> torch.Tensor:

312 + """XPU fused custom kernel path.

313 +

314 + On XPU, fuse the whole rescale + normalise into a single custom

315 + kernel. The eager path materializes an fp32 intermediate and then

316 + casts back, which adds device-side compute that cancels the

317 + bandwidth saving of transferring uint8 pixel_values. The fused

318 + kernel reads uint8 directly and writes ``visual_dtype`` in one pass.

319 + """

320 + # The out-of-tree XPU kernel only supports the uint8 input that

321 + # device-side normalisation guarantees; fail loudly otherwise.

322 + assert pixel_values.dtype == torch.uint8, (

323 + f"xpu_fused_input_norm requires uint8 input, got {pixel_values.dtype}"

324 + )

325 + return torch.ops.vllm.xpu_fused_input_norm(

326 + pixel_values, self.weight, self.bias, visual_dtype

327 + )

328 +

329 + def forward_oot(

330 + self, pixel_values: torch.Tensor, visual_dtype: torch.dtype

331 + ) -> torch.Tensor:

332 + """Out-of-tree platform override entrypoint."""

333 + return self.forward_native(pixel_values, visual_dtype)

334 +

335 +

336 +def build_mm_input_norm(model_config: ModelConfig) -> nn.Module:

337 + """Build the input normalisation module for a model.

338 +

339 + Returns an ``IdentityInputNorm`` when device-side normalisation is

340 + disabled or the processor's rescale/normalise is numerically the identity

341 + transform; otherwise a ``FusedMMInputNorm`` built from the processor's

342 + parameters.

343 + """

344 + mm_config = getattr(model_config, "multimodal_config", None)

345 + if not getattr(mm_config, "mm_device_do_normalize", False):

346 + return IdentityInputNorm()

347 +

348 + params = _load_norm_params(model_config)

349 +

350 + # Flags off are already defaulted to identity values by

351 + # ``_load_norm_params``; this also catches configs that request

352 + # rescale/normalise but are numerically the identity transform.

353 + if params.is_identity:

354 + return IdentityInputNorm()

355 +

356 + return FusedMMInputNorm(

357 + image_mean=params.image_mean,

358 + image_std=params.image_std,

359 + rescale_factor=params.rescale_factor,

360 + channel=len(params.image_mean),

361 + )

vllm/model_executor/models/qwen2_5_vl.py

+6 -5

@@ -55,6 +55,10 @@ from vllm.logger import init_logger

55 from vllm.model_executor.layers.activation import get_act_and_mul_fn

56 from vllm.model_executor.layers.attention import MMEncoderAttention

57 from vllm.model_executor.layers.conv import Conv3dLayer

58 +from vllm.model_executor.layers.fusion.mm_input_norm import (

59 + IdentityInputNorm,

60 + build_mm_input_norm,

61 +)

62 from vllm.model_executor.layers.layernorm import RMSNorm

63 from vllm.model_executor.layers.linear import (

64 ColumnParallelLinear,

@@ -120,7 +124,6 @@ from .utils import (

124 maybe_prefix,

125 )

126 from .vision import (

123 - FusedInputNorm,

127 get_fp8_padded_hidden_size,

128 get_vit_attn_backend,

129 is_vit_use_data_parallel,

@@ -664,9 +667,7 @@ class Qwen2_5_VisionTransformer(nn.Module):

667 in_channels=in_channels,

668 hidden_size=self.hidden_size,

669 )

667 - self.input_norm = (

668 - input_norm if input_norm is not None else FusedInputNorm.identity()

669 - )

670 + self.input_norm = input_norm if input_norm is not None else IdentityInputNorm()

671

672 norm_layer = partial(RMSNorm, eps=norm_eps)

673 head_dim = self.hidden_size // self.num_heads

@@ -1425,7 +1426,7 @@ class Qwen2_5_VLForConditionalGeneration(

1426 vision_config=config.vision_config,

1427 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1428 quant_config=self.quant_config,

1428 - input_norm=FusedInputNorm.from_model_config(self.model_config),

1429 + input_norm=build_mm_input_norm(self.model_config),

1430 prefix=maybe_prefix(prefix, "visual"),

1431 )

1432

vllm/model_executor/models/qwen2_vl.py

+6 -5

@@ -51,6 +51,10 @@ from vllm.logger import init_logger

51 from vllm.model_executor.layers.activation import QuickGELU

52 from vllm.model_executor.layers.attention import MMEncoderAttention

53 from vllm.model_executor.layers.conv import Conv3dLayer

54 +from vllm.model_executor.layers.fusion.mm_input_norm import (

55 + IdentityInputNorm,

56 + build_mm_input_norm,

57 +)

58 from vllm.model_executor.layers.linear import (

59 ColumnParallelLinear,

60 RowParallelLinear,

@@ -105,7 +109,6 @@ from .utils import (

109 maybe_prefix,

110 )

111 from .vision import (

108 - FusedInputNorm,

112 get_vit_attn_backend,

113 is_vit_use_data_parallel,

114 run_dp_sharded_mrope_vision_model,

@@ -562,9 +565,7 @@ class Qwen2VisionTransformer(nn.Module):

565 in_channels=in_channels,

566 embed_dim=embed_dim,

567 )

565 - self.input_norm = (

566 - input_norm if input_norm is not None else FusedInputNorm.identity()

567 - )

568 + self.input_norm = input_norm if input_norm is not None else IdentityInputNorm()

569

570 norm_layer = partial(nn.LayerNorm, eps=norm_eps)

571 head_dim = embed_dim // num_heads

@@ -1295,7 +1296,7 @@ class Qwen2VLForConditionalGeneration(

1296 config.vision_config,

1297 norm_eps=getattr(config, "rms_norm_eps", 1e-6),

1298 quant_config=quant_config,

1298 - input_norm=FusedInputNorm.from_model_config(self.model_config),

1299 + input_norm=build_mm_input_norm(self.model_config),

1300 prefix=maybe_prefix(prefix, "visual"),

1301 )

1302

vllm/model_executor/models/vision.py

+1 -165

@@ -8,10 +8,9 @@ from collections.abc import Callable

8 from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar

10 import torch

11 -import torch.nn as nn

11 from transformers import PreTrainedConfig

14 -from vllm.config import ModelConfig, MultiModalConfig, get_current_vllm_config_or_none

13 +from vllm.config import MultiModalConfig, get_current_vllm_config_or_none

14 from vllm.distributed import (

15 get_tensor_model_parallel_rank,

16 get_tensor_model_parallel_world_size,

@@ -19,7 +18,6 @@ from vllm.distributed import (

18 )

19 from vllm.logger import init_logger

20 from vllm.platforms import current_platform

22 -from vllm.transformers_utils.processor import get_processor, get_processor_config

21 from vllm.utils.math_utils import round_up

22 from vllm.v1.attention.backends.registry import AttentionBackendEnum

@@ -578,165 +576,3 @@ def run_dp_sharded_mrope_vision_model(

576 "Found unassigned embeddings"

577 )

578 return out_embeddings

581 -

582 -

583 -class FusedInputNorm(nn.Module):

584 - """Module that applies rescaling and normalization to input images.

585 - Equivalent to: output = (input * rescale_factor - mean) / std

586 - """

587 -

588 - def __init__(

589 - self,

590 - image_mean: list[float],

591 - image_std: list[float],

592 - rescale_factor: float,

593 - channel: int = 3,

594 - dtype: torch.dtype = torch.float32,

595 - ):

596 - super().__init__()

597 -

598 - self.channel = channel

599 -

600 - # Model construction can set the accelerator as PyTorch's default

601 - # device. Determine whether the normalization is an identity on CPU

602 - # so torch.allclose does not introduce a device synchronization while

603 - # the model is being initialized. The actual buffers below still use

604 - # the caller's default device.

605 - image_mean_cpu = torch.tensor(image_mean, dtype=dtype, device="cpu") * (

606 - 1.0 / rescale_factor

607 - )

608 - image_std_cpu = torch.tensor(image_std, dtype=dtype, device="cpu") * (

609 - 1.0 / rescale_factor

610 - )

611 - weight_cpu = 1.0 / image_std_cpu

612 - bias_cpu = -image_mean_cpu / image_std_cpu

613 - self.is_identity = bool(

614 - torch.allclose(weight_cpu, torch.ones_like(weight_cpu))

615 - and torch.allclose(bias_cpu, torch.zeros_like(bias_cpu))

616 - )

617 -

618 - if not self.is_identity:

619 - image_mean_tensor = torch.tensor(image_mean, dtype=dtype) * (

620 - 1.0 / rescale_factor

621 - )

622 - image_std_tensor = torch.tensor(image_std, dtype=dtype) * (

623 - 1.0 / rescale_factor

624 - )

625 - weight = 1.0 / image_std_tensor

626 - bias = -image_mean_tensor / image_std_tensor

627 - self.register_buffer("weight", weight)

628 - self.register_buffer("bias", bias)

629 - else:

630 - self.register_buffer("weight", None)

631 - self.register_buffer("bias", None)

632 -

633 - @property

634 - def dtype(self) -> torch.dtype:

635 - return self.weight.dtype

636 -

637 - @classmethod

638 - def identity(

639 - cls, channel: int = 3, dtype: torch.dtype = torch.float32

640 - ) -> "FusedInputNorm":

641 - return cls(

642 - image_mean=[0.0, 0.0, 0.0],

643 - image_std=[1.0, 1.0, 1.0],

644 - rescale_factor=1.0,

645 - channel=channel,

646 - dtype=dtype,

647 - )

648 -

649 - @classmethod

650 - def from_model_config(cls, model_config: "ModelConfig") -> nn.Module:

651 - if not model_config.multimodal_config.mm_device_do_normalize:

652 - return cls.identity()

653 -

654 - model = model_config.model

655 - revision = model_config.revision

656 -

657 - # Try to read parameters from the processor config

658 - config = get_processor_config(model, revision=revision)

659 - do_rescale = config.get("do_rescale", None)

660 - do_normalize = config.get("do_normalize", None)

661 - image_mean = config.get("image_mean", None)

662 - image_std = config.get("image_std", None)

663 - rescale_factor = config.get("rescale_factor", None)

664 -

665 - # Fallback to the image_processor object if any parameter is missing

666 - if None in [do_rescale, do_normalize, image_mean, image_std, rescale_factor]:

667 - image_processor = get_processor(model, revision=revision).image_processor

668 -

669 - if do_rescale is None:

670 - do_rescale = getattr(image_processor, "do_rescale", None)

671 - if do_normalize is None:

672 - do_normalize = getattr(image_processor, "do_normalize", None)

673 - if image_mean is None:

674 - image_mean = getattr(image_processor, "image_mean", None)

675 - if image_std is None:

676 - image_std = getattr(image_processor, "image_std", None)

677 - if rescale_factor is None:

678 - rescale_factor = getattr(image_processor, "rescale_factor", None)

679 -

680 - # Apply defaults based on flags

681 - if not do_rescale:

682 - rescale_factor = 1.0

683 - if not do_normalize:

684 - image_mean = [0.0, 0.0, 0.0]

685 - image_std = [1.0, 1.0, 1.0]

686 -

687 - # Ensure all required parameters are resolved

688 - assert None not in [

689 - do_rescale,

690 - do_normalize,

691 - image_mean,

692 - image_std,

693 - rescale_factor,

694 - ], "Some normalization parameters are still None after resolution."

695 -

696 - # If no processing is needed, return an identity module

697 - if not do_rescale and not do_normalize:

698 - return cls.identity()

699 -

700 - return cls(

701 - image_mean=image_mean, image_std=image_std, rescale_factor=rescale_factor

702 - )

703 -

704 - def forward(

705 - self,

706 - grid_thw: torch.Tensor,

707 - visual_dtype: torch.dtype,

708 - ) -> torch.Tensor:

709 - if self.is_identity:

710 - return grid_thw.to(visual_dtype)

711 -

712 - assert grid_thw.ndim == 2

713 - patches, size = grid_thw.shape

714 - patch_size = size // self.channel

715 -

716 - # On XPU, fuse the whole rescale + normalize into a single custom

717 - # kernel. The eager path below materializes an fp32 intermediate and

718 - # then casts back, which adds device-side compute that cancels the

719 - # bandwidth saving of transferring uint8 pixel_values. The fused

720 - # kernel reads uint8 directly and writes ``visual_dtype`` in one pass.

721 - if (

722 - current_platform.is_xpu()

723 - and grid_thw.dtype == torch.uint8

724 - and self.weight.dtype == torch.float32

725 - ):

726 - return torch.ops.vllm.xpu_fused_input_norm(

727 - grid_thw, self.weight, self.bias, visual_dtype

728 - )

729 -

730 - # Apply the per-channel affine transform directly instead of via

731 - # F.batch_norm. batch_norm dispatches to cuDNN, whose batch-norm

732 - # kernels cap the batch dimension near the CUDA grid limit (~65535);

733 - # here that dimension is the number of patches, which grows unbounded

734 - # with image resolution and batch size and overflows the cap on large

735 - # image-heavy requests (CUDNN_STATUS_INTERNAL_ERROR). The plain

736 - # broadcasted multiply-add is numerically identical and has no such

737 - # limit.

738 - x = grid_thw.to(self.dtype).view(patches, self.channel, patch_size)

739 - x = x * self.weight.view(1, self.channel, 1) + self.bias.view(

740 - 1, self.channel, 1

741 - )

742 - return x.view(patches, size).to(visual_dtype)