`vllm.models.deepseek_v4.common.vision` ¶

DeepSeek-V4 vision tower (ViT + aligner) with TP-sharded linears.

Ported from the official reference implementation (deepseek-ai/DeepSeek-V4-Flash-Vision-Exp). Weight names match the HF checkpoint so no renaming is needed at load time. Attention and MLP weights are tensor-parallel sharded (replicated when the vision head count is not divisible by TP size, or under --mm-encoder-tp-mode data); the patch embed and norms are replicated, so the residual stream is full-width on every rank.

Classes:

DeepseekV4Aligner –

Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.
DeepseekV4ViT –

DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE.

Functions:

run_dp_sharded_vision_tower –

Run the ViT + aligner with images sharded across TP ranks.

`DeepseekV4Aligner` ¶

Bases: Module

Spatial merge (downsample_ratio x downsample_ratio) + MLP projector.

Source code in vllm/models/deepseek_v4/common/vision.py

class DeepseekV4Aligner(nn.Module):
    """Spatial merge (downsample_ratio x downsample_ratio) + MLP projector."""

    def __init__(self, config):
        super().__init__()
        use_data_parallel = is_vit_use_data_parallel(config.vision_n_heads)
        self.downsample_ratio = config.vision_downsample_ratio
        self.out_dim = config.hidden_size
        in_dim = config.vision_dim * self.downsample_ratio**2
        self.w1 = ColumnParallelLinear(
            in_dim,
            config.hidden_size,
            bias=True,
            quant_config=None,
            disable_tp=use_data_parallel,
        )
        self.w2 = RowParallelLinear(
            config.hidden_size,
            config.hidden_size,
            bias=True,
            quant_config=None,
            disable_tp=use_data_parallel,
        )

    def forward(self, x: torch.Tensor, n_vit_h: int, n_vit_w: int) -> torch.Tensor:
        r = self.downsample_ratio
        x = x.view(n_vit_h, n_vit_w, -1).permute(2, 0, 1)
        x = F.pad(x, (0, -n_vit_w % r, 0, -n_vit_h % r))
        x = F.unfold(x.unsqueeze(0), r, stride=r).squeeze(0).transpose(0, 1)
        hidden, _ = self.w1(x)
        out, _ = self.w2(F.gelu(hidden))
        return out

`DeepseekV4ViT` ¶

Bases: Module

DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE.

Source code in vllm/models/deepseek_v4/common/vision.py

class DeepseekV4ViT(nn.Module):
    """DeepSeek-V4 ViT: full bidirectional attention per image, 2D RoPE."""

    def __init__(self, config):
        super().__init__()
        self.rope_dim = config.vision_dim // config.vision_n_heads // 2
        self.rope_theta = config.vision_rope_theta
        self.patch_embed = DeepseekV4PatchEmbed(config)
        self.blocks = nn.ModuleList(
            [
                DeepseekV4VisionBlock(config, prefix=f"blocks.{i}")
                for i in range(config.vision_n_layers)
            ]
        )
        self.norm = DeepseekV4RMSNorm(config.vision_dim)

    def forward(
        self, patches: torch.Tensor, n_vit_h: int, n_vit_w: int
    ) -> torch.Tensor:
        x = self.patch_embed(patches)
        cos, sin = get_vision_cos_sin(n_vit_h, n_vit_w, self.rope_dim, self.rope_theta)
        cos = cos.to(device=x.device)
        sin = sin.to(device=x.device)
        for block in self.blocks:
            x = block(x, cos, sin)
        return self.norm(x)

`run_dp_sharded_vision_tower(vision_model, aligner, patches, vit_grid)` ¶

Run the ViT + aligner with images sharded across TP ranks.

Every rank holds the full tower weights (--mm-encoder-tp-mode data) and receives the full patches batch. Images are assigned to ranks by patch count (greedy load balancing), each rank encodes only its share, and per-image embeddings are exchanged with one padded all-gather and returned in the original image order.

Parameters:

vision_model ¶
(DeepseekV4ViT) –

The (weight-replicated) ViT tower.
aligner ¶
(DeepseekV4Aligner) –

The (weight-replicated) spatial-merge projector.
patches ¶
(Tensor) –

(sum(n_vit_h * n_vit_w), 3, p, p) patches of all images.
vit_grid ¶
(list[list[int]]) –

[n_vit_h, n_vit_w] per image.

Returns:

list[Tensor] –

One (n_aligner_rows, hidden_size) embedding tensor per image.

Source code in vllm/models/deepseek_v4/common/vision.py

def run_dp_sharded_vision_tower(
    vision_model: DeepseekV4ViT,
    aligner: DeepseekV4Aligner,
    patches: torch.Tensor,
    vit_grid: list[list[int]],
) -> list[torch.Tensor]:
    """Run the ViT + aligner with images sharded across TP ranks.

    Every rank holds the full tower weights (``--mm-encoder-tp-mode data``)
    and receives the full ``patches`` batch. Images are assigned to ranks by
    patch count (greedy load balancing), each rank encodes only its share,
    and per-image embeddings are exchanged with one padded all-gather and
    returned in the original image order.

    Args:
        vision_model: The (weight-replicated) ViT tower.
        aligner: The (weight-replicated) spatial-merge projector.
        patches: ``(sum(n_vit_h * n_vit_w), 3, p, p)`` patches of all images.
        vit_grid: ``[n_vit_h, n_vit_w]`` per image.

    Returns:
        One ``(n_aligner_rows, hidden_size)`` embedding tensor per image.
    """
    tp_size = get_tensor_model_parallel_world_size()
    tp_rank = get_tensor_model_parallel_rank()

    sizes = [h * w for h, w in vit_grid]
    cum_patches = [0, *itertools.accumulate(sizes)]
    image_to_tp_rank, gpu_sample_counts, _ = get_load_balance_assignment(sizes, tp_size)
    cum_sample_counts = [0, *itertools.accumulate(gpu_sample_counts)]

    # Rows the aligner emits per image: each grid dim is padded up to a
    # multiple of the merge ratio before r x r patches fold into one row.
    r = aligner.downsample_ratio
    rows_per_image = [-(-h // r) * -(-w // r) for h, w in vit_grid]

    def rank_image_idxs(g: int) -> list[int]:
        return image_to_tp_rank[cum_sample_counts[g] : cum_sample_counts[g + 1]]

    # All-gather needs one shape on every rank; pad each rank's packed
    # output to the largest per-rank row count (computable locally).
    max_rows = max(
        sum(rows_per_image[i] for i in rank_image_idxs(g)) for g in range(tp_size)
    )

    local_embeds = [
        aligner(
            vision_model(
                patches[cum_patches[i] : cum_patches[i + 1]],
                vit_grid[i][0],
                vit_grid[i][1],
            ),
            vit_grid[i][0],
            vit_grid[i][1],
        )
        for i in rank_image_idxs(tp_rank)
    ]
    if local_embeds:
        embeds_local = torch.cat(local_embeds, dim=0)
    else:
        embeds_local = patches.new_zeros((0, aligner.out_dim))
    if embeds_local.shape[0] < max_rows:
        embeds_local = torch.cat(
            [
                embeds_local,
                embeds_local.new_zeros(
                    (max_rows - embeds_local.shape[0], embeds_local.shape[1])
                ),
            ],
            dim=0,
        )
    gathered = tensor_model_parallel_all_gather(embeds_local.contiguous(), dim=0)

    out: list[torch.Tensor] = [None] * len(vit_grid)  # type: ignore[list-item]
    for g in range(tp_size):
        offset = g * max_rows
        for i in rank_image_idxs(g):
            out[i] = gathered[offset : offset + rows_per_image[i]]
            offset += rows_per_image[i]
    return out

`vllm.models.deepseek_v4.common.vision` ¶

`DeepseekV4Aligner` ¶

`DeepseekV4ViT` ¶

`run_dp_sharded_vision_tower(vision_model, aligner, patches, vit_grid)` ¶

`vision_model` ¶

`aligner` ¶

`patches` ¶

`vit_grid` ¶

vllm.models.deepseek_v4.common.vision ¶

DeepseekV4Aligner ¶

DeepseekV4ViT ¶

run_dp_sharded_vision_tower(vision_model, aligner, patches, vit_grid) ¶

vision_model ¶

aligner ¶

patches ¶

vit_grid ¶

`vllm.models.deepseek_v4.common.vision` ¶

`DeepseekV4Aligner` ¶

`DeepseekV4ViT` ¶

`run_dp_sharded_vision_tower(vision_model, aligner, patches, vit_grid)` ¶

`vision_model` ¶

`aligner` ¶

`patches` ¶

`vit_grid` ¶