b8199f6049 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -31,29 +31,34 @@ import torch

31 import torch.nn as nn

32 import torch.nn.functional as F

33 from packaging.version import Version

34 -from transformers import PretrainedConfig

35 -from transformers import __version__ as TRANSFORMERS_VERSION

34 from transformers.feature_extraction_utils import BatchFeature

35 from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (

36 + Qwen3OmniMoeAudioEncoderConfig,

37 Qwen3OmniMoeConfig,

38 Qwen3OmniMoeThinkerConfig,

39 )

41 -from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (

42 - Qwen3OmniMoeAudioEncoder,

43 -)

40 from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (

41 Qwen3OmniMoeProcessor,

42 )

43 from transformers.models.whisper import WhisperFeatureExtractor

44

45 +# isort: off

46 +from transformers import PretrainedConfig

47 +from transformers import __version__ as TRANSFORMERS_VERSION

48 +# isort: on

49 +

50 from vllm.compilation.decorators import support_torch_compile

51 from vllm.config import MultiModalConfig, VllmConfig

51 -from vllm.distributed import get_pp_group

52 +from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size

53 from vllm.logger import init_logger

54 from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY

55 +from vllm.model_executor.layers.attention.mm_encoder_attention import (

56 + MMEncoderAttention,

57 +)

58 from vllm.model_executor.layers.conv import Conv3dLayer

59 from vllm.model_executor.layers.linear import (

60 ColumnParallelLinear,

61 + QKVParallelLinear,

62 RowParallelLinear,

63 )

64 from vllm.model_executor.layers.logits_processor import LogitsProcessor

@@ -104,11 +109,6 @@ from .vision import (

109 get_vit_attn_backend,

110 )

111

107 -try:

108 - import flash_attn

109 -except (ImportError, ModuleNotFoundError):

110 - flash_attn = None

111 -

112 logger = init_logger(__name__)

113

114

@@ -121,6 +121,415 @@ def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):

121 return output_lengths

122

123

124 +# ============= Audio Encoder Components =============

125 +

126 +

127 +class SinusoidsPositionEmbedding(nn.Module):

128 + """Sinusoidal position embedding for audio encoder."""

129 +

130 + def __init__(self, length: int, channels: int, max_timescale: int = 10000):

131 + super().__init__()

132 + self.length = length

133 + self.channels = channels

134 + self.max_timescale = max_timescale

135 +

136 + if channels % 2 != 0:

137 + raise ValueError("SinusoidsPositionEmbedding needs even channels input")

138 +

139 + log_timescale_increment = np.log(max_timescale) / (channels // 2 - 1)

140 + inv_timescales = torch.exp(

141 + -log_timescale_increment * torch.arange(channels // 2).float()

142 + )

143 + scaled_time = (

144 + torch.arange(length)[:, np.newaxis] * inv_timescales[np.newaxis, :]

145 + )

146 + positional_embedding = torch.cat(

147 + [torch.sin(scaled_time), torch.cos(scaled_time)], dim=1

148 + )

149 + self.register_buffer(

150 + "positional_embedding", positional_embedding, persistent=False

151 + )

152 +

153 + def forward(self, seqlen: int) -> torch.Tensor:

154 + return self.positional_embedding[:seqlen, :]

155 +

156 +

157 +class Qwen3OmniMoeAudioAttention(nn.Module):

158 + """Multi-headed attention for Qwen3-Omni Audio Encoder using MMEncoderAttention."""

159 +

160 + def __init__(

161 + self,

162 + config: Qwen3OmniMoeAudioEncoderConfig,

163 + multimodal_config: MultiModalConfig | None = None,

164 + prefix: str = "",

165 + ):

166 + super().__init__()

167 + self.embed_dim = config.d_model

168 + self.num_heads = config.encoder_attention_heads

169 + self.head_dim = self.embed_dim // self.num_heads

170 + tp_size = get_tensor_model_parallel_world_size()

171 + self.num_local_heads = self.num_heads // tp_size

172 +

173 + if (self.head_dim * self.num_heads) != self.embed_dim:

174 + raise ValueError(

175 + f"embed_dim must be divisible by num_heads (got `embed_dim`: "

176 + f"{self.embed_dim} and `num_heads`: {self.num_heads})."

177 + )

178 +

179 + self.scaling = self.head_dim**-0.5

180 +

181 + self.qkv = QKVParallelLinear(

182 + hidden_size=self.embed_dim,

183 + head_size=self.head_dim,

184 + total_num_heads=self.num_heads,

185 + total_num_kv_heads=self.num_heads,

186 + bias=True,

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

188 + )

189 +

190 + self.out_proj = RowParallelLinear(

191 + input_size=self.embed_dim,

192 + output_size=self.embed_dim,

193 + bias=True,

194 + prefix=f"{prefix}.out_proj",

195 + )

196 +

197 + self.attn = MMEncoderAttention(

198 + num_heads=self.num_local_heads,

199 + head_size=self.head_dim,

200 + scale=self.scaling,

201 + multimodal_config=multimodal_config,

202 + )

203 +

204 + def forward(

205 + self,

206 + hidden_states: torch.Tensor,

207 + cu_seqlens: torch.Tensor,

208 + max_seqlen: torch.Tensor | None,

209 + ) -> torch.Tensor:

210 + seq_length, _ = hidden_states.size()

211 + qkv, _ = self.qkv(hidden_states)

212 + q, k, v = qkv.chunk(3, dim=-1)

213 + q = q.view(1, seq_length, -1, self.head_dim)

214 + k = k.view(1, seq_length, -1, self.head_dim)

215 + v = v.view(1, seq_length, -1, self.head_dim)

216 +

217 + attn_output = self.attn(

218 + query=q,

219 + key=k,

220 + value=v,

221 + cu_seqlens=cu_seqlens,

222 + max_seqlen=max_seqlen,

223 + )

224 +

225 + attn_output = attn_output.view(seq_length, -1)

226 + output, _ = self.out_proj(attn_output)

227 + return output

228 +

229 +

230 +class Qwen3OmniMoeAudioEncoderLayer(nn.Module):

231 + """Transformer encoder layer for Qwen3-Omni Audio Encoder."""

232 +

233 + def __init__(

234 + self,

235 + config: Qwen3OmniMoeAudioEncoderConfig,

236 + multimodal_config: MultiModalConfig | None = None,

237 + prefix: str = "",

238 + ):

239 + super().__init__()

240 + self.embed_dim = config.d_model

241 + self.self_attn = Qwen3OmniMoeAudioAttention(

242 + config, multimodal_config=multimodal_config, prefix=f"{prefix}.self_attn"

243 + )

244 + self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)

245 + self.activation_fn = _ACTIVATION_REGISTRY[config.activation_function]

246 + self.fc1 = ColumnParallelLinear(

247 + self.embed_dim,

248 + config.encoder_ffn_dim,

249 + bias=True,

250 + prefix=f"{prefix}.fc1",

251 + )

252 + self.fc2 = RowParallelLinear(

253 + config.encoder_ffn_dim,

254 + self.embed_dim,

255 + bias=True,

256 + prefix=f"{prefix}.fc2",

257 + )

258 + self.final_layer_norm = nn.LayerNorm(self.embed_dim)

259 +

260 + def forward(

261 + self,

262 + hidden_states: torch.Tensor,

263 + cu_seqlens: torch.Tensor,

264 + max_seqlen: torch.Tensor | None,

265 + ) -> torch.Tensor:

266 + """

267 + Args:

268 + hidden_states: Input tensor of shape (seq_len, hidden_size)

269 + cu_seqlens: Cumulative sequence lengths

270 + max_seqlen: Maximum sequence length in the batch

271 + """

272 + residual = hidden_states

273 + hidden_states = self.self_attn_layer_norm(hidden_states)

274 + hidden_states = self.self_attn(

275 + hidden_states=hidden_states,

276 + cu_seqlens=cu_seqlens,

277 + max_seqlen=max_seqlen,

278 + )

279 + hidden_states = residual + hidden_states

280 +

281 + residual = hidden_states

282 + hidden_states = self.final_layer_norm(hidden_states)

283 + hidden_states, _ = self.fc1(hidden_states)

284 + hidden_states = self.activation_fn(hidden_states)

285 + hidden_states, _ = self.fc2(hidden_states)

286 + hidden_states = residual + hidden_states

287 +

288 + # Clamp for numerical stability with fp16

289 + if hidden_states.dtype == torch.float16:

290 + clamp_value = torch.finfo(hidden_states.dtype).max - 1000

291 + hidden_states = torch.clamp(

292 + hidden_states, min=-clamp_value, max=clamp_value

293 + )

294 +

295 + return hidden_states

296 +

297 +

298 +class Qwen3OmniMoeAudioEncoder(nn.Module):

299 + """vLLM-native Qwen3-Omni Audio Encoder."""

300 +

301 + def __init__(

302 + self,

303 + config: Qwen3OmniMoeAudioEncoderConfig,

304 + multimodal_config: MultiModalConfig | None = None,

305 + prefix: str = "",

306 + ):

307 + super().__init__()

308 +

309 + embed_dim = config.d_model

310 + self.num_mel_bins = config.num_mel_bins

311 + self.max_source_positions = config.max_source_positions

312 + self.n_window = config.n_window

313 + self.n_window_infer = config.n_window_infer

314 + self.conv_chunksize = config.conv_chunksize

315 +

316 + # Position embedding

317 + self.positional_embedding = SinusoidsPositionEmbedding(

318 + self.max_source_positions, embed_dim

319 + )

320 +

321 + # Convolutional layers for mel-spectrogram processing

322 + self.conv2d1 = nn.Conv2d(1, config.downsample_hidden_size, 3, 2, padding=1)

323 + self.conv2d2 = nn.Conv2d(

324 + config.downsample_hidden_size,

325 + config.downsample_hidden_size,

326 + 3,

327 + 2,

328 + padding=1,

329 + )

330 + self.conv2d3 = nn.Conv2d(

331 + config.downsample_hidden_size,

332 + config.downsample_hidden_size,

333 + 3,

334 + 2,

335 + padding=1,

336 + )

337 +

338 + conv_out_dim = config.downsample_hidden_size * (

339 + (((config.num_mel_bins + 1) // 2 + 1) // 2 + 1) // 2

340 + )

341 + self.conv_out = nn.Linear(conv_out_dim, config.d_model, bias=False)

342 +

343 + # Transformer encoder layers

344 + self.layers = nn.ModuleList(

345 + [

346 + Qwen3OmniMoeAudioEncoderLayer(

347 + config,

348 + multimodal_config=multimodal_config,

349 + prefix=f"{prefix}.layers.{i}",

350 + )

351 + for i in range(config.encoder_layers)

352 + ]

353 + )

354 +

355 + # Output layers

356 + self.ln_post = nn.LayerNorm(config.d_model)

357 + self.proj1 = nn.Linear(config.d_model, config.d_model)

358 + self.act = _ACTIVATION_REGISTRY[config.activation_function]

359 + self.proj2 = nn.Linear(config.d_model, config.output_dim)

360 +

361 + # Get attention backend

362 + attn_backend_override = (

363 + multimodal_config.mm_encoder_attn_backend

364 + if multimodal_config is not None

365 + else None

366 + )

367 + self.attn_backend = get_vit_attn_backend(

368 + head_size=config.d_model // config.encoder_attention_heads,

369 + dtype=torch.get_default_dtype(),

370 + attn_backend_override=attn_backend_override,

371 + )

372 +

373 + def compute_attn_mask_seqlen(self, cu_seqlens: torch.Tensor) -> torch.Tensor | None:

374 + """Compute max_seqlen only for flash attention backends."""

375 + max_seqlen = None

376 + if self.attn_backend in {

377 + AttentionBackendEnum.FLASH_ATTN,

378 + AttentionBackendEnum.ROCM_AITER_FA,

379 + }:

380 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max()

381 + return max_seqlen

382 +

383 + @property

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

385 + return self.conv2d1.weight.dtype

386 +

387 + @property

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

389 + return self.conv2d1.weight.device

390 +

391 + def forward(

392 + self,

393 + input_features: torch.Tensor,

394 + feature_lens: torch.Tensor,

395 + aftercnn_lens: torch.Tensor,

396 + ):

397 + # Compute chunk information

398 + chunk_num = torch.ceil(feature_lens / (self.n_window * 2)).long()

399 +

400 + chunk_lengths = torch.tensor(

401 + [self.n_window * 2] * chunk_num.sum(),

402 + dtype=torch.long,

403 + device=feature_lens.device,

404 + )

405 + tail_chunk_index = F.pad(chunk_num, (1, 0), value=-1).cumsum(0)[1:]

406 + chunk_lengths[tail_chunk_index] = feature_lens % (self.n_window * 2)

407 + chunk_lengths[chunk_lengths == 0] = self.n_window * 2

408 +

409 + # Split input features into chunks and pad

410 + chunk_list = input_features.T.split(chunk_lengths.tolist(), dim=0)

411 + padded_feature = nn.utils.rnn.pad_sequence(

412 + chunk_list, batch_first=True

413 + ).transpose(1, 2)

414 +

415 + # Compute feature lengths after CNN

416 + feature_lens_after_cnn = self._get_cnn_output_lengths(chunk_lengths)

417 + # Vectorized mask creation: avoid creating many small tensors

418 + max_len_after_cnn = feature_lens_after_cnn.max().item()

419 + indices = torch.arange(max_len_after_cnn, device=padded_feature.device)

420 + padded_mask_after_cnn = indices.unsqueeze(0) < feature_lens_after_cnn.unsqueeze(

421 + 1

422 + )

423 +

424 + # Add channel dimension for conv2d

425 + padded_feature = padded_feature.unsqueeze(1)

426 +

427 + # Apply convolutional layers (chunk if needed to avoid OOM)

428 + if padded_feature.size(0) <= self.conv_chunksize:

429 + # Fast path: no chunking needed

430 + padded_embed = F.gelu(self.conv2d1(padded_feature))

431 + padded_embed = F.gelu(self.conv2d2(padded_embed))

432 + padded_embed = F.gelu(self.conv2d3(padded_embed))

433 + else:

434 + # Chunked processing to avoid OOM

435 + padded_embeds = []

436 + for chunk in padded_feature.split(self.conv_chunksize, dim=0):

437 + padded_embed = F.gelu(self.conv2d1(chunk))

438 + padded_embed = F.gelu(self.conv2d2(padded_embed))

439 + padded_embed = F.gelu(self.conv2d3(padded_embed))

440 + padded_embeds.append(padded_embed)

441 + padded_embed = torch.cat(padded_embeds, dim=0)

442 +

443 + # (batch, channels, freq, time) -> (batch, time, channels*freq)

444 + b, c, f, t = padded_embed.size()

445 + padded_embed = self.conv_out(

446 + padded_embed.permute(0, 3, 1, 2).contiguous().view(b, t, c * f)

447 + )

448 +

449 + # Add positional embedding

450 + positional_embedding = (

451 + self.positional_embedding.positional_embedding[: padded_embed.shape[1], :]

452 + .unsqueeze(0)

453 + .to(padded_embed.dtype)

454 + )

455 + padded_embed = padded_embed + positional_embedding

456 +

457 + # Extract valid hidden states and compute cu_seqlens

458 + hidden_states = padded_embed[padded_mask_after_cnn]

459 +

460 + # Compute cumulative sequence lengths for chunked attention

461 + cu_chunk_lens = [0]

462 + window_aftercnn = padded_mask_after_cnn.shape[-1] * (

463 + self.n_window_infer // (self.n_window * 2)

464 + )

465 + # Use tolist() for efficient batch conversion from tensor to Python

466 + for cnn_len in aftercnn_lens.tolist():

467 + num_full_chunks = cnn_len // window_aftercnn

468 + remainder = cnn_len % window_aftercnn

469 + cu_chunk_lens.extend([window_aftercnn] * num_full_chunks)

470 + if remainder:

471 + cu_chunk_lens.append(remainder)

472 + cu_seqlens = torch.tensor(cu_chunk_lens, device=aftercnn_lens.device).cumsum(

473 + -1, dtype=torch.int32

474 + )

475 +

476 + max_seqlen = self.compute_attn_mask_seqlen(cu_seqlens)

477 +

478 + # Apply transformer layers

479 + for encoder_layer in self.layers:

480 + hidden_states = encoder_layer(

481 + hidden_states,

482 + cu_seqlens,

483 + max_seqlen,

484 + )

485 +

486 + # Apply output layers

487 + hidden_states = self.ln_post(hidden_states)

488 + hidden_states = self.proj1(hidden_states)

489 + hidden_states = self.act(hidden_states)

490 + hidden_states = self.proj2(hidden_states)

491 +

492 + return hidden_states

493 +

494 + def _get_cnn_output_lengths(self, input_lengths: torch.Tensor) -> torch.Tensor:

495 + """Compute output lengths after the three conv2d layers."""

496 + lengths = input_lengths

497 + for _ in range(3):

498 + lengths = (lengths - 1) // 2 + 1

499 + return lengths

500 +

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

502 + """Load weights with mapping from HuggingFace format."""

503 + stacked_params_mapping = [

504 + # (param_name, shard_name, shard_id)

505 + ("self_attn.qkv.", "self_attn.q_proj.", "q"),

506 + ("self_attn.qkv.", "self_attn.k_proj.", "k"),

507 + ("self_attn.qkv.", "self_attn.v_proj.", "v"),

508 + ]

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

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

511 +

512 + for name, loaded_weight in weights:

513 + for param_name, weight_name, shard_id in stacked_params_mapping:

514 + if weight_name not in name:

515 + continue

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

517 +

518 + param = params_dict[name]

519 + weight_loader = param.weight_loader

520 + weight_loader(param, loaded_weight, shard_id)

521 + break

522 + else:

523 + param = params_dict.get(name)

524 + if param is not None:

525 + weight_loader = getattr(

526 + param, "weight_loader", default_weight_loader

527 + )

528 + weight_loader(param, loaded_weight)

529 + loaded_params.add(name)

530 + return loaded_params

531 +

532 +

533 class Qwen3_VisionPatchEmbed(nn.Module):

534 def __init__(

535 self,

@@ -144,7 +553,7 @@ class Qwen3_VisionPatchEmbed(nn.Module):

553 )

554

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

147 - L, C = x.shape

556 + L, _ = x.shape

557 x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)

558 x = self.proj(x).view(L, self.hidden_size)

559 return x

@@ -224,7 +633,7 @@ class Qwen3_VisionBlock(nn.Module):

633 cu_seqlens: torch.Tensor,

634 rotary_pos_emb_cos: torch.Tensor,

635 rotary_pos_emb_sin: torch.Tensor,

227 - max_seqlen: torch.Tensor, # Only used for Flash Attention

636 + max_seqlen: torch.Tensor | None, # Only used for Flash Attention

637 ) -> torch.Tensor:

638 x = x + self.attn(

639 self.norm1(x),

@@ -1142,12 +1551,11 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix

1551

1552 audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths)

1553

1145 - audio_outputs = self.audio_tower(

1554 + audio_features = self.audio_tower(

1555 input_features.to(self.audio_tower.dtype),

1556 feature_lens=audio_feature_lengths,

1557 aftercnn_lens=audio_output_lengths,

1558 )

1150 - audio_features = audio_outputs.last_hidden_state

1559 return audio_features.split(audio_output_lengths.tolist())

1560

1561

@@ -1205,20 +1613,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(

1613 self.config = thinker_config

1614 self.multimodal_config = multimodal_config

1615

1208 - # force "use_flash_attention_2=True" to audio tower to align

1209 - # the results.

1210 - if flash_attn is not None:

1211 - audio_config = thinker_config.audio_config

1212 - audio_config._attn_implementation_autoset = True

1213 - audio_config._attn_implementation = "flash_attention_2"

1214 - else:

1215 - logger.warning(

1216 - "flash_attn is not available, the model may not yield the "

1217 - "exactly same result as the transformers implementation "

1218 - "in the audio tower part."

1219 - )

1220 -

1221 - self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)

1616 + self.audio_tower = Qwen3OmniMoeAudioEncoder(

1617 + thinker_config.audio_config,

1618 + multimodal_config=multimodal_config,

1619 + prefix=maybe_prefix(prefix, "audio_tower"),

1620 + )

1621

1622 self.visual = Qwen3Omni_VisionTransformer(

1623 vision_config=thinker_config.vision_config,

[Model] Re-implement Qwen3Omni Audio Encoder (#32167)