replace batch_norm to numerically identical without cudnn (#51734)
Signed-off-by: khushali9 <khushali.desai9@gmail.com> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Khushali Desai committed
Aug 10, 2026 at 21:43 UTC
dc5101fb1b2723fc6264249569f453a803304ca5
3 files changed
+80
-19
docs/design/mm_processing.md
+8
-4
@@ -72,10 +72,14 @@ To accelerate the multi‑modal data pipeline (decoding, resizing, normalisation
72
73
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.
74
75
-- **How it works**: We use a dedicated `FusedInputNorm` module — essentially a frozen `BatchNorm1d` with the rescale factor baked into its statistics — and bake the rescale factor (typically 1/255) directly into the effective mean and standard deviation:
76
- - Effective mean = `image_mean * (1/rescale_factor)`
77
- - Effective std = `image_std * (1/rescale_factor)`
78
-- **At runtime**: The `FusedInputNorm` layer takes raw `uint8` pixel values (0–255) and performs the full normalised mapping in a single GPU kernel—no CPU involvement.
75
+**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`.
76
+
77
+The parameters are set as follows:
78
+
79
+- `weight` controls both the standard deviation and the rescale factor
80
+- `bias` centers the data using the mean and the same rescale factor
81
+
82
+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.
83
84
#### Optimized Data Path for Fused Normalisation
85
tests/models/test_vision.py
+60
@@ -13,6 +13,7 @@ from vllm.distributed.parallel_state import (
13
initialize_model_parallel,
14
)
15
from vllm.model_executor.models.vision import (
16
+ FusedInputNorm,
17
get_load_balance_assignment,
18
resolve_visual_encoder_outputs,
19
run_dp_sharded_mrope_vision_model,
@@ -492,3 +493,62 @@ def test_simple_mrope_vision_model_spatial_merge(spatial_merge_size: int):
493
494
assert output.shape[0] == expected_output_patches
495
assert output.shape[1] == vision_model.out_hidden_size
496
+
497
+
498
+def _reference_input_norm(
499
+ pixel_values: torch.Tensor,
500
+ image_mean: list[float],
501
+ image_std: list[float],
502
+ rescale_factor: float,
503
+ channel: int,
504
+ out_dtype: torch.dtype,
505
+) -> torch.Tensor:
506
+ """Straightforward per-channel affine: (x * rescale - mean) / std."""
507
+ patches, size = pixel_values.shape
508
+ patch_size = size // channel
509
+ mean = torch.tensor(image_mean, dtype=torch.float32).view(1, channel, 1)
510
+ std = torch.tensor(image_std, dtype=torch.float32).view(1, channel, 1)
511
+ x = pixel_values.to(torch.float32).view(patches, channel, patch_size)
512
+ x = (x * rescale_factor - mean) / std
513
+ return x.view(patches, size).to(out_dtype)
514
+
515
+
516
+@pytest.mark.parametrize("num_patches", [1, 37, 70000])
517
+def test_fused_input_norm_matches_reference(num_patches: int):
518
+ """FusedInputNorm must equal the plain affine, including for num_patches
519
+ above the cuDNN batch-norm grid limit (~65535) that previously raised
520
+ CUDNN_STATUS_INTERNAL_ERROR (issue #51717)."""
521
+ channel = 3
522
+ patch_size = 14 * 14
523
+ image_mean = [0.48145466, 0.4578275, 0.40821073]
524
+ image_std = [0.26862954, 0.26130258, 0.27577711]
525
+ rescale_factor = 1.0 / 255.0
526
+
527
+ set_random_seed(0)
528
+ pixel_values = torch.randint(
529
+ 0, 256, (num_patches, channel * patch_size), dtype=torch.float32
530
+ )
531
+
532
+ norm = FusedInputNorm(
533
+ image_mean=image_mean,
534
+ image_std=image_std,
535
+ rescale_factor=rescale_factor,
536
+ channel=channel,
537
+ )
538
+ assert not norm.is_identity
539
+
540
+ out = norm(pixel_values, visual_dtype=torch.float32)
541
+ expected = _reference_input_norm(
542
+ pixel_values, image_mean, image_std, rescale_factor, channel, torch.float32
543
+ )
544
+ torch.testing.assert_close(out, expected)
545
+
546
+
547
+def test_fused_input_norm_identity_passthrough():
548
+ """The identity configuration returns the input unchanged (cast only)."""
549
+ norm = FusedInputNorm.identity()
550
+ assert norm.is_identity
551
+
552
+ pixel_values = torch.randn(8, 3 * 196, dtype=torch.float32)
553
+ out = norm(pixel_values, visual_dtype=torch.bfloat16)
554
+ torch.testing.assert_close(out, pixel_values.to(torch.bfloat16))
vllm/model_executor/models/vision.py
+12
-15
@@ -9,7 +9,6 @@ from typing import Final, Generic, Literal, Protocol, TypeAlias, TypeVar
9
10
import torch
11
import torch.nn as nn
12
-import torch.nn.functional as F
12
from transformers import PretrainedConfig
13
14
from vllm.config import ModelConfig, MultiModalConfig, get_current_vllm_config_or_none
@@ -632,13 +631,9 @@ class FusedInputNorm(nn.Module):
631
bias = -image_mean_tensor / image_std_tensor
632
self.register_buffer("weight", weight)
633
self.register_buffer("bias", bias)
635
- self.register_buffer("running_mean", torch.zeros_like(image_mean_tensor))
636
- self.register_buffer("running_var", torch.ones_like(image_mean_tensor))
634
else:
635
self.register_buffer("weight", None)
636
self.register_buffer("bias", None)
640
- self.register_buffer("running_mean", None)
641
- self.register_buffer("running_var", None)
637
638
@property
639
def dtype(self) -> torch.dtype:
@@ -723,14 +718,16 @@ class FusedInputNorm(nn.Module):
718
patches, size = grid_thw.shape
719
patch_size = size // self.channel
720
726
- grid_thw = grid_thw.view(patches, self.channel, patch_size)
727
- grid_thw = F.batch_norm(
728
- grid_thw.to(self.dtype),
729
- running_mean=self.running_mean,
730
- running_var=self.running_var,
731
- weight=self.weight,
732
- bias=self.bias,
733
- training=False,
734
- eps=0.0,
721
+ # Apply the per-channel affine transform directly instead of via
722
+ # F.batch_norm. batch_norm dispatches to cuDNN, whose batch-norm
723
+ # kernels cap the batch dimension near the CUDA grid limit (~65535);
724
+ # here that dimension is the number of patches, which grows unbounded
725
+ # with image resolution and batch size and overflows the cap on large
726
+ # image-heavy requests (CUDNN_STATUS_INTERNAL_ERROR). The plain
727
+ # broadcasted multiply-add is numerically identical and has no such
728
+ # limit.
729
+ x = grid_thw.to(self.dtype).view(patches, self.channel, patch_size)
730
+ x = x * self.weight.view(1, self.channel, 1) + self.bias.view(
731
+ 1, self.channel, 1
732
)
736
- return grid_thw.view(patches, size).to(visual_dtype)
733
+ return x.view(patches, size).to(visual_dtype)