| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """DeepSeek-V4 vision tower (ViT + aligner) with TP-sharded linears. |
| 4 | |
| 5 | Ported from the official reference implementation |
| 6 | (deepseek-ai/DeepSeek-V4-Flash-Vision-Exp). Weight names match the HF |
| 7 | checkpoint so no renaming is needed at load time. Attention and MLP weights |
| 8 | are tensor-parallel sharded (replicated when the vision head count is not |
| 9 | divisible by TP size, or under ``--mm-encoder-tp-mode data``); the patch |
| 10 | embed and norms are replicated, so the residual stream is full-width on |
| 11 | every rank. |
| 12 | """ |
| 13 | |
| 14 | import itertools |
| 15 | from functools import lru_cache |
| 16 | |
| 17 | import torch |
| 18 | import torch.nn.functional as F |
| 19 | from torch import nn |
| 20 | |
| 21 | from vllm.distributed import ( |
| 22 | get_tensor_model_parallel_rank, |
| 23 | get_tensor_model_parallel_world_size, |
| 24 | tensor_model_parallel_all_gather, |
| 25 | ) |
| 26 | from vllm.model_executor.layers.activation import SiluAndMul |
| 27 | from vllm.model_executor.layers.attention.mm_encoder_attention import ( |
| 28 | MMEncoderAttention, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.linear import ( |
| 31 | ColumnParallelLinear, |
| 32 | MergedColumnParallelLinear, |
| 33 | QKVParallelLinear, |
| 34 | ReplicatedLinear, |
| 35 | RowParallelLinear, |
| 36 | ) |
| 37 | from vllm.model_executor.models.vision import ( |
| 38 | get_load_balance_assignment, |
| 39 | is_vit_use_data_parallel, |
| 40 | ) |
| 41 | from vllm.utils.torch_utils import async_tensor_h2d |
| 42 | |
| 43 | |
| 44 | def _compute_vision_cos_sin( |
| 45 | n_h: int, n_w: int, dim: int, theta: float |
| 46 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 47 | inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) |
| 48 | hpos = torch.arange(n_h).unsqueeze(1).expand(n_h, n_w) |
| 49 | wpos = torch.arange(n_w).unsqueeze(0).expand(n_h, n_w) |
| 50 | freqs = torch.stack([hpos, wpos], dim=-1).reshape(-1, 2, 1).float() |
| 51 | freqs = (freqs * inv_freq).flatten(1) |
| 52 | return freqs.cos().unsqueeze(1), freqs.sin().unsqueeze(1) |
| 53 | |
| 54 | |
| 55 | @lru_cache(8) |
| 56 | def get_vision_cos_sin( |
| 57 | n_h: int, n_w: int, dim: int, theta: float |
| 58 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 59 | return _compute_vision_cos_sin(n_h, n_w, dim, theta) |
| 60 | |
| 61 | |
| 62 | def apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: |
| 63 | dtype = x.dtype |
| 64 | x1, x2 = x.float().chunk(2, dim=-1) |
| 65 | return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1).to(dtype) |
| 66 | |
| 67 | |
| 68 | class DeepseekV4RMSNorm(nn.Module): |
| 69 | def __init__(self, dim: int, eps: float = 1e-6): |
| 70 | super().__init__() |
| 71 | self.eps = eps |
| 72 | self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32)) |
| 73 | |
| 74 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 75 | dtype = x.dtype |
| 76 | x = x.float() |
| 77 | x = x * torch.rsqrt(x.square().mean(-1, keepdim=True) + self.eps) |
| 78 | return (self.weight * x).to(dtype) |
| 79 | |
| 80 | |
| 81 | class DeepseekV4PatchEmbed(nn.Module): |
| 82 | def __init__(self, config): |
| 83 | super().__init__() |
| 84 | # Replicated: the residual stream is full-width on every rank. |
| 85 | self.proj = ReplicatedLinear( |
| 86 | 3 * config.vision_patch_size**2, |
| 87 | config.vision_dim, |
| 88 | bias=True, |
| 89 | quant_config=None, |
| 90 | ) |
| 91 | |
| 92 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 93 | out, _ = self.proj(x.flatten(1)) |
| 94 | return out |
| 95 | |
| 96 | |
| 97 | class DeepseekV4VisionAttention(nn.Module): |
| 98 | def __init__(self, config, prefix: str = ""): |
| 99 | super().__init__() |
| 100 | # Convention from minimax_m3/siglip: compute the TP/DP choice locally |
| 101 | # instead of threading a flag through every constructor. |
| 102 | use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) |
| 103 | self.tp_size = ( |
| 104 | 1 if use_data_parallel else get_tensor_model_parallel_world_size() |
| 105 | ) |
| 106 | self.n_heads = config.vision_n_heads // self.tp_size |
| 107 | self.head_dim = config.vision_dim // config.vision_n_heads |
| 108 | self.hidden_size = config.vision_dim |
| 109 | self.wqkv = QKVParallelLinear( |
| 110 | config.vision_dim, |
| 111 | self.head_dim, |
| 112 | config.vision_n_heads, |
| 113 | bias=True, |
| 114 | quant_config=None, |
| 115 | prefix=f"{prefix}.wqkv", |
| 116 | disable_tp=use_data_parallel, |
| 117 | ) |
| 118 | self.wo = RowParallelLinear( |
| 119 | config.vision_dim, |
| 120 | config.vision_dim, |
| 121 | bias=True, |
| 122 | quant_config=None, |
| 123 | prefix=f"{prefix}.wo", |
| 124 | disable_tp=use_data_parallel, |
| 125 | ) |
| 126 | self.attn = MMEncoderAttention( |
| 127 | num_heads=self.n_heads, |
| 128 | head_size=self.head_dim, |
| 129 | prefix=f"{prefix}.attn", |
| 130 | ) |
| 131 | |
| 132 | def forward( |
| 133 | self, |
| 134 | x: torch.Tensor, |
| 135 | cos: torch.Tensor, |
| 136 | sin: torch.Tensor, |
| 137 | cu_seqlens: torch.Tensor | None = None, |
| 138 | max_seqlen: torch.Tensor | None = None, |
| 139 | ) -> torch.Tensor: |
| 140 | n = x.size(0) |
| 141 | qkv, _ = self.wqkv(x) |
| 142 | q, k, v = (t.view(n, self.n_heads, self.head_dim) for t in qkv.chunk(3, -1)) |
| 143 | q = apply_rotary(q, cos, sin).unsqueeze(0) # (b=1, n, h, d) |
| 144 | k = apply_rotary(k, cos, sin).unsqueeze(0) |
| 145 | # Dense (cu_seqlens=None): one image per call. Varlen (cu_seqlens |
| 146 | # set): packed multi-image batch for encoder CUDA graph replay. |
| 147 | o = self.attn( |
| 148 | q, k, v.unsqueeze(0), cu_seqlens=cu_seqlens, max_seqlen=max_seqlen |
| 149 | ) |
| 150 | out, _ = self.wo(o.reshape(n, -1)) |
| 151 | return out |
| 152 | |
| 153 | |
| 154 | class DeepseekV4VisionMLP(nn.Module): |
| 155 | def __init__(self, config, prefix: str = ""): |
| 156 | super().__init__() |
| 157 | use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) |
| 158 | self.w1 = MergedColumnParallelLinear( |
| 159 | config.vision_dim, |
| 160 | [config.vision_inter_dim] * 2, |
| 161 | bias=False, |
| 162 | quant_config=None, |
| 163 | prefix=f"{prefix}.w1", |
| 164 | disable_tp=use_data_parallel, |
| 165 | ) |
| 166 | self.w2 = RowParallelLinear( |
| 167 | config.vision_inter_dim, |
| 168 | config.vision_dim, |
| 169 | bias=False, |
| 170 | quant_config=None, |
| 171 | prefix=f"{prefix}.w2", |
| 172 | disable_tp=use_data_parallel, |
| 173 | ) |
| 174 | self.act_fn = SiluAndMul() |
| 175 | |
| 176 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 177 | gate_up, _ = self.w1(x) |
| 178 | out, _ = self.w2(self.act_fn(gate_up)) |
| 179 | return out |
| 180 | |
| 181 | |
| 182 | class DeepseekV4VisionBlock(nn.Module): |
| 183 | def __init__(self, config, prefix: str = ""): |
| 184 | super().__init__() |
| 185 | self.norm1 = DeepseekV4RMSNorm(config.vision_dim) |
| 186 | self.attn = DeepseekV4VisionAttention(config, prefix=f"{prefix}.attn") |
| 187 | self.norm2 = DeepseekV4RMSNorm(config.vision_dim) |
| 188 | self.mlp = DeepseekV4VisionMLP(config, prefix=f"{prefix}.mlp") |
| 189 | |
| 190 | def forward( |
| 191 | self, |
| 192 | x: torch.Tensor, |
| 193 | cos: torch.Tensor, |
| 194 | sin: torch.Tensor, |
| 195 | cu_seqlens: torch.Tensor | None = None, |
| 196 | max_seqlen: torch.Tensor | None = None, |
| 197 | ) -> torch.Tensor: |
| 198 | x = x + self.attn(self.norm1(x), cos, sin, cu_seqlens, max_seqlen) |
| 199 | return x + self.mlp(self.norm2(x)) |
| 200 | |
| 201 | |
| 202 | class DeepseekV4ViT(nn.Module): |
| 203 | """DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE.""" |
| 204 | |
| 205 | def __init__(self, config): |
| 206 | super().__init__() |
| 207 | self.rope_dim = config.vision_dim // config.vision_n_heads // 2 |
| 208 | self.rope_theta = config.vision_rope_theta |
| 209 | self.patch_embed = DeepseekV4PatchEmbed(config) |
| 210 | self.blocks = nn.ModuleList( |
| 211 | [ |
| 212 | DeepseekV4VisionBlock(config, prefix=f"blocks.{i}") |
| 213 | for i in range(config.vision_n_layers) |
| 214 | ] |
| 215 | ) |
| 216 | self.norm = DeepseekV4RMSNorm(config.vision_dim) |
| 217 | |
| 218 | def forward( |
| 219 | self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int |
| 220 | ) -> torch.Tensor: |
| 221 | x = self.patch_embed(patches) |
| 222 | cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta) |
| 223 | cos = async_tensor_h2d(cos, x.device) |
| 224 | sin = async_tensor_h2d(sin, x.device) |
| 225 | for block in self.blocks: |
| 226 | x = block(x, cos, sin) |
| 227 | return self.norm(x) |
| 228 | |
| 229 | def forward_packed( |
| 230 | self, |
| 231 | patches: torch.Tensor, |
| 232 | cos: torch.Tensor, |
| 233 | sin: torch.Tensor, |
| 234 | cu_seqlens: torch.Tensor, |
| 235 | max_seqlen: torch.Tensor, |
| 236 | ) -> torch.Tensor: |
| 237 | """Varlen path for encoder CUDA graphs: multiple images packed along |
| 238 | rows, per-image attention via ``cu_seqlens``, RoPE tables precomputed |
| 239 | on device. Pure-tensor (no host reads, no H2D), safe to capture.""" |
| 240 | x = self.patch_embed(patches) |
| 241 | for block in self.blocks: |
| 242 | x = block(x, cos, sin, cu_seqlens, max_seqlen) |
| 243 | return self.norm(x) |
| 244 | |
| 245 | |
| 246 | def build_packed_vit_metadata( |
| 247 | grids: list[list[int]] | list[tuple[int, int]], |
| 248 | *, |
| 249 | rope_dim: int, |
| 250 | rope_theta: float, |
| 251 | device: torch.device, |
| 252 | max_seqlen_override: int | None = None, |
| 253 | cached: bool = True, |
| 254 | ) -> dict[str, torch.Tensor]: |
| 255 | """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``. |
| 256 | |
| 257 | Args: |
| 258 | grids: ``[n_vit_h, n_vit_w]`` per image, in packing order. |
| 259 | rope_dim: RoPE embedding dimension of the ViT. |
| 260 | rope_theta: RoPE base frequency of the ViT. |
| 261 | device: Device for the returned tensors (except ``max_seqlen``, |
| 262 | which stays on the host). |
| 263 | max_seqlen_override: Worst-case value baked in at CUDA graph capture |
| 264 | (the attention wrapper reads ``max_seqlen`` on the host, so the |
| 265 | capture-time value becomes a graph constant). |
| 266 | cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time |
| 267 | dummy grids pass ``False`` to avoid evicting real entries. |
| 268 | |
| 269 | """ |
| 270 | cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin |
| 271 | cos_list: list[torch.Tensor] = [] |
| 272 | sin_list: list[torch.Tensor] = [] |
| 273 | cu_seqlens = [0] |
| 274 | max_seqlen = 0 |
| 275 | for n_h, n_w in grids: |
| 276 | cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta) |
| 277 | cos_list.append(cos) |
| 278 | sin_list.append(sin) |
| 279 | n = n_h * n_w |
| 280 | cu_seqlens.append(cu_seqlens[-1] + n) |
| 281 | max_seqlen = max(max_seqlen, n) |
| 282 | if cos_list: |
| 283 | cos = async_tensor_h2d(torch.cat(cos_list), device) |
| 284 | sin = async_tensor_h2d(torch.cat(sin_list), device) |
| 285 | else: |
| 286 | cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) |
| 287 | sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device) |
| 288 | if max_seqlen_override is not None: |
| 289 | max_seqlen = max_seqlen_override |
| 290 | return { |
| 291 | "vit_cos": cos, |
| 292 | "vit_sin": sin, |
| 293 | "cu_seqlens": async_tensor_h2d(cu_seqlens, device, dtype=torch.int32), |
| 294 | # Read on the host by the attention wrapper; keep on CPU. |
| 295 | "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32), |
| 296 | } |
| 297 | |
| 298 | |
| 299 | def build_packed_merge_metadata( |
| 300 | grids: list[list[int]] | list[tuple[int, int]], |
| 301 | downsample_ratio: int, |
| 302 | *, |
| 303 | device: torch.device, |
| 304 | dtype: torch.dtype, |
| 305 | ) -> dict[str, torch.Tensor]: |
| 306 | """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``. |
| 307 | |
| 308 | For each aligner output row, the ``r x r`` source patch rows in the |
| 309 | packed ViT output; positions past the image edge (the eager path pads |
| 310 | them with zeros) get mask 0 and a clamped in-range index. |
| 311 | """ |
| 312 | r = downsample_ratio |
| 313 | idx_list: list[torch.Tensor] = [] |
| 314 | mask_list: list[torch.Tensor] = [] |
| 315 | offset = 0 |
| 316 | for n_h, n_w in grids: |
| 317 | rows = -(-n_h // r) |
| 318 | cols = -(-n_w // r) |
| 319 | bi = torch.arange(rows).repeat_interleave(cols) |
| 320 | bj = torch.arange(cols).repeat(rows) |
| 321 | di = torch.arange(r).repeat_interleave(r) |
| 322 | dj = torch.arange(r).repeat(r) |
| 323 | hi = bi[:, None] * r + di[None, :] |
| 324 | wj = bj[:, None] * r + dj[None, :] |
| 325 | valid = (hi < n_h) & (wj < n_w) |
| 326 | idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1) |
| 327 | idx_list.append(idx) |
| 328 | mask_list.append(valid) |
| 329 | offset += n_h * n_w |
| 330 | if idx_list: |
| 331 | merge_idx = async_tensor_h2d(torch.cat(idx_list), device) |
| 332 | merge_mask = async_tensor_h2d( |
| 333 | torch.cat(mask_list).unsqueeze(-1), device, dtype=dtype |
| 334 | ) |
| 335 | else: |
| 336 | merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device) |
| 337 | merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device) |
| 338 | return {"merge_idx": merge_idx, "merge_mask": merge_mask} |
| 339 | |
| 340 | |
| 341 | class DeepseekV4Aligner(nn.Module): |
| 342 | """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.""" |
| 343 | |
| 344 | def __init__(self, config): |
| 345 | super().__init__() |
| 346 | use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads) |
| 347 | self.downsample_ratio = config.vision_downsample_ratio |
| 348 | self.out_dim = config.hidden_size |
| 349 | in_dim = config.vision_dim * self.downsample_ratio**2 |
| 350 | self.w1 = ColumnParallelLinear( |
| 351 | in_dim, |
| 352 | config.hidden_size, |
| 353 | bias=True, |
| 354 | quant_config=None, |
| 355 | disable_tp=use_data_parallel, |
| 356 | ) |
| 357 | self.w2 = RowParallelLinear( |
| 358 | config.hidden_size, |
| 359 | config.hidden_size, |
| 360 | bias=True, |
| 361 | quant_config=None, |
| 362 | disable_tp=use_data_parallel, |
| 363 | ) |
| 364 | |
| 365 | def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor: |
| 366 | r = self.downsample_ratio |
| 367 | x = x.view(n_vit_h, n_vit_w, -1).permute(2, 0, 1) |
| 368 | x = F.pad(x, (0, -n_vit_w % r, 0, -n_vit_h % r)) |
| 369 | x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1) |
| 370 | hidden, _ = self.w1(x) |
| 371 | out, _ = self.w2(F.gelu(hidden)) |
| 372 | return out |
| 373 | |
| 374 | def forward_packed( |
| 375 | self, |
| 376 | x: torch.Tensor, |
| 377 | merge_idx: torch.Tensor, |
| 378 | merge_mask: torch.Tensor, |
| 379 | ) -> torch.Tensor: |
| 380 | """Packed multi-image path for encoder CUDA graphs. |
| 381 | |
| 382 | ``merge_idx``/``merge_mask`` come from ``build_packed_merge_metadata``; |
| 383 | gather+mask reproduces the eager ``F.pad`` + ``F.unfold`` merge |
| 384 | exactly (unfold rows are channel-major, so the gathered patch rows |
| 385 | are transposed before flattening). |
| 386 | """ |
| 387 | r2 = self.downsample_ratio**2 |
| 388 | m = merge_idx.shape[0] |
| 389 | gathered = (x[merge_idx] * merge_mask).view(m, r2, -1) |
| 390 | gathered = gathered.permute(0, 2, 1).reshape(m, -1) |
| 391 | hidden, _ = self.w1(gathered) |
| 392 | out, _ = self.w2(F.gelu(hidden)) |
| 393 | return out |
| 394 | |
| 395 | |
| 396 | def run_dp_sharded_vision_tower( |
| 397 | vision_model: DeepseekV4ViT, |
| 398 | aligner: DeepseekV4Aligner, |
| 399 | patches: torch.Tensor, |
| 400 | vit_grid: list[list[int]], |
| 401 | ) -> list[torch.Tensor]: |
| 402 | """Run the ViT + aligner with images sharded across TP ranks. |
| 403 | |
| 404 | Every rank holds the full tower weights (``--mm-encoder-tp-mode data``) |
| 405 | and receives the full ``patches`` batch. Images are assigned to ranks by |
| 406 | patch count (greedy load balancing), each rank encodes only its share, |
| 407 | and per-image embeddings are exchanged with one padded all-gather and |
| 408 | returned in the original image order. |
| 409 | |
| 410 | Args: |
| 411 | vision_model: The (weight-replicated) ViT tower. |
| 412 | aligner: The (weight-replicated) spatial-merge projector. |
| 413 | patches: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` patches of all images. |
| 414 | vit_grid: ``[n_vit_h, n_vit_w]`` per image. |
| 415 | |
| 416 | Returns: |
| 417 | One ``(n_aligner_rows, hidden_size)`` embedding tensor per image. |
| 418 | |
| 419 | """ |
| 420 | tp_size = get_tensor_model_parallel_world_size() |
| 421 | tp_rank = get_tensor_model_parallel_rank() |
| 422 | |
| 423 | sizes = [h * w for h, w in vit_grid] |
| 424 | cum_patches = [0, *itertools.accumulate(sizes)] |
| 425 | image_to_tp_rank, gpu_sample_counts, _ = get_load_balance_assignment(sizes, tp_size) |
| 426 | cum_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)] |
| 427 | |
| 428 | # Rows the aligner emits per image: each grid dim is padded up to a |
| 429 | # multiple of the merge ratio before r x r patches fold into one row. |
| 430 | r = aligner.downsample_ratio |
| 431 | rows_per_image = [-(-h // r) * -(-w // r) for h, w in vit_grid] |
| 432 | |
| 433 | def rank_image_idxs(g: int) -> list[int]: |
| 434 | return image_to_tp_rank[cum_sample_counts[g] : cum_sample_counts[g + 1]] |
| 435 | |
| 436 | # All-gather needs one shape on every rank; pad each rank's packed |
| 437 | # output to the largest per-rank row count (computable locally). |
| 438 | max_rows = max( |
| 439 | sum(rows_per_image[i] for i in rank_image_idxs(g)) for g in range(tp_size) |
| 440 | ) |
| 441 | |
| 442 | local_embeds = [ |
| 443 | aligner( |
| 444 | vision_model( |
| 445 | patches[cum_patches[i] : cum_patches[i + 1]], |
| 446 | vit_grid[i][0], |
| 447 | vit_grid[i][1], |
| 448 | ), |
| 449 | vit_grid[i][0], |
| 450 | vit_grid[i][1], |
| 451 | ) |
| 452 | for i in rank_image_idxs(tp_rank) |
| 453 | ] |
| 454 | if local_embeds: |
| 455 | embeds_local = torch.cat(local_embeds, dim=0) |
| 456 | else: |
| 457 | embeds_local = patches.new_zeros((0, aligner.out_dim)) |
| 458 | if embeds_local.shape[0] < max_rows: |
| 459 | embeds_local = torch.cat( |
| 460 | [ |
| 461 | embeds_local, |
| 462 | embeds_local.new_zeros( |
| 463 | (max_rows - embeds_local.shape[0], embeds_local.shape[1]) |
| 464 | ), |
| 465 | ], |
| 466 | dim=0, |
| 467 | ) |
| 468 | gathered = tensor_model_parallel_all_gather(embeds_local.contiguous(), dim=0) |
| 469 | |
| 470 | out: list[torch.Tensor] = [None] * len(vit_grid) # type: ignore[list-item] |
| 471 | for g in range(tp_size): |
| 472 | offset = g * max_rows |
| 473 | for i in rank_image_idxs(g): |
| 474 | out[i] = gathered[offset : offset + rows_per_image[i]] |
| 475 | offset += rows_per_image[i] |
| 476 | return out |