@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 3127274d02

Fork 0

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

[MM Encoder] Apply DP ViT for Qwen3-VL model series (#24955)

Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Huang Jie <92386084+JJJYmmm@users.noreply.github.com> Co-authored-by: 松灵 <26085463+wulipc@users.noreply.github.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Roger Wang committed Sep 17, 2025 at 21:04 UTC 3127274d022b0bc8ff6ba9ceef41a99a6f01ad2d

2 files changed +77 -19

vllm/model_executor/models/qwen3_vl.py

+75 -19

@@ -126,20 +126,23 @@ class Qwen3_VisionMLP(nn.Module):

126 bias: bool = False,

127 act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

128 quant_config: Optional[QuantizationConfig] = None,

129 - prefix: str = ""):

129 + prefix: str = "",

130 + use_data_parallel: bool = False):

131 super().__init__()

132 self.linear_fc1 = ColumnParallelLinear(in_features,

133 hidden_features,

134 bias=bias,

135 quant_config=quant_config,

136 return_bias=False,

136 - prefix=f"{prefix}.linear_fc1")

137 + prefix=f"{prefix}.linear_fc1",

138 + disable_tp=use_data_parallel)

139 self.linear_fc2 = RowParallelLinear(hidden_features,

140 in_features,

141 bias=bias,

142 quant_config=quant_config,

143 return_bias=False,

142 - prefix=f"{prefix}.linear_fc2")

144 + prefix=f"{prefix}.linear_fc2",

145 + disable_tp=use_data_parallel)

146 self.act_fn = act_fn

147

148 def forward(self, x: torch.Tensor):

@@ -158,23 +161,27 @@ class Qwen3_VisionBlock(nn.Module):

161 norm_layer: Optional[Callable[[int], nn.Module]] = None,

162 quant_config: Optional[QuantizationConfig] = None,

163 prefix: str = "",

164 + use_data_parallel: bool = False,

165 ) -> None:

166 super().__init__()

167 if norm_layer is None:

168 norm_layer = partial(nn.LayerNorm, eps=1e-6)

169 self.norm1 = norm_layer(dim)

170 self.norm2 = norm_layer(dim)

167 - self.attn = Qwen2_5_VisionAttention(embed_dim=dim,

168 - num_heads=num_heads,

169 - projection_size=dim,

170 - quant_config=quant_config,

171 - prefix=f"{prefix}.attn")

171 + self.attn = Qwen2_5_VisionAttention(

172 + embed_dim=dim,

173 + num_heads=num_heads,

174 + projection_size=dim,

175 + quant_config=quant_config,

176 + prefix=f"{prefix}.attn",

177 + use_data_parallel=use_data_parallel)

178 self.mlp = Qwen3_VisionMLP(dim,

179 mlp_hidden_dim,

180 act_fn=act_fn,

181 bias=True,

182 quant_config=quant_config,

177 - prefix=f"{prefix}.mlp")

183 + prefix=f"{prefix}.mlp",

184 + use_data_parallel=use_data_parallel)

185

186 def forward(

187 self,

@@ -205,6 +212,7 @@ class Qwen3_VisionPatchMerger(nn.Module):

212 use_postshuffle_norm: bool = False,

213 quant_config: Optional[QuantizationConfig] = None,

214 prefix: str = "",

215 + use_data_parallel: bool = False,

216 ) -> None:

217 super().__init__()

218 self.hidden_size = context_dim * (spatial_merge_size**2)

@@ -222,13 +230,15 @@ class Qwen3_VisionPatchMerger(nn.Module):

230 self.hidden_size,

231 bias=True,

232 quant_config=quant_config,

225 - prefix=f"{prefix}.linear_fc1")

233 + prefix=f"{prefix}.linear_fc1",

234 + disable_tp=use_data_parallel)

235 self.act_fn = nn.GELU()

236 self.linear_fc2 = RowParallelLinear(self.hidden_size,

237 d_model,

238 bias=True,

239 quant_config=quant_config,

231 - prefix=f"{prefix}.linear_fc2")

240 + prefix=f"{prefix}.linear_fc2",

241 + disable_tp=use_data_parallel)

242

243 def forward(self, x: torch.Tensor) -> torch.Tensor:

244 if self.use_postshuffle_norm:

@@ -250,6 +260,7 @@ class Qwen3_VisionTransformer(nn.Module):

260 norm_eps: float = 1e-6,

261 quant_config: Optional[QuantizationConfig] = None,

262 prefix: str = "",

263 + use_data_parallel: bool = False,

264 ) -> None:

265 super().__init__()

266 self.hidden_size = vision_config.hidden_size

@@ -260,6 +271,12 @@ class Qwen3_VisionTransformer(nn.Module):

271 self.spatial_merge_unit = self.spatial_merge_size**2

272 self.temporal_patch_size = vision_config.temporal_patch_size

273 self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes

274 + self.use_data_parallel = use_data_parallel

275 +

276 + # NOTE: This is used for creating empty tensor for all_gather for

277 + # DP ViT. Here out_hidden_size is enlarged due to deepstack

278 + self.out_hidden_size = (vision_config.out_hidden_size *

279 + (1 + len(self.deepstack_visual_indexes)))

280

281 self.patch_embed = Qwen3_VisionPatchEmbed(

282 patch_size=self.patch_size,

@@ -283,7 +300,8 @@ class Qwen3_VisionTransformer(nn.Module):

300 act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],

301 norm_layer=norm_layer,

302 quant_config=quant_config,

286 - prefix=f"{prefix}.blocks.{layer_idx}")

303 + prefix=f"{prefix}.blocks.{layer_idx}",

304 + use_data_parallel=use_data_parallel)

305 for layer_idx in range(vision_config.depth)

306 ])

307

@@ -294,6 +312,7 @@ class Qwen3_VisionTransformer(nn.Module):

312 spatial_merge_size=self.spatial_merge_size,

313 quant_config=quant_config,

314 prefix=f"{prefix}.merger",

315 + use_data_parallel=use_data_parallel,

316 )

317

318 self.deepstack_merger_list = nn.ModuleList([

@@ -304,7 +323,8 @@ class Qwen3_VisionTransformer(nn.Module):

323 use_postshuffle_norm=True,

324 norm_layer=norm_layer,

325 quant_config=quant_config,

307 - prefix=f"{prefix}.deepstack_merger_list.{layer_idx}")

326 + prefix=f"{prefix}.deepstack_merger_list.{layer_idx}",

327 + use_data_parallel=use_data_parallel)

328 for layer_idx in range(len(self.deepstack_visual_indexes))

329 ])

330

@@ -325,7 +345,14 @@ class Qwen3_VisionTransformer(nn.Module):

345

346 def rot_pos_emb(self, grid_thw):

347 pos_ids = []

328 - for t, h, w in grid_thw:

348 + # Support both Tensor and list inputs for DP path

349 + if isinstance(grid_thw, list):

350 + grid_list = grid_thw

351 + max_grid_size = max(max(h, w) for _, h, w in grid_list)

352 + else:

353 + grid_list = grid_thw.tolist()

354 + max_grid_size = int(grid_thw[:, 1:].max().item())

355 + for t, h, w in grid_list:

356 hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)

357 hpos_ids = hpos_ids.reshape(

358 h // self.spatial_merge_size,

@@ -348,7 +375,6 @@ class Qwen3_VisionTransformer(nn.Module):

375 pos_ids.append(

376 torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

377 pos_ids = torch.cat(pos_ids, dim=0)

351 - max_grid_size = grid_thw[:, 1:].max()

378 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

379 rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

380 return rotary_pos_emb

@@ -453,10 +479,18 @@ class Qwen3_VisionTransformer(nn.Module):

479 hidden_states = hidden_states + pos_embeds

480 rotary_pos_emb = self.rot_pos_emb(grid_thw)

481

482 + if isinstance(grid_thw, list):

483 + grid_thw_tensor = torch.tensor(grid_thw,

484 + device=hidden_states.device,

485 + dtype=torch.int32)

486 + else:

487 + grid_thw_tensor = grid_thw

488 +

489 cu_seqlens = torch.repeat_interleave(

457 - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum(

490 + grid_thw_tensor[:, 1] * grid_thw_tensor[:, 2],

491 + grid_thw_tensor[:, 0]).cumsum(

492 dim=0,

459 - dtype=grid_thw.dtype

493 + dtype=grid_thw_tensor.dtype

494 if torch.jit.is_tracing() else torch.int32,

495 )

496 cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

@@ -984,6 +1018,9 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1018 "up_proj",

1019 ],

1020 }

1021 +

1022 + supports_encoder_tp_data = True

1023 +

1024 # To ensure correct weight loading and mapping.

1025 hf_to_vllm_mapper = WeightsMapper(

1026 orig_to_new_prefix={

@@ -1009,12 +1046,14 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1046

1047 self.config = config

1048 self.multimodal_config = multimodal_config

1049 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

1050

1051 self.visual = Qwen3_VisionTransformer(

1052 config.vision_config,

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

1054 quant_config=self._maybe_ignore_quant_config(quant_config),

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

1056 + use_data_parallel=self.use_data_parallel,

1057 )

1058

1059 self.language_model = Qwen3LLMForCausalLM(vllm_config=vllm_config,

@@ -1177,7 +1216,15 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1216 image_embeds = image_input["image_embeds"].type(self.visual.dtype)

1217 else:

1218 pixel_values = image_input["pixel_values"].type(self.visual.dtype)

1180 - image_embeds = self.visual(pixel_values, grid_thw=grid_thw)

1219 + if self.use_data_parallel:

1220 + from vllm.multimodal.utils import (

1221 + run_dp_sharded_mrope_vision_model)

1222 + return run_dp_sharded_mrope_vision_model(self.visual,

1223 + pixel_values,

1224 + grid_thw_list,

1225 + rope_type="rope_3d")

1226 + else:

1227 + image_embeds = self.visual(pixel_values, grid_thw=grid_thw)

1228

1229 # Split concatenated embeddings for each image item.

1230 # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync

@@ -1199,7 +1246,16 @@ class Qwen3VLForConditionalGeneration(nn.Module, SupportsMultiModal,

1246 else:

1247 pixel_values_videos = video_input["pixel_values_videos"].type(

1248 self.visual.dtype)

1202 - video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)

1249 + if self.use_data_parallel:

1250 + from vllm.multimodal.utils import (

1251 + run_dp_sharded_mrope_vision_model)

1252 + return run_dp_sharded_mrope_vision_model(self.visual,

1253 + pixel_values_videos,

1254 + grid_thw_list,

1255 + rope_type="rope_3d")

1256 + else:

1257 + video_embeds = self.visual(pixel_values_videos,

1258 + grid_thw=grid_thw)

1259

1260 # Split concatenated embeddings for each video item.

1261 # Using prod on grid_thw_list instead of grid_thw.prod avoids CUDA sync

vllm/model_executor/models/qwen3_vl_moe.py

@@ -315,12 +315,14 @@ class Qwen3VLMoeForConditionalGeneration(Qwen3VLForConditionalGeneration):

315

316 self.config = config

317 self.multimodal_config = multimodal_config

318 + self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

319

320 self.visual = Qwen3_VisionTransformer(

321 config.vision_config,

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

323 quant_config=self._maybe_ignore_quant_config(quant_config),

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

325 + use_data_parallel=self.use_data_parallel,

326 )

327

328 self.language_model = Qwen3MoeLLMForCausalLM(vllm_config=vllm_config,