diff --git a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
index 81c622db00..5755580142 100644
--- a/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
+++ b/tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
@@ -127,7 +127,14 @@ def _run_verify_block():
# stub with the attribute is enough for metadata construction.
layer_name = "placeholder"
vllm_config.compilation_config.static_forward_context[layer_name] = (
- types.SimpleNamespace(prefill_backend=torch.empty((1,)))
+ types.SimpleNamespace(
+ prefill_backend=torch.empty((1,)),
+ q_lora_rank=None,
+ kv_lora_rank=KV_LORA_RANK,
+ qk_nope_head_dim=QK_NOPE_HEAD_DIM,
+ qk_rope_head_dim=QK_ROPE_HEAD_DIM,
+ v_head_dim=V_HEAD_DIM,
+ )
)
init_workspace_manager(device)
diff --git a/tests/models/registry.py b/tests/models/registry.py
index e209d973b9..ed1786fdab 100644
--- a/tests/models/registry.py
+++ b/tests/models/registry.py
@@ -839,6 +839,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
"DeepseekOCR2ForCausalLM": _HfExamplesInfo(
"deepseek-ai/DeepSeek-OCR-2",
),
+ "Dots3NoteForCausalLM": _HfExamplesInfo(
+ "rednote-hilab/dots3.note",
+ trust_remote_code=True,
+ is_available_online=False,
+ ),
"UnlimitedOCRForCausalLM": _HfExamplesInfo(
"baidu/Unlimited-OCR",
),
@@ -1650,6 +1655,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
trust_remote_code=True,
is_available_online=False,
),
+ "Dots3NoteMTPModel": _HfExamplesInfo(
+ "rednote-hilab/dots3.note",
+ speculative_model="rednote-hilab/dots3.note",
+ trust_remote_code=True,
+ is_available_online=False,
+ ),
"Gemma4MTPModel": _HfExamplesInfo(
"google/gemma-4-E4B-it",
speculative_model="google/gemma-4-E4B-it-assistant",
diff --git a/tests/models/test_registry.py b/tests/models/test_registry.py
index 45db16e879..6eda6075d2 100644
--- a/tests/models/test_registry.py
+++ b/tests/models/test_registry.py
@@ -54,6 +54,11 @@ def test_registry_imports(model_arch):
):
pytest.skip("DSparkDraftModel is only supported on CUDA and ROCm")
+ if model_arch in ("Dots3NoteForCausalLM", "Dots3NoteMTPModel") and not (
+ current_platform.is_cuda()
+ ):
+ pytest.skip("Dots3 NOTE is only supported on CUDA")
+
# Ensure all model classes can be imported successfully
model_cls = ModelRegistry._try_load_model_cls(model_arch)
assert model_cls is not None
diff --git a/tests/tool_parsers/test_dots_tool_parser.py b/tests/tool_parsers/test_dots_tool_parser.py
new file mode 100644
index 0000000000..151bca5def
--- /dev/null
+++ b/tests/tool_parsers/test_dots_tool_parser.py
@@ -0,0 +1,275 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from collections.abc import Iterable
+from unittest.mock import MagicMock
+
+import pytest
+
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionRequest,
+ ChatCompletionToolsParam,
+)
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
+from vllm.tool_parsers import ToolParserManager
+from vllm.tool_parsers.dots_tool_parser import DotsToolParser
+
+
+def _tool(name: str, properties: dict) -> ChatCompletionToolsParam:
+ return ChatCompletionToolsParam(
+ type="function",
+ function={
+ "name": name,
+ "description": "test tool",
+ "parameters": {"type": "object", "properties": properties},
+ },
+ )
+
+
+def _request(tools: list[ChatCompletionToolsParam]) -> ChatCompletionRequest:
+ return ChatCompletionRequest(
+ model="test-model",
+ messages=[{"role": "user", "content": "hi"}],
+ tools=tools,
+ tool_choice="auto",
+ )
+
+
+def _stream(
+ parser: DotsToolParser,
+ chunks: Iterable[str],
+ request: ChatCompletionRequest,
+) -> list[DeltaMessage]:
+ previous_text = ""
+ messages: list[DeltaMessage] = []
+ for chunk in chunks:
+ current_text = previous_text + chunk
+ delta = parser.extract_tool_calls_streaming(
+ previous_text,
+ current_text,
+ chunk,
+ [],
+ [],
+ [],
+ request,
+ )
+ if delta is not None:
+ messages.append(delta)
+ previous_text = current_text
+ return messages
+
+
+@pytest.fixture
+def parser() -> DotsToolParser:
+ return DotsToolParser(MagicMock())
+
+
+def test_registered_in_tool_parser_manager() -> None:
+ assert ToolParserManager.get_tool_parser("dots") is DotsToolParser
+
+
+def test_non_stream_xml_converts_schema_types_and_resolves_ref(
+ parser: DotsToolParser,
+) -> None:
+ tool = ChatCompletionToolsParam(
+ type="function",
+ function={
+ "name": "set_location",
+ "description": "Set location",
+ "parameters": {
+ "type": "object",
+ "properties": {
+ "location": {"$ref": "#/$defs/Location"},
+ "days": {"type": "integer"},
+ "include_weather": {"type": "boolean"},
+ },
+ "$defs": {
+ "Location": {
+ "type": "object",
+ "properties": {"city": {"type": "string"}},
+ }
+ },
+ },
+ },
+ )
+ request = _request([tool])
+ text = (
+ "ok"
+ ''
+ '{"city": "Shanghai"}'
+ '3'
+ 'true'
+ ""
+ ""
+ )
+
+ result = parser.extract_tool_calls(text, request)
+
+ assert result.content == "ok"
+ assert result.tools_called
+ assert len(result.tool_calls) == 1
+ assert result.tool_calls[0].function.name == "set_location"
+ assert json.loads(result.tool_calls[0].function.arguments) == {
+ "location": {"city": "Shanghai"},
+ "days": 3,
+ "include_weather": True,
+ }
+
+
+def test_non_stream_supports_multiple_invokes_and_json_fallback(
+ parser: DotsToolParser,
+) -> None:
+ tools = [
+ _tool("search", {"query": {"type": "string"}}),
+ _tool("open", {"id": {"type": "integer"}}),
+ ]
+ request = _request(tools)
+ text = (
+ ""
+ 'chairs'
+ '7'
+ ""
+ '{"name":"search","arguments":{"query":"tables"}}'
+ ""
+ )
+
+ result = parser.extract_tool_calls(text, request)
+
+ assert [call.function.name for call in result.tool_calls] == [
+ "search",
+ "open",
+ "search",
+ ]
+ assert [json.loads(call.function.arguments) for call in result.tool_calls] == [
+ {"query": "chairs"},
+ {"id": 7},
+ {"query": "tables"},
+ ]
+
+
+def test_non_stream_unknown_tool_is_left_as_content(parser: DotsToolParser) -> None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+ text = (
+ ""
+ 'chairs'
+ ""
+ )
+
+ result = parser.extract_tool_calls(text, request)
+
+ assert not result.tools_called
+ assert result.tool_calls == []
+ assert result.content == text
+
+
+def test_streaming_buffers_partial_marker_and_emits_all_complete_calls(
+ parser: DotsToolParser,
+) -> None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+ chunks = [
+ "visible"
+ 'chairs'
+ ""
+ ""
+ 'tables'
+ ""
+ ),
+ ]
+
+ messages = _stream(parser, chunks, request)
+
+ assert "".join(message.content or "" for message in messages) == "visible"
+ calls = [call for message in messages for call in message.tool_calls]
+ assert [call.index for call in calls] == [0, 1]
+ assert [json.loads(call.function.arguments or "") for call in calls] == [
+ {"query": "chairs"},
+ {"query": "tables"},
+ ]
+
+
+def test_streaming_filters_unknown_tools_and_surfaces_content(
+ parser: DotsToolParser,
+) -> None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+ text = (
+ ""
+ 'chairs'
+ ""
+ )
+
+ messages = _stream(parser, [text], request)
+
+ assert len(messages) == 1
+ assert messages[0].tool_calls == []
+ assert "ghost" in (messages[0].content or "")
+ assert parser._buffer == ""
+
+
+def test_streaming_malformed_block_does_not_block_later_valid_call(
+ parser: DotsToolParser,
+) -> None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+ messages = _stream(
+ parser,
+ [
+ "garbage",
+ (
+ ""
+ ''
+ 'chairs'
+ ""
+ ),
+ ],
+ request,
+ )
+
+ assert messages[0].content == "garbage"
+ assert messages[0].tool_calls == []
+ assert messages[1].tool_calls[0].function.name == "search"
+
+
+def test_streaming_strips_stray_end_marker(parser: DotsToolParser) -> None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+ messages = _stream(parser, ["some text "], request)
+
+ assert len(messages) == 1
+ assert messages[0].content == "some text "
+ assert messages[0].tool_calls == []
+
+
+def test_streaming_flushes_partial_opening_marker_at_eof(
+ parser: DotsToolParser,
+) -> None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+
+ messages = _stream(parser, ["answer None:
+ request = _request([_tool("search", {"query": {"type": "string"}})])
+
+ messages = _stream(
+ parser,
+ [
+ "",
+ '{"name":"search","arguments":{"query":"chairs"}}',
+ "",
+ ],
+ request,
+ )
+
+ calls = [call for message in messages for call in message.tool_calls]
+ assert [call.function.name for call in calls] == ["search", None]
+ assert "".join(call.function.arguments or "" for call in calls) == (
+ '{"query": "chairs"}'
+ )
+ assert messages[-1].tool_calls == calls
diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py
index 9ab7d68b89..cd29193724 100644
--- a/tests/v1/attention/test_mla_backends.py
+++ b/tests/v1/attention/test_mla_backends.py
@@ -78,6 +78,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
kv_cache_dtype=cache_dtype,
head_size=576,
non_causal_multi_token_decode=False,
+ sliding_window=None,
)
vllm_config = SimpleNamespace(
cache_config=SimpleNamespace(block_size=64), model_config=None
@@ -460,6 +461,7 @@ class MockSparseMLAAttentionLayer:
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.v_head_dim = v_head_dim
+ self.q_lora_rank = None
self.kv_lora_rank = kv_lora_rank
# Compute weight matrices in the format expected by forward_impl
@@ -583,6 +585,7 @@ class MockMLAAttentionLayer(MLAAttention):
self.qk_nope_head_dim = qk_nope_head_dim
self.qk_rope_head_dim = qk_rope_head_dim
self.v_head_dim = v_head_dim
+ self.q_lora_rank = None
self.kv_lora_rank = kv_lora_rank
# Compute weight matrices from kv_b_proj (like MLAAttention does)
diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py
index fa8c282007..177f5a262b 100644
--- a/vllm/config/speculative.py
+++ b/vllm/config/speculative.py
@@ -36,6 +36,7 @@ logger = init_logger(__name__)
MTPModelTypes = Literal[
"deepseek_mtp",
+ "dots3_note_mtp",
"mimo_mtp",
"mimo_v2_mtp",
"glm4_moe_mtp",
@@ -336,6 +337,21 @@ class SpeculativeConfig:
@staticmethod
def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
initial_architecture = hf_config.architectures[0]
+ if hf_config.model_type == "dots3_note":
+ n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
+ mtp_layer_types = getattr(hf_config, "mtp_layer_types", None)
+ if mtp_layer_types is None:
+ mtp_layer_types = ["sliding_attention"] * n_predict
+ if len(mtp_layer_types) != n_predict:
+ raise ValueError(
+ "mtp_layer_types must have one entry per MTP layer: "
+ f"got {len(mtp_layer_types)} for {n_predict} layers"
+ )
+ hf_config.layer_types = [*hf_config.layer_types, *mtp_layer_types]
+ hf_config.model_type = "dots3_note_mtp"
+ hf_config.update(
+ {"n_predict": n_predict, "architectures": ["Dots3NoteMTPModel"]}
+ )
if hf_config.model_type in (
"deepseek_v3",
"deepseek_v32",
diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py
index 0b59956543..33bcbf8aaf 100644
--- a/vllm/model_executor/layers/attention/mla_attention.py
+++ b/vllm/model_executor/layers/attention/mla_attention.py
@@ -299,6 +299,7 @@ from vllm.v1.kv_cache_interface import (
AttentionSpec,
KVCacheSpec,
MLAAttentionSpec,
+ SlidingWindowMLASpec,
get_kv_quant_mode,
)
@@ -406,6 +407,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
indexer: object | None = None,
topk_indices_buffer: torch.Tensor | None = None,
non_causal_multi_token_decode: bool = False,
+ sliding_window: int | None = None,
+ prefill_backend_cls: type[MLAPrefillBackend] | None = None,
**extra_impl_args,
):
super().__init__()
@@ -423,6 +426,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
self.layer_name = prefix
self.indexer = indexer
self.non_causal_multi_token_decode = non_causal_multi_token_decode
+ self.sliding_window = sliding_window
self.num_kv_heads = 1
self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
@@ -515,7 +519,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
scale=self.scale,
num_kv_heads=1,
alibi_slopes=None,
- sliding_window=None,
+ sliding_window=sliding_window,
kv_cache_dtype=self.kv_cache_dtype,
logits_soft_cap=None,
attn_type=AttentionType.DECODER,
@@ -551,7 +555,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
self.prefill_backend = None
else:
try:
- prefill_backend_cls = get_mla_prefill_backend(vllm_config)
+ prefill_backend_cls = prefill_backend_cls or get_mla_prefill_backend(
+ vllm_config
+ )
except ValueError:
if (
not self.impl.is_sparse
@@ -1139,13 +1145,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):
kv_cache_dtype = kv_cache_dtype_str_to_dtype(
self.kv_cache_dtype, vllm_config.model_config
)
- return MLAAttentionSpec(
+ common_kwargs = dict(
block_size=vllm_config.cache_config.block_size,
num_kv_heads=1,
head_size=self.head_size,
dtype=kv_cache_dtype,
cache_dtype_str=self.kv_cache_dtype,
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
+ )
+ if self.sliding_window is not None:
+ return SlidingWindowMLASpec(
+ **common_kwargs,
+ sliding_window=self.sliding_window,
+ )
+ return MLAAttentionSpec(
+ **common_kwargs,
non_causal_multi_token_decode=self.non_causal_multi_token_decode,
)
@@ -2056,7 +2070,15 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
self.num_heads = get_num_attention_heads_from_layers(
vllm_config, layer_names
) or self.model_config.get_num_attention_heads(parallel_config)
- self.mla_dims = get_mla_dims(self.model_config)
+ # Hybrid MLA models may use different latent dimensions per KV group.
+ layer = self.compilation_config.static_forward_context[layer_names[0]]
+ self.mla_dims = MLADims(
+ q_lora_rank=layer.q_lora_rank,
+ kv_lora_rank=layer.kv_lora_rank,
+ qk_nope_head_dim=layer.qk_nope_head_dim,
+ qk_rope_head_dim=layer.qk_rope_head_dim,
+ v_head_dim=layer.v_head_dim,
+ )
self.aot_schedule = current_platform.is_cuda()
self.kv_cache_spec = kv_cache_spec
@@ -2092,7 +2114,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
(
self.chunked_prefill_workspace_size
+ self.chunked_prefill_workspace_size // self.dcp_world_size,
- self.model_config.get_head_size(),
+ self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
),
dtype=torch.bfloat16
if use_packed_fp8_cache
@@ -2109,7 +2131,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
self.chunked_prefill_workspace = torch.empty(
(
self.chunked_prefill_workspace_size,
- self.model_config.get_head_size(),
+ self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
),
dtype=torch.bfloat16 if use_packed_fp8_cache else self.q_data_type,
device=device,
diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py
index 7237e924c5..87545b0925 100644
--- a/vllm/model_executor/models/registry.py
+++ b/vllm/model_executor/models/registry.py
@@ -371,6 +371,10 @@ _MULTIMODAL_MODELS = {
"DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
"DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
"DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
+ "Dots3NoteForCausalLM": (
+ "vllm.models.dots3_note",
+ "Dots3NoteForCausalLM",
+ ),
"UnlimitedOCRForCausalLM": ("unlimited_ocr", "UnlimitedOCRForCausalLM"),
"DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
"Eagle2_5_VLForConditionalGeneration": (
@@ -645,6 +649,7 @@ _SPECULATIVE_DECODING_MODELS = {
"Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
"EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
+ "Dots3NoteMTPModel": ("vllm.models.dots3_note", "Dots3NoteMTP"),
"DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
"BailingMoeV3MTPModel": ("bailing_moe_v3_mtp", "BailingMoeV3MTPModel"),
"MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
diff --git a/vllm/models/dots3_note/__init__.py b/vllm/models/dots3_note/__init__.py
new file mode 100644
index 0000000000..3c466e1aea
--- /dev/null
+++ b/vllm/models/dots3_note/__init__.py
@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dots3Note model entry point."""
+
+from .nvidia.mtp import Dots3NoteMTP
+from .nvidia.multimodal import Dots3NoteForCausalLM
+
+__all__ = ["Dots3NoteForCausalLM", "Dots3NoteMTP"]
diff --git a/vllm/models/dots3_note/common/__init__.py b/vllm/models/dots3_note/common/__init__.py
new file mode 100644
index 0000000000..ede52dd76d
--- /dev/null
+++ b/vllm/models/dots3_note/common/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Platform-independent Dots3Note preprocessing."""
diff --git a/vllm/models/dots3_note/common/processor.py b/vllm/models/dots3_note/common/processor.py
new file mode 100644
index 0000000000..81748d253b
--- /dev/null
+++ b/vllm/models/dots3_note/common/processor.py
@@ -0,0 +1,811 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Multimodal preprocessing for Dots3Note checkpoints."""
+
+import math
+from collections.abc import Mapping, Sequence
+from functools import cached_property
+from typing import Any, cast
+
+import numpy as np
+import torch
+from PIL import Image
+from transformers import BatchFeature
+
+from vllm.config.multimodal import (
+ AudioDummyOptions,
+ BaseDummyOptions,
+ ImageDummyOptions,
+ VideoDummyOptions,
+)
+from vllm.inputs import MultiModalDataDict
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
+from vllm.multimodal.media import MediaWithBytes
+from vllm.multimodal.parse import (
+ MultiModalDataItems,
+ MultiModalDataParser,
+ VideoProcessorItems,
+)
+from vllm.multimodal.processing import (
+ BaseDummyInputsBuilder,
+ BaseMultiModalProcessor,
+ BaseProcessingInfo,
+ PromptReplacement,
+ PromptUpdate,
+ PromptUpdateDetails,
+)
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
+
+from .video import preprocess_dots3_note_video
+
+IMAGE_START = "<|img|>"
+IMAGE_PAD = "<|imgpad|>"
+IMAGE_END = "<|endofimg|>"
+AUDIO_START = "<|audio_comp_start|>"
+AUDIO_PAD = "<|audio_comp_pad|>"
+AUDIO_END = "<|audio_comp_end|>"
+VIDEO_PLACEHOLDER = "<|video_pad|>"
+
+_HOP_LENGTH = 160
+_DEFAULT_SAMPLE_RATE = 16000
+
+
+def load_note_config_section(
+ model: str,
+ revision: str | None,
+ section: str,
+) -> dict[str, Any] | None:
+ config = get_hf_file_to_dict("config.json", model, revision)
+ value = (config or {}).get(section)
+ return value if isinstance(value, dict) else None
+
+
+class Dots3NoteImageProcessor:
+ """CPU image preprocessing matching the native Cybertron ViT."""
+
+ def __init__(
+ self,
+ config: Mapping[str, Any],
+ image_details: Mapping[str, Any] | None = None,
+ ) -> None:
+ self.min_pixels = int(config["min_pixels"])
+ self.max_pixels = int(config["max_pixels"])
+ self.patch_size = int(config["patch_size"])
+ self.temporal_patch_size = int(config["temporal_patch_size"])
+ self.merge_size = int(config["merge_size"])
+ self.pre_pixel_shuffle = bool(config.get("pre_pixel_shuffle", True))
+ self.image_mean = np.asarray(config["image_mean"], dtype=np.float32)
+ self.image_std = np.asarray(config["image_std"], dtype=np.float32)
+ self.image_details = dict(image_details or {})
+
+ @property
+ def factor(self) -> int:
+ return self.patch_size * self.merge_size
+
+ @staticmethod
+ def _round_by_factor(value: int, factor: int) -> int:
+ return round(value / factor) * factor
+
+ @staticmethod
+ def _ceil_by_factor(value: float, factor: int) -> int:
+ return math.ceil(value / factor) * factor
+
+ @staticmethod
+ def _floor_by_factor(value: float, factor: int) -> int:
+ return math.floor(value / factor) * factor
+
+ def resized_size(
+ self,
+ width: int,
+ height: int,
+ *,
+ detail: str = "auto",
+ min_pixels: int | None = None,
+ max_pixels: int | None = None,
+ target_height: int | None = None,
+ target_width: int | None = None,
+ ) -> tuple[int, int]:
+ detail_config = self.image_details.get(detail, {})
+ min_pixels = int(
+ min_pixels
+ if min_pixels is not None
+ else detail_config.get("min_pixels", self.min_pixels)
+ )
+ max_pixels = int(
+ max_pixels
+ if max_pixels is not None
+ else detail_config.get("max_pixels", self.max_pixels)
+ )
+ height = int(target_height or detail_config.get("target_height") or height)
+ width = int(target_width or detail_config.get("target_width") or width)
+
+ factor = self.factor
+ if min(height, width) < factor // 4:
+ raise ValueError(
+ f"Image height and width must be at least {factor // 4}, "
+ f"got {height}x{width}"
+ )
+ if max(height, width) / min(height, width) > 200:
+ raise ValueError("Image aspect ratio must be smaller than 200")
+
+ resized_h = max(factor, self._round_by_factor(height, factor))
+ resized_w = max(factor, self._round_by_factor(width, factor))
+ if resized_h * resized_w > max_pixels:
+ beta = math.sqrt(height * width / max_pixels)
+ resized_h = max(factor, self._floor_by_factor(height / beta, factor))
+ resized_w = max(factor, self._floor_by_factor(width / beta, factor))
+ elif resized_h * resized_w < min_pixels:
+ beta = math.sqrt(min_pixels / (height * width))
+ resized_h = self._ceil_by_factor(height * beta, factor)
+ resized_w = self._ceil_by_factor(width * beta, factor)
+ if resized_h * resized_w > max_pixels:
+ beta = math.sqrt(resized_h * resized_w / max_pixels)
+ resized_h = max(factor, self._floor_by_factor(resized_h / beta, factor))
+ resized_w = max(factor, self._floor_by_factor(resized_w / beta, factor))
+ return resized_h, resized_w
+
+ def preprocess(
+ self,
+ image: Image.Image,
+ *,
+ detail: str = "auto",
+ min_pixels: int | None = None,
+ max_pixels: int | None = None,
+ target_height: int | None = None,
+ target_width: int | None = None,
+ ) -> tuple[torch.Tensor, torch.Tensor]:
+ if not isinstance(image, Image.Image):
+ raise TypeError(f"Expected a PIL image, got {type(image)}")
+ if image.mode == "RGBA":
+ background = Image.new("RGB", image.size, (255, 255, 255))
+ background.paste(image, mask=image.getchannel("A"))
+ image = background
+ elif image.mode != "RGB":
+ image = image.convert("RGB")
+
+ resized_h, resized_w = self.resized_size(
+ *image.size,
+ detail=detail,
+ min_pixels=min_pixels,
+ max_pixels=max_pixels,
+ target_height=target_height,
+ target_width=target_width,
+ )
+ image = image.resize((resized_w, resized_h), Image.Resampling.BICUBIC)
+ array = np.asarray(image, dtype=np.float32) / 255.0
+ array = (array - self.image_mean) / self.image_std
+ patches = array.transpose(2, 0, 1)[None]
+ if patches.shape[0] == 1:
+ patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))
+
+ channel = patches.shape[1]
+ grid_t = patches.shape[0] // self.temporal_patch_size
+ grid_h = resized_h // self.patch_size
+ grid_w = resized_w // self.patch_size
+ if self.pre_pixel_shuffle:
+ patches = patches.reshape(
+ grid_t,
+ self.temporal_patch_size,
+ channel,
+ grid_h // self.merge_size,
+ self.merge_size,
+ self.patch_size,
+ grid_w // self.merge_size,
+ self.merge_size,
+ self.patch_size,
+ )
+ patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
+ else:
+ patches = patches.reshape(
+ grid_t,
+ self.temporal_patch_size,
+ channel,
+ grid_h,
+ self.patch_size,
+ grid_w,
+ self.patch_size,
+ )
+ patches = patches.transpose(0, 3, 5, 2, 1, 4, 6)
+
+ pixel_values = torch.from_numpy(
+ patches.reshape(
+ grid_t * grid_h * grid_w,
+ channel * self.temporal_patch_size * self.patch_size * self.patch_size,
+ )
+ )
+ grid_thw = torch.tensor([grid_t, grid_h, grid_w], dtype=torch.long)
+ return pixel_values, grid_thw
+
+
+class Dots3NoteProcessor:
+ """Small HF-like processor used by vLLM's multimodal frontend."""
+
+ def __init__(
+ self,
+ tokenizer,
+ image_processor: Dots3NoteImageProcessor | None,
+ *,
+ max_model_len: int,
+ video_audio_enabled: bool,
+ audio_token_stride: int,
+ ) -> None:
+ self.tokenizer = tokenizer
+ self.image_processor = image_processor
+ self.max_model_len = max_model_len
+ self.video_audio_enabled = video_audio_enabled
+ self.audio_token_stride = audio_token_stride
+
+ def _token_ids(self, text: str) -> list[int]:
+ if hasattr(self.tokenizer, "encode"):
+ return self.tokenizer.encode(text, add_special_tokens=False)
+ token_ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
+ if token_ids and isinstance(token_ids[0], list):
+ return token_ids[0]
+ return token_ids
+
+ def _process_video(
+ self,
+ video,
+ *,
+ detail: str,
+ size_overrides: Mapping[str, int | None],
+ seq: int,
+ output_reserve: int | None,
+ audio_cap: float,
+ audio_sample_rate: int,
+ k_mode: str,
+ max_new_tokens: int,
+ question: str,
+ ) -> dict[str, torch.Tensor]:
+ if self.image_processor is None:
+ raise ValueError("This NOTE checkpoint has no vision encoder")
+ if not self.video_audio_enabled:
+ audio_cap = 0.0
+ parts = preprocess_dots3_note_video(
+ video,
+ tokenizer=self.tokenizer,
+ question=question,
+ seq=seq,
+ output_reserve=output_reserve,
+ audio_cap=audio_cap,
+ audio_sample_rate=audio_sample_rate,
+ k_mode=k_mode,
+ max_new_tokens=max_new_tokens,
+ )
+
+ fragments: list[str] = []
+ pixel_values: list[torch.Tensor] = []
+ image_grids: list[torch.Tensor] = []
+ audio_values: list[torch.Tensor] = []
+ modalities: list[int] = []
+ image_pad_count = 0
+ audio_pad_count = 0
+
+ for part in parts:
+ if part.kind == "text":
+ fragments.append(cast(str, part.value))
+ continue
+ if part.kind == "image":
+ pixels, grid = self.image_processor.preprocess(
+ cast(Image.Image, part.value),
+ detail=detail,
+ min_pixels=size_overrides.get("min_pixels"),
+ max_pixels=size_overrides.get("max_pixels"),
+ target_height=size_overrides.get("target_height"),
+ target_width=size_overrides.get("target_width"),
+ )
+ num_tokens = int(grid.prod()) // self.image_processor.merge_size**2
+ fragments.append(f"{IMAGE_START}{IMAGE_PAD * num_tokens}{IMAGE_END}")
+ pixel_values.append(pixels)
+ image_grids.append(grid)
+ modalities.append(0)
+ image_pad_count += num_tokens
+ continue
+
+ waveform = torch.from_numpy(
+ np.ascontiguousarray(part.value, dtype=np.float32)
+ ).view(torch.int32)
+ num_tokens = math.ceil(waveform.numel() / self.audio_token_stride)
+ fragments.append(f"{AUDIO_START}{AUDIO_PAD * num_tokens}{AUDIO_END}")
+ audio_values.append(waveform)
+ modalities.append(1)
+ audio_pad_count += num_tokens
+
+ if not pixel_values:
+ raise ValueError("NOTE video preprocessing produced no frames")
+ fragment_ids = self._token_ids("".join(fragments))
+ vocab = self.tokenizer.get_vocab()
+ embed_mask = torch.isin(
+ torch.tensor(fragment_ids, dtype=torch.long),
+ torch.tensor([vocab[IMAGE_PAD], vocab[AUDIO_PAD]], dtype=torch.long),
+ )
+ expected_embeds = image_pad_count + audio_pad_count
+ if int(embed_mask.sum()) != expected_embeds:
+ raise ValueError(
+ "NOTE video placeholder expansion produced an invalid embedding mask: "
+ f"expected={expected_embeds}, actual={int(embed_mask.sum())}"
+ )
+
+ empty_audio = torch.empty(0, dtype=torch.int32)
+ return {
+ "video_pixel_values": torch.cat(pixel_values),
+ "video_image_grid_thw": torch.stack(image_grids),
+ "video_audio_values": (
+ torch.cat(audio_values) if audio_values else empty_audio
+ ),
+ "video_audio_lengths": torch.tensor(
+ [value.numel() for value in audio_values],
+ dtype=torch.long,
+ ),
+ "video_modalities": torch.tensor(modalities, dtype=torch.uint8),
+ "video_input_ids": torch.tensor(fragment_ids, dtype=torch.long),
+ "video_embed_mask": embed_mask,
+ "video_frame_counts": torch.tensor([len(image_grids)], dtype=torch.long),
+ "video_audio_counts": torch.tensor([len(audio_values)], dtype=torch.long),
+ "video_emission_counts": torch.tensor([len(modalities)], dtype=torch.long),
+ "video_prompt_lengths": torch.tensor([len(fragment_ids)], dtype=torch.long),
+ "video_patch_counts": torch.tensor(
+ [sum(int(grid.prod()) for grid in image_grids)], dtype=torch.long
+ ),
+ "video_audio_sample_counts": torch.tensor(
+ [sum(value.numel() for value in audio_values)], dtype=torch.long
+ ),
+ }
+
+ def __call__(self, text: str, **kwargs: object) -> BatchFeature:
+ modality_order = [
+ key for key in kwargs if key in ("images", "audios", "videos")
+ ]
+ raw_images = kwargs.pop("images", None)
+ raw_audios = kwargs.pop("audios", None)
+ raw_videos = kwargs.pop("videos", None)
+ images = (
+ list(cast(Sequence[Image.Image], raw_images))
+ if raw_images is not None
+ else []
+ )
+ audios = (
+ list(cast(Sequence[np.ndarray], raw_audios))
+ if raw_audios is not None
+ else []
+ )
+ videos = (
+ list(cast(Sequence[object], raw_videos)) if raw_videos is not None else []
+ )
+ if videos and (images or audios):
+ raise ValueError(
+ "Dots3Note does not support mixing a native video with "
+ "separate image/audio inputs"
+ )
+ if len(videos) > 1:
+ raise ValueError("Dots3Note supports one video per request")
+ detail = kwargs.pop("image_detail", "auto")
+ details = list(detail) if isinstance(detail, (list, tuple)) else None
+ size_overrides: dict[str, int | None] = {}
+ for key in ("min_pixels", "max_pixels", "target_height", "target_width"):
+ if key in kwargs:
+ value = kwargs.pop(key)
+ size_overrides[key] = None if value is None else int(cast(Any, value))
+ video_seq = int(cast(Any, kwargs.pop("seq", self.max_model_len)))
+ reserve_value = kwargs.pop("output_reserve", None)
+ output_reserve = (
+ None if reserve_value is None else int(cast(Any, reserve_value))
+ )
+ audio_cap = float(cast(Any, kwargs.pop("audio_cap", 1.0)))
+ audio_sample_rate = int(cast(Any, kwargs.pop("audio_sr", _DEFAULT_SAMPLE_RATE)))
+ k_mode = str(kwargs.pop("k_mode", "eval_ek"))
+ max_new_tokens = int(cast(Any, kwargs.pop("max_new_tokens", 0)))
+ question = str(kwargs.pop("video_question", text))
+
+ tokenized = self.tokenizer(text, **kwargs)
+ input_ids = tokenized["input_ids"]
+ if isinstance(input_ids, torch.Tensor):
+ if input_ids.ndim == 1:
+ input_ids = input_ids.unsqueeze(0)
+ elif input_ids and isinstance(input_ids[0], int):
+ input_ids = [input_ids]
+ data: dict[str, object] = {"input_ids": input_ids}
+
+ for modality in modality_order:
+ if modality == "images" and images:
+ if self.image_processor is None:
+ raise ValueError("This NOTE checkpoint has no vision encoder")
+ pixel_values = []
+ grids = []
+ for idx, image in enumerate(images):
+ image_detail = details[idx] if details is not None else str(detail)
+ pixels, grid = self.image_processor.preprocess(
+ image,
+ detail=image_detail,
+ min_pixels=size_overrides.get("min_pixels"),
+ max_pixels=size_overrides.get("max_pixels"),
+ target_height=size_overrides.get("target_height"),
+ target_width=size_overrides.get("target_width"),
+ )
+ pixel_values.append(pixels)
+ grids.append(grid)
+ data["pixel_values"] = torch.cat(pixel_values)
+ data["image_grid_thw"] = torch.stack(grids)
+ elif modality == "audios" and audios:
+ waveforms = [
+ torch.from_numpy(
+ np.ascontiguousarray(audio, dtype=np.float32)
+ ).view(torch.int32)
+ for audio in audios
+ ]
+ data["audio_values"] = torch.cat(waveforms)
+ data["audio_lengths"] = torch.tensor(
+ [waveform.numel() for waveform in waveforms],
+ dtype=torch.long,
+ )
+ elif modality == "videos" and videos:
+ video_detail = details[0] if details is not None else str(detail)
+ data.update(
+ self._process_video(
+ videos[0],
+ detail=video_detail,
+ size_overrides=size_overrides,
+ seq=video_seq,
+ output_reserve=output_reserve,
+ audio_cap=audio_cap,
+ audio_sample_rate=audio_sample_rate,
+ k_mode=k_mode,
+ max_new_tokens=max_new_tokens,
+ question=question,
+ )
+ )
+ return BatchFeature(data=data)
+
+
+class Dots3NoteProcessingInfo(BaseProcessingInfo):
+ @cached_property
+ def vision_config(self) -> dict[str, Any] | None:
+ model_config = self.ctx.model_config
+ return load_note_config_section(
+ model_config.model,
+ model_config.revision,
+ "vision_config",
+ )
+
+ @cached_property
+ def audio_config(self) -> dict[str, Any] | None:
+ model_config = self.ctx.model_config
+ return load_note_config_section(
+ model_config.model,
+ model_config.revision,
+ "audio_config",
+ )
+
+ @cached_property
+ def image_processor(self) -> Dots3NoteImageProcessor | None:
+ if self.vision_config is None:
+ return None
+ model_config = self.ctx.model_config
+ processor_config = get_hf_file_to_dict(
+ "preprocessor_config.json",
+ model_config.model,
+ model_config.revision,
+ )
+ preprocessor_config = (processor_config or {}).get("vision_config")
+ if not isinstance(preprocessor_config, dict):
+ raise ValueError(
+ "NOTE vision checkpoint is missing preprocessor_config.json"
+ )
+ image_details = preprocessor_config.get("image_details", {})
+ return Dots3NoteImageProcessor(preprocessor_config, image_details)
+
+ @cached_property
+ def processor(self) -> Dots3NoteProcessor:
+ mm_config = self.ctx.get_mm_config()
+ audio_config = self.audio_config or {}
+ audio_token_stride = (
+ _HOP_LENGTH
+ * (8 if audio_config.get("use_conv2d_stem", True) else 2)
+ * int(audio_config.get("merge_factor", 1))
+ )
+ return Dots3NoteProcessor(
+ self.get_tokenizer(),
+ self.image_processor,
+ max_model_len=self.ctx.model_config.max_model_len,
+ video_audio_enabled=(
+ self.audio_config is not None
+ and (
+ mm_config.get_limit_per_prompt("audio") > 0
+ or mm_config.get_limit_per_prompt("video") > 0
+ )
+ ),
+ audio_token_stride=audio_token_stride,
+ )
+
+ def get_hf_processor(self, **kwargs: object) -> Dots3NoteProcessor:
+ return self.processor
+
+ def get_data_parser(self) -> MultiModalDataParser:
+ sample_rate = int((self.audio_config or {}).get("sampling_rate", 16000))
+ return MultiModalDataParser(target_sr=float(sample_rate), target_channels=1)
+
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
+ limits: dict[str, int | None] = {}
+ if self.vision_config is not None:
+ limits["image"] = 512
+ limits["video"] = 1
+ if self.audio_config is not None:
+ limits["audio"] = 128
+ return limits
+
+ def get_max_image_size(self) -> tuple[int, int]:
+ image_processor = self.image_processor
+ if image_processor is None:
+ return 0, 0
+ factor = image_processor.factor
+ max_tokens = image_processor.max_pixels // (factor * factor)
+ for height_factor in range(math.isqrt(max_tokens), 0, -1):
+ if max_tokens % height_factor == 0:
+ width_factor = max_tokens // height_factor
+ if width_factor / height_factor <= 200:
+ return factor * width_factor, factor * height_factor
+ return factor * max_tokens, factor
+
+ def get_mm_max_tokens_per_item(
+ self,
+ seq_len: int,
+ mm_counts: Mapping[str, int],
+ ) -> Mapping[str, int]:
+ result: dict[str, int] = {}
+ if self.image_processor is not None:
+ factor = self.image_processor.factor
+ result["image"] = self.image_processor.max_pixels // (factor * factor)
+ if self.audio_config is not None:
+ stride = (
+ _HOP_LENGTH
+ * (8 if self.audio_config.get("use_conv2d_stem", True) else 2)
+ * int(self.audio_config.get("merge_factor", 1))
+ )
+ chunk_samples = int(
+ self.audio_config.get("chunk_seconds", 60)
+ * self.audio_config.get("sampling_rate", _DEFAULT_SAMPLE_RATE)
+ )
+ result["audio"] = min(seq_len, math.ceil(chunk_samples / stride))
+ if self.image_processor is not None:
+ result["video"] = max(1, seq_len - seq_len // 4)
+ return result
+
+
+class Dots3NoteDummyInputsBuilder(BaseDummyInputsBuilder[Dots3NoteProcessingInfo]):
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
+ return (
+ f"{IMAGE_START}{IMAGE_PAD}{IMAGE_END}" * mm_counts.get("image", 0)
+ + f"{AUDIO_START}{AUDIO_PAD}{AUDIO_END}" * mm_counts.get("audio", 0)
+ + VIDEO_PLACEHOLDER * mm_counts.get("video", 0)
+ )
+
+ def get_dummy_mm_data(
+ self,
+ seq_len: int,
+ mm_counts: Mapping[str, int],
+ mm_options: Mapping[str, BaseDummyOptions],
+ ) -> MultiModalDataDict:
+ data: dict[str, Any] = {}
+ num_images = mm_counts.get("image", 0)
+ if num_images:
+ width, height = self.info.get_max_image_size()
+ data["image"] = self._get_dummy_images(
+ width=width,
+ height=height,
+ num_images=num_images,
+ overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
+ )
+ num_audios = mm_counts.get("audio", 0)
+ if num_audios:
+ audio_config = self.info.audio_config or {}
+ stride = (
+ _HOP_LENGTH
+ * (8 if audio_config.get("use_conv2d_stem", True) else 2)
+ * int(audio_config.get("merge_factor", 1))
+ )
+ chunk_samples = int(
+ audio_config.get("chunk_seconds", 60)
+ * audio_config.get("sampling_rate", _DEFAULT_SAMPLE_RATE)
+ )
+ data["audio"] = self._get_dummy_audios(
+ length=min(chunk_samples, max(1, seq_len) * stride),
+ num_audios=num_audios,
+ overrides=cast(AudioDummyOptions | None, mm_options.get("audio")),
+ )
+ num_videos = mm_counts.get("video", 0)
+ if num_videos:
+ width, height = self.info.get_max_image_size()
+ factor = self.info.image_processor.factor # type: ignore[union-attr]
+ frame_tokens = max(1, width * height // (factor * factor))
+ video_tokens = max(1, seq_len - seq_len // 4)
+ num_frames = max(4, math.ceil(video_tokens / frame_tokens))
+ data["video"] = self._get_dummy_videos(
+ width=width,
+ height=height,
+ num_frames=num_frames,
+ num_videos=num_videos,
+ overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
+ )
+ return data
+
+
+class Dots3NoteMultiModalProcessor(BaseMultiModalProcessor[Dots3NoteProcessingInfo]):
+ def _get_hf_mm_data(
+ self,
+ mm_items: MultiModalDataItems,
+ ) -> tuple[Mapping[str, object], Mapping[str, object]]:
+ processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
+ if "video" not in mm_items:
+ return processor_data, passthrough_data
+
+ videos = mm_items.get_items("video", VideoProcessorItems)
+ raw_videos: list[object] = []
+ for index, item in enumerate(videos.data):
+ if isinstance(item, MediaWithBytes):
+ raw_videos.append(item.original_bytes)
+ else:
+ raw_videos.append(videos.get(index))
+ processor_data = dict(processor_data)
+ processor_data["videos"] = raw_videos
+ return processor_data, passthrough_data
+
+ def _call_hf_processor(
+ self,
+ prompt: str,
+ mm_data: Mapping[str, object],
+ mm_kwargs: Mapping[str, object],
+ tok_kwargs: Mapping[str, object],
+ ) -> BatchFeature:
+ return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
+
+ def _hf_processor_applies_updates(
+ self,
+ prompt_text: str,
+ mm_items: MultiModalDataItems,
+ hf_processor_mm_kwargs: Mapping[str, object],
+ tokenization_kwargs: Mapping[str, object],
+ ) -> bool:
+ return False
+
+ def _get_mm_fields_config(
+ self,
+ hf_inputs: BatchFeature,
+ hf_processor_mm_kwargs: Mapping[str, object],
+ ) -> Mapping[str, MultiModalFieldConfig]:
+ fields: dict[str, MultiModalFieldConfig] = {}
+ if "pixel_values" in hf_inputs:
+ grid_thw = hf_inputs["image_grid_thw"]
+ fields["pixel_values"] = MultiModalFieldConfig.flat_from_sizes(
+ "image", grid_thw.prod(-1)
+ )
+ fields["image_grid_thw"] = MultiModalFieldConfig.batched(
+ "image", keep_on_cpu=True
+ )
+ if "audio_values" in hf_inputs:
+ fields["audio_values"] = MultiModalFieldConfig.flat_from_sizes(
+ "audio", hf_inputs["audio_lengths"]
+ )
+ fields["audio_lengths"] = MultiModalFieldConfig.batched(
+ "audio", keep_on_cpu=True
+ )
+ if "video_input_ids" in hf_inputs:
+ prompt_lengths = hf_inputs["video_prompt_lengths"]
+ frame_counts = hf_inputs["video_frame_counts"]
+ audio_counts = hf_inputs["video_audio_counts"]
+ emission_counts = hf_inputs["video_emission_counts"]
+ fields["video_pixel_values"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", hf_inputs["video_patch_counts"]
+ )
+ fields["video_image_grid_thw"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", frame_counts, keep_on_cpu=True
+ )
+ fields["video_audio_values"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", hf_inputs["video_audio_sample_counts"]
+ )
+ fields["video_audio_lengths"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", audio_counts, keep_on_cpu=True
+ )
+ fields["video_modalities"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", emission_counts, keep_on_cpu=True
+ )
+ fields["video_input_ids"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", prompt_lengths, keep_on_cpu=True
+ )
+ fields["video_embed_mask"] = MultiModalFieldConfig.flat_from_sizes(
+ "video", prompt_lengths, keep_on_cpu=True
+ )
+ for key in (
+ "video_frame_counts",
+ "video_audio_counts",
+ "video_emission_counts",
+ ):
+ fields[key] = MultiModalFieldConfig.batched("video", keep_on_cpu=True)
+ return fields
+
+ def _get_prompt_updates(
+ self,
+ mm_items: MultiModalDataItems,
+ hf_processor_mm_kwargs: Mapping[str, object],
+ out_mm_kwargs: MultiModalKwargsItems,
+ ) -> Sequence[PromptUpdate]:
+ vocab = self.info.get_tokenizer().get_vocab()
+ updates: list[PromptUpdate] = []
+
+ if "image" in out_mm_kwargs:
+ image_start_id = vocab[IMAGE_START]
+ image_pad_id = vocab[IMAGE_PAD]
+ image_end_id = vocab[IMAGE_END]
+ merge_size = self.info.image_processor.merge_size # type: ignore[union-attr]
+
+ def image_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
+ grid = out_mm_kwargs["image"][item_idx]["image_grid_thw"].data
+ assert isinstance(grid, torch.Tensor)
+ num_tokens = int(grid.prod()) // merge_size**2
+ return PromptUpdateDetails.select_token_id(
+ [image_start_id] + [image_pad_id] * num_tokens + [image_end_id],
+ image_pad_id,
+ )
+
+ updates.append(
+ PromptReplacement(
+ modality="image",
+ target=[image_start_id, image_pad_id, image_end_id],
+ replacement=image_replacement,
+ )
+ )
+
+ if "audio" in out_mm_kwargs:
+ audio_start_id = vocab[AUDIO_START]
+ audio_pad_id = vocab[AUDIO_PAD]
+ audio_end_id = vocab[AUDIO_END]
+ config = self.info.audio_config or {}
+ stride = (
+ _HOP_LENGTH
+ * (8 if config.get("use_conv2d_stem", True) else 2)
+ * int(config.get("merge_factor", 1))
+ )
+
+ def audio_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
+ length = out_mm_kwargs["audio"][item_idx]["audio_lengths"].data
+ assert isinstance(length, torch.Tensor)
+ num_tokens = math.ceil(int(length.item()) / stride)
+ return PromptUpdateDetails.select_token_id(
+ [audio_start_id] + [audio_pad_id] * num_tokens + [audio_end_id],
+ audio_pad_id,
+ )
+
+ updates.append(
+ PromptReplacement(
+ modality="audio",
+ target=[audio_start_id, audio_pad_id, audio_end_id],
+ replacement=audio_replacement,
+ )
+ )
+ if "video" in out_mm_kwargs:
+
+ def video_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
+ item = out_mm_kwargs["video"][item_idx]
+ input_ids_data = item["video_input_ids"].data
+ embed_mask_data = item["video_embed_mask"].data
+ assert isinstance(input_ids_data, torch.Tensor)
+ assert isinstance(embed_mask_data, torch.Tensor)
+ input_ids = input_ids_data.tolist()
+ embed_mask = embed_mask_data.bool()
+
+ def select_video_embeds(tokenizer, full) -> torch.Tensor:
+ del tokenizer, full
+ return embed_mask
+
+ return PromptUpdateDetails(
+ full=input_ids,
+ is_embed=select_video_embeds,
+ )
+
+ updates.append(
+ PromptReplacement(
+ modality="video",
+ target=VIDEO_PLACEHOLDER,
+ replacement=video_replacement,
+ )
+ )
+ return updates
diff --git a/vllm/models/dots3_note/common/video.py b/vllm/models/dots3_note/common/video.py
new file mode 100644
index 0000000000..907bb75738
--- /dev/null
+++ b/vllm/models/dots3_note/common/video.py
@@ -0,0 +1,497 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Train-consistent video preprocessing for Dots3Note."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import math
+import os
+import random
+from dataclasses import dataclass
+from typing import Literal
+
+import numpy as np
+from PIL import Image
+
+_ALIGN = 28
+_MIN_FRAMES = 4
+_PF_FLOOR = 128
+_PF_CEIL = 1024
+_FPS_CAP = 1.0
+_FPS_MIN = 0.2
+_FRAME_OVERHEAD = 15
+_BUDGET_OVERHEAD = 2240
+_INTERLEAVE_MIN_SECONDS = 1.0
+_AUDIO_SAMPLE_RATE = 16000
+_AUDIO_SAMPLES_PER_TOKEN = 1280
+_AUDIO_CHUNK_SECONDS = 30
+
+
+@dataclass(frozen=True)
+class Dots3NoteVideoPart:
+ kind: Literal["text", "image", "audio"]
+ value: str | Image.Image | np.ndarray
+
+
+def _token_len(tokenizer, text: str) -> int:
+ if not text:
+ return 0
+ if hasattr(tokenizer, "encode"):
+ return len(tokenizer.encode(text, add_special_tokens=False))
+ return len(tokenizer(text, add_special_tokens=False)["input_ids"])
+
+
+def _compute_target_size(
+ orig_h: int,
+ orig_w: int,
+ min_pixels: int,
+ max_pixels: int,
+) -> tuple[int, int]:
+ height = max(_ALIGN, round(orig_h / _ALIGN) * _ALIGN)
+ width = max(_ALIGN, round(orig_w / _ALIGN) * _ALIGN)
+ if height * width > max_pixels:
+ beta = math.sqrt(orig_h * orig_w / max_pixels)
+ height = max(_ALIGN, math.floor(orig_h / beta / _ALIGN) * _ALIGN)
+ width = max(_ALIGN, math.floor(orig_w / beta / _ALIGN) * _ALIGN)
+ elif height * width < min_pixels:
+ beta = math.sqrt(min_pixels / max(1, orig_h * orig_w))
+ height = math.ceil(orig_h * beta / _ALIGN) * _ALIGN
+ width = math.ceil(orig_w * beta / _ALIGN) * _ALIGN
+ if height * width > max_pixels:
+ beta = math.sqrt(height * width / max_pixels)
+ height = max(_ALIGN, math.floor(height / beta / _ALIGN) * _ALIGN)
+ width = max(_ALIGN, math.floor(width / beta / _ALIGN) * _ALIGN)
+ return int(height), int(width)
+
+
+def _real_patches_at(
+ orig_h: int,
+ orig_w: int,
+ patch_cap: int,
+) -> int:
+ target_h, target_w = _compute_target_size(
+ orig_h,
+ orig_w,
+ _PF_FLOOR * _ALIGN * _ALIGN,
+ max(_PF_FLOOR, patch_cap) * _ALIGN * _ALIGN,
+ )
+ return (target_h // _ALIGN) * (target_w // _ALIGN)
+
+
+def _frame_hard_cap(seq_length: int) -> int:
+ need = max(1, (seq_length - _BUDGET_OVERHEAD) // (_PF_FLOOR + 15))
+ if need <= 1024:
+ return 1024
+ cap = 1
+ while cap < need:
+ cap <<= 1
+ return cap
+
+
+def _solve_degrade(
+ visual_budget: int,
+ duration: float,
+ orig_h: int,
+ orig_w: int,
+ orig_fps: float,
+ seq_length: int,
+) -> tuple[int, int]:
+ aligned_h = max(_ALIGN, round(orig_h / _ALIGN) * _ALIGN)
+ aligned_w = max(_ALIGN, round(orig_w / _ALIGN) * _ALIGN)
+ orig_max_pf = (aligned_h // _ALIGN) * (aligned_w // _ALIGN)
+ fps_cap = min(_FPS_CAP, max(orig_fps, 1e-6))
+ pf_cap = min(_PF_CEIL, max(orig_max_pf, _PF_FLOOR))
+ frame_cap = _frame_hard_cap(seq_length)
+
+ def usage(scale: float) -> tuple[int, float, int, int]:
+ fps = _FPS_MIN + scale * (fps_cap - _FPS_MIN)
+ patch_cap = _PF_FLOOR + scale * (pf_cap - _PF_FLOOR)
+ num_frames = max(
+ _MIN_FRAMES,
+ min(int(round(duration * fps)), frame_cap),
+ )
+ patches = _real_patches_at(
+ orig_h,
+ orig_w,
+ int(round(patch_cap)),
+ )
+ return (
+ num_frames * (patches + _FRAME_OVERHEAD),
+ fps,
+ int(round(patch_cap)),
+ num_frames,
+ )
+
+ if usage(1.0)[0] <= visual_budget:
+ _, _, patch_cap, num_frames = usage(1.0)
+ return num_frames, patch_cap
+
+ floor_cost = _real_patches_at(orig_h, orig_w, _PF_FLOOR) + _FRAME_OVERHEAD
+ if usage(0.0)[0] > visual_budget:
+ return max(
+ _MIN_FRAMES,
+ min(visual_budget // floor_cost, frame_cap),
+ ), _PF_FLOOR
+
+ low, high = 0.0, 1.0
+ for _ in range(50):
+ mid = (low + high) / 2
+ if usage(mid)[0] <= visual_budget:
+ low = mid
+ else:
+ high = mid
+ _, _, patch_cap, num_frames = usage(low)
+ return num_frames, patch_cap
+
+
+def _audio_tokens(duration: float, sample_rate: int) -> int:
+ if duration <= 0:
+ return 0
+ total_samples = int(duration * sample_rate)
+ chunk_samples = _AUDIO_CHUNK_SECONDS * sample_rate
+ num_tokens = 0
+ position = 0
+ while position < total_samples:
+ length = min(chunk_samples, total_samples - position)
+ num_tokens += math.ceil(length / _AUDIO_SAMPLES_PER_TOKEN)
+ position += chunk_samples
+ return num_tokens + 2
+
+
+def _decode_audio(
+ video_bytes: bytes,
+ sample_rate: int,
+) -> tuple[np.ndarray | None, float]:
+ from torchcodec.decoders import AudioDecoder
+
+ try:
+ decoder = AudioDecoder(
+ video_bytes,
+ sample_rate=sample_rate,
+ num_channels=1,
+ )
+ samples = decoder.get_all_samples()
+ except Exception:
+ return None, 0.0
+
+ waveform = samples.data
+ if waveform is None or waveform.numel() == 0:
+ return None, 0.0
+ if waveform.ndim == 2:
+ waveform = waveform.mean(dim=0) if waveform.shape[0] > 1 else waveform[0]
+ array = waveform.cpu().numpy()
+ pcm = (np.clip(array, -1.0, 1.0) * 32767.0).astype(np.int16)
+ return pcm, float(pcm.shape[0]) / sample_rate
+
+
+def _open_video(video_bytes: bytes):
+ from torchcodec.decoders import VideoDecoder
+
+ try:
+ return VideoDecoder(
+ video_bytes,
+ dimension_order="NHWC",
+ num_ffmpeg_threads=1,
+ seek_mode="approximate",
+ )
+ except TypeError:
+ return VideoDecoder(
+ video_bytes,
+ dimension_order="NHWC",
+ num_ffmpeg_threads=1,
+ )
+
+
+def _jpeg_roundtrip(image: Image.Image, quality: int) -> Image.Image:
+ buffer = io.BytesIO()
+ image.save(buffer, format="JPEG", quality=quality)
+ buffer.seek(0)
+ with Image.open(buffer) as decoded:
+ return decoded.convert("RGB").copy()
+
+
+def _decode_frames(
+ video_bytes: bytes,
+ visual_budget: int,
+ seq_length: int,
+ jpeg_quality: int,
+) -> tuple[list[tuple[float, Image.Image]], float]:
+ decoder = _open_video(video_bytes)
+ metadata = decoder.metadata
+ duration = float(metadata.duration_seconds or 0)
+ orig_h = int(metadata.height)
+ orig_w = int(metadata.width)
+ total_frames = int(getattr(metadata, "num_frames", 0) or 0)
+ orig_fps = float(getattr(metadata, "average_fps", 0) or 0) or 25.0
+ if duration <= 0 or orig_h <= 0 or orig_w <= 0:
+ raise ValueError(
+ "Invalid video metadata: "
+ f"duration={duration}, height={orig_h}, width={orig_w}"
+ )
+ if total_frames <= 0:
+ total_frames = max(1, int(duration * orig_fps))
+
+ num_frames, patch_cap = _solve_degrade(
+ visual_budget,
+ duration,
+ orig_h,
+ orig_w,
+ orig_fps,
+ seq_length,
+ )
+ aligned_h = max(_ALIGN, round(orig_h / _ALIGN) * _ALIGN)
+ aligned_w = max(_ALIGN, round(orig_w / _ALIGN) * _ALIGN)
+ orig_max_pf = (aligned_h // _ALIGN) * (aligned_w // _ALIGN)
+ max_pixels = min(patch_cap, orig_max_pf) * _ALIGN * _ALIGN
+ target_h, target_w = _compute_target_size(
+ orig_h,
+ orig_w,
+ _PF_FLOOR * _ALIGN * _ALIGN,
+ max_pixels,
+ )
+ num_frames = max(_MIN_FRAMES, min(num_frames, total_frames))
+ if num_frames == 1:
+ indices = [0]
+ else:
+ step = (total_frames - 1) / (num_frames - 1)
+ indices = [int(round(i * step)) for i in range(num_frames)]
+ indices = sorted({max(0, min(index, total_frames - 1)) for index in indices})
+
+ try:
+ decoded = decoder.get_frames_at(indices=indices).data
+ except (IndexError, RuntimeError):
+ safe = list(indices)
+ while safe and safe[-1] > 0:
+ safe.pop()
+ try:
+ decoded = decoder.get_frames_at(indices=safe).data
+ indices = safe
+ break
+ except (IndexError, RuntimeError):
+ continue
+ else:
+ raise
+
+ frames: list[tuple[float, Image.Image]] = []
+ for index, frame in zip(indices, decoded):
+ array = frame.cpu().numpy() if hasattr(frame, "cpu") else np.asarray(frame)
+ image = Image.fromarray(array)
+ if image.size != (target_w, target_h):
+ image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
+ frames.append(
+ (round(index / orig_fps, 3), _jpeg_roundtrip(image, jpeg_quality))
+ )
+ return frames, duration
+
+
+def _decoded_frames(video) -> tuple[np.ndarray, float]:
+ metadata = None
+ if isinstance(video, tuple):
+ video, metadata = video
+ frames = np.asarray(video)
+ if frames.ndim != 4 or frames.shape[-1] not in (3, 4):
+ raise TypeError(
+ "Decoded NOTE video must have shape (frames, height, width, channels)"
+ )
+ fps = float((metadata or {}).get("fps", 1.0)) if metadata else 1.0
+ return frames, max(fps, 1e-6)
+
+
+def _prepare_decoded_frames(
+ video,
+ visual_budget: int,
+ seq_length: int,
+ jpeg_quality: int,
+) -> tuple[list[tuple[float, Image.Image]], float]:
+ frames, fps = _decoded_frames(video)
+ duration = len(frames) / fps
+ orig_h, orig_w = frames.shape[1:3]
+ num_frames, patch_cap = _solve_degrade(
+ visual_budget,
+ duration,
+ orig_h,
+ orig_w,
+ fps,
+ seq_length,
+ )
+ num_frames = min(max(1, num_frames), len(frames))
+ indices = np.linspace(0, len(frames) - 1, num_frames).round().astype(int)
+ target_h, target_w = _compute_target_size(
+ orig_h,
+ orig_w,
+ _PF_FLOOR * _ALIGN * _ALIGN,
+ patch_cap * _ALIGN * _ALIGN,
+ )
+ output = []
+ for index in sorted(set(indices.tolist())):
+ image = Image.fromarray(frames[index]).convert("RGB")
+ if image.size != (target_w, target_h):
+ image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
+ output.append((round(index / fps, 3), _jpeg_roundtrip(image, jpeg_quality)))
+ return output, duration
+
+
+def _format_timestamp(seconds: float) -> str:
+ total_centiseconds = int(round(max(seconds, 0.0) * 100))
+ hours = total_centiseconds // 360000
+ minutes = (total_centiseconds // 6000) % 60
+ secs = (total_centiseconds // 100) % 60
+ centiseconds = total_centiseconds % 100
+ return f"{hours:02d}:{minutes:02d}:{secs:02d}.{centiseconds:02d}"
+
+
+def _group_bounds(
+ num_frames: int,
+ duration: float,
+ mode: str,
+ rng: random.Random,
+) -> list[int]:
+ if num_frames <= 1 or duration <= 0:
+ return [0, num_frames]
+ max_groups = min(
+ num_frames,
+ max(1, int(duration // _INTERLEAVE_MIN_SECONDS)),
+ )
+ if mode == "whole" or max_groups <= 1:
+ groups = 1
+ elif mode == "eval30":
+ groups = round(math.sqrt(max_groups))
+ elif mode == "eval_ek":
+ groups = round((max_groups - 1) / math.log(max_groups))
+ elif mode == "logk":
+ groups = round(math.exp(rng.uniform(0.0, math.log(max_groups))))
+ else:
+ raise ValueError(f"Unsupported video k_mode: {mode}")
+ groups = max(1, min(max_groups, groups))
+ if groups == 1:
+ return [0, num_frames]
+ if mode == "logk":
+ cuts = sorted(rng.sample(range(1, num_frames), groups - 1))
+ else:
+ cuts = sorted(
+ {
+ round(index * num_frames / groups)
+ for index in range(1, groups)
+ if 0 < round(index * num_frames / groups) < num_frames
+ }
+ )
+ return [0, *cuts, num_frames]
+
+
+def preprocess_dots3_note_video(
+ video,
+ *,
+ tokenizer,
+ question: str,
+ seq: int,
+ output_reserve: int | None = None,
+ audio_cap: float = 1.0,
+ audio_sample_rate: int = _AUDIO_SAMPLE_RATE,
+ k_mode: str = "eval_ek",
+ max_new_tokens: int = 0,
+) -> list[Dots3NoteVideoPart]:
+ """Expand one video into timestamped, interleaved image/audio parts."""
+ if seq <= 0:
+ raise ValueError(f"seq must be positive, got {seq}")
+ configured_reserve = seq // 4 if output_reserve is None else output_reserve
+ effective_reserve = max(configured_reserve, max_new_tokens)
+ if effective_reserve >= seq:
+ raise ValueError(
+ "output_reserve/max_new_tokens must leave room for video input"
+ )
+ if audio_cap < 0:
+ raise ValueError(f"audio_cap must be non-negative, got {audio_cap}")
+ if audio_sample_rate <= 0:
+ raise ValueError(f"audio_sample_rate must be positive, got {audio_sample_rate}")
+
+ seq_length = seq - effective_reserve
+ jpeg_quality = int(os.environ.get("XHS_VIDEO_JPEG_QUALITY", "85"))
+ pcm: np.ndarray | None = None
+ audio_duration = 0.0
+ if isinstance(video, bytes) and audio_cap > 0:
+ pcm, audio_duration = _decode_audio(video, audio_sample_rate)
+
+ audio_token_count = (
+ _audio_tokens(audio_duration, audio_sample_rate) if pcm is not None else 0
+ )
+ estimated_frames = max(1, int(audio_duration * _FPS_CAP))
+ max_groups = min(
+ estimated_frames,
+ max(1, int(audio_duration // _INTERLEAVE_MIN_SECONDS)),
+ )
+ reserved_audio_tokens = audio_token_count + 3 * max_groups
+ min_visual_tokens = _MIN_FRAMES * (_PF_FLOOR + _FRAME_OVERHEAD)
+ if (
+ audio_token_count > audio_cap * seq_length
+ or reserved_audio_tokens + min_visual_tokens + _BUDGET_OVERHEAD > seq_length
+ ):
+ pcm = None
+ audio_duration = 0.0
+ reserved_audio_tokens = 0
+
+ overhead = (
+ _token_len(
+ tokenizer,
+ "<|system|>You are a helpful assistant.<|endofsystem|>\n",
+ )
+ + 2
+ + _token_len(tokenizer, "")
+ + 64
+ )
+ visual_budget = max(
+ _PF_FLOOR + _FRAME_OVERHEAD,
+ seq_length - overhead - reserved_audio_tokens,
+ )
+ if isinstance(video, bytes):
+ frames, video_duration = _decode_frames(
+ video,
+ visual_budget,
+ seq_length,
+ jpeg_quality,
+ )
+ else:
+ frames, video_duration = _prepare_decoded_frames(
+ video,
+ visual_budget,
+ seq_length,
+ jpeg_quality,
+ )
+
+ if pcm is None:
+ parts: list[Dots3NoteVideoPart] = []
+ for timestamp, image in frames:
+ parts.append(
+ Dots3NoteVideoPart("text", f"<{_format_timestamp(timestamp)}>")
+ )
+ parts.append(Dots3NoteVideoPart("image", image))
+ return parts
+
+ video_id = hashlib.sha1(video).hexdigest()
+ record_key = hashlib.sha1(f"{video_id}|{question}".encode()).hexdigest()
+ seed_hex = hashlib.sha1(f"42|flatten|{record_key}".encode()).hexdigest()
+ rng = random.Random(int(seed_hex[:8], 16))
+ bounds = _group_bounds(len(frames), audio_duration, k_mode, rng)
+ parts = []
+ for group in range(len(bounds) - 1):
+ start, end = bounds[group], bounds[group + 1]
+ if end <= start:
+ continue
+ start_time = 0.0 if group == 0 else frames[start][0]
+ end_time = audio_duration if group == len(bounds) - 2 else frames[end][0]
+ if end_time <= start_time:
+ end_time = start_time + audio_duration / max(1, len(bounds) - 1)
+ for timestamp, image in frames[start:end]:
+ parts.append(
+ Dots3NoteVideoPart("text", f"<{_format_timestamp(timestamp)}>")
+ )
+ parts.append(Dots3NoteVideoPart("image", image))
+ sample_start = max(0, int(round(start_time * audio_sample_rate)))
+ sample_end = min(len(pcm), int(round(end_time * audio_sample_rate)))
+ if sample_end > sample_start:
+ segment = np.ascontiguousarray(
+ pcm[sample_start:sample_end].astype(np.float32) / 32768.0
+ )
+ parts.append(Dots3NoteVideoPart("audio", segment))
+ return parts
diff --git a/vllm/models/dots3_note/nvidia/__init__.py b/vllm/models/dots3_note/nvidia/__init__.py
new file mode 100644
index 0000000000..2912b340db
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/__init__.py
@@ -0,0 +1,3 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA-specific Dots3Note model implementations."""
diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py
new file mode 100644
index 0000000000..d0ef7dbfd6
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/attention.py
@@ -0,0 +1,807 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Dots3 NOTE sliding-window MLA attention backends for Hopper.
+
+Prefill and mixed batches expand the latent cache and use FlashAttention-3
+varlen MHA. Decode-only batches use the Triton absorbed-MQA kernel.
+"""
+
+from dataclasses import dataclass
+
+import torch
+
+import vllm._custom_ops as ops
+from vllm.model_executor.layers.attention.mla_attention import (
+ MLACommonDecodeMetadata,
+ MLACommonMetadata,
+ QueryLenSupport,
+)
+from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
+from vllm.utils.torch_utils import is_quantized_kv_cache
+from vllm.v1.attention.backend import AttentionLayer, CommonAttentionMetadata
+from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
+ FlashAttnMLASparseBackend,
+ FlashAttnMLASparseImpl,
+ FlashAttnMLASparseMetadata,
+)
+from vllm.v1.attention.backends.mla.prefill.base import MLADimensions
+from vllm.v1.attention.backends.mla.prefill.flash_attn import (
+ FlashAttnPrefillBackend,
+)
+from vllm.v1.attention.backends.mla.sparse_utils import (
+ triton_convert_req_index_to_global_index,
+)
+from vllm.v1.attention.backends.mla.triton_mla import (
+ TritonMLABackend,
+ TritonMLAImpl,
+ TritonMLAMetadataBuilder,
+)
+from vllm.v1.worker.workspace import (
+ current_workspace_manager,
+ is_workspace_manager_initialized,
+)
+from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
+
+
+@triton.jit
+def _gather_swa_kv_kernel(
+ kv_cache_ptr,
+ block_table_ptr,
+ seq_lens_ptr,
+ kv_out_ptr,
+ valid_out_ptr,
+ k_scale_ptr,
+ kv_cache_stride_block,
+ kv_cache_stride_token,
+ kv_cache_stride_dim,
+ block_table_stride_req,
+ kv_out_stride_req,
+ kv_out_stride_token,
+ kv_out_stride_dim,
+ valid_out_stride_req,
+ valid_out_stride_token,
+ GATHER_LEN: tl.constexpr,
+ KV_DIM: tl.constexpr,
+ BLOCK_T: tl.constexpr,
+ BLOCK_D: tl.constexpr,
+ PAGE_SIZE: tl.constexpr,
+):
+ req_idx = tl.program_id(0)
+ token_block_idx = tl.program_id(1)
+ dim_block_idx = tl.program_id(2)
+
+ seq_len = tl.load(seq_lens_ptr + req_idx).to(tl.int32)
+ gather_start = tl.maximum(seq_len - GATHER_LEN, 0)
+ token_offsets = token_block_idx * BLOCK_T + tl.arange(0, BLOCK_T)
+ logical_tokens = gather_start + token_offsets
+ logical_valid = (token_offsets < GATHER_LEN) & (logical_tokens < seq_len)
+
+ logical_pages = logical_tokens // PAGE_SIZE
+ page_offsets = logical_tokens - logical_pages * PAGE_SIZE
+ physical_pages = tl.load(
+ block_table_ptr + req_idx * block_table_stride_req + logical_pages,
+ mask=logical_valid,
+ other=-1,
+ ).to(tl.int64)
+ physical_valid = logical_valid & (physical_pages >= 0)
+
+ dim_offsets = dim_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
+ values = tl.load(
+ kv_cache_ptr
+ + physical_pages[:, None] * kv_cache_stride_block
+ + page_offsets[:, None] * kv_cache_stride_token
+ + dim_offsets[None, :] * kv_cache_stride_dim,
+ mask=physical_valid[:, None] & (dim_offsets[None, :] < KV_DIM),
+ other=0.0,
+ )
+ if values.dtype.is_fp8():
+ values = values.to(tl.float32) * tl.load(k_scale_ptr)
+ tl.store(
+ kv_out_ptr
+ + req_idx * kv_out_stride_req
+ + token_offsets[:, None] * kv_out_stride_token
+ + dim_offsets[None, :] * kv_out_stride_dim,
+ values,
+ mask=(token_offsets[:, None] < GATHER_LEN) & (dim_offsets[None, :] < KV_DIM),
+ )
+ tl.store(
+ valid_out_ptr
+ + req_idx * valid_out_stride_req
+ + token_offsets * valid_out_stride_token,
+ physical_valid,
+ mask=(dim_block_idx == 0) & (token_offsets < GATHER_LEN),
+ )
+
+
+@triton.jit
+def _apply_swa_score_mask_kernel(
+ scores_ptr,
+ seq_lens_ptr,
+ valid_ptr,
+ scores_stride_req,
+ scores_stride_head,
+ scores_stride_query,
+ scores_stride_token,
+ valid_stride_req,
+ valid_stride_token,
+ GATHER_LEN: tl.constexpr,
+ WINDOW_SIZE: tl.constexpr,
+ QUERY_LEN: tl.constexpr,
+ BLOCK_T: tl.constexpr,
+):
+ req_idx = tl.program_id(0)
+ head_idx = tl.program_id(1)
+ query_idx = tl.program_id(2)
+
+ token_offsets = tl.arange(0, BLOCK_T)
+ seq_len = tl.load(seq_lens_ptr + req_idx).to(tl.int32)
+ gather_start = tl.maximum(seq_len - GATHER_LEN, 0)
+ kv_positions = gather_start + token_offsets
+ query_position = seq_len - QUERY_LEN + query_idx
+ page_valid = tl.load(
+ valid_ptr + req_idx * valid_stride_req + token_offsets * valid_stride_token,
+ mask=token_offsets < GATHER_LEN,
+ other=0,
+ ).to(tl.int1)
+ valid = (
+ (token_offsets < GATHER_LEN)
+ & page_valid
+ & (kv_positions <= query_position)
+ & (kv_positions >= query_position - WINDOW_SIZE + 1)
+ & (query_position >= 0)
+ )
+ tl.store(
+ scores_ptr
+ + req_idx * scores_stride_req
+ + head_idx * scores_stride_head
+ + query_idx * scores_stride_query
+ + token_offsets * scores_stride_token,
+ tl.full((BLOCK_T,), -3.4028234663852886e38, tl.float32),
+ mask=(token_offsets < GATHER_LEN) & ~valid,
+ )
+
+
+@dataclass
+class Dots3NoteDecodeMetadata(MLACommonDecodeMetadata):
+ query_len: int
+
+
+@dataclass
+class _SlidingWindowChunk:
+ req_start: int
+ req_end: int
+ query_start: int
+ query_end: int
+ cu_seq_lens_q: torch.Tensor
+ cu_seq_lens_k: torch.Tensor
+ starts: torch.Tensor
+ token_to_seq: torch.Tensor
+ num_kv_tokens: int
+ max_seq_len_q: int
+ max_seq_len_k: int
+
+
+@dataclass
+class _SlidingWindowMetadata:
+ chunks: list[_SlidingWindowChunk]
+ workspace: torch.Tensor
+
+
+def _build_sliding_window_metadata(
+ *,
+ seq_lens_cpu: torch.Tensor,
+ query_start_loc_cpu: torch.Tensor,
+ sliding_window: int,
+ workspace: torch.Tensor,
+ workspace_size: int,
+ device: torch.device,
+) -> _SlidingWindowMetadata:
+ """Plan per-request latent-cache gathers for SWA varlen attention."""
+ query_lens_cpu = (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).to(
+ dtype=torch.int32
+ )
+ seq_lens_cpu = seq_lens_cpu.to(dtype=torch.int32)
+ kv_lens_cpu = torch.minimum(seq_lens_cpu, query_lens_cpu + sliding_window - 1)
+ starts_cpu = seq_lens_cpu - kv_lens_cpu
+
+ chunks: list[_SlidingWindowChunk] = []
+ req_start = 0
+ while req_start < query_lens_cpu.numel():
+ req_end = req_start
+ num_kv_tokens = 0
+ while req_end < query_lens_cpu.numel():
+ next_len = int(kv_lens_cpu[req_end].item())
+ if num_kv_tokens and num_kv_tokens + next_len > workspace_size:
+ break
+ if next_len > workspace_size:
+ raise ValueError(
+ "Dots3 NOTE SWA prefill window exceeds the MLA workspace: "
+ f"{next_len} > {workspace_size}"
+ )
+ num_kv_tokens += next_len
+ req_end += 1
+
+ query_lens = query_lens_cpu[req_start:req_end]
+ kv_lens = kv_lens_cpu[req_start:req_end]
+ num_reqs = req_end - req_start
+ cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
+ cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
+ torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:])
+ torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:])
+ token_to_seq_cpu = torch.repeat_interleave(
+ torch.arange(num_reqs, dtype=torch.int32), kv_lens
+ )
+ query_start = int(query_start_loc_cpu[req_start].item())
+ query_end = int(query_start_loc_cpu[req_end].item())
+ chunks.append(
+ _SlidingWindowChunk(
+ req_start=req_start,
+ req_end=req_end,
+ query_start=query_start,
+ query_end=query_end,
+ cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True),
+ cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True),
+ starts=starts_cpu[req_start:req_end].to(device, non_blocking=True),
+ token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
+ num_kv_tokens=num_kv_tokens,
+ max_seq_len_q=int(query_lens.max().item()),
+ max_seq_len_k=int(kv_lens.max().item()),
+ )
+ )
+ req_start = req_end
+
+ return _SlidingWindowMetadata(chunks=chunks, workspace=workspace)
+
+
+class Dots3NoteFlashAttnPrefillBackend(FlashAttnPrefillBackend):
+ """FA3 varlen prefill for the NOTE SWA MLA dimensions."""
+
+ @classmethod
+ def supports_mla_dimensions(cls, dims: MLADimensions) -> bool:
+ return dims == MLADimensions(
+ qk_nope_head_dim=192,
+ qk_rope_head_dim=64,
+ v_head_dim=128,
+ )
+
+ def __init__(self, *args, **kwargs) -> None:
+ super().__init__(*args, **kwargs)
+ # FA3 on SM90 does not support NOTE's Q/K=256, V=128 combination via
+ # the different-head-dimension path. Padding V selects its supported
+ # equal-dimension varlen kernel.
+ self.requires_v_padding = True
+
+ def supports_quant_output(self, quant_key) -> bool:
+ return False
+
+ def run_sliding_window(
+ self,
+ q: torch.Tensor,
+ k: torch.Tensor,
+ v: torch.Tensor,
+ cu_seq_lens_q: torch.Tensor,
+ cu_seq_lens_k: torch.Tensor,
+ max_seq_len_q: int,
+ max_seq_len_k: int,
+ sliding_window: int,
+ ) -> torch.Tensor:
+ output = self._flash_attn_varlen_diff_headdims(
+ q=q,
+ k=k,
+ v=v,
+ cu_seqlens_q=cu_seq_lens_q,
+ cu_seqlens_k=cu_seq_lens_k,
+ max_seqlen_q=max_seq_len_q,
+ max_seqlen_k=max_seq_len_k,
+ softmax_scale=self.scale,
+ causal=True,
+ window_size=(sliding_window - 1, 0),
+ return_softmax_lse=False,
+ )
+ assert isinstance(output, torch.Tensor)
+ return output
+
+
+class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder):
+ """Keep decode on MQA and route prefill/mixed batches through FA3."""
+
+ query_len_support = QueryLenSupport.UNIFORM
+
+ def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
+ self.sliding_window = kv_cache_spec.sliding_window
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
+
+ def _reserve_attn_logits_workspace(self) -> None:
+ if not is_workspace_manager_initialized():
+ return
+ batch = self.vllm_config.scheduler_config.max_num_seqs
+ max_query_len = self.reorder_batch_threshold or 1
+ gather_len = (self.sliding_window + max_query_len - 1 + 7) // 8 * 8
+ current_workspace_manager().get_simultaneous(
+ (
+ (
+ batch,
+ gather_len,
+ self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
+ ),
+ self.model_config.dtype,
+ ),
+ ((batch, gather_len), torch.bool),
+ )
+
+ def _build_decode(
+ self,
+ block_table_tensor: torch.Tensor,
+ seq_lens_device: torch.Tensor,
+ max_seq_len: int,
+ query_start_loc_cpu: torch.Tensor,
+ query_start_loc_device: torch.Tensor,
+ num_decode_tokens: int,
+ dcp_tot_seq_lens_device: torch.Tensor | None,
+ ) -> Dots3NoteDecodeMetadata:
+ del max_seq_len, query_start_loc_device, num_decode_tokens
+ query_len = int(query_start_loc_cpu[1] - query_start_loc_cpu[0])
+ return Dots3NoteDecodeMetadata(
+ block_table=block_table_tensor,
+ seq_lens=seq_lens_device,
+ dcp_tot_seq_lens=dcp_tot_seq_lens_device,
+ query_len=query_len,
+ )
+
+ def build(
+ self,
+ common_prefix_len: int,
+ common_attn_metadata: CommonAttentionMetadata,
+ fast_build: bool = False,
+ ) -> MLACommonMetadata:
+ metadata = super().build(
+ common_prefix_len,
+ common_attn_metadata,
+ fast_build=fast_build,
+ )
+ if metadata.prefill is None:
+ return metadata
+ if self.dcp_world_size > 1:
+ raise NotImplementedError(
+ "Dots3 NOTE SWA prefill does not support decode context parallelism"
+ )
+
+ reqs_start = metadata.num_decodes
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
+ assert seq_lens_cpu is not None
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[reqs_start:]
+ query_start_loc_cpu = query_start_loc_cpu - query_start_loc_cpu[0]
+ sliding_metadata = _build_sliding_window_metadata(
+ seq_lens_cpu=seq_lens_cpu[reqs_start:],
+ query_start_loc_cpu=query_start_loc_cpu,
+ sliding_window=self.sliding_window,
+ workspace=self.chunked_prefill_workspace,
+ workspace_size=self.chunked_prefill_workspace_size,
+ device=self.device,
+ )
+ metadata.prefill.chunked_context = None
+ metadata.prefill.sliding_window = sliding_metadata # type: ignore[attr-defined]
+ if metadata.num_decodes > 0 and metadata.num_prefills > 0:
+ query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
+ for chunk in sliding_metadata.chunks:
+ req_start, req_end = chunk.req_start, chunk.req_end
+ seq_lens = common_attn_metadata.seq_lens[
+ reqs_start + req_start : reqs_start + req_end
+ ]
+ query_lens = query_lens_cpu[req_start:req_end].to(
+ self.device, non_blocking=True
+ )
+ kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1)
+ chunk.cu_seq_lens_k.zero_()
+ torch.cumsum(
+ kv_lens,
+ dim=0,
+ out=chunk.cu_seq_lens_k[1:],
+ dtype=torch.int32,
+ )
+ chunk.starts = seq_lens - kv_lens
+ token_ids = torch.arange(
+ chunk.num_kv_tokens,
+ dtype=torch.int32,
+ device=self.device,
+ )
+ chunk.token_to_seq = torch.searchsorted(
+ chunk.cu_seq_lens_k[1:], token_ids, right=True
+ ).to(torch.int32)
+ chunk.token_to_seq.clamp_max_(req_end - req_start - 1)
+ assert metadata.prefill.prefill_backend is not None
+ metadata.prefill.prefill_backend.prepare_metadata(metadata.prefill)
+ return metadata
+
+
+class Dots3NoteTritonMLABackend(TritonMLABackend):
+ """Internal NOTE SWA specialization; not a user-selectable backend."""
+
+ @classmethod
+ def get_supported_head_sizes(cls) -> list[int]:
+ return [1088]
+
+ @classmethod
+ def supports_sliding_window(cls) -> bool:
+ return True
+
+ @staticmethod
+ def get_impl_cls() -> type["Dots3NoteTritonMLAImpl"]:
+ return Dots3NoteTritonMLAImpl
+
+ @staticmethod
+ def get_builder_cls() -> type[Dots3NoteMLAMetadataBuilder]:
+ return Dots3NoteMLAMetadataBuilder
+
+
+class Dots3NoteTritonMLAImpl(TritonMLAImpl):
+ def __init__(
+ self,
+ num_heads: int,
+ head_size: int,
+ scale: float,
+ num_kv_heads: int,
+ alibi_slopes: list[float] | None,
+ sliding_window: int | None,
+ kv_cache_dtype: str,
+ logits_soft_cap: float | None,
+ attn_type: str,
+ kv_sharing_target_layer_name: str | None,
+ **mla_args,
+ ) -> None:
+ assert sliding_window is not None
+ super().__init__(
+ num_heads=num_heads,
+ head_size=head_size,
+ scale=scale,
+ num_kv_heads=num_kv_heads,
+ alibi_slopes=alibi_slopes,
+ sliding_window=None,
+ kv_cache_dtype=kv_cache_dtype,
+ logits_soft_cap=logits_soft_cap,
+ attn_type=attn_type,
+ kv_sharing_target_layer_name=kv_sharing_target_layer_name,
+ **mla_args,
+ )
+ self.sliding_window = sliding_window
+
+ def _forward_swa_mqa(
+ self,
+ q: torch.Tensor,
+ kv_cache: torch.Tensor,
+ decode: Dots3NoteDecodeMetadata,
+ k_scale: torch.Tensor,
+ ) -> tuple[torch.Tensor, None]:
+ query_len = decode.query_len
+ num_tokens, num_heads, head_size = q.shape
+ assert num_tokens % query_len == 0
+ assert head_size == self.head_size
+ num_reqs = num_tokens // query_len
+ block_table = decode.block_table[:num_reqs]
+ seq_lens = decode.seq_lens[:num_reqs]
+ gather_len = (self.sliding_window + query_len - 1 + 7) // 8 * 8
+
+ workspace_specs = (
+ ((num_reqs, gather_len, head_size), q.dtype),
+ ((num_reqs, gather_len), torch.bool),
+ )
+ if is_workspace_manager_initialized():
+ kv_latent, kv_valid = current_workspace_manager().get_simultaneous(
+ *workspace_specs
+ )
+ else:
+ kv_latent = torch.empty(
+ workspace_specs[0][0], dtype=q.dtype, device=q.device
+ )
+ kv_valid = torch.empty(
+ workspace_specs[1][0], dtype=torch.bool, device=q.device
+ )
+
+ block_t = 8
+ block_d = 128
+ _gather_swa_kv_kernel[
+ (
+ num_reqs,
+ triton.cdiv(gather_len, block_t),
+ triton.cdiv(head_size, block_d),
+ )
+ ](
+ kv_cache,
+ block_table,
+ seq_lens,
+ kv_latent,
+ kv_valid,
+ k_scale,
+ kv_cache.stride(0),
+ kv_cache.stride(1),
+ kv_cache.stride(2),
+ block_table.stride(0),
+ kv_latent.stride(0),
+ kv_latent.stride(1),
+ kv_latent.stride(2),
+ kv_valid.stride(0),
+ kv_valid.stride(1),
+ GATHER_LEN=gather_len,
+ KV_DIM=head_size,
+ BLOCK_T=block_t,
+ BLOCK_D=block_d,
+ PAGE_SIZE=kv_cache.shape[1],
+ num_warps=4,
+ )
+
+ q = q.view(num_reqs, query_len, num_heads, head_size)
+ scores = torch.bmm(
+ q.reshape(num_reqs, query_len * num_heads, head_size),
+ kv_latent.transpose(1, 2),
+ ).view(num_reqs, query_len, num_heads, gather_len)
+ scores = scores.float()
+ scores.mul_(self.scale)
+ scores_by_head = scores.transpose(1, 2)
+ _apply_swa_score_mask_kernel[(num_reqs, num_heads, query_len)](
+ scores_by_head,
+ seq_lens,
+ kv_valid,
+ scores_by_head.stride(0),
+ scores_by_head.stride(1),
+ scores_by_head.stride(2),
+ scores_by_head.stride(3),
+ kv_valid.stride(0),
+ kv_valid.stride(1),
+ GATHER_LEN=gather_len,
+ WINDOW_SIZE=self.sliding_window,
+ QUERY_LEN=query_len,
+ BLOCK_T=triton.next_power_of_2(gather_len),
+ num_warps=4,
+ )
+ probs = torch.softmax(scores, dim=-1).to(q.dtype)
+ output = torch.bmm(
+ probs.reshape(num_reqs, query_len * num_heads, gather_len),
+ kv_latent[..., : self.kv_lora_rank],
+ )
+ return output.view(num_tokens, num_heads, self.kv_lora_rank), None
+
+ def forward_mha(
+ self,
+ q: torch.Tensor,
+ kv_c_normed: torch.Tensor,
+ k_pe: torch.Tensor,
+ kv_c_and_k_pe_cache: torch.Tensor,
+ attn_metadata: MLACommonMetadata,
+ k_scale: torch.Tensor,
+ output: torch.Tensor,
+ output_scale: torch.Tensor | None = None,
+ ) -> None:
+ prefill = attn_metadata.prefill
+ sliding = getattr(prefill, "sliding_window", None)
+ if prefill is None or sliding is None:
+ return super().forward_mha(
+ q,
+ kv_c_normed,
+ k_pe,
+ kv_c_and_k_pe_cache,
+ attn_metadata,
+ k_scale,
+ output,
+ output_scale,
+ )
+ assert output_scale is None
+ assert isinstance(prefill.prefill_backend, Dots3NoteFlashAttnPrefillBackend)
+ use_fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype()
+ output = output.view(-1, self.num_heads, self.v_head_dim)
+
+ for chunk in sliding.chunks:
+ toks = chunk.num_kv_tokens
+ workspace = sliding.workspace
+ block_table = prefill.block_table[chunk.req_start : chunk.req_end]
+ if is_quantized_kv_cache(self.kv_cache_dtype) and not use_fp8_prefill:
+ ops.gather_and_maybe_dequant_cache(
+ src_cache=kv_c_and_k_pe_cache,
+ dst=workspace,
+ block_table=block_table,
+ cu_seq_lens=chunk.cu_seq_lens_k,
+ token_to_seq=chunk.token_to_seq,
+ num_tokens=toks,
+ kv_cache_dtype=self.kv_cache_dtype,
+ scale=k_scale,
+ seq_starts=chunk.starts,
+ )
+ else:
+ ops.cp_gather_cache(
+ src_cache=kv_c_and_k_pe_cache,
+ dst=workspace,
+ block_table=block_table,
+ cu_seq_lens=chunk.cu_seq_lens_k,
+ batch_size=chunk.req_end - chunk.req_start,
+ seq_starts=chunk.starts,
+ )
+
+ # Slicing the latent part from the packed latent+RoPE cache keeps
+ # the full cache-entry row stride. FP8 linear kernels require
+ # contiguous rows and would otherwise read them at wrong offsets.
+ kv_c = workspace[:toks, : self.kv_lora_rank].contiguous()
+ weight_dtype = (
+ self.kv_b_proj.weight.dtype
+ if hasattr(self.kv_b_proj, "weight")
+ else self.kv_b_proj.params_dtype
+ )
+ if (
+ use_fp8_prefill or weight_dtype != current_platform.fp8_dtype()
+ ) and weight_dtype != torch.uint8:
+ kv_c = kv_c.to(weight_dtype)
+ k_pe_chunk = workspace[:toks, self.kv_lora_rank :].unsqueeze(1)
+ kv_nope = self.kv_b_proj(kv_c)[0].view(
+ -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
+ )
+ if use_fp8_prefill:
+ kv_nope = kv_nope.to(prefill.q_data_type)
+ k_pe_chunk = k_pe_chunk.to(prefill.q_data_type)
+ k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
+ k = self._concat_k_nope_k_pe(k_nope, k_pe_chunk)
+ chunk_output = prefill.prefill_backend.run_sliding_window(
+ q=q[chunk.query_start : chunk.query_end],
+ k=k,
+ v=v,
+ cu_seq_lens_q=chunk.cu_seq_lens_q,
+ cu_seq_lens_k=chunk.cu_seq_lens_k,
+ max_seq_len_q=chunk.max_seq_len_q,
+ max_seq_len_k=chunk.max_seq_len_k,
+ sliding_window=self.sliding_window,
+ )
+ output[chunk.query_start : chunk.query_end].copy_(
+ chunk_output[..., : self.v_head_dim]
+ )
+
+ def forward_mqa(
+ self,
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+ kv_c_and_k_pe_cache: torch.Tensor,
+ attn_metadata: MLACommonMetadata,
+ layer: AttentionLayer,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ assert kv_c_and_k_pe_cache.numel() > 0
+ assert attn_metadata.decode is not None
+ if isinstance(q, tuple):
+ q = torch.cat(q, dim=-1)
+
+ decode = attn_metadata.decode
+ assert isinstance(decode, Dots3NoteDecodeMetadata)
+ return self._forward_swa_mqa(
+ q,
+ kv_c_and_k_pe_cache,
+ decode,
+ layer._k_scale,
+ )
+
+
+class Dots3NotePaddedSparseBackend(FlashAttnMLASparseBackend):
+ """NOTE DSA backend for cache rows padded to the SWA latent width."""
+
+ @staticmethod
+ def get_name() -> str:
+ return "DOTS3_NOTE_PADDED_MLA_SPARSE"
+
+ @staticmethod
+ def get_impl_cls() -> type["Dots3NotePaddedSparseImpl"]:
+ return Dots3NotePaddedSparseImpl
+
+
+class Dots3NotePaddedSparseImpl(FlashAttnMLASparseImpl):
+ """Read top-k KV directly from uniformly padded cache rows."""
+
+ def _logical_cache(self, kv_cache: torch.Tensor) -> torch.Tensor:
+ assert kv_cache.shape[-1] >= self.head_size
+ return kv_cache[..., : self.head_size]
+
+ def do_kv_cache_update(
+ self,
+ kv_c_normed: torch.Tensor,
+ k_pe: torch.Tensor,
+ kv_cache: torch.Tensor,
+ slot_mapping: torch.Tensor,
+ kv_cache_dtype: str,
+ k_scale: torch.Tensor,
+ ) -> None:
+ super().do_kv_cache_update(
+ kv_c_normed,
+ k_pe,
+ self._logical_cache(kv_cache),
+ slot_mapping,
+ kv_cache_dtype,
+ k_scale,
+ )
+
+ def forward_mha( # type: ignore[override]
+ self,
+ q: torch.Tensor,
+ kv_c_normed: torch.Tensor,
+ k_pe: torch.Tensor,
+ kv_c_and_k_pe_cache: torch.Tensor,
+ attn_metadata: FlashAttnMLASparseMetadata,
+ k_scale: torch.Tensor,
+ output: torch.Tensor,
+ output_scale: torch.Tensor | None = None,
+ ) -> None:
+ super().forward_mha(
+ q,
+ kv_c_normed,
+ k_pe,
+ self._logical_cache(kv_c_and_k_pe_cache),
+ attn_metadata,
+ k_scale,
+ output,
+ output_scale,
+ )
+
+ def forward_mqa(
+ self,
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
+ kv_c_and_k_pe_cache: torch.Tensor,
+ attn_metadata: FlashAttnMLASparseMetadata,
+ layer: AttentionLayer,
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
+ if not isinstance(q, tuple):
+ raise NotImplementedError(
+ "Dots3NotePaddedSparseImpl expects split (q_nope, q_rope) input."
+ )
+ q_nope, q_rope = q
+ num_actual_toks = q_rope.shape[0]
+
+ assert self.topk_indices_buffer is not None
+ topk_indices = self.topk_indices_buffer[:num_actual_toks]
+ topk_indices, valid_counts = triton_convert_req_index_to_global_index(
+ attn_metadata.req_id_per_token[:num_actual_toks],
+ attn_metadata.block_table,
+ topk_indices,
+ BLOCK_SIZE=attn_metadata.block_size,
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
+ return_valid_counts=True,
+ )
+ physical_head_size = kv_c_and_k_pe_cache.shape[-1]
+ block_stride, token_stride, head_stride = kv_c_and_k_pe_cache.stride()
+ if (
+ token_stride != physical_head_size
+ or head_stride != 1
+ or block_stride % physical_head_size != 0
+ ):
+ raise RuntimeError(
+ "Dots3 NOTE padded DSA cache must contain contiguous token rows; "
+ f"got shape={tuple(kv_c_and_k_pe_cache.shape)} and "
+ f"stride={kv_c_and_k_pe_cache.stride()}"
+ )
+ block_row_stride = block_stride // physical_head_size
+ block_indices = torch.div(
+ topk_indices,
+ attn_metadata.block_size,
+ rounding_mode="floor",
+ )
+ block_indices.clamp_min_(0).mul_(block_row_stride - attn_metadata.block_size)
+ topk_indices.add_(block_indices)
+
+ num_cache_rows = (
+ kv_c_and_k_pe_cache.shape[0] - 1
+ ) * block_row_stride + attn_metadata.block_size
+ kv_cache = kv_c_and_k_pe_cache.as_strided(
+ (num_cache_rows, physical_head_size),
+ (physical_head_size, 1),
+ )
+ cu_seqlens_q = torch.arange(
+ num_actual_toks + 1,
+ dtype=torch.int32,
+ device=q_rope.device,
+ )
+ output = flash_attn_varlen_func(
+ q=q_rope,
+ k=kv_cache[:, self.kv_lora_rank : self.head_size].unsqueeze(1).unsqueeze(1),
+ v=kv_cache[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1),
+ q_v=q_nope,
+ max_seqlen_q=1,
+ cu_seqlens_q=cu_seqlens_q,
+ max_seqlen_k=topk_indices.shape[1],
+ seqused_k=valid_counts,
+ block_table=topk_indices,
+ softmax_scale=self.scale,
+ causal=True,
+ fa_version=3,
+ )
+ return output, None
diff --git a/vllm/models/dots3_note/nvidia/audio.py b/vllm/models/dots3_note/nvidia/audio.py
new file mode 100644
index 0000000000..26ba964470
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/audio.py
@@ -0,0 +1,305 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import json
+import math
+from functools import lru_cache
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from transformers import WhisperConfig
+from transformers.audio_utils import mel_filter_bank
+
+from .audio_encoder import DotsSpeechEncoder
+
+SAMPLE_RATE = 16000
+N_FFT = 400
+HOP_LENGTH = 160
+DEFAULT_CHUNK_LENGTH_S = 60
+DEFAULT_CONV_TEMPORAL_STRIDE = 8
+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
+
+
+class Dots3NoteAudioConfig:
+ def __init__(self, **kwargs):
+ self.encoder_type = kwargs.get("encoder_type", "dots")
+ self.whisper_config = kwargs.get("whisper_config", {})
+ self.whisper_adapter_in_dim = kwargs.get(
+ "whisper_adapter_in_dim", kwargs.get("adapter_in_dim", 1280)
+ )
+ self.whisper_adapter_out_dim = kwargs.get(
+ "whisper_adapter_out_dim", kwargs.get("adapter_out_dim", 2048)
+ )
+ self.sampling_rate = kwargs.get("sampling_rate", SAMPLE_RATE)
+ self.audio_comp_start = kwargs.get("audio_comp_start", "<|audio_comp_start|>")
+ self.audio_comp_span = kwargs.get("audio_comp_span", "<|audio_comp_pad|>")
+ self.audio_comp_end = kwargs.get("audio_comp_end", "<|audio_comp_end|>")
+ self.merge_factor = kwargs.get("merge_factor", DEFAULT_MERGE_FACTOR)
+ self.chunk_seconds = kwargs.get("chunk_seconds", DEFAULT_CHUNK_LENGTH_S)
+ self.use_conv2d_stem = kwargs.get("use_conv2d_stem", True)
+ self.use_rope = kwargs.get("use_rope", True)
+ self.use_rms_norm = kwargs.get("use_rms_norm", True)
+ self.use_causal = kwargs.get("use_causal", False)
+ self.downsample_hidden_size = kwargs.get("downsample_hidden_size", 480)
+ self.conv_chunksize = kwargs.get("conv_chunksize", 500)
+ self.conv_stem_gradient_checkpointing = kwargs.get(
+ "conv_stem_gradient_checkpointing", False
+ )
+ self.conv_bucket_step = kwargs.get("conv_bucket_step")
+ self.conv_bucket_max_elements = kwargs.get("conv_bucket_max_elements")
+ self.rope_parameters = kwargs.get(
+ "rope_parameters",
+ {
+ "partial_rotary_factor": 0.5,
+ "rope_theta": 10000.0,
+ "rope_type": "default",
+ },
+ )
+
+ @classmethod
+ def from_json(cls, path):
+ with open(path) as f:
+ return cls(**json.load(f))
+
+ @property
+ def conv_temporal_stride(self) -> int:
+ return 8 if self.use_conv2d_stem else 2
+
+ @property
+ def token_stride(self) -> int:
+ return HOP_LENGTH * self.conv_temporal_stride * self.merge_factor
+
+ @property
+ def chunk_samples(self) -> int:
+ return int(self.chunk_seconds * SAMPLE_RATE)
+
+ @property
+ def chunk_mel_frames(self) -> int:
+ return int(self.chunk_seconds * 100)
+
+
+def pad_or_trim(array, length=N_SAMPLES, axis=-1):
+ if array.shape[axis] > length:
+ array = array.index_select(
+ dim=axis, index=torch.arange(length, device=array.device)
+ )
+ if array.shape[axis] < length:
+ pad_widths = [(0, 0)] * array.ndim
+ pad_widths[axis] = (0, length - array.shape[axis])
+ array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
+ return array
+
+
+@lru_cache(maxsize=4)
+def _mel_filters(device, n_mels=128):
+ filters = mel_filter_bank(
+ num_frequency_bins=1 + N_FFT // 2,
+ num_mel_filters=n_mels,
+ min_frequency=0.0,
+ max_frequency=float(SAMPLE_RATE) / 2.0,
+ sampling_rate=SAMPLE_RATE,
+ norm="slaney",
+ mel_scale="slaney",
+ )
+ return torch.from_numpy(filters).T.contiguous().float().to(device)
+
+
+@lru_cache(maxsize=4)
+def _hann_window(device):
+ # Generate on CPU (default) then move, to preserve the exact reference
+ # window values. Direct on-device generation changes numerics slightly.
+ return torch.hann_window(N_FFT).to(device)
+
+
+def 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.max() - 8.0)
+ log_spec = (log_spec + 4.0) / 4.0
+ return log_spec
+
+
+def compute_audio_token_length(
+ num_samples,
+ *,
+ sample_rate: int = SAMPLE_RATE,
+ chunk_seconds: int = DEFAULT_CHUNK_LENGTH_S,
+ hop_length: int = HOP_LENGTH,
+ conv_temporal_stride: int = DEFAULT_CONV_TEMPORAL_STRIDE,
+ merge_factor: int = DEFAULT_MERGE_FACTOR,
+):
+ stride = hop_length * conv_temporal_stride * merge_factor
+ total = 0
+ time_step = 0
+ while time_step * sample_rate < num_samples:
+ chunk_len = min(
+ num_samples - time_step * sample_rate, chunk_seconds * sample_rate
+ )
+ total += math.ceil(chunk_len / stride)
+ time_step += chunk_seconds
+ return total
+
+
+class DotsEncoderWithMask(nn.Module):
+ def __init__(self, config: Dots3NoteAudioConfig):
+ super().__init__()
+ whisper_config = WhisperConfig(**config.whisper_config)
+ whisper_config.use_rope = config.use_rope
+ whisper_config.rope_parameters = config.rope_parameters
+ whisper_config.use_rms_norm = config.use_rms_norm
+ whisper_config.use_causal = config.use_causal
+ whisper_config.use_conv2d_stem = config.use_conv2d_stem
+ whisper_config.downsample_hidden_size = config.downsample_hidden_size
+ whisper_config.conv_chunksize = config.conv_chunksize
+ whisper_config.conv_stem_gradient_checkpointing = (
+ config.conv_stem_gradient_checkpointing
+ )
+ whisper_config.conv_bucket_step = config.conv_bucket_step
+ whisper_config.conv_bucket_max_elements = config.conv_bucket_max_elements
+
+ self.speech_encoder = DotsSpeechEncoder(whisper_config)
+ self.merge_factor = config.merge_factor
+ self.chunk_seconds = config.chunk_seconds
+ self.chunk_samples = config.chunk_samples
+ self.chunk_mel_frames = config.chunk_mel_frames
+ self.conv_temporal_stride = config.conv_temporal_stride
+
+ @property
+ def device(self):
+ return next(self.speech_encoder.parameters()).device
+
+ def _forward_speech_encoder(
+ self,
+ mel_features: torch.Tensor,
+ 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
+
+ def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor:
+ segments = []
+ time_step = 0
+ while time_step * SAMPLE_RATE < audio_waveform.shape[0]:
+ segments.append(
+ audio_waveform[
+ time_step * SAMPLE_RATE : (time_step + self.chunk_seconds)
+ * SAMPLE_RATE
+ ]
+ )
+ time_step += self.chunk_seconds
+
+ mel_features = []
+ token_lens = []
+ audio_sample_lens = []
+ for audio_segment in segments:
+ segment_length = audio_segment.shape[0]
+ token_len = (segment_length - 1) // (
+ HOP_LENGTH * self.conv_temporal_stride * self.merge_factor
+ ) + 1
+ pad_audio = pad_or_trim(audio_segment.flatten(), length=self.chunk_samples)
+ mel = log_mel_spectrogram(pad_audio)
+ assert mel.shape[1] == self.chunk_mel_frames
+ mel_features.append(mel)
+ token_lens.append(token_len)
+ audio_sample_lens.append(segment_length)
+
+ mel_features = torch.stack(mel_features, dim=0)
+ # Keep input_seq_lens on CPU: the conv2d bucket path reads it via
+ # ``.item()`` and CPU scalars avoid device->host syncs. The encoder's
+ # varlen path moves it to the device itself.
+ input_seq_lens = torch.tensor(token_lens, dtype=torch.long) * self.merge_factor
+ audio_embedding = self._forward_speech_encoder(
+ mel_features, input_seq_lens, audio_sample_lens
+ )
+
+ chunk_embeddings = []
+ for idx, token_len in enumerate(token_lens):
+ chunk_embeddings.append(
+ audio_embedding[idx, : token_len * self.merge_factor, :]
+ )
+ return torch.cat(chunk_embeddings, dim=0).unsqueeze(0)
+
+
+class AudioAdapter(nn.Module):
+ def __init__(self, in_dim, out_dim):
+ super().__init__()
+ self.proj = nn.Sequential(
+ nn.LayerNorm(in_dim),
+ nn.Linear(in_dim, out_dim),
+ nn.GELU(),
+ nn.Linear(out_dim, out_dim),
+ )
+
+ def forward(self, x):
+ return self.proj(x)
+
+
+class Dots3NoteAudioModel(nn.Module):
+ def __init__(self, config: Dots3NoteAudioConfig):
+ super().__init__()
+ self.config = config
+ if config.encoder_type != "dots":
+ raise ValueError("Dots3Note only supports encoder_type='dots'")
+ self.merge_factor = config.merge_factor
+ self.audio_adapter = AudioAdapter(
+ config.whisper_adapter_in_dim,
+ config.whisper_adapter_out_dim,
+ )
+
+ self.dots_encoder = DotsEncoderWithMask(config)
+
+ @property
+ def device(self):
+ return self.dots_encoder.device
+
+ def _merge_embeddings(self, embedding: torch.Tensor) -> torch.Tensor:
+ if self.merge_factor <= 1:
+ return embedding
+ return embedding.reshape(
+ embedding.shape[0],
+ embedding.shape[1] // self.merge_factor,
+ embedding.shape[2] * self.merge_factor,
+ )
+
+ def _encode_single_audio(self, audio_waveform):
+ embedding = self.dots_encoder.encode_waveform(audio_waveform)
+ embedding = self._merge_embeddings(embedding)
+ embedding = self.audio_adapter(embedding)
+ return embedding.squeeze(0)
+
+ def _split_waveforms(self, audio_inputs, lengths):
+ # Convert lengths to CPU ints once so the per-audio slicing below stays
+ # on the host and does not trigger a device->host sync per audio.
+ if isinstance(audio_inputs, list):
+ return audio_inputs
+ lengths_list = lengths.tolist()
+ waveforms = []
+ audio_start = 0
+ for length in lengths_list:
+ waveforms.append(audio_inputs[audio_start : audio_start + length])
+ audio_start += length
+ return waveforms
+
+ 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
diff --git a/vllm/models/dots3_note/nvidia/audio_encoder.py b/vllm/models/dots3_note/nvidia/audio_encoder.py
new file mode 100644
index 0000000000..7faf01987c
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/audio_encoder.py
@@ -0,0 +1,736 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+"""Dots-path speech encoder for inference only (single GPU).
+
+Ported from cybertron_alm ``dots_audio_encoder/modeling_whisper.py``.
+Upstream ``WhisperEncoder`` is exposed as :class:`DotsSpeechEncoder`.
+"""
+
+import math
+from typing import Any, ClassVar
+
+import torch
+from torch import nn
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import BaseModelOutput
+from transformers.modeling_utils import PreTrainedModel
+from transformers.models.whisper.configuration_whisper import WhisperConfig
+from transformers.utils import logging
+
+from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported
+
+_FLASH_ATTN_VERSION = 3 if is_fa_version_supported(3) else 2
+
+logger = logging.get_logger(__name__)
+
+__all__ = ["DotsSpeechEncoder"]
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(dim))
+ self.eps = eps
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ var = x.pow(2).mean(-1, keepdim=True)
+ x = x * torch.rsqrt(var + self.eps)
+ return self.weight * x
+
+
+def swiglu(x: torch.Tensor) -> torch.Tensor:
+ x1, x2 = x.chunk(2, dim=-1)
+ return torch.nn.functional.silu(x1) * x2
+
+
+class RotaryEmbedding(nn.Module):
+ def __init__(
+ self,
+ head_dim: int,
+ rope_parameters: dict,
+ base_seq_len: int = 0,
+ ):
+ super().__init__()
+ self.partial_rotary_factor = float(
+ rope_parameters.get("partial_rotary_factor", 1.0)
+ )
+ self.rope_theta = float(rope_parameters.get("rope_theta", 10000.0))
+ self.rope_type = rope_parameters.get("rope_type", "default")
+ rotary_dim = int(head_dim * self.partial_rotary_factor)
+ # keep even rotary dimension to align cos/sin pairs
+ self.rotary_dim = (rotary_dim // 2) * 2
+ self.attention_scaling = 1.0
+ if self.rope_type != "default":
+ logger.warning(
+ "RoPE type %s is not implemented in WhisperLv3; using default",
+ self.rope_type,
+ )
+ if self.rotary_dim > 0:
+ inv_freq = 1.0 / (
+ self.rope_theta
+ ** (
+ torch.arange(0, self.rotary_dim, 2, dtype=torch.float)
+ / max(self.rotary_dim, 1)
+ )
+ )
+ else:
+ inv_freq = torch.tensor([])
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self.base_seq_len = max(0, int(base_seq_len))
+ self._cache: (
+ tuple[int, torch.dtype, torch.device, torch.Tensor, torch.Tensor] | None
+ ) = None
+
+ @torch.no_grad()
+ def get_cos_sin(
+ self, position_ids: torch.Tensor, dtype: torch.dtype, device: torch.device
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
+ if self.rotary_dim == 0:
+ return None, None
+ seq_len = position_ids.shape[-1]
+ if position_ids.shape[0] == 1 and self._cache is not None:
+ cached_seq_len, cached_dtype, cached_device, cached_cos, cached_sin = (
+ self._cache
+ )
+ if (
+ cached_seq_len >= seq_len
+ and cached_dtype == dtype
+ and cached_device == device
+ ):
+ return cached_cos[:, :seq_len, :], cached_sin[:, :seq_len, :]
+
+ if position_ids.shape[0] == 1:
+ cache_seq_len = max(seq_len, self.base_seq_len)
+ position_ids = torch.arange(cache_seq_len, device=device)[None, :]
+ # [1, D/2, 1]
+ if self.inv_freq.dtype != torch.float32:
+ inv_freq = 1.0 / (
+ self.rope_theta
+ ** (
+ torch.arange(
+ 0, self.rotary_dim, 2, dtype=torch.float, device=device
+ )
+ / max(self.rotary_dim, 1)
+ )
+ )
+ else:
+ inv_freq = self.inv_freq.to(device=device)
+ inv_freq_expanded = inv_freq[None, :, None]
+ # [B, 1, T]
+ position_ids_expanded = position_ids[:, None, :].float()
+ # Force float32 since bfloat16 loses precision on long contexts
+ # See https://github.com/huggingface/transformers/pull/29285
+ device_type = device.type if isinstance(device, torch.device) else "cpu"
+ with torch.autocast(device_type=device_type, enabled=False):
+ freqs = (
+ inv_freq_expanded.float() @ position_ids_expanded.float()
+ ).transpose(1, 2)
+ emb = torch.cat((freqs, freqs), dim=-1)
+ cos = emb.cos() * self.attention_scaling
+ sin = emb.sin() * self.attention_scaling
+ cos = cos.to(dtype)
+ sin = sin.to(dtype)
+ if position_ids.shape[0] == 1:
+ self._cache = (position_ids.shape[-1], dtype, device, cos, sin)
+ return cos[:, :seq_len, :], sin[:, :seq_len, :]
+ return cos, sin
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(
+ q: torch.Tensor,
+ k: torch.Tensor,
+ cos: torch.Tensor | None,
+ sin: torch.Tensor | None,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ if cos is None or sin is None:
+ return q, k
+
+ if q.dim() == 4:
+ # [B, T, H, D]
+ rotary_dim = cos.shape[-1]
+ cos = cos.unsqueeze(2) # [B, T, 1, rotary_dim]
+ sin = sin.unsqueeze(2)
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+ return torch.cat([q_embed, q_pass], dim=-1), torch.cat(
+ [k_embed, k_pass], dim=-1
+ )
+
+ if q.dim() == 3:
+ # [tokens, H, D] (varlen flattened)
+ rotary_dim = cos.shape[-1]
+ cos = cos.unsqueeze(1) # [tokens, 1, rotary_dim]
+ sin = sin.unsqueeze(1)
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
+ return torch.cat([q_embed, q_pass], dim=-1), torch.cat(
+ [k_embed, k_pass], dim=-1
+ )
+
+ return q, k
+
+
+class WhisperPositionalEmbedding(nn.Embedding):
+ def __init__(
+ self, num_positions: int, embedding_dim: int, padding_idx: int | None = None
+ ):
+ super().__init__(num_positions, embedding_dim)
+
+ def forward(self, input_ids, past_key_values_length=0):
+ return self.weight[
+ past_key_values_length : past_key_values_length + input_ids.shape[1]
+ ]
+
+
+class WhisperAttention(nn.Module):
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
+
+ def __init__(
+ self,
+ embed_dim: int,
+ num_heads: int,
+ dropout: float = 0.0,
+ is_decoder: bool = False,
+ bias: bool = True,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.num_heads = num_heads
+ self.dropout = dropout
+ self.head_dim = embed_dim // num_heads
+
+ if (self.head_dim * num_heads) != self.embed_dim:
+ raise ValueError(
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
+ f" and `num_heads`: {num_heads})."
+ )
+ self.scaling = self.head_dim**-0.5
+ self.is_decoder = is_decoder
+
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+ def forward_flash_attn(
+ self,
+ hidden_states,
+ cu_seqlens_q=None,
+ cu_seqlens_kv=None,
+ max_seqlen_q=None,
+ max_seqlen_kv=None,
+ output_attentions=False,
+ rotary_cos=None,
+ rotary_sin=None,
+ ):
+ """Dense eager attention with SGLang FA3 for packed variable-length input."""
+ query_states = self.q_proj(hidden_states)
+ key_states = self.k_proj(hidden_states)
+ value_states = self.v_proj(hidden_states)
+ if cu_seqlens_q is None:
+ bsz, tgt_len, _ = hidden_states.size()
+ query_states = query_states.view(
+ bsz, tgt_len, self.num_heads, self.head_dim
+ )
+ key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
+ value_states = value_states.view(
+ bsz, tgt_len, self.num_heads, self.head_dim
+ )
+ if rotary_cos is not None and rotary_sin is not None:
+ cos = rotary_cos[:, :tgt_len, :]
+ sin = rotary_sin[:, :tgt_len, :]
+ query_states, key_states = apply_rotary_pos_emb(
+ query_states, key_states, cos, sin
+ )
+ attn_output, attn_probs = self._eager_attention(
+ query_states, key_states, value_states, output_attentions
+ )
+ attn_output = attn_output.view(bsz, tgt_len, self.embed_dim)
+ else:
+ query_states = query_states.view(-1, self.num_heads, self.head_dim)
+ key_states = key_states.view(-1, self.num_heads, self.head_dim)
+ value_states = value_states.view(-1, self.num_heads, self.head_dim)
+ if rotary_cos is not None and rotary_sin is not None:
+ query_states, key_states = apply_rotary_pos_emb(
+ query_states, key_states, rotary_cos, rotary_sin
+ )
+ result = flash_attn_varlen_func(
+ query_states,
+ key_states,
+ value_states,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_k=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_k=max_seqlen_kv,
+ softmax_scale=self.scaling,
+ causal=self.is_decoder,
+ return_softmax_lse=output_attentions,
+ fa_version=_FLASH_ATTN_VERSION,
+ )
+ if isinstance(result, tuple):
+ attn_output = result[0]
+ attn_probs = result[1] if output_attentions else None
+ else:
+ attn_output = result
+ attn_probs = None
+ attn_output = attn_output.view(-1, self.embed_dim)
+
+ return self.out_proj(attn_output), attn_probs
+
+ def _eager_attention(self, query, key, value, output_attentions):
+ query = query.transpose(1, 2)
+ key = key.transpose(1, 2)
+ value = value.transpose(1, 2)
+ scores = torch.matmul(query, key.transpose(-2, -1)) * self.scaling
+ if self.is_decoder:
+ seq_len = scores.shape[-1]
+ causal_mask = torch.ones(
+ seq_len, seq_len, dtype=torch.bool, device=scores.device
+ ).triu(1)
+ scores.masked_fill_(causal_mask, torch.finfo(scores.dtype).min)
+ probabilities = torch.softmax(scores, dim=-1, dtype=torch.float32).to(
+ query.dtype
+ )
+ output = torch.matmul(probabilities, value).transpose(1, 2)
+ return output, probabilities if output_attentions else None
+
+
+# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Whisper
+class WhisperEncoderLayer(nn.Module):
+ def __init__(self, config: WhisperConfig):
+ super().__init__()
+ self.embed_dim = config.d_model
+ use_causal = getattr(config, "use_causal", False)
+ self.self_attn = WhisperAttention(
+ embed_dim=self.embed_dim,
+ num_heads=config.encoder_attention_heads,
+ dropout=config.attention_dropout,
+ is_decoder=use_causal, # enable causal attention when use_causal=True
+ )
+ self.use_causal = use_causal
+ norm_cls = RMSNorm if getattr(config, "use_rms_norm", False) else nn.LayerNorm
+ self.self_attn_layer_norm = norm_cls(self.embed_dim)
+ self.dropout = config.dropout
+ self.activation_fn = (
+ ACT2FN[config.activation_function]
+ if config.activation_function != "swiglu"
+ else swiglu
+ )
+ self.use_swiglu = config.activation_function == "swiglu"
+ self.activation_dropout = config.activation_dropout
+ ffn_dim = config.encoder_ffn_dim
+ fc1_out = ffn_dim * 2 if self.use_swiglu else ffn_dim
+ self.fc1 = nn.Linear(self.embed_dim, fc1_out)
+ self.fc2 = nn.Linear(ffn_dim, self.embed_dim)
+ self.final_layer_norm = norm_cls(self.embed_dim)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens_q: torch.Tensor = None,
+ cu_seqlens_kv: torch.Tensor = None,
+ max_seqlen_q: int | None = None,
+ max_seqlen_kv: int | None = None,
+ output_attentions: bool = False,
+ rotary_cos: torch.Tensor | None = None,
+ rotary_sin: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ """
+ Args:
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
+ attention_mask (`torch.FloatTensor`): attention mask of size
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+ layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
+ `(encoder_attention_heads,)`.
+ output_attentions (`bool`, *optional*):
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
+ returned tensors for more detail.
+ """
+ residual = hidden_states
+ hidden_states = self.self_attn_layer_norm(hidden_states)
+
+ hidden_states, attn_weights = self.self_attn.forward_flash_attn(
+ hidden_states=hidden_states,
+ cu_seqlens_q=cu_seqlens_q,
+ cu_seqlens_kv=cu_seqlens_kv,
+ max_seqlen_q=max_seqlen_q,
+ max_seqlen_kv=max_seqlen_kv,
+ output_attentions=output_attentions,
+ rotary_cos=rotary_cos,
+ rotary_sin=rotary_sin,
+ )
+ hidden_states = nn.functional.dropout(
+ hidden_states, p=self.dropout, training=self.training
+ )
+ hidden_states = residual + hidden_states
+
+ residual = hidden_states
+ hidden_states = self.final_layer_norm(hidden_states)
+ if self.use_swiglu:
+ hidden_states = swiglu(self.fc1(hidden_states))
+ else:
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
+ hidden_states = nn.functional.dropout(
+ hidden_states, p=self.activation_dropout, training=self.training
+ )
+ hidden_states = self.fc2(hidden_states)
+ hidden_states = nn.functional.dropout(
+ hidden_states, p=self.dropout, training=self.training
+ )
+ hidden_states = residual + hidden_states
+
+ if hidden_states.dtype == torch.float16 and (
+ torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
+ ):
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
+ hidden_states = torch.clamp(
+ hidden_states, min=-clamp_value, max=clamp_value
+ )
+
+ outputs: tuple[Any, ...] = (hidden_states,)
+
+ if output_attentions:
+ outputs += (attn_weights,)
+
+ return outputs
+
+
+class DotsSpeechPreTrainedModel(PreTrainedModel):
+ config_class = WhisperConfig
+ base_model_prefix = "model"
+ main_input_name = "input_features"
+ supports_gradient_checkpointing = False
+ _no_split_modules: ClassVar[list[str]] = ["WhisperEncoderLayer"]
+
+ def _init_weights(self, module):
+ std = self.config.init_std
+ if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d)):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.bias is not None:
+ module.bias.data.zero_()
+ elif isinstance(module, nn.Embedding):
+ module.weight.data.normal_(mean=0.0, std=std)
+ if module.padding_idx is not None:
+ module.weight.data[module.padding_idx].zero_()
+
+
+class DotsSpeechEncoder(DotsSpeechPreTrainedModel):
+ """
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
+ [`WhisperEncoderLayer`].
+
+ Args:
+ config: WhisperConfig
+ """
+
+ def __init__(self, config: WhisperConfig):
+ super().__init__(config)
+ self.dropout = config.dropout
+
+ embed_dim = config.d_model
+ self.use_latent_input = getattr(config, "use_latent_input", False)
+ self.use_causal = getattr(config, "use_causal", False)
+ self.use_conv2d_stem = getattr(config, "use_conv2d_stem", False)
+ latent_dim = getattr(config, "latent_dim", None)
+ if self.use_latent_input and latent_dim is None:
+ raise ValueError(
+ "DotsSpeechEncoder: use_latent_input=True requires config.latent_dim"
+ )
+ if self.use_conv2d_stem and self.use_latent_input:
+ raise ValueError(
+ "DotsSpeechEncoder: use_conv2d_stem and use_latent_input are mutually exclusive"
+ )
+ self.num_mel_bins = (
+ latent_dim if self.use_latent_input and latent_dim is not None else 128
+ )
+ self.padding_idx = config.pad_token_id
+ self.max_source_positions = config.max_source_positions
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
+
+ if self.use_conv2d_stem:
+ # Conv2D stem: 3 layers of stride-2 for 8x downsampling
+ dhs = getattr(config, "downsample_hidden_size", 480)
+ # Causal: keep freq padding, remove time padding (handled by pad(14,0) in forward)
+ conv_padding = (1, 0) if self.use_causal else 1
+ self.conv2d1 = nn.Conv2d(
+ 1, dhs, kernel_size=3, stride=2, padding=conv_padding
+ )
+ self.conv2d2 = nn.Conv2d(
+ dhs, dhs, kernel_size=3, stride=2, padding=conv_padding
+ )
+ self.conv2d3 = nn.Conv2d(
+ dhs, dhs, kernel_size=3, stride=2, padding=conv_padding
+ )
+ # After 3x stride-2 on freq=128: 128→64→32→16; linear projects dhs*16 → embed_dim
+ freq_after = self.num_mel_bins
+ for _ in range(3):
+ freq_after = (freq_after + 1) // 2
+ self.conv_out = nn.Linear(dhs * freq_after, embed_dim, bias=False)
+ self.conv1 = None
+ self.conv2 = None
+ elif self.use_latent_input:
+ # Latent path keeps length (stride=1) and applies GLU after conv1
+ self.conv1 = nn.Conv1d(
+ self.num_mel_bins, embed_dim * 2, kernel_size=3, stride=1, padding=1
+ )
+ self.conv2 = nn.Conv1d(
+ embed_dim, embed_dim, kernel_size=3, stride=1, padding=1
+ )
+ else:
+ self.conv1 = nn.Conv1d(
+ self.num_mel_bins, embed_dim, kernel_size=3, padding=1
+ )
+ self.conv2 = nn.Conv1d(
+ embed_dim, embed_dim, kernel_size=3, stride=2, padding=1
+ )
+
+ self.use_rope = getattr(config, "use_rope", False)
+ rope_parameters = getattr(config, "rope_parameters", {}) or {}
+ self.rope_parameters = rope_parameters
+ if self.use_rope:
+ head_dim = embed_dim // config.encoder_attention_heads
+ self.rotary_embedding = RotaryEmbedding(
+ head_dim,
+ rope_parameters,
+ base_seq_len=self.max_source_positions,
+ )
+ self.embed_positions = None
+ else:
+ self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
+
+ self.layers = nn.ModuleList(
+ [WhisperEncoderLayer(config) for _ in range(config.encoder_layers)]
+ )
+ norm_cls = RMSNorm if getattr(config, "use_rms_norm", False) else nn.LayerNorm
+ self.layer_norm = norm_cls(config.d_model)
+
+ self.post_init()
+
+ def get_input_embeddings(self) -> nn.Module:
+ if self.use_conv2d_stem:
+ return self.conv2d1
+ return self.conv1
+
+ def set_input_embeddings(self, value: nn.Module):
+ self.conv1 = value
+
+ @staticmethod
+ def _temporal_mask(feat, valid_lens):
+ """Zero out temporal positions >= valid_lens. feat: [B,C,F,T], valid_lens: [B] tensor."""
+ T = feat.shape[-1]
+ mask = torch.arange(T, device=feat.device)[None, :] < valid_lens[:, None]
+ return mask[:, None, None, :]
+
+ def _conv2d_stem_one_chunk(self, chunk, chunk_valid_mel_lens=None):
+ """Run 3x Conv2d(stride=2) + GELU with per-layer masking."""
+ # Causal: left-pad time by 14 = 2 + 4 + 8 (receptive field of the 3-layer
+ # stride-2 stack mapped back to input time). Combined with conv padding=(1,0),
+ # this makes the whole conv2d stem strictly causal.
+ if self.use_causal:
+ chunk = nn.functional.pad(chunk, (14, 0))
+ if chunk_valid_mel_lens is not None:
+ chunk_valid_mel_lens = chunk_valid_mel_lens + 14
+ # Step 0: mask mel — silence_mel(-1.5) → 0
+ if chunk_valid_mel_lens is not None:
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
+ chunk = nn.functional.gelu(self.conv2d1(chunk))
+ # Step 1: mask conv1 output — gelu(conv(0)+bias) → 0
+ if chunk_valid_mel_lens is not None:
+ chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
+ chunk = nn.functional.gelu(self.conv2d2(chunk))
+ # Step 2: mask conv2 output — gelu(bias) → 0
+ if chunk_valid_mel_lens is not None:
+ chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
+ chunk = nn.functional.gelu(self.conv2d3(chunk))
+ # Step 3: mask conv3 output — gelu(bias) → 0
+ if chunk_valid_mel_lens is not None:
+ chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
+ return chunk
+
+ def _forward_conv2d_stem(
+ self, input_features, input_seq_lens=None, audio_sample_lens=None
+ ):
+ """Conv2D stem: 3x Conv2d(stride=2) → 8x downsample. Returns [B, T/8, embed_dim]."""
+ x = input_features.unsqueeze(1) # [B, 1, 128, T]
+
+ if audio_sample_lens is not None:
+ hop_length = 160
+ if not isinstance(audio_sample_lens, torch.Tensor):
+ audio_sample_lens = torch.tensor(audio_sample_lens, device=x.device)
+ valid_mel_lens = audio_sample_lens.to(x.device) // hop_length
+ else:
+ valid_mel_lens = None
+ x = self._conv2d_stem_one_chunk(x, valid_mel_lens)
+ # [B, dhs, F_out, T/8] → [B, T/8, dhs*freq_after] → [B, T/8, embed_dim]
+ batch, channels, frequency, time = x.shape
+ x = x.permute(0, 3, 1, 2).reshape(batch, time, channels * frequency)
+ inputs_embeds = self.conv_out(x)
+
+ return inputs_embeds
+
+ def _forward_conv1d_stem(self, input_features):
+ """Standard Conv1d stem: conv1 + conv2(stride=2), 2x downsample. Returns [B, T/2, embed_dim]."""
+ assert self.conv1 is not None and self.conv2 is not None
+ if self.use_causal:
+ x = nn.functional.pad(input_features, (2, 0))
+ x = nn.functional.gelu(self.conv1(x))
+ x = nn.functional.pad(x, (2, 0))
+ x = nn.functional.gelu(self.conv2(x))
+ else:
+ x = nn.functional.gelu(self.conv1(input_features))
+ x = nn.functional.gelu(self.conv2(x))
+ return x.permute(0, 2, 1)
+
+ def _forward_latent_stem(self, input_features):
+ """Latent input stem: conv1+GLU + conv2, stride=1, no downsample. Returns [B, T, embed_dim]."""
+ assert self.conv1 is not None and self.conv2 is not None
+ if self.use_causal:
+ x = nn.functional.pad(input_features, (2, 0))
+ x = nn.functional.glu(self.conv1(x), dim=1)
+ x = nn.functional.pad(x, (2, 0))
+ x = nn.functional.gelu(self.conv2(x))
+ else:
+ x = nn.functional.glu(self.conv1(input_features), dim=1)
+ x = nn.functional.gelu(self.conv2(x))
+ return x.permute(0, 2, 1)
+
+ def forward(
+ self,
+ input_features,
+ input_seq_lens=None,
+ audio_sample_lens=None,
+ output_attentions=None,
+ output_hidden_states=None,
+ return_dict=None,
+ ):
+ """Mel `[B, n_mels, T]` → encoder hidden states. Inference-only."""
+ output_attentions = (
+ output_attentions
+ if output_attentions is not None
+ else self.config.output_attentions
+ )
+ output_hidden_states = (
+ output_hidden_states
+ if output_hidden_states is not None
+ else self.config.output_hidden_states
+ )
+ return_dict = (
+ return_dict if return_dict is not None else self.config.use_return_dict
+ )
+ if self.use_conv2d_stem:
+ inputs_embeds = self._forward_conv2d_stem(
+ input_features, input_seq_lens, audio_sample_lens
+ )
+ elif self.use_latent_input:
+ inputs_embeds = self._forward_latent_stem(input_features)
+ else:
+ inputs_embeds = self._forward_conv1d_stem(input_features)
+ rotary_cos = None
+ rotary_sin = None
+ position_ids = torch.arange(
+ inputs_embeds.shape[1], device=inputs_embeds.device
+ )[None, :]
+ if self.use_rope:
+ rotary_cos, rotary_sin = self.rotary_embedding.get_cos_sin(
+ position_ids, inputs_embeds.dtype, inputs_embeds.device
+ )
+ hidden_states = inputs_embeds
+ else:
+ assert self.embed_positions is not None
+ embed_pos = self.embed_positions.weight[: inputs_embeds.shape[1]]
+ hidden_states = inputs_embeds + embed_pos
+
+ encoder_states: tuple[torch.Tensor, ...] | None = (
+ () if output_hidden_states else None
+ )
+ all_attentions: tuple[torch.Tensor, ...] | None = (
+ () if output_attentions else None
+ )
+
+ if input_seq_lens is not None:
+ # Build varlen metadata and pack valid tokens without per-sample cat loops.
+ # Read max on whatever device the caller provided: when it is a CPU
+ # tensor this avoids a device->host sync (one per forward).
+ max_seqlen_q = int(input_seq_lens.max().item())
+ input_seq_lens = input_seq_lens.to(
+ device=hidden_states.device, dtype=torch.long
+ )
+ B, S, D = hidden_states.shape
+ max_seqlen_kv = max_seqlen_q
+ cu_seqlens_q = torch.nn.functional.pad(
+ input_seq_lens.cumsum(0, dtype=torch.int32), (1, 0)
+ )
+ cu_seqlens_kv = cu_seqlens_q
+
+ token_positions = torch.arange(S, device=hidden_states.device)[None, :]
+ valid_token_mask = token_positions < input_seq_lens[:, None]
+ hidden_states = hidden_states[valid_token_mask]
+ if rotary_cos is not None and rotary_sin is not None:
+ packed_positions = token_positions.expand(B, S)[valid_token_mask]
+ rotary_cos = rotary_cos.squeeze(0).index_select(0, packed_positions)
+ rotary_sin = rotary_sin.squeeze(0).index_select(0, packed_positions)
+ else:
+ cu_seqlens_q = None
+ cu_seqlens_kv = None
+ max_seqlen_q = None
+ max_seqlen_kv = None
+
+ for encoder_layer in self.layers:
+ if output_hidden_states:
+ assert encoder_states is not None
+ encoder_states = encoder_states + (hidden_states,)
+ layer_outputs = encoder_layer(
+ hidden_states,
+ cu_seqlens_q,
+ cu_seqlens_kv,
+ max_seqlen_q,
+ max_seqlen_kv,
+ output_attentions=output_attentions,
+ rotary_cos=rotary_cos,
+ rotary_sin=rotary_sin,
+ )
+ hidden_states = layer_outputs[0]
+ if output_attentions:
+ assert all_attentions is not None
+ all_attentions = all_attentions + (layer_outputs[1],)
+
+ if input_seq_lens is not None:
+ # Recover packed varlen output directly on device.
+ recover_positions = torch.arange(max_seqlen_q, device=hidden_states.device)[
+ None, :
+ ]
+ recover_mask = recover_positions < input_seq_lens[:, None]
+ recovered = hidden_states.new_zeros(B, max_seqlen_q, D)
+ recovered[recover_mask] = hidden_states
+ hidden_states = recovered
+
+ hidden_states = self.layer_norm(hidden_states)
+ if output_hidden_states:
+ assert encoder_states is not None
+ encoder_states = encoder_states + (hidden_states,)
+
+ if not return_dict:
+ return tuple(
+ v
+ for v in [hidden_states, encoder_states, all_attentions]
+ if v is not None
+ )
+ return BaseModelOutput(
+ last_hidden_state=hidden_states,
+ hidden_states=encoder_states,
+ attentions=all_attentions,
+ )
diff --git a/vllm/models/dots3_note/nvidia/model.py b/vllm/models/dots3_note/nvidia/model.py
new file mode 100644
index 0000000000..83e675c469
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/model.py
@@ -0,0 +1,684 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA implementation of the Dots3Note language model."""
+
+import copy
+from collections.abc import Iterable, Iterator
+from dataclasses import replace
+
+import torch
+from torch import nn
+
+import vllm.envs as envs
+from vllm.config import VllmConfig
+from vllm.distributed import (
+ get_pp_group,
+ get_tensor_model_parallel_rank,
+ get_tensor_model_parallel_world_size,
+ tensor_model_parallel_all_gather,
+ tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.attention import MLAAttention
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import (
+ ColumnParallelLinear,
+ DCPGroupColumnParallelLinear,
+ ReplicatedLinear,
+ RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+ GroupShape,
+ scaled_dequantize,
+)
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.models.deepseek_v2 import (
+ DeepSeekV2FusedQkvAProjLinear,
+ DeepseekV2MLAAttention,
+ DeepseekV2MLP,
+ DeepseekV2MoE,
+)
+from vllm.model_executor.models.utils import (
+ PPMissingLayer,
+ make_empty_intermediate_tensors_factory,
+ make_layers,
+ sequence_parallel_chunk,
+)
+from vllm.models.deepseek_v32.nvidia.model import (
+ DeepseekV32DecoderLayer,
+ DeepseekV32ForCausalLM,
+ DeepseekV32Model,
+)
+from vllm.platforms import current_platform
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
+
+from .attention import (
+ Dots3NoteFlashAttnPrefillBackend,
+ Dots3NotePaddedSparseBackend,
+ Dots3NoteTritonMLABackend,
+)
+
+
+def _padded_mlp_size(
+ intermediate_size: int,
+ quant_config: QuantizationConfig | None,
+ is_sequence_parallel: bool = False,
+) -> int:
+ block_size = getattr(quant_config, "weight_block_size", None)
+ if block_size is None or is_sequence_parallel:
+ return intermediate_size
+ tp_size = get_tensor_model_parallel_world_size()
+ blocks = (intermediate_size + block_size[0] - 1) // block_size[0]
+ return ((blocks + tp_size - 1) // tp_size) * block_size[0] * tp_size
+
+
+class Dots3NoteMoE(DeepseekV2MoE):
+ """DeepSeek MoE with NOTE's block-padded shared expert."""
+
+ def __init__(
+ self,
+ config,
+ parallel_config,
+ quant_config: QuantizationConfig | None = None,
+ reduce_results: bool = True,
+ prefix: str = "",
+ apply_routed_scale_to_output: bool = False,
+ ) -> None:
+ num_shared_experts = config.n_shared_experts
+ routed_config = copy.copy(config)
+ object.__setattr__(routed_config, "n_shared_experts", None)
+ super().__init__(
+ config=routed_config,
+ parallel_config=parallel_config,
+ quant_config=quant_config,
+ reduce_results=False,
+ prefix=prefix,
+ apply_routed_scale_to_output=apply_routed_scale_to_output,
+ )
+ self.n_shared_experts = num_shared_experts
+ self.reduce_results = reduce_results
+ self.shared_experts = DeepseekV2MLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=_padded_mlp_size(
+ config.moe_intermediate_size * num_shared_experts,
+ quant_config,
+ self.is_sequence_parallel,
+ ),
+ hidden_act=config.hidden_act,
+ quant_config=quant_config,
+ is_sequence_parallel=self.is_sequence_parallel,
+ reduce_results=False,
+ prefix=f"{prefix}.shared_experts",
+ )
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ already_sequence_parallel: bool = False,
+ ) -> torch.Tensor:
+ num_tokens = hidden_states.shape[0]
+ gather_output = self.is_sequence_parallel and not already_sequence_parallel
+ if gather_output:
+ hidden_states = sequence_parallel_chunk(hidden_states)
+ assert self.shared_experts is not None
+ output = super().forward(
+ hidden_states, already_sequence_parallel=True
+ ) + self.shared_experts(hidden_states)
+ if gather_output:
+ return tensor_model_parallel_all_gather(output, 0)[:num_tokens]
+ if self.reduce_results:
+ return tensor_model_parallel_all_reduce(output)
+ return output
+
+
+def _forward_note_mla(
+ attention,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ *,
+ g_proj: nn.Module,
+ k_rope_only_layernorm: nn.Module,
+ attention_gate_type: str,
+ q_lora_scale: float,
+ kv_lora_scale: float,
+ llama_4_scaling: torch.Tensor | None = None,
+) -> torch.Tensor:
+ qkv_lora = attention.fused_qkv_a_proj(hidden_states)[0]
+ q_c, kv_lora = qkv_lora.split(
+ [
+ attention.q_lora_rank,
+ attention.kv_lora_rank + attention.qk_rope_head_dim,
+ ],
+ dim=-1,
+ )
+ q_c = attention.q_a_layernorm(q_c) * q_lora_scale
+ kv_c, k_pe = kv_lora.split(
+ [attention.kv_lora_rank, attention.qk_rope_head_dim], dim=-1
+ )
+ kv_c_normed = attention.kv_a_layernorm(kv_c) * kv_lora_scale
+ k_pe = k_rope_only_layernorm(k_pe).unsqueeze(1)
+
+ q = attention.q_b_proj(q_c)[0]
+ heads = attention.num_heads
+ if attention.dcp_q_replicate:
+ heads *= attention.q_b_proj.group_size
+ q = q.view(-1, heads, attention.qk_head_dim)
+ q[..., attention.qk_nope_head_dim :], k_pe = attention.rotary_emb(
+ positions, q[..., attention.qk_nope_head_dim :], k_pe
+ )
+
+ if attention.indexer and attention.is_sparse and not attention.skip_topk:
+ attention.indexer(hidden_states, q_c, positions, attention.indexer_rope_emb)
+ if llama_4_scaling is not None:
+ q *= llama_4_scaling
+
+ q_dcp_replicated = None
+ if attention.dcp_q_replicate:
+ q_dcp_replicated, q = q, attention.q_b_proj._local_view(q)
+ attn_out = attention.mla_attn(
+ q,
+ kv_c_normed,
+ k_pe,
+ output_shape=(
+ hidden_states.shape[0],
+ attention.num_heads * attention.v_head_dim,
+ ),
+ q_dcp_replicated=q_dcp_replicated,
+ )
+
+ gate = g_proj(hidden_states)[0]
+ if attention_gate_type == "headwise":
+ if gate.shape[-1] != attention.num_heads:
+ rank = get_tensor_model_parallel_rank()
+ gate = gate.narrow(-1, rank * attention.num_heads, attention.num_heads)
+ attn_out = attn_out.view(-1, attention.num_heads, attention.v_head_dim)
+ gate = torch.sigmoid(gate.float()).to(attn_out.dtype)
+ attn_out = (attn_out * gate.unsqueeze(-1)).flatten(-2)
+ else:
+ gate = torch.sigmoid(gate.float()).to(attn_out.dtype)
+ attn_out = attn_out * gate
+ return attention.o_proj(attn_out)[0]
+
+
+class Dots3NotePaddedMLAAttention(MLAAttention):
+ """MLA layer whose physical cache rows match NOTE's SWA rows."""
+
+ def __init__(self, *args, physical_head_size: int, **kwargs) -> None:
+ kwargs["attn_backend"] = Dots3NotePaddedSparseBackend
+ super().__init__(*args, **kwargs)
+ assert physical_head_size >= self.head_size
+ self.physical_head_size = physical_head_size
+
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MLAAttentionSpec:
+ spec = super().get_kv_cache_spec(vllm_config)
+ assert isinstance(spec, MLAAttentionSpec)
+ return replace(spec, head_size=self.physical_head_size)
+
+
+class Dots3NoteFullAttention(DeepseekV2MLAAttention):
+ """NOTE DSA on the existing DeepSeek sparse MLA attention."""
+
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ config,
+ prefix: str,
+ topk_indices_buffer: torch.Tensor | None = None,
+ ) -> None:
+ quant_config = vllm_config.quant_config
+ local_config = copy.copy(config)
+ object.__setattr__(
+ local_config,
+ "rope_parameters",
+ {
+ "rope_type": "default",
+ "rope_theta": config.rope_parameters["rope_theta"],
+ },
+ )
+ super().__init__(
+ vllm_config=vllm_config,
+ config=local_config,
+ hidden_size=config.hidden_size,
+ num_heads=config.num_attention_heads,
+ qk_nope_head_dim=config.qk_nope_head_dim,
+ qk_rope_head_dim=config.qk_rope_head_dim,
+ v_head_dim=config.v_head_dim,
+ q_lora_rank=config.q_lora_rank,
+ kv_lora_rank=config.kv_lora_rank,
+ max_position_embeddings=config.max_position_embeddings,
+ cache_config=vllm_config.cache_config,
+ quant_config=quant_config,
+ prefix=prefix,
+ topk_indices_buffer=topk_indices_buffer,
+ reduce_results=False,
+ )
+ wrapper = self._modules.pop("mla_attn")
+ inner_attention = wrapper.mla_attn
+ vllm_config.compilation_config.static_forward_context.pop(
+ inner_attention.layer_name
+ )
+ prefill_backend_cls = (
+ type(inner_attention.prefill_backend)
+ if inner_attention.prefill_backend is not None
+ else None
+ )
+ wrapper.mla_attn = Dots3NotePaddedMLAAttention(
+ num_heads=inner_attention.num_heads,
+ scale=inner_attention.scale,
+ q_lora_rank=inner_attention.q_lora_rank,
+ kv_lora_rank=inner_attention.kv_lora_rank,
+ qk_nope_head_dim=inner_attention.qk_nope_head_dim,
+ qk_rope_head_dim=inner_attention.qk_rope_head_dim,
+ v_head_dim=inner_attention.v_head_dim,
+ kv_b_proj=inner_attention.kv_b_proj,
+ dcp_q_replicate=inner_attention.dcp_q_replicate,
+ cache_config=vllm_config.cache_config,
+ quant_config=quant_config,
+ prefix=inner_attention.layer_name,
+ use_sparse=True,
+ indexer=wrapper.indexer,
+ topk_indices_buffer=topk_indices_buffer,
+ prefill_backend_cls=prefill_backend_cls,
+ physical_head_size=(config.swa_kv_lora_rank + config.swa_qk_rope_head_dim),
+ )
+ gate_type = config.attention_gate_type
+ gate_cls = ReplicatedLinear if gate_type == "headwise" else ColumnParallelLinear
+ gate_size = (
+ config.num_attention_heads
+ if gate_type == "headwise"
+ else config.num_attention_heads * config.v_head_dim
+ )
+ self.g_proj = gate_cls(
+ config.hidden_size,
+ gate_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.g_proj",
+ )
+ self.k_rope_only_layernorm = RMSNorm(
+ config.qk_rope_head_dim, eps=config.rms_norm_eps
+ )
+ self.attention_gate_type = gate_type
+ self.q_lora_scale = 1.0
+ self.kv_lora_scale = 1.0
+ if config.apply_mla_qkv_lora_rescale:
+ self.q_lora_scale = (config.hidden_size / config.q_lora_rank) ** 0.5
+ self.kv_lora_scale = (config.hidden_size / config.kv_lora_rank) ** 0.5
+ self.mla_attn = wrapper
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ llama_4_scaling: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ return _forward_note_mla(
+ self.mla_attn,
+ positions,
+ hidden_states,
+ g_proj=self.g_proj,
+ k_rope_only_layernorm=self.k_rope_only_layernorm,
+ attention_gate_type=self.attention_gate_type,
+ q_lora_scale=self.q_lora_scale,
+ kv_lora_scale=self.kv_lora_scale,
+ llama_4_scaling=llama_4_scaling,
+ )
+
+
+class Dots3NoteSlidingAttention(nn.Module):
+ """NOTE SWA constructed directly as dense sliding-window MLA."""
+
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ config,
+ prefix: str,
+ topk_indices_buffer: torch.Tensor | None = None,
+ ) -> None:
+ super().__init__()
+ del topk_indices_buffer
+ num_heads = config.swa_num_attention_heads
+ qk_nope_head_dim = config.swa_qk_nope_head_dim
+ qk_rope_head_dim = config.swa_qk_rope_head_dim
+ v_head_dim = config.swa_v_head_dim
+ q_lora_rank = config.swa_q_lora_rank
+ kv_lora_rank = config.swa_kv_lora_rank
+ qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+ tp_size = get_tensor_model_parallel_world_size()
+ assert num_heads % tp_size == 0
+ quant_config = vllm_config.quant_config
+
+ fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
+ config.hidden_size,
+ [q_lora_rank, kv_lora_rank + qk_rope_head_dim],
+ quant_config=quant_config,
+ prefix=f"{prefix}.fused_qkv_a_proj",
+ )
+ qrep_enabled = (
+ envs.VLLM_DCP_Q_REPLICATE
+ and vllm_config.parallel_config.decode_context_parallel_size > 1
+ and vllm_config.parallel_config.prefill_context_parallel_size <= 1
+ )
+ q_proj_cls = (
+ DCPGroupColumnParallelLinear if qrep_enabled else ColumnParallelLinear
+ )
+ q_a_layernorm = RMSNorm(q_lora_rank, eps=config.rms_norm_eps)
+ q_b_proj = q_proj_cls(
+ q_lora_rank,
+ num_heads * qk_head_dim,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.q_b_proj",
+ )
+ kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps)
+ kv_b_proj = ColumnParallelLinear(
+ kv_lora_rank,
+ num_heads * (qk_nope_head_dim + v_head_dim),
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.kv_b_proj",
+ )
+ o_proj = RowParallelLinear(
+ num_heads * v_head_dim,
+ config.hidden_size,
+ bias=False,
+ reduce_results=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.o_proj",
+ )
+ gate_type = config.swa_attention_gate_type
+ gate_cls = ReplicatedLinear if gate_type == "headwise" else ColumnParallelLinear
+ gate_size = num_heads if gate_type == "headwise" else num_heads * v_head_dim
+ g_proj = gate_cls(
+ config.hidden_size,
+ gate_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.g_proj",
+ )
+ k_rope_only_layernorm = RMSNorm(qk_rope_head_dim, eps=config.rms_norm_eps)
+ rotary_emb = get_rope(
+ qk_rope_head_dim,
+ max_position=config.max_position_embeddings,
+ rope_parameters={
+ "rope_type": "default",
+ "rope_theta": config.swa_rope_theta,
+ },
+ is_neox_style=False,
+ )
+ cache_config = vllm_config.cache_config
+ if cache_config.cache_dtype == "fp8_ds_mla":
+ cache_config = copy.copy(cache_config)
+ cache_config.cache_dtype = "fp8"
+ apply_rescale = config.apply_mla_qkv_lora_rescale
+ self.hidden_size = config.hidden_size
+ self.num_heads = num_heads // tp_size
+ self.qk_nope_head_dim = qk_nope_head_dim
+ self.qk_rope_head_dim = qk_rope_head_dim
+ self.qk_head_dim = qk_head_dim
+ self.v_head_dim = v_head_dim
+ self.q_lora_rank = q_lora_rank
+ self.kv_lora_rank = kv_lora_rank
+ self.fused_qkv_a_proj = fused_qkv_a_proj
+ self.q_a_layernorm = q_a_layernorm
+ self.q_b_proj = q_b_proj
+ self.kv_a_layernorm = kv_a_layernorm
+ self.kv_b_proj = kv_b_proj
+ self.rotary_emb = rotary_emb
+ self.o_proj = o_proj
+ self.g_proj = g_proj
+ self.k_rope_only_layernorm = k_rope_only_layernorm
+ self.indexer = None
+ self.indexer_rope_emb = None
+ self.is_sparse = False
+ self.skip_topk = False
+ self.dcp_q_replicate = getattr(q_b_proj, "qrep_active", False)
+ self.attention_gate_type = gate_type
+ self.q_lora_scale = (
+ (config.hidden_size / q_lora_rank) ** 0.5 if apply_rescale else 1.0
+ )
+ self.kv_lora_scale = (
+ (config.hidden_size / kv_lora_rank) ** 0.5 if apply_rescale else 1.0
+ )
+ self.mla_attn = MLAAttention(
+ num_heads=self.num_heads,
+ scale=qk_head_dim**-0.5,
+ qk_nope_head_dim=qk_nope_head_dim,
+ qk_rope_head_dim=qk_rope_head_dim,
+ v_head_dim=v_head_dim,
+ q_lora_rank=q_lora_rank,
+ kv_lora_rank=kv_lora_rank,
+ kv_b_proj=kv_b_proj,
+ dcp_q_replicate=self.dcp_q_replicate,
+ cache_config=cache_config,
+ quant_config=quant_config,
+ prefix=f"{prefix}.attn",
+ sliding_window=config.sliding_window_size,
+ attn_backend=Dots3NoteTritonMLABackend,
+ prefill_backend_cls=Dots3NoteFlashAttnPrefillBackend,
+ )
+
+ def forward(
+ self,
+ positions: torch.Tensor,
+ hidden_states: torch.Tensor,
+ llama_4_scaling: torch.Tensor | None = None,
+ ) -> torch.Tensor:
+ return _forward_note_mla(
+ self,
+ positions,
+ hidden_states,
+ g_proj=self.g_proj,
+ k_rope_only_layernorm=self.k_rope_only_layernorm,
+ attention_gate_type=self.attention_gate_type,
+ q_lora_scale=self.q_lora_scale,
+ kv_lora_scale=self.kv_lora_scale,
+ llama_4_scaling=llama_4_scaling,
+ )
+
+
+class Dots3NoteDecoderLayer(DeepseekV32DecoderLayer):
+ """DeepSeek-V3.2 decoder orchestration with NOTE-local modules."""
+
+ def __init__(
+ self,
+ vllm_config: VllmConfig,
+ prefix: str,
+ config=None,
+ topk_indices_buffer: torch.Tensor | None = None,
+ ) -> None:
+ nn.Module.__init__(self)
+ if config is None:
+ config = vllm_config.model_config.hf_config
+ quant_config = vllm_config.quant_config
+ parallel_config = vllm_config.parallel_config
+ self.hidden_size = config.hidden_size
+ layer_idx = int(prefix.split(sep=".")[-1])
+ self.layer_idx = layer_idx
+ self.use_mha = False
+ attention_cls = (
+ Dots3NoteSlidingAttention
+ if config.layer_types[layer_idx] == "sliding_attention"
+ else Dots3NoteFullAttention
+ )
+ self.self_attn = attention_cls(
+ vllm_config=vllm_config,
+ config=config,
+ prefix=f"{prefix}.self_attn",
+ topk_indices_buffer=topk_indices_buffer,
+ )
+
+ moe_layer_freq = getattr(config, "moe_layer_freq", 1)
+ is_moe = (
+ layer_idx < config.num_hidden_layers
+ and config.n_routed_experts is not None
+ and layer_idx >= config.first_k_dense_replace
+ and layer_idx % moe_layer_freq == 0
+ )
+ if is_moe:
+ self.mlp = Dots3NoteMoE(
+ config=config,
+ parallel_config=parallel_config,
+ quant_config=quant_config,
+ reduce_results=False,
+ prefix=f"{prefix}.mlp",
+ apply_routed_scale_to_output=False,
+ )
+ else:
+ self.mlp = DeepseekV2MLP(
+ hidden_size=config.hidden_size,
+ intermediate_size=_padded_mlp_size(
+ config.intermediate_size, quant_config
+ ),
+ hidden_act=config.hidden_act,
+ quant_config=quant_config,
+ prefix=f"{prefix}.mlp",
+ reduce_results=False,
+ )
+ self.use_sequence_parallel_moe = (
+ parallel_config.use_sequence_parallel_moe
+ and parallel_config.pipeline_parallel_size == 1
+ and isinstance(self.mlp, DeepseekV2MoE)
+ )
+ self.tp_size = parallel_config.tensor_parallel_size
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.post_attention_layernorm = RMSNorm(
+ config.hidden_size, eps=config.rms_norm_eps
+ )
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
+
+
+class Dots3NoteModel(DeepseekV32Model):
+ """DeepSeek-V3.2 runtime shell with NOTE-local decoder layers."""
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ nn.Module.__init__(self)
+ config = vllm_config.model_config.hf_config
+ quant_config = vllm_config.quant_config
+ self.config = config
+ self.device = current_platform.device_type
+ self.vocab_size = config.vocab_size
+ self.is_v32 = True
+ self._weight_block_size = getattr(quant_config, "weight_block_size", None)
+ topk_indices_buffer = torch.empty(
+ vllm_config.scheduler_config.max_num_batched_tokens,
+ config.index_topk,
+ dtype=torch.int32,
+ device=self.device,
+ )
+ if get_pp_group().is_first_rank:
+ self.embed_tokens = VocabParallelEmbedding(
+ config.vocab_size,
+ config.hidden_size,
+ quant_config=quant_config,
+ prefix=f"{prefix}.embed_tokens",
+ )
+ else:
+ self.embed_tokens = PPMissingLayer()
+ self.start_layer, self.end_layer, self.layers = make_layers(
+ config.num_hidden_layers,
+ lambda prefix: Dots3NoteDecoderLayer(
+ vllm_config=vllm_config,
+ prefix=prefix,
+ topk_indices_buffer=topk_indices_buffer,
+ ),
+ prefix=f"{prefix}.layers",
+ )
+ if get_pp_group().is_last_rank:
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ else:
+ self.norm = PPMissingLayer()
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
+ ["hidden_states", "residual"], config.hidden_size
+ )
+ self.aux_hidden_state_layers = tuple[int, ...]()
+ self.num_redundant_experts = (
+ vllm_config.parallel_config.eplb_config.num_redundant_experts
+ )
+
+ def _pad_dense_mlp_weight(
+ self, name: str, loaded_weight: torch.Tensor
+ ) -> torch.Tensor:
+ block_size = self._weight_block_size
+ if block_size is None or ".mlp.experts." in name:
+ return loaded_weight
+ if not any(
+ proj_name in name
+ for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
+ ):
+ return loaded_weight
+ dim = 1 if ".down_proj." in name else 0
+ block_step = 1 if name.endswith("weight_scale_inv") else block_size[0]
+ multiple = get_tensor_model_parallel_world_size() * block_step
+ pad = (-loaded_weight.shape[dim]) % multiple
+ if pad == 0:
+ return loaded_weight
+ pad_shape = list(loaded_weight.shape)
+ pad_shape[dim] = pad
+ return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
+
+ def _adapt_weights(
+ self, weights: Iterable[tuple[str, torch.Tensor]]
+ ) -> Iterator[tuple[str, torch.Tensor]]:
+ pending_indexer_fp8: dict[tuple[str, str], dict[str, torch.Tensor]] = {}
+ for name, weight in weights:
+ if name.startswith("mtp."):
+ continue
+ if ".indexer." in name:
+ parts = name.split(".")
+ if "layers" in parts:
+ layer_idx = int(parts[parts.index("layers") + 1])
+ if (
+ layer_idx < self.config.num_hidden_layers
+ and self.config.layer_types[layer_idx] == "sliding_attention"
+ ):
+ continue
+
+ projection = next(
+ (
+ candidate
+ for candidate in ("wk", "weights_proj")
+ if f".indexer.{candidate}." in name
+ ),
+ None,
+ )
+ is_fp8_weight = (
+ name.endswith(".weight") and weight.dtype == torch.float8_e4m3fn
+ )
+ is_scale = "weight_scale" in name
+ if projection is not None and (is_fp8_weight or is_scale):
+ layer_prefix = name.rsplit(f".{projection}.", 1)[0]
+ key = (layer_prefix, projection)
+ entry = pending_indexer_fp8.setdefault(key, {})
+ entry["weight" if is_fp8_weight else "scale"] = weight
+ if "weight" not in entry or "scale" not in entry:
+ continue
+ weight_fp8 = entry["weight"]
+ scale_inv = entry["scale"]
+ del pending_indexer_fp8[key]
+ group_shape = GroupShape(
+ weight_fp8.shape[0] // scale_inv.shape[0],
+ weight_fp8.shape[1] // scale_inv.shape[1],
+ )
+ dequantized = scaled_dequantize(
+ weight_fp8,
+ scale_inv,
+ group_shape=group_shape,
+ out_dtype=torch.bfloat16,
+ )
+ yield f"{layer_prefix}.{projection}.weight", dequantized
+ continue
+ yield name, self._pad_dense_mlp_weight(name, weight)
+ if pending_indexer_fp8:
+ missing = ", ".join(
+ f"{prefix}.{proj}" for prefix, proj in pending_indexer_fp8
+ )
+ raise ValueError(f"Incomplete FP8 indexer weights: {missing}")
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ return super().load_weights(self._adapt_weights(weights))
+
+
+class Dots3NoteLanguageModelForCausalLM(DeepseekV32ForCausalLM):
+ model_cls = Dots3NoteModel
diff --git a/vllm/models/dots3_note/nvidia/mtp.py b/vllm/models/dots3_note/nvidia/mtp.py
new file mode 100644
index 0000000000..d2e31dfe8e
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/mtp.py
@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NVIDIA multi-token predictor for Dots3Note."""
+
+from collections.abc import Iterable, Iterator
+
+import torch
+import torch.nn as nn
+
+from vllm.config import VllmConfig
+from vllm.distributed import (
+ get_tensor_model_parallel_world_size,
+ tensor_model_parallel_all_gather,
+ tensor_model_parallel_all_reduce,
+)
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.linear import ReplicatedLinear
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
+from vllm.model_executor.models.deepseek_mtp import SharedHead
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
+from vllm.models.deepseek_v32.nvidia.mtp import (
+ DeepseekV32MTP,
+ DeepseekV32MultiTokenPredictor,
+)
+
+from .model import Dots3NoteDecoderLayer
+
+
+class Dots3NoteMultiTokenPredictorLayer(nn.Module):
+ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+ super().__init__()
+ assert vllm_config.speculative_config is not None
+ config = vllm_config.speculative_config.draft_model_config.hf_config
+ quant_config = vllm_config.quant_config
+
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+ self.eh_proj = ReplicatedLinear(
+ config.hidden_size * 2,
+ config.hidden_size,
+ bias=False,
+ quant_config=quant_config,
+ prefix=f"{prefix}.eh_proj",
+ )
+ self.shared_head = SharedHead(
+ config=config, prefix=prefix, quant_config=quant_config
+ )
+ self.mtp_block = Dots3NoteDecoderLayer(
+ vllm_config=vllm_config,
+ config=config,
+ prefix=prefix,
+ )
+
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ positions: torch.Tensor,
+ previous_hidden_states: torch.Tensor,
+ inputs_embeds: torch.Tensor | None = None,
+ spec_step_index: int = 0,
+ ) -> torch.Tensor:
+ del input_ids, spec_step_index
+ assert inputs_embeds is not None
+ eh_input = fused_eh_norm(
+ positions,
+ inputs_embeds,
+ previous_hidden_states,
+ self.enorm.weight,
+ self.hnorm.weight,
+ self.enorm.variance_epsilon,
+ )
+ hidden_states = self.eh_proj(eh_input)[0]
+ hidden_states, residual = self.mtp_block(
+ positions=positions, hidden_states=hidden_states, residual=None
+ )
+ is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe
+ if not is_sequence_parallel:
+ hidden_states = tensor_model_parallel_all_reduce(hidden_states)
+ hidden_states, _ = self.shared_head.norm(hidden_states, residual)
+ if is_sequence_parallel:
+ hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
+ hidden_states = hidden_states[: positions.shape[0]]
+ return hidden_states
+
+
+class Dots3NoteMultiTokenPredictor(DeepseekV32MultiTokenPredictor):
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ nn.Module.__init__(self)
+ assert vllm_config.speculative_config is not None
+ config = vllm_config.speculative_config.draft_model_config.hf_config
+ self.mtp_start_layer_idx = config.num_hidden_layers
+ self.num_mtp_layers = config.num_nextn_predict_layers
+ self.layers = nn.ModuleDict(
+ {
+ str(idx): Dots3NoteMultiTokenPredictorLayer(
+ vllm_config, f"{prefix}.layers.{idx}"
+ )
+ for idx in range(
+ self.mtp_start_layer_idx,
+ self.mtp_start_layer_idx + self.num_mtp_layers,
+ )
+ }
+ )
+ self.embed_tokens = VocabParallelEmbedding(
+ config.vocab_size,
+ config.hidden_size,
+ prefix=maybe_prefix(prefix, "embed_tokens"),
+ )
+ self.register_buffer(
+ "max_token_id",
+ torch.tensor(config.vocab_size - 1, dtype=torch.int64),
+ persistent=False,
+ )
+ self.logits_processor = LogitsProcessor(config.vocab_size)
+
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+ input_ids = torch.minimum(input_ids, self.max_token_id)
+ return self.embed_tokens(input_ids)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor,
+ positions: torch.Tensor,
+ previous_hidden_states: torch.Tensor,
+ inputs_embeds: torch.Tensor | None = None,
+ spec_step_idx: int = 0,
+ ) -> torch.Tensor:
+ if inputs_embeds is None:
+ inputs_embeds = self.embed_input_ids(input_ids)
+ return super().forward(
+ input_ids,
+ positions,
+ previous_hidden_states,
+ inputs_embeds,
+ spec_step_idx,
+ )
+
+
+class Dots3NoteMTP(DeepseekV32MTP):
+ has_own_embed_tokens = True
+ has_own_lm_head = False
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ nn.Module.__init__(self)
+ self.config = vllm_config.model_config.hf_config
+ self.quant_config = vllm_config.quant_config
+ self.model = Dots3NoteMultiTokenPredictor(
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+ )
+ self.set_moe_parameters()
+
+ def _pad_dense_mlp_weight(
+ self, name: str, loaded_weight: torch.Tensor
+ ) -> torch.Tensor:
+ block_size = getattr(self.quant_config, "weight_block_size", None)
+ if block_size is None or ".mlp.experts." in name:
+ return loaded_weight
+ if not any(
+ proj_name in name
+ for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
+ ):
+ return loaded_weight
+ dim = 1 if ".down_proj." in name else 0
+ block_step = 1 if name.endswith("weight_scale_inv") else block_size[0]
+ multiple = get_tensor_model_parallel_world_size() * block_step
+ pad = (-loaded_weight.shape[dim]) % multiple
+ if pad == 0:
+ return loaded_weight
+ pad_shape = list(loaded_weight.shape)
+ pad_shape[dim] = pad
+ return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
+
+ def _adapt_weights(
+ self, weights: Iterable[tuple[str, torch.Tensor]]
+ ) -> Iterator[tuple[str, torch.Tensor]]:
+ mtp_layer = self.config.num_hidden_layers
+ for name, weight in weights:
+ if name.startswith("model.mtp.embed_tokens."):
+ name = name.replace(
+ "model.mtp.embed_tokens.",
+ f"model.layers.{mtp_layer}.embed_tokens.",
+ 1,
+ )
+ yield name, self._pad_dense_mlp_weight(name, weight)
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ return super().load_weights(self._adapt_weights(weights))
diff --git a/vllm/models/dots3_note/nvidia/multimodal.py b/vllm/models/dots3_note/nvidia/multimodal.py
new file mode 100644
index 0000000000..cd4bfeae6a
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/multimodal.py
@@ -0,0 +1,304 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""vLLM composition layer for Dots3Note image and audio encoders."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm.config import VllmConfig
+from vllm.model_executor.models.interfaces import (
+ MultiModalEmbeddings,
+ SupportsMultiModal,
+ SupportsPP,
+)
+from vllm.model_executor.models.module_mapping import MultiModelKeys
+from vllm.model_executor.models.utils import (
+ AutoWeightsLoader,
+ IntermediateTensors,
+ WeightsMapper,
+ maybe_prefix,
+)
+from vllm.multimodal import MULTIMODAL_REGISTRY
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
+
+from ..common.processor import (
+ AUDIO_END,
+ AUDIO_PAD,
+ AUDIO_START,
+ IMAGE_END,
+ IMAGE_PAD,
+ IMAGE_START,
+ VIDEO_PLACEHOLDER,
+ Dots3NoteDummyInputsBuilder,
+ Dots3NoteMultiModalProcessor,
+ Dots3NoteProcessingInfo,
+ load_note_config_section,
+)
+from .audio import Dots3NoteAudioConfig, Dots3NoteAudioModel
+from .model import Dots3NoteLanguageModelForCausalLM
+from .vision import DotsMoEVitConfig, DotsMoEVitModel
+
+
+@MULTIMODAL_REGISTRY.register_processor(
+ Dots3NoteMultiModalProcessor,
+ info=Dots3NoteProcessingInfo,
+ dummy_inputs=Dots3NoteDummyInputsBuilder,
+)
+class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
+ """Dots3Note model with optional image and audio towers."""
+
+ supports_encoder_tp_data = True
+
+ hf_to_vllm_mapper = WeightsMapper(
+ orig_to_new_prefix={
+ "model.": "language_model.model.",
+ "lm_head.": "language_model.lm_head.",
+ "mtp.": "language_model.mtp.",
+ "vision_encoder.": "visual.",
+ "audio_encoder.": "audio_tower.",
+ }
+ )
+
+ @classmethod
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
+ if modality.startswith("image"):
+ return f"{IMAGE_START}{IMAGE_PAD}{IMAGE_END}"
+ if modality.startswith("audio"):
+ return f"{AUDIO_START}{AUDIO_PAD}{AUDIO_END}"
+ if modality.startswith("video"):
+ return VIDEO_PLACEHOLDER
+ raise ValueError(f"Unsupported modality: {modality}")
+
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+ super().__init__()
+ model_config = vllm_config.model_config
+ self.config = model_config.hf_config
+ self.quant_config = vllm_config.quant_config
+ self.multimodal_config = model_config.multimodal_config
+ assert self.multimodal_config is not None
+
+ added_tokens = get_hf_file_to_dict(
+ "added_tokens.json",
+ model_config.model,
+ model_config.revision,
+ )
+ if added_tokens is None or IMAGE_PAD not in added_tokens:
+ raise ValueError("NOTE tokenizer is missing the image padding token")
+ self.config.image_token_index = int(added_tokens[IMAGE_PAD])
+
+ vision_config_dict = load_note_config_section(
+ model_config.model,
+ model_config.revision,
+ "vision_config",
+ )
+ audio_config_dict = load_note_config_section(
+ model_config.model,
+ model_config.revision,
+ "audio_config",
+ )
+ video_enabled = (
+ vision_config_dict is not None
+ and self.multimodal_config.get_limit_per_prompt("video") > 0
+ )
+ image_enabled = vision_config_dict is not None and (
+ self.multimodal_config.get_limit_per_prompt("image") > 0 or video_enabled
+ )
+ audio_enabled = audio_config_dict is not None and (
+ self.multimodal_config.get_limit_per_prompt("audio") > 0 or video_enabled
+ )
+
+ with self._mark_tower_model(vllm_config, {"image", "audio", "video"}):
+ self.visual: DotsMoEVitModel | None
+ if image_enabled:
+ assert vision_config_dict is not None
+ self.visual = DotsMoEVitModel(DotsMoEVitConfig(**vision_config_dict))
+ else:
+ self.visual = None
+ self.audio_tower: Dots3NoteAudioModel | None
+ if audio_enabled:
+ assert audio_config_dict is not None
+ self.audio_tower = Dots3NoteAudioModel(
+ Dots3NoteAudioConfig(**audio_config_dict)
+ )
+ else:
+ self.audio_tower = None
+ # The native encoder service casts each complete tower before
+ # loading its checkpoint. This also converts explicitly-created
+ # floating buffers (for example RoPE tables and router state),
+ # which merely constructing under vLLM's default dtype does not.
+ if self.visual is not None:
+ self.visual.to(dtype=model_config.dtype)
+ if self.audio_tower is not None:
+ self.audio_tower.to(dtype=model_config.dtype)
+ with self._mark_language_model(vllm_config):
+ self.language_model = Dots3NoteLanguageModelForCausalLM(
+ vllm_config=vllm_config,
+ prefix=maybe_prefix(prefix, "language_model"),
+ )
+ self.make_empty_intermediate_tensors = (
+ self.language_model.make_empty_intermediate_tensors
+ )
+
+ def _process_image_input(
+ self,
+ pixel_values: torch.Tensor,
+ image_grid_thw: torch.Tensor,
+ ) -> 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()
+ return image_embeds.split(sizes)
+
+ def _process_audio_input(
+ self,
+ audio_values: torch.Tensor,
+ audio_lengths: torch.Tensor,
+ ) -> tuple[torch.Tensor, ...]:
+ if self.audio_tower is None:
+ return ()
+ if audio_values.dtype != torch.int32:
+ raise TypeError(
+ "NOTE audio values must carry float32 waveform bits as int32, "
+ f"got {audio_values.dtype}"
+ )
+ waveforms = audio_values.contiguous().view(torch.float32)
+ audio_embeds, item_lengths = self.audio_tower(waveforms, audio_lengths)
+ return audio_embeds.split(item_lengths)
+
+ def _process_video_input(
+ self,
+ pixel_values: torch.Tensor,
+ image_grid_thw: torch.Tensor,
+ audio_values: torch.Tensor,
+ audio_lengths: torch.Tensor,
+ modalities: torch.Tensor,
+ frame_counts: torch.Tensor,
+ audio_counts: torch.Tensor,
+ emission_counts: torch.Tensor,
+ ) -> tuple[torch.Tensor, ...]:
+ image_embeds = self._process_image_input(pixel_values, image_grid_thw)
+ audio_embeds = (
+ self._process_audio_input(audio_values, audio_lengths)
+ if audio_lengths.numel()
+ else ()
+ )
+ orders = modalities.tolist()
+ frame_counts_list = frame_counts.tolist()
+ audio_counts_list = audio_counts.tolist()
+ emission_counts_list = emission_counts.tolist()
+ outputs: list[torch.Tensor] = []
+ image_idx = audio_idx = order_idx = 0
+ for num_frames, num_audios, num_emissions in zip(
+ frame_counts_list,
+ audio_counts_list,
+ emission_counts_list,
+ ):
+ video_order = orders[order_idx : order_idx + num_emissions]
+ video_parts: list[torch.Tensor] = []
+ video_image_start = image_idx
+ video_audio_start = audio_idx
+ for modality in video_order:
+ if modality == 0:
+ video_parts.append(image_embeds[image_idx])
+ image_idx += 1
+ elif modality == 1:
+ if audio_idx >= len(audio_embeds):
+ raise ValueError("NOTE video audio tower output is missing")
+ video_parts.append(audio_embeds[audio_idx])
+ audio_idx += 1
+ else:
+ raise ValueError(f"Unknown NOTE video modality id: {modality}")
+ if image_idx - video_image_start != num_frames:
+ raise ValueError("NOTE video frame order/count mismatch")
+ if audio_idx - video_audio_start != num_audios:
+ raise ValueError("NOTE video audio order/count mismatch")
+ outputs.append(torch.cat(video_parts))
+ order_idx += num_emissions
+ if image_idx != len(image_embeds) or audio_idx != len(audio_embeds):
+ raise ValueError("NOTE video encoder outputs were not fully consumed")
+ return tuple(outputs)
+
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
+ multimodal_embeddings: list[torch.Tensor] = []
+ handled: set[str] = set()
+ for input_key in kwargs:
+ if input_key == "pixel_values" and "image" not in handled:
+ pixel_values = kwargs.get("pixel_values")
+ image_grid_thw = kwargs.get("image_grid_thw")
+ if isinstance(pixel_values, torch.Tensor) and isinstance(
+ image_grid_thw, torch.Tensor
+ ):
+ multimodal_embeddings.extend(
+ self._process_image_input(pixel_values, image_grid_thw)
+ )
+ handled.add("image")
+ elif input_key == "audio_values" and "audio" not in handled:
+ audio_values = kwargs.get("audio_values")
+ audio_lengths = kwargs.get("audio_lengths")
+ if isinstance(audio_values, torch.Tensor) and isinstance(
+ audio_lengths, torch.Tensor
+ ):
+ multimodal_embeddings.extend(
+ self._process_audio_input(audio_values, audio_lengths)
+ )
+ handled.add("audio")
+ elif input_key == "video_pixel_values" and "video" not in handled:
+ video_inputs = (
+ kwargs.get("video_pixel_values"),
+ kwargs.get("video_image_grid_thw"),
+ kwargs.get("video_audio_values"),
+ kwargs.get("video_audio_lengths"),
+ kwargs.get("video_modalities"),
+ kwargs.get("video_frame_counts"),
+ kwargs.get("video_audio_counts"),
+ kwargs.get("video_emission_counts"),
+ )
+ if all(isinstance(value, torch.Tensor) for value in video_inputs):
+ multimodal_embeddings.extend(
+ self._process_video_input(*video_inputs) # type: ignore[arg-type]
+ )
+ handled.add("video")
+ return tuple(multimodal_embeddings)
+
+ def forward(
+ self,
+ input_ids: torch.Tensor | None,
+ positions: torch.Tensor,
+ intermediate_tensors: IntermediateTensors | None = None,
+ inputs_embeds: torch.Tensor | None = None,
+ **kwargs: object,
+ ) -> torch.Tensor | IntermediateTensors:
+ return self.language_model(
+ input_ids,
+ positions,
+ intermediate_tensors,
+ inputs_embeds,
+ )
+
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+ return self.language_model.compute_logits(hidden_states)
+
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+ skip_prefixes = []
+ if self.visual is None:
+ skip_prefixes.append("visual.")
+ if self.audio_tower is None:
+ skip_prefixes.append("audio_tower.")
+ return AutoWeightsLoader(self, skip_prefixes=skip_prefixes).load_weights(
+ weights,
+ mapper=self.hf_to_vllm_mapper,
+ )
+
+ def process_weights_after_loading(self) -> None:
+ if self.visual is not None:
+ self.visual.process_weights_after_loading()
+
+ def get_mm_mapping(self) -> MultiModelKeys:
+ return MultiModelKeys.from_string_field(
+ language_model="language_model",
+ tower_model=["visual", "audio_tower"],
+ )
diff --git a/vllm/models/dots3_note/nvidia/vision.py b/vllm/models/dots3_note/nvidia/vision.py
new file mode 100644
index 0000000000..c6f560fcf2
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/vision.py
@@ -0,0 +1,677 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+import math
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+from torch.nn import LayerNorm
+from transformers.configuration_utils import PretrainedConfig
+from transformers.modeling_utils import PreTrainedModel
+
+from vllm.third_party.deep_gemm import per_block_cast_to_fp8
+
+from .vision_attention import (
+ VisionRotaryEmbedding,
+ apply_vision_attention_residual,
+ attn_uses_seqlens,
+ build_vision_attention,
+ prepare_seqlens_for_attention,
+ resolve_attn_implementation,
+)
+from .vision_moe import note_vision_fused_moe_fp8
+
+
+class DotsMoEVitConfig(PretrainedConfig):
+ model_type: str = "dots_moe_vit"
+
+ def __init__(
+ self,
+ embed_dim: int = 1536,
+ hidden_size: int = 2048,
+ intermediate_size: int = 4224,
+ moe_intermediate_size: int = 2112,
+ num_hidden_layers: int = 42,
+ num_attention_heads: int = 24,
+ num_channels: int = 3,
+ patch_size: int = 14,
+ spatial_merge_size: int = 2,
+ temporal_patch_size: int = 1,
+ rms_norm_eps: float = 1e-5,
+ use_bias: bool = False,
+ use_qk_norm: bool = True,
+ attn_implementation="flash_attention_3", # "eager","eager_v2","sdpa","flash_attention_2","flash_attention_3"
+ initializer_range=0.02,
+ is_causal=False,
+ post_norm=True,
+ gradient_checkpointing=False,
+ pyramid_num_routed: list[int] | None = None,
+ capacity_factor: float = 2.0,
+ router_scoring_func: str = "sigmoid",
+ router_scale: float = 1.0,
+ adapter_in_dim: int = 1536,
+ adapter_out_dim: int = 2048,
+ adapter_merge_size: int = 2,
+ # ``pixel_shuffle_mlp`` reshapes 2x2 spatial neighbours via NHWC
+ # pixel-shuffle, then runs layer norm and a two-layer MLP.
+ # ``patch_merger`` (cybertron PatchMerger, e.g. fireall_iter02275) skips the pixel-shuffle permutation and instead views every 4 consecutive 2x2-grouped tokens as one row.
+ adapter_type: str = "pixel_shuffle_mlp",
+ # If True the preprocessor already emits patches in 2x2-grouped order (qwen ``merge_size=2`` flatten path), so RoPE positions must also be regrouped
+ # to match the in-block layout. Older checkpoints were trained with
+ # ``pre_pixel_shuffle=False``, so the default keeps row-major RoPE.
+ pre_pixel_shuffle: bool = False,
+ # Keep the native encoder's post-load compile behavior. Compiling in
+ # ``__init__`` would rename checkpoint keys through ``_orig_mod``.
+ enable_torch_compile: bool = True,
+ # If True, use FP8 MoE implementation
+ enable_fp8_moe: bool = True,
+ **kwargs: Any,
+ ):
+ super().__init__(**kwargs)
+ self.embed_dim = embed_dim
+ self.hidden_size = hidden_size
+ self.intermediate_size = intermediate_size
+ self.moe_intermediate_size = moe_intermediate_size
+ self.num_hidden_layers = num_hidden_layers
+ self.num_attention_heads = num_attention_heads
+ self.num_channels = num_channels
+ self.patch_size = patch_size
+ self.spatial_merge_size = spatial_merge_size
+ self.temporal_patch_size = temporal_patch_size
+ self.rms_norm_eps = rms_norm_eps
+ self.use_bias = use_bias
+ self.use_qk_norm = use_qk_norm
+ self.attn_implementation = attn_implementation
+ self.initializer_range = initializer_range
+ self.is_causal = is_causal
+ self.post_norm = post_norm
+ self.gradient_checkpointing = gradient_checkpointing
+ self.pyramid_num_routed = pyramid_num_routed or []
+ self.capacity_factor = capacity_factor
+ self.router_scoring_func = router_scoring_func
+ self.router_scale = router_scale
+ self.adapter_in_dim = adapter_in_dim
+ self.adapter_out_dim = adapter_out_dim
+ self.adapter_merge_size = adapter_merge_size
+ if adapter_type not in ("pixel_shuffle_mlp", "patch_merger"):
+ raise ValueError(
+ f"adapter_type must be 'pixel_shuffle_mlp' or 'patch_merger', got {adapter_type!r}"
+ )
+ self.adapter_type = adapter_type
+ self.pre_pixel_shuffle = pre_pixel_shuffle
+ self.enable_torch_compile = enable_torch_compile
+ self.enable_fp8_moe = enable_fp8_moe
+
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(dim))
+ self.eps = eps
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+ def extra_repr(self) -> str:
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
+
+ def _norm(self, x: torch.Tensor) -> torch.Tensor:
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+
+# ---- FFN modules ----
+
+
+class DotsSwiGLUFFN(nn.Module):
+ def __init__(self, in_features, hidden_features, bias=False):
+ super().__init__()
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
+ self.fc2 = nn.Linear(hidden_features, in_features, bias=bias)
+ self.fc3 = nn.Linear(in_features, hidden_features, bias=bias)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ return self.fc2(F.silu(self.fc1(x)) * self.fc3(x))
+
+
+class MoESwiGLUFFN(nn.Module):
+ """MoE FFN with per-expert SwiGLU experts, sigmoid/softmax gating, top-k routing."""
+
+ def __init__(self, config: DotsMoEVitConfig, layer_number: int):
+ super().__init__()
+ self.config = config
+ self.layer_number = layer_number
+ self.hidden_size = config.embed_dim
+ self.num_routed = config.pyramid_num_routed[layer_number]
+ self.capacity_factor = config.capacity_factor
+ self.router_scoring_func = config.router_scoring_func
+ self.router_scale = config.router_scale
+
+ self.register_buffer(
+ "router_bias", torch.zeros(self.num_routed, dtype=torch.float32)
+ )
+
+ self.experts = nn.ModuleList(
+ [
+ DotsSwiGLUFFN(
+ self.hidden_size, config.moe_intermediate_size, bias=config.use_bias
+ )
+ for _ in range(self.num_routed)
+ ]
+ )
+
+ self.gate_weight = nn.Parameter(
+ torch.empty((self.num_routed, self.hidden_size), dtype=torch.float32)
+ )
+ nn.init.kaiming_uniform_(self.gate_weight, a=math.sqrt(5))
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ # Mirror cybertron's ``AIMv2MoEGEMMSwiGLUFFN.forward``: keep ``gating_prob`` /
+ # router-bias add in fp32 and run topk in fp32. The legacy ``.type_as(x_flat)``
+ # path on bf16 made expert routing diverge whenever two routed-expert scores
+ # tied within bf16 precision, which is the dominant source of numerical drift
+ # observed against latest cybertron checkpoints.
+ epsilon = 1e-9
+ x_flat = x.contiguous()
+ num_tokens = x_flat.shape[0]
+
+ gate_logits = F.linear(x_flat.float(), self.gate_weight.float())
+
+ if self.router_scoring_func == "sigmoid":
+ gating_prob = torch.sigmoid(gate_logits)
+ else:
+ gating_prob = torch.softmax(gate_logits, dim=-1, dtype=torch.float32)
+
+ aggregated_output = torch.zeros_like(x_flat)
+ aggregated_gate = torch.zeros(num_tokens, dtype=x.dtype, device=x.device)
+
+ topk = min(int(self.capacity_factor), self.num_routed)
+
+ gating_with_bias = gating_prob + self.router_bias.to(torch.float32).unsqueeze(0)
+ _, topk_indices = torch.topk(gating_with_bias, k=topk, dim=-1, sorted=False)
+
+ routed_weights = gating_prob.gather(1, topk_indices)
+ if self.router_scoring_func == "sigmoid" and topk > 1:
+ routed_weights = routed_weights / (
+ routed_weights.sum(dim=-1, keepdim=True) + epsilon
+ )
+ routed_weights = (routed_weights * self.router_scale).to(x_flat.dtype)
+
+ for expert_idx in range(self.num_routed):
+ selected_mask = topk_indices == expert_idx
+ if selected_mask.sum() == 0:
+ continue
+ n_idx, top = torch.where(selected_mask)
+ # Fancy indexing can yield non-contiguous rows; cuBLAS bf16 GEMM may then fail
+ # with ``CUBLAS_STATUS_INVALID_VALUE`` inside ``F.linear``.
+ x_selected = x_flat[n_idx].contiguous()
+ expert_output = self.experts[expert_idx](x_selected)
+ contrib = expert_output * routed_weights[n_idx, top].unsqueeze(-1)
+ aggregated_output[n_idx] = aggregated_output[n_idx] + contrib
+ aggregated_gate[n_idx] = aggregated_gate[n_idx] + routed_weights[n_idx, top]
+
+ aggregated_output = aggregated_output / (
+ aggregated_gate.unsqueeze(-1) + epsilon
+ )
+ return aggregated_output
+
+
+def _ceil_to_multiple(value: int, multiple: int) -> int:
+ return ((value + multiple - 1) // multiple) * multiple
+
+
+def _per_block_cast_to_fp8_padded(
+ weight: torch.Tensor,
+ block_size: int = 128,
+) -> tuple[torch.Tensor, torch.Tensor]:
+ rows, columns = weight.shape
+ padded = weight.new_zeros(
+ _ceil_to_multiple(rows, block_size),
+ _ceil_to_multiple(columns, block_size),
+ )
+ padded[:rows, :columns] = weight
+ return per_block_cast_to_fp8(
+ padded.contiguous(),
+ use_ue8m0=False,
+ gran_k=block_size,
+ )
+
+
+class MoESwiGLUFFNFP8(MoESwiGLUFFN):
+ """NOTE vision MoE using the checkpoint's local block-FP8 semantics."""
+
+ @torch.no_grad()
+ def process_weights_after_loading(self) -> None:
+ if hasattr(self, "_fused_w13_fp8"):
+ return
+
+ w13_weights = []
+ w13_scales = []
+ w2_weights = []
+ w2_scales = []
+ for expert in self.experts:
+ w1, s1 = _per_block_cast_to_fp8_padded(expert.fc1.weight)
+ w3, s3 = _per_block_cast_to_fp8_padded(expert.fc3.weight)
+ w2, s2 = _per_block_cast_to_fp8_padded(expert.fc2.weight)
+ w13_weights.append(torch.cat((w1, w3), dim=0))
+ w13_scales.append(torch.cat((s1, s3), dim=0))
+ w2_weights.append(w2)
+ w2_scales.append(s2)
+
+ self.register_buffer(
+ "_fused_w13_fp8",
+ torch.stack(w13_weights).contiguous(),
+ persistent=False,
+ )
+ self.register_buffer(
+ "_fused_w13_scale",
+ torch.stack(w13_scales).contiguous(),
+ persistent=False,
+ )
+ self.register_buffer(
+ "_fused_w2_fp8",
+ torch.stack(w2_weights).contiguous(),
+ persistent=False,
+ )
+ self.register_buffer(
+ "_fused_w2_scale",
+ torch.stack(w2_scales).contiguous(),
+ persistent=False,
+ )
+ del self.experts
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ if not hasattr(self, "_fused_w13_fp8"):
+ raise RuntimeError("NOTE vision FP8 weights were not initialized")
+
+ gate_logits = F.linear(x.float(), self.gate_weight.float())
+ if self.router_scoring_func == "sigmoid":
+ scores = torch.sigmoid(gate_logits)
+ else:
+ scores = torch.softmax(gate_logits, dim=-1, dtype=torch.float32)
+
+ topk = min(int(self.capacity_factor), self.num_routed)
+ biased_scores = scores + self.router_bias.float().unsqueeze(0)
+ topk_ids = torch.topk(biased_scores, k=topk, dim=-1, sorted=False).indices
+ topk_weights = scores.gather(1, topk_ids)
+ if self.router_scoring_func == "sigmoid" and topk > 1:
+ topk_weights = topk_weights / (
+ topk_weights.sum(dim=-1, keepdim=True) + 1e-9
+ )
+ topk_weights = topk_weights * self.router_scale
+
+ output = note_vision_fused_moe_fp8(
+ x.contiguous(),
+ self._fused_w13_fp8,
+ self._fused_w2_fp8,
+ topk_weights,
+ topk_ids.to(torch.int32),
+ self._fused_w13_scale,
+ self._fused_w2_scale,
+ )
+ denominator = topk_weights.sum(dim=-1, keepdim=True).clamp_min(1e-9)
+ return (output / denominator).type_as(x)
+
+
+# ---- PatchEmbed ----
+
+
+class DotsPatchEmbed(nn.Module):
+ def __init__(self, config: DotsMoEVitConfig):
+ super().__init__()
+ self.num_channels = config.num_channels
+ self.patch_size = config.patch_size
+ self.temporal_patch_size = config.temporal_patch_size
+ self.embed_dim = config.embed_dim
+ self.proj = nn.Conv2d(
+ config.num_channels,
+ config.embed_dim,
+ kernel_size=(config.patch_size, config.patch_size),
+ stride=(config.patch_size, config.patch_size),
+ )
+ self.norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ x = x.view(
+ -1,
+ self.num_channels,
+ self.temporal_patch_size,
+ self.patch_size,
+ self.patch_size,
+ )[:, :, 0]
+ x = self.proj(x).view(-1, self.embed_dim)
+ x = self.norm(x)
+ return x
+
+
+# ---- Block ----
+
+
+class MoEVisionBlock(nn.Module):
+ def __init__(self, config: DotsMoEVitConfig, layer_number: int):
+ super().__init__()
+ attn_impl = resolve_attn_implementation(
+ config.attn_implementation, eager_fallback="eager_v2"
+ )
+ self.attn = build_vision_attention(attn_impl, config, eager_fallback="eager_v2")
+ self._attn_uses_seqlens = attn_uses_seqlens(attn_impl)
+ self.norm_1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
+ self.norm_2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
+
+ is_moe = (
+ config.pyramid_num_routed
+ and layer_number < len(config.pyramid_num_routed)
+ and config.pyramid_num_routed[layer_number] > 0
+ )
+ if is_moe:
+ mlp_cls = MoESwiGLUFFNFP8 if config.enable_fp8_moe else MoESwiGLUFFN
+ self.mlp = mlp_cls(config, layer_number)
+ else:
+ self.mlp = DotsSwiGLUFFN(
+ config.embed_dim, config.intermediate_size, bias=config.use_bias
+ )
+
+ def forward(
+ self,
+ hidden_states,
+ cu_seqlens,
+ rotary_pos_emb,
+ max_seqlen: int,
+ seqlens: list[int] | None = None,
+ ) -> torch.Tensor:
+ hidden_states = apply_vision_attention_residual(
+ self.attn,
+ self.norm_1,
+ hidden_states,
+ cu_seqlens,
+ max_seqlen,
+ rotary_pos_emb,
+ seqlens=seqlens,
+ uses_seqlens=self._attn_uses_seqlens,
+ )
+ hidden_states = hidden_states + self.mlp(self.norm_2(hidden_states))
+ return hidden_states
+
+
+# ---- Adapter (pixel_shuffle + MLP) ----
+
+
+def _pixel_shuffle(x, scale_factor=0.5):
+ if x.size(1) % 2 == 1:
+ x = torch.cat([x[:, :1], x], dim=1)
+ if x.size(2) % 2 == 1:
+ x = torch.cat([x[:, :, :1], x], dim=2)
+ n, h, w, c = x.size()
+ x = x.reshape(n, h, int(w * scale_factor), int(c / scale_factor))
+ x = x.permute(0, 2, 1, 3).contiguous()
+ x = x.reshape(
+ n,
+ int(w * scale_factor),
+ int(h * scale_factor),
+ int(c / (scale_factor * scale_factor)),
+ )
+ x = x.permute(0, 2, 1, 3).contiguous()
+ return x
+
+
+class PixelShuffleAdapter(nn.Module):
+ """Legacy adapter: NHWC pixel-shuffle spatial merge + LayerNorm + 2-layer MLP.
+
+ Mirrors ``cybertron`` ``FCAdapter(pool_kind='pixel_shuffle', proj_kind='mlp2x_ln_gelu')``.
+ State-dict keys: ``proj.0`` (LayerNorm of in_dim*merge**2), ``proj.1`` / ``proj.3`` (Linear).
+ """
+
+ def __init__(self, config: DotsMoEVitConfig):
+ super().__init__()
+ in_dim = config.adapter_in_dim
+ out_dim = config.adapter_out_dim
+ merge_size = config.adapter_merge_size
+ merged_dim = in_dim * merge_size**2
+ self.proj = nn.Sequential(
+ LayerNorm(merged_dim),
+ nn.Linear(merged_dim, out_dim),
+ nn.GELU(),
+ nn.Linear(out_dim, out_dim),
+ )
+
+ def forward(
+ self,
+ patch_embed: torch.Tensor,
+ grid_thw: torch.Tensor,
+ ) -> torch.Tensor:
+ assert patch_embed.dim() == 2 and grid_thw is not None
+ image_features = []
+ token_index = 0
+ for i in range(grid_thw.shape[0]):
+ grid_t, grid_h, grid_w = grid_thw[i]
+ images_token_length = grid_t * grid_h * grid_w
+ _pe = patch_embed[token_index : token_index + images_token_length]
+ token_index += images_token_length
+ if grid_t == 1:
+ _pe = _pe.reshape(int(grid_h), int(grid_w), -1).unsqueeze(0)
+ else:
+ _pe = _pe.reshape(int(grid_t), int(grid_h), int(grid_w), -1)
+ _pe = _pixel_shuffle(_pe, scale_factor=0.5)
+ _pe = _pe.squeeze(0) if grid_t == 1 else _pe.reshape(-1, _pe.shape[-1])
+ image_features.append(_pe.reshape(-1, _pe.shape[-1]))
+ out = torch.cat(image_features, dim=0)
+ out = self.proj(out)
+ return out
+
+
+class PatchMergerAdapter(nn.Module):
+ """Cybertron ``PatchMerger`` (``pool_kind='patch_merger', proj_kind='identity'``).
+
+ Assumes the encoder output is already laid out in ``merge_size``x``merge_size`` groups
+ (qwen ``pre_pixel_shuffle`` preprocessor + RoPE grouped accordingly), so merging is a
+ simple ``view(-1, merge**2 * in_dim)`` of consecutive tokens. State-dict layout matches
+ cybertron's ``PatchMerger`` (``ln_q`` over the per-token dim, ``mlp.0`` / ``mlp.2`` Linear).
+ """
+
+ def __init__(self, config: DotsMoEVitConfig):
+ super().__init__()
+ in_dim = config.adapter_in_dim
+ out_dim = config.adapter_out_dim
+ merge_size = config.adapter_merge_size
+ merged_dim = in_dim * merge_size**2
+ self.merge_size = merge_size
+ self.merged_dim = merged_dim
+ self.ln_q = LayerNorm(in_dim, eps=1e-6)
+ self.mlp = nn.Sequential(
+ nn.Linear(merged_dim, merged_dim),
+ nn.GELU(),
+ nn.Linear(merged_dim, out_dim),
+ )
+
+ def forward(
+ self,
+ patch_embed: torch.Tensor,
+ grid_thw: torch.Tensor,
+ ) -> torch.Tensor:
+ assert patch_embed.dim() == 2 and grid_thw is not None
+ x = self.ln_q(patch_embed)
+ x = x.reshape(-1, self.merged_dim)
+ return self.mlp(x)
+
+
+_ADAPTER_CLASSES: dict[str, type[nn.Module]] = {
+ "pixel_shuffle_mlp": PixelShuffleAdapter,
+ "patch_merger": PatchMergerAdapter,
+}
+
+
+# ---- Full Model ----
+
+
+class DotsMoEVitModel(PreTrainedModel):
+ config_class = DotsMoEVitConfig
+
+ def __init__(self, config: DotsMoEVitConfig) -> None:
+ super().__init__(config)
+ self.config = config
+ self.spatial_merge_size = config.spatial_merge_size
+
+ self.patch_embed = DotsPatchEmbed(config)
+
+ head_dim = config.embed_dim // config.num_attention_heads
+ self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2, cache_seq_len=100000)
+
+ self.blocks = nn.ModuleList(
+ [MoEVisionBlock(config, i) for i in range(config.num_hidden_layers)]
+ )
+
+ if config.post_norm:
+ self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
+
+ adapter_cls = _ADAPTER_CLASSES.get(config.adapter_type)
+ if adapter_cls is None:
+ raise ValueError(f"Unknown adapter_type {config.adapter_type!r}")
+ self.adapter = adapter_cls(config)
+
+ self.gradient_checkpointing = False
+ self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
+
+ @property
+ def dtype(self) -> torch.dtype:
+ mlp = self.blocks[0].mlp
+ if hasattr(mlp, "fc13"):
+ return mlp.fc13.weight.dtype
+ if hasattr(mlp, "fc1"):
+ return mlp.fc1.weight.dtype
+ expert = mlp.experts[0]
+ if hasattr(expert, "fc13"):
+ return expert.fc13.weight.dtype
+ return expert.fc1.weight.dtype
+
+ @property
+ def device(self) -> torch.device:
+ return self.patch_embed.proj.weight.device
+
+ def process_weights_after_loading(self) -> None:
+ for block in self.blocks:
+ if isinstance(block.mlp, MoESwiGLUFFNFP8):
+ block.mlp.process_weights_after_loading()
+ if self.config.enable_torch_compile:
+ self.compile_block_modules()
+
+ def compile_block_modules(self, **compile_kwargs: Any) -> None:
+ for block in self.blocks:
+ block.attn = torch.compile(block.attn, **compile_kwargs)
+ 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):
+ # 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``.
+ if self.config.pre_pixel_shuffle:
+ rope_merge_size = (
+ self.spatial_merge_size if self.spatial_merge_size > 1 else 2
+ )
+ 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)
+ )
+ hpos_ids = hpos_ids.reshape(
+ h // rope_merge_size,
+ rope_merge_size,
+ w // rope_merge_size,
+ rope_merge_size,
+ )
+ 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 = wpos_ids.reshape(
+ h // rope_merge_size,
+ rope_merge_size,
+ w // rope_merge_size,
+ rope_merge_size,
+ )
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten()
+ 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()
+ 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
+
+ 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 _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
+ )
+ 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
+ ) -> 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
+ )
+ 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",
+ )
+
+ 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,
+ )
+ else:
+ hidden_states = blk(
+ hidden_states,
+ cu_seqlens,
+ rotary_pos_emb,
+ max_seqlen,
+ seqlens=seqlens,
+ )
+
+ if self.config.post_norm:
+ hidden_states = self.post_trunk_norm(hidden_states)
+
+ hidden_states = self.adapter(hidden_states, grid_thw)
+ return hidden_states
diff --git a/vllm/models/dots3_note/nvidia/vision_attention.py b/vllm/models/dots3_note/nvidia/vision_attention.py
new file mode 100644
index 0000000000..38c9c5a146
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/vision_attention.py
@@ -0,0 +1,477 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+# ruff: noqa: E501
+"""Shared vision attention stack for Dots dense / MoE ViT encoders.
+
+Exports every attention backend (eager, eager_v2, sdpa, flash_attention_2,
+flash_attention_3), RoPE helpers, and block wiring utilities.
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported
+
+_flash_attn_available = is_fa_version_supported(2)
+_flash_attn3_available = is_fa_version_supported(3)
+_flash_attn3_varlen_func = flash_attn_varlen_func
+_flash_attn3_use_positional_args = False
+
+
+# ---------------------------------------------------------------------------
+# RoPE
+# ---------------------------------------------------------------------------
+
+
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
+ x1 = x[..., : x.shape[-1] // 2]
+ x2 = x[..., x.shape[-1] // 2 :]
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb_vision(
+ tensor: torch.Tensor, freqs: torch.Tensor
+) -> 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()
+ output = (tensor * cos) + (rotate_half(tensor) * sin)
+ return output.to(orig_dtype)
+
+
+class VisionRotaryEmbedding(nn.Module):
+ """2D vision RoPE frequency table.
+
+ When ``cache_seq_len`` is set, precomputes and reuses frequencies up to that
+ length (MoE default). Otherwise frequencies are computed on each forward
+ (dense default).
+ """
+
+ def __init__(
+ self,
+ dim: int,
+ theta: float = 10000.0,
+ cache_seq_len: int | None = None,
+ ) -> None:
+ super().__init__()
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
+ self._cache_seq_len = cache_seq_len
+ self.freqs_cache: torch.Tensor
+ if cache_seq_len is not None:
+ self.register_buffer(
+ "freqs_cache",
+ self._compute_freqs(cache_seq_len),
+ persistent=False,
+ )
+
+ def _compute_freqs(self, seqlen: int) -> torch.Tensor:
+ seq = torch.arange(
+ seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
+ )
+ return torch.outer(seq, self.inv_freq)
+
+ def forward(self, seqlen: int) -> torch.Tensor:
+ if self._cache_seq_len is None:
+ return self._compute_freqs(seqlen)
+ if seqlen > self.freqs_cache.shape[0]:
+ self.freqs_cache = self._compute_freqs(seqlen)
+ return self.freqs_cache[:seqlen]
+
+
+# ---------------------------------------------------------------------------
+# Attention modules
+# ---------------------------------------------------------------------------
+
+
+class _RMSNorm(nn.Module):
+ """Q/K norm inside attention; matches Dots ViT RMSNorm (fp32 reduce, then cast back)."""
+
+ def __init__(self, dim: int, eps: float = 1e-6):
+ super().__init__()
+ self.weight = nn.Parameter(torch.ones(dim))
+ self.eps = eps
+
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
+ output = self._norm(x.float()).type_as(x)
+ return output * self.weight
+
+ def _norm(self, x: torch.Tensor) -> torch.Tensor:
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+
+@dataclass(frozen=True)
+class VisionAttentionParams:
+ dim: int
+ num_heads: int
+ bias: bool = True
+ is_causal: bool = False
+ rms_norm_eps: float = 1e-5
+ use_qk_norm: bool = False
+
+
+def vision_attention_params_from_config(config: Any) -> VisionAttentionParams:
+ return VisionAttentionParams(
+ dim=config.embed_dim,
+ num_heads=config.num_attention_heads,
+ bias=getattr(config, "use_bias", True),
+ is_causal=config.is_causal,
+ rms_norm_eps=config.rms_norm_eps,
+ use_qk_norm=getattr(config, "use_qk_norm", False),
+ )
+
+
+class _VisionAttentionBase(nn.Module):
+ """QKV projection, optional Q/K norm, and output projection."""
+
+ def __init__(self, params: VisionAttentionParams) -> None:
+ super().__init__()
+ self.num_heads = params.num_heads
+ self.head_dim = params.dim // params.num_heads
+ self.is_causal = params.is_causal
+ self.use_qk_norm = params.use_qk_norm
+ self.qkv = nn.Linear(params.dim, params.dim * 3, bias=params.bias)
+ self.proj = nn.Linear(params.dim, params.dim, bias=params.bias)
+ if self.use_qk_norm:
+ self.q_norm = _RMSNorm(self.head_dim, eps=params.rms_norm_eps)
+ self.k_norm = _RMSNorm(self.head_dim, eps=params.rms_norm_eps)
+
+ def _qkv_with_rope(
+ self,
+ hidden_states: torch.Tensor,
+ rotary_pos_emb: torch.Tensor | None,
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+ seq_length = hidden_states.shape[0]
+ q, k, v = (
+ self.qkv(hidden_states)
+ .reshape(seq_length, 3, self.num_heads, -1)
+ .permute(1, 0, 2, 3)
+ .unbind(0)
+ )
+ if self.use_qk_norm:
+ q = self.q_norm(q)
+ k = self.k_norm(k)
+ 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
+
+
+class VisionAttention(_VisionAttentionBase):
+ """Eager attention with a dense block-diagonal mask (cu_seqlens boundaries)."""
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ rotary_pos_emb: torch.Tensor | None = None,
+ seqlens: list[int] | None = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
+
+ attention_mask = torch.full(
+ [1, seq_length, seq_length],
+ torch.finfo(q.dtype).min,
+ device=q.device,
+ dtype=q.dtype,
+ )
+ for i in range(1, len(cu_seqlens)):
+ attention_mask[
+ ...,
+ cu_seqlens[i - 1] : cu_seqlens[i],
+ cu_seqlens[i - 1] : cu_seqlens[i],
+ ] = 0
+
+ q = q.transpose(0, 1)
+ k = k.transpose(0, 1)
+ v = v.transpose(0, 1)
+ attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)
+ attn_weights = attn_weights + attention_mask
+ attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
+ attn_output = torch.matmul(attn_weights, v)
+ attn_output = attn_output.transpose(0, 1).reshape(seq_length, -1)
+ return self.proj(attn_output)
+
+
+class VisionAttentionV2(_VisionAttentionBase):
+ """Eager attention per varlen segment (lower peak memory than :class:`VisionAttention`)."""
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ rotary_pos_emb: torch.Tensor | None = None,
+ seqlens: list[int] | None = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
+
+ if seqlens is None:
+ seqlens = torch.diff(cu_seqlens).tolist()
+
+ outputs = []
+ for q_i, k_i, v_i in zip(
+ torch.split(q, seqlens, 0),
+ torch.split(k, seqlens, 0),
+ torch.split(v, seqlens, 0),
+ ):
+ q_i = q_i.transpose(0, 1)
+ k_i = k_i.transpose(0, 1)
+ v_i = v_i.transpose(0, 1)
+ out = torch.matmul(q_i, k_i.transpose(1, 2)) / math.sqrt(self.head_dim)
+ out = F.softmax(out, dim=-1, dtype=torch.float32).to(q.dtype)
+ out = torch.matmul(out, v_i).transpose(0, 1)
+ outputs.append(out)
+
+ attn_output = torch.concat(outputs, dim=0).reshape(seq_length, -1)
+ return self.proj(attn_output)
+
+
+# MoE checkpoints historically named this implementation ``VisionEagerAttention``.
+VisionEagerAttention = VisionAttentionV2
+
+
+class VisionFlashAttention2(_VisionAttentionBase):
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ rotary_pos_emb: torch.Tensor | None = None,
+ seqlens: list[int] | None = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
+ attn_output = flash_attn_varlen_func(
+ q,
+ k,
+ v,
+ cu_seqlens_q=cu_seqlens,
+ cu_seqlens_k=cu_seqlens,
+ max_seqlen_q=max_seqlen,
+ max_seqlen_k=max_seqlen,
+ causal=self.is_causal,
+ fa_version=2,
+ ).reshape(seq_length, -1)
+ return self.proj(attn_output)
+
+
+class VisionFlashAttention3(_VisionAttentionBase):
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ rotary_pos_emb: torch.Tensor | None = None,
+ seqlens: list[int] | None = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
+ if _flash_attn3_use_positional_args:
+ attn_output = _flash_attn3_varlen_func(
+ q,
+ k,
+ v,
+ cu_seqlens,
+ cu_seqlens,
+ max_seqlen,
+ max_seqlen,
+ causal=self.is_causal,
+ )
+ else:
+ result = _flash_attn3_varlen_func(
+ q,
+ k,
+ v,
+ cu_seqlens_q=cu_seqlens,
+ cu_seqlens_k=cu_seqlens,
+ max_seqlen_q=max_seqlen,
+ max_seqlen_k=max_seqlen,
+ causal=self.is_causal,
+ fa_version=3,
+ )
+ attn_output = result[0] if isinstance(result, tuple) else result
+ attn_output = attn_output.reshape(seq_length, -1)
+ return self.proj(attn_output)
+
+
+class VisionSdpaAttention(_VisionAttentionBase):
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ rotary_pos_emb: torch.Tensor | None = None,
+ seqlens: list[int] | None = None,
+ ) -> torch.Tensor:
+ seq_length = hidden_states.shape[0]
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
+
+ attention_mask = torch.zeros(
+ [1, seq_length, seq_length], device=q.device, dtype=torch.bool
+ )
+ for i in range(1, len(cu_seqlens)):
+ attention_mask[
+ ...,
+ cu_seqlens[i - 1] : cu_seqlens[i],
+ cu_seqlens[i - 1] : cu_seqlens[i],
+ ] = True
+
+ q = q.transpose(0, 1).unsqueeze(0)
+ k = k.transpose(0, 1).unsqueeze(0)
+ v = v.transpose(0, 1).unsqueeze(0)
+
+ if attention_mask.stride(-1) != 1:
+ attention_mask = torch.empty_like(
+ attention_mask, memory_format=torch.contiguous_format
+ ).copy_(attention_mask)
+
+ from torch.nn.attention import SDPBackend, sdpa_kernel
+
+ with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
+ attn_output = F.scaled_dot_product_attention(
+ q, k, v, attention_mask, dropout_p=0.0
+ )
+
+ attn_output = attn_output.squeeze(0).transpose(0, 1).reshape(seq_length, -1)
+ return self.proj(attn_output)
+
+
+DOTS_VISION_ATTENTION_CLASSES: dict[str, type[nn.Module]] = {
+ "eager": VisionAttention,
+ "eager_v2": VisionAttentionV2,
+ "flash_attention_2": VisionFlashAttention2,
+ "flash_attention_3": VisionFlashAttention3,
+ "sdpa": VisionSdpaAttention,
+}
+
+
+def resolve_attn_implementation(
+ attn_implementation: str,
+ *,
+ eager_fallback: str = "eager",
+) -> str:
+ """Apply FA3 → FA2 → eager fallback when backends are missing.
+
+ ``eager_fallback`` selects which eager backend is used when flash-attn is
+ unavailable. MoE historically used per-segment eager (``eager_v2``); dense
+ uses full-mask ``eager``.
+ """
+ if attn_implementation == "flash_attention_3" and not _flash_attn3_available:
+ print(
+ "flash attention 3 not available! fallback to flash attention 2 implementation"
+ )
+ attn_implementation = "flash_attention_2"
+ if attn_implementation == "flash_attention_2" and not _flash_attn_available:
+ print("flash attention 2 not available! fallback to eager implementation")
+ attn_implementation = eager_fallback
+ return attn_implementation
+
+
+def build_vision_attention(
+ attn_implementation: str,
+ config: Any,
+ *,
+ dim: int | None = None,
+ num_heads: int | None = None,
+ bias: bool | None = None,
+ eager_fallback: str = "eager",
+) -> nn.Module:
+ """Instantiate a vision attention module from ``attn_implementation`` and config."""
+ attn_implementation = resolve_attn_implementation(
+ attn_implementation, eager_fallback=eager_fallback
+ )
+ if attn_implementation not in DOTS_VISION_ATTENTION_CLASSES:
+ raise ValueError(
+ f"Unknown attn_implementation {attn_implementation!r}; "
+ f"expected one of {sorted(DOTS_VISION_ATTENTION_CLASSES)}"
+ )
+ if dim is None:
+ params = vision_attention_params_from_config(config)
+ else:
+ params = VisionAttentionParams(
+ dim=dim,
+ num_heads=num_heads
+ if num_heads is not None
+ else config.num_attention_heads,
+ bias=bias if bias is not None else getattr(config, "use_bias", True),
+ is_causal=config.is_causal,
+ rms_norm_eps=config.rms_norm_eps,
+ use_qk_norm=getattr(config, "use_qk_norm", False),
+ )
+ return DOTS_VISION_ATTENTION_CLASSES[attn_implementation](params)
+
+
+def attn_uses_seqlens(attn_implementation: str) -> bool:
+ return attn_implementation == "eager_v2"
+
+
+def prepare_seqlens_for_attention(
+ attn_implementation: str,
+ cu_seqlens: torch.Tensor,
+ *,
+ eager_fallback: str = "eager",
+) -> list[int] | None:
+ """Return per-segment lengths when the resolved backend needs ``seqlens``."""
+ resolved = resolve_attn_implementation(
+ attn_implementation, eager_fallback=eager_fallback
+ )
+ if attn_uses_seqlens(resolved):
+ return torch.diff(cu_seqlens).tolist()
+ return None
+
+
+def apply_vision_attention_residual(
+ attn: nn.Module,
+ norm: nn.Module,
+ hidden_states: torch.Tensor,
+ cu_seqlens: torch.Tensor,
+ max_seqlen: int,
+ rotary_pos_emb: torch.Tensor,
+ *,
+ seqlens: list[int] | None = None,
+ uses_seqlens: bool = False,
+) -> torch.Tensor:
+ """Pre-norm residual attention used by dense / MoE vision blocks."""
+ attn_kwargs: dict[str, Any] = {"rotary_pos_emb": rotary_pos_emb}
+ if uses_seqlens:
+ attn_kwargs["seqlens"] = seqlens
+ return hidden_states + attn(
+ norm(hidden_states),
+ cu_seqlens,
+ max_seqlen,
+ **attn_kwargs,
+ )
+
+
+__all__ = [
+ "DOTS_VISION_ATTENTION_CLASSES",
+ "VisionAttention",
+ "VisionAttentionParams",
+ "VisionAttentionV2",
+ "VisionEagerAttention",
+ "VisionFlashAttention2",
+ "VisionFlashAttention3",
+ "VisionRotaryEmbedding",
+ "VisionSdpaAttention",
+ "apply_rotary_pos_emb_vision",
+ "apply_vision_attention_residual",
+ "attn_uses_seqlens",
+ "build_vision_attention",
+ "prepare_seqlens_for_attention",
+ "resolve_attn_implementation",
+ "rotate_half",
+ "vision_attention_params_from_config",
+]
diff --git a/vllm/models/dots3_note/nvidia/vision_moe.py b/vllm/models/dots3_note/nvidia/vision_moe.py
new file mode 100644
index 0000000000..5f9564351c
--- /dev/null
+++ b/vllm/models/dots3_note/nvidia/vision_moe.py
@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""NOTE vision MoE execution matching the native encoder's FP8 semantics."""
+
+import torch
+
+from vllm import _custom_ops as ops
+from vllm.model_executor.layers.fused_moe.activation import (
+ MoEActivation,
+ apply_moe_activation,
+)
+from vllm.model_executor.layers.fused_moe.fused_moe import (
+ _prepare_expert_assignment,
+ dispatch_fused_moe_kernel,
+ try_get_optimal_moe_config,
+)
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
+ per_token_group_quant_fp8,
+)
+from vllm.triton_utils import tl
+
+_BLOCK_SHAPE = [128, 128]
+
+
+def note_vision_fused_moe_fp8(
+ hidden_states: torch.Tensor,
+ w13: torch.Tensor,
+ w2: torch.Tensor,
+ topk_weights: torch.Tensor,
+ topk_ids: torch.Tensor,
+ w13_scale: torch.Tensor,
+ w2_scale: torch.Tensor,
+) -> torch.Tensor:
+ """Run top-k SwiGLU MoE with the encoder's block-FP8 formula.
+
+ Unlike vLLM's generic functional MoE path, the native NOTE encoder keeps
+ dynamic activation scales as FP32 instead of rounding them to E8M0. Weight
+ scales are block 128x128 and activations are quantized per token/group-128.
+ """
+ if hidden_states.dtype != torch.bfloat16:
+ raise TypeError(
+ "NOTE vision FP8 MoE requires bfloat16 activations, got "
+ f"{hidden_states.dtype}"
+ )
+
+ num_tokens = hidden_states.shape[0]
+ num_experts, intermediate_size, _ = w13.shape
+ output_size = w2.shape[1]
+ topk = topk_ids.shape[1]
+ config = try_get_optimal_moe_config(
+ w13.shape,
+ w2.shape,
+ topk,
+ "fp8_w8a8",
+ num_tokens,
+ block_shape=_BLOCK_SHAPE,
+ )
+
+ cache = torch.empty(
+ num_tokens * topk * max(intermediate_size, output_size),
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
+ )
+ first_gemm = cache[: num_tokens * topk * intermediate_size].view(
+ num_tokens, topk, intermediate_size
+ )
+ second_gemm = cache[: num_tokens * topk * output_size].view(
+ num_tokens, topk, output_size
+ )
+ activated_size = intermediate_size // 2
+ activated = torch.empty(
+ (num_tokens * topk, activated_size),
+ dtype=hidden_states.dtype,
+ device=hidden_states.device,
+ )
+
+ quantized_input, input_scale = per_token_group_quant_fp8(
+ hidden_states,
+ _BLOCK_SHAPE[1],
+ use_ue8m0=False,
+ )
+ sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
+ topk_ids,
+ config,
+ num_tokens,
+ topk,
+ num_experts,
+ None,
+ block_shape=_BLOCK_SHAPE,
+ )
+ dispatch_fused_moe_kernel(
+ quantized_input,
+ w13,
+ first_gemm,
+ input_scale,
+ w13_scale,
+ None,
+ topk_weights,
+ sorted_token_ids,
+ expert_ids,
+ num_tokens_post_padded,
+ False,
+ topk,
+ config,
+ tl.bfloat16,
+ True,
+ False,
+ False,
+ False,
+ False,
+ _BLOCK_SHAPE,
+ )
+
+ apply_moe_activation(
+ MoEActivation.SILU,
+ activated,
+ first_gemm.view(-1, intermediate_size),
+ )
+ quantized_activated, activated_scale = per_token_group_quant_fp8(
+ activated,
+ _BLOCK_SHAPE[1],
+ use_ue8m0=False,
+ )
+ dispatch_fused_moe_kernel(
+ quantized_activated,
+ w2,
+ second_gemm,
+ activated_scale,
+ w2_scale,
+ None,
+ topk_weights,
+ sorted_token_ids,
+ expert_ids,
+ num_tokens_post_padded,
+ True,
+ 1,
+ config,
+ tl.bfloat16,
+ True,
+ False,
+ False,
+ False,
+ False,
+ _BLOCK_SHAPE,
+ )
+
+ output = torch.empty_like(hidden_states)
+ ops.moe_sum(second_gemm, output)
+ return output
diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py
index 2292f25a85..535e8e2781 100644
--- a/vllm/tool_parsers/__init__.py
+++ b/vllm/tool_parsers/__init__.py
@@ -22,6 +22,10 @@ Example:
_TOOL_PARSERS_TO_REGISTER = {
+ "dots": (
+ "dots_tool_parser",
+ "DotsToolParser",
+ ),
"deepseek_v3": ( # name
"deepseekv3_tool_parser", # filename
"DeepSeekV3ToolParser", # class_name
diff --git a/vllm/tool_parsers/dots_tool_parser.py b/vllm/tool_parsers/dots_tool_parser.py
new file mode 100644
index 0000000000..987d891d8a
--- /dev/null
+++ b/vllm/tool_parsers/dots_tool_parser.py
@@ -0,0 +1,434 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import json
+from collections.abc import Sequence
+from typing import Any
+
+import regex as re
+
+from vllm.entrypoints.chat_utils import make_tool_call_id
+from vllm.entrypoints.openai.chat_completion.protocol import (
+ ChatCompletionRequest,
+ ChatCompletionToolsParam,
+)
+from vllm.entrypoints.openai.engine.protocol import (
+ DeltaFunctionCall,
+ DeltaMessage,
+ DeltaToolCall,
+ ExtractedToolCallInformation,
+ FunctionCall,
+ ToolCall,
+)
+from vllm.logger import init_logger
+from vllm.tokenizers import TokenizerLike
+from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
+from vllm.tool_parsers.utils import is_complete_json, partial_tag_overlap
+
+logger = init_logger(__name__)
+
+
+class DotsToolParser(ToolParser):
+ """Parse Dots tool calls in their XML wrapper format.
+
+ The canonical body contains one or more ``invoke`` elements::
+
+
+
+ weather in Shanghai
+
+
+
+ A JSON object with ``name`` and ``arguments`` is also accepted as a
+ fallback. Multiple wrapper blocks and multiple invokes per block are
+ supported.
+ """
+
+ supports_required_and_named = False
+
+ tool_call_start_token = ""
+ tool_call_end_token = ""
+
+ _block_regex = re.compile(
+ rf"{re.escape(tool_call_start_token)}\s*(.*?)\s*"
+ rf"{re.escape(tool_call_end_token)}",
+ re.DOTALL,
+ )
+ _invoke_regex = re.compile(
+ r"[^>]+)>(?P.*?)",
+ re.DOTALL,
+ )
+ _parameter_regex = re.compile(
+ r"[^>]+)>(?P.*?)",
+ re.DOTALL,
+ )
+
+ def __init__(
+ self,
+ tokenizer: TokenizerLike,
+ tools: list[Tool] | None = None,
+ ) -> None:
+ super().__init__(tokenizer, tools)
+ self._buffer = ""
+
+ @staticmethod
+ def _extract_name(value: str) -> str:
+ value = value.strip()
+ if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
+ return value[1:-1]
+ return value
+
+ @staticmethod
+ def _convert_param_value(value: str, param_type: Any) -> Any:
+ if value.lower() == "null":
+ return None
+
+ if isinstance(param_type, list):
+ param_type = next((item for item in param_type if item != "null"), "string")
+ if not isinstance(param_type, str):
+ param_type = str(param_type)
+ param_type = param_type.lower()
+
+ if param_type in {"string", "str", "text"}:
+ return value
+ if param_type in {"integer", "int"}:
+ try:
+ return int(value)
+ except (TypeError, ValueError):
+ return value
+ if param_type in {"number", "float"}:
+ try:
+ number = float(value)
+ return int(number) if number.is_integer() else number
+ except (TypeError, ValueError):
+ return value
+ if param_type in {"boolean", "bool"}:
+ return value.lower() in {"true", "1"}
+
+ try:
+ return json.loads(value)
+ except (json.JSONDecodeError, TypeError, ValueError):
+ return value
+
+ def _resolve_param_type(
+ self,
+ schema: Any,
+ defs: dict[str, Any],
+ depth: int = 0,
+ ) -> Any | None:
+ if not isinstance(schema, dict) or depth > 10:
+ return None
+ if "type" in schema:
+ return schema["type"]
+
+ ref = schema.get("$ref")
+ if isinstance(ref, str) and ref.startswith("#/$defs/"):
+ return self._resolve_param_type(
+ defs.get(ref.rsplit("/", 1)[-1]), defs, depth + 1
+ )
+
+ for keyword in ("anyOf", "oneOf", "allOf"):
+ alternatives = schema.get(keyword)
+ if not isinstance(alternatives, list):
+ continue
+ for alternative in alternatives:
+ if isinstance(alternative, dict) and alternative.get("type") == "null":
+ continue
+ resolved = self._resolve_param_type(alternative, defs, depth + 1)
+ if resolved is not None:
+ return resolved
+ return None
+
+ @staticmethod
+ def _tool_schema(
+ name: str,
+ tools: list[ChatCompletionToolsParam] | None,
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
+ for tool in tools or []:
+ if tool.function.name != name:
+ continue
+ schema = tool.function.parameters
+ if not isinstance(schema, dict):
+ break
+ properties = schema.get("properties", {})
+ defs = schema.get("$defs", {})
+ return (
+ properties if isinstance(properties, dict) else {},
+ defs if isinstance(defs, dict) else {},
+ )
+ return {}, {}
+
+ def _parse_xml_invoke(
+ self,
+ match: re.Match[str],
+ tools: list[ChatCompletionToolsParam] | None,
+ ) -> dict[str, Any]:
+ name = self._extract_name(match.group("name"))
+ properties, defs = self._tool_schema(name, tools)
+ arguments: dict[str, Any] = {}
+
+ for parameter in self._parameter_regex.finditer(match.group("body")):
+ param_name = self._extract_name(parameter.group("name"))
+ value = parameter.group("value").strip()
+ param_type: Any = "string"
+ if param_name in properties:
+ param_type = (
+ self._resolve_param_type(properties[param_name], defs) or "string"
+ )
+ arguments[param_name] = self._convert_param_value(value, param_type)
+
+ return {"name": name, "arguments": arguments}
+
+ def _parse_block(
+ self,
+ content: str,
+ tools: list[ChatCompletionToolsParam] | None,
+ ) -> list[dict[str, Any]]:
+ content = content.strip()
+ if content.startswith(" set[str]:
+ return {tool.function.name for tool in tools or []}
+
+ def _validated_call(
+ self,
+ parsed: dict[str, Any],
+ tools: list[ChatCompletionToolsParam] | None,
+ ) -> tuple[str, dict[str, Any]] | None:
+ name = parsed.get("name")
+ if not isinstance(name, str) or name not in self._known_tool_names(tools):
+ return None
+ arguments = parsed.get("arguments", parsed.get("parameters", {})) or {}
+ if not isinstance(arguments, dict):
+ return None
+ return name, arguments
+
+ def extract_tool_calls(
+ self,
+ model_output: str,
+ request: ChatCompletionRequest,
+ ) -> ExtractedToolCallInformation:
+ marker_index = model_output.find(self.tool_call_start_token)
+ if marker_index == -1:
+ return ExtractedToolCallInformation(
+ tools_called=False,
+ tool_calls=[],
+ content=model_output,
+ )
+
+ tool_calls: list[ToolCall] = []
+ for block in self._block_regex.finditer(model_output):
+ try:
+ for parsed in self._parse_block(block.group(1), request.tools):
+ validated = self._validated_call(parsed, request.tools)
+ if validated is None:
+ continue
+ name, arguments = validated
+ tool_calls.append(
+ ToolCall(
+ function=FunctionCall(
+ name=name,
+ arguments=json.dumps(arguments, ensure_ascii=False),
+ )
+ )
+ )
+ except (json.JSONDecodeError, TypeError, ValueError) as exc:
+ logger.warning("Failed to parse Dots tool call: %s", exc)
+
+ normal_text = model_output[:marker_index].strip()
+ if not tool_calls:
+ return ExtractedToolCallInformation(
+ tools_called=False,
+ tool_calls=[],
+ content=model_output,
+ )
+ return ExtractedToolCallInformation(
+ tools_called=True,
+ tool_calls=tool_calls,
+ content=normal_text or None,
+ )
+
+ def _append_complete_stream_call(
+ self,
+ name: str,
+ arguments: dict[str, Any],
+ tool_calls: list[DeltaToolCall],
+ ) -> None:
+ self.current_tool_id += 1
+ serialized = json.dumps(arguments, ensure_ascii=False)
+ self.prev_tool_call_arr.append({"name": name, "arguments": arguments})
+ self.streamed_args_for_tool.append(serialized)
+ tool_calls.append(
+ DeltaToolCall(
+ index=self.current_tool_id,
+ id=make_tool_call_id(),
+ type="function",
+ function=DeltaFunctionCall(name=name, arguments=serialized),
+ )
+ )
+
+ def _stream_complete_json_body(
+ self,
+ tools: list[ChatCompletionToolsParam] | None,
+ tool_calls: list[DeltaToolCall],
+ ) -> None:
+ content = self._buffer[len(self.tool_call_start_token) :].strip()
+ if not content or not is_complete_json(content):
+ return
+
+ try:
+ parsed = json.loads(content)
+ except (json.JSONDecodeError, TypeError, ValueError):
+ return
+ if not isinstance(parsed, dict):
+ return
+ validated = self._validated_call(parsed, tools)
+ if validated is None:
+ return
+
+ name, arguments = validated
+ serialized = json.dumps(arguments, ensure_ascii=False)
+ if not self.current_tool_name_sent:
+ self.current_tool_id += 1
+ tool_calls.append(
+ DeltaToolCall(
+ index=self.current_tool_id,
+ id=make_tool_call_id(),
+ type="function",
+ function=DeltaFunctionCall(name=name, arguments=""),
+ )
+ )
+ self.prev_tool_call_arr.append({"name": name, "arguments": arguments})
+ self.streamed_args_for_tool.append("")
+ self.current_tool_name_sent = True
+
+ streamed = self.streamed_args_for_tool[self.current_tool_id]
+ if serialized.startswith(streamed):
+ argument_diff = serialized[len(streamed) :]
+ if argument_diff:
+ tool_calls.append(
+ DeltaToolCall(
+ index=self.current_tool_id,
+ function=DeltaFunctionCall(arguments=argument_diff),
+ )
+ )
+ self.streamed_args_for_tool[self.current_tool_id] += argument_diff
+
+ def extract_tool_calls_streaming(
+ self,
+ previous_text: str,
+ current_text: str,
+ delta_text: str,
+ previous_token_ids: Sequence[int],
+ current_token_ids: Sequence[int],
+ delta_token_ids: Sequence[int],
+ request: ChatCompletionRequest,
+ ) -> DeltaMessage | None:
+ del current_text, previous_token_ids, current_token_ids, delta_token_ids
+ if not previous_text:
+ self._buffer = ""
+ self.prev_tool_call_arr = []
+ self.current_tool_id = -1
+ self.current_tool_name_sent = False
+ self.streamed_args_for_tool = []
+
+ self._buffer += delta_text
+ normal_parts: list[str] = []
+ tool_calls: list[DeltaToolCall] = []
+
+ while self._buffer:
+ marker_index = self._buffer.find(self.tool_call_start_token)
+ if marker_index == -1:
+ partial_len = partial_tag_overlap(
+ self._buffer, self.tool_call_start_token
+ )
+ if partial_len:
+ normal_parts.append(self._buffer[:-partial_len])
+ self._buffer = self._buffer[-partial_len:]
+ else:
+ normal_parts.append(self._buffer)
+ self._buffer = ""
+ normal_parts = [
+ part.replace(self.tool_call_end_token, "") for part in normal_parts
+ ]
+ break
+
+ if marker_index > 0:
+ normal_parts.append(self._buffer[:marker_index])
+ self._buffer = self._buffer[marker_index:]
+
+ end_index = self._buffer.find(
+ self.tool_call_end_token, len(self.tool_call_start_token)
+ )
+ if end_index == -1:
+ self._stream_complete_json_body(request.tools, tool_calls)
+ break
+
+ content = self._buffer[len(self.tool_call_start_token) : end_index]
+ self._buffer = self._buffer[end_index + len(self.tool_call_end_token) :]
+ try:
+ parsed_calls = self._parse_block(content, request.tools)
+ if not parsed_calls:
+ raise ValueError("Dots tool-call block contains no invoke")
+
+ block_had_streamed_call = self.current_tool_name_sent
+ valid_calls = [
+ validated
+ for parsed in parsed_calls
+ if (validated := self._validated_call(parsed, request.tools))
+ is not None
+ ]
+ if self.current_tool_name_sent and valid_calls:
+ name, arguments = valid_calls.pop(0)
+ serialized = json.dumps(arguments, ensure_ascii=False)
+ streamed = self.streamed_args_for_tool[self.current_tool_id]
+ if serialized.startswith(streamed):
+ remaining = serialized[len(streamed) :]
+ if remaining:
+ tool_calls.append(
+ DeltaToolCall(
+ index=self.current_tool_id,
+ function=DeltaFunctionCall(arguments=remaining),
+ )
+ )
+ self.prev_tool_call_arr[self.current_tool_id] = {
+ "name": name,
+ "arguments": arguments,
+ }
+ self.streamed_args_for_tool[self.current_tool_id] = serialized
+
+ for name, arguments in valid_calls:
+ self._append_complete_stream_call(name, arguments, tool_calls)
+
+ if not valid_calls and not block_had_streamed_call:
+ normal_parts.append(content.strip())
+ except (json.JSONDecodeError, TypeError, ValueError) as exc:
+ logger.warning("Failed to parse streamed Dots tool call: %s", exc)
+ normal_parts.append(content.strip())
+
+ self.current_tool_name_sent = False
+
+ content_delta = "".join(normal_parts)
+ if not content_delta and not tool_calls:
+ return None
+ return DeltaMessage(content=content_delta or None, tool_calls=tool_calls)
+
+ def flush_pending_normal_text(self) -> str:
+ """Return a partial opening marker as text when generation ends."""
+ if not self._buffer or self.tool_call_start_token in self._buffer:
+ return ""
+ normal_text = self._buffer.replace(self.tool_call_end_token, "")
+ self._buffer = ""
+ return normal_text
diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py
index ad8f0535b2..8bbba8df24 100644
--- a/vllm/transformers_utils/config.py
+++ b/vllm/transformers_utils/config.py
@@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
deepseek_vl_v2="DeepseekVLV2Config",
deepseek_v32="DeepseekV3Config",
deepseek_v4="DeepseekV4Config",
+ dots3_note="Dots3NoteConfig",
k3_dspark="K3DSparkConfig",
flex_olmo="FlexOlmoConfig",
fireredlid="FireRedLIDConfig",
diff --git a/vllm/transformers_utils/configs/__init__.py b/vllm/transformers_utils/configs/__init__.py
index 6685ee4f78..a6d477831d 100644
--- a/vllm/transformers_utils/configs/__init__.py
+++ b/vllm/transformers_utils/configs/__init__.py
@@ -35,6 +35,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
"DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma",
"DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
"DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
+ "Dots3NoteConfig": "vllm.transformers_utils.configs.dots3_note",
"K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark",
"DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
"EAGLEConfig": "vllm.transformers_utils.configs.eagle",
@@ -127,6 +128,7 @@ __all__ = [
"DeepseekVLV2Config",
"DeepseekV3Config",
"DeepseekV4Config",
+ "Dots3NoteConfig",
"K3DSparkConfig",
"DotsOCRConfig",
"EAGLEConfig",
diff --git a/vllm/transformers_utils/configs/dots3_note.py b/vllm/transformers_utils/configs/dots3_note.py
new file mode 100644
index 0000000000..832761544a
--- /dev/null
+++ b/vllm/transformers_utils/configs/dots3_note.py
@@ -0,0 +1,25 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from transformers import DeepseekV3Config
+
+
+class Dots3NoteConfig(DeepseekV3Config):
+ """Configuration for the Dots3Note hybrid DSA/SWA model."""
+
+ model_type = "dots3_note"
+
+ def __init__(self, **kwargs):
+ # These fields are absent from existing Note config.json files. Do
+ # not inherit DeepSeek-V3's 8-group/4-group router defaults: Note was
+ # trained with an ungrouped (1/1) noaux_tc router, as defined by the
+ # reference Dots3Config in SGLang haimain. A different grouping
+ # changes the selected experts at every MoE layer.
+ kwargs.setdefault("n_group", 1)
+ kwargs.setdefault("topk_group", 1)
+ # Dots3 projects the indexer RoPE coordinates in adjacent (GPT-J)
+ # pairs. DeepSeek-V3.2 defaults to the split-half NeoX layout when
+ # this flag is absent, which rotates different learned coordinates.
+ kwargs.setdefault("indexer_rope_interleave", True)
+ kwargs.setdefault("num_nextn_predict_layers", 1)
+ super().__init__(**kwargs)
diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py
index 2e6bd05cd9..4b65e2584c 100644
--- a/vllm/transformers_utils/model_arch_config_convertor.py
+++ b/vllm/transformers_utils/model_arch_config_convertor.py
@@ -317,6 +317,7 @@ class ModelArchConfigConvertorBase:
"deepseek_v3",
"deepseek_v32",
"deepseek_v4",
+ "dots3_note",
"deepseek_mtp",
"k3_dspark",
"glm_moe_dsa",