vllm.multimodal.video_prune.vidcom2 ¶

def _multi_scale_gaussian(x: torch.Tensor, center: torch.Tensor) -> torch.Tensor:
    """Sum Gaussian kernels over `_ALPHAS`; `(T, N, C) -> (T, N)` scores."""
    dist_sq = ((x - center) ** 2).sum(dim=-1)
    return sum(torch.exp(-dist_sq / (2.0 * a)) for a in _ALPHAS)

def compute_retained_tokens_count(
    tokens_per_frame: int, num_frames: int, q: float
) -> int:
    """Number of video tokens retained after VidCom2 pruning.

    The target is `(1 - q) * total_tokens`, i.e. a retention ratio of
    `1 - q` averaged across frames. Because the per-frame budget is floored
    at one token, the global minimum is `num_frames` (one token per frame).
    """
    total_tokens = tokens_per_frame * num_frames
    base_num = int(total_tokens * (1.0 - q))
    return max(num_frames, min(base_num, total_tokens))

def compute_retention_mask(
    video_embeds: torch.Tensor,
    video_size_thw: torch.LongTensor | tuple[int, int, int],
    spatial_merge_size: int,
    q: float,
) -> torch.Tensor:
    """Compute the VidCom2 retention mask for a single video.

    Args:
        video_embeds: `(T*H*W/merge^2, hidden_size)` post-ViT token features.
        video_size_thw: `(T, H, W)` grid dimensions.
        spatial_merge_size: ViT spatial merge factor (e.g. 2).
        q: Pruning fraction in `[0, 1)`; retention ratio is `1 - q`.

    Returns:
        Flat bool tensor of shape `(T*H*W/merge^2,)`, True for retained
        tokens. The True count equals `compute_retained_tokens_count` so
        placeholders sized at prompt-processing time match exactly.
    """
    T, H, W = map(int, video_size_thw)
    rows = H // spatial_merge_size
    cols = W // spatial_merge_size
    tokens_per_frame = rows * cols
    total_tokens = T * tokens_per_frame

    device = video_embeds.device
    if tokens_per_frame == 0 or total_tokens == 0:
        return torch.ones(0, dtype=torch.bool, device=device)

    target_retained = compute_retained_tokens_count(
        tokens_per_frame=tokens_per_frame, num_frames=T, q=q
    )
    target_retained = min(target_retained, total_tokens)

    # 1. Score in the lowest-variance half of channels.
    variances = video_embeds.var(dim=0, unbiased=False)
    k_channels = max(1, int(video_embeds.size(-1) * _LOW_VAR_CHANNEL_RATIO))
    _, low_var_idx = torch.topk(variances, k=k_channels, largest=False)
    sel = video_embeds.index_select(-1, low_var_idx)

    # 2. Multi-scale Gaussian similarity to video and per-frame centers.
    frames = sel.view(T, tokens_per_frame, sel.size(-1))
    frames = F.normalize(frames, dim=-1)
    vid_center = frames.mean(dim=(0, 1), keepdim=True)  # (1, 1, C)
    frame_center = frames.mean(dim=1, keepdim=True)  # (T, 1, C)
    v_score = _multi_scale_gaussian(frames, vid_center)
    f_score = _multi_scale_gaussian(frames, frame_center)
    # Higher similarity = more redundant; lowest-similarity tokens are kept.
    similarity = v_score + f_score  # (T, tpf)

    # 3. Per-frame dynamic budget: distinctive frames get a larger share.
    base = 1.0 - q
    frame_scores = -v_score.mean(dim=-1)  # (T,)
    probs = F.softmax((frame_scores - frame_scores.max()) / _SOFTMAX_TEMPERATURE, dim=0)
    scales = (base * (1.0 + probs - probs.mean())).clamp(max=1.0)
    ks = (scales * tokens_per_frame).round().long().clamp(min=1, max=tokens_per_frame)

    # 4. Retain the smallest-similarity tokens per frame.
    mask_2d = torch.zeros(T, tokens_per_frame, dtype=torch.bool, device=device)
    for i in range(T):
        k_i = int(ks[i].item())
        if k_i <= 0:
            continue
        _, idx = torch.topk(similarity[i], k=k_i, largest=False, sorted=False)
        mask_2d[i].scatter_(0, idx, True)

    # 5. Reconcile rounding/clamp drift to the exact target count by score.
    flat_mask = mask_2d.view(-1)
    flat_sim = similarity.view(-1)
    current = int(flat_mask.sum().item())
    if current > target_retained:
        drop_n = current - target_retained
        retained_idx = flat_mask.nonzero(as_tuple=False).squeeze(-1)
        retained_sim = flat_sim[retained_idx]
        _, worst = torch.topk(retained_sim, k=drop_n, largest=True, sorted=False)
        flat_mask[retained_idx[worst]] = False
    elif current < target_retained:
        add_n = target_retained - current
        available_idx = (~flat_mask).nonzero(as_tuple=False).squeeze(-1)
        if available_idx.numel() > 0:
            available_sim = flat_sim[available_idx]
            add_n = min(add_n, available_idx.numel())
            _, best = torch.topk(available_sim, k=add_n, largest=False, sorted=False)
            flat_mask[available_idx[best]] = True

    return flat_mask