iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Model] Fix KV cache layout and optimize Dots3 NOTE Omni encoders (#53460)

Browse Source 
Signed-off-by: KurodaKanbei <mistergalahad@gmail.com>
This commit is contained in:
范裕达
2026-08-23 10:02:35 -07:00
committed by GitHub
parent 6cd9713463
commit 185cada36b
5 changed files with 210 additions and 89 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/dots3_note/nvidia/attention.py
+8
View File
@@ -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"
vllm/models/dots3_note/nvidia/audio.py
+96 -14
View File
@@ -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)
vllm/models/dots3_note/nvidia/multimodal.py
+22 -2
View File
@@ -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(
vllm/models/dots3_note/nvidia/vision.py
+55 -61
View File
@@ -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
vllm/models/dots3_note/nvidia/vision_attention.py
+29 -12
View File
@@ -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",