Skip to content

vllm.models.inkling.common.towers

Inkling Titan vision + audio towers.

The vision tower (InklingVision / HMLPPatchEncoder) emits one token per image patch; the audio tower (InklingAudio) emits one token per audio frame. Both use vLLM's standard RMSNorm (CPU-friendly, with a native fallback).

Functions:

_prime_factors(n)

Return the prime factors of n in ascending order.

Source code in vllm/models/inkling/common/towers.py 
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def _prime_factors(n: int) -> list[int]:
    """Return the prime factors of *n* in ascending order."""
    if n < 1:
        raise ValueError("n must be a positive integer")

    factors: list[int] = []
    while n % 2 == 0:
        factors.append(2)
        n //= 2
    p = 3
    while p * p <= n:
        while n % p == 0:
            factors.append(p)
            n //= p
        p += 2
    if n > 1:
        factors.append(n)
    return factors

fold_timespace_to_depth(vision_patches_bthwc, t_fold, hw_fold)

(B, T, H, W, C) -> (B, T//t, H//hw, W//hw, C(thw**2)).

Source code in vllm/models/inkling/common/towers.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
def fold_timespace_to_depth(
    vision_patches_bthwc: torch.Tensor, t_fold: int, hw_fold: int
) -> torch.Tensor:
    """(B, T, H, W, C) -> (B, T//t, H//hw, W//hw, C*(t*hw**2))."""
    B, T, H, W, C = vision_patches_bthwc.shape

    assert T % t_fold == 0, f"Temporal dimension {T} must be divisible by {t_fold}"
    assert H % hw_fold == 0, f"Height dimension {H} must be divisible by {hw_fold}"
    assert W % hw_fold == 0, f"Width dimension {W} must be divisible by {hw_fold}"

    t_new = T // t_fold
    h_new = H // hw_fold
    w_new = W // hw_fold

    x = vision_patches_bthwc.reshape(
        B, t_new, t_fold, h_new, hw_fold, w_new, hw_fold, C
    )
    x = x.permute(0, 1, 3, 5, 2, 4, 6, 7)
    x = x.reshape(B, t_new, h_new, w_new, t_fold * hw_fold * hw_fold * C)
    return x

plan_out_scales(temporal_patch_size, patch_size, n_layers, n_channels=3)

Plan the (t, h, w, c) dimensions for each HMLP layer.

Spatial dims expand first, then temporal; channel counts round up to multiples of 64.

Source code in vllm/models/inkling/common/towers.py 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
def plan_out_scales(
    temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3
) -> list[tuple[int, int, int, int]]:
    """Plan the (t, h, w, c) dimensions for each HMLP layer.

    Spatial dims expand first, then temporal; channel counts round up to
    multiples of 64.
    """
    if patch_size <= 1:
        raise ValueError(
            "patch_size must be greater than 1, otherwise this doesn't make sense"
        )

    def _round_up(x: int) -> int:
        return int(np.ceil(x / 64)) * 64

    last_h_scale = 1
    scales: list[tuple[int, int, int, int]] = [(1, 1, 1, n_channels)]
    for pscale in _prime_factors(patch_size)[::-1]:
        last_h_scale *= pscale
        scales.append(
            (
                1,
                last_h_scale,
                last_h_scale,
                _round_up((last_h_scale**2) * n_channels),
            )
        )
    last_t_scale = 1
    for tscale in _prime_factors(temporal_patch_size)[::-1]:
        last_t_scale *= tscale
        scales.append(
            (
                last_t_scale,
                last_h_scale,
                last_h_scale,
                _round_up((last_h_scale**2) * n_channels * last_t_scale),
            )
        )

    size_reduction = np.prod(np.array(scales)[:, :-1], 1)

    log_ideal_scales = np.linspace(
        0,
        np.log(patch_size * patch_size * temporal_patch_size * n_channels),
        n_layers + 1,
    )
    cost_matrix = np.abs(log_ideal_scales[:, None] - np.log(size_reduction)[None])

    if n_layers >= len(scales):
        idxs = np.argmin(cost_matrix, axis=1)
    else:
        from scipy.optimize import linear_sum_assignment

        idxs = linear_sum_assignment(cost_matrix)[1]

    assert len(idxs) >= 2
    idxs[0] = 0
    idxs[-1] = len(scales) - 1

    return [scales[i] for i in idxs]