@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c245d35ff4

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 15:00 UTC

[Model] Add MiMo-V2.5 support (#40967)

Signed-off-by: Jee Jee Li <pandaleefree@gmail.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: Isotr0py <Isotr0py@outlook.com> Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: zjy0516 <zhujiangyun@inferact.ai> Co-authored-by: yasong <yasong.wang@inferact.ai> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: Copilot <copilot@github.com>

Isotr0py committed Apr 27, 2026 at 21:26 UTC c245d35ff467bb3e9a73fcb3c4b02e6c7a3d2964

16 files changed +4737 -5

docs/models/supported_models.md

@@ -439,6 +439,7 @@ th {

440 | `MiMoForCausalLM` | MiMo | `XiaomiMiMo/MiMo-7B-RL`, etc. | ✅︎ | ✅︎ |

443 | `MiniCPMForCausalLM` | MiniCPM | `openbmb/MiniCPM-2B-sft-bf16`, `openbmb/MiniCPM-2B-dpo-bf16`, `openbmb/MiniCPM-S-1B-sft`, etc. | ✅︎ | ✅︎ |

444 | `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ |

@@ -590,6 +591,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

591 | `LlavaNextVideoForConditionalGeneration` | LLaVA-NeXT-Video | T + V | `llava-hf/LLaVA-NeXT-Video-7B-hf`, etc. | | ✅︎ |

592 | `LlavaOnevisionForConditionalGeneration` | LLaVA-Onevision | T + I+ + V+ | `llava-hf/llava-onevision-qwen2-7b-ov-hf`, `llava-hf/llava-onevision-qwen2-0.5b-ov-hf`, etc. | | ✅︎ |

593 | `MiDashengLMModel` | MiDashengLM | T + A+ | `mispeech/midashenglm-7b` | | ✅︎ |

595 | `MiniCPMO` | MiniCPM-O | T + IE+ + VE+ + AE+ | `openbmb/MiniCPM-o-2_6`, etc. | ✅︎ | ✅︎ |

597 | `MiniMaxVL01ForConditionalGeneration` | MiniMax-VL | T + IE+ | `MiniMaxAI/MiniMax-VL-01`, etc. | | ✅︎ |

tests/models/registry.py

+18

@@ -594,6 +594,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

594 "MiMoV2FlashForCausalLM": _HfExamplesInfo(

595 "XiaomiMiMo/MiMo-V2-Flash", trust_remote_code=True

596 ),

597 + "MiMoV2ProForCausalLM": _HfExamplesInfo(

598 + "XiaomiMiMo/MiMo-V2.5-Pro", trust_remote_code=True, is_available_online=False

599 + ),

600 "Dots1ForCausalLM": _HfExamplesInfo("rednote-hilab/dots.llm1.inst"),

601 }

602

@@ -1069,6 +1072,9 @@ _MULTIMODAL_EXAMPLE_MODELS = {

1072 "MiDashengLMModel": _HfExamplesInfo(

1073 "mispeech/midashenglm-7b", trust_remote_code=True

1074 ),

1075 + "MiMoV2OmniForCausalLM": _HfExamplesInfo(

1076 + "XiaomiMiMo/MiMo-V2.5-Omni", trust_remote_code=True, is_available_online=False

1077 + ),

1078 "MiniCPMO": _HfExamplesInfo(

1079 "openbmb/MiniCPM-o-2_6",

1080 trust_remote_code=True,

@@ -1552,6 +1558,18 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {

1558 trust_remote_code=True,

1559 speculative_model="XiaomiMiMo/MiMo-7B-RL",

1560 ),

1561 + "MiMoV2MTPModel": _HfExamplesInfo(

1562 + "XiaomiMiMo/MiMo-V2.5-Pro",

1563 + trust_remote_code=True,

1564 + speculative_model="XiaomiMiMo/MiMo-V2.5-Pro",

1565 + is_available_online=False,

1566 + ),

1567 + "MiMoV2OmniMTPModel": _HfExamplesInfo(

1568 + "XiaomiMiMo/MiMo-V2.5-Omni",

1569 + trust_remote_code=True,

1570 + speculative_model="XiaomiMiMo/MiMo-V2.5-Omni",

1571 + is_available_online=False,

1572 + ),

1573 "NemotronHMTPModel": _HfExamplesInfo(

1574 "nvidia/Nemotron-Super-Placeholder",

1575 speculative_model="nvidia/Nemotron-Super-Placeholder",

vllm/config/model.py

+1 -1

@@ -521,6 +521,7 @@ class ModelConfig:

521 if dict_overrides:

522 self._apply_dict_overrides(hf_config, dict_overrides)

523 self.hf_text_config = get_hf_text_config(self.hf_config)

524 + self.model_arch_config = self.get_model_arch_config()

525 self.attention_chunk_size = getattr(

526 self.hf_text_config, "attention_chunk_size", None

527 )

@@ -528,7 +529,6 @@ class ModelConfig:

529 self.hf_image_processor_config = get_hf_image_processor_config(

530 self.model, hf_token=self.hf_token, revision=self.revision

531 )

531 - self.model_arch_config = self.get_model_arch_config()

532

533 architectures = self.architectures

534 registry = self.registry

vllm/config/speculative.py

+43

@@ -34,6 +34,7 @@ logger = init_logger(__name__)

34 MTPModelTypes = Literal[

35 "deepseek_mtp",

36 "mimo_mtp",

37 + "mimo_v2_mtp",

38 "glm4_moe_mtp",

39 "glm4_moe_lite_mtp",

40 "glm_ocr_mtp",

@@ -332,6 +333,48 @@ class SpeculativeConfig:

333 }

334 )

335

336 + if (arch := hf_config.architectures[0]) in (

337 + "MiMoV2ProForCausalLM",

338 + "MiMoV2OmniForCausalLM",

339 + ):

340 + from vllm.model_executor.models.mimo_v2_mtp import (

341 + _MIMO_V2_PRO_NUM_MTP_LAYERS,

342 + )

343 +

344 + mtp_arch_maps = {

345 + "MiMoV2ProForCausalLM": "MiMoV2MTPModel",

346 + "MiMoV2OmniForCausalLM": "MiMoV2OmniMTPModel",

347 + }

348 +

349 + hf_config.model_type = "mimo_v2_mtp"

350 + # vLLM currently supports only the first MiMo-V2 MTP layer.

351 + n_predict = _MIMO_V2_PRO_NUM_MTP_LAYERS

352 + hf_config.update(

353 + {

354 + "num_hidden_layers": 0,

355 + "n_predict": n_predict,

356 + "num_nextn_predict_layers": n_predict,

357 + "architectures": [mtp_arch_maps[arch]],

358 + }

359 + )

360 +

361 + if hf_config.architectures[0] == "MiMoV2FlashForCausalLM":

362 + from vllm.model_executor.models.mimo_v2_mtp import (

363 + _MIMO_V2_FLASH_NUM_MTP_LAYERS,

364 + )

365 +

366 + hf_config.model_type = "mimo_v2_mtp"

367 + # vLLM currently supports only the first MiMo-V2 MTP layer.

368 + n_predict = _MIMO_V2_FLASH_NUM_MTP_LAYERS

369 + hf_config.update(

370 + {

371 + "num_hidden_layers": 0,

372 + "n_predict": n_predict,

373 + "num_nextn_predict_layers": n_predict,

374 + "architectures": ["MiMoV2MTPModel"],

375 + }

376 + )

377 +

378 if hf_config.architectures[0] == "Glm4MoeForCausalLM":

379 hf_config.model_type = "glm4_moe_mtp"

380 n_predict = getattr(hf_config, "num_nextn_predict_layers", None)

vllm/model_executor/layers/linear.py

+6 -1

@@ -1120,7 +1120,12 @@ class QKVParallelLinear(ColumnParallelLinear):

1120 # Special case for Quantization.

1121 # If quantized, we need to adjust the offset and size to account

1122 # for the packing.

1123 - if (

1123 + if isinstance(param, BlockQuantScaleParameter):

1124 + weight_block_size = getattr(self, "weight_block_size", None)

1125 + shard_size, shard_offset = adjust_block_scale_shard(

1126 + weight_block_size, shard_size, shard_offset

1127 + )

1128 + elif (

1129 isinstance(param, (PackedColumnParameter, PackedvLLMParameter))

1130 and param.packed_dim == param.output_dim

1131 ):

vllm/model_executor/models/mimo_audio.py new

+1389

@@ -0,0 +1,1389 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +"""MiMo audio: tokenizer, encoding utilities, and audio encoder.

4 +

5 +Ported from SGLang's mimo_audio.py.

6 +Audio tokenizer adapted from https://github.com/XiaomiMiMo/MiMo-Audio-Tokenizer.git

7 +"""

8 +

9 +import dataclasses

10 +import json

11 +import logging

12 +import math

13 +import os

14 +import typing as tp

15 +from dataclasses import dataclass

16 +from functools import wraps

17 +

18 +import torch

19 +import torch.distributed as dist

20 +import torch.nn as nn

21 +import torch.nn.functional as F

22 +from einops import rearrange, repeat

23 +from transformers.activations import ACT2FN

24 +from transformers.configuration_utils import PretrainedConfig

25 +from transformers.modeling_utils import PreTrainedModel

26 +from transformers.models.qwen2.configuration_qwen2 import Qwen2Config

27 +from transformers.models.qwen2.modeling_qwen2 import Qwen2Model

28 +

29 +logger = logging.getLogger(__name__)

30 +

31 +

32 +# ---------------------------------------------------------------------------

33 +# Vector quantization (from MiMo-Audio-Tokenizer)

34 +# ---------------------------------------------------------------------------

35 +

36 +

37 +def _vq_default(val: tp.Any, d: tp.Any) -> tp.Any:

38 + return val if val is not None else d

39 +

40 +

41 +def _ema_inplace(moving_avg, new, decay: float):

42 + if dist.is_initialized():

43 + dist.all_reduce(new, op=dist.ReduceOp.SUM)

44 + moving_avg.data.mul_(decay).add_(new, alpha=(1 - decay))

45 +

46 +

47 +def _laplace_smoothing(x, n_categories: int, epsilon: float = 1e-5):

48 + return (x + epsilon) / (x.sum() + n_categories * epsilon)

49 +

50 +

51 +def _uniform_init(*shape: int):

52 + t = torch.empty(shape)

53 + nn.init.kaiming_uniform_(t)

54 + return t

55 +

56 +

57 +def _sample_vectors(samples, num: int):

58 + num_samples, device = samples.shape[0], samples.device

59 +

60 + if num_samples >= num:

61 + indices = torch.randperm(num_samples, device=device)[:num]

62 + else:

63 + indices = torch.randint(0, num_samples, (num,), device=device)

64 +

65 + selected_samples = samples[indices]

66 +

67 + if dist.is_initialized():

68 + dist.broadcast(selected_samples, src=0)

69 +

70 + return selected_samples

71 +

72 +

73 +def _kmeans(samples, num_clusters: int, num_iters: int = 10):

74 + dim, dtype = samples.shape[-1], samples.dtype

75 +

76 + means = _sample_vectors(samples, num_clusters)

77 +

78 + for _ in range(num_iters):

79 + dists = -(

80 + samples.pow(2).sum(1, keepdim=True)

81 + - 2 * samples @ means.t()

82 + + means.t().pow(2).sum(0, keepdim=True)

83 + )

84 +

85 + buckets = dists.max(dim=-1).indices

86 + bins = torch.bincount(buckets, minlength=num_clusters)

87 +

88 + new_means = buckets.new_zeros(num_clusters, dim, dtype=dtype)

89 + new_means = new_means.scatter_add_(

90 + 0, repeat(buckets, "n -> n d", d=dim), samples

91 + )

92 +

93 + if dist.is_initialized():

94 + dist.all_reduce(bins, op=dist.ReduceOp.SUM)

95 + dist.all_reduce(new_means, op=dist.ReduceOp.SUM)

96 +

97 + zero_mask = bins == 0

98 + bins_min_clamped = bins.masked_fill(zero_mask, 1)

99 +

100 + new_means = new_means / bins_min_clamped[..., None]

101 +

102 + means = torch.where(zero_mask[..., None], means, new_means)

103 +

104 + return means, bins

105 +

106 +

107 +def _rotate_half(x):

108 + x1 = x[..., : x.shape[-1] // 2]

109 + x2 = x[..., x.shape[-1] // 2 :]

110 + return torch.cat((-x2, x1), dim=-1)

111 +

112 +

113 +def apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1):

114 + cos = cos.unsqueeze(unsqueeze_dim)

115 + sin = sin.unsqueeze(unsqueeze_dim)

116 + q_embed = (q * cos) + (_rotate_half(q) * sin)

117 + k_embed = (k * cos) + (_rotate_half(k) * sin)

118 + return q_embed, k_embed

119 +

120 +

121 +def _compute_default_rope_parameters(

122 + config=None, device=None, seq_len=None, **rope_kwargs

123 +):

124 + if config is not None and len(rope_kwargs) > 0:

125 + raise ValueError(

126 + "Unexpected arguments: `**rope_kwargs` and `config` are mutually exclusive"

127 + )

128 + if len(rope_kwargs) > 0:

129 + base = rope_kwargs["base"]

130 + dim = rope_kwargs["dim"]

131 + elif config is not None:

132 + base = config.rope_theta

133 + partial_rotary_factor = (

134 + config.partial_rotary_factor

135 + if hasattr(config, "partial_rotary_factor")

136 + else 1.0

137 + )

138 + head_dim = (

139 + getattr(config, "head_dim", None)

140 + or config.hidden_size // config.num_attention_heads

141 + )

142 + dim = int(head_dim * partial_rotary_factor)

143 + attention_factor = 1.0

144 + inv_freq = 1.0 / (

145 + base

146 + ** (

147 + torch.arange(0, dim, 2, dtype=torch.int64).to(

148 + device=device, dtype=torch.float

149 + )

150 + / dim

151 + )

152 + )

153 + return inv_freq, attention_factor

154 +

155 +

156 +_ROPE_INIT_FUNCTIONS = {

157 + "default": _compute_default_rope_parameters,

158 +}

159 +

160 +

161 +def _dynamic_rope_update(rope_forward):

162 + def dynamic_frequency_update(self, position_ids, device):

163 + seq_len = torch.max(position_ids) + 1

164 + if seq_len > self.max_seq_len_cached:

165 + inv_freq, self.attention_scaling = self.rope_init_fn(

166 + self.config, device, seq_len=seq_len

167 + )

168 + self.register_buffer("inv_freq", inv_freq, persistent=False)

169 + self.max_seq_len_cached = seq_len

170 +

171 + if (

172 + seq_len < self.original_max_seq_len

173 + and self.max_seq_len_cached > self.original_max_seq_len

174 + ):

175 + self.original_inv_freq = self.original_inv_freq.to(device)

176 + self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)

177 + self.max_seq_len_cached = self.original_max_seq_len

178 +

179 + @wraps(rope_forward)

180 + def wrapper(self, x, position_ids):

181 + if "dynamic" in self.rope_type:

182 + dynamic_frequency_update(self, position_ids, device=x.device)

183 + return rope_forward(self, x, position_ids)

184 +

185 + return wrapper

186 +

187 +

188 +class AudioRotaryEmbedding(nn.Module):

189 + def __init__(self, base, dim, max_seq_len, rope_type="default", device=None):

190 + super().__init__()

191 + self.max_seq_len = max_seq_len

192 + self.rope_type = rope_type

193 + self.rope_init_fn = _ROPE_INIT_FUNCTIONS[self.rope_type]

194 + inv_freq, self.attention_scaling = self.rope_init_fn(

195 + device=device, base=base, dim=dim

196 + )

197 + self.register_buffer("inv_freq", inv_freq, persistent=False)

198 + self.original_inv_freq = self.inv_freq

199 +

200 + @torch.no_grad()

201 + @_dynamic_rope_update

202 + def forward(self, x, position_ids):

203 + inv_freq_expanded = self.inv_freq[:, None].float().expand(-1, 1).to(x.device)

204 + position_ids_expanded = position_ids[None, :].float()

205 + device_type = (

206 + x.device.type

207 + if isinstance(x.device.type, str) and x.device.type != "mps"

208 + else "cpu"

209 + )

210 + with torch.autocast(device_type=device_type, enabled=False):

211 + freqs = (

212 + inv_freq_expanded.float() @ position_ids_expanded.float()

213 + ).transpose(0, 1)

214 + emb = torch.cat((freqs, freqs), dim=-1)

215 + cos = emb.cos() * self.attention_scaling

216 + sin = emb.sin() * self.attention_scaling

217 + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)

218 +

219 +

220 +class EuclideanCodebook(nn.Module):

221 + def __init__(

222 + self,

223 + dim: int,

224 + codebook_size: int,

225 + kmeans_init: int = False,

226 + kmeans_iters: int = 10,

227 + decay: float = 0.99,

228 + epsilon: float = 1e-5,

229 + threshold_ema_dead_code: int = 2,

230 + ):

231 + super().__init__()

232 + self.decay = decay

233 + init_fn: tp.Callable[..., torch.Tensor] | tp.Any = (

234 + _uniform_init if not kmeans_init else torch.zeros

235 + )

236 + embed = init_fn(codebook_size, dim)

237 +

238 + self.codebook_size = codebook_size

239 + self.kmeans_iters = kmeans_iters

240 + self.epsilon = epsilon

241 + self.threshold_ema_dead_code = threshold_ema_dead_code

242 +

243 + self.register_buffer("inited", torch.Tensor([not kmeans_init]))

244 + self.register_buffer("cluster_size", torch.zeros(codebook_size))

245 + self.register_buffer("embed", embed)

246 + self.register_buffer("embed_avg", embed.clone())

247 +

248 + @torch.jit.ignore

249 + def init_embed_(self, data):

250 + if self.inited:

251 + return

252 +

253 + embed, cluster_size = _kmeans(data, self.codebook_size, self.kmeans_iters)

254 + self.embed.data.copy_(embed)

255 + self.embed_avg.data.copy_(embed.clone())

256 + self.cluster_size.data.copy_(cluster_size)

257 + self.inited.data.copy_(torch.Tensor([True]))

258 +

259 + def replace_(self, samples, mask):

260 + replace_num = mask.sum()

261 + modified_codebook = self.embed.clone()

262 + modified_codebook[mask] = _sample_vectors(samples, replace_num)

263 + self.embed.data.copy_(modified_codebook)

264 +

265 + def expire_codes_(self, batch_samples):

266 + if self.threshold_ema_dead_code == 0:

267 + return

268 +

269 + expired_codes = self.cluster_size < self.threshold_ema_dead_code

270 + if not torch.any(expired_codes):

271 + return

272 +

273 + batch_samples = rearrange(batch_samples, "... d -> (...) d")

274 + self.replace_(batch_samples, mask=expired_codes)

275 +

276 + def preprocess(self, x):

277 + x = rearrange(x, "... d -> (...) d")

278 + return x

279 +

280 + def quantize(self, x):

281 + embed = self.embed.t()

282 + dist_val = -(

283 + x.pow(2).sum(1, keepdim=True)

284 + - 2 * x @ embed

285 + + embed.pow(2).sum(0, keepdim=True)

286 + )

287 + embed_ind = dist_val.max(dim=-1).indices

288 + return embed_ind

289 +

290 + def postprocess_emb(self, embed_ind, shape):

291 + return embed_ind.view(*shape[:-1])

292 +

293 + def dequantize(self, embed_ind):

294 + quantize = F.embedding(embed_ind, self.embed)

295 + return quantize

296 +

297 + def encode(self, x):

298 + shape = x.shape

299 + x = self.preprocess(x)

300 + embed_ind = self.quantize(x)

301 + embed_ind = self.postprocess_emb(embed_ind, shape)

302 + return embed_ind

303 +

304 + def decode(self, embed_ind):

305 + quantize = self.dequantize(embed_ind)

306 + return quantize

307 +

308 + def forward(self, x):

309 + shape, dtype = x.shape, x.dtype

310 + x = self.preprocess(x)

311 +

312 + self.init_embed_(x)

313 +

314 + embed_ind = self.quantize(x)

315 + embed_onehot = F.one_hot(embed_ind, self.codebook_size).type(dtype)

316 + embed_ind = self.postprocess_emb(embed_ind, shape)

317 + quantize = self.dequantize(embed_ind)

318 +

319 + if self.training:

320 + self.expire_codes_(x)

321 + _ema_inplace(self.cluster_size, embed_onehot.sum(0), self.decay)

322 + embed_sum = x.t() @ embed_onehot

323 + _ema_inplace(self.embed_avg, embed_sum.t().contiguous(), self.decay)

324 + cluster_size = (

325 + _laplace_smoothing(self.cluster_size, self.codebook_size, self.epsilon)

326 + * self.cluster_size.sum()

327 + )

328 + embed_normalized = self.embed_avg / cluster_size.unsqueeze(1)

329 + self.embed.data.copy_(embed_normalized)

330 +

331 + return quantize, embed_ind

332 +

333 +

334 +class VectorQuantization(nn.Module):

335 + def __init__(

336 + self,

337 + dim: int,

338 + codebook_size: int,

339 + codebook_dim: int | None = None,

340 + decay: float = 0.99,

341 + epsilon: float = 1e-5,

342 + kmeans_init: bool = True,

343 + kmeans_iters: int = 50,

344 + threshold_ema_dead_code: int = 2,

345 + commitment_weight: float = 1.0,

346 + ):

347 + super().__init__()

348 + _codebook_dim: int = _vq_default(codebook_dim, dim)

349 +

350 + requires_projection = _codebook_dim != dim

351 + self.project_in = (

352 + nn.Linear(dim, _codebook_dim) if requires_projection else nn.Identity()

353 + )

354 + self.project_out = (

355 + nn.Linear(_codebook_dim, dim) if requires_projection else nn.Identity()

356 + )

357 +

358 + self.epsilon = epsilon

359 + self.commitment_weight = commitment_weight

360 +

361 + self._codebook = EuclideanCodebook(

362 + dim=_codebook_dim,

363 + codebook_size=codebook_size,

364 + kmeans_init=kmeans_init,

365 + kmeans_iters=kmeans_iters,

366 + decay=decay,

367 + epsilon=epsilon,

368 + threshold_ema_dead_code=threshold_ema_dead_code,

369 + )

370 + self.codebook_size = codebook_size

371 +

372 + @property

373 + def codebook(self):

374 + return self._codebook.embed

375 +

376 + def encode(self, x):

377 + x = self.project_in(x)

378 + embed_in = self._codebook.encode(x)

379 + return embed_in

380 +

381 + def decode(self, embed_ind):

382 + quantize = self._codebook.decode(embed_ind)

383 + quantize = self.project_out(quantize)

384 + return quantize

385 +

386 + def forward(self, x):

387 + device = x.device

388 + x = self.project_in(x)

389 +

390 + quantize, embed_ind = self._codebook(x)

391 +

392 + if self.training:

393 + quantize = x + (quantize - x).detach()

394 +

395 + loss = torch.tensor([0.0], device=device, requires_grad=self.training)

396 +

397 + quantize = self.project_out(quantize)

398 + return quantize, embed_ind, loss

399 +

400 +

401 +class ResidualVectorQuantization(nn.Module):

402 + def __init__(self, *, num_quantizers, codebook_size, **kwargs):

403 + super().__init__()

404 + if isinstance(codebook_size, int):

405 + codebook_size = [codebook_size] * num_quantizers

406 + elif len(codebook_size) < num_quantizers:

407 + codebook_size += [codebook_size[-1]] * (num_quantizers - len(codebook_size))

408 + self.layers = nn.ModuleList(

409 + [

410 + VectorQuantization(codebook_size=codebook_size[i], **kwargs)

411 + for i in range(num_quantizers)

412 + ]

413 + )

414 +

415 + def forward(self, x, n_q: int | None = None, layers: list | None = None):

416 + quantized_out = 0.0

417 + residual = x

418 +

419 + all_losses = []

420 + all_indices = []

421 + out_quantized = []

422 +

423 + n_q = n_q or len(self.layers)

424 +

425 + for i, layer in enumerate(self.layers[:n_q]):

426 + quantized, indices, loss = layer(residual)

427 + residual = residual - quantized

428 + quantized_out = quantized_out + quantized

429 +

430 + all_indices.append(indices)

431 + all_losses.append(loss)

432 + if layers and i in layers:

433 + out_quantized.append(quantized_out)

434 +

435 + out_losses, out_indices = map(torch.stack, (all_losses, all_indices))

436 + return quantized_out, out_indices, out_losses, out_quantized

437 +

438 + def encode(

439 + self, x: torch.Tensor, n_q: int | None = None, st: int | None = None

440 + ) -> torch.Tensor:

441 + residual = x

442 + all_indices = []

443 + n_q = len(self.layers) if n_q is None else n_q

444 + st = 0 if st is None else st

445 + for layer in self.layers[st:n_q]:

446 + indices = layer.encode(residual)

447 + quantized = layer.decode(indices)

448 + residual = residual - quantized

449 + all_indices.append(indices)

450 + out_indices = torch.stack(all_indices)

451 + return out_indices

452 +

453 + def decode(self, q_indices: torch.Tensor, st: int = 0) -> torch.Tensor:

454 + quantized_out = self.layers[st].decode(q_indices[0])

455 + for i in range(1, len(q_indices)):

456 + layer = self.layers[st + i]

457 + quantized = layer.decode(q_indices[i])

458 + quantized_out = quantized_out + quantized

459 + return quantized_out

460 +

461 +

462 +class ResidualVectorQuantizer(nn.Module):

463 + def __init__(

464 + self,

465 + dimension: int = 256,

466 + n_q: int = 8,

467 + bins: int | list = 1024,

468 + decay: float = 0.99,

469 + kmeans_init: bool = True,

470 + kmeans_iters: int = 50,

471 + threshold_ema_dead_code: int = 2,

472 + ):

473 + super().__init__()

474 + self.n_q = n_q

475 + self.dimension = dimension

476 + self.bins = bins

477 + self.decay = decay

478 + self.kmeans_init = kmeans_init

479 + self.kmeans_iters = kmeans_iters

480 + self.threshold_ema_dead_code = threshold_ema_dead_code

481 + self.vq = ResidualVectorQuantization(

482 + dim=self.dimension,

483 + codebook_size=self.bins,

484 + num_quantizers=self.n_q,

485 + decay=self.decay,

486 + kmeans_init=self.kmeans_init,

487 + kmeans_iters=self.kmeans_iters,

488 + threshold_ema_dead_code=self.threshold_ema_dead_code,

489 + )

490 +

491 + def forward(

492 + self,

493 + x: torch.Tensor,

494 + n_q: int | None = None,

495 + layers: list | None = None,

496 + ):

497 + n_q = n_q if n_q else self.n_q

498 + quantized, codes, commit_loss, quantized_list = self.vq(

499 + x, n_q=n_q, layers=layers

500 + )

501 + return quantized, codes, torch.mean(commit_loss), quantized_list

502 +

503 + def encode(

504 + self, x: torch.Tensor, n_q: int | None = None, st: int | None = None

505 + ) -> torch.Tensor:

506 + n_q = n_q if n_q else self.n_q

507 + st = st or 0

508 + codes = self.vq.encode(x, n_q=n_q, st=st)

509 + return codes

510 +

511 + def decode(self, codes: torch.Tensor, st: int = 0) -> torch.Tensor:

512 + quantized = self.vq.decode(codes, st=st)

513 + return quantized

514 +

515 +

516 +# ---------------------------------------------------------------------------

517 +# Audio tokenizer

518 +# ---------------------------------------------------------------------------

519 +

520 +

521 +class MiMoAudioTokenizerConfig(PretrainedConfig):

522 + model_type = "mimo_audio_tokenizer"

523 +

524 + def __init__(

525 + self,

526 + max_audio_seconds: int = 1800,

527 + stride_size: int = 2,

528 + avg_pooler: int = 1,

529 + d_model: int = 768,

530 + scale_embedding: bool = True,

531 + kernel_size: int = 3,

532 + activation_function: str = "gelu",

533 + encoder_layers: int = 8,

534 + encoder_skip_layer_id: int = None,

535 + encoder_attention_heads: int = 12,

536 + encoder_ffn_dim: int = 3072,

537 + encoder_causal: bool = False,

538 + encoder_attn_window_size: list = None,

539 + decoder_layers: int = 8,

540 + decoder_attention_heads: int = 12,

541 + decoder_ffn_dim: int = 3072,

542 + decoder_kernel_size: int = 3,

543 + decoder_stride_size: int = 2,

544 + decoder_causal: bool = True,

545 + decoder_attn_window_size: list = None,

546 + nfft: int = 1024,

547 + vocoder_dim: int = 512,

548 + vocoder_intermediate_dim: int = 4096,

549 + vocoder_num_layers: int = 30,

550 + n_mels: int = 80,

551 + sampling_rate: int = 24000,

552 + hop_length: int = 240,

553 + window_size: int = 1024,

554 + vocoder_padding: str = "same",

555 + fmin: int = 0,

556 + fmax: int = None,

557 + num_quantizers: int = 12,

558 + codebook_size: list = None,

559 + threshold_ema_dead_code: int = 10,

560 + position_embedding_type: str = "rope",

561 + rope_theta: int = 10000,

562 + rope_type: str = "default",

563 + ln_type: str = "LayerNorm",

564 + vocoder_attention_heads: int = 4,

565 + vocoder_attn_window_size: list = None,

566 + use_istft_only: bool = False,

567 + hybrid_attention: bool = False,

568 + hybrid_block_size: int = 8,

569 + swa_per_block: int = 2,

570 + **kwargs,

571 + ):

572 + super().__init__(**kwargs)

573 + self.max_audio_seconds = max_audio_seconds

574 + self.stride_size = stride_size

575 + self.avg_pooler = avg_pooler

576 + self.d_model = d_model

577 + self.scale_embedding = scale_embedding

578 + self.kernel_size = kernel_size

579 + self.activation_function = activation_function

580 + self.encoder_layers = encoder_layers

581 + self.encoder_skip_layer_id = encoder_skip_layer_id

582 + self.encoder_attention_heads = encoder_attention_heads

583 + self.encoder_ffn_dim = encoder_ffn_dim

584 + self.encoder_causal = encoder_causal

585 + self.encoder_attn_window_size = (

586 + encoder_attn_window_size

587 + if encoder_attn_window_size is not None

588 + else [-1, -1]

589 + )

590 + self.decoder_layers = decoder_layers

591 + self.decoder_attention_heads = decoder_attention_heads

592 + self.decoder_ffn_dim = decoder_ffn_dim

593 + self.decoder_kernel_size = decoder_kernel_size

594 + self.decoder_stride_size = decoder_stride_size

595 + self.decoder_causal = decoder_causal

596 + self.decoder_attn_window_size = (

597 + decoder_attn_window_size

598 + if decoder_attn_window_size is not None

599 + else [-1, -1]

600 + )

601 + self.nfft = nfft

602 + self.vocoder_dim = vocoder_dim

603 + self.vocoder_intermediate_dim = vocoder_intermediate_dim

604 + self.vocoder_num_layers = vocoder_num_layers

605 + self.n_mels = n_mels

606 + self.sampling_rate = sampling_rate

607 + self.hop_length = hop_length

608 + self.window_size = window_size

609 + self.vocoder_padding = vocoder_padding

610 + self.fmin = fmin

611 + self.fmax = fmax

612 + self.num_quantizers = num_quantizers

613 + self.codebook_size = codebook_size if codebook_size is not None else [1024]

614 + self.threshold_ema_dead_code = threshold_ema_dead_code

615 + self.position_embedding_type = position_embedding_type

616 + self.rope_theta = rope_theta

617 + self.rope_type = rope_type

618 + self.ln_type = ln_type

619 + self.vocoder_attention_heads = vocoder_attention_heads

620 + self.vocoder_attn_window_size = (

621 + vocoder_attn_window_size

622 + if vocoder_attn_window_size is not None

623 + else [40, 10]

624 + )

625 + self.use_istft_only = use_istft_only

626 + self.hybrid_attention = hybrid_attention

627 + self.hybrid_block_size = hybrid_block_size

628 + self.swa_per_block = swa_per_block

629 +

630 +

631 +def get_sequence_mask(inputs, inputs_length):

632 + if inputs.dim() == 3:

633 + bsz, tgt_len, _ = inputs.size()

634 + else:

635 + bsz, tgt_len = inputs_length.shape[0], torch.max(inputs_length)

636 + sequence_mask = torch.arange(0, tgt_len).to(inputs.device)

637 + sequence_mask = torch.lt(sequence_mask, inputs_length.reshape(bsz, 1)).view(

638 + bsz, tgt_len, 1

639 + )

640 + unpacking_index = torch.cumsum(sequence_mask.to(torch.int64).view(-1), dim=0) - 1

641 + return sequence_mask, unpacking_index

642 +

643 +

644 +def unpack_hidden_states(

645 + hidden_states, lengths, sequence_mask=None, unpacking_index=None

646 +):

647 + bsz = lengths.shape[0]

648 + if sequence_mask is None or unpacking_index is None:

649 + sequence_mask, unpacking_index = get_sequence_mask(hidden_states, lengths)

650 + hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(

651 + bsz, torch.max(lengths), hidden_states.shape[-1]

652 + )

653 + return torch.where(sequence_mask, hidden_states, 0)

654 +

655 +

656 +def get_position_ids(lengths):

657 + total_len = lengths.sum()

658 + offset = torch.cat([torch.zeros(1).to(lengths), lengths[:-1].cumsum(dim=0)])

659 + offset = torch.repeat_interleave(offset, lengths)

660 + return torch.arange(0, total_len).to(offset) - offset

661 +

662 +

663 +LAYER_NORM = {"LayerNorm": nn.LayerNorm}

664 +

665 +

666 +class AudioEncoderAttention(nn.Module):

667 + def __init__(

668 + self,

669 + embed_dim: int,

670 + num_heads: int,

671 + window_size: tuple[int, int] = (-1, -1),

672 + causal: bool = False,

673 + ):

674 + super().__init__()

675 + self.embed_dim = embed_dim

676 + self.num_heads = num_heads

677 + self.head_dim = embed_dim // num_heads

678 + self.window_size = window_size

679 + self.causal = causal

680 +

681 + self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)

682 + self.v_proj = nn.Linear(embed_dim, embed_dim, bias=True)

683 + self.q_proj = nn.Linear(embed_dim, embed_dim, bias=True)

684 + self.out_proj = nn.Linear(embed_dim, embed_dim, bias=True)

685 +

686 + def forward(

687 + self,

688 + hidden_states: torch.Tensor,

689 + cu_seqlens: torch.Tensor,

690 + max_seqlen: int,

691 + rope_position_embeddings=None,

692 + ):

693 + from vllm.vllm_flash_attn import flash_attn_varlen_func

694 +

695 + bsz, _ = hidden_states.size()

696 +

697 + query_states = self.q_proj(hidden_states).view(

698 + bsz, self.num_heads, self.head_dim

699 + )

700 + key_states = self.k_proj(hidden_states).view(bsz, self.num_heads, self.head_dim)

701 + value_states = self.v_proj(hidden_states).view(

702 + bsz, self.num_heads, self.head_dim

703 + )

704 +

705 + if rope_position_embeddings is not None:

706 + cos, sin = rope_position_embeddings

707 + query_states, key_states = apply_rotary_pos_emb(

708 + query_states, key_states, cos, sin

709 + )

710 +

711 + attn_output = flash_attn_varlen_func(

712 + query_states,

713 + key_states,

714 + value_states,

715 + cu_seqlens_q=cu_seqlens,

716 + cu_seqlens_k=cu_seqlens,

717 + max_seqlen_q=max_seqlen,

718 + max_seqlen_k=max_seqlen,

719 + causal=self.causal,

720 + window_size=list(self.window_size),

721 + )

722 +

723 + attn_output = attn_output.reshape(bsz, self.embed_dim)

724 + attn_output = self.out_proj(attn_output)

725 + return attn_output

726 +

727 +

728 +class AudioEncoderTransformerLayer(nn.Module):

729 + def __init__(

730 + self,

731 + config: MiMoAudioTokenizerConfig,

732 + causal: bool,

733 + attn_window_size: tuple[int, int] = (-1, -1),

734 + ):

735 + super().__init__()

736 + self.embed_dim = config.d_model

737 +

738 + self.self_attn = AudioEncoderAttention(

739 + embed_dim=self.embed_dim,

740 + num_heads=config.encoder_attention_heads,

741 + window_size=attn_window_size,

742 + causal=causal,

743 + )

744 + self.self_attn_layer_norm = LAYER_NORM[config.ln_type](self.embed_dim)

745 +

746 + self.activation_fn = ACT2FN[config.activation_function]

747 + self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)

748 + self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)

749 + self.final_layer_norm = LAYER_NORM[config.ln_type](self.embed_dim)

750 +

751 + def forward(

752 + self,

753 + hidden_states: torch.Tensor,

754 + cu_seqlens: torch.Tensor,

755 + max_seqlen: int,

756 + rope_position_embeddings: tuple[torch.Tensor, torch.Tensor],

757 + ) -> torch.Tensor:

758 + residual = hidden_states

759 + hidden_states = self.self_attn_layer_norm(hidden_states)

760 + hidden_states = self.self_attn(

761 + hidden_states,

762 + cu_seqlens,

763 + max_seqlen,

764 + rope_position_embeddings=rope_position_embeddings,

765 + )

766 + hidden_states = residual + hidden_states

767 +

768 + residual = hidden_states

769 + hidden_states = self.final_layer_norm(hidden_states)

770 + hidden_states = self.activation_fn(self.fc1(hidden_states))

771 + hidden_states = self.fc2(hidden_states)

772 + hidden_states = residual + hidden_states

773 +

774 + return hidden_states

775 +

776 +

777 +class AudioEncoder(nn.Module):

778 + def __init__(

779 + self,

780 + config: MiMoAudioTokenizerConfig,

781 + ):

782 + super().__init__()

783 + self.config = config

784 + self.max_source_positions = (

785 + config.max_audio_seconds * config.sampling_rate // config.hop_length

786 + ) // config.stride_size

787 + self.embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0

788 + self.skip_layer_idx = config.encoder_skip_layer_id

789 +

790 + self.conv1 = nn.Conv1d(

791 + config.n_mels,

792 + config.d_model,

793 + kernel_size=config.kernel_size,

794 + padding=1,

795 + )

796 + self.conv2 = nn.Conv1d(

797 + config.d_model,

798 + config.d_model,

799 + kernel_size=config.kernel_size,

800 + stride=config.stride_size,

801 + padding=1,

802 + )

803 +

804 + self.position_embedding = AudioRotaryEmbedding(

805 + config.rope_theta,

806 + config.d_model // config.encoder_attention_heads,

807 + self.max_source_positions,

808 + config.rope_type,

809 + )

810 +

811 + attn_window_sizes = []

812 + if config.hybrid_attention:

813 + for i in range(config.encoder_layers):

814 + if i % config.swa_per_block < config.swa_per_block - 1:

815 + attn_window_sizes.append(tuple(config.encoder_attn_window_size))

816 + else:

817 + attn_window_sizes.append((-1, -1))

818 + else:

819 + attn_window_sizes = [

820 + tuple(config.encoder_attn_window_size)

821 + ] * config.encoder_layers

822 +

823 + self.layers = nn.ModuleList(

824 + [

825 + AudioEncoderTransformerLayer(

826 + config=config,

827 + causal=config.encoder_causal,

828 + attn_window_size=attn_window_sizes[i],

829 + )

830 + for i in range(config.encoder_layers)

831 + ]

832 + )

833 +

834 + self.layer_norm = LAYER_NORM[config.ln_type](config.d_model)

835 +

836 + if config.avg_pooler != 1:

837 + self.down_sample_layer = nn.Sequential(

838 + nn.Conv1d(

839 + config.d_model,

840 + config.d_model,

841 + config.avg_pooler,

842 + config.avg_pooler,

843 + bias=False,

844 + ),

845 + nn.GELU(),

846 + )

847 + self.down_sample_norm = LAYER_NORM[config.ln_type](config.d_model)

848 + else:

849 + self.down_sample_layer = None

850 +

851 + if config.num_quantizers != 0:

852 + self.quantizer = ResidualVectorQuantizer(

853 + dimension=config.d_model,

854 + n_q=config.num_quantizers,

855 + bins=config.codebook_size,

856 + threshold_ema_dead_code=config.threshold_ema_dead_code,

857 + )

858 + else:

859 + self.quantizer = None

860 +

861 + def get_features(self, input_features, output_length):

862 + input_features = input_features.to(self.conv1.weight)

863 + inputs_embeds = nn.functional.gelu(self.conv1(input_features))

864 + inputs_embeds = nn.functional.gelu(self.conv2(inputs_embeds))

865 + inputs_embeds = inputs_embeds.permute(0, 2, 1)

866 + bsz, tgt_len, _ = inputs_embeds.size()

867 + hidden_states = inputs_embeds

868 +

869 + position_ids = get_position_ids(output_length).long().to(input_features.device)

870 + rope_position_embeddings = self.position_embedding(input_features, position_ids)

871 +

872 + attention_mask, unpacking_index = get_sequence_mask(

873 + hidden_states, output_length

874 + )

875 + hidden_states = torch.masked_select(hidden_states, attention_mask).view(

876 + torch.sum(output_length), self.config.d_model

877 + )

878 +

879 + cu_seqlens = F.pad(

880 + torch.cumsum(output_length, dim=0), (1, 0), "constant", 0

881 + ).to(device=hidden_states.device, dtype=torch.int32)

882 + max_seqlen = torch.max(output_length).to(torch.int32).item()

883 +

884 + skip_connect_hidden_states = 0.0

885 + for idx, encoder_layer in enumerate(self.layers):

886 + hidden_states = encoder_layer(

887 + hidden_states,

888 + cu_seqlens,

889 + max_seqlen,

890 + rope_position_embeddings=rope_position_embeddings,

891 + )

892 + if (self.skip_layer_idx is not None) and idx == self.skip_layer_idx - 1:

893 + skip_connect_hidden_states = hidden_states.clone()

894 +

895 + hidden_states += skip_connect_hidden_states

896 + hidden_states = self.layer_norm(hidden_states)

897 +

898 + if self.down_sample_layer is not None:

899 + hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(

900 + bsz, tgt_len, self.config.d_model

901 + )

902 + if hidden_states.size(1) % self.config.avg_pooler:

903 + pad_len = (

904 + self.config.avg_pooler

905 + - hidden_states.size(1) % self.config.avg_pooler

906 + )

907 + hidden_states = torch.nn.functional.pad(

908 + hidden_states, (0, 0, 0, pad_len), mode="constant", value=0.0

909 + )

910 + tgt_len += pad_len

911 + tgt_len = tgt_len // self.config.avg_pooler

912 + hidden_states = self.down_sample_layer(hidden_states.transpose(1, 2))

913 + output_length = (

914 + output_length // self.config.avg_pooler

915 + + (output_length % self.config.avg_pooler != 0).int()

916 + )

917 + hidden_states = hidden_states.transpose(1, 2)

918 + attention_mask, unpacking_index = get_sequence_mask(

919 + hidden_states, output_length

920 + )

921 + hidden_states = torch.masked_select(hidden_states, attention_mask).view(

922 + torch.sum(output_length), self.config.d_model

923 + )

924 + hidden_states = self.down_sample_norm(hidden_states)

925 +

926 + return (

927 + hidden_states,

928 + output_length,

929 + attention_mask,

930 + unpacking_index,

931 + tgt_len,

932 + bsz,

933 + )

934 +

935 + def get_output_length(self, mel_len):

936 + tgt_len = mel_len + 3 - self.config.kernel_size

937 + return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1

938 +

939 + @torch.no_grad()

940 + def encode(

941 + self,

942 + input_features,

943 + input_lens=None,

944 + output_length=None,

945 + return_codes_only=False,

946 + n_q=None,

947 + use_quantizer=True,

948 + ):

949 + if output_length is None:

950 + output_length = self.get_output_length(input_lens)

951 + input_features = unpack_hidden_states(input_features, input_lens)

952 + hidden_states, output_length, attention_mask, unpacking_index, tgt_len, bsz = (

953 + self.get_features(

954 + input_features=input_features.transpose(1, 2),

955 + output_length=output_length,

956 + )

957 + )

958 +

959 + dtype = hidden_states.dtype

960 + if use_quantizer and self.quantizer is not None:

961 + self.quantizer.float()

962 + codes = self.quantizer.encode(hidden_states.float(), n_q=n_q)

963 + if return_codes_only:

964 + return codes, output_length

965 + hidden_states = self.quantizer.decode(codes)

966 + hidden_states = hidden_states.to(dtype)

967 + else:

968 + codes = None

969 +

970 + hidden_states_packed = hidden_states.clone()

971 + hidden_states = torch.index_select(hidden_states, 0, unpacking_index).view(

972 + bsz, tgt_len, self.config.d_model

973 + )

974 + hidden_states = torch.where(attention_mask, hidden_states, 0)

975 + return hidden_states, hidden_states_packed, output_length, codes

976 +

977 + @torch.no_grad()

978 + def decode_vq(self, codes):

979 + self.quantizer.float()

980 + return self.quantizer.decode(codes)

981 +

982 +

983 +class MiMoAudioTokenizer(PreTrainedModel):

984 + config_class = MiMoAudioTokenizerConfig

985 +

986 + def __init__(self, config: MiMoAudioTokenizerConfig):

987 + super().__init__(config)

988 + self.config = config

989 + self.sampling_rate = config.sampling_rate

990 + self.encoder = AudioEncoder(config=config)

991 + self.downsample_rate = int(config.hop_length * 2 * config.avg_pooler)

992 +

993 + def get_output_length(self, mel_len):

994 + tgt_len = mel_len + 3 - self.config.kernel_size

995 + return (tgt_len + 2 - self.config.kernel_size) // self.config.stride_size + 1

996 +

997 + @torch.no_grad()

998 + def encode(self, mels, input_lens, use_quantizer=True):

999 + input_features = mels

1000 + encoder_output_length = self.get_output_length(input_lens)

1001 + hidden_states, hidden_states_packed, encoder_output_length, codes = (

1002 + self.encoder.encode(

1003 + input_features, input_lens=input_lens, use_quantizer=use_quantizer

1004 + )

1005 + )

1006 + return hidden_states, hidden_states_packed, encoder_output_length, codes

1007 +

1008 +

1009 +# ---------------------------------------------------------------------------

1010 +# Audio encoding utilities

1011 +# ---------------------------------------------------------------------------

1012 +

1013 +

1014 +def group_by_length(features: torch.Tensor, lengths: torch.Tensor, max_length: int):

1015 + if features.size(0) != lengths.sum().item():

1016 + raise ValueError(

1017 + f"Feature size mismatch: {features.size(0)} vs {lengths.sum().item()}"

1018 + )

1019 +

1020 + split_points = []

1021 + current_sum = 0

1022 +

1023 + for i, seq_len in enumerate(lengths):

1024 + if current_sum + seq_len > max_length and current_sum > 0:

1025 + split_points.append(i)

1026 + current_sum = seq_len.item()

1027 + else:

1028 + current_sum += seq_len.item()

1029 +

1030 + group_sizes = []

1031 + prev = 0

1032 + for point in split_points:

1033 + group_sizes.append(point - prev)

1034 + prev = point

1035 + if prev < len(lengths):

1036 + group_sizes.append(len(lengths) - prev)

1037 +

1038 + len_groups = torch.split(lengths, group_sizes)

1039 + feature_sizes = [group.sum().item() for group in len_groups]

1040 + feature_groups = torch.split(features, feature_sizes)

1041 +

1042 + return feature_groups, len_groups

1043 +

1044 +

1045 +@torch.no_grad()

1046 +def encode_batch(

1047 + audio_tokenizer_encoder,

1048 + input_features: torch.Tensor,

1049 + input_lens: torch.Tensor,

1050 + max_length: int = 256000,

1051 +):

1052 + feature_groups, len_groups = group_by_length(input_features, input_lens, max_length)

1053 +

1054 + encoded_parts = []

1055 + for features, lengths in zip(feature_groups, len_groups):

1056 + codes, _ = audio_tokenizer_encoder.encode(

1057 + input_features=features, input_lens=lengths, return_codes_only=True

1058 + )

1059 + encoded_parts.append(codes)

1060 +

1061 + return torch.cat(encoded_parts, dim=-1)

1062 +

1063 +

1064 +def _segment_lengths_for_mel(mel: torch.Tensor, segment_size: int):

1065 + """Split mel into segments of segment_size with a possible shorter remainder."""

1066 + input_len = mel.size(0)

1067 + segs = [segment_size] * (input_len // segment_size)

1068 + if input_len % segment_size > 0:

1069 + segs.append(input_len % segment_size)

1070 + return segs

1071 +

1072 +

1073 +@torch.no_grad()

1074 +def tokenize_audio_batch(mels, audio_tokenizer_encoder, segment_size=6000, device=None):

1075 + """Tokenize multiple mels in one encode_batch call.

1076 +

1077 + Returns list of code tensors, each [T_i, C] for that mel.

1078 + """

1079 + if not mels:

1080 + return []

1081 + if device is None:

1082 + device = next(audio_tokenizer_encoder.parameters()).device

1083 + input_len_seg_per_mel = [_segment_lengths_for_mel(m, segment_size) for m in mels]

1084 + input_lens_flat = [s for segs in input_len_seg_per_mel for s in segs]

1085 + input_features = torch.cat([m.to(device) for m in mels], dim=0)

1086 + input_lens_t = torch.tensor(input_lens_flat, dtype=torch.long, device=device)

1087 + codes_packed = encode_batch(

1088 + audio_tokenizer_encoder,

1089 + input_features=input_features,

1090 + input_lens=input_lens_t,

1091 + )

1092 + codes = codes_packed.transpose(0, 1).detach() # [total_code_T, C]

1093 + code_lengths = []

1094 + for segs in input_len_seg_per_mel:

1095 + out_len = audio_tokenizer_encoder.get_output_length(

1096 + torch.tensor(segs, dtype=torch.long, device=device)

1097 + )

1098 + if getattr(audio_tokenizer_encoder, "down_sample_layer", None) is not None:

1099 + avg = audio_tokenizer_encoder.config.avg_pooler

1100 + out_len = out_len // avg + (out_len % avg != 0).long()

1101 + code_lengths.append(out_len.sum().item())

1102 + code_list = torch.split(codes, code_lengths)

1103 + return list(code_list)

1104 +

1105 +

1106 +# ---------------------------------------------------------------------------

1107 +# MimoAudioEncoderConfig

1108 +# ---------------------------------------------------------------------------

1109 +

1110 +

1111 +@dataclass

1112 +class MimoAudioEncoderConfig:

1113 + """Config for MimoAudioEncoder.

1114 +

1115 + Field names match the audio_config dict in the model checkpoint.

1116 + """

1117 +

1118 + speech_vocab_size: str = "1025-1025-129-129-129-129-129-129"

1119 + speech_zeroemb_idx: str = "1024-1024-128-128-128-128-128-128"

1120 + group_size: int = 4

1121 + audio_channels: int = 8

1122 + input_local_layers: int = 6

1123 + input_local_dim: int = 1024

1124 + input_full_attention: bool = True

1125 + input_local_attn_heads: int = 64

1126 + input_local_head_dim: int = 16

1127 + input_local_intermediate_size: int = 4096

1128 + input_local_hidden_dropout: float = 0.0

1129 + out_hidden_size: int = 4096

1130 + rope_theta: float = 640000.0

1131 + partial_rotary_factor: float = 0.334

1132 + projection_layers: int = 1

1133 + add_post_norm: bool = False

1134 + audio_segment_size: int = 6000

1135 +

1136 + @classmethod

1137 + def from_dict(cls, d: dict) -> "MimoAudioEncoderConfig":

1138 + known = {f.name for f in dataclasses.fields(cls)}

1139 + return cls(**{k: v for k, v in d.items() if k in known})

1140 +

1141 +

1142 +# ---------------------------------------------------------------------------

1143 +# AudioProjection

1144 +# ---------------------------------------------------------------------------

1145 +

1146 +

1147 +class AudioProjection(nn.Module):

1148 + def __init__(

1149 + self,

1150 + input_size: int,

1151 + hidden_size: int,

1152 + output_size: int,

1153 + ) -> None:

1154 + super().__init__()

1155 + self.mlp = nn.Sequential(

1156 + nn.Linear(input_size, hidden_size, bias=False),

1157 + nn.GELU(),

1158 + nn.Linear(hidden_size, output_size, bias=False),

1159 + )

1160 +

1161 + def forward(self, x: torch.Tensor) -> torch.Tensor:

1162 + return self.mlp(x)

1163 +

1164 +

1165 +# ---------------------------------------------------------------------------

1166 +# MimoAudioEncoder

1167 +# ---------------------------------------------------------------------------

1168 +

1169 +

1170 +class MimoAudioEncoder(nn.Module):

1171 + """Audio encoder for MiMo-V2-Omni.

1172 +

1173 + Encodes mel spectrograms into LLM-compatible embeddings via:

1174 + 1. Audio tokenizer (VQ codes)

1175 + 2. Speech embeddings lookup

1176 + 3. Local Qwen2 transformer

1177 + 4. Linear projection

1178 + """

1179 +

1180 + def __init__(self, config, model_path: str = "") -> None:

1181 + super().__init__()

1182 + if isinstance(config, dict):

1183 + config = MimoAudioEncoderConfig.from_dict(config)

1184 + self.config = config

1185 + self.audio_channels = config.audio_channels

1186 + self.audio_group_size = config.group_size

1187 + self.audio_segment_size = config.audio_segment_size

1188 +

1189 + speech_vocab_sizes = self._parse_maybe_list(

1190 + config.speech_vocab_size, config.audio_channels

1191 + )

1192 + speech_empty_ids = self._parse_maybe_list(

1193 + config.speech_zeroemb_idx, config.audio_channels

1194 + )

1195 +

1196 + input_local_config = Qwen2Config(

1197 + hidden_size=config.input_local_dim,

1198 + num_hidden_layers=config.input_local_layers,

1199 + num_attention_heads=config.input_local_attn_heads,

1200 + num_key_value_heads=config.input_local_attn_heads,

1201 + intermediate_size=config.input_local_intermediate_size,

1202 + attention_dropout=config.input_local_hidden_dropout,

1203 + rope_theta=config.rope_theta,

1204 + partial_rotary_factor=config.partial_rotary_factor,

1205 + )

1206 +

1207 + self.input_local_transformer = Qwen2Model(input_local_config)

1208 +

1209 + if not config.add_post_norm:

1210 + self.input_local_transformer.norm = nn.Identity()

1211 +

1212 + self.speech_embeddings = nn.ModuleList(

1213 + [

1214 + nn.Embedding(

1215 + speech_vocab_sizes[i],

1216 + config.input_local_dim,

1217 + padding_idx=speech_empty_ids[i],

1218 + )

1219 + for i in range(config.audio_channels)

1220 + ]

1221 + )

1222 +

1223 + if config.projection_layers == 1:

1224 + self.projection = nn.Linear(

1225 + config.input_local_dim * config.group_size,

1226 + config.out_hidden_size,

1227 + bias=False,

1228 + )

1229 + elif config.projection_layers == 2:

1230 + self.projection = AudioProjection(

1231 + config.input_local_dim * config.group_size,

1232 + config.input_local_dim * config.group_size * 4,

1233 + config.out_hidden_size,

1234 + )

1235 + else:

1236 + raise ValueError(f"Invalid projection_layers: {config.projection_layers}")

1237 +

1238 + self.audio_tokenizer: MiMoAudioTokenizer | None = None

1239 + if model_path:

1240 + audio_tokenizer_path = os.path.join(model_path, "audio_tokenizer")

1241 + if os.path.exists(audio_tokenizer_path):

1242 + dev = torch.get_default_device()

1243 + self.audio_tokenizer = self._load_audio_tokenizer(

1244 + audio_tokenizer_path, dev

1245 + )

1246 + else:

1247 + logger.warning(

1248 + "Audio tokenizer not found at %s, audio encoding disabled",

1249 + audio_tokenizer_path,

1250 + )

1251 +

1252 + @staticmethod

1253 + def _load_audio_tokenizer(path: str, device: torch.device) -> MiMoAudioTokenizer:

1254 + """Load MiMoAudioTokenizer from directory."""

1255 + from safetensors.torch import load_file

1256 +

1257 + config_path = os.path.join(path, "config.json")

1258 + with open(config_path) as f:

1259 + config_dict = json.load(f)

1260 + config = MiMoAudioTokenizer.config_class(**config_dict)

1261 + model = MiMoAudioTokenizer(config)

1262 + safetensors_path = os.path.join(path, "model.safetensors")

1263 + bin_path = os.path.join(path, "pytorch_model.bin")

1264 + if os.path.exists(safetensors_path):

1265 + state_dict = load_file(safetensors_path, device="cpu")

1266 + elif os.path.exists(bin_path):

1267 + state_dict = torch.load(bin_path, map_location="cpu", weights_only=True)

1268 + else:

1269 + raise FileNotFoundError(

1270 + f"No model weights found in {path} "

1271 + "(expected model.safetensors or pytorch_model.bin)"

1272 + )

1273 + model.load_state_dict(state_dict, strict=False)

1274 + model = model.to(device=device, dtype=torch.bfloat16)

1275 + model.eval()

1276 + model.requires_grad_(False)

1277 + return model

1278 +

1279 + def _parse_maybe_list(self, value, length: int) -> list[int]:

1280 + if isinstance(value, str) and "-" in value:

1281 + return [int(s) for s in value.split("-")]

1282 + return [int(value)] * length

1283 +

1284 + def apply_input_local_transformer(self, speech_embeddings: torch.Tensor):

1285 + output = self.input_local_transformer(

1286 + inputs_embeds=speech_embeddings,

1287 + return_dict=True,

1288 + is_causal=not self.config.input_full_attention,

1289 + )

1290 + return output.last_hidden_state

1291 +

1292 + def apply_speech_embeddings(self, audio_codes: torch.Tensor) -> torch.Tensor:

1293 + num_segments = audio_codes.shape[0]

1294 + _audio_embeddings = torch.zeros(

1295 + (num_segments, self.config.group_size, self.config.input_local_dim),

1296 + dtype=next(self.speech_embeddings[0].parameters()).dtype,

1297 + device=audio_codes.device,

1298 + )

1299 + for i in range(self.config.audio_channels):

1300 + _audio_embeddings.add_(self.speech_embeddings[i](audio_codes[:, :, i]))

1301 + return _audio_embeddings

1302 +

1303 + def process_audio(self, audio: torch.Tensor) -> torch.Tensor:

1304 + """Pad audio codes to group_size boundary.

1305 +

1306 + Args:

1307 + audio: [T, audio_channels] code tensor

1308 +

1309 + Returns:

1310 + [T//group_size, group_size, audio_channels]

1311 + """

1312 + T = audio.shape[0]

1313 + audio = audio[:, : self.audio_channels]

1314 + padded_T = (

1315 + (T + self.audio_group_size - 1)

1316 + // self.audio_group_size

1317 + * self.audio_group_size

1318 + )

1319 + padded_audio = torch.cat(

1320 + [

1321 + audio,

1322 + torch.zeros(

1323 + padded_T - T,

1324 + self.audio_channels,

1325 + dtype=torch.int32,

1326 + device=audio.device,

1327 + )

1328 + + audio[-1, :],

1329 + ],

1330 + dim=0,

1331 + )

1332 + padded_audio = padded_audio.reshape(

1333 + padded_T // self.audio_group_size,

1334 + self.audio_group_size,

1335 + self.audio_channels,

1336 + )

1337 + return padded_audio

1338 +

1339 + def get_audio_feature(

1340 + self, mel_specs: list[torch.Tensor]

1341 + ) -> tuple[torch.Tensor, list[int]]:

1342 + """Encode mel spectrograms into LLM embedding space.

1343 +

1344 + Args:

1345 + mel_specs: list of mel spectrogram tensors, each [T, n_mels]

1346 +

1347 + Returns:

1348 + Tuple of:

1349 + - audio_embeds: [total_tokens, out_hidden_size] concatenated embeddings

1350 + - item_token_lens: list of int, number of tokens per input item

1351 + """

1352 + if self.audio_tokenizer is None:

1353 + raise RuntimeError(

1354 + "audio_tokenizer is not loaded. "

1355 + "Ensure model_path points to a directory containing audio_tokenizer/."

1356 + )

1357 +

1358 + if not mel_specs:

1359 + device = next(self.projection.parameters()).device

1360 + dtype = next(self.projection.parameters()).dtype

1361 + return (

1362 + torch.empty(0, self.config.out_hidden_size, device=device, dtype=dtype),

1363 + [],

1364 + )

1365 +

1366 + device = next(self.audio_tokenizer.encoder.parameters()).device

1367 + code_list = tokenize_audio_batch(

1368 + mel_specs,

1369 + self.audio_tokenizer.encoder,

1370 + segment_size=self.audio_segment_size,

1371 + device=device,

1372 + )

1373 +

1374 + item_token_lens: list[int] = []

1375 + codecs_to_concat = []

1376 + for codecs in code_list:

1377 + padded_codes = self.process_audio(codecs)

1378 + codecs_to_concat.append(padded_codes)

1379 + item_token_lens.append(padded_codes.shape[0])

1380 +

1381 + audio_codes = torch.cat(

1382 + codecs_to_concat, dim=0

1383 + ) # [total_T//group_size, group_size, audio_channels]

1384 +

1385 + _audio_embeddings = self.apply_speech_embeddings(audio_codes)

1386 + audio_embeds = self.apply_input_local_transformer(_audio_embeddings)

1387 + B = audio_embeds.shape[0]

1388 + audio_embeds = self.projection(audio_embeds.reshape(B, -1))

1389 + return audio_embeds, item_token_lens

vllm/model_executor/models/mimo_v2.py renamed

+22 -2

@@ -6,6 +6,7 @@ from itertools import islice

6 import torch

7 from torch import nn

9 +from vllm.compilation.decorators import support_torch_compile

10 from vllm.config import (

11 CacheConfig,

12 VllmConfig,

@@ -268,7 +269,7 @@ class MiMoV2Attention(nn.Module):

269 self.total_num_heads * self.v_head_dim,

270 hidden_size,

271 bias=False,

271 - quant_config=quant_config,

272 + quant_config=quant_config if "mtp.layers" not in prefix else None,

273 reduce_results=True,

274 prefix=f"{prefix}.o_proj",

275 )

@@ -440,6 +441,7 @@ class MiMoV2FlashDecoderLayer(nn.Module):

441 return self.config.hybrid_layer_pattern[self.layer_id] == 1

442

443

444 +@support_torch_compile

445 class MiMoV2Model(nn.Module):

446 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

447 super().__init__()

@@ -603,7 +605,13 @@ class MiMoV2Model(nn.Module):

605

606 if expert_matched:

607 continue

606 -

608 + # Support fused qkv_proj checkpoint (Pro format)

609 + if "qkv_proj" in name:

610 + if name in params_dict:

611 + param = params_dict[name]

612 + loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]

613 + default_weight_loader(param, loaded_weight)

614 + continue

615 stacked_matched = False

616 for param_name, weight_name, shard_id in stacked_params_mapping:

617 if weight_name not in name:

@@ -662,6 +670,11 @@ class MiMoV2Model(nn.Module):

670

671

672 class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

673 + packed_modules_mapping = {

674 + "qkv_proj": ["q_proj", "k_proj", "v_proj"],

675 + "gate_up_proj": ["gate_proj", "up_proj"],

676 + }

677 +

678 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

679 super().__init__()

680 config = vllm_config.model_config.hf_config

@@ -718,3 +731,10 @@ class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts):

731 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

732 loader = AutoWeightsLoader(self)

733 return loader.load_weights(weights)

734 +

735 +

736 +class MiMoV2ProForCausalLM(MiMoV2FlashForCausalLM):

737 + packed_modules_mapping = {

738 + "qkv_proj": ["qkv_proj"],

739 + "gate_up_proj": ["gate_proj", "up_proj"],

740 + }

vllm/model_executor/models/mimo_v2_mtp.py new

+373

@@ -0,0 +1,373 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +"""Inference-only MiMo-V2 MTP (Multi-Token Prediction) draft model.

5 +

6 +Supports both MiMo-V2-Pro and MiMo-V2-Flash checkpoints.

7 +

8 +Checkpoint weight layout (model.mtp.layers.{idx}.*):

9 + enorm - RMSNorm for token embeddings

10 + hnorm - RMSNorm for previous hidden states

11 + eh_proj - ReplicatedLinear(hidden*2 -> hidden)

12 + input_layernorm - pre-attention RMSNorm

13 + self_attn.* - attention weights; format differs by variant:

14 + Pro: fused qkv_proj [Q;K;V] concatenated

15 + Flash: separate q_proj, k_proj, v_proj

16 + pre_mlp_layernorm - post-attention / pre-MLP RMSNorm

17 + mlp.* - dense MLP (gate_proj / up_proj / down_proj)

18 + final_layernorm - norm applied before logit computation

19 +"""

20 +

21 +from collections.abc import Iterable

22 +

23 +import torch

24 +import torch.nn as nn

25 +from transformers import PretrainedConfig

26 +

27 +from vllm.config import VllmConfig

28 +from vllm.distributed import (

29 + get_tensor_model_parallel_rank,

30 + get_tensor_model_parallel_world_size,

31 +)

32 +from vllm.model_executor.layers.layernorm import RMSNorm

33 +from vllm.model_executor.layers.linear import ReplicatedLinear

34 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

35 +from vllm.model_executor.layers.quantization import QuantizationConfig

36 +from vllm.model_executor.layers.vocab_parallel_embedding import (

37 + ParallelLMHead,

38 + VocabParallelEmbedding,

39 +)

40 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

41 +from vllm.sequence import IntermediateTensors

42 +

43 +from .interfaces import (

44 + MultiModalEmbeddings,

45 + SupportsMultiModal,

46 + _require_is_multimodal,

47 +)

48 +from .mimo_v2 import MiMoV2Attention, MiMoV2MLP

49 +from .utils import _merge_multimodal_embeddings, maybe_prefix

50 +

51 +# MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports

52 +# only the first layer and only one speculative token.

53 +_MIMO_V2_PRO_NUM_MTP_LAYERS = 1

54 +_MIMO_V2_FLASH_NUM_MTP_LAYERS = 1

55 +

56 +

57 +class MiMoV2MTPLayer(nn.Module):

58 + """Single MTP predictor layer for MiMo-V2 (Pro and Flash).

59 +

60 + Mirrors the single-layer MiMo-V2 nextn reference implementation.

61 + """

62 +

63 + def __init__(

64 + self,

65 + config: PretrainedConfig,

66 + prefix: str,

67 + quant_config: QuantizationConfig | None = None,

68 + ) -> None:

69 + super().__init__()

70 +

71 + # Predictor head components

72 + self.enorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)

73 + self.hnorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)

74 + self.eh_proj = ReplicatedLinear(

75 + config.hidden_size * 2, config.hidden_size, bias=False

76 + )

77 +

78 + # MTP uses the SWA attention configuration

79 + # implementation.

80 + swa_rope_theta = getattr(

81 + config,

82 + "swa_rope_theta",

83 + getattr(config, "rope_theta", 1000000),

84 + )

85 + sliding_window_size = getattr(config, "sliding_window_size", -1)

86 +

87 + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)

88 + self.self_attn = MiMoV2Attention(

89 + hidden_size=config.hidden_size,

90 + num_heads=config.swa_num_attention_heads,

91 + num_kv_heads=config.swa_num_key_value_heads,

92 + head_dim=config.swa_head_dim,

93 + v_head_dim=getattr(config, "swa_v_head_dim", None),

94 + v_scale=getattr(config, "attention_value_scale", None),

95 + sliding_window_size=sliding_window_size,

96 + attention_bias=config.attention_bias,

97 + add_swa_attention_sink_bias=getattr(

98 + config, "add_swa_attention_sink_bias", False

99 + ),

100 + layer_id=0,

101 + rope_theta=swa_rope_theta,

102 + max_position_embeddings=getattr(config, "max_position_embeddings", 32768),

103 + quant_config=quant_config,

104 + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0),

105 + prefix=f"{prefix}.self_attn",

106 + )

107 + self.pre_mlp_layernorm = RMSNorm(

108 + config.hidden_size, eps=config.layernorm_epsilon

109 + )

110 + self.mlp = MiMoV2MLP(

111 + hidden_size=config.hidden_size,

112 + intermediate_size=config.intermediate_size,

113 + hidden_act=config.hidden_act,

114 + quant_config=quant_config,

115 + prefix=f"{prefix}.mlp",

116 + )

117 + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon)

118 +

119 + def forward(

120 + self,

121 + inputs_embeds: torch.Tensor,

122 + positions: torch.Tensor,

123 + previous_hidden_states: torch.Tensor,

124 + ) -> torch.Tensor:

125 + # Combine token embedding and previous hidden state

126 + h, _ = self.eh_proj(

127 + torch.cat(

128 + [self.enorm(inputs_embeds), self.hnorm(previous_hidden_states)], dim=-1

129 + )

130 + )

131 +

132 + # Transformer block with fused residual norms

133 + residual = h

134 + h = self.input_layernorm(h)

135 + h = self.self_attn(positions=positions, hidden_states=h)

136 + h, residual = self.pre_mlp_layernorm(h, residual)

137 + h = self.mlp(h)

138 + h = h + residual

139 +

140 + return self.final_layernorm(h)

141 +

142 +

143 +class _MiMoV2MTPLayers(nn.Module):

144 + """Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*"""

145 +

146 + def __init__(

147 + self,

148 + config: PretrainedConfig,

149 + num_mtp_layers: int,

150 + quant_config: QuantizationConfig | None,

151 + prefix: str,

152 + ) -> None:

153 + super().__init__()

154 + self.layers = nn.ModuleDict(

155 + {

156 + str(i): MiMoV2MTPLayer(

157 + config=config,

158 + prefix=f"{prefix}.{i}",

159 + quant_config=quant_config,

160 + )

161 + for i in range(num_mtp_layers)

162 + }

163 + )

164 +

165 +

166 +class MiMoV2MultiTokenPredictor(nn.Module):

167 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

168 + super().__init__()

169 +

170 + config = vllm_config.model_config.hf_config

171 + spec_cfg = vllm_config.speculative_config

172 + assert spec_cfg is not None

173 + if spec_cfg.num_speculative_tokens != 1:

174 + raise ValueError(

175 + "MiMo-V2 MTP in vLLM only supports num_speculative_tokens=1."

176 + )

177 + num_mtp_layers = 1

178 +

179 + self.num_mtp_layers = num_mtp_layers

180 +

181 + self.embed_tokens = VocabParallelEmbedding(

182 + config.vocab_size,

183 + config.hidden_size,

184 + )

185 +

186 + self.mtp = _MiMoV2MTPLayers(

187 + config=config,

188 + num_mtp_layers=num_mtp_layers,

189 + quant_config=vllm_config.quant_config,

190 + prefix=maybe_prefix(prefix, "mtp.layers"),

191 + )

192 +

193 + self.logits_processor = LogitsProcessor(config.vocab_size)

194 +

195 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

196 + return self.embed_tokens(input_ids)

197 +

198 + def forward(

199 + self,

200 + input_ids: torch.Tensor,

201 + positions: torch.Tensor,

202 + previous_hidden_states: torch.Tensor,

203 + inputs_embeds: torch.Tensor | None = None,

204 + spec_step_idx: int = 0,

205 + ) -> torch.Tensor:

206 + assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."

207 + if inputs_embeds is None:

208 + inputs_embeds = self.embed_input_ids(input_ids)

209 + return self.mtp.layers[str(spec_step_idx)](

210 + inputs_embeds, positions, previous_hidden_states

211 + )

212 +

213 + def compute_logits(

214 + self,

215 + hidden_states: torch.Tensor,

216 + lm_head: ParallelLMHead,

217 + spec_step_idx: int = 0,

218 + ) -> torch.Tensor:

219 + assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."

220 + return self.logits_processor(lm_head, hidden_states)

221 +

222 +

223 +class MiMoV2MTP(nn.Module):

224 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:

225 + super().__init__()

226 + self.config = vllm_config.model_config.hf_config

227 + self.model = MiMoV2MultiTokenPredictor(

228 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

229 + )

230 + self.lm_head = ParallelLMHead(

231 + self.config.vocab_size,

232 + self.config.hidden_size,

233 + prefix=maybe_prefix(prefix, "lm_head"),

234 + )

235 +

236 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

237 + return self.model.embed_input_ids(input_ids)

238 +

239 + def forward(

240 + self,

241 + input_ids: torch.Tensor | None,

242 + positions: torch.Tensor,

243 + hidden_states: torch.Tensor,

244 + intermediate_tensors: IntermediateTensors | None = None,

245 + inputs_embeds: torch.Tensor | None = None,

246 + spec_step_idx: int = 0,

247 + ) -> torch.Tensor:

248 + assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."

249 + return self.model(

250 + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx

251 + )

252 +

253 + def compute_logits(

254 + self,

255 + hidden_states: torch.Tensor,

256 + spec_step_idx: int = 0,

257 + ) -> torch.Tensor | None:

258 + assert spec_step_idx == 0, "MiMo-V2 MTP only supports one speculative token."

259 + return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx)

260 +

261 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

262 + tp_rank = get_tensor_model_parallel_rank()

263 + tp_size = get_tensor_model_parallel_world_size()

264 +

265 + stacked_params_mapping = [

266 + ("gate_up_proj", "gate_proj", 0),

267 + ("gate_up_proj", "up_proj", 1),

268 + # Flash format: separate projections → fused qkv_proj

269 + ("qkv_proj", "q_proj", "q"),

270 + ("qkv_proj", "k_proj", "k"),

271 + ("qkv_proj", "v_proj", "v"),

272 + ]

273 +

274 + params_dict = dict(self.named_parameters())

275 + loaded_params: set[str] = set()

276 +

277 + for name, loaded_weight in weights:

278 + if "rotary_emb.inv_freq" in name:

279 + continue

280 +

281 + # Only load MTP-related weights, shared embeddings, and lm_head

282 + if (

283 + "model.mtp" not in name

284 + and "model.embed_tokens" not in name

285 + and not name.startswith("lm_head")

286 + ):

287 + continue

288 +

289 + # Support fused qkv_proj checkpoint (Pro format).

290 + # The checkpoint is stored pre-sharded for TP=8 as

291 + # [Q_rank0, K_rank0, V_rank0, Q_rank1, ...], so splitting along

292 + # dim 0 with chunk(tp_size) gives each rank its Q+K+V slice for

293 + # both the FP8 weight and the block weight_scale_inv. This matches

294 + # how the main model loads the same layout.

295 + if "qkv_proj" in name:

296 + if name in params_dict:

297 + param = params_dict[name]

298 + loaded_weight = loaded_weight.chunk(tp_size, dim=0)[tp_rank]

299 + default_weight_loader(param, loaded_weight)

300 + loaded_params.add(name)

301 + continue

302 +

303 + # gate_proj/up_proj → gate_up_proj stacking (both formats);

304 + # Flash: q_proj/k_proj/v_proj → qkv_proj merging.

305 + stacked_matched = False

306 + for param_name, weight_name, shard_id in stacked_params_mapping:

307 + if weight_name not in name:

308 + continue

309 + name_rewritten = name.replace(weight_name, param_name)

310 + if (

311 + name_rewritten.endswith(".bias")

312 + and name_rewritten not in params_dict

313 + ):

314 + continue

315 + if name_rewritten not in params_dict:

316 + continue

317 + param = params_dict[name_rewritten]

318 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

319 + weight_loader(param, loaded_weight, shard_id)

320 + loaded_params.add(name_rewritten)

321 + stacked_matched = True

322 + break

323 +

324 + if stacked_matched:

325 + continue

326 +

327 + if name.endswith(".bias") and name not in params_dict:

328 + continue

329 + if name not in params_dict:

330 + continue

331 +

332 + param = params_dict[name]

333 + # attention_sink_bias is head-parallel; slice by tp

334 + if "attention_sink_bias" in name:

335 + total_heads = loaded_weight.shape[0]

336 + heads_per_rank = total_heads // tp_size

337 + loaded_weight = loaded_weight.narrow(

338 + 0, tp_rank * heads_per_rank, heads_per_rank

339 + )

340 +

341 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

342 + weight_loader(param, loaded_weight)

343 + loaded_params.add(name)

344 +

345 + return loaded_params

346 +

347 +

348 +class MiMoV2OmniMTP(MiMoV2MTP, SupportsMultiModal):

349 + def embed_input_ids(

350 + self,

351 + input_ids: torch.Tensor,

352 + multimodal_embeddings: MultiModalEmbeddings | None = None,

353 + *,

354 + is_multimodal: torch.Tensor | None = None,

355 + ) -> torch.Tensor:

356 + inputs_embeds = self._embed_text_input_ids(

357 + input_ids,

358 + self.model.embed_input_ids,

359 + is_multimodal=is_multimodal,

360 + )

361 +

362 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

363 + return inputs_embeds

364 +

365 + is_multimodal = _require_is_multimodal(is_multimodal)

366 +

367 + inputs_embeds = _merge_multimodal_embeddings(

368 + inputs_embeds=inputs_embeds,

369 + multimodal_embeddings=multimodal_embeddings,

370 + is_multimodal=is_multimodal,

371 + )

372 +

373 + return inputs_embeds

vllm/model_executor/models/mimo_v2_omni.py new

+1488

@@ -0,0 +1,1488 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +import math

4 +from collections.abc import Callable, Iterable, Mapping, Sequence

5 +from functools import partial

6 +from typing import Any

7 +

8 +import einops

9 +import numpy as np

10 +import torch

11 +import torch.nn as nn

12 +import torch.nn.functional as F

13 +from transformers import BatchFeature, PretrainedConfig

14 +from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize

15 +

16 +from vllm.config import VllmConfig

17 +from vllm.config.multimodal import BaseDummyOptions

18 +from vllm.distributed import parallel_state

19 +from vllm.distributed import utils as dist_utils

20 +from vllm.inputs import MultiModalDataDict

21 +from vllm.model_executor.layers.activation import get_act_and_mul_fn

22 +from vllm.model_executor.layers.attention import MMEncoderAttention

23 +from vllm.model_executor.layers.layernorm import RMSNorm

24 +from vllm.model_executor.layers.linear import (

25 + ColumnParallelLinear,

26 + QKVParallelLinear,

27 + RowParallelLinear,

28 +)

29 +from vllm.model_executor.layers.quantization import QuantizationConfig

30 +from vllm.model_executor.layers.rotary_embedding import get_rope

31 +from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb

32 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

33 +from vllm.model_executor.models.vision import is_vit_use_data_parallel

34 +from vllm.multimodal import MULTIMODAL_REGISTRY

35 +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems

36 +from vllm.multimodal.parse import ImageSize, MultiModalDataItems

37 +from vllm.multimodal.processing import (

38 + BaseDummyInputsBuilder,

39 + BaseMultiModalProcessor,

40 + BaseProcessingInfo,

41 + PromptReplacement,

42 + PromptUpdate,

43 + PromptUpdateDetails,

44 +)

45 +from vllm.transformers_utils.configs.mimo_v2_omni import Mimo_VLVisionConfig

46 +from vllm.transformers_utils.processors.mimo_v2_omni import (

47 + MiMoOmniProcessor,

48 + VideoAudioInput,

49 + _format_timestamp,

50 +)

51 +

52 +from .interfaces import (

53 + MultiModalEmbeddings,

54 + SupportsMultiModal,

55 + SupportsPP,

56 + SupportsQuant,

57 +)

58 +from .mimo_audio import MimoAudioEncoder

59 +from .mimo_v2 import MiMoV2FlashForCausalLM

60 +from .qwen2_5_vl import (

61 + Qwen2_5_VisionMLP,

62 + Qwen2_5_VisionPatchEmbed,

63 + Qwen2_5_VLImageEmbeddingInputs,

64 + Qwen2_5_VLImageInputs,

65 + Qwen2_5_VLImagePixelInputs,

66 + Qwen2_5_VLVideoEmbeddingInputs,

67 + Qwen2_5_VLVideoInputs,

68 + Qwen2_5_VLVideoPixelInputs,

69 +)

70 +from .qwen2_vl import _create_qwen2vl_field_factory

71 +from .utils import AutoWeightsLoader, IntermediateTensors, WeightsMapper, maybe_prefix

72 +

73 +

74 +class MiMoVisionMLP(Qwen2_5_VisionMLP):

75 + pass

76 +

77 +

78 +class MiMoVisionPatchEmbed(Qwen2_5_VisionPatchEmbed):

79 + pass

80 +

81 +

82 +class MiMoVisionPatchMerger(nn.Module):

83 + def __init__(

84 + self,

85 + d_model: int,

86 + context_dim: int,

87 + norm_layer: Callable[[int], nn.Module] | None = None,

88 + spatial_merge_size: int = 2,

89 + quant_config: QuantizationConfig | None = None,

90 + prefix: str = "",

91 + ) -> None:

92 + super().__init__()

93 + use_data_parallel = is_vit_use_data_parallel()

94 + self.hidden_size = context_dim * (spatial_merge_size**2)

95 + if norm_layer is None:

96 + norm_layer = partial(nn.LayerNorm, eps=1e-6)

97 + self.ln_q = norm_layer(context_dim)

98 +

99 + self.mlp = nn.Sequential(

100 + ColumnParallelLinear(

101 + self.hidden_size,

102 + self.hidden_size,

103 + bias=False,

104 + quant_config=quant_config,

105 + prefix=f"{prefix}.mlp.0",

106 + return_bias=False,

107 + disable_tp=use_data_parallel,

108 + ),

109 + nn.GELU(),

110 + RowParallelLinear(

111 + self.hidden_size,

112 + d_model,

113 + bias=False,

114 + quant_config=quant_config,

115 + prefix=f"{prefix}.mlp.2",

116 + return_bias=False,

117 + disable_tp=use_data_parallel,

118 + ),

119 + )

120 +

121 + def forward(self, x: torch.Tensor) -> torch.Tensor:

122 + x = self.ln_q(x)

123 + x = x.view(-1, self.hidden_size)

124 + out = self.mlp(x)

125 + return out

126 +

127 +

128 +class MiMoVisionAttention(nn.Module):

129 + def __init__(

130 + self,

131 + embed_dim: int,

132 + num_heads: int,

133 + num_kv_heads: int,

134 + qk_channels: int,

135 + kv_channels: int,

136 + use_sink: bool = False,

137 + visual_token_window_size: int = 64,

138 + quant_config: QuantizationConfig | None = None,

139 + prefix: str = "",

140 + ) -> None:

141 + super().__init__()

142 + use_data_parallel = is_vit_use_data_parallel()

143 + self.tp_size = (

144 + 1

145 + if use_data_parallel

146 + else parallel_state.get_tensor_model_parallel_world_size()

147 + )

148 + self.tp_rank = parallel_state.get_tensor_model_parallel_rank()

149 +

150 + self.num_heads = num_heads

151 + self.num_kv_heads = num_kv_heads

152 + self.qk_channels = qk_channels

153 + self.kv_channels = kv_channels

154 + self.embed_dim = embed_dim

155 +

156 + self.num_heads_per_partition = dist_utils.divide(num_heads, self.tp_size)

157 + self.num_kv_heads_per_partition = dist_utils.divide(num_kv_heads, self.tp_size)

158 +

159 + # Attention scale uses the Q/K head dimension (qk_channels)

160 + self.scale = qk_channels**-0.5

161 +

162 + # QKV: Q is (num_heads * qk_channels), KV are (num_kv_heads * kv_channels)

163 + self.qkv = QKVParallelLinear(

164 + hidden_size=embed_dim,

165 + head_size=qk_channels,

166 + total_num_heads=num_heads,

167 + total_num_kv_heads=num_kv_heads,

168 + v_head_size=kv_channels,

169 + bias=True,

170 + quant_config=quant_config,

171 + prefix=f"{prefix}.qkv",

172 + disable_tp=use_data_parallel,

173 + )

174 +

175 + # Output projection: input is (num_heads * kv_channels) after attention

176 + self.proj = RowParallelLinear(

177 + input_size=num_heads * kv_channels,

178 + output_size=embed_dim,

179 + bias=True,

180 + quant_config=quant_config,

181 + prefix=f"{prefix}.proj",

182 + disable_tp=use_data_parallel,

183 + )

184 +

185 + # For full attention (non-window blocks)

186 + self.attn = MMEncoderAttention(

187 + num_heads=self.num_heads_per_partition,

188 + head_size=kv_channels,

189 + scale=self.scale,

190 + num_kv_heads=self.num_kv_heads_per_partition,

191 + prefix=f"{prefix}.attn",

192 + )

193 +

194 + # Rotary embeddings applied separately to Q and K

195 + self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True)

196 +

197 + # Sink attention weights (loaded but not used in vLLM flash_attn)

198 + # The checkpoint stores these only for non-full-attention blocks

199 + self.use_sink = use_sink

200 + if use_sink:

201 + self.sinks = nn.Parameter(

202 + torch.empty(num_heads),

203 + requires_grad=False,

204 + )

205 + else:

206 + self.sinks = None

207 +

208 + self.visual_token_window_size = visual_token_window_size

209 +

210 + def _forward_window_attn(

211 + self,

212 + q: torch.Tensor,

213 + k: torch.Tensor,

214 + v: torch.Tensor,

215 + cu_seqlens: torch.Tensor,

216 + max_seqlen: torch.Tensor,

217 + ) -> torch.Tensor:

218 + """Window attention via flash_attn_varlen_func with window_size."""

219 + from vllm.vllm_flash_attn import flash_attn_varlen_func

220 +

221 + w = self.visual_token_window_size

222 + output = flash_attn_varlen_func(

223 + q,

224 + k,

225 + v,

226 + cu_seqlens_q=cu_seqlens,

227 + cu_seqlens_k=cu_seqlens,

228 + max_seqlen_q=max_seqlen,

229 + max_seqlen_k=max_seqlen,

230 + softmax_scale=self.scale,

231 + causal=False,

232 + window_size=[w, w],

233 + )

234 + return output

235 +

236 + def forward(

237 + self,

238 + x: torch.Tensor,

239 + cu_seqlens: torch.Tensor,

240 + rotary_pos_emb_cos: torch.Tensor,

241 + rotary_pos_emb_sin: torch.Tensor,

242 + max_seqlen: torch.Tensor,

243 + full_attn: bool = True,

244 + ) -> torch.Tensor:

245 + """

246 + Args:

247 + x: [seq_len, batch=1, embed_dim] (seq-first convention)

248 + cu_seqlens: cumulative sequence lengths [num_seqs+1], int32

249 + rotary_pos_emb_cos: [seq_len, qk_channels // 2]

250 + rotary_pos_emb_sin: [seq_len, qk_channels // 2]

251 + max_seqlen: maximum sequence length

252 + full_attn: if True, full attention; if False, window attention

253 + """

254 + # [seq_len, 1, embed_dim] -> QKV projection

255 + qkv, _ = self.qkv(x) # [seq_len, 1, q_size + kv_size + kv_size]

256 + seq_len, batch_size, _ = qkv.shape

257 +

258 + q_size = self.num_heads_per_partition * self.qk_channels

259 + kv_size = self.num_kv_heads_per_partition * self.kv_channels

260 + q, k, v = qkv.split([q_size, kv_size, kv_size], dim=-1)

261 +

262 + # Rearrange to [batch, seq, head, head_dim] for rotary application

263 + q = einops.rearrange(q, "s b (h d) -> b s h d", h=self.num_heads_per_partition)

264 + k = einops.rearrange(

265 + k, "s b (h d) -> b s h d", h=self.num_kv_heads_per_partition

266 + )

267 + v = einops.rearrange(

268 + v, "s b (h d) -> b s h d", h=self.num_kv_heads_per_partition

269 + )

270 +

271 + # Apply rotary embeddings to Q and K independently (handles GQA)

272 + if rotary_pos_emb_cos is not None and rotary_pos_emb_sin is not None:

273 + q = self.apply_rotary_emb(q, rotary_pos_emb_cos, rotary_pos_emb_sin)

274 + k = self.apply_rotary_emb(k, rotary_pos_emb_cos, rotary_pos_emb_sin)

275 +

276 + if full_attn:

277 + # Full attention via MMEncoderAttention

278 + # Flatten to [batch, seq, heads * head_dim]

279 + q_flat = q.reshape(batch_size, seq_len, -1)

280 + k_flat = k.reshape(batch_size, seq_len, -1)

281 + v_flat = v.reshape(batch_size, seq_len, -1)

282 + context_layer = self.attn(

283 + query=q_flat,

284 + key=k_flat,

285 + value=v_flat,

286 + cu_seqlens=cu_seqlens,

287 + max_seqlen=max_seqlen,

288 + )

289 + # context_layer: [batch, seq, num_heads, head_dim] or [batch, seq, hidden]

290 + # Ensure shape is [seq, batch, num_heads * kv_channels]

291 + if context_layer.dim() == 4:

292 + context_layer = einops.rearrange(

293 + context_layer, "b s h d -> s b (h d)"

294 + ).contiguous()

295 + else:

296 + context_layer = einops.rearrange(

297 + context_layer, "b s d -> s b d"

298 + ).contiguous()

299 + else:

300 + # Window attention via flash_attn_varlen_func with window_size

301 + # Flatten batch dimension: [seq, head, head_dim]

302 + q_varlen = einops.rearrange(q, "b s h d -> (b s) h d")

303 + k_varlen = einops.rearrange(k, "b s h d -> (b s) h d")

304 + v_varlen = einops.rearrange(v, "b s h d -> (b s) h d")

305 + output = self._forward_window_attn(

306 + q_varlen, k_varlen, v_varlen, cu_seqlens, max_seqlen

307 + )

308 + # output: [total_tokens, num_heads, kv_channels]

309 + context_layer = einops.rearrange(

310 + output, "(b s) h d -> s b (h d)", b=batch_size

311 + ).contiguous()

312 +

313 + output, _ = self.proj(context_layer)

314 + return output

315 +

316 +

317 +class MiMoVisionBlock(nn.Module):

318 + def __init__(

319 + self,

320 + dim: int,

321 + num_heads: int,

322 + num_kv_heads: int,

323 + qk_channels: int,

324 + kv_channels: int,

325 + mlp_hidden_dim: int,

326 + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

327 + norm_eps: float = 1e-6,

328 + use_sink: bool = False,

329 + visual_token_window_size: int = 64,

330 + quant_config: QuantizationConfig | None = None,

331 + prefix: str = "",

332 + ) -> None:

333 + super().__init__()

334 + self.norm1 = RMSNorm(dim, eps=norm_eps)

335 + self.norm2 = RMSNorm(dim, eps=norm_eps)

336 + self.attn = MiMoVisionAttention(

337 + embed_dim=dim,

338 + num_heads=num_heads,

339 + num_kv_heads=num_kv_heads,

340 + qk_channels=qk_channels,

341 + kv_channels=kv_channels,

342 + use_sink=use_sink,

343 + visual_token_window_size=visual_token_window_size,

344 + quant_config=quant_config,

345 + prefix=f"{prefix}.attn",

346 + )

347 + self.mlp = MiMoVisionMLP(

348 + in_features=dim,

349 + hidden_features=mlp_hidden_dim,

350 + act_fn=act_fn,

351 + bias=True,

352 + quant_config=quant_config,

353 + prefix=f"{prefix}.mlp",

354 + )

355 +

356 + def forward(

357 + self,

358 + x: torch.Tensor,

359 + cu_seqlens: torch.Tensor,

360 + rotary_pos_emb_cos: torch.Tensor,

361 + rotary_pos_emb_sin: torch.Tensor,

362 + max_seqlen: torch.Tensor,

363 + full_attn: bool = True,

364 + ) -> torch.Tensor:

365 + # x: [seq_len, batch=1, dim]

366 + x_attn = self.attn(

367 + self.norm1(x),

368 + cu_seqlens=cu_seqlens,

369 + rotary_pos_emb_cos=rotary_pos_emb_cos,

370 + rotary_pos_emb_sin=rotary_pos_emb_sin,

371 + max_seqlen=max_seqlen,

372 + full_attn=full_attn,

373 + )

374 + # Fused residual add + norm2

375 + x_norm, residual = self.norm2(x, residual=x_attn)

376 + x = residual + self.mlp(x_norm)

377 + return x

378 +

379 +

380 +class MiMoVisionTransformer(nn.Module):

381 + def __init__(

382 + self,

383 + vision_cfg: PretrainedConfig,

384 + *,

385 + norm_eps: float = 1e-6,

386 + quant_config: QuantizationConfig | None = None,

387 + prefix: str = "",

388 + ):

389 + super().__init__()

390 + self.spatial_merge_size = vision_cfg.spatial_merge_size

391 + self.spatial_merge_unit = self.spatial_merge_size**2

392 + self.fullatt_block_indexes = vision_cfg.fullatt_block_indexes

393 + self.vit_window_attn_types = vision_cfg.vit_window_attn_types

394 + self.visual_token_window_size = vision_cfg.visual_token_window_size

395 + self.hidden_size = vision_cfg.hidden_size

396 + self.num_heads = vision_cfg.num_heads

397 + self.num_kv_heads = vision_cfg.num_key_value_heads

398 + self.qk_channels = vision_cfg.qk_channels

399 + self.kv_channels = vision_cfg.kv_channels

400 +

401 + self.patch_embed = MiMoVisionPatchEmbed(

402 + patch_size=vision_cfg.patch_size,

403 + temporal_patch_size=vision_cfg.temporal_patch_size,

404 + in_channels=vision_cfg.in_channels,

405 + hidden_size=vision_cfg.hidden_size,

406 + )

407 +

408 + norm_layer = partial(RMSNorm, eps=norm_eps)

409 +

410 + # Rotary embedding for 2D positions.

411 + # With partial_rotary_factor=0.5 and head_size=qk_channels:

412 + # rotary_dim = qk_channels // 2

413 + # get_cos_sin returns cos, sin each of shape [pos, rotary_dim // 2]

414 + # After indexing with 2D pos_ids and flattening:

415 + # result shape = [tokens, rotary_dim] = [tokens, qk_channels // 2]

416 + # which is what ApplyRotaryEmb expects as cos/sin input.

417 + self.rotary_pos_emb = get_rope(

418 + head_size=vision_cfg.qk_channels,

419 + max_position=8192,

420 + is_neox_style=True,

421 + rope_parameters={"partial_rotary_factor": 0.5},

422 + )

423 +

424 + self.blocks = nn.ModuleList(

425 + [

426 + MiMoVisionBlock(

427 + dim=vision_cfg.hidden_size,

428 + num_heads=vision_cfg.num_heads,

429 + num_kv_heads=vision_cfg.num_key_value_heads,

430 + qk_channels=vision_cfg.qk_channels,

431 + kv_channels=vision_cfg.kv_channels,

432 + mlp_hidden_dim=vision_cfg.intermediate_size,

433 + act_fn=get_act_and_mul_fn(vision_cfg.hidden_act),

434 + norm_eps=norm_eps,

435 + use_sink=(

436 + vision_cfg.use_sink

437 + and i not in vision_cfg.fullatt_block_indexes

438 + ),

439 + visual_token_window_size=vision_cfg.visual_token_window_size,

440 + quant_config=quant_config,

441 + prefix=f"{prefix}.blocks.{i}",

442 + )

443 + for i in range(vision_cfg.depth)

444 + ]

445 + )

446 +

447 + self.merger = MiMoVisionPatchMerger(

448 + d_model=vision_cfg.out_hidden_size,

449 + context_dim=vision_cfg.hidden_size,

450 + norm_layer=norm_layer,

451 + spatial_merge_size=vision_cfg.spatial_merge_size,

452 + quant_config=quant_config,

453 + prefix=f"{prefix}.merger",

454 + )

455 +

456 + @property

457 + def dtype(self) -> torch.dtype:

458 + return self.patch_embed.proj.weight.dtype

459 +

460 + @property

461 + def device(self) -> torch.device:

462 + return self.patch_embed.proj.weight.device

463 +

464 + def apply_index(self, tensor: torch.Tensor, index: torch.Tensor) -> torch.Tensor:

465 + """Reindex tensor at the spatial_merge_unit granularity."""

466 + tensor = tensor.unflatten(0, (-1, self.spatial_merge_unit))

467 + tensor = tensor[index]

468 + tensor = tensor.flatten(0, 1)

469 + return tensor

470 +

471 + def get_window_index_1d(

472 + self, grid_thw: torch.Tensor, col: bool = True

473 + ) -> torch.Tensor:

474 + """Compute 1D window indices for col-based or row-based SWA reordering."""

475 + window_index: list[torch.Tensor] = []

476 + window_index_id = 0

477 + for grid_t, grid_h, grid_w in grid_thw:

478 + llm_grid_h = grid_h // self.spatial_merge_size

479 + llm_grid_w = grid_w // self.spatial_merge_size

480 + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(

481 + grid_t, llm_grid_h, llm_grid_w

482 + )

483 + index_new = index.transpose(1, 2).reshape(-1) if col else index.reshape(-1)

484 + window_index.append(index_new + window_index_id)

485 + window_index_id += int((grid_t * llm_grid_h * llm_grid_w).item())

486 + return torch.cat(window_index, dim=0)

487 +

488 + def rot_pos_emb(self, grid_thw: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

489 + """Compute 2D rotary position embedding cos/sin for given grid sizes.

490 +

491 + Returns:

492 + cos: [total_tokens, qk_channels // 2]

493 + sin: [total_tokens, qk_channels // 2]

494 + """

495 + cos_list, sin_list = [], []

496 + for i in range(grid_thw.size(0)):

497 + t, h, w = int(grid_thw[i, 0]), int(grid_thw[i, 1]), int(grid_thw[i, 2])

498 +

499 + # Build 2D position IDs with spatial_merge_size interleaving

500 + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)

501 + hpos_ids = (

502 + hpos_ids.reshape(

503 + h // self.spatial_merge_size,

504 + self.spatial_merge_size,

505 + w // self.spatial_merge_size,

506 + self.spatial_merge_size,

507 + )

508 + .permute(0, 2, 1, 3)

509 + .flatten()

510 + )

511 + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

512 + wpos_ids = (

513 + wpos_ids.reshape(

514 + h // self.spatial_merge_size,

515 + self.spatial_merge_size,

516 + w // self.spatial_merge_size,

517 + self.spatial_merge_size,

518 + )

519 + .permute(0, 2, 1, 3)

520 + .flatten()

521 + )

522 + pos_ids = torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)

523 + # pos_ids: [t*h*w, 2]

524 +

525 + max_grid_size = max(h, w)

526 + # get_cos_sin returns cos, sin each of shape [max_grid_size, rotary_dim//2]

527 + # where rotary_dim = qk_channels // 2 (from partial_rotary_factor=0.5)

528 + cos, sin = self.rotary_pos_emb.get_cos_sin(max_grid_size)

529 +

530 + # [t*h*w, 2, rotary_dim//2] -> [t*h*w, rotary_dim] (= qk_channels // 2)

531 + cos_img = cos[pos_ids].flatten(1)

532 + sin_img = sin[pos_ids].flatten(1)

533 + cos_list.append(cos_img)

534 + sin_list.append(sin_img)

535 +

536 + return torch.cat(cos_list, dim=0), torch.cat(sin_list, dim=0)

537 +

538 + def forward(self, x: torch.Tensor, grid_thw: torch.Tensor) -> torch.Tensor:

539 + """

540 + Args:

541 + x: [total_tokens, C] pre-flattened patches

542 + grid_thw: [num_images, 3] tensor of (t, h, w) for each image/video

543 + Returns:

544 + [merged_tokens, out_hidden_size]

545 + """

546 + # Ensure grid_thw is a tensor

547 + if not isinstance(grid_thw, torch.Tensor):

548 + grid_thw = torch.tensor(grid_thw, dtype=torch.long)

549 +

550 + # Move to visual model device/dtype

551 + x = x.to(device=self.device, dtype=self.dtype)

552 +

553 + # Patch embedding: [total_tokens, hidden_size]

554 + x = self.patch_embed(x)

555 +

556 + # Compute 2D rotary positional embeddings

557 + # cos, sin: [total_tokens, qk_channels // 2]

558 + rotary_cos, rotary_sin = self.rot_pos_emb(grid_thw)

559 + rotary_cos = rotary_cos.to(device=x.device)

560 + rotary_sin = rotary_sin.to(device=x.device)

561 +

562 + # Compute cu_seqlens for flash_attn (per-image/video sequence lengths)

563 + seqlens = torch.repeat_interleave(

564 + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

565 + )

566 + cu_seqlens = torch.cat(

567 + [

568 + torch.tensor([0], device=x.device, dtype=torch.int32),

569 + seqlens.cumsum(dim=0).to(device=x.device, dtype=torch.int32),

570 + ]

571 + )

572 + max_seqlen = seqlens.max()

573 +

574 + # Precompute col-based window index for type=1 (col SWA) layers

575 + window_index_1d_col = self.get_window_index_1d(grid_thw, col=True).to(

576 + device=x.device

577 + )

578 + reverse_window_index_1d_col = torch.argsort(window_index_1d_col)

579 +

580 + # Col-based rotary embeddings (reordered at spatial_merge_unit granularity).

581 + # apply_index reorders groups of spatial_merge_unit tokens, just like x.

582 + col_cos = self.apply_index(rotary_cos, window_index_1d_col)

583 + col_sin = self.apply_index(rotary_sin, window_index_1d_col)

584 +

585 + # Add batch dimension: [total_tokens, 1, hidden_size]

586 + x = x.unsqueeze(1)

587 +

588 + for i, blk in enumerate(self.blocks):

589 + window_attn_type = self.vit_window_attn_types[i]

590 +

591 + # Reorder tokens to col-based layout when entering col-SWA region

592 + if window_attn_type == 1 and (

593 + i == 0 or self.vit_window_attn_types[i - 1] != 1

594 + ):

595 + x = self.apply_index(x, window_index_1d_col)

596 +

597 + # Restore row-based order when leaving col-SWA region

598 + if (

599 + i > 0

600 + and window_attn_type != 1

601 + and self.vit_window_attn_types[i - 1] == 1

602 + ):

603 + x = self.apply_index(x, reverse_window_index_1d_col)

604 +

605 + # Use col-based embeddings for col-SWA layers

606 + cos_now = col_cos if window_attn_type == 1 else rotary_cos

607 + sin_now = col_sin if window_attn_type == 1 else rotary_sin

608 +

609 + full_attn = i in self.fullatt_block_indexes

610 + x = blk(

611 + x,

612 + cu_seqlens=cu_seqlens,

613 + rotary_pos_emb_cos=cos_now,

614 + rotary_pos_emb_sin=sin_now,

615 + max_seqlen=max_seqlen,

616 + full_attn=full_attn,

617 + )

618 +

619 + # Restore row-based order if last block was col-SWA

620 + if self.vit_window_attn_types[-1] == 1:

621 + x = self.apply_index(x, reverse_window_index_1d_col)

622 +

623 + # Remove batch dim and merge spatial tokens

624 + # x: [total_tokens, 1, hidden_size] -> [total_tokens, hidden_size]

625 + x = x.squeeze(1)

626 + x = self.merger(x)

627 + return x

628 +

629 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

630 + stacked_params_mapping = [

631 + ("mlp.gate_up_proj", "mlp.gate_proj", 0),

632 + ("mlp.gate_up_proj", "mlp.up_proj", 1),

633 + ]

634 + params_dict = dict(self.named_parameters(remove_duplicate=False))

635 + loaded_params: set[str] = set()

636 +

637 + for name, loaded_weight in weights:

638 + for param_name, weight_name, shard_id in stacked_params_mapping:

639 + if weight_name not in name:

640 + continue

641 + name = name.replace(weight_name, param_name)

642 + param = params_dict[name]

643 + weight_loader = param.weight_loader

644 + weight_loader(param, loaded_weight, shard_id)

645 + break

646 + else:

647 + param = params_dict[name]

648 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

649 + weight_loader(param, loaded_weight)

650 + loaded_params.add(name)

651 + return loaded_params

652 +

653 +

654 +class MiMoV2OmniProcessingInfo(BaseProcessingInfo):

655 + def get_supported_mm_limits(self) -> Mapping[str, int | None]:

656 + return {"audio": None, "image": None, "video": None}

657 +

658 + def get_hf_config(self):

659 + config = self.ctx.get_hf_config()

660 + if isinstance(config.vision_config, dict):

661 + config.vision_config = Mimo_VLVisionConfig.from_dict(config.vision_config)

662 + return config

663 +

664 + def get_hf_processor(self, **kwargs: object) -> MiMoOmniProcessor:

665 + hf_config = self.get_hf_config()

666 + tokenizer = self.get_tokenizer()

667 + return MiMoOmniProcessor.from_hf_config(tokenizer, hf_config)

668 +

669 + def get_image_processor(self, **kwargs: object):

670 + return self.get_hf_processor(**kwargs).image_processor

671 +

672 + def get_data_parser(self):

673 + from vllm.multimodal.parse import MultiModalDataParser

674 +

675 + return MultiModalDataParser(target_sr=24000.0)

676 +

677 + def get_mm_max_tokens_per_item(

678 + self,

679 + seq_len: int,

680 + mm_counts: Mapping[str, int],

681 + ) -> Mapping[str, int]:

682 + return {

683 + "image": self.get_max_image_tokens(),

684 + "video": self.get_max_video_tokens(seq_len, mm_counts),

685 + }

686 +

687 + def _get_vision_info(

688 + self,

689 + *,

690 + image_width: int,

691 + image_height: int,

692 + num_frames: int = 1,

693 + do_resize: bool = True,

694 + image_processor,

695 + mm_kwargs: Mapping[str, object],

696 + ) -> tuple[ImageSize, int]:

697 + hf_config = self.get_hf_config()

698 + vision_config = hf_config.vision_config

699 + patch_size = vision_config.patch_size

700 + merge_size = vision_config.spatial_merge_size

701 + temporal_patch_size = vision_config.temporal_patch_size

702 + tokens_per_second = vision_config.tokens_per_second

703 +

704 + mm_kwargs = self.ctx.get_merged_mm_kwargs(mm_kwargs)

705 + size = image_processor.size

706 + if override_size := mm_kwargs.get("size"):

707 + size = size | override_size

708 + if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None:

709 + size = size | {"shortest_edge": override_min_pixels}

710 + if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None:

711 + size = size | {"longest_edge": override_max_pixels}

712 +

713 + if do_resize:

714 + resized_height, resized_width = smart_resize(

715 + height=image_height,

716 + width=image_width,

717 + factor=patch_size * merge_size,

718 + min_pixels=size["shortest_edge"],

719 + max_pixels=size["longest_edge"],

720 + )

721 + preprocessed_size = ImageSize(width=resized_width, height=resized_height)

722 + else:

723 + preprocessed_size = ImageSize(width=image_width, height=image_height)

724 +

725 + # For video, MiMo resamples to tokens_per_second fps before temporal patching,

726 + # effective tokens = num_frames * tokens_per_second / temporal_patch_size.

727 + # For images (num_frames == 1) no resampling is applied.

728 + if num_frames > 1:

729 + effective_frames = num_frames * tokens_per_second

730 + else:

731 + effective_frames = num_frames

732 + padded_num_frames = effective_frames + effective_frames % temporal_patch_size

733 + grid_t = max(padded_num_frames // temporal_patch_size, 1)

734 + grid_h = preprocessed_size.height // patch_size

735 + grid_w = preprocessed_size.width // patch_size

736 + num_patches = grid_t * grid_h * grid_w

737 + num_vision_tokens = num_patches // (merge_size**2)

738 + return preprocessed_size, num_vision_tokens

739 +

740 + def get_num_image_tokens(

741 + self,

742 + *,

743 + image_width: int,

744 + image_height: int,

745 + image_processor,

746 + mm_kwargs: Mapping[str, object],

747 + ) -> int:

748 + _, num_image_tokens = self._get_vision_info(

749 + image_width=image_width,

750 + image_height=image_height,

751 + num_frames=1,

752 + image_processor=image_processor,

753 + mm_kwargs=mm_kwargs,

754 + )

755 + return num_image_tokens

756 +

757 + def get_num_video_tokens(

758 + self,

759 + *,

760 + image_width: int,

761 + image_height: int,

762 + num_frames: int,

763 + image_processor,

764 + mm_kwargs: Mapping[str, object],

765 + ) -> int:

766 + _, num_video_tokens = self._get_vision_info(

767 + image_width=image_width,

768 + image_height=image_height,

769 + num_frames=num_frames,

770 + image_processor=image_processor,

771 + mm_kwargs=mm_kwargs,

772 + )

773 + return num_video_tokens

774 +

775 + def get_image_size_with_most_features(

776 + self, max_pixels: int | None = None

777 + ) -> ImageSize:

778 + hf_config = self.get_hf_config()

779 + vision_config = hf_config.vision_config

780 + patch_size = vision_config.patch_size

781 + merge_size = vision_config.spatial_merge_size

782 +

783 + if max_pixels is None:

784 + image_processor = self.get_image_processor()

785 + mm_kwargs = self.ctx.get_merged_mm_kwargs({})

786 + size = image_processor.size

787 + if override_size := mm_kwargs.get("size"):

788 + size = size | override_size

789 + if (override_min_pixels := mm_kwargs.get("min_pixels")) is not None:

790 + size = size | {"shortest_edge": override_min_pixels}

791 + if (override_max_pixels := mm_kwargs.get("max_pixels")) is not None:

792 + size = size | {"longest_edge": override_max_pixels}

793 + max_pixels = size["longest_edge"]

794 +

795 + unit = patch_size * merge_size

796 + max_seq_len = max_pixels // (unit * unit)

797 +

798 + def closest_factor_pair(n: int) -> tuple[int, int]:

799 + for d in range(math.isqrt(n), 0, -1):

800 + if n % d == 0:

801 + return d, n // d

802 + return 1, n

803 +

804 + height_factor, width_factor = 1, max_seq_len

805 + for seq_len in range(max_seq_len, 0, -1):

806 + height_factor, width_factor = closest_factor_pair(seq_len)

807 + if width_factor / height_factor <= 200:

808 + break

809 +

810 + return ImageSize(width=unit * width_factor, height=unit * height_factor)

811 +

812 + def get_max_image_tokens(self) -> int:

813 + image_processor = self.get_image_processor()

814 + target_width, target_height = self.get_image_size_with_most_features()

815 + return self.get_num_image_tokens(

816 + image_width=target_width,

817 + image_height=target_height,

818 + image_processor=image_processor,

819 + mm_kwargs={},

820 + )

821 +

822 + def _get_max_video_frames(self, max_tokens: int, start_num_frames: int = 1) -> int:

823 + image_processor = self.get_image_processor()

824 + target_width, target_height = self.get_image_size_with_most_features()

825 + num_frames = start_num_frames

826 + while True:

827 + next_num_frames = num_frames + 1

828 + next_max_tokens = self.get_num_video_tokens(

829 + image_width=target_width,

830 + image_height=target_height,

831 + num_frames=next_num_frames,

832 + image_processor=image_processor,

833 + mm_kwargs={},

834 + )

835 + if next_max_tokens > max_tokens:

836 + break

837 + num_frames = next_num_frames

838 + return num_frames

839 +

840 + def get_num_frames_with_most_features(

841 + self,

842 + seq_len: int,

843 + mm_counts: Mapping[str, int],

844 + max_frames_per_video: int = 14,

845 + ) -> int:

846 + max_videos = mm_counts.get("video", 0)

847 + max_total_frames = self._get_max_video_frames(seq_len)

848 + max_frames_per_video = min(

849 + max_total_frames // max(max_videos, 1), max_frames_per_video

850 + )

851 + return max(max_frames_per_video, 1)

852 +

853 + def get_max_video_tokens(

854 + self,

855 + seq_len: int,

856 + mm_counts: Mapping[str, int],

857 + ) -> int:

858 + image_processor = self.get_image_processor()

859 + target_width, target_height = self.get_image_size_with_most_features()

860 + return self.get_num_video_tokens(

861 + image_width=target_width,

862 + image_height=target_height,

863 + num_frames=self.get_num_frames_with_most_features(seq_len, mm_counts),

864 + image_processor=image_processor,

865 + mm_kwargs={},

866 + )

867 +

868 +

869 +class MiMoV2OmniMultiModalProcessor(BaseMultiModalProcessor[MiMoV2OmniProcessingInfo]):

870 + """vLLM multimodal processor for MiMo-Omni (image + video).

871 +

872 + Key differences from Qwen2.5-VL:

873 + - Videos use timestamp tokens between temporal grid positions.

874 + - The HF processor expects ``(TCHW_tensor, timestamps_T_tensor)`` video

875 + tuples rather than plain numpy arrays.

876 + - ``video_start_times`` is tracked so prompt-update reconstruction can

877 + regenerate the exact same timestamp token IDs.

878 + """

879 +

880 + # fps assumed for vllm-decoded video (numpy T,H,W,C arrays).

881 + # The video loader samples ~32 frames; treat each frame as 1 s apart so

882 + # MiMoVLProcessor sees 1 fps input and resamples internally.

883 + _INPUT_FPS: float = 1.0

884 +

885 + def _get_mm_fields_config(

886 + self,

887 + hf_inputs: BatchFeature,

888 + hf_processor_mm_kwargs: Mapping[str, object],

889 + ) -> Mapping[str, MultiModalFieldConfig]:

890 + merge_size = self.info.get_hf_config().vision_config.spatial_merge_size

891 + fields: dict[str, MultiModalFieldConfig] = dict(

892 + **_create_qwen2vl_field_factory(merge_size)(hf_inputs),

893 + second_per_grid_ts=MultiModalFieldConfig.batched("video"),

894 + video_start_times=MultiModalFieldConfig.batched("video"),

895 + audio_features=MultiModalFieldConfig.batched("audio"),

896 + audio_token_lens=MultiModalFieldConfig.batched("audio"),

897 + )

898 + # video_audio fields: only present when video_audio content was processed

899 + if "video_audio_n_segs" in hf_inputs:

900 + fields["video_audio_n_segs"] = MultiModalFieldConfig.batched("video")

901 + # video_audio_seg_lens: list of per-video 1D tensors, batched("video")

902 + if "video_audio_seg_lens" in hf_inputs:

903 + fields["video_audio_seg_lens"] = MultiModalFieldConfig.batched("video")

904 + if "va_audio_features" in hf_inputs:

905 + fields["va_audio_features"] = MultiModalFieldConfig.batched("va_audio")

906 + return fields

907 +

908 + def _call_hf_processor(

909 + self,

910 + prompt: str,

911 + mm_data: Mapping[str, object],

912 + mm_kwargs: Mapping[str, object],

913 + tok_kwargs: Mapping[str, object],

914 + ) -> BatchFeature:

915 + """Convert numpy video arrays to (TCHW, timestamps) tuples for MiMo.

916 + Also remap 'audios' → 'audio' since MiMoOmniProcessor.__call__ uses

917 + the singular form.

918 + """

919 + # Remap audios → audio (MiMoOmniProcessor uses singular param name)

920 + if "audios" in mm_data:

921 + mm_data = {**mm_data, "audio": mm_data["audios"]}

922 + mm_data = {k: v for k, v in mm_data.items() if k != "audios"}

923 +

924 + # Handle video_audio items: convert video part to (TCHW, timestamps) tuple

925 + if "video_audio" in mm_data:

926 + va_converted: list[VideoAudioInput] = []

927 + for va_item in mm_data["video_audio"]:

928 + if isinstance(va_item, VideoAudioInput):

929 + vid = va_item.video

930 + else:

931 + # Expect (video_frames, audio_source) tuple

932 + vid, audio_src = va_item

933 + va_item = VideoAudioInput(video=vid, audio=audio_src)

934 + vid = vid

935 + # Convert video frames to (TCHW, timestamps) if needed

936 + if (

937 + isinstance(vid, tuple)

938 + and len(vid) == 2

939 + and isinstance(vid[0], torch.Tensor)

940 + and isinstance(vid[1], torch.Tensor)

941 + ):

942 + va_converted.append(va_item)

943 + else:

944 + if isinstance(vid, np.ndarray):

945 + frames = torch.from_numpy(vid)

946 + elif isinstance(vid, torch.Tensor):

947 + frames = vid

948 + else:

949 + frames = torch.tensor(np.array(vid))

950 + if frames.ndim == 4 and frames.shape[-1] in (1, 3, 4):

951 + frames = frames.permute(0, 3, 1, 2).float()

952 + else:

953 + frames = frames.float()

954 + T = frames.shape[0]

955 + timestamps = torch.arange(T, dtype=torch.float32) / self._INPUT_FPS

956 + va_converted.append(

957 + VideoAudioInput(

958 + video=(frames, timestamps),

959 + audio=va_item.audio,

960 + )

961 + )

962 + mm_data = {**mm_data, "video_audio": va_converted}

963 +

964 + if "videos" in mm_data:

965 + converted: list[tuple[torch.Tensor, torch.Tensor]] = []

966 + for video in mm_data["videos"]:

967 + if (

968 + isinstance(video, tuple)

969 + and len(video) == 2

970 + and isinstance(video[0], torch.Tensor)

971 + and isinstance(video[1], torch.Tensor)

972 + ):

973 + # already in MiMo format

974 + converted.append(video)

975 + else:

976 + # numpy (T, H, W, C) or torch (T, H, W, C) / (T, C, H, W)

977 + if isinstance(video, np.ndarray):

978 + frames = torch.from_numpy(video)

979 + elif isinstance(video, torch.Tensor):

980 + frames = video

981 + else:

982 + frames = torch.tensor(np.array(video))

983 +

984 + if frames.ndim == 4 and frames.shape[-1] in (1, 3, 4):

985 + # THWC → TCHW

986 + frames = frames.permute(0, 3, 1, 2).float()

987 + else:

988 + frames = frames.float()

989 +

990 + T = frames.shape[0]

991 + timestamps = torch.arange(T, dtype=torch.float32) / self._INPUT_FPS

992 + converted.append((frames, timestamps))

993 +

994 + mm_data = {**mm_data, "videos": converted}

995 +

996 + return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)

997 +

998 + def _get_prompt_updates(

999 + self,

1000 + mm_items: MultiModalDataItems,

1001 + hf_processor_mm_kwargs: Mapping[str, Any],

1002 + out_mm_kwargs: MultiModalKwargsItems,

1003 + ) -> Sequence[PromptUpdate]:

1004 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

1005 + hf_config = self.info.get_hf_config()

1006 + tokenizer = self.info.get_tokenizer()

1007 + vocab = tokenizer.get_vocab()

1008 +

1009 + merge_size = hf_config.vision_config.spatial_merge_size

1010 + p = hf_processor.mimo_processor

1011 +

1012 + image_pad_id = vocab[hf_processor.image_token]

1013 + video_pad_id = vocab[hf_processor.video_token]

1014 + audio_pad_id = vocab.get("<|audio_pad|>")

1015 + vision_start_id = p.vision_start_token_id

1016 + vision_end_id = p.vision_end_token_id

1017 + video_start_id = p.video_start_token_id

1018 + video_end_id = p.video_end_token_id

1019 + audio_start_id = p.audio_start_token_id

1020 + audio_end_id = p.audio_end_token_id

1021 +

1022 + def get_image_replacement(item_idx: int) -> PromptUpdateDetails:

1023 + out_item = out_mm_kwargs["image"][item_idx]

1024 + grid_thw = out_item["image_grid_thw"].data

1025 + n_tokens = int(grid_thw.prod()) // merge_size**2

1026 + return [image_pad_id] * n_tokens

1027 +

1028 + def get_video_replacement(item_idx: int) -> PromptUpdateDetails:

1029 + out_item = out_mm_kwargs["video"][item_idx]

1030 + grid_thw = out_item["video_grid_thw"].data

1031 + spt = float(out_item["second_per_grid_ts"].data)

1032 + start = float(out_item["video_start_times"].data)

1033 +

1034 + T, H, W = map(int, grid_thw)

1035 + n_per_grid = H * W // (merge_size * merge_size)

1036 +

1037 + # Check if this is a video_audio item

1038 + n_segs_field = out_item.get("video_audio_n_segs")

1039 + n_segs_val = int(n_segs_field.data) if n_segs_field is not None else 0

1040 + va_seg_lens: list[int] | None = None

1041 + if n_segs_val > 0:

1042 + seg_lens_field = out_item.get("video_audio_seg_lens")

1043 + if seg_lens_field is not None:

1044 + va_seg_lens = seg_lens_field.data[:n_segs_val].tolist()

1045 +

1046 + full: list[int] = [video_start_id]

1047 + is_embed_mask: list[bool] = [False]

1048 +

1049 + if va_seg_lens is None:

1050 + # Regular video: timestamp + vision tokens per grid

1051 + for j in range(T):

1052 + ts_text = _format_timestamp(start + j * spt)

1053 + ts_ids = tokenizer.encode(ts_text, add_special_tokens=False)

1054 + full.extend(ts_ids)

1055 + is_embed_mask.extend([False] * len(ts_ids))

1056 + full.append(vision_start_id)

1057 + is_embed_mask.append(False)

1058 + full.extend([video_pad_id] * n_per_grid)

1059 + is_embed_mask.extend([True] * n_per_grid)

1060 + full.append(vision_end_id)

1061 + is_embed_mask.append(False)

1062 + else:

1063 + # video_audio: interleaved vision+audio per group

1064 + n_groups = len(va_seg_lens)

1065 + frames_per_group = T // n_groups # 1 for il=0, T for il=-1

1066 + for g in range(n_groups):

1067 + # Timestamp for first frame of this group

1068 + frame0 = g * frames_per_group

1069 + ts_text = _format_timestamp(start + frame0 * spt)

1070 + ts_ids = tokenizer.encode(ts_text, add_special_tokens=False)

1071 + full.extend(ts_ids)

1072 + is_embed_mask.extend([False] * len(ts_ids))

1073 + # Vision tokens for all frames in this group

1074 + for f in range(frames_per_group):

1075 + full.append(vision_start_id)

1076 + is_embed_mask.append(False)

1077 + full.extend([video_pad_id] * n_per_grid)

1078 + is_embed_mask.extend([True] * n_per_grid)

1079 + full.append(vision_end_id)

1080 + is_embed_mask.append(False)

1081 + # Audio tokens for this group

1082 + seg_len = va_seg_lens[g]

1083 + full.append(audio_start_id)

1084 + is_embed_mask.append(False)

1085 + full.extend([audio_pad_id] * seg_len)

1086 + is_embed_mask.extend([True] * seg_len)

1087 + full.append(audio_end_id)

1088 + is_embed_mask.append(False)

1089 +

1090 + full.append(video_end_id)

1091 + is_embed_mask.append(False)

1092 +

1093 + embed_t = torch.tensor(is_embed_mask)

1094 + return PromptUpdateDetails(

1095 + full=full,

1096 + is_embed=lambda _tok, _seq: embed_t,

1097 + )

1098 +

1099 + def get_audio_replacement(item_idx: int) -> PromptUpdateDetails:

1100 + out_item = out_mm_kwargs["audio"][item_idx]

1101 + tok_len = int(out_item["audio_token_lens"].data)

1102 + return [audio_pad_id] * tok_len

1103 +

1104 + updates: list[PromptUpdate] = [

1105 + PromptReplacement(

1106 + modality="image",

1107 + target=[image_pad_id],

1108 + replacement=get_image_replacement,

1109 + ),

1110 + PromptReplacement(

1111 + modality="video",

1112 + target=[video_pad_id],

1113 + replacement=get_video_replacement,

1114 + ),

1115 + ]

1116 + if audio_pad_id is not None and audio_start_id is not None:

1117 + updates.append(

1118 + PromptReplacement(

1119 + modality="audio",

1120 + target=[audio_pad_id],

1121 + replacement=get_audio_replacement,

1122 + )

1123 + )

1124 + return updates

1125 +

1126 +

1127 +class MiMoV2OmniDummyInputsBuilder(BaseDummyInputsBuilder[MiMoV2OmniProcessingInfo]):

1128 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

1129 + num_images = mm_counts.get("image", 0)

1130 + num_videos = mm_counts.get("video", 0)

1131 + num_audios = mm_counts.get("audio", 0)

1132 + image_ph = "<|vision_start|><|image_pad|><|vision_end|>"

1133 + video_ph = "<|vision_start|><|video_pad|><|vision_end|>"

1134 + audio_ph = "<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>"

1135 + return image_ph * num_images + video_ph * num_videos + audio_ph * num_audios

1136 +

1137 + def get_dummy_mm_data(

1138 + self,

1139 + seq_len: int,

1140 + mm_counts: Mapping[str, int],

1141 + mm_options: Mapping[str, BaseDummyOptions],

1142 + ) -> MultiModalDataDict:

1143 + num_images = mm_counts.get("image", 0)

1144 + num_videos = mm_counts.get("video", 0)

1145 +

1146 + target_width, target_height = self.info.get_image_size_with_most_features()

1147 + target_num_frames = self.info.get_num_frames_with_most_features(

1148 + seq_len, mm_counts

1149 + )

1150 +

1151 + return {

1152 + "image": self._get_dummy_images(

1153 + width=target_width,

1154 + height=target_height,

1155 + num_images=num_images,

1156 + overrides=mm_options.get("image"),

1157 + ),

1158 + "video": self._get_dummy_videos(

1159 + width=target_width,

1160 + height=target_height,

1161 + num_frames=target_num_frames,

1162 + num_videos=num_videos,

1163 + overrides=mm_options.get("video"),

1164 + ),

1165 + }

1166 +

1167 +

1168 +@MULTIMODAL_REGISTRY.register_processor(

1169 + MiMoV2OmniMultiModalProcessor,

1170 + info=MiMoV2OmniProcessingInfo,

1171 + dummy_inputs=MiMoV2OmniDummyInputsBuilder,

1172 +)

1173 +class MiMoV2OmniForCausalLM(nn.Module, SupportsMultiModal, SupportsPP, SupportsQuant):

1174 + # To ensure correct weight loading and mapping.

1175 + hf_to_vllm_mapper = WeightsMapper(

1176 + orig_to_new_prefix={

1177 + # audio encoder

1178 + "speech_embeddings.": "audio_encoder.speech_embeddings.",

1179 + # mapping for new names in checkpoint saved after transformers v4.52

1180 + "model.language_model.": "language_model.model.",

1181 + "model.visual.": "visual.",

1182 + # mapping for original checkpoint

1183 + "lm_head.": "language_model.lm_head.",

1184 + "model.": "language_model.model.",

1185 + }

1186 + )

1187 +

1188 + @classmethod

1189 + def get_placeholder_str(cls, modality: str, i: int) -> str | None:

1190 + if modality.startswith("image"):

1191 + return "<|vision_start|><|image_pad|><|vision_end|>"

1192 + if modality.startswith("video"):

1193 + return "<|vision_start|><|video_pad|><|vision_end|>"

1194 + if modality.startswith("audio"):

1195 + return "<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>"

1196 +

1197 + raise ValueError(f"Unsupported modality: {modality}")

1198 +

1199 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1200 + super().__init__()

1201 + config = vllm_config.model_config.hf_config

1202 + self.config = config

1203 + # Omni ViT/Audio Encoder BF16

1204 + vision_config = (

1205 + Mimo_VLVisionConfig.from_dict(config.vision_config)

1206 + if isinstance(config.vision_config, dict)

1207 + else config.vision_config

1208 + )

1209 + with self._mark_tower_model(vllm_config, {"image", "video"}):

1210 + self.visual = MiMoVisionTransformer(

1211 + vision_config,

1212 + norm_eps=getattr(vllm_config, "rms_norm_eps", 1e-6),

1213 + quant_config=None,

1214 + prefix=maybe_prefix(prefix, "visual"),

1215 + )

1216 + audio_config = getattr(config, "audio_config", None)

1217 + model_path = vllm_config.model_config.model

1218 + if audio_config is not None:

1219 + with self._mark_tower_model(vllm_config, "audio"):

1220 + self.audio_encoder = MimoAudioEncoder(

1221 + audio_config, model_path=model_path

1222 + )

1223 + else:

1224 + self.audio_encoder = None

1225 + with self._mark_language_model(vllm_config):

1226 + self.language_model = MiMoV2FlashForCausalLM(

1227 + vllm_config=vllm_config,

1228 + prefix=maybe_prefix(prefix, "language_model"),

1229 + )

1230 +

1231 + self.make_empty_intermediate_tensors = (

1232 + self.language_model.make_empty_intermediate_tensors

1233 + )

1234 +

1235 + def _parse_and_validate_image_input(

1236 + self, **kwargs: object

1237 + ) -> Qwen2_5_VLImageInputs | None:

1238 + pixel_values = kwargs.pop("pixel_values", None)

1239 + image_embeds = kwargs.pop("image_embeds", None)

1240 + image_grid_thw = kwargs.pop("image_grid_thw", None)

1241 +

1242 + if pixel_values is None and image_embeds is None:

1243 + return None

1244 +

1245 + if pixel_values is not None:

1246 + return Qwen2_5_VLImagePixelInputs(

1247 + type="pixel_values",

1248 + pixel_values=pixel_values,

1249 + image_grid_thw=image_grid_thw,

1250 + )

1251 +

1252 + if image_embeds is not None:

1253 + return Qwen2_5_VLImageEmbeddingInputs(

1254 + type="image_embeds",

1255 + image_embeds=image_embeds,

1256 + image_grid_thw=image_grid_thw,

1257 + )

1258 +

1259 + def _parse_and_validate_video_input(

1260 + self, **kwargs: object

1261 + ) -> Qwen2_5_VLVideoInputs | None:

1262 + pixel_values_videos = kwargs.pop("pixel_values_videos", None)

1263 + video_embeds = kwargs.pop("video_embeds", None)

1264 + video_grid_thw = kwargs.pop("video_grid_thw", None)

1265 + second_per_grid_ts = kwargs.pop("second_per_grid_ts", None)

1266 +

1267 + if pixel_values_videos is None and video_embeds is None:

1268 + return None

1269 +

1270 + if pixel_values_videos is not None:

1271 + return Qwen2_5_VLVideoPixelInputs(

1272 + type="pixel_values_videos",

1273 + pixel_values_videos=pixel_values_videos,

1274 + video_grid_thw=video_grid_thw,

1275 + second_per_grid_ts=second_per_grid_ts,

1276 + )

1277 +

1278 + if video_embeds is not None:

1279 + return Qwen2_5_VLVideoEmbeddingInputs(

1280 + type="video_embeds",

1281 + video_embeds=video_embeds,

1282 + video_grid_thw=video_grid_thw,

1283 + second_per_grid_ts=second_per_grid_ts,

1284 + )

1285 +

1286 + def _process_image_input(

1287 + self, image_input: Qwen2_5_VLImageInputs

1288 + ) -> tuple[torch.Tensor, ...]:

1289 + grid_thw = image_input["image_grid_thw"]

1290 + assert grid_thw.ndim == 2

1291 + grid_thw_list = grid_thw.tolist()

1292 +

1293 + if image_input["type"] == "image_embeds":

1294 + image_embeds = image_input["image_embeds"].type(self.visual.dtype)

1295 + else:

1296 + pixel_values = image_input["pixel_values"]

1297 + image_embeds = self.visual(pixel_values, grid_thw=grid_thw_list)

1298 +

1299 + # Split concatenated embeddings for each image item.

1300 + merge_size = self.visual.spatial_merge_size

1301 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1302 + return image_embeds.split(sizes)

1303 +

1304 + def _process_video_input(

1305 + self, video_input: Qwen2_5_VLVideoInputs

1306 + ) -> tuple[torch.Tensor, ...]:

1307 + grid_thw = video_input["video_grid_thw"]

1308 + assert grid_thw.ndim == 2

1309 + grid_thw_list = grid_thw.tolist()

1310 +

1311 + if video_input["type"] == "video_embeds":

1312 + video_embeds = video_input["video_embeds"].type(self.visual.dtype)

1313 + else:

1314 + pixel_values_videos = video_input["pixel_values_videos"]

1315 + video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw_list)

1316 +

1317 + # Split concatenated embeddings for each video item.

1318 + merge_size = self.visual.spatial_merge_size

1319 + sizes = (grid_thw.prod(-1) // merge_size // merge_size).tolist()

1320 + return video_embeds.split(sizes)

1321 +

1322 + def _parse_and_validate_audio_input(self, **kwargs: object) -> dict | None:

1323 + audio_features = kwargs.pop("audio_features", None)

1324 + audio_token_lens = kwargs.pop("audio_token_lens", None)

1325 + if audio_features is None:

1326 + return None

1327 + return {

1328 + "type": "audio",

1329 + "audio_features": audio_features,

1330 + "audio_token_lens": audio_token_lens,

1331 + }

1332 +

1333 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1334 + mm_input_by_modality = {}

1335 +

1336 + # Preserve the order of modalities if there are multiple of them

1337 + # from the order of kwargs.

1338 + for input_key in kwargs:

1339 + if (

1340 + input_key in ("pixel_values", "image_embeds")

1341 + and "image" not in mm_input_by_modality

1342 + ):

1343 + mm_input_by_modality["image"] = self._parse_and_validate_image_input(

1344 + **kwargs

1345 + )

1346 + if (

1347 + input_key in ("pixel_values_videos", "video_embeds")

1348 + and "video" not in mm_input_by_modality

1349 + ):

1350 + mm_input_by_modality["video"] = self._parse_and_validate_video_input(

1351 + **kwargs

1352 + )

1353 + if input_key == "audio_features" and "audio" not in mm_input_by_modality:

1354 + mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(

1355 + **kwargs

1356 + )

1357 + return mm_input_by_modality

1358 +

1359 + def _process_audio_input(self, audio_input: dict) -> tuple[torch.Tensor, ...]:

1360 + mel_specs = audio_input["audio_features"]

1361 + if self.audio_encoder is None:

1362 + return ()

1363 + # Normalize to List[2D-Tensor].

1364 + # MultiModalBatchedField._reduce_data either wraps a single [T, 128]

1365 + # into [1, T, 128] via unsqueeze(0) or stacks N same-T items into

1366 + # [N, T, 128]. Indexing along dim-0 extracts the per-item [T, 128].

1367 + if isinstance(mel_specs, torch.Tensor):

1368 + mel_specs = list(mel_specs) # [1,T,128] or [N,T,128] → [[T,128],...]

1369 + if not mel_specs:

1370 + return ()

1371 + audio_embeds, item_token_lens = self.audio_encoder.get_audio_feature(mel_specs)

1372 + return tuple(audio_embeds.split(item_token_lens))

1373 +

1374 + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:

1375 + # Pop video_audio-specific fields before main mm parsing

1376 + video_audio_n_segs = kwargs.pop("video_audio_n_segs", None)

1377 + video_audio_seg_lens = kwargs.pop("video_audio_seg_lens", None)

1378 + va_audio_features = kwargs.pop("va_audio_features", None)

1379 +

1380 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

1381 + if not mm_input_by_modality and va_audio_features is None:

1382 + return []

1383 +

1384 + # The result multimodal_embeddings is tuple of tensors, with each

1385 + # tensor corresponding to a multimodal data item (image, video, or audio).

1386 + multimodal_embeddings: list[torch.Tensor] = []

1387 +

1388 + # Pre-process va audio: one mel spec per va video → per-video audio embeddings

1389 + # keyed by va video index (0-based among va videos only)

1390 + va_audio_embs_list: list[tuple[torch.Tensor, ...]] = []

1391 + if va_audio_features is not None and self.audio_encoder is not None:

1392 + mel_list = (

1393 + list(va_audio_features)

1394 + if isinstance(va_audio_features, torch.Tensor)

1395 + else list(va_audio_features)

1396 + )

1397 + for mel_spec in mel_list:

1398 + embs, tok_lens = self.audio_encoder.get_audio_feature([mel_spec])

1399 + # tok_lens is a list/tensor with one entry (total tokens for this mel)

1400 + va_audio_embs_list.append(embs) # shape (total_tok, hidden)

1401 +

1402 + va_cursor = 0 # index into va_audio_embs_list

1403 +

1404 + # NOTE: Iterate in dict insertion order to preserve token sequence order.

1405 + for modality in mm_input_by_modality:

1406 + multimodal_input = mm_input_by_modality[modality]

1407 + if modality == "image":

1408 + multimodal_embeddings.extend(

1409 + self._process_image_input(multimodal_input)

1410 + )

1411 + elif modality == "video":

1412 + video_embs_tuple = self._process_video_input(multimodal_input)

1413 + if video_audio_n_segs is None:

1414 + multimodal_embeddings.extend(video_embs_tuple)

1415 + else:

1416 + grid_thw = multimodal_input["video_grid_thw"]

1417 + for i, vid_embs in enumerate(video_embs_tuple):

1418 + n_segs = int(video_audio_n_segs[i])

1419 + if n_segs == 0 or not va_audio_embs_list:

1420 + multimodal_embeddings.append(vid_embs)

1421 + else:

1422 + T = int(grid_thw[i][0])

1423 + n_per_grid = vid_embs.shape[0] // T

1424 + frames = list(vid_embs.split(n_per_grid, dim=0))

1425 + frames_per_group = T // n_segs

1426 + # Per-group audio token lengths for this va video

1427 + # video_audio_seg_lens is (num_videos, max_T); row i

1428 + # has valid values in [:n_segs], rest are zeros.

1429 + seg_lens = video_audio_seg_lens[i][:n_segs].tolist()

1430 + # Split full audio embs for this va video by group lengths

1431 + full_va_embs = va_audio_embs_list[va_cursor]

1432 + va_cursor += 1

1433 + group_audio_embs = full_va_embs.split(seg_lens)

1434 + # Interleave: all vid frames in group, then audio for group

1435 + for g in range(n_segs):

1436 + for f in range(frames_per_group):

1437 + multimodal_embeddings.append(

1438 + frames[g * frames_per_group + f]

1439 + )

1440 + multimodal_embeddings.append(group_audio_embs[g])

1441 + elif modality == "audio":

1442 + multimodal_embeddings.extend(

1443 + self._process_audio_input(multimodal_input)

1444 + )

1445 + return tuple(multimodal_embeddings)

1446 +

1447 + def forward(

1448 + self,

1449 + input_ids: torch.Tensor | None,

1450 + positions: torch.Tensor,

1451 + intermediate_tensors: IntermediateTensors | None = None,

1452 + inputs_embeds: torch.Tensor | None = None,

1453 + **kwargs: object,

1454 + ) -> torch.Tensor | IntermediateTensors:

1455 + """Run forward pass for Qwen2.5-VL.

1456 +

1457 + Args:

1458 + input_ids: Flattened (concatenated) input_ids corresponding to a

1459 + batch.

1460 + positions: Flattened (concatenated) position ids corresponding to a

1461 + batch. **NOTE**: If mrope is enabled (default setting for

1462 + Qwen2.5-VL opensource models), the shape will be `(3, seq_len)`,

1463 + otherwise it will be `(seq_len,).

1464 + """

1465 +

1466 + if intermediate_tensors is not None:

1467 + inputs_embeds = None

1468 +

1469 + hidden_states = self.language_model.model(

1470 + input_ids=input_ids,

1471 + positions=positions,

1472 + intermediate_tensors=intermediate_tensors,

1473 + inputs_embeds=inputs_embeds,

1474 + )

1475 + return hidden_states

1476 +

1477 + def compute_logits(

1478 + self,

1479 + hidden_states: torch.Tensor,

1480 + ) -> torch.Tensor | None:

1481 + return self.language_model.compute_logits(hidden_states)

1482 +

1483 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1484 + audio_loaded: set[str] = set()

1485 +

1486 + loader = AutoWeightsLoader(self, skip_prefixes=["audio_tokenizer."])

1487 + auto_loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1488 + return audio_loaded | auto_loaded

vllm/model_executor/models/registry.py

+5 -1

@@ -171,7 +171,8 @@ _TEXT_GENERATION_MODELS = {

171 "MptForCausalLM": ("mpt", "MPTForCausalLM"),

172 "MPTForCausalLM": ("mpt", "MPTForCausalLM"),

173 "MiMoForCausalLM": ("mimo", "MiMoForCausalLM"),

174 - "MiMoV2FlashForCausalLM": ("mimo_v2_flash", "MiMoV2FlashForCausalLM"),

174 + "MiMoV2FlashForCausalLM": ("mimo_v2", "MiMoV2FlashForCausalLM"),

175 + "MiMoV2ProForCausalLM": ("mimo_v2", "MiMoV2ProForCausalLM"),

176 "NemotronForCausalLM": ("nemotron", "NemotronForCausalLM"),

177 "NemotronHForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),

178 "NemotronHPuzzleForCausalLM": ("nemotron_h", "NemotronHForCausalLM"),

@@ -468,6 +469,7 @@ _MULTIMODAL_MODELS = {

469 ),

470 "MantisForConditionalGeneration": ("llava", "MantisForConditionalGeneration"),

471 "MiDashengLMModel": ("midashenglm", "MiDashengLMModel"),

472 + "MiMoV2OmniForCausalLM": ("mimo_v2_omni", "MiMoV2OmniForCausalLM"),

473 "MiniMaxVL01ForConditionalGeneration": (

474 "minimax_vl_01",

475 "MiniMaxVL01ForConditionalGeneration",

@@ -570,6 +572,8 @@ _MULTIMODAL_MODELS = {

572 _SPECULATIVE_DECODING_MODELS = {

573 "ExtractHiddenStatesModel": ("extract_hidden_states", "ExtractHiddenStatesModel"),

574 "MiMoMTPModel": ("mimo_mtp", "MiMoMTP"),

575 + "MiMoV2MTPModel": ("mimo_v2_mtp", "MiMoV2MTP"),

576 + "MiMoV2OmniMTPModel": ("mimo_v2_mtp", "MiMoV2OmniMTP"),

577 "EagleLlamaForCausalLM": ("llama_eagle", "EagleLlamaForCausalLM"),

578 "EagleLlama4ForCausalLM": ("llama4_eagle", "EagleLlama4ForCausalLM"),

579 "EagleMiniCPMForCausalLM": ("minicpm_eagle", "EagleMiniCPMForCausalLM"),

vllm/transformers_utils/configs/mimo_v2_omni.py new

+65

@@ -0,0 +1,65 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +

5 +from transformers import PretrainedConfig

6 +

7 +

8 +class Mimo_VLVisionConfig(PretrainedConfig):

9 + model_type = "mimovl"

10 + base_config_key = "vision_config"

11 +

12 + def __init__(

13 + self,

14 + depth=28,

15 + hidden_size=1280,

16 + hidden_act="silu",

17 + intermediate_size=4608,

18 + num_heads=32,

19 + in_channels=3,

20 + patch_size=16,

21 + spatial_merge_size=2,

22 + temporal_patch_size=2,

23 + tokens_per_second=2,

24 + window_size=128,

25 + out_hidden_size=2048,

26 + fullatt_block_indexes=None,

27 + initializer_range=0.02,

28 + kv_channels=64, # HACK

29 + qk_channels=64,

30 + num_query_groups=4,

31 + num_key_value_heads=8,

32 + vit_window_attn_types=None,

33 + visual_token_window_size=64,

34 + **kwargs,

35 + ):

36 + super().__init__(**kwargs)

37 +

38 + self.depth = depth

39 + self.hidden_size = hidden_size

40 + self.hidden_act = hidden_act

41 + self.intermediate_size = intermediate_size

42 + self.num_heads = num_heads

43 + # Support GQA: if num_key_value_heads is not provided,

44 + # default to num_heads (MHA)

45 + if num_key_value_heads is None:

46 + num_key_value_heads = num_heads

47 + self.num_key_value_heads = num_key_value_heads

48 + self.in_channels = in_channels

49 + self.patch_size = patch_size

50 + self.spatial_merge_size = spatial_merge_size

51 + self.temporal_patch_size = temporal_patch_size

52 + self.tokens_per_second = tokens_per_second

53 + self.window_size = window_size

54 + self.fullatt_block_indexes = (

55 + fullatt_block_indexes

56 + if fullatt_block_indexes is not None

57 + else [7, 15, 23, 31]

58 + )

59 + self.out_hidden_size = out_hidden_size

60 + self.initializer_range = initializer_range

61 + self.kv_channels = kv_channels

62 + self.qk_channels = qk_channels

63 + self.num_query_groups = num_query_groups

64 + self.vit_window_attn_types = vit_window_attn_types or [-1] * depth

65 + self.visual_token_window_size = visual_token_window_size

vllm/transformers_utils/model_arch_config_convertor.py

+35

@@ -445,6 +445,37 @@ class MimoMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

445 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)

446

447

448 +def _strip_mimo_v2_attention_chunk_size(

449 + hf_config: PretrainedConfig, hf_text_config: PretrainedConfig

450 +) -> None:

451 + # MiMo-V2-Flash's config.json sets `attention_chunk_size=128` but the

452 + # architecture does not actually use chunked local attention. Leaving it

453 + # set makes vLLM disable the hybrid KV cache manager

454 + for cfg in (hf_text_config, hf_config):

455 + if cfg is not None and hasattr(cfg, "attention_chunk_size"):

456 + delattr(cfg, "attention_chunk_size")

457 +

458 +

459 +class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase):

460 + def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):

461 + super().__init__(hf_config, hf_text_config)

462 + _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)

463 +

464 +

465 +class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

466 + def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig):

467 + super().__init__(hf_config, hf_text_config)

468 + _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config)

469 +

470 + def get_num_hidden_layers(self) -> int:

471 + n = getattr(self.hf_text_config, "num_nextn_predict_layers", None)

472 + if n is not None:

473 + return n

474 + # Fall back to n_predict set by hf_config_override

475 + n = getattr(self.hf_text_config, "n_predict", None)

476 + return n if n is not None else 0

477 +

478 +

479 class GLM4MoeMTPModelArchConfigConvertor(ModelArchConfigConvertorBase):

480 def get_num_hidden_layers(self) -> int:

481 return getattr(self.hf_text_config, "num_nextn_predict_layers", 0)

@@ -511,6 +542,10 @@ MODEL_ARCH_CONFIG_CONVERTORS = {

542 "qwen3_next_mtp": Qwen3NextMTPModelArchConfigConvertor,

543 "qwen3_5_mtp": Qwen3_5MTPModelArchConfigConvertor,

544 "mimo_mtp": MimoMTPModelArchConfigConvertor,

545 + "mimo_v2_pro": MimoV2ModelArchConfigConvertor,

546 + "mimo_v2_flash": MimoV2ModelArchConfigConvertor,

547 + "mimo_v2_mtp": MimoV2MTPModelArchConfigConvertor,

548 + "mimo_v2_omni_mtp": MimoV2MTPModelArchConfigConvertor,

549 "glm4_moe_mtp": GLM4MoeMTPModelArchConfigConvertor,

550 "glm_ocr_mtp": GLM4MoeMTPModelArchConfigConvertor,

551 "ernie_mtp": ErnieMTPModelArchConfigConvertor,

vllm/transformers_utils/processors/__init__.py

@@ -27,6 +27,7 @@ __all__ = [

27 "IsaacProcessor",

28 "KimiAudioProcessor",

29 "KimiK25Processor",

30 + "MiMoOmniProcessor",

31 "MistralCommonPixtralProcessor",

32 "MistralCommonVoxtralProcessor",

33 "NanoNemotronVLProcessor",

@@ -57,6 +58,7 @@ _CLASS_TO_MODULE: dict[str, str] = {

58 "IsaacProcessor": "vllm.transformers_utils.processors.isaac",

59 "KimiAudioProcessor": "vllm.transformers_utils.processors.kimi_audio",

60 "KimiK25Processor": "vllm.transformers_utils.processors.kimi_k25",

61 + "MiMoOmniProcessor": "vllm.transformers_utils.processors.mimo_v2_omni",

62 "MistralCommonPixtralProcessor": "vllm.transformers_utils.processors.pixtral",

63 "MistralCommonVoxtralProcessor": "vllm.transformers_utils.processors.voxtral",

64 "NanoNemotronVLProcessor": "vllm.transformers_utils.processors.nano_nemotron_vl",

vllm/transformers_utils/processors/mimo_v2_omni.py new

+1285

@@ -0,0 +1,1285 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +# mypy: ignore-errors

4 +"""MiMo-Omni multimodal processor for vLLM.

5 +

6 +Ported from SGLang's MiMoV2OmniProcessor / MiMoVLProcessor implementations.

7 +"""

8 +

9 +import contextlib

10 +import copy

11 +import io

12 +import logging

13 +import math

14 +from collections import OrderedDict

15 +from concurrent.futures import ThreadPoolExecutor, as_completed

16 +from dataclasses import dataclass, field

17 +from io import BytesIO

18 +from typing import Any, Literal

19 +

20 +import numpy as np

21 +import regex as re

22 +import requests

23 +import torch

24 +import torch.nn.functional as F

25 +from PIL import Image

26 +from transformers import BatchFeature, TensorType

27 +from transformers.processing_utils import ProcessorMixin

28 +

29 +try:

30 + from torchcodec.decoders import AudioDecoder

31 +

32 + _HAS_TORCHCODEC = True

33 +except ImportError:

34 + AudioDecoder = None

35 + _HAS_TORCHCODEC = False

36 +

37 +try:

38 + import torchaudio

39 + from torchaudio.transforms import MelSpectrogram as _MelSpectrogram

40 +

41 + _HAS_TORCHAUDIO = True

42 +except ImportError:

43 + torchaudio = None # type: ignore[assignment]

44 + _MelSpectrogram = None # type: ignore[assignment,misc]

45 + _HAS_TORCHAUDIO = False

46 +

47 +logger = logging.getLogger(__name__)

48 +

49 +# ---------------------------------------------------------------------------

50 +# Constants

51 +# ---------------------------------------------------------------------------

52 +

53 +_PIXEL_MEAN = [123.675, 116.28, 103.53]

54 +_PIXEL_STD = [58.395, 57.12, 57.375]

55 +_mean_std_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = {}

56 +

57 +

58 +# ---------------------------------------------------------------------------

59 +# Data classes

60 +# ---------------------------------------------------------------------------

61 +

62 +

63 +@dataclass

64 +class ImageInput:

65 + # PIL.Image | str (path/url/base64) | bytes | torch.Tensor (C,H,W)

66 + image: Any

67 + max_pixels: int | None = None

68 + min_pixels: int | None = None

69 +

70 +

71 +@dataclass

72 +class VideoInput:

73 + # tuple[frames_TCHW: torch.Tensor, timestamps_T: torch.Tensor]

74 + video: Any

75 + min_pixels: int | None = None

76 + max_pixels: int | None = None

77 + total_max_pixels: int | None = None

78 + fps: float | None = None

79 + num_frames: int | None = None

80 + max_frames: int | None = None

81 + min_frames: int | None = None

82 + do_include_last_frame: bool | None = False

83 + start_time: float | None = None

84 + end_time: float | None = None

85 + segment_type: Literal["individual", "partial"] = "individual"

86 +

87 +

88 +@dataclass

89 +class AudioInput:

90 + # str (path/url/base64) | bytes | tuple[waveform_1D, sr]

91 + # | np.ndarray | torch.Tensor (T,n_vq)

92 + audio: Any

93 +

94 +

95 +@dataclass

96 +class VideoAudioInput:

97 + video: Any # same as VideoInput.video

98 + audio: Any # same as AudioInput.audio

99 + min_pixels: int | None = None

100 + max_pixels: int | None = None

101 + total_max_pixels: int | None = None

102 + fps: float | None = None

103 + num_frames: int | None = None

104 + max_frames: int | None = None

105 + min_frames: int | None = None

106 + do_include_last_frame: bool | None = False

107 + start_time: float | None = None

108 + end_time: float | None = None

109 + segment_type: Literal["individual", "partial"] = "individual"

110 +

111 +

112 +@dataclass

113 +class Content:

114 + type: Literal["text", "image", "video", "audio", "video_audio"]

115 + content: Any

116 + is_target: bool | None = None

117 +

118 +

119 +@dataclass

120 +class MiMoVLInputSample:

121 + input_ids: torch.Tensor

122 + labels: torch.Tensor | None

123 + pixel_values: list[torch.Tensor]

124 + pixel_values_videos: list[torch.Tensor]

125 + image_thw_grids: list[torch.Tensor]

126 + video_thw_grids: list[torch.Tensor]

127 + audio_inputs: list[torch.Tensor]

128 + second_per_grid_ts: list[float] = field(default_factory=list)

129 + video_start_times: list[float] = field(default_factory=list)

130 + audio_token_lens: list[int] = field(default_factory=list)

131 + va_audio_inputs: list[torch.Tensor] = field(default_factory=list)

132 + video_audio_n_segs: list[int] = field(default_factory=list)

133 + video_audio_seg_lens: list[int] = field(default_factory=list)

134 + position_ids: torch.Tensor | None = None

135 + rope_deltas: torch.Tensor | None = None

136 + extra: dict = field(default_factory=dict)

137 +

138 +

139 +# ---------------------------------------------------------------------------

140 +# Vision utilities

141 +# ---------------------------------------------------------------------------

142 +

143 +

144 +def _format_timestamp(ts: float) -> str:

145 + return f"{int(ts // 60):02d}:{int(ts % 60):02d}"

146 +

147 +

148 +def _smart_resize(

149 + h: int, w: int, factor: int, min_px: int, max_px: int

150 +) -> tuple[int, int]:

151 + if min(h, w) < factor:

152 + if h < w:

153 + h, w = factor, int(w * factor / h)

154 + else:

155 + w, h = factor, int(h * factor / w)

156 + elif max(h, w) / min(h, w) > 200:

157 + raise ValueError(f"Aspect ratio > 200 not allowed: {h}x{w}")

158 + h_bar = round(h / factor) * factor

159 + w_bar = round(w / factor) * factor

160 + if h_bar * w_bar > max_px:

161 + beta = math.sqrt((h * w) / max_px)

162 + h_bar = math.floor(h / beta / factor) * factor

163 + w_bar = math.floor(w / beta / factor) * factor

164 + elif h_bar * w_bar < min_px:

165 + beta = math.sqrt(min_px / (h * w))

166 + h_bar = math.ceil(h * beta / factor) * factor

167 + w_bar = math.ceil(w * beta / factor) * factor

168 + return int(h_bar), int(w_bar)

169 +

170 +

171 +def _to_rgb(img: Image.Image) -> Image.Image:

172 + if img.mode == "RGBA":

173 + bg = Image.new("RGB", img.size, (255, 255, 255))

174 + bg.paste(img, mask=img.split()[3])

175 + return bg

176 + return img.convert("RGB")

177 +

178 +

179 +def _standardize(images: torch.Tensor) -> torch.Tensor:

180 + key = str(images.device)

181 + if key not in _mean_std_cache:

182 + mean = torch.tensor(_PIXEL_MEAN, device=images.device).view(1, -1, 1, 1)

183 + std = torch.tensor(_PIXEL_STD, device=images.device).view(1, -1, 1, 1)

184 + _mean_std_cache[key] = (mean, std)

185 + mean, std = _mean_std_cache[key]

186 + return (images - mean) / std

187 +

188 +

189 +def _transform_batch(

190 + frames: torch.Tensor,

191 + factor: int,

192 + min_px: int,

193 + max_px: int,

194 + device: torch.device | None = None,

195 +) -> tuple[torch.Tensor, int, int]:

196 + if device is not None:

197 + frames = frames.to(device)

198 + _, _, h, w = frames.shape

199 + h_bar, w_bar = _smart_resize(h, w, factor, min_px, max_px)

200 + resized = F.interpolate(

201 + frames.float(), (h_bar, w_bar), mode="bilinear", align_corners=False

202 + )

203 + return _standardize(resized), w_bar, h_bar

204 +

205 +

206 +def _transform_single(

207 + img: Any,

208 + factor: int,

209 + min_px: int,

210 + max_px: int,

211 + device: torch.device | None = None,

212 +) -> tuple[torch.Tensor, int, int]:

213 + if isinstance(img, torch.Tensor):

214 + t = img.float()

215 + _, h, w = t.shape

216 + elif isinstance(img, Image.Image):

217 + img = img.convert("RGB")

218 + w, h = img.size

219 + t = torch.from_numpy(np.array(img)).permute(2, 0, 1).float()

220 + else:

221 + raise TypeError(f"Expected Tensor or PIL.Image, got {type(img)}")

222 + if device is not None:

223 + t = t.to(device)

224 + h_bar, w_bar = _smart_resize(h, w, factor, min_px, max_px)

225 + out = F.interpolate(

226 + t.unsqueeze(0), (h_bar, w_bar), mode="bilinear", align_corners=False

227 + )

228 + return _standardize(out).squeeze(0), w_bar, h_bar

229 +

230 +

231 +def _fetch_image(src: Any) -> Image.Image:

232 + if isinstance(src, Image.Image):

233 + return _to_rgb(src)

234 + if isinstance(src, bytes):

235 + return _to_rgb(copy.deepcopy(Image.open(BytesIO(src))))

236 + if isinstance(src, str):

237 + if src.startswith(("http://", "https://")):

238 + r = requests.get(src, timeout=30)

239 + r.raise_for_status()

240 + return _to_rgb(copy.deepcopy(Image.open(BytesIO(r.content))))

241 + if src.startswith("file://"):

242 + return _to_rgb(Image.open(src[7:]))

243 + if src.startswith("data:image"):

244 + import pybase64 as _b64

245 +

246 + _, b64 = src.split("base64,", 1)

247 + return _to_rgb(copy.deepcopy(Image.open(BytesIO(_b64.b64decode(b64)))))

248 + return _to_rgb(Image.open(src))

249 + raise ValueError(f"Unrecognized image source: {type(src)}")

250 +

251 +

252 +# ---------------------------------------------------------------------------

253 +# Core processor

254 +# ---------------------------------------------------------------------------

255 +

256 +

257 +class MiMoVLProcessor:

258 + """Core MiMo-VL multimodal processor.

259 +

260 + Handles image/video/audio preprocessing and token sequence construction.

261 + Ported from SGLang's MiMoVLProcessor.

262 + """

263 +

264 + def __init__(

265 + self,

266 + tokenizer: Any,

267 + patch_size: int = 14,

268 + merge_size: int = 2,

269 + temporal_patch_size: int = 2,

270 + temporal_compression_ratio: int = 1,

271 + use_video_timestamps: bool = True,

272 + video_audio_interleave_length: int = 0,

273 + audio_kernel_size: int = 3,

274 + audio_stride_size: int = 2,

275 + audio_avg_pooler: int = 2,

276 + audio_sampling_rate: int = 24000,

277 + audio_nfft: int = 960,

278 + audio_hop_length: int = 240,

279 + audio_window_size: int = 960,

280 + audio_fmin: float = 0.0,

281 + audio_fmax: float | None = None,

282 + audio_n_mels: int = 128,

283 + audio_segment_size: int = 6000,

284 + audio_channels: int = 8,

285 + audio_group_size: int = 4,

286 + audio_input_id_per_second: float = 25.0,

287 + audio_zeroemb_idx: int = 4096,

288 + image_min_pixels: int | None = None,

289 + image_max_pixels: int | None = None,

290 + video_min_pixels: int | None = None,

291 + video_max_pixels: int | None = None,

292 + video_total_max_pixels: int | None = None,

293 + fps: float | None = None,

294 + num_frames: int | None = None,

295 + max_frames: int | None = None,

296 + min_frames: int | None = None,

297 + image_token_id: int | None = None,

298 + video_token_id: int | None = None,

299 + audio_token_id: int | None = None,

300 + vision_start_token_id: int | None = None,

301 + vision_end_token_id: int | None = None,

302 + audio_start_token_id: int | None = None,

303 + audio_end_token_id: int | None = None,

304 + video_start_token_id: int | None = None,

305 + video_end_token_id: int | None = None,

306 + pad_token_id: int | None = None,

307 + rope_type: str = "rope",

308 + video_process_num_threads: int = 16,

309 + device: Any | None = None,

310 + **kwargs: Any,

311 + ) -> None:

312 + self.tokenizer = tokenizer

313 + self.video_process_num_threads = video_process_num_threads

314 + self.device = torch.device(device) if isinstance(device, str) else device

315 +

316 + self.rope_type = "rope" if rope_type == "1d" else rope_type

317 + assert self.rope_type in ("rope", "mrope"), (

318 + f"Unknown rope_type: {self.rope_type}"

319 + )

320 +

321 + # video timestamps require 1-D rope

322 + assert use_video_timestamps, "use_video_timestamps must be True"

323 + assert self.rope_type == "rope", (

324 + "use_video_timestamps requires rope_type='rope'"

325 + )

326 + self.use_video_timestamps = use_video_timestamps

327 + self.video_audio_interleave_length = video_audio_interleave_length

328 +

329 + self.image_token_id = image_token_id

330 + self.video_token_id = video_token_id

331 + self.audio_token_id = audio_token_id

332 + self.vision_start_token_id = vision_start_token_id

333 + self.vision_end_token_id = vision_end_token_id

334 + self.audio_start_token_id = audio_start_token_id

335 + self.audio_end_token_id = audio_end_token_id

336 + self.video_start_token_id = video_start_token_id

337 + self.video_end_token_id = video_end_token_id

338 + self.pad_token_id = pad_token_id

339 +

340 + self.patch_size = patch_size

341 + self.merge_size = merge_size

342 + self.temporal_patch_size = temporal_patch_size

343 + self.temporal_compression_ratio = temporal_compression_ratio

344 +

345 + self.audio_sampling_rate = audio_sampling_rate

346 + self.audio_nfft = audio_nfft

347 + self.audio_hop_length = audio_hop_length

348 + self.audio_window_size = audio_window_size

349 + self.audio_fmin = audio_fmin

350 + self.audio_fmax = audio_fmax

351 + self.audio_n_mels = audio_n_mels

352 + self.audio_segment_size = audio_segment_size

353 + self.audio_kernel_size = audio_kernel_size

354 + self.audio_stride_size = audio_stride_size

355 + self.audio_avg_pooler = audio_avg_pooler

356 + self.audio_channels = audio_channels

357 + self.audio_group_size = audio_group_size

358 + self.audio_input_id_per_second = audio_input_id_per_second

359 +

360 + self._mel_spec_kwargs = dict(

361 + sample_rate=audio_sampling_rate,

362 + n_fft=audio_nfft,

363 + hop_length=audio_hop_length,

364 + win_length=audio_window_size,

365 + f_min=audio_fmin,

366 + f_max=audio_fmax,

367 + n_mels=audio_n_mels,

368 + power=1.0,

369 + center=True,

370 + )

371 + self._mel_spectrogram: Any | None = None

372 + self._resamplers: OrderedDict = OrderedDict()

373 + self._resamplers_max = 16

374 +

375 + if isinstance(audio_zeroemb_idx, int):

376 + self.audio_zeroemb_idxs = torch.tensor(

377 + [audio_zeroemb_idx] * audio_channels, dtype=torch.int32

378 + )

379 + else:

380 + self.audio_zeroemb_idxs = torch.tensor(audio_zeroemb_idx, dtype=torch.int32)

381 +

382 + assert image_min_pixels is not None, "image_min_pixels must be set"

383 + assert image_max_pixels is not None, "image_max_pixels must be set"

384 + assert video_min_pixels is not None, "video_min_pixels must be set"

385 + assert video_max_pixels is not None, "video_max_pixels must be set"

386 + assert video_total_max_pixels is not None, "video_total_max_pixels must be set"

387 + assert fps is not None or num_frames is not None, (

388 + "fps or num_frames must be set"

389 + )

390 +

391 + self._img_kw = {"min_pixels": image_min_pixels, "max_pixels": image_max_pixels}

392 + self._vid_kw = {

393 + "min_pixels": video_min_pixels,

394 + "max_pixels": video_max_pixels,

395 + "total_max_pixels": video_total_max_pixels,

396 + "fps": fps,

397 + "num_frames": num_frames,

398 + "max_frames": max_frames,

399 + "min_frames": min_frames,

400 + }

401 +

402 + @property

403 + def mel_spectrogram(self) -> Any:

404 + if self._mel_spectrogram is None:

405 + if _MelSpectrogram is None:

406 + raise RuntimeError(

407 + "torchaudio is required for audio. "

408 + "Install with: pip install torchaudio"

409 + )

410 + self._mel_spectrogram = _MelSpectrogram(**self._mel_spec_kwargs)

411 + return self._mel_spectrogram

412 +

413 + def _resolve_img_kw(self, img: ImageInput) -> dict:

414 + return {

415 + "min_px": (

416 + img.min_pixels

417 + if img.min_pixels is not None

418 + else self._img_kw["min_pixels"]

419 + ),

420 + "max_px": (

421 + img.max_pixels

422 + if img.max_pixels is not None

423 + else self._img_kw["max_pixels"]

424 + ),

425 + }

426 +

427 + def _resolve_vid_kw(self, vid: VideoInput) -> dict:

428 + kw: dict = {}

429 + for k in ("min_pixels", "max_pixels", "total_max_pixels"):

430 + kw[k] = getattr(vid, k) or self._vid_kw[k]

431 + if vid.num_frames is not None:

432 + kw["num_frames"] = vid.num_frames

433 + elif vid.fps is not None:

434 + kw["fps"] = vid.fps

435 + if vid.max_frames is not None:

436 + kw["max_frames"] = vid.max_frames

437 + if vid.min_frames is not None:

438 + kw["min_frames"] = vid.min_frames

439 + elif self._vid_kw["num_frames"] is not None:

440 + kw["num_frames"] = self._vid_kw["num_frames"]

441 + elif self._vid_kw["fps"] is not None:

442 + kw["fps"] = self._vid_kw["fps"]

443 + if self._vid_kw["max_frames"] is not None:

444 + kw["max_frames"] = self._vid_kw["max_frames"]

445 + if self._vid_kw["min_frames"] is not None:

446 + kw["min_frames"] = self._vid_kw["min_frames"]

447 + else:

448 + raise ValueError(

449 + "No video sampling strategy specified (fps or num_frames)."

450 + )

451 + return kw

452 +

453 + def preprocess_audio(self, audio: Any) -> tuple[torch.Tensor, int]:

454 + """Decode audio bytes/path/tuple → (mel_spec (T, n_mels), token_len)."""

455 + if isinstance(audio, tuple):

456 + waveform, original_sr = audio

457 + else:

458 + if AudioDecoder is None:

459 + raise RuntimeError(

460 + "torchcodec is required for audio. "

461 + "Install with: pip install torchcodec"

462 + )

463 + if isinstance(audio, bytes):

464 + file_obj: Any = io.BytesIO(audio)

465 + elif isinstance(audio, str):

466 + if audio.startswith("data:"):

467 + import pybase64 as _b64

468 +

469 + file_obj = io.BytesIO(_b64.b64decode(audio.split(",")[1]))

470 + elif audio.startswith(("http://", "https://")):

471 + r = requests.get(audio, timeout=30)

472 + r.raise_for_status()

473 + file_obj = io.BytesIO(r.content)

474 + else:

475 + file_obj = audio

476 + else:

477 + raise ValueError(f"Unsupported audio source type: {type(audio)}")

478 + samples = AudioDecoder(file_obj).get_all_samples()

479 + waveform = samples.data

480 + original_sr = samples.sample_rate

481 +

482 + if original_sr != self.audio_sampling_rate:

483 + if original_sr not in self._resamplers:

484 + if len(self._resamplers) >= self._resamplers_max:

485 + self._resamplers.popitem(last=False)

486 + self._resamplers[original_sr] = torchaudio.transforms.Resample(

487 + orig_freq=original_sr, new_freq=self.audio_sampling_rate

488 + )

489 + self._resamplers.move_to_end(original_sr)

490 + waveform = self._resamplers[original_sr](waveform)

491 +

492 + if waveform.ndim == 2:

493 + waveform = waveform.mean(dim=0)

494 + spec = self.mel_spectrogram(waveform[None, :])

495 + spec = torch.log(torch.clip(spec, min=1e-7)).squeeze().transpose(0, 1)

496 +

497 + n = spec.shape[0]

498 + n = n + 3 - self.audio_kernel_size

499 + n = (n + 2 - self.audio_kernel_size) // self.audio_stride_size + 1

500 + n = n // self.audio_avg_pooler + int(n % self.audio_avg_pooler != 0)

501 + token_len = math.ceil(n / self.audio_group_size)

502 + return spec, token_len

503 +

504 + def process_image(self, image: ImageInput) -> torch.Tensor:

505 + kw = self._resolve_img_kw(image)

506 + src = image.image

507 + if isinstance(src, (str, bytes)):

508 + src = _fetch_image(src)

509 + tensor, _, _ = _transform_single(

510 + src,

511 + factor=self.patch_size * self.merge_size,

512 + device=self.device,

513 + **kw,

514 + )

515 + return tensor

516 +

517 + def process_video(

518 + self, video_input: VideoInput

519 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, dict]:

520 + kw = self._resolve_vid_kw(video_input)

521 + video = video_input.video

522 + if not isinstance(video, tuple):

523 + raise ValueError(

524 + f"video must be a (frames_TCHW, timestamps_T) tuple, "

525 + f"got {type(video)}. "

526 + "Decode the video before calling the processor."

527 + )

528 + frames, timestamps = video

529 +

530 + fps = (

531 + 1.0

532 + if len(timestamps) < 2

533 + else float(1.0 / (float(timestamps[1]) - float(timestamps[0])))

534 + )

535 + start = (

536 + video_input.start_time

537 + if video_input.start_time is not None

538 + else float(timestamps[0])

539 + )

540 + end = (

541 + video_input.end_time

542 + if video_input.end_time is not None

543 + else float(timestamps[-1]) + 1.0 / fps

544 + )

545 +

546 + if video_input.segment_type != "individual":

547 + mask = (timestamps >= start) & (timestamps < end)

548 + idxs = torch.where(mask)[0]

549 + if len(idxs) == 0:

550 + idxs = torch.where(timestamps <= start)[0][-1:]

551 + frames, timestamps = frames[idxs], timestamps[idxs]

552 +

553 + tp = self.temporal_patch_size * self.temporal_compression_ratio

554 + n = frames.shape[0]

555 + total_px = kw["total_max_pixels"]

556 + max_px = max(

557 + kw["min_pixels"], min(total_px * tp // max(n, 1), kw["max_pixels"])

558 + )

559 +

560 + if n % tp != 0:

561 + pad = tp - n % tp

562 + frames = torch.cat(

563 + [frames, frames[-1:].repeat(pad, *([1] * (frames.ndim - 1)))],

564 + dim=0,

565 + )

566 + timestamps = torch.cat([timestamps, timestamps[-1:].repeat(pad)], dim=0)

567 +

568 + transformed, _, _ = _transform_batch(

569 + frames,

570 + factor=self.patch_size * self.merge_size,

571 + min_px=kw["min_pixels"],

572 + max_px=max_px,

573 + device=self.device,

574 + )

575 + patches, thw = self._flatten_visual(transformed, "video")

576 + meta = {

577 + "fps_sampled": fps,

578 + "segment_start_time": start,

579 + "segment_end_time": end,

580 + }

581 + return patches, thw, timestamps, meta

582 +

583 + def process_audio(self, audio: AudioInput) -> Any:

584 + src = audio.audio

585 + if isinstance(src, np.ndarray):

586 + src = (torch.from_numpy(src).float(), self.audio_sampling_rate)

587 + if isinstance(src, (str, bytes, tuple)):

588 + return self.preprocess_audio(src)

589 + # Pre-tokenized tensor (T, n_vq)

590 + assert isinstance(src, torch.Tensor) and src.ndim == 2

591 + T = src.shape[0]

592 + src = src[:, : self.audio_channels].to(torch.long)

593 + pad_T = (

594 + (T + self.audio_group_size - 1)

595 + // self.audio_group_size

596 + * self.audio_group_size

597 + )

598 + padding = (

599 + torch.zeros(pad_T - T, self.audio_channels, dtype=torch.long) + src[-1]

600 + )

601 + src = torch.cat([src, padding], dim=0)

602 + return src.reshape(

603 + pad_T // self.audio_group_size, self.audio_group_size, self.audio_channels

604 + )

605 +

606 + def _flatten_visual(

607 + self, visual: torch.Tensor, kind: str

608 + ) -> tuple[torch.Tensor, torch.Tensor]:

609 + if kind == "image":

610 + h, w = visual.shape[-2:]

611 + patches = visual.unsqueeze(0).repeat(self.temporal_patch_size, 1, 1, 1)

612 + else: # video / video_audio

613 + temporal_stride = self.temporal_compression_ratio * self.temporal_patch_size

614 + assert visual.shape[0] % temporal_stride == 0

615 + patches = visual

616 + h, w = patches.shape[-2:]

617 +

618 + C = patches.shape[1]

619 + grid_t = patches.shape[0] // self.temporal_patch_size

620 + grid_h, grid_w = h // self.patch_size, w // self.patch_size

621 +

622 + patches = (

623 + patches.contiguous()

624 + .view(

625 + grid_t,

626 + self.temporal_patch_size,

627 + C,

628 + grid_h // self.merge_size,

629 + self.merge_size,

630 + self.patch_size,

631 + grid_w // self.merge_size,

632 + self.merge_size,

633 + self.patch_size,

634 + )

635 + .permute(0, 3, 6, 4, 7, 2, 1, 5, 8)

636 + .contiguous()

637 + .view(

638 + grid_t * grid_h * grid_w,

639 + C * self.temporal_patch_size * self.patch_size * self.patch_size,

640 + )

641 + )

642 + thw = torch.tensor([grid_t, grid_h, grid_w], dtype=torch.int32)

643 + return patches, thw

644 +

645 + def process(

646 + self, contents: list[Content], verbose: bool = False

647 + ) -> MiMoVLInputSample:

648 + input_ids: list[int] = []

649 + labels: list[int] = []

650 + img_pv: list[torch.Tensor] = []

651 + img_grids: list[torch.Tensor] = []

652 + vid_pv: list[torch.Tensor] = []

653 + vid_grids: list[torch.Tensor] = []

654 + audio_inputs: list[torch.Tensor] = []

655 + is_audio_tokenized: list[bool] = []

656 + audio_token_lens: list[int] = []

657 + second_per_grid_ts: list[float] = []

658 + video_start_times: list[float] = []

659 + va_audio_inputs: list[torch.Tensor] = []

660 + video_audio_n_segs: list[int] = []

661 + video_audio_seg_lens: list[int] = []

662 +

663 + # Pre-decode videos in parallel

664 + vid_info = [

665 + (i, c.content, c.type == "video_audio")

666 + for i, c in enumerate(contents)

667 + if c.type in ("video", "video_audio")

668 + ]

669 + vid_results: dict[int, tuple] = {}

670 + if vid_info:

671 + n_t = min(self.video_process_num_threads, len(vid_info))

672 + if n_t > 1 and len(vid_info) > 1:

673 + with ThreadPoolExecutor(max_workers=n_t) as ex:

674 + fut_map = {

675 + ex.submit(self.process_video, vi): idx

676 + for idx, vi, _ in vid_info

677 + }

678 + for fut in as_completed(fut_map):

679 + vid_results[fut_map[fut]] = fut.result()

680 + else:

681 + for idx, vi, _ in vid_info:

682 + vid_results[idx] = self.process_video(vi)

683 +

684 + for ci, content in enumerate(contents):

685 + _ids: list[int] = []

686 + _lbls: list[int] | None = None

687 +

688 + if content.type == "text":

689 + _ids = (

690 + self.tokenizer.encode(content.content)

691 + if isinstance(content.content, str)

692 + else list(content.content)

693 + )

694 + if content.is_target:

695 + _lbls = _ids

696 +

697 + elif content.type == "image":

698 + tensor = self.process_image(content.content)

699 + patches, thw = self._flatten_visual(tensor, "image")

700 + t, h, w = thw.tolist()

701 + n_tok = (t * h * w) // (self.merge_size**2)

702 + img_pv.append(patches)

703 + img_grids.append(thw)

704 + _ids = (

705 + [self.vision_start_token_id]

706 + + [self.image_token_id] * n_tok

707 + + [self.vision_end_token_id]

708 + )

709 +

710 + elif content.type == "video":

711 + patches, thw, ts, meta = vid_results[ci]

712 + t, h, w = thw.tolist()

713 + n_per_grid = h * w // (self.merge_size**2)

714 + vid_pv.append(patches)

715 + vid_grids.append(thw)

716 + second_per_grid_ts.append(

717 + self.temporal_patch_size / meta["fps_sampled"]

718 + )

719 + video_start_times.append(float(ts[0]))

720 + video_audio_n_segs.append(0)

721 +

722 + stride = self.temporal_patch_size * self.temporal_compression_ratio

723 + ts_texts = [_format_timestamp(float(x)) for x in ts[::stride]]

724 + ts_ids_list = [self.tokenizer.encode(s) for s in ts_texts]

725 +

726 + _ids = [self.video_start_token_id]

727 + for ts_ids in ts_ids_list:

728 + _ids += (

729 + ts_ids

730 + + [self.vision_start_token_id]

731 + + [self.video_token_id] * n_per_grid

732 + + [self.vision_end_token_id]

733 + )

734 + _ids += [self.video_end_token_id]

735 +

736 + elif content.type == "audio":

737 + processed = self.process_audio(content.content)

738 + if isinstance(processed, tuple):

739 + is_audio_tokenized.append(False)

740 + spec, tok_len = processed

741 + audio_inputs.append(spec)

742 + else:

743 + is_audio_tokenized.append(True)

744 + tok_len = processed.shape[0]

745 + audio_inputs.append(processed)

746 + audio_token_lens.append(tok_len)

747 + _ids = (

748 + [self.audio_start_token_id]

749 + + [self.audio_token_id] * tok_len

750 + + [self.audio_end_token_id]

751 + )

752 +

753 + elif content.type == "video_audio":

754 + patches, thw, ts, meta = vid_results[ci]

755 + second_per_grid_ts.append(

756 + self.temporal_patch_size / meta["fps_sampled"]

757 + )

758 + video_start_times.append(float(ts[0]))

759 + processed_audio = self.process_audio(content.content)

760 + tok_per_sec = self.audio_input_id_per_second / self.audio_group_size

761 +

762 + t, h, w = thw.tolist()

763 + vid_pv.append(patches)

764 + vid_grids.append(thw)

765 +

766 + if isinstance(processed_audio, tuple):

767 + # Mel spec (not pre-tokenized): store in va_audio_inputs separately

768 + spec, total_atok = processed_audio

769 + va_audio_inputs.append(spec)

770 + _va_is_tokenized = False

771 + else:

772 + # Pre-tokenized: not expected in vLLM, but handle defensively

773 + total_atok = processed_audio.shape[0]

774 + _va_is_tokenized = True

775 +

776 + n_per_grid = h * w // (self.merge_size**2)

777 + stride = self.temporal_patch_size * self.temporal_compression_ratio

778 + grid_ts = ts[::stride]

779 + ts_texts = [_format_timestamp(float(x)) for x in grid_ts]

780 + ts_ids_list = [self.tokenizer.encode(s) for s in ts_texts]

781 +

782 + units: list[tuple] = []

783 + for i in range(len(grid_ts)):

784 + a_start = int(float(grid_ts[i]) * tok_per_sec)

785 + a_end = (

786 + int(float(grid_ts[i + 1]) * tok_per_sec)

787 + if i < len(grid_ts) - 1

788 + else int(meta["segment_end_time"] * tok_per_sec)

789 + )

790 + seg_len = min(a_end, total_atok) - a_start

791 + assert seg_len > 0, f"Zero-length audio segment at grid index {i}"

792 + seg = (

793 + processed_audio[a_start : a_start + seg_len]

794 + if _va_is_tokenized

795 + else None

796 + )

797 + units.append(

798 + (

799 + float(grid_ts[i]),

800 + ts_texts[i],

801 + ts_ids_list[i],

802 + n_per_grid,

803 + seg_len,

804 + seg,

805 + )

806 + )

807 +

808 + il = self.video_audio_interleave_length

809 + if il == -1:

810 + groups: list[list] = [list(enumerate(units))]

811 + elif il == 0:

812 + groups = [[(i, u)] for i, u in enumerate(units)]

813 + else:

814 + groups, cur, t_ptr = [], [], 0.0

815 + for i, u in enumerate(units):

816 + while u[0] >= t_ptr + il:

817 + if cur:

818 + groups.append(cur)

819 + cur = []

820 + t_ptr += il

821 + cur.append((i, u))

822 + if cur:

823 + groups.append(cur)

824 +

825 + # Track n_segs (= num groups) and per-group audio token counts

826 + video_audio_n_segs.append(len(groups))

827 + for group in groups:

828 + group_seg_len = sum(u[4] for _, u in group)

829 + video_audio_seg_lens.append(group_seg_len)

830 +

831 + _ids = [self.video_start_token_id]

832 + for group in groups:

833 + _ids += group[0][1][2] # first-unit timestamp token ids

834 + _vid_tok: list[int] = []

835 + _aud_tok: list[int] = []

836 + for _, u in group:

837 + _, _, _, vid_n, seg_n, seg_audio = u

838 + _vid_tok += (

839 + [self.vision_start_token_id]

840 + + [self.video_token_id] * vid_n

841 + + [self.vision_end_token_id]

842 + )

843 + _aud_tok += [self.audio_token_id] * seg_n

844 + if seg_audio is not None:

845 + # Pre-tokenized per-frame segments (rare in vLLM)

846 + audio_inputs.append(seg_audio)

847 + _ids += (

848 + _vid_tok

849 + + [self.audio_start_token_id]

850 + + _aud_tok

851 + + [self.audio_end_token_id]

852 + )

853 + _ids += [self.video_end_token_id]

854 +

855 + input_ids.extend(_ids)

856 + labels.extend(

857 + _lbls if _lbls is not None else [self.pad_token_id] * len(_ids)

858 + )

859 +

860 + ids_t = torch.tensor(input_ids)

861 + lbl_arr = np.roll(labels, shift=-1)

862 + lbl_arr[-1] = self.pad_token_id

863 + lbl_t = torch.tensor(lbl_arr)

864 +

865 + extra: dict = {}

866 + if is_audio_tokenized:

867 + assert all(is_audio_tokenized) or not any(is_audio_tokenized)

868 + extra["is_audio_tokenized"] = is_audio_tokenized[0]

869 +

870 + position_ids = torch.arange(ids_t.shape[0]).expand(3, -1)

871 + rope_deltas = torch.zeros((1, 1), dtype=torch.int32)

872 +

873 + return MiMoVLInputSample(

874 + input_ids=ids_t,

875 + labels=lbl_t,

876 + pixel_values=img_pv,

877 + pixel_values_videos=vid_pv,

878 + image_thw_grids=img_grids,

879 + video_thw_grids=vid_grids,

880 + audio_inputs=audio_inputs,

881 + second_per_grid_ts=second_per_grid_ts,

882 + video_start_times=video_start_times,

883 + audio_token_lens=audio_token_lens,

884 + va_audio_inputs=va_audio_inputs,

885 + video_audio_n_segs=video_audio_n_segs,

886 + video_audio_seg_lens=video_audio_seg_lens,

887 + position_ids=position_ids,

888 + rope_deltas=rope_deltas,

889 + extra=extra,

890 + )

891 +

892 +

893 +# ---------------------------------------------------------------------------

894 +# vLLM ProcessorMixin wrapper

895 +# ---------------------------------------------------------------------------

896 +

897 +

898 +class MiMoOmniProcessor(ProcessorMixin):

899 + """HuggingFace-compatible ProcessorMixin wrapper for MiMo-Omni.

900 +

901 + Accepts PIL images, pre-decoded video tuples (frames_TCHW, timestamps_T),

902 + and audio (file path / bytes / (waveform, sr) tuple / numpy array).

903 + """

904 +

905 + attributes = ["tokenizer"]

906 + tokenizer_class = "AutoTokenizer"

907 +

908 + # Single or multi-pad placeholders produced by the chat template / prior expansion

909 + _IMG_RE = re.compile(r"<\|vision_start\|>(?:<\|image_pad\|>)+<\|vision_end\|>")

910 + _VID_RE = re.compile(r"<\|vision_start\|>(?:<\|video_pad\|>)+<\|vision_end\|>")

911 + _AUD_RE = re.compile(

912 + r"<\|mimo_audio_start\|>(?:<\|audio_pad\|>)+<\|mimo_audio_end\|>"

913 + )

914 +

915 + _MM_RE = re.compile(

916 + r"(<\|vision_start\|>(?:<\|image_pad\|>)+<\|vision_end\|>"

917 + r"|<\|vision_start\|>(?:<\|video_pad\|>)+<\|vision_end\|>"

918 + r"|<\|mimo_audio_start\|>(?:<\|audio_pad\|>)+<\|mimo_audio_end\|>)"

919 + )

920 +

921 + def __init__(

922 + self,

923 + tokenizer: Any,

924 + *,

925 + patch_size: int = 14,

926 + merge_size: int = 2,

927 + temporal_patch_size: int = 2,

928 + temporal_compression_ratio: int = 1,

929 + image_min_pixels: int | None = None,

930 + image_max_pixels: int | None = None,

931 + video_min_pixels: int | None = None,

932 + video_max_pixels: int | None = None,

933 + video_total_max_pixels: int | None = None,

934 + fps: float = 2.0,

935 + num_frames: int | None = None,

936 + max_frames: int = 256,

937 + min_frames: int = 8,

938 + video_audio_interleave_length: int = 0,

939 + audio_sampling_rate: int = 24000,

940 + audio_nfft: int = 960,

941 + audio_hop_length: int = 240,

942 + audio_window_size: int = 960,

943 + audio_fmin: float = 0.0,

944 + audio_fmax: float | None = None,

945 + audio_n_mels: int = 128,

946 + audio_segment_size: int = 6000,

947 + audio_kernel_size: int = 3,

948 + audio_stride_size: int = 2,

949 + audio_avg_pooler: int = 2,

950 + audio_channels: int = 8,

951 + audio_group_size: int = 4,

952 + audio_input_id_per_second: float = 25.0,

953 + audio_zeroemb_idx: int = 4096,

954 + image_token_id: int | None = None,

955 + video_token_id: int | None = None,

956 + audio_token_id: int | None = None,

957 + vision_start_token_id: int | None = None,

958 + vision_end_token_id: int | None = None,

959 + audio_start_token_id: int | None = None,

960 + audio_end_token_id: int | None = None,

961 + video_start_token_id: int | None = None,

962 + video_end_token_id: int | None = None,

963 + rope_type: str = "rope",

964 + ) -> None:

965 + self.tokenizer = tokenizer

966 +

967 + unit = patch_size * merge_size

968 + self.mimo_processor = MiMoVLProcessor(

969 + tokenizer=tokenizer,

970 + patch_size=patch_size,

971 + merge_size=merge_size,

972 + temporal_patch_size=temporal_patch_size,

973 + temporal_compression_ratio=temporal_compression_ratio,

974 + use_video_timestamps=True,

975 + video_audio_interleave_length=video_audio_interleave_length,

976 + audio_sampling_rate=audio_sampling_rate,

977 + audio_nfft=audio_nfft,

978 + audio_hop_length=audio_hop_length,

979 + audio_window_size=audio_window_size,

980 + audio_fmin=audio_fmin,

981 + audio_fmax=audio_fmax,

982 + audio_n_mels=audio_n_mels,

983 + audio_segment_size=audio_segment_size,

984 + audio_kernel_size=audio_kernel_size,

985 + audio_stride_size=audio_stride_size,

986 + audio_avg_pooler=audio_avg_pooler,

987 + audio_channels=audio_channels,

988 + audio_group_size=audio_group_size,

989 + audio_input_id_per_second=audio_input_id_per_second,

990 + audio_zeroemb_idx=audio_zeroemb_idx,

991 + image_min_pixels=image_min_pixels or (4 * unit * unit),

992 + image_max_pixels=image_max_pixels or (4096 * unit * unit),

993 + video_min_pixels=video_min_pixels or (4 * unit * unit),

994 + video_max_pixels=video_max_pixels or (4096 * unit * unit),

995 + video_total_max_pixels=video_total_max_pixels or (16384 * unit * unit),

996 + fps=fps,

997 + num_frames=num_frames,

998 + max_frames=max_frames,

999 + min_frames=min_frames,

1000 + image_token_id=image_token_id,

1001 + video_token_id=video_token_id,

1002 + audio_token_id=audio_token_id,

1003 + vision_start_token_id=vision_start_token_id,

1004 + vision_end_token_id=vision_end_token_id,

1005 + audio_start_token_id=audio_start_token_id,

1006 + audio_end_token_id=audio_end_token_id,

1007 + video_start_token_id=video_start_token_id,

1008 + video_end_token_id=video_end_token_id,

1009 + pad_token_id=tokenizer.pad_token_id,

1010 + rope_type=rope_type,

1011 + )

1012 +

1013 + @classmethod

1014 + def from_hf_config(cls, tokenizer: Any, hf_config: Any) -> "MiMoOmniProcessor":

1015 + """Convenience factory: instantiate directly from an HF model config object."""

1016 + vc = hf_config.vision_config

1017 + if isinstance(vc, dict):

1018 + patch_size = vc.get("patch_size", 14)

1019 + merge_size = vc.get("spatial_merge_size", 2)

1020 + temporal_patch_size = vc.get("temporal_patch_size", 2)

1021 + else:

1022 + patch_size = getattr(vc, "patch_size", 14)

1023 + merge_size = getattr(vc, "spatial_merge_size", 2)

1024 + temporal_patch_size = getattr(vc, "temporal_patch_size", 2)

1025 +

1026 + pc: dict = getattr(hf_config, "processor_config", {}) or {}

1027 + ac = getattr(hf_config, "audio_config", None)

1028 + audio_sr: int | None = pc.get("audio_sampling_rate")

1029 + if audio_sr is None and ac is not None:

1030 + if isinstance(ac, dict):

1031 + audio_sr = ac.get("sampling_rate") or ac.get("sample_rate")

1032 + else:

1033 + audio_sr = getattr(ac, "sampling_rate", None) or getattr(

1034 + ac, "sample_rate", None

1035 + )

1036 +

1037 + rope_type = "rope"

1038 + rs = getattr(hf_config, "rope_scaling", None)

1039 + if rs and rs.get("type") == "default" and rs.get("mrope_section") is not None:

1040 + rope_type = "mrope"

1041 +

1042 + unit = patch_size * merge_size

1043 + return cls(

1044 + tokenizer,

1045 + patch_size=patch_size,

1046 + merge_size=merge_size,

1047 + temporal_patch_size=temporal_patch_size,

1048 + image_min_pixels=pc.get("image_min_pixels") or (4 * unit * unit),

1049 + image_max_pixels=pc.get("image_max_pixels") or (4096 * unit * unit),

1050 + video_min_pixels=pc.get("video_min_pixels") or (4 * unit * unit),

1051 + video_max_pixels=pc.get("video_max_pixels") or (4096 * unit * unit),

1052 + video_total_max_pixels=(

1053 + pc.get("video_total_max_pixels") or (16384 * unit * unit)

1054 + ),

1055 + fps=pc.get("fps") or 2.0,

1056 + num_frames=pc.get("num_frames"),

1057 + max_frames=pc.get("max_frames") or 256,

1058 + min_frames=pc.get("min_frames") or 8,

1059 + video_audio_interleave_length=pc.get("video_audio_interleave_length", 0),

1060 + audio_sampling_rate=audio_sr or 24000,

1061 + image_token_id=pc.get("image_token_id"),

1062 + video_token_id=pc.get("video_token_id"),

1063 + audio_token_id=pc.get("audio_token_id"),

1064 + vision_start_token_id=pc.get("vision_start_token_id"),

1065 + vision_end_token_id=pc.get("vision_end_token_id"),

1066 + audio_start_token_id=pc.get("audio_start_token_id"),

1067 + audio_end_token_id=pc.get("audio_end_token_id"),

1068 + video_start_token_id=pc.get("video_start_token_id"),

1069 + video_end_token_id=pc.get("video_end_token_id"),

1070 + rope_type=rope_type,

1071 + )

1072 +

1073 + @property

1074 + def image_token(self) -> str:

1075 + """Token string used as image placeholder (for vLLM integration)."""

1076 + return "<|image_pad|>"

1077 +

1078 + @property

1079 + def video_token(self) -> str:

1080 + """Token string used as video placeholder (for vLLM integration)."""

1081 + return "<|video_pad|>"

1082 +

1083 + @property

1084 + def image_processor(self) -> Any:

1085 + """Minimal image-processor-like object for vLLM processing-info compat."""

1086 + p = self.mimo_processor

1087 +

1088 + class _ImageProcessor:

1089 + merge_size = p.merge_size

1090 + size = {

1091 + "shortest_edge": p._img_kw["min_pixels"],

1092 + "longest_edge": p._img_kw["max_pixels"],

1093 + }

1094 +

1095 + return _ImageProcessor()

1096 +

1097 + def _modality(self, token: str) -> str:

1098 + if self._IMG_RE.fullmatch(token):

1099 + return "image"

1100 + if self._VID_RE.fullmatch(token):

1101 + return "video"

1102 + if self._AUD_RE.fullmatch(token):

1103 + return "audio"

1104 + return "unknown"

1105 +

1106 + def __call__(

1107 + self,

1108 + text: str | list[str] | None = None,

1109 + images: Any = None,

1110 + videos: Any = None,

1111 + audio: Any = None,

1112 + video_audio: Any = None,

1113 + return_tensors: str | TensorType | None = None,

1114 + **kwargs: Any,

1115 + ) -> BatchFeature:

1116 + """Process multimodal inputs into model-ready tensors.

1117 +

1118 + Args:

1119 + text: Prompt string(s) containing multimodal placeholders

1120 + ``<|vision_start|><|image_pad|><|vision_end|>``,

1121 + ``<|vision_start|><|video_pad|><|vision_end|>``, or

1122 + ``<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>``.

1123 + images: PIL.Image or list[PIL.Image].

1124 + videos: list of ``(frames_TCHW: torch.Tensor, timestamps_T: torch.Tensor)``

1125 + tuples (pre-decoded).

1126 + audio: list of ``str`` (path/url/base64), ``bytes``,

1127 + ``(waveform_1D, sample_rate)`` tuples, or ``np.ndarray``.

1128 + return_tensors: Passed to :class:`BatchFeature`.

1129 +

1130 + Returns:

1131 + :class:`BatchFeature` with keys:

1132 + - ``input_ids``

1133 + - ``pixel_values`` + ``image_grid_thw``

1134 + - ``pixel_values_videos`` + ``video_grid_thw`` + ``second_per_grid_ts``

1135 + - ``audio_features``

1136 + """

1137 + if isinstance(text, list):

1138 + text = text[0] if len(text) == 1 else "\n".join(text)

1139 +

1140 + imgs: list = (

1141 + ([images] if isinstance(images, Image.Image) else list(images))

1142 + if images is not None

1143 + else []

1144 + )

1145 + vids: list = list(videos) if videos is not None else []

1146 + auds: list = list(audio) if audio is not None else []

1147 + va_items: list = list(video_audio) if video_audio is not None else []

1148 +

1149 + # If audio exists but text has no audio placeholder, prepend one

1150 + _aud_placeholder = "<|mimo_audio_start|><|audio_pad|><|mimo_audio_end|>"

1151 + if auds and text is not None and not self._AUD_RE.search(text):

1152 + text = _aud_placeholder + text

1153 +

1154 + # Build Content list

1155 + contents: list[Content] = []

1156 +

1157 + if text and (imgs or vids or auds or va_items):

1158 + parts = self._MM_RE.split(text)

1159 + img_it = iter(imgs)

1160 + vid_it = iter(vids)

1161 + aud_it = iter(auds)

1162 + va_it = iter(va_items)

1163 + for part in parts:

1164 + if self._MM_RE.fullmatch(part):

1165 + mod = self._modality(part)

1166 + if mod == "image":

1167 + with contextlib.suppress(StopIteration):

1168 + contents.append(

1169 + Content(

1170 + type="image",

1171 + content=ImageInput(image=next(img_it)),

1172 + )

1173 + )

1174 + elif mod == "video":

1175 + # Try regular video first, fall back to video_audio

1176 + vid_item = None

1177 + vid_type = "video"

1178 + with contextlib.suppress(StopIteration):

1179 + vid_item = next(vid_it)

1180 + if vid_item is None:

1181 + with contextlib.suppress(StopIteration):

1182 + vid_item = next(va_it)

1183 + vid_type = "video_audio"

1184 + if vid_item is not None:

1185 + if vid_type == "video":

1186 + contents.append(

1187 + Content(

1188 + type="video",

1189 + content=VideoInput(video=vid_item),

1190 + )

1191 + )

1192 + else:

1193 + contents.append(

1194 + Content(

1195 + type="video_audio",

1196 + content=vid_item,

1197 + )

1198 + )

1199 + elif mod == "audio":

1200 + with contextlib.suppress(StopIteration):

1201 + contents.append(

1202 + Content(

1203 + type="audio",

1204 + content=AudioInput(audio=next(aud_it)),

1205 + )

1206 + )

1207 + elif part:

1208 + contents.append(Content(type="text", content=part))

1209 + elif text:

1210 + contents.append(Content(type="text", content=text))

1211 + else:

1212 + for img in imgs:

1213 + contents.append(Content(type="image", content=ImageInput(image=img)))

1214 + for vid in vids:

1215 + contents.append(Content(type="video", content=VideoInput(video=vid)))

1216 + for aud in auds:

1217 + contents.append(Content(type="audio", content=AudioInput(audio=aud)))

1218 + for va_item in va_items:

1219 + contents.append(Content(type="video_audio", content=va_item))

1220 +

1221 + if not contents:

1222 + ids = self.tokenizer(text or "", return_tensors=return_tensors)["input_ids"]

1223 + return BatchFeature(data={"input_ids": ids}, tensor_type=return_tensors)

1224 +

1225 + sample = self.mimo_processor.process(contents, verbose=False)

1226 +

1227 + # vLLM expects input_ids to have a batch dimension [1, seq_len].

1228 + data: dict = {"input_ids": sample.input_ids.unsqueeze(0)}

1229 +

1230 + if sample.pixel_values:

1231 + data["pixel_values"] = torch.cat(sample.pixel_values, dim=0)

1232 + data["image_grid_thw"] = torch.stack(sample.image_thw_grids)

1233 +

1234 + if sample.pixel_values_videos:

1235 + data["pixel_values_videos"] = torch.cat(sample.pixel_values_videos, dim=0)

1236 + data["video_grid_thw"] = torch.stack(sample.video_thw_grids)

1237 + if sample.second_per_grid_ts:

1238 + data["second_per_grid_ts"] = torch.tensor(

1239 + sample.second_per_grid_ts, dtype=torch.float32

1240 + )

1241 + if sample.video_start_times:

1242 + data["video_start_times"] = torch.tensor(

1243 + sample.video_start_times, dtype=torch.float32

1244 + )

1245 + if sample.video_audio_n_segs:

1246 + data["video_audio_n_segs"] = torch.tensor(

1247 + sample.video_audio_n_segs, dtype=torch.long

1248 + )

1249 + # video_audio_seg_lens: 2D padded tensor (num_videos, max_T).

1250 + # Row i has the per-group audio token lengths for video i

1251 + # (zeros for regular videos; valid values for video_audio videos).

1252 + n_segs_list = sample.video_audio_n_segs

1253 + max_segs = max(n_segs_list) if n_segs_list else 0

1254 + if max_segs > 0:

1255 + seg_lens_2d = torch.zeros(len(n_segs_list), max_segs, dtype=torch.long)

1256 + flat_cursor = 0

1257 + for vi, n in enumerate(n_segs_list):

1258 + if n > 0:

1259 + seg_lens_2d[vi, :n] = torch.tensor(

1260 + sample.video_audio_seg_lens[flat_cursor : flat_cursor + n],

1261 + dtype=torch.long,

1262 + )

1263 + flat_cursor += n

1264 + data["video_audio_seg_lens"] = seg_lens_2d

1265 +

1266 + # audio_features is a list of variable-length mel-spec tensors; pop it

1267 + # before BatchFeature conversion to avoid "batched tensors of the same

1268 + # length" errors, then re-attach it after.

1269 + audio_features = None

1270 + if sample.audio_inputs:

1271 + audio_features = sample.audio_inputs

1272 + if "is_audio_tokenized" in sample.extra:

1273 + data["is_audio_tokenized"] = sample.extra["is_audio_tokenized"]

1274 + if sample.audio_token_lens:

1275 + data["audio_token_lens"] = torch.tensor(

1276 + sample.audio_token_lens, dtype=torch.long

1277 + )

1278 +

1279 + bf = BatchFeature(data=data, tensor_type=return_tensors)

1280 + if audio_features is not None:

1281 + bf["audio_features"] = audio_features

1282 + # va_audio_features: list of mel-spec tensors (one per video_audio item)

1283 + if sample.va_audio_inputs:

1284 + bf["va_audio_features"] = sample.va_audio_inputs

1285 + return bf

vllm/v1/core/kv_cache_utils.py

@@ -1381,7 +1381,9 @@ def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):

1381 block_size=spec.block_size,

1382 num_kv_heads=spec.num_kv_heads,

1383 head_size=spec.head_size,

1384 + head_size_v=spec.head_size_v,

1385 dtype=spec.dtype,

1386 + kv_quant_mode=spec.kv_quant_mode,

1387 sliding_window=spec.sliding_window,

1388 page_size_padded=spec.page_size_padded,

1389 )

vllm/v1/spec_decode/llm_base_proposer.py

@@ -1344,6 +1344,7 @@ class SpecDecodeBaseProposer:

1344 "Exaone4_5_ForConditionalGeneration",

1345 "GlmOcrForConditionalGeneration",

1346 "HunYuanVLForConditionalGeneration",

1347 + "MiMoV2OmniForCausalLM",

1348 "Qwen2_5_VLForConditionalGeneration",

1349 "Qwen3_5ForConditionalGeneration",

1350 "Qwen3_5MoeForConditionalGeneration",