| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling Titan vision + audio towers. |
| 4 | |
| 5 | The vision tower (``InklingVision`` / ``HMLPPatchEncoder``) emits one token per |
| 6 | image patch; the audio tower (``InklingAudio``) emits one token per audio frame. |
| 7 | Both use vLLM's standard ``RMSNorm`` (CPU-friendly, with a native fallback). |
| 8 | """ |
| 9 | |
| 10 | from __future__ import annotations |
| 11 | |
| 12 | from itertools import combinations |
| 13 | from typing import cast |
| 14 | |
| 15 | import numpy as np |
| 16 | import torch |
| 17 | import torch.nn.functional as F |
| 18 | from torch import nn |
| 19 | |
| 20 | from vllm.model_executor.layers.layernorm import RMSNorm |
| 21 | |
| 22 | from ..configs import InklingAudioConfig, InklingVisionConfig |
| 23 | |
| 24 | # =========================================================================== |
| 25 | # Vision tower (HMLPPatchEncoder / InklingVision) |
| 26 | # =========================================================================== |
| 27 | |
| 28 | |
| 29 | def _prime_factors(n: int) -> list[int]: |
| 30 | """Return the prime factors of *n* in ascending order.""" |
| 31 | if n < 1: |
| 32 | raise ValueError("n must be a positive integer") |
| 33 | |
| 34 | factors: list[int] = [] |
| 35 | while n % 2 == 0: |
| 36 | factors.append(2) |
| 37 | n //= 2 |
| 38 | p = 3 |
| 39 | while p * p <= n: |
| 40 | while n % p == 0: |
| 41 | factors.append(p) |
| 42 | n //= p |
| 43 | p += 2 |
| 44 | if n > 1: |
| 45 | factors.append(n) |
| 46 | return factors |
| 47 | |
| 48 | |
| 49 | def linear_sum_assignment( |
| 50 | cost_matrix: np.ndarray, |
| 51 | ) -> tuple[np.ndarray, np.ndarray]: |
| 52 | """Implement SciPy's assignment for Inkling's ordered L1 cost matrix.""" |
| 53 | rows = np.arange(cost_matrix.shape[0]) |
| 54 | cols = np.array( |
| 55 | min( |
| 56 | combinations(range(cost_matrix.shape[1]), len(rows)), |
| 57 | key=lambda candidate: cost_matrix[rows, candidate].sum(), |
| 58 | ) |
| 59 | ) |
| 60 | return rows, cols |
| 61 | |
| 62 | |
| 63 | def plan_out_scales( |
| 64 | temporal_patch_size: int, patch_size: int, n_layers: int, n_channels: int = 3 |
| 65 | ) -> list[tuple[int, int, int, int]]: |
| 66 | """Plan the (t, h, w, c) dimensions for each HMLP layer. |
| 67 | |
| 68 | Spatial dims expand first, then temporal; channel counts round up to |
| 69 | multiples of 64. |
| 70 | """ |
| 71 | if patch_size <= 1: |
| 72 | raise ValueError( |
| 73 | "patch_size must be greater than 1, otherwise this doesn't make sense" |
| 74 | ) |
| 75 | |
| 76 | def _round_up(x: int) -> int: |
| 77 | return int(np.ceil(x / 64)) * 64 |
| 78 | |
| 79 | last_h_scale = 1 |
| 80 | scales: list[tuple[int, int, int, int]] = [(1, 1, 1, n_channels)] |
| 81 | for pscale in _prime_factors(patch_size)[::-1]: |
| 82 | last_h_scale *= pscale |
| 83 | scales.append( |
| 84 | ( |
| 85 | 1, |
| 86 | last_h_scale, |
| 87 | last_h_scale, |
| 88 | _round_up((last_h_scale**2) * n_channels), |
| 89 | ) |
| 90 | ) |
| 91 | last_t_scale = 1 |
| 92 | for tscale in _prime_factors(temporal_patch_size)[::-1]: |
| 93 | last_t_scale *= tscale |
| 94 | scales.append( |
| 95 | ( |
| 96 | last_t_scale, |
| 97 | last_h_scale, |
| 98 | last_h_scale, |
| 99 | _round_up((last_h_scale**2) * n_channels * last_t_scale), |
| 100 | ) |
| 101 | ) |
| 102 | |
| 103 | size_reduction = np.prod(np.array(scales)[:, :-1], 1) |
| 104 | |
| 105 | log_ideal_scales = np.linspace( |
| 106 | 0, |
| 107 | np.log(patch_size * patch_size * temporal_patch_size * n_channels), |
| 108 | n_layers + 1, |
| 109 | ) |
| 110 | cost_matrix = np.abs(log_ideal_scales[:, None] - np.log(size_reduction)[None]) |
| 111 | |
| 112 | if n_layers >= len(scales): |
| 113 | idxs = np.argmin(cost_matrix, axis=1) |
| 114 | else: |
| 115 | idxs = linear_sum_assignment(cost_matrix)[1] |
| 116 | |
| 117 | assert len(idxs) >= 2 |
| 118 | idxs[0] = 0 |
| 119 | idxs[-1] = len(scales) - 1 |
| 120 | |
| 121 | return [scales[i] for i in idxs] |
| 122 | |
| 123 | |
| 124 | def fold_timespace_to_depth( |
| 125 | vision_patches_bthwc: torch.Tensor, t_fold: int, hw_fold: int |
| 126 | ) -> torch.Tensor: |
| 127 | """(B, T, H, W, C) -> (B, T//t, H//hw, W//hw, C*(t*hw**2)).""" |
| 128 | B, T, H, W, C = vision_patches_bthwc.shape |
| 129 | |
| 130 | assert T % t_fold == 0, f"Temporal dimension {T} must be divisible by {t_fold}" |
| 131 | assert H % hw_fold == 0, f"Height dimension {H} must be divisible by {hw_fold}" |
| 132 | assert W % hw_fold == 0, f"Width dimension {W} must be divisible by {hw_fold}" |
| 133 | |
| 134 | t_new = T // t_fold |
| 135 | h_new = H // hw_fold |
| 136 | w_new = W // hw_fold |
| 137 | |
| 138 | x = vision_patches_bthwc.reshape( |
| 139 | B, t_new, t_fold, h_new, hw_fold, w_new, hw_fold, C |
| 140 | ) |
| 141 | x = x.permute(0, 1, 3, 5, 2, 4, 6, 7) |
| 142 | x = x.reshape(B, t_new, h_new, w_new, t_fold * hw_fold * hw_fold * C) |
| 143 | return x |
| 144 | |
| 145 | |
| 146 | class HMLPPatchEncoder(nn.Module): |
| 147 | def __init__(self, config: InklingVisionConfig): |
| 148 | super().__init__() |
| 149 | self.decoder_dmodel = config.decoder_dmodel |
| 150 | self.patch_size = config.patch_size |
| 151 | self.temporal_patch_size = config.temporal_patch_size |
| 152 | self.n_channels = config.n_channels |
| 153 | self.n_layers = config.n_layers |
| 154 | self.use_vision_norm = config.use_vision_norm |
| 155 | |
| 156 | self.scales: list[tuple[int, int, int, int]] = plan_out_scales( |
| 157 | self.temporal_patch_size, self.patch_size, self.n_layers, self.n_channels |
| 158 | ) |
| 159 | self.layers: nn.ModuleDict = nn.ModuleDict() |
| 160 | for i, (start_scale, end_scale) in enumerate( |
| 161 | zip(self.scales[:-1], self.scales[1:]) |
| 162 | ): |
| 163 | shuffle_mult = ( |
| 164 | (end_scale[0] // start_scale[0]) |
| 165 | * (end_scale[1] // start_scale[1]) |
| 166 | * (end_scale[2] // start_scale[2]) |
| 167 | ) |
| 168 | if i == self.n_layers - 1: |
| 169 | self.layers[f"linear_{i}"] = nn.Linear( |
| 170 | start_scale[3] * shuffle_mult, self.decoder_dmodel, bias=False |
| 171 | ) |
| 172 | else: |
| 173 | self.layers[f"linear_{i}"] = nn.Linear( |
| 174 | start_scale[3] * shuffle_mult, end_scale[3], bias=False |
| 175 | ) |
| 176 | self.layers[f"norm_{i}"] = RMSNorm(end_scale[3]) |
| 177 | |
| 178 | self.final_norm: RMSNorm | None = None |
| 179 | if self.use_vision_norm: |
| 180 | assert self.decoder_dmodel is not None |
| 181 | self.final_norm = RMSNorm(self.decoder_dmodel) |
| 182 | |
| 183 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 184 | # Fused norm+gelu on CUDA (same fp32-accum/bf16-rounding structure as |
| 185 | # the generic path below; differs only by reduction order). |
| 186 | fused = None |
| 187 | if x.is_cuda and x.dtype == torch.bfloat16: |
| 188 | from vllm.models.inkling.nvidia.ops.mm_towers import rmsnorm_gelu |
| 189 | |
| 190 | fused = rmsnorm_gelu |
| 191 | |
| 192 | num_patches, T, H, W, C = x.shape |
| 193 | prefolded = False |
| 194 | for i, (start_scale, end_scale) in enumerate( |
| 195 | zip(self.scales[:-1], self.scales[1:]) |
| 196 | ): |
| 197 | t_fold = end_scale[0] // start_scale[0] |
| 198 | hw_fold = end_scale[1] // start_scale[1] |
| 199 | if (hw_fold > 1 or t_fold > 1) and not prefolded: |
| 200 | x = fold_timespace_to_depth(x, t_fold, hw_fold) |
| 201 | prefolded = False |
| 202 | assert x.shape[1:-1] == ( |
| 203 | T // end_scale[0], |
| 204 | H // end_scale[1], |
| 205 | W // end_scale[2], |
| 206 | ) |
| 207 | x = self.layers[f"linear_{i}"](x) |
| 208 | if i < self.n_layers - 1: |
| 209 | norm = cast(RMSNorm, self.layers[f"norm_{i}"]) |
| 210 | if fused is not None: |
| 211 | # If the NEXT layer starts with a copying fold (spatial |
| 212 | # dims still > 1 after folding), store this layer's |
| 213 | # output directly in the folded layout instead. |
| 214 | nxt = self.scales[i + 2] |
| 215 | ntf = nxt[0] // end_scale[0] |
| 216 | nhf = nxt[1] // end_scale[1] |
| 217 | copy_fold = (ntf > 1 or nhf > 1) and ( |
| 218 | x.shape[2] // nhf > 1 or x.shape[3] // nhf > 1 |
| 219 | ) |
| 220 | x = fused( |
| 221 | x, |
| 222 | norm.weight, |
| 223 | norm.variance_epsilon, |
| 224 | gelu=True, |
| 225 | fold=(ntf, nhf) if copy_fold else None, |
| 226 | ) |
| 227 | prefolded = copy_fold |
| 228 | else: |
| 229 | # rms_norm kernel only supports rank 2-4; x is 5-D here. |
| 230 | orig = x.shape |
| 231 | x = norm(x.reshape(-1, x.shape[-1])).reshape(orig) |
| 232 | x = F.gelu(x) |
| 233 | |
| 234 | if self.final_norm is not None: |
| 235 | if fused is not None: |
| 236 | x = fused( |
| 237 | x, |
| 238 | self.final_norm.weight, |
| 239 | self.final_norm.variance_epsilon, |
| 240 | gelu=False, |
| 241 | ) |
| 242 | else: |
| 243 | orig = x.shape |
| 244 | x = self.final_norm(x.reshape(-1, x.shape[-1])).reshape(orig) |
| 245 | |
| 246 | x = x.reshape(num_patches, -1) |
| 247 | return x |
| 248 | |
| 249 | |
| 250 | class InklingVision(nn.Module): |
| 251 | def __init__(self, config: InklingVisionConfig, prefix: str = ""): |
| 252 | del prefix |
| 253 | super().__init__() |
| 254 | assert config.vision_encoder_type == "hmlp" |
| 255 | self.vision_encoder = HMLPPatchEncoder(config) |
| 256 | |
| 257 | @property |
| 258 | def dtype(self) -> torch.dtype: |
| 259 | return next(self.parameters()).dtype |
| 260 | |
| 261 | @property |
| 262 | def device(self) -> torch.device: |
| 263 | return next(self.parameters()).device |
| 264 | |
| 265 | def forward(self, vision_features: torch.Tensor) -> torch.Tensor: |
| 266 | return self.vision_encoder(vision_features) |
| 267 | |
| 268 | |
| 269 | # =========================================================================== |
| 270 | # Audio tower (InklingAudio) |
| 271 | # =========================================================================== |
| 272 | |
| 273 | |
| 274 | class InklingAudio(nn.Module): |
| 275 | def __init__(self, config: InklingAudioConfig, prefix: str = ""): |
| 276 | del prefix |
| 277 | super().__init__() |
| 278 | assert config.audio_mode == "dmel" |
| 279 | self.n_mel_bins = config.n_mel_bins |
| 280 | self.mel_vocab_size = config.mel_vocab_size |
| 281 | self.use_audio_norm = config.use_audio_norm |
| 282 | self.encoder = nn.Embedding( |
| 283 | config.n_mel_bins * config.mel_vocab_size, config.decoder_dmodel |
| 284 | ) |
| 285 | self.final_norm: RMSNorm | None = None |
| 286 | if self.use_audio_norm: |
| 287 | assert config.decoder_dmodel is not None |
| 288 | self.final_norm = RMSNorm(config.decoder_dmodel, eps=1e-6) |
| 289 | |
| 290 | @property |
| 291 | def dtype(self) -> torch.dtype: |
| 292 | return self.encoder.weight.dtype |
| 293 | |
| 294 | @property |
| 295 | def device(self) -> torch.device: |
| 296 | return self.encoder.weight.device |
| 297 | |
| 298 | def forward(self, audio_features: torch.Tensor) -> torch.Tensor: |
| 299 | assert audio_features.shape[1] == self.n_mel_bins |
| 300 | |
| 301 | # dMel bins are integer indices; cast once to int32 on the right device |
| 302 | # (no float round-trip). |
| 303 | audio_features = audio_features.to( |
| 304 | device=self.encoder.weight.device, dtype=torch.int32 |
| 305 | ) |
| 306 | |
| 307 | weight = self.encoder.weight |
| 308 | if audio_features.is_cuda and weight.dtype == torch.bfloat16: |
| 309 | # One kernel: per-bin offset + embedding gather + fp32 sum + norm. |
| 310 | # Skips the [T, n_mel_bins, D] intermediate entirely (bit-exact). |
| 311 | from vllm.models.inkling.nvidia.ops.mm_towers import dmel_embed_sum_norm |
| 312 | |
| 313 | return dmel_embed_sum_norm( |
| 314 | audio_features.contiguous(), |
| 315 | weight, |
| 316 | self.final_norm.weight if self.final_norm is not None else None, |
| 317 | self.final_norm.variance_epsilon if self.final_norm else 0.0, |
| 318 | ) |
| 319 | |
| 320 | embedding_indices = ( |
| 321 | torch.arange(self.n_mel_bins, device=audio_features.device) |
| 322 | * self.mel_vocab_size |
| 323 | ).unsqueeze(0) + audio_features |
| 324 | |
| 325 | hidden_states = ( |
| 326 | self.encoder(embedding_indices.reshape(-1)) |
| 327 | .reshape(audio_features.shape[0], audio_features.shape[1], -1) |
| 328 | .sum(axis=1) |
| 329 | ) |
| 330 | |
| 331 | if self.final_norm is not None: |
| 332 | hidden_states = self.final_norm(hidden_states) |
| 333 | |
| 334 | return hidden_states |