vllm.models.deepseek_v4.common.vision ¶

def build_packed_merge_metadata(
    grids: list[list[int]] | list[tuple[int, int]],
    downsample_ratio: int,
    *,
    device: torch.device,
    dtype: torch.dtype,
) -> dict[str, torch.Tensor]:
    """Gather indices/mask for ``DeepseekV4Aligner.forward_packed``.

    For each aligner output row, the ``r x r`` source patch rows in the
    packed ViT output; positions past the image edge (the eager path pads
    them with zeros) get mask 0 and a clamped in-range index.
    """
    r = downsample_ratio
    idx_list: list[torch.Tensor] = []
    mask_list: list[torch.Tensor] = []
    offset = 0
    for n_h, n_w in grids:
        rows = -(-n_h // r)
        cols = -(-n_w // r)
        bi = torch.arange(rows).repeat_interleave(cols)
        bj = torch.arange(cols).repeat(rows)
        di = torch.arange(r).repeat_interleave(r)
        dj = torch.arange(r).repeat(r)
        hi = bi[:, None] * r + di[None, :]
        wj = bj[:, None] * r + dj[None, :]
        valid = (hi < n_h) & (wj < n_w)
        idx = offset + hi.clamp(max=n_h - 1) * n_w + wj.clamp(max=n_w - 1)
        idx_list.append(idx)
        mask_list.append(valid)
        offset += n_h * n_w
    if idx_list:
        merge_idx = torch.cat(idx_list).to(device)
        merge_mask = torch.cat(mask_list).unsqueeze(-1).to(device=device, dtype=dtype)
    else:
        merge_idx = torch.zeros((0, r * r), dtype=torch.int64, device=device)
        merge_mask = torch.zeros((0, r * r, 1), dtype=dtype, device=device)
    return {"merge_idx": merge_idx, "merge_mask": merge_mask}

def build_packed_vit_metadata(
    grids: list[list[int]] | list[tuple[int, int]],
    *,
    rope_dim: int,
    rope_theta: float,
    device: torch.device,
    max_seqlen_override: int | None = None,
    cached: bool = True,
) -> dict[str, torch.Tensor]:
    """Precompute packed-batch ViT metadata for ``DeepseekV4ViT.forward_packed``.

    Args:
        grids: ``[n_vit_h, n_vit_w]`` per image, in packing order.
        rope_dim: RoPE embedding dimension of the ViT.
        rope_theta: RoPE base frequency of the ViT.
        device: Device for the returned tensors (except ``max_seqlen``,
            which stays on the host).
        max_seqlen_override: Worst-case value baked in at CUDA graph capture
            (the attention wrapper reads ``max_seqlen`` on the host, so the
            capture-time value becomes a graph constant).
        cached: Use the shared ``lru_cache`` for RoPE tables. Capture-time
            dummy grids pass ``False`` to avoid evicting real entries.

    """
    cos_sin = get_vision_cos_sin if cached else _compute_vision_cos_sin
    cos_list: list[torch.Tensor] = []
    sin_list: list[torch.Tensor] = []
    cu_seqlens = [0]
    max_seqlen = 0
    for n_h, n_w in grids:
        cos, sin = cos_sin(n_h, n_w, rope_dim, rope_theta)
        cos_list.append(cos)
        sin_list.append(sin)
        n = n_h * n_w
        cu_seqlens.append(cu_seqlens[-1] + n)
        max_seqlen = max(max_seqlen, n)
    if cos_list:
        cos = torch.cat(cos_list).to(device)
        sin = torch.cat(sin_list).to(device)
    else:
        cos = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
        sin = torch.zeros((0, 1, rope_dim), dtype=torch.float32, device=device)
    if max_seqlen_override is not None:
        max_seqlen = max_seqlen_override
    return {
        "vit_cos": cos,
        "vit_sin": sin,
        "cu_seqlens": torch.tensor(cu_seqlens, dtype=torch.int32, device=device),
        # Read on the host by the attention wrapper; keep on CPU.
        "max_seqlen": torch.tensor(max_seqlen, dtype=torch.int32),
    }

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