A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
Add RADIO Vision Encoder Support to vLLM (#24595)
Signed-off-by: Daniel Afrimi <danielafrimi8@gmail.com> Co-authored-by: root <root@cw-dfw-h100-001-305-026.cm.cluster>
danielafrimi committed
Sep 17, 2025 at 15:53 UTC
252ada5559808783d6a23b489156b3705cea0417
5 files changed
+828
-58
tests/models/multimodal/pooling/test_radio.py
new
+86
@@ -0,0 +1,86 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+import torch
5
+import torch.nn as nn
6
+from huggingface_hub import snapshot_download
7
+from transformers import AutoConfig, AutoModel, CLIPImageProcessor
8
+
9
+from vllm.distributed import cleanup_dist_env_and_memory
10
+from vllm.model_executor.models.radio import RadioModel
11
+from vllm.transformers_utils.configs.radio import RadioConfig
12
+from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE
13
+
14
+from ....conftest import ImageTestAssets
15
+
16
+# we use snapshot_download to prevent conflicts between
17
+# dynamic_module and trust_remote_code for hf_runner
18
+DOWNLOAD_PATTERN = ["*.json", "*.py", "*.safetensors", "*.txt", "*.model"]
19
+
20
+
21
+@torch.inference_mode()
22
+def run_radio_test(
23
+ image_assets: ImageTestAssets,
24
+ model_id: str,
25
+ *,
26
+ dtype: str,
27
+):
28
+ model = snapshot_download(model_id, allow_patterns=DOWNLOAD_PATTERN)
29
+ torch_dtype = STR_DTYPE_TO_TORCH_DTYPE[dtype]
30
+
31
+ img_processor = CLIPImageProcessor.from_pretrained(model)
32
+ images = [asset.pil_image for asset in image_assets]
33
+ # Input resolution must be a multiple of `self.min_resolution_step`.
34
+ # Using `self.get_nearest_supported_resolution`, for assets 432x642 the
35
+ # nearest supported resolution is 432x640.
36
+ pixel_values = [
37
+ img_processor(
38
+ image,
39
+ return_tensors='pt').pixel_values.to(torch_dtype)[:, :, :, :640]
40
+ for image in images
41
+ ]
42
+
43
+ config = AutoConfig.from_pretrained(model_id, trust_remote_code=True)
44
+
45
+ hf_model = AutoModel.from_pretrained(
46
+ model_id,
47
+ config=config,
48
+ torch_dtype=torch_dtype,
49
+ trust_remote_code=True,
50
+ ).to("cuda")
51
+ hf_model.eval()
52
+
53
+ hf_outputs_per_image = [
54
+ hf_model(pixel_value.to("cuda")).features
55
+ for pixel_value in pixel_values
56
+ ]
57
+
58
+ radio_config = RadioConfig(model_name=config.args["model"],
59
+ reg_tokens=config.args["register_multiple"])
60
+ vllm_model = RadioModel(radio_config)
61
+ vllm_model.load_weights(hf_model.state_dict())
62
+ vllm_model = vllm_model.to("cuda", torch_dtype)
63
+
64
+ vllm_outputs_per_image = [
65
+ vllm_model(pixel_values=pixel_value.to("cuda"))
66
+ for pixel_value in pixel_values
67
+ ]
68
+ del vllm_model, hf_model
69
+ cleanup_dist_env_and_memory()
70
+
71
+ cos_similar = nn.CosineSimilarity(dim=-1)
72
+ for vllm_output, hf_output in zip(vllm_outputs_per_image,
73
+ hf_outputs_per_image):
74
+ assert cos_similar(vllm_output, hf_output).mean() > 0.99
75
+
76
+
77
+@pytest.mark.parametrize("model_id", [
78
+ "nvidia/C-RADIOv2-H",
79
+])
80
+@pytest.mark.parametrize("dtype", ["half"])
81
+def test_radio(dist_init, image_assets, model_id, dtype: str) -> None:
82
+ run_radio_test(
83
+ image_assets,
84
+ model_id,
85
+ dtype=dtype,
86
+ )
vllm/model_executor/models/nano_nemotron_vl.py
+60
-58
@@ -18,8 +18,8 @@ import torch
18
import torch.nn as nn
19
import torchvision.transforms as T
20
from PIL import Image
21
-from transformers import (AutoModel, BatchEncoding, BatchFeature,
22
- PretrainedConfig, TensorType)
21
+from transformers import (BatchEncoding, BatchFeature, PretrainedConfig,
22
+ TensorType)
23
24
from vllm.config import VllmConfig
25
from vllm.model_executor.layers.activation import ReLUSquaredActivation
@@ -32,6 +32,7 @@ from vllm.model_executor.models.internvl import (calculate_internvl_targets,
32
get_internvl_target_ratios)
33
from vllm.model_executor.models.module_mapping import MultiModelKeys
34
from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
35
+from vllm.model_executor.models.radio import RadioModel
36
from vllm.model_executor.models.utils import (flatten_bn,
37
init_vllm_registered_model,
38
maybe_prefix,
@@ -48,6 +49,7 @@ from vllm.multimodal.processing import (BaseMultiModalProcessor,
49
PromptUpdate, PromptUpdateDetails)
50
from vllm.multimodal.profiling import BaseDummyInputsBuilder
51
from vllm.sequence import IntermediateTensors
52
+from vllm.transformers_utils.configs.radio import RadioConfig
53
from vllm.transformers_utils.tokenizer import AnyTokenizer
54
from vllm.utils.tensor_schema import TensorSchema, TensorShape
55
@@ -122,11 +124,6 @@ NanoNemotronVLVideoInputs = Union[NanoNemotronVLVideoPixelInputs,
124
NanoNemotronVLVideoEmbeddingInputs]
125
126
125
-def input_conditioner(x, norm_mean, norm_std):
126
- y = (x - norm_mean) / norm_std
127
- return y
128
-
129
-
127
def dynamic_preprocess(image,
128
*,
129
image_size=512,
@@ -305,8 +302,7 @@ class BaseNanoNemotronVLProcessor(ABC):
302
images, max_num_tiles)
303
image_inputs: dict[str, NestedTensors] = {
304
"pixel_values_flat":
308
- input_conditioner(torch.cat(pixel_values_lst), self.norm_mean,
309
- self.norm_std),
305
+ torch.cat(pixel_values_lst),
306
"image_num_patches":
307
torch.tensor([len(item) for item in pixel_values_lst]),
308
}
@@ -428,8 +424,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):
424
425
video_inputs: dict[str, NestedTensors] = {
426
"pixel_values_flat_video":
431
- input_conditioner(torch.cat(pixel_values_lst_video),
432
- self.norm_mean, self.norm_std),
427
+ torch.cat(pixel_values_lst_video),
428
"video_num_patches":
429
torch.tensor([len(item) for item in pixel_values_lst_video]),
430
}
@@ -905,18 +900,9 @@ class NemotronH_Nano_VL(nn.Module, HasInnerState, IsHybrid,
900
hf_config=config.text_config,
901
prefix=maybe_prefix(prefix, "language_model"),
902
)
908
- self.vision_model = AutoModel.from_config(config.vision_config,
909
- trust_remote_code=True)
910
- self.vision_model.model._initialize_weights = (
911
- self.vision_model.model._init_weights)
912
- # Move input normalization to processor to mirror original HF
913
- # implementation where normalization is done in fp32
914
- self.vision_model.radio_model.make_preprocessor_external()
915
- self.vision_model = self.vision_model.to(
903
+ self.vision_model = self.get_vit_model_from_radio_config(config).to(
904
self.language_model.config.torch_dtype)
905
918
- self.drop_vision_class_token = True
919
-
906
# Construct the vision projection.
907
vit_hidden_size = config.vit_hidden_size
908
vision_projection_hidden_size = config.projector_hidden_size
@@ -972,7 +958,7 @@ class NemotronH_Nano_VL(nn.Module, HasInnerState, IsHybrid,
958
return x
959
960
def extract_feature(self, pixel_values):
975
- vit_embeds = self.vision_model(pixel_values).features
961
+ vit_embeds = self.vision_model(pixel_values)
962
vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
963
h = w = int(vit_embeds.shape[1]**0.5)
964
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
@@ -1212,47 +1198,39 @@ class NemotronH_Nano_VL(nn.Module, HasInnerState, IsHybrid,
1198
sampling_metadata)
1199
1200
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
1201
+ adapter_dict = dict(self.mlp1.named_parameters())
1202
1216
- def is_vision_model_weights(weight: tuple[str, torch.Tensor]):
1217
- return weight[0].startswith("vision_model")
1203
+ def is_llm(name: str) -> bool:
1204
+ return name.startswith("language_model")
1205
1206
def is_adapter_weights(weight: tuple[str, torch.Tensor]):
1207
return weight[0].startswith("mlp1")
1208
1222
- # Get references to parameters for direct loading
1223
- vision_model_dict = dict(self.vision_model.named_parameters())
1224
- vision_model_buffers = dict(self.vision_model.named_buffers())
1225
- adapter_dict = dict(self.mlp1.named_parameters())
1226
-
1227
- def llm_weights_generator():
1228
- # Single pass over weights
1229
- for name, w in weights:
1230
- if is_vision_model_weights((name, w)):
1231
- # Load vision encoder weights directly
1232
- trimmed_name = ".".join(name.split(".")[1:])
1233
- if "input_conditioner" in trimmed_name:
1234
- continue
1235
- if trimmed_name in vision_model_buffers:
1236
- param = vision_model_buffers[trimmed_name]
1237
- else:
1238
- param = vision_model_dict[trimmed_name]
1239
- with torch.no_grad():
1240
- default_weight_loader(param, w)
1241
- elif is_adapter_weights((name, w)):
1242
- # Load vision-language adapter weights directly
1243
- trimmed_name = ".".join(name.split(".")[1:])
1244
- param = adapter_dict[trimmed_name]
1245
- with torch.no_grad():
1246
- default_weight_loader(param, w)
1247
- else:
1248
- # LLM weights: yield them to be loaded
1249
- # by language_model.load_weights
1250
- assert name.startswith("language_model")
1251
- trimmed_name = ".".join(name.split(".")[1:])
1252
- yield (trimmed_name, w)
1253
-
1254
- # Now we call the language model load with the generator
1255
- self.language_model.load_weights(llm_weights_generator())
1209
+ def is_vision_weights(name: str) -> bool:
1210
+ return name.startswith("vision_model.radio_model.")
1211
+
1212
+ # Separate weights by component
1213
+ llm_weights = []
1214
+ vision_weights = []
1215
+
1216
+ for name, w in weights:
1217
+ if is_llm(name):
1218
+ # Strip 'language_model.' prefix for LLM weights
1219
+ llm_weights.append((".".join(name.split(".")[1:]), w))
1220
+ elif is_adapter_weights((name, w)):
1221
+ # Load vision-language adapter weights directly
1222
+ trimmed_name = ".".join(name.split(".")[1:])
1223
+ param = adapter_dict[trimmed_name]
1224
+ with torch.no_grad():
1225
+ default_weight_loader(param, w)
1226
+ elif is_vision_weights(name):
1227
+ # Convert: vision_model.radio_model.* → radio_model.*
1228
+ hf_key = name[len(
1229
+ "vision_model."):] # Remove "vision_model." prefix
1230
+ vision_weights.append((hf_key, w))
1231
+
1232
+ self.language_model.load_weights(llm_weights)
1233
+ self.vision_model.load_weights(vision_weights)
1234
1235
def print_architecture(self,
1236
detailed: bool = True,
@@ -1370,6 +1348,30 @@ class NemotronH_Nano_VL(nn.Module, HasInnerState, IsHybrid,
1348
},
1349
}
1350
1351
+ def get_vit_model_from_radio_config(self, hf_config):
1352
+ hf_config_vision = hf_config.vision_config
1353
+ model_name = hf_config_vision.args.get("model")
1354
+ if model_name is None:
1355
+ raise ValueError(f'Unsupported vit model type: {model_name}')
1356
+
1357
+ preferred_resolution = getattr(hf_config_vision,
1358
+ "preferred_resolution", None)
1359
+ image_size = preferred_resolution[0] if preferred_resolution else 224
1360
+ patch_size = getattr(hf_config_vision, "patch_size", 16)
1361
+
1362
+ radio_config = RadioConfig(
1363
+ model_name=model_name,
1364
+ image_size=image_size,
1365
+ patch_size=patch_size,
1366
+ norm_mean=hf_config.norm_mean,
1367
+ norm_std=hf_config.norm_std,
1368
+ reg_tokens=(hf_config_vision.args.get("register_multiple")
1369
+ if hasattr(hf_config_vision, "args")
1370
+ and isinstance(hf_config_vision.args, dict) else None),
1371
+ )
1372
+
1373
+ return RadioModel(config=radio_config)
1374
+
1375
def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
1376
return self.language_model.mamba_cache.copy_inputs_before_cuda_graphs(
1377
input_buffers, **kwargs)
vllm/model_executor/models/radio.py
new
+576
@@ -0,0 +1,576 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# Copyright (c) 2023-2024, NVIDIA CORPORATION. All rights reserved.
4
+#
5
+# NVIDIA CORPORATION and its licensors retain all intellectual property
6
+# and proprietary rights in and to this software, related documentation
7
+# and any modifications thereto. Any use, reproduction, disclosure or
8
+# distribution of this software and related documentation without an express
9
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
10
+
11
+import math
12
+from collections.abc import Iterable
13
+from itertools import repeat
14
+from typing import Optional, Union
15
+
16
+import torch
17
+import torch.nn as nn
18
+import torch.nn.functional as F
19
+from einops import rearrange
20
+from transformers import PretrainedConfig
21
+
22
+from vllm.model_executor.layers.quantization import QuantizationConfig
23
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
24
+from vllm.model_executor.models.intern_vit import InternVisionEncoder
25
+
26
+input_dim_t = Union[int, tuple[int, int]]
27
+norm_t = Union[tuple[float, float, float], torch.Tensor]
28
+
29
+
30
+def _ntuple(n):
31
+
32
+ def parse(x):
33
+ if isinstance(x, Iterable) and not isinstance(x, str):
34
+ return tuple(x)
35
+ return tuple(repeat(x, n))
36
+
37
+ return parse
38
+
39
+
40
+to_1tuple = _ntuple(1)
41
+to_2tuple = _ntuple(2)
42
+to_3tuple = _ntuple(3)
43
+to_4tuple = _ntuple(4)
44
+to_ntuple = _ntuple
45
+
46
+
47
+class InputConditioner(nn.Module):
48
+
49
+ def __init__(
50
+ self,
51
+ input_scale: float,
52
+ norm_mean: norm_t,
53
+ norm_std: norm_t,
54
+ dtype: torch.dtype = None,
55
+ ):
56
+ super().__init__()
57
+
58
+ self.dtype = dtype
59
+
60
+ self.register_buffer("norm_mean", _to_tensor(norm_mean) / input_scale)
61
+ self.register_buffer("norm_std", _to_tensor(norm_std) / input_scale)
62
+
63
+ def forward(self, x: torch.Tensor):
64
+ y = (x - self.norm_mean) / self.norm_std
65
+ if self.dtype is not None:
66
+ y = y.to(self.dtype)
67
+ return y
68
+
69
+
70
+def _to_tensor(v: norm_t):
71
+ return torch.as_tensor(v, dtype=torch.float32).view(-1, 1, 1)
72
+
73
+
74
+class ClsToken(nn.Module):
75
+
76
+ def __init__(
77
+ self,
78
+ ndim: int,
79
+ num_tokens: int = 1,
80
+ enabled: bool = True,
81
+ register_multiple: Optional[int] = None,
82
+ num_registers: Optional[int] = None,
83
+ ):
84
+ super().__init__()
85
+
86
+ self.ndim = ndim
87
+ self.enabled = enabled
88
+ self.num_registers = 0
89
+ self.num_tokens = num_tokens
90
+ if enabled:
91
+ if num_registers:
92
+ self.num_registers = num_registers
93
+ elif register_multiple:
94
+ self.num_registers = register_multiple - (num_tokens %
95
+ register_multiple)
96
+
97
+ scale = ndim**-0.5
98
+ self.token = nn.Parameter(
99
+ torch.randn(num_tokens + self.num_registers, ndim) * scale)
100
+
101
+ else:
102
+ self.token = None
103
+
104
+ self.num_patches = self.num_tokens + self.num_registers
105
+
106
+ def forward(self, x: torch.Tensor):
107
+ if self.token is None:
108
+ return x
109
+
110
+ token = self.token.unsqueeze(0).expand(x.shape[0], -1, -1)
111
+ x = torch.cat([
112
+ token,
113
+ x,
114
+ ], dim=1)
115
+
116
+ return x
117
+
118
+
119
+class ViTPatchGenerator(nn.Module):
120
+
121
+ def __init__(
122
+ self,
123
+ # config: PretrainedConfig,
124
+ patch_size: int,
125
+ embed_dim: int,
126
+ input_dims: input_dim_t,
127
+ abs_pos: bool = True,
128
+ normalize_patches: bool = False,
129
+ cls_token: bool = False,
130
+ max_input_dims: Optional[input_dim_t] = None,
131
+ pos_dropout: float = 0.0,
132
+ return_pos_enc: bool = False,
133
+ num_cls_tokens: int = 1,
134
+ register_multiple: Optional[int] = None,
135
+ num_registers: Optional[int] = None,
136
+ patch_bias: bool = False,
137
+ device=None,
138
+ dtype=None,
139
+ ):
140
+ super().__init__()
141
+ if isinstance(input_dims, int):
142
+ input_dims = (input_dims, input_dims)
143
+
144
+ if max_input_dims is None:
145
+ max_input_dims = input_dims
146
+ if isinstance(max_input_dims, int):
147
+ max_input_dims = (max_input_dims, max_input_dims)
148
+
149
+ max_input_dims = tuple(
150
+ int(math.ceil(d / patch_size) * patch_size)
151
+ for d in max_input_dims)
152
+
153
+ self.cpe_mode = max_input_dims != input_dims
154
+ self.pos_dropout = pos_dropout
155
+ self.return_pos_enc = return_pos_enc
156
+
157
+ factory = dict(device=device, dtype=dtype)
158
+
159
+ self.patch_size = patch_size
160
+ self.abs_pos = abs_pos
161
+ self.embed_dim = embed_dim
162
+
163
+ self.num_rows = max_input_dims[0] // patch_size
164
+ self.num_cols = max_input_dims[1] // patch_size
165
+ self.input_dims = tuple(d // patch_size for d in input_dims)
166
+ self.num_patches = self.num_rows * self.num_cols
167
+ self.max_input_dims = max_input_dims
168
+
169
+ self.im_to_patches = Im2Patches(patch_size)
170
+ self.embedder = ViTPatchLinear(patch_size,
171
+ embed_dim,
172
+ bias=patch_bias,
173
+ **factory)
174
+
175
+ if abs_pos:
176
+ scale = embed_dim**-0.5
177
+ self.pos_embed = nn.Parameter(
178
+ torch.randn(1, self.num_patches, embed_dim, **factory) * scale)
179
+
180
+ self.cls_token = ClsToken(
181
+ embed_dim,
182
+ num_tokens=num_cls_tokens,
183
+ enabled=cls_token,
184
+ register_multiple=register_multiple,
185
+ num_registers=num_registers,
186
+ )
187
+
188
+ self.patch_normalizer = nn.LayerNorm(
189
+ embed_dim) if normalize_patches else nn.Identity()
190
+
191
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
192
+ patches = self.embed_patches(x)
193
+ patches, pos_enc = self.apply_pos_enc(patches, input_size=x.shape[2:])
194
+ patches = self.cls_token(patches)
195
+ patches = self.patch_normalizer(patches)
196
+ if self.return_pos_enc:
197
+ return patches, pos_enc
198
+ return patches
199
+
200
+ @property
201
+ def apply_cls_token(self):
202
+ return self.cls_token.enabled
203
+
204
+ @property
205
+ def num_cls_tokens(self):
206
+ return self.cls_token.num_tokens
207
+
208
+ @property
209
+ def num_cls_patches(self):
210
+ return self.cls_token.num_patches
211
+
212
+ @property
213
+ def num_registers(self):
214
+ return self.cls_token.num_registers
215
+
216
+ @property
217
+ def num_skip(self):
218
+ return self.num_cls_tokens + self.num_registers
219
+
220
+ def _load_embed(self, src_embed: torch.Tensor, targ_embed: nn.Parameter):
221
+ if src_embed.shape != targ_embed.shape:
222
+ src_size = int(math.sqrt(src_embed.shape[1]))
223
+
224
+ assert src_size**2 == src_embed.shape[
225
+ 1], 'Unable to interpolate non-square embedding'
226
+
227
+ src_embed = rearrange(src_embed,
228
+ 'b (h w) c -> b c h w',
229
+ h=src_size,
230
+ w=src_size)
231
+ src_embed = F.interpolate(src_embed,
232
+ size=(self.num_rows, self.num_cols),
233
+ mode='bicubic',
234
+ align_corners=True,
235
+ antialias=False)
236
+ src_embed = rearrange(src_embed, 'b c h w -> b (h w) c')
237
+ targ_embed.data.copy_(src_embed)
238
+
239
+ def _load_projection(self, src_proj_weight: torch.Tensor,
240
+ targ_proj_weight: torch.Tensor):
241
+ if src_proj_weight.shape != targ_proj_weight.shape:
242
+ src_patch_size = int(math.sqrt(src_proj_weight.shape[1] // 3))
243
+
244
+ assert (src_patch_size**2) * 3 == src_proj_weight.shape[
245
+ 1], 'Unable to interpolate non-square patch size'
246
+
247
+ src_proj_weight = rearrange(src_proj_weight,
248
+ 'b (c h w) -> b c h w',
249
+ c=3,
250
+ h=src_patch_size,
251
+ w=src_patch_size)
252
+ src_proj_weight = F.interpolate(src_proj_weight,
253
+ size=(self.patch_size,
254
+ self.patch_size),
255
+ mode='bicubic',
256
+ align_corners=True,
257
+ antialias=False)
258
+ src_proj_weight = rearrange(src_proj_weight,
259
+ 'b c h w -> b (c h w)')
260
+ targ_proj_weight.data.copy_(src_proj_weight)
261
+
262
+ def embed_patches(self, x: torch.Tensor) -> torch.Tensor:
263
+ patches = self.im_to_patches(x)
264
+ patches = self.embedder(patches)
265
+ return patches
266
+
267
+ def apply_pos_enc(
268
+ self,
269
+ patches: torch.Tensor,
270
+ patch_idxs: Optional[torch.Tensor] = None,
271
+ input_size: Optional[tuple[int, int]] = None,
272
+ ) -> torch.Tensor:
273
+ if not self.abs_pos:
274
+ return patches
275
+
276
+ pos_enc = self.get_pos_enc(patches.shape[0], patch_idxs, input_size)
277
+
278
+ if self.training and self.pos_dropout > 0:
279
+ keeps = torch.rand(patches.shape[0],
280
+ 1,
281
+ 1,
282
+ dtype=pos_enc.dtype,
283
+ device=pos_enc.device) > self.pos_dropout
284
+ pos_enc_drop = torch.where(keeps, pos_enc, 0)
285
+ else:
286
+ pos_enc_drop = pos_enc
287
+
288
+ return patches + pos_enc_drop, pos_enc
289
+
290
+ def get_pos_enc(
291
+ self,
292
+ batch_size: int,
293
+ patch_idxs: Optional[torch.Tensor] = None,
294
+ input_size: Optional[tuple[int, int]] = None,
295
+ ) -> torch.Tensor:
296
+ if input_size is None:
297
+ input_dims = self.input_dims
298
+ else:
299
+ input_dims = tuple(d // self.patch_size for d in input_size)
300
+
301
+ pos_embed = self._get_pos_embeddings(batch_size, input_dims)
302
+
303
+ if patch_idxs is None:
304
+ return pos_embed
305
+
306
+ exp_patch_idxs = patch_idxs.unsqueeze(-1).expand(
307
+ -1, -1, pos_embed.shape[-1])
308
+
309
+ pos_embed = torch.gather(pos_embed.expand(patch_idxs.shape[0], -1, -1),
310
+ dim=1,
311
+ index=exp_patch_idxs)
312
+ return pos_embed
313
+
314
+ def _get_pos_embeddings(self, batch_size: int, input_dims: tuple[int,
315
+ int]):
316
+ if (self.num_rows, self.num_cols) == input_dims:
317
+ return self.pos_embed
318
+
319
+ pos_embed = self.pos_embed.reshape(1, self.num_rows, self.num_cols,
320
+ -1).permute(0, 3, 1, 2)
321
+
322
+ def window_select(pos_embed):
323
+ if input_dims[0] < pos_embed.shape[-2]:
324
+ pos_embed = pos_embed[..., :input_dims[0], :]
325
+ if input_dims[1] < pos_embed.shape[-1]:
326
+ pos_embed = pos_embed[..., :, :input_dims[1]]
327
+ return pos_embed
328
+
329
+ if self.cpe_mode:
330
+ if self.training:
331
+ min_scale = math.sqrt(0.1)
332
+ scale = torch.rand(batch_size, 1, 1, device=pos_embed.device
333
+ ) * (1 - min_scale) + min_scale
334
+ aspect_min = math.log(3 / 4)
335
+ aspect_max = -aspect_min
336
+ aspect = torch.exp(
337
+ torch.rand(batch_size, 1, 1, device=pos_embed.device) *
338
+ (aspect_max - aspect_min) + aspect_min)
339
+
340
+ scale_x = scale * aspect
341
+ scale_y = scale * (1 / aspect)
342
+ scale_xy = torch.stack([scale_x, scale_y], dim=-1).clamp_(0, 1)
343
+
344
+ pos_xy = torch.rand(
345
+ batch_size, 1, 1, 2,
346
+ device=pos_embed.device) * (1 - scale_xy)
347
+
348
+ lin_x = torch.linspace(
349
+ 0, 1, steps=input_dims[1],
350
+ device=pos_embed.device)[None, None].expand(
351
+ batch_size, input_dims[0], -1)
352
+ lin_y = torch.linspace(
353
+ 0, 1, steps=input_dims[0],
354
+ device=pos_embed.device)[None, :, None].expand(
355
+ batch_size, -1, input_dims[1])
356
+
357
+ lin_xy = torch.stack([lin_x, lin_y], dim=-1)
358
+
359
+ grid_xy = lin_xy * scale_xy + pos_xy
360
+
361
+ # Convert to [-1, 1] range
362
+ grid_xy.mul_(2).sub_(1)
363
+
364
+ pos_embed = F.grid_sample(
365
+ pos_embed.float().expand(batch_size, -1, -1, -1),
366
+ grid=grid_xy,
367
+ mode='bilinear',
368
+ padding_mode='zeros',
369
+ align_corners=True,
370
+ ).to(pos_embed.dtype)
371
+ else:
372
+ max_dim = max(input_dims)
373
+ pos_embed = F.interpolate(pos_embed.float(),
374
+ size=(max_dim, max_dim),
375
+ align_corners=True,
376
+ mode='bilinear').to(pos_embed.dtype)
377
+
378
+ pos_embed = window_select(pos_embed)
379
+ else:
380
+ pos_embed = window_select(pos_embed)
381
+
382
+ if pos_embed.shape[-2:] != input_dims:
383
+ pos_embed = F.interpolate(pos_embed.float(),
384
+ size=input_dims,
385
+ align_corners=True,
386
+ mode='bilinear').to(pos_embed.dtype)
387
+
388
+ pos_embed = pos_embed.flatten(2).permute(0, 2, 1)
389
+
390
+ return pos_embed
391
+
392
+
393
+class Im2Patches(nn.Module):
394
+
395
+ def __init__(self, patch_size: int):
396
+ super().__init__()
397
+ self.patch_size = patch_size
398
+
399
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
400
+ if self.patch_size == 1:
401
+ patches = x.flatten(2)
402
+ patches = patches.permute(0, 2, 1)
403
+ return patches
404
+
405
+ py = x.shape[-2] // self.patch_size
406
+ px = x.shape[-1] // self.patch_size
407
+ patches = rearrange(
408
+ x,
409
+ 'b c (py yy) (px xx) -> b (py px) (c yy xx)',
410
+ py=py,
411
+ yy=self.patch_size,
412
+ px=px,
413
+ xx=self.patch_size,
414
+ )
415
+ return patches
416
+
417
+
418
+class ViTPatchLinear(nn.Linear):
419
+
420
+ def __init__(self,
421
+ patch_size: int,
422
+ embed_dim: int,
423
+ bias: bool = False,
424
+ **factory):
425
+ super().__init__(3 * (patch_size**2), embed_dim, bias=bias, **factory)
426
+ self.patch_size = patch_size
427
+
428
+
429
+class RadioInternVisionModel(nn.Module):
430
+ packed_modules_mapping = {
431
+ "qkv": ["qkv"],
432
+ }
433
+
434
+ def __init__(
435
+ self,
436
+ config: PretrainedConfig = None,
437
+ quant_config: Optional[QuantizationConfig] = None,
438
+ *,
439
+ num_hidden_layers_override: Optional[int] = None,
440
+ num_dummy_heads: int = 0,
441
+ prefix: str = "",
442
+ ) -> None:
443
+ super().__init__()
444
+
445
+ self.config = config
446
+ self.img_size, self.grid_size, self.num_patches = self._init_img_size(
447
+ to_2tuple(config.patch_size), config.image_size)
448
+ max_img_size = int(
449
+ round(config.max_img_size / config.patch_size) * config.patch_size)
450
+ self.patch_generator = ViTPatchGenerator(
451
+ config.patch_size,
452
+ config.hidden_size,
453
+ input_dims=self.img_size,
454
+ max_input_dims=max_img_size,
455
+ cls_token=True,
456
+ register_multiple=config.reg_tokens)
457
+
458
+ self.encoder = InternVisionEncoder(
459
+ config=config,
460
+ quant_config=quant_config,
461
+ num_hidden_layers_override=num_hidden_layers_override,
462
+ num_dummy_heads=num_dummy_heads,
463
+ prefix=f"{prefix}.encoder",
464
+ )
465
+
466
+ def _init_img_size(self, patch_size, img_size: Union[int, tuple[int,
467
+ int]]):
468
+ if img_size is None:
469
+ return None, None, None
470
+ img_size = to_2tuple(img_size)
471
+ grid_size = tuple([s // p for s, p in zip(img_size, patch_size)])
472
+ num_patches = grid_size[0] * grid_size[1]
473
+ return img_size, grid_size, num_patches
474
+
475
+ def get_input_embeddings(self):
476
+ return self.embeddings
477
+
478
+ def forward(self, x: torch.Tensor) -> torch.FloatTensor:
479
+ assert self.patch_generator is not None
480
+ hidden_states = self.patch_generator(x)
481
+ encoder_outputs = self.encoder(inputs_embeds=hidden_states)
482
+ return encoder_outputs
483
+
484
+
485
+class RadioModel(nn.Module):
486
+ packed_modules_mapping = {
487
+ "qkv": ["qkv"],
488
+ }
489
+
490
+ def __init__(
491
+ self,
492
+ config: PretrainedConfig,
493
+ quant_config: Optional[QuantizationConfig] = None,
494
+ *,
495
+ num_hidden_layers_override: Optional[int] = None,
496
+ num_dummy_heads: int = 0,
497
+ prefix: str = "",
498
+ ) -> None:
499
+ super().__init__()
500
+
501
+ self.config = config
502
+ self.input_conditioner = InputConditioner(
503
+ input_scale=1.0,
504
+ norm_mean=config.norm_mean,
505
+ norm_std=config.norm_std,
506
+ )
507
+ self.model = RadioInternVisionModel(
508
+ config=config,
509
+ quant_config=quant_config,
510
+ num_hidden_layers_override=num_hidden_layers_override,
511
+ num_dummy_heads=num_dummy_heads,
512
+ prefix=prefix)
513
+
514
+ def forward(
515
+ self,
516
+ pixel_values: Optional[torch.Tensor] = None,
517
+ pixel_embeds: Optional[torch.Tensor] = None,
518
+ ) -> torch.FloatTensor:
519
+ x = self.input_conditioner(pixel_values)
520
+ y = self.model(x)
521
+ return self._extract_final(y)
522
+
523
+ def load_weights(self, weights) -> set[str]:
524
+ loaded_params: set[str] = set()
525
+ params_dict = dict(self.named_parameters())
526
+
527
+ if isinstance(weights, dict):
528
+ weights_list = list(weights.items())
529
+ else:
530
+ weights_list = list(weights)
531
+
532
+ for name, weight in weights_list:
533
+ if not name.startswith("radio_model."):
534
+ # Skip non-radio weights
535
+ continue
536
+
537
+ sub = name[len("radio_model."):] # drop "radio_model." prefix
538
+
539
+ # Skip buffers not used in vLLM
540
+ if sub in {"summary_idxs"}:
541
+ continue
542
+
543
+ vllm_key = None
544
+ if sub.startswith("model.patch_generator."):
545
+ vllm_key = f"model.patch_generator.{sub.split('.', 2)[-1]}"
546
+ elif sub.startswith("input_conditioner."):
547
+ vllm_key = f"input_conditioner.{sub.split('.', 1)[-1]}"
548
+ elif sub.startswith("model.blocks."):
549
+ # Encoder blocks: HF 'model.blocks.{i}.' ->
550
+ # vLLM 'model.encoder.layers.{i}.'
551
+ parts = sub.split(".")
552
+ if len(parts) >= 4:
553
+ layer_idx = parts[2]
554
+ suffix = ".".join(parts[3:])
555
+ # Skip layer-scale entries that vLLM doesn't use
556
+ if suffix in {"ls1", "ls2"} or suffix.startswith(
557
+ ("ls1.", "ls2.")):
558
+ continue
559
+ vllm_key = f"model.encoder.layers.{layer_idx}.{suffix}"
560
+
561
+ if vllm_key and vllm_key in params_dict:
562
+ param = params_dict[vllm_key]
563
+ weight_loader = getattr(param, "weight_loader",
564
+ default_weight_loader)
565
+ weight_loader(param, weight)
566
+ loaded_params.add(vllm_key)
567
+
568
+ return loaded_params
569
+
570
+ def _extract_final(self, y: torch.Tensor):
571
+ # Remove CLS + REGISTERS tokens
572
+ patch_gen = getattr(self.model, "patch_generator", None)
573
+ if patch_gen is not None:
574
+ all_feat = y[:, patch_gen.num_skip:]
575
+
576
+ return all_feat
vllm/transformers_utils/configs/__init__.py
+2
@@ -26,6 +26,7 @@ from vllm.transformers_utils.configs.nemotron_vl import Nemotron_Nano_VL_Config
26
from vllm.transformers_utils.configs.olmo3 import Olmo3Config
27
from vllm.transformers_utils.configs.ovis import OvisConfig
28
from vllm.transformers_utils.configs.qwen3_next import Qwen3NextConfig
29
+from vllm.transformers_utils.configs.radio import RadioConfig
30
from vllm.transformers_utils.configs.speculators.base import SpeculatorsConfig
31
from vllm.transformers_utils.configs.step3_vl import (Step3TextConfig,
32
Step3VisionEncoderConfig,
@@ -48,6 +49,7 @@ __all__ = [
49
"Nemotron_Nano_VL_Config",
50
"Olmo3Config",
51
"OvisConfig",
52
+ "RadioConfig",
53
"SpeculatorsConfig",
54
"UltravoxConfig",
55
"Step3VLConfig",
vllm/transformers_utils/configs/radio.py
new
+104
@@ -0,0 +1,104 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Radio vision model configuration"""
4
+
5
+from typing import Optional, Union
6
+
7
+from transformers.configuration_utils import PretrainedConfig
8
+from transformers.utils import logging
9
+
10
+logger = logging.get_logger(__name__)
11
+
12
+VIT_TIMM_DIM_BY_NAME: dict[str, tuple[int, int, int, int]] = {
13
+ "vit_small_patch16_224": (384, 12, 6, 1536),
14
+ "vit_base_patch16_224": (768, 12, 12, 3072),
15
+ "vit_large_patch16_224": (1024, 24, 16, 4096),
16
+ "vit_huge_patch16_224": (1280, 32, 16, 5120),
17
+}
18
+
19
+OPENAI_CLIP_MEAN = (0.48145466, 0.4578275, 0.40821073)
20
+OPENAI_CLIP_STD = (0.26862954, 0.26130258, 0.27577711)
21
+
22
+
23
+class RadioConfig(PretrainedConfig):
24
+ r"""
25
+ This is the configuration class to store the configuration of a Radio
26
+ vision model. It is used to instantiate a Radio model according to the
27
+ specified arguments, defining the model architecture.
28
+
29
+ Args:
30
+ model_name (`str`, *optional*, defaults to "vit_base_patch16_224"):
31
+ Name of the vision transformer model (e.g., "vit_base_patch16_224").
32
+ Used to determine architecture dimensions from
33
+ `VIT_TIMM_DIM_BY_NAME`.
34
+ image_size (`int`, *optional*, defaults to 224):
35
+ The size (resolution) of each image.
36
+ patch_size (`int`, *optional*, defaults to 16):
37
+ The size (resolution) of each patch.
38
+ qkv_bias (`bool`, *optional*, defaults to True):
39
+ Whether to add a bias to the queries, keys and values.
40
+ qk_normalization (`bool`, *optional*, defaults to False):
41
+ Whether to apply normalization to queries and keys.
42
+ norm_type (`str`, *optional*, defaults to "layer_norm"):
43
+ The normalization type to use.
44
+ layer_norm_eps (`float`, *optional*, defaults to 1e-6):
45
+ The epsilon used by the layer normalization layers.
46
+ initializer_factor (`float`, *optional*, defaults to 1.0):
47
+ A factor for initializing all weight matrices.
48
+ hidden_act (`str`, *optional*, defaults to "gelu"):
49
+ The non-linear activation function in the encoder.
50
+ max_img_size (`int`, *optional*, defaults to 2048):
51
+ Maximum image size for position embeddings.
52
+ norm_mean (`tuple` or `list`, *optional*,
53
+ defaults to (0.48145466, 0.4578275, 0.40821073)):
54
+ Mean values for image normalization (RGB channels).
55
+ norm_std (`tuple` or `list`, *optional*,
56
+ defaults to (0.26862954, 0.26130258, 0.27577711)):
57
+ Standard deviation values for image normalization (RGB channels).
58
+ reg_tokens (`int`, *optional*):
59
+ Number of register tokens to use.
60
+ """
61
+
62
+ model_type = "radio"
63
+
64
+ def __init__(
65
+ self,
66
+ model_name: str,
67
+ image_size: int = 224,
68
+ patch_size: int = 16,
69
+ qkv_bias: bool = True,
70
+ qk_normalization: bool = False,
71
+ norm_type: str = "layer_norm",
72
+ layer_norm_eps: float = 1e-6,
73
+ initializer_factor: float = 1.0,
74
+ hidden_act: str = "gelu",
75
+ max_img_size: int = 2048,
76
+ norm_mean: Union[tuple[float, float, float], list] = OPENAI_CLIP_MEAN,
77
+ norm_std: Union[tuple[float, float, float], list] = OPENAI_CLIP_STD,
78
+ reg_tokens: Optional[int] = None,
79
+ **kwargs,
80
+ ):
81
+ self.model_name = model_name
82
+ (
83
+ self.hidden_size,
84
+ self.num_hidden_layers,
85
+ self.num_attention_heads,
86
+ self.intermediate_size,
87
+ ) = VIT_TIMM_DIM_BY_NAME[model_name]
88
+ self.image_size = image_size
89
+ self.patch_size = patch_size
90
+ self.qkv_bias = qkv_bias
91
+ self.qk_normalization = qk_normalization
92
+ self.norm_type = norm_type
93
+ self.layer_norm_eps = layer_norm_eps
94
+ self.initializer_factor = initializer_factor
95
+ self.hidden_act = hidden_act
96
+ self.max_img_size = max_img_size
97
+ self.norm_mean = list(norm_mean) if isinstance(norm_mean,
98
+ (tuple,
99
+ list)) else norm_mean
100
+ self.norm_std = list(norm_std) if isinstance(norm_std,
101
+ (tuple,
102
+ list)) else norm_std
103
+ self.reg_tokens = reg_tokens
104
+ super().__init__(**kwargs)