diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py index d0ef7dbfd6..af17094d4d 100644 --- a/vllm/models/dots3_note/nvidia/attention.py +++ b/vllm/models/dots3_note/nvidia/attention.py @@ -37,6 +37,7 @@ from vllm.v1.attention.backends.mla.triton_mla import ( TritonMLAImpl, TritonMLAMetadataBuilder, ) +from vllm.v1.kv_cache_interface import KVCacheLayout from vllm.v1.worker.workspace import ( current_workspace_manager, is_workspace_manager_initialized, @@ -677,6 +678,13 @@ class Dots3NoteTritonMLAImpl(TritonMLAImpl): class Dots3NotePaddedSparseBackend(FlashAttnMLASparseBackend): """NOTE DSA backend for cache rows padded to the SWA latent width.""" + @classmethod + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: + # NOTE packs the smaller DSA index pages beside the padded MLA/SWA pages. + # Keep the layer dimension inside the block dimension so mixed page sizes + # remain representable by the shared KV-cache allocation. + return (KVCacheLayout.BLHNC,) + @staticmethod def get_name() -> str: return "DOTS3_NOTE_PADDED_MLA_SPARSE" diff --git a/vllm/models/dots3_note/nvidia/audio.py b/vllm/models/dots3_note/nvidia/audio.py index 26ba964470..8c276e84b7 100644 --- a/vllm/models/dots3_note/nvidia/audio.py +++ b/vllm/models/dots3_note/nvidia/audio.py @@ -21,6 +21,7 @@ DEFAULT_MERGE_FACTOR = 1 N_SAMPLES = DEFAULT_CHUNK_LENGTH_S * SAMPLE_RATE STRIDE = HOP_LENGTH * DEFAULT_CONV_TEMPORAL_STRIDE * DEFAULT_MERGE_FACTOR ENCODER_SEQ_LEN = N_SAMPLES // HOP_LENGTH // 2 +_AUDIO_FORWARD_MAX_SEGMENTS = 16 class Dots3NoteAudioConfig: @@ -126,6 +127,17 @@ def log_mel_spectrogram(audio, n_mels=128): return log_spec +def batched_log_mel_spectrogram(audio, n_mels=128): + window = _hann_window(audio.device) + stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True) + magnitudes = stft[..., :-1].abs() ** 2 + filters = _mel_filters(audio.device, n_mels) + mel_spec = filters @ magnitudes + log_spec = torch.clamp(mel_spec, min=1e-10).log10() + log_spec = torch.maximum(log_spec, log_spec.amax(dim=(-2, -1), keepdim=True) - 8.0) + return (log_spec + 4.0) / 4.0 + + def compute_audio_token_length( num_samples, *, @@ -181,14 +193,34 @@ class DotsEncoderWithMask(nn.Module): input_seq_lens: torch.Tensor, audio_sample_lens: list[int], ) -> torch.Tensor: - """Run the eager speech encoder without server-side slicing/batching.""" mel_features = mel_features.to(dtype=torch.bfloat16, device=self.device) - return self.speech_encoder( - mel_features, - return_dict=True, - input_seq_lens=input_seq_lens, - audio_sample_lens=audio_sample_lens, - ).last_hidden_state + if mel_features.shape[0] <= _AUDIO_FORWARD_MAX_SEGMENTS: + return self.speech_encoder( + mel_features, + return_dict=True, + input_seq_lens=input_seq_lens, + audio_sample_lens=audio_sample_lens, + ).last_hidden_state + + chunks = [] + for start in range(0, mel_features.shape[0], _AUDIO_FORWARD_MAX_SEGMENTS): + end = start + _AUDIO_FORWARD_MAX_SEGMENTS + chunks.append( + self.speech_encoder( + mel_features[start:end], + return_dict=True, + input_seq_lens=input_seq_lens[start:end], + audio_sample_lens=audio_sample_lens[start:end], + ).last_hidden_state + ) + max_seq_len = max(chunk.shape[1] for chunk in chunks) + chunks = [ + F.pad(chunk, (0, 0, 0, max_seq_len - chunk.shape[1])) + if chunk.shape[1] < max_seq_len + else chunk + for chunk in chunks + ] + return torch.cat(chunks, dim=0) def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor: segments = [] @@ -233,6 +265,53 @@ class DotsEncoderWithMask(nn.Module): ) return torch.cat(chunk_embeddings, dim=0).unsqueeze(0) + def encode_waveforms(self, audio_waveforms) -> list[torch.Tensor]: + stride = HOP_LENGTH * self.conv_temporal_stride * self.merge_factor + segment_token_lengths = [] + segment_sample_lengths = [] + audio_segment_counts = [] + padded_segments = [] + + for audio_waveform in audio_waveforms: + segment_count = 0 + time_step = 0 + while time_step * SAMPLE_RATE < audio_waveform.shape[0]: + segment = audio_waveform[ + time_step * SAMPLE_RATE : (time_step + self.chunk_seconds) + * SAMPLE_RATE + ] + segment_length = segment.shape[0] + padded_segments.append( + pad_or_trim(segment.flatten(), length=self.chunk_samples) + ) + segment_token_lengths.append((segment_length - 1) // stride + 1) + segment_sample_lengths.append(segment_length) + segment_count += 1 + time_step += self.chunk_seconds + audio_segment_counts.append(segment_count) + + mel_features = batched_log_mel_spectrogram(torch.stack(padded_segments)) + assert mel_features.shape[-1] == self.chunk_mel_frames + input_seq_lens = ( + torch.tensor(segment_token_lengths, dtype=torch.long) * self.merge_factor + ) + audio_embedding = self._forward_speech_encoder( + mel_features, input_seq_lens, segment_sample_lengths + ) + + embeddings = [] + segment_idx = 0 + for segment_count in audio_segment_counts: + chunks = [] + for _ in range(segment_count): + token_length = segment_token_lengths[segment_idx] + chunks.append( + audio_embedding[segment_idx, : token_length * self.merge_factor, :] + ) + segment_idx += 1 + embeddings.append(torch.cat(chunks).unsqueeze(0)) + return embeddings + class AudioAdapter(nn.Module): def __init__(self, in_dim, out_dim): @@ -294,12 +373,15 @@ class Dots3NoteAudioModel(nn.Module): audio_start += length return waveforms + def _forward_batched(self, waveforms): + raw_embeddings = self.dots_encoder.encode_waveforms(waveforms) + merged = [ + self._merge_embeddings(embedding).squeeze(0) for embedding in raw_embeddings + ] + token_lengths = [embedding.shape[0] for embedding in merged] + packed = self.audio_adapter(torch.cat(merged)) + return packed, token_lengths + def forward(self, audio_inputs, lengths): waveforms = self._split_waveforms(audio_inputs, lengths) - all_embeddings = [] - token_lengths = [] - for waveform in waveforms: - embedding = self._encode_single_audio(waveform) - all_embeddings.append(embedding) - token_lengths.append(embedding.shape[0]) - return torch.cat(all_embeddings, dim=0), token_lengths + return self._forward_batched(waveforms) diff --git a/vllm/models/dots3_note/nvidia/multimodal.py b/vllm/models/dots3_note/nvidia/multimodal.py index 00f5091822..421d185bb2 100644 --- a/vllm/models/dots3_note/nvidia/multimodal.py +++ b/vllm/models/dots3_note/nvidia/multimodal.py @@ -3,6 +3,7 @@ """vLLM composition layer for Dots3Note image and audio encoders.""" from collections.abc import Iterable +from contextlib import contextmanager import torch from torch import nn @@ -41,6 +42,21 @@ from .model import Dots3NoteLanguageModelForCausalLM from .vision import DotsMoEVitConfig, DotsMoEVitModel +@contextmanager +def _skip_linear_init(): + """Skip initialization for VE Linear weights loaded immediately afterward.""" + original_reset_parameters = nn.Linear.reset_parameters + + def _noop_reset_parameters(_linear: nn.Linear) -> None: + return None + + nn.Linear.reset_parameters = _noop_reset_parameters + try: + yield + finally: + nn.Linear.reset_parameters = original_reset_parameters + + @MULTIMODAL_REGISTRY.register_processor( Dots3NoteMultiModalProcessor, info=Dots3NoteProcessingInfo, @@ -113,7 +129,10 @@ class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): self.visual: DotsMoEVitModel | None if image_enabled: assert vision_config_dict is not None - self.visual = DotsMoEVitModel(DotsMoEVitConfig(**vision_config_dict)) + with _skip_linear_init(): + self.visual = DotsMoEVitModel( + DotsMoEVitConfig(**vision_config_dict) + ) else: self.visual = None self.audio_tower: Dots3NoteAudioModel | None @@ -158,9 +177,10 @@ class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP): ) -> tuple[torch.Tensor, ...]: if self.visual is None: return () - image_embeds = self.visual(pixel_values, image_grid_thw) merge_size = self.visual.spatial_merge_size sizes = (image_grid_thw.prod(-1) // merge_size**2).tolist() + meta = self.visual.prepare_meta(image_grid_thw) + image_embeds = self.visual(pixel_values, meta) return image_embeds.split(sizes) def _process_audio_input( diff --git a/vllm/models/dots3_note/nvidia/vision.py b/vllm/models/dots3_note/nvidia/vision.py index c6f560fcf2..70ea31cb6f 100644 --- a/vllm/models/dots3_note/nvidia/vision.py +++ b/vllm/models/dots3_note/nvidia/vision.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # ruff: noqa: E501 import math +from dataclasses import dataclass from typing import Any import torch @@ -15,10 +16,11 @@ from vllm.third_party.deep_gemm import per_block_cast_to_fp8 from .vision_attention import ( VisionRotaryEmbedding, + VisionRotaryPositionEmbedding, apply_vision_attention_residual, attn_uses_seqlens, build_vision_attention, - prepare_seqlens_for_attention, + prepare_rotary_pos_emb_vision, resolve_attn_implementation, ) from .vision_moe import note_vision_fused_moe_fp8 @@ -502,6 +504,15 @@ _ADAPTER_CLASSES: dict[str, type[nn.Module]] = { } +@dataclass(frozen=True) +class VitForwardMeta: + cu_seqlens: torch.Tensor + max_seqlen: int + rotary_pos_emb: VisionRotaryPositionEmbedding + grid_thw_cpu: torch.Tensor + seqlens: list[int] | None = None + + # ---- Full Model ---- @@ -562,7 +573,7 @@ class DotsMoEVitModel(PreTrainedModel): block.norm_1 = torch.compile(block.norm_1, **compile_kwargs) block.norm_2 = torch.compile(block.norm_2, **compile_kwargs) - def get_pos_ids_by_grid(self, grid_thw): + def get_pos_ids_by_grid(self, grid_thw_cpu: torch.Tensor): # Mirrors ``cybertron`` ``AIMv2NativeModel.rot_pos_emb``: when ``pre_pixel_shuffle`` # is set, RoPE positions follow the qwen ``merge_size`` grouped layout (default 2x2); # otherwise positions are flat row-major regardless of ``spatial_merge_size``. @@ -573,10 +584,8 @@ class DotsMoEVitModel(PreTrainedModel): else: rope_merge_size = 1 pos_ids = [] - for t, h, w in grid_thw: - hpos_ids = ( - torch.arange(h, device=grid_thw.device).unsqueeze(1).expand(-1, w) - ) + for t, h, w in grid_thw_cpu: + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) hpos_ids = hpos_ids.reshape( h // rope_merge_size, rope_merge_size, @@ -585,9 +594,7 @@ class DotsMoEVitModel(PreTrainedModel): ) hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten() - wpos_ids = ( - torch.arange(w, device=grid_thw.device).unsqueeze(0).expand(h, -1) - ) + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) wpos_ids = wpos_ids.reshape( h // rope_merge_size, rope_merge_size, @@ -598,80 +605,67 @@ class DotsMoEVitModel(PreTrainedModel): pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) return pos_ids - def rot_pos_emb(self, grid_thw): - pos_ids = self.get_pos_ids_by_grid(grid_thw) - pos_ids = torch.cat(pos_ids, dim=0) - max_grid_size = grid_thw[:, 1:].max() + def rot_pos_emb(self, grid_thw_cpu: torch.Tensor): + pos_ids = torch.cat(self.get_pos_ids_by_grid(grid_thw_cpu), dim=0) + pos_ids = pos_ids.to(self.device, non_blocking=True) + max_grid_size = int(grid_thw_cpu[:, 1:].max()) rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) - rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) - return rotary_pos_emb + return rotary_pos_emb_full[pos_ids].flatten(1) - def _build_cu_seqlens_from_grid(self, grid_thw: torch.Tensor): - cu_seqlens = torch.repeat_interleave( - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] - ).cumsum( - dim=0, - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - # Same as (cu_seqlens[1:] - cu_seqlens[:-1]).max(); computed here to avoid D2H inside attention. - max_seqlen = int((grid_thw[:, 1] * grid_thw[:, 2]).max().item()) - return cu_seqlens, max_seqlen + def prepare_meta(self, grid_thw_cpu: torch.Tensor) -> VitForwardMeta: + if grid_thw_cpu.device.type != "cpu": + raise ValueError("NOTE vision grid_thw must remain on CPU") - def _build_single_temporal_cu_seqlens_from_grid(self, grid_thw: torch.Tensor): - seq_lens = grid_thw[:, 1] * grid_thw[:, 2] - cu_seqlens = torch.empty( - (grid_thw.shape[0] + 1,), device=grid_thw.device, dtype=torch.int32 + seq_per_frame = grid_thw_cpu[:, 1] * grid_thw_cpu[:, 2] + frame_lengths = torch.repeat_interleave(seq_per_frame, grid_thw_cpu[:, 0]) + cu_seqlens_cpu = F.pad( + frame_lengths.cumsum(dim=0, dtype=torch.int32), (1, 0), value=0 + ) + resolved_attn = resolve_attn_implementation( + self.config.attn_implementation, eager_fallback="eager_v2" + ) + return VitForwardMeta( + cu_seqlens=cu_seqlens_cpu.to(self.device, non_blocking=True), + max_seqlen=int(seq_per_frame.max()), + rotary_pos_emb=prepare_rotary_pos_emb_vision( + self.rot_pos_emb(grid_thw_cpu) + ), + grid_thw_cpu=grid_thw_cpu, + seqlens=( + frame_lengths.tolist() if attn_uses_seqlens(resolved_attn) else None + ), ) - cu_seqlens[0] = 0 - torch.cumsum(seq_lens, dim=0, dtype=torch.int32, out=cu_seqlens[1:]) - max_seqlen = int(seq_lens.max().item()) - return cu_seqlens, max_seqlen def forward( - self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True + self, hidden_states: torch.Tensor, meta: VitForwardMeta, bf16=True ) -> torch.Tensor: - grid_thw = grid_thw.to(device=self.device) if bf16: - hidden_states = hidden_states.to(device=self.device, dtype=self.dtype) - hidden_states = self.patch_embed(hidden_states) - - rotary_pos_emb = self.rot_pos_emb(grid_thw) - - if grid_thw[:, 0].sum().item() == grid_thw.shape[0]: - cu_seqlens, max_seqlen = self._build_single_temporal_cu_seqlens_from_grid( - grid_thw + hidden_states = hidden_states.to( + device=self.device, dtype=self.dtype, non_blocking=True ) - else: - cu_seqlens, max_seqlen = self._build_cu_seqlens_from_grid(grid_thw) - - seqlens = prepare_seqlens_for_attention( - self.config.attn_implementation, - cu_seqlens, - eager_fallback="eager_v2", - ) + hidden_states = self.patch_embed(hidden_states) for blk in self.blocks: if self.gradient_checkpointing and self.training: hidden_states = self._gradient_checkpointing_func( blk.__call__, hidden_states, - cu_seqlens, - rotary_pos_emb, - max_seqlen, - seqlens, + meta.cu_seqlens, + meta.rotary_pos_emb, + meta.max_seqlen, + meta.seqlens, ) else: hidden_states = blk( hidden_states, - cu_seqlens, - rotary_pos_emb, - max_seqlen, - seqlens=seqlens, + meta.cu_seqlens, + meta.rotary_pos_emb, + meta.max_seqlen, + seqlens=meta.seqlens, ) if self.config.post_norm: hidden_states = self.post_trunk_norm(hidden_states) - hidden_states = self.adapter(hidden_states, grid_thw) + hidden_states = self.adapter(hidden_states, meta.grid_thw_cpu) return hidden_states diff --git a/vllm/models/dots3_note/nvidia/vision_attention.py b/vllm/models/dots3_note/nvidia/vision_attention.py index 38c9c5a146..4d5821094b 100644 --- a/vllm/models/dots3_note/nvidia/vision_attention.py +++ b/vllm/models/dots3_note/nvidia/vision_attention.py @@ -36,15 +36,28 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) +VisionRotaryPositionEmbedding = tuple[torch.Tensor, torch.Tensor] + + +def prepare_rotary_pos_emb_vision( + freqs: torch.Tensor, +) -> VisionRotaryPositionEmbedding: + """Materialize vision RoPE cos/sin once for every encoder layer.""" + cos = freqs.cos().unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() + sin = freqs.sin().unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() + return cos, sin + + def apply_rotary_pos_emb_vision( - tensor: torch.Tensor, freqs: torch.Tensor + tensor: torch.Tensor, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding, ) -> torch.Tensor: orig_dtype = tensor.dtype tensor = tensor.float() - cos = freqs.cos() - sin = freqs.sin() - cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() - sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float() + if isinstance(rotary_pos_emb, torch.Tensor): + cos, sin = prepare_rotary_pos_emb_vision(rotary_pos_emb) + else: + cos, sin = rotary_pos_emb output = (tensor * cos) + (rotate_half(tensor) * sin) return output.to(orig_dtype) @@ -149,7 +162,7 @@ class _VisionAttentionBase(nn.Module): def _qkv_with_rope( self, hidden_states: torch.Tensor, - rotary_pos_emb: torch.Tensor | None, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: seq_length = hidden_states.shape[0] q, k, v = ( @@ -161,6 +174,8 @@ class _VisionAttentionBase(nn.Module): if self.use_qk_norm: q = self.q_norm(q) k = self.k_norm(k) + if rotary_pos_emb is None: + raise ValueError("Vision rotary position embedding is required") q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0) k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0) return q, k, v @@ -174,7 +189,7 @@ class VisionAttention(_VisionAttentionBase): hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, - rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None, seqlens: list[int] | None = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] @@ -212,7 +227,7 @@ class VisionAttentionV2(_VisionAttentionBase): hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, - rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None, seqlens: list[int] | None = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] @@ -249,7 +264,7 @@ class VisionFlashAttention2(_VisionAttentionBase): hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, - rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None, seqlens: list[int] | None = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] @@ -274,7 +289,7 @@ class VisionFlashAttention3(_VisionAttentionBase): hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, - rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None, seqlens: list[int] | None = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] @@ -313,7 +328,7 @@ class VisionSdpaAttention(_VisionAttentionBase): hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, - rotary_pos_emb: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None, seqlens: list[int] | None = None, ) -> torch.Tensor: seq_length = hidden_states.shape[0] @@ -439,7 +454,7 @@ def apply_vision_attention_residual( hidden_states: torch.Tensor, cu_seqlens: torch.Tensor, max_seqlen: int, - rotary_pos_emb: torch.Tensor, + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding, *, seqlens: list[int] | None = None, uses_seqlens: bool = False, @@ -465,11 +480,13 @@ __all__ = [ "VisionFlashAttention2", "VisionFlashAttention3", "VisionRotaryEmbedding", + "VisionRotaryPositionEmbedding", "VisionSdpaAttention", "apply_rotary_pos_emb_vision", "apply_vision_attention_residual", "attn_uses_seqlens", "build_vision_attention", + "prepare_rotary_pos_emb_vision", "prepare_seqlens_for_attention", "resolve_attn_implementation", "rotate_half",