@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 185cada36b

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

Signed-off-by: KurodaKanbei <mistergalahad@gmail.com>

范裕达 committed Aug 24, 2026 at 01:02 UTC 185cada36bb25aa55f762d004d54c5ca1e3fc753

5 files changed +211 -90

vllm/models/dots3_note/nvidia/attention.py

@@ -37,6 +37,7 @@ from vllm.v1.attention.backends.mla.triton_mla import (

37 TritonMLAImpl,

38 TritonMLAMetadataBuilder,

39 )

40 +from vllm.v1.kv_cache_interface import KVCacheLayout

41 from vllm.v1.worker.workspace import (

42 current_workspace_manager,

43 is_workspace_manager_initialized,

@@ -677,6 +678,13 @@ class Dots3NoteTritonMLAImpl(TritonMLAImpl):

678 class Dots3NotePaddedSparseBackend(FlashAttnMLASparseBackend):

679 """NOTE DSA backend for cache rows padded to the SWA latent width."""

680

681 + @classmethod

682 + def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]:

683 + # NOTE packs the smaller DSA index pages beside the padded MLA/SWA pages.

684 + # Keep the layer dimension inside the block dimension so mixed page sizes

685 + # remain representable by the shared KV-cache allocation.

686 + return (KVCacheLayout.BLHNC,)

687 +

688 @staticmethod

689 def get_name() -> str:

690 return "DOTS3_NOTE_PADDED_MLA_SPARSE"

vllm/models/dots3_note/nvidia/audio.py

+96 -14

@@ -21,6 +21,7 @@ DEFAULT_MERGE_FACTOR = 1

21 N_SAMPLES = DEFAULT_CHUNK_LENGTH_S * SAMPLE_RATE

22 STRIDE = HOP_LENGTH * DEFAULT_CONV_TEMPORAL_STRIDE * DEFAULT_MERGE_FACTOR

23 ENCODER_SEQ_LEN = N_SAMPLES // HOP_LENGTH // 2

24 +_AUDIO_FORWARD_MAX_SEGMENTS = 16

27 class Dots3NoteAudioConfig:

@@ -126,6 +127,17 @@ def log_mel_spectrogram(audio, n_mels=128):

127 return log_spec

128

129

130 +def batched_log_mel_spectrogram(audio, n_mels=128):

131 + window = _hann_window(audio.device)

132 + stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)

133 + magnitudes = stft[..., :-1].abs() ** 2

134 + filters = _mel_filters(audio.device, n_mels)

135 + mel_spec = filters @ magnitudes

136 + log_spec = torch.clamp(mel_spec, min=1e-10).log10()

137 + log_spec = torch.maximum(log_spec, log_spec.amax(dim=(-2, -1), keepdim=True) - 8.0)

138 + return (log_spec + 4.0) / 4.0

139 +

140 +

141 def compute_audio_token_length(

142 num_samples,

143 *,

@@ -181,14 +193,34 @@ class DotsEncoderWithMask(nn.Module):

193 input_seq_lens: torch.Tensor,

194 audio_sample_lens: list[int],

195 ) -> torch.Tensor:

184 - """Run the eager speech encoder without server-side slicing/batching."""

196 mel_features = mel_features.to(dtype=torch.bfloat16, device=self.device)

186 - return self.speech_encoder(

187 - mel_features,

188 - return_dict=True,

189 - input_seq_lens=input_seq_lens,

190 - audio_sample_lens=audio_sample_lens,

191 - ).last_hidden_state

197 + if mel_features.shape[0] <= _AUDIO_FORWARD_MAX_SEGMENTS:

198 + return self.speech_encoder(

199 + mel_features,

200 + return_dict=True,

201 + input_seq_lens=input_seq_lens,

202 + audio_sample_lens=audio_sample_lens,

203 + ).last_hidden_state

204 +

205 + chunks = []

206 + for start in range(0, mel_features.shape[0], _AUDIO_FORWARD_MAX_SEGMENTS):

207 + end = start + _AUDIO_FORWARD_MAX_SEGMENTS

208 + chunks.append(

209 + self.speech_encoder(

210 + mel_features[start:end],

211 + return_dict=True,

212 + input_seq_lens=input_seq_lens[start:end],

213 + audio_sample_lens=audio_sample_lens[start:end],

214 + ).last_hidden_state

215 + )

216 + max_seq_len = max(chunk.shape[1] for chunk in chunks)

217 + chunks = [

218 + F.pad(chunk, (0, 0, 0, max_seq_len - chunk.shape[1]))

219 + if chunk.shape[1] < max_seq_len

220 + else chunk

221 + for chunk in chunks

222 + ]

223 + return torch.cat(chunks, dim=0)

224

225 def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor:

226 segments = []

@@ -233,6 +265,53 @@ class DotsEncoderWithMask(nn.Module):

265 )

266 return torch.cat(chunk_embeddings, dim=0).unsqueeze(0)

267

268 + def encode_waveforms(self, audio_waveforms) -> list[torch.Tensor]:

269 + stride = HOP_LENGTH * self.conv_temporal_stride * self.merge_factor

270 + segment_token_lengths = []

271 + segment_sample_lengths = []

272 + audio_segment_counts = []

273 + padded_segments = []

274 +

275 + for audio_waveform in audio_waveforms:

276 + segment_count = 0

277 + time_step = 0

278 + while time_step * SAMPLE_RATE < audio_waveform.shape[0]:

279 + segment = audio_waveform[

280 + time_step * SAMPLE_RATE : (time_step + self.chunk_seconds)

281 + * SAMPLE_RATE

282 + ]

283 + segment_length = segment.shape[0]

284 + padded_segments.append(

285 + pad_or_trim(segment.flatten(), length=self.chunk_samples)

286 + )

287 + segment_token_lengths.append((segment_length - 1) // stride + 1)

288 + segment_sample_lengths.append(segment_length)

289 + segment_count += 1

290 + time_step += self.chunk_seconds

291 + audio_segment_counts.append(segment_count)

292 +

293 + mel_features = batched_log_mel_spectrogram(torch.stack(padded_segments))

294 + assert mel_features.shape[-1] == self.chunk_mel_frames

295 + input_seq_lens = (

296 + torch.tensor(segment_token_lengths, dtype=torch.long) * self.merge_factor

297 + )

298 + audio_embedding = self._forward_speech_encoder(

299 + mel_features, input_seq_lens, segment_sample_lengths

300 + )

301 +

302 + embeddings = []

303 + segment_idx = 0

304 + for segment_count in audio_segment_counts:

305 + chunks = []

306 + for _ in range(segment_count):

307 + token_length = segment_token_lengths[segment_idx]

308 + chunks.append(

309 + audio_embedding[segment_idx, : token_length * self.merge_factor, :]

310 + )

311 + segment_idx += 1

312 + embeddings.append(torch.cat(chunks).unsqueeze(0))

313 + return embeddings

314 +

315

316 class AudioAdapter(nn.Module):

317 def __init__(self, in_dim, out_dim):

@@ -294,12 +373,15 @@ class Dots3NoteAudioModel(nn.Module):

373 audio_start += length

374 return waveforms

375

376 + def _forward_batched(self, waveforms):

377 + raw_embeddings = self.dots_encoder.encode_waveforms(waveforms)

378 + merged = [

379 + self._merge_embeddings(embedding).squeeze(0) for embedding in raw_embeddings

380 + ]

381 + token_lengths = [embedding.shape[0] for embedding in merged]

382 + packed = self.audio_adapter(torch.cat(merged))

383 + return packed, token_lengths

384 +

385 def forward(self, audio_inputs, lengths):

386 waveforms = self._split_waveforms(audio_inputs, lengths)

299 - all_embeddings = []

300 - token_lengths = []

301 - for waveform in waveforms:

302 - embedding = self._encode_single_audio(waveform)

303 - all_embeddings.append(embedding)

304 - token_lengths.append(embedding.shape[0])

305 - return torch.cat(all_embeddings, dim=0), token_lengths

387 + return self._forward_batched(waveforms)

vllm/models/dots3_note/nvidia/multimodal.py

+22 -2

@@ -3,6 +3,7 @@

3 """vLLM composition layer for Dots3Note image and audio encoders."""

5 from collections.abc import Iterable

6 +from contextlib import contextmanager

8 import torch

9 from torch import nn

@@ -41,6 +42,21 @@ from .model import Dots3NoteLanguageModelForCausalLM

42 from .vision import DotsMoEVitConfig, DotsMoEVitModel

45 +@contextmanager

46 +def _skip_linear_init():

47 + """Skip initialization for VE Linear weights loaded immediately afterward."""

48 + original_reset_parameters = nn.Linear.reset_parameters

49 +

50 + def _noop_reset_parameters(_linear: nn.Linear) -> None:

51 + return None

52 +

53 + nn.Linear.reset_parameters = _noop_reset_parameters

54 + try:

55 + yield

56 + finally:

57 + nn.Linear.reset_parameters = original_reset_parameters

58 +

59 +

60 @MULTIMODAL_REGISTRY.register_processor(

61 Dots3NoteMultiModalProcessor,

62 info=Dots3NoteProcessingInfo,

@@ -113,7 +129,10 @@ class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

129 self.visual: DotsMoEVitModel | None

130 if image_enabled:

131 assert vision_config_dict is not None

116 - self.visual = DotsMoEVitModel(DotsMoEVitConfig(**vision_config_dict))

132 + with _skip_linear_init():

133 + self.visual = DotsMoEVitModel(

134 + DotsMoEVitConfig(**vision_config_dict)

135 + )

136 else:

137 self.visual = None

138 self.audio_tower: Dots3NoteAudioModel | None

@@ -158,9 +177,10 @@ class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):

177 ) -> tuple[torch.Tensor, ...]:

178 if self.visual is None:

179 return ()

161 - image_embeds = self.visual(pixel_values, image_grid_thw)

180 merge_size = self.visual.spatial_merge_size

181 sizes = (image_grid_thw.prod(-1) // merge_size**2).tolist()

182 + meta = self.visual.prepare_meta(image_grid_thw)

183 + image_embeds = self.visual(pixel_values, meta)

184 return image_embeds.split(sizes)

185

186 def _process_audio_input(

vllm/models/dots3_note/nvidia/vision.py

+56 -62

@@ -2,6 +2,7 @@

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

3 # ruff: noqa: E501

4 import math

5 +from dataclasses import dataclass

6 from typing import Any

8 import torch

@@ -15,10 +16,11 @@ from vllm.third_party.deep_gemm import per_block_cast_to_fp8

17 from .vision_attention import (

18 VisionRotaryEmbedding,

19 + VisionRotaryPositionEmbedding,

20 apply_vision_attention_residual,

21 attn_uses_seqlens,

22 build_vision_attention,

21 - prepare_seqlens_for_attention,

23 + prepare_rotary_pos_emb_vision,

24 resolve_attn_implementation,

25 )

26 from .vision_moe import note_vision_fused_moe_fp8

@@ -502,6 +504,15 @@ _ADAPTER_CLASSES: dict[str, type[nn.Module]] = {

504 }

505

506

507 +@dataclass(frozen=True)

508 +class VitForwardMeta:

509 + cu_seqlens: torch.Tensor

510 + max_seqlen: int

511 + rotary_pos_emb: VisionRotaryPositionEmbedding

512 + grid_thw_cpu: torch.Tensor

513 + seqlens: list[int] | None = None

514 +

515 +

516 # ---- Full Model ----

517

518

@@ -562,7 +573,7 @@ class DotsMoEVitModel(PreTrainedModel):

573 block.norm_1 = torch.compile(block.norm_1, **compile_kwargs)

574 block.norm_2 = torch.compile(block.norm_2, **compile_kwargs)

575

565 - def get_pos_ids_by_grid(self, grid_thw):

576 + def get_pos_ids_by_grid(self, grid_thw_cpu: torch.Tensor):

577 # Mirrors ``cybertron`` ``AIMv2NativeModel.rot_pos_emb``: when ``pre_pixel_shuffle``

578 # is set, RoPE positions follow the qwen ``merge_size`` grouped layout (default 2x2);

579 # otherwise positions are flat row-major regardless of ``spatial_merge_size``.

@@ -573,10 +584,8 @@ class DotsMoEVitModel(PreTrainedModel):

584 else:

585 rope_merge_size = 1

586 pos_ids = []

576 - for t, h, w in grid_thw:

577 - hpos_ids = (

578 - torch.arange(h, device=grid_thw.device).unsqueeze(1).expand(-1, w)

579 - )

587 + for t, h, w in grid_thw_cpu:

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

589 hpos_ids = hpos_ids.reshape(

590 h // rope_merge_size,

591 rope_merge_size,

@@ -585,9 +594,7 @@ class DotsMoEVitModel(PreTrainedModel):

594 )

595 hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten()

596

588 - wpos_ids = (

589 - torch.arange(w, device=grid_thw.device).unsqueeze(0).expand(h, -1)

590 - )

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

598 wpos_ids = wpos_ids.reshape(

599 h // rope_merge_size,

600 rope_merge_size,

@@ -598,80 +605,67 @@ class DotsMoEVitModel(PreTrainedModel):

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

606 return pos_ids

607

601 - def rot_pos_emb(self, grid_thw):

602 - pos_ids = self.get_pos_ids_by_grid(grid_thw)

603 - pos_ids = torch.cat(pos_ids, dim=0)

604 - max_grid_size = grid_thw[:, 1:].max()

608 + def rot_pos_emb(self, grid_thw_cpu: torch.Tensor):

609 + pos_ids = torch.cat(self.get_pos_ids_by_grid(grid_thw_cpu), dim=0)

610 + pos_ids = pos_ids.to(self.device, non_blocking=True)

611 + max_grid_size = int(grid_thw_cpu[:, 1:].max())

612 rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

606 - rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

607 - return rotary_pos_emb

608 -

609 - def _build_cu_seqlens_from_grid(self, grid_thw: torch.Tensor):

610 - cu_seqlens = torch.repeat_interleave(

611 - grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

612 - ).cumsum(

613 - dim=0,

614 - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,

613 + return rotary_pos_emb_full[pos_ids].flatten(1)

614 +

615 + def prepare_meta(self, grid_thw_cpu: torch.Tensor) -> VitForwardMeta:

616 + if grid_thw_cpu.device.type != "cpu":

617 + raise ValueError("NOTE vision grid_thw must remain on CPU")

618 +

619 + seq_per_frame = grid_thw_cpu[:, 1] * grid_thw_cpu[:, 2]

620 + frame_lengths = torch.repeat_interleave(seq_per_frame, grid_thw_cpu[:, 0])

621 + cu_seqlens_cpu = F.pad(

622 + frame_lengths.cumsum(dim=0, dtype=torch.int32), (1, 0), value=0

623 )

616 - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

617 - # Same as (cu_seqlens[1:] - cu_seqlens[:-1]).max(); computed here to avoid D2H inside attention.

618 - max_seqlen = int((grid_thw[:, 1] * grid_thw[:, 2]).max().item())

619 - return cu_seqlens, max_seqlen

620 -

621 - def _build_single_temporal_cu_seqlens_from_grid(self, grid_thw: torch.Tensor):

622 - seq_lens = grid_thw[:, 1] * grid_thw[:, 2]

623 - cu_seqlens = torch.empty(

624 - (grid_thw.shape[0] + 1,), device=grid_thw.device, dtype=torch.int32

624 + resolved_attn = resolve_attn_implementation(

625 + self.config.attn_implementation, eager_fallback="eager_v2"

626 + )

627 + return VitForwardMeta(

628 + cu_seqlens=cu_seqlens_cpu.to(self.device, non_blocking=True),

629 + max_seqlen=int(seq_per_frame.max()),

630 + rotary_pos_emb=prepare_rotary_pos_emb_vision(

631 + self.rot_pos_emb(grid_thw_cpu)

632 + ),

633 + grid_thw_cpu=grid_thw_cpu,

634 + seqlens=(

635 + frame_lengths.tolist() if attn_uses_seqlens(resolved_attn) else None

636 + ),

637 )

626 - cu_seqlens[0] = 0

627 - torch.cumsum(seq_lens, dim=0, dtype=torch.int32, out=cu_seqlens[1:])

628 - max_seqlen = int(seq_lens.max().item())

629 - return cu_seqlens, max_seqlen

638

639 def forward(

632 - self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True

640 + self, hidden_states: torch.Tensor, meta: VitForwardMeta, bf16=True

641 ) -> torch.Tensor:

634 - grid_thw = grid_thw.to(device=self.device)

642 if bf16:

636 - hidden_states = hidden_states.to(device=self.device, dtype=self.dtype)

637 - hidden_states = self.patch_embed(hidden_states)

638 -

639 - rotary_pos_emb = self.rot_pos_emb(grid_thw)

640 -

641 - if grid_thw[:, 0].sum().item() == grid_thw.shape[0]:

642 - cu_seqlens, max_seqlen = self._build_single_temporal_cu_seqlens_from_grid(

643 - grid_thw

643 + hidden_states = hidden_states.to(

644 + device=self.device, dtype=self.dtype, non_blocking=True

645 )

645 - else:

646 - cu_seqlens, max_seqlen = self._build_cu_seqlens_from_grid(grid_thw)

647 -

648 - seqlens = prepare_seqlens_for_attention(

649 - self.config.attn_implementation,

650 - cu_seqlens,

651 - eager_fallback="eager_v2",

652 - )

646 + hidden_states = self.patch_embed(hidden_states)

647

648 for blk in self.blocks:

649 if self.gradient_checkpointing and self.training:

650 hidden_states = self._gradient_checkpointing_func(

651 blk.__call__,

652 hidden_states,

659 - cu_seqlens,

660 - rotary_pos_emb,

661 - max_seqlen,

662 - seqlens,

653 + meta.cu_seqlens,

654 + meta.rotary_pos_emb,

655 + meta.max_seqlen,

656 + meta.seqlens,

657 )

658 else:

659 hidden_states = blk(

660 hidden_states,

667 - cu_seqlens,

668 - rotary_pos_emb,

669 - max_seqlen,

670 - seqlens=seqlens,

661 + meta.cu_seqlens,

662 + meta.rotary_pos_emb,

663 + meta.max_seqlen,

664 + seqlens=meta.seqlens,

665 )

666

667 if self.config.post_norm:

668 hidden_states = self.post_trunk_norm(hidden_states)

669

676 - hidden_states = self.adapter(hidden_states, grid_thw)

670 + hidden_states = self.adapter(hidden_states, meta.grid_thw_cpu)

671 return hidden_states

vllm/models/dots3_note/nvidia/vision_attention.py

+29 -12

@@ -36,15 +36,28 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor:

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

39 +VisionRotaryPositionEmbedding = tuple[torch.Tensor, torch.Tensor]

40 +

41 +

42 +def prepare_rotary_pos_emb_vision(

43 + freqs: torch.Tensor,

44 +) -> VisionRotaryPositionEmbedding:

45 + """Materialize vision RoPE cos/sin once for every encoder layer."""

46 + cos = freqs.cos().unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()

47 + sin = freqs.sin().unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()

48 + return cos, sin

49 +

50 +

51 def apply_rotary_pos_emb_vision(

40 - tensor: torch.Tensor, freqs: torch.Tensor

52 + tensor: torch.Tensor,

53 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding,

54 ) -> torch.Tensor:

55 orig_dtype = tensor.dtype

56 tensor = tensor.float()

44 - cos = freqs.cos()

45 - sin = freqs.sin()

46 - cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()

47 - sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()

57 + if isinstance(rotary_pos_emb, torch.Tensor):

58 + cos, sin = prepare_rotary_pos_emb_vision(rotary_pos_emb)

59 + else:

60 + cos, sin = rotary_pos_emb

61 output = (tensor * cos) + (rotate_half(tensor) * sin)

62 return output.to(orig_dtype)

@@ -149,7 +162,7 @@ class _VisionAttentionBase(nn.Module):

162 def _qkv_with_rope(

163 self,

164 hidden_states: torch.Tensor,

152 - rotary_pos_emb: torch.Tensor | None,

165 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None,

166 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

167 seq_length = hidden_states.shape[0]

168 q, k, v = (

@@ -161,6 +174,8 @@ class _VisionAttentionBase(nn.Module):

174 if self.use_qk_norm:

175 q = self.q_norm(q)

176 k = self.k_norm(k)

177 + if rotary_pos_emb is None:

178 + raise ValueError("Vision rotary position embedding is required")

179 q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)

180 k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)

181 return q, k, v

@@ -174,7 +189,7 @@ class VisionAttention(_VisionAttentionBase):

189 hidden_states: torch.Tensor,

190 cu_seqlens: torch.Tensor,

191 max_seqlen: int,

177 - rotary_pos_emb: torch.Tensor | None = None,

192 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None,

193 seqlens: list[int] | None = None,

194 ) -> torch.Tensor:

195 seq_length = hidden_states.shape[0]

@@ -212,7 +227,7 @@ class VisionAttentionV2(_VisionAttentionBase):

227 hidden_states: torch.Tensor,

228 cu_seqlens: torch.Tensor,

229 max_seqlen: int,

215 - rotary_pos_emb: torch.Tensor | None = None,

230 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None,

231 seqlens: list[int] | None = None,

232 ) -> torch.Tensor:

233 seq_length = hidden_states.shape[0]

@@ -249,7 +264,7 @@ class VisionFlashAttention2(_VisionAttentionBase):

264 hidden_states: torch.Tensor,

265 cu_seqlens: torch.Tensor,

266 max_seqlen: int,

252 - rotary_pos_emb: torch.Tensor | None = None,

267 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None,

268 seqlens: list[int] | None = None,

269 ) -> torch.Tensor:

270 seq_length = hidden_states.shape[0]

@@ -274,7 +289,7 @@ class VisionFlashAttention3(_VisionAttentionBase):

289 hidden_states: torch.Tensor,

290 cu_seqlens: torch.Tensor,

291 max_seqlen: int,

277 - rotary_pos_emb: torch.Tensor | None = None,

292 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None,

293 seqlens: list[int] | None = None,

294 ) -> torch.Tensor:

295 seq_length = hidden_states.shape[0]

@@ -313,7 +328,7 @@ class VisionSdpaAttention(_VisionAttentionBase):

328 hidden_states: torch.Tensor,

329 cu_seqlens: torch.Tensor,

330 max_seqlen: int,

316 - rotary_pos_emb: torch.Tensor | None = None,

331 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding | None = None,

332 seqlens: list[int] | None = None,

333 ) -> torch.Tensor:

334 seq_length = hidden_states.shape[0]

@@ -439,7 +454,7 @@ def apply_vision_attention_residual(

454 hidden_states: torch.Tensor,

455 cu_seqlens: torch.Tensor,

456 max_seqlen: int,

442 - rotary_pos_emb: torch.Tensor,

457 + rotary_pos_emb: torch.Tensor | VisionRotaryPositionEmbedding,

458 *,

459 seqlens: list[int] | None = None,

460 uses_seqlens: bool = False,

@@ -465,11 +480,13 @@ __all__ = [

480 "VisionFlashAttention2",

481 "VisionFlashAttention3",

482 "VisionRotaryEmbedding",

483 + "VisionRotaryPositionEmbedding",

484 "VisionSdpaAttention",

485 "apply_rotary_pos_emb_vision",

486 "apply_vision_attention_residual",

487 "attn_uses_seqlens",

488 "build_vision_attention",

489 + "prepare_rotary_pos_emb_vision",

490 "prepare_seqlens_for_attention",

491 "resolve_attn_implementation",

492 "rotate_half",