[Model] Add native Dots3 NOTE multimodal support (#51255)
Signed-off-by: KurodaKanbei <mistergalahad@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
范裕达 committed
Aug 12, 2026 at 23:13 UTC
9035151d6c9fb726181469f9e6aa9ccbf9a5dacb
28 files changed
+6468
-7
tests/kernels/attention/test_rocm_aiter_mla_causal_verify_mask.py
+8
-1
@@ -127,7 +127,14 @@ def _run_verify_block():
127
# stub with the attribute is enough for metadata construction.
128
layer_name = "placeholder"
129
vllm_config.compilation_config.static_forward_context[layer_name] = (
130
- types.SimpleNamespace(prefill_backend=torch.empty((1,)))
130
+ types.SimpleNamespace(
131
+ prefill_backend=torch.empty((1,)),
132
+ q_lora_rank=None,
133
+ kv_lora_rank=KV_LORA_RANK,
134
+ qk_nope_head_dim=QK_NOPE_HEAD_DIM,
135
+ qk_rope_head_dim=QK_ROPE_HEAD_DIM,
136
+ v_head_dim=V_HEAD_DIM,
137
+ )
138
)
139
140
init_workspace_manager(device)
tests/models/registry.py
+11
@@ -839,6 +839,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {
839
"DeepseekOCR2ForCausalLM": _HfExamplesInfo(
840
"deepseek-ai/DeepSeek-OCR-2",
841
),
842
+ "Dots3NoteForCausalLM": _HfExamplesInfo(
843
+ "rednote-hilab/dots3.note",
844
+ trust_remote_code=True,
845
+ is_available_online=False,
846
+ ),
847
"UnlimitedOCRForCausalLM": _HfExamplesInfo(
848
"baidu/Unlimited-OCR",
849
),
@@ -1650,6 +1655,12 @@ _SPECULATIVE_DECODING_EXAMPLE_MODELS = {
1655
trust_remote_code=True,
1656
is_available_online=False,
1657
),
1658
+ "Dots3NoteMTPModel": _HfExamplesInfo(
1659
+ "rednote-hilab/dots3.note",
1660
+ speculative_model="rednote-hilab/dots3.note",
1661
+ trust_remote_code=True,
1662
+ is_available_online=False,
1663
+ ),
1664
"Gemma4MTPModel": _HfExamplesInfo(
1665
"google/gemma-4-E4B-it",
1666
speculative_model="google/gemma-4-E4B-it-assistant",
tests/models/test_registry.py
+5
@@ -54,6 +54,11 @@ def test_registry_imports(model_arch):
54
):
55
pytest.skip("DSparkDraftModel is only supported on CUDA and ROCm")
56
57
+ if model_arch in ("Dots3NoteForCausalLM", "Dots3NoteMTPModel") and not (
58
+ current_platform.is_cuda()
59
+ ):
60
+ pytest.skip("Dots3 NOTE is only supported on CUDA")
61
+
62
# Ensure all model classes can be imported successfully
63
model_cls = ModelRegistry._try_load_model_cls(model_arch)
64
assert model_cls is not None
tests/tool_parsers/test_dots_tool_parser.py
new
+275
@@ -0,0 +1,275 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+from collections.abc import Iterable
6
+from unittest.mock import MagicMock
7
+
8
+import pytest
9
+
10
+from vllm.entrypoints.openai.chat_completion.protocol import (
11
+ ChatCompletionRequest,
12
+ ChatCompletionToolsParam,
13
+)
14
+from vllm.entrypoints.openai.engine.protocol import DeltaMessage
15
+from vllm.tool_parsers import ToolParserManager
16
+from vllm.tool_parsers.dots_tool_parser import DotsToolParser
17
+
18
+
19
+def _tool(name: str, properties: dict) -> ChatCompletionToolsParam:
20
+ return ChatCompletionToolsParam(
21
+ type="function",
22
+ function={
23
+ "name": name,
24
+ "description": "test tool",
25
+ "parameters": {"type": "object", "properties": properties},
26
+ },
27
+ )
28
+
29
+
30
+def _request(tools: list[ChatCompletionToolsParam]) -> ChatCompletionRequest:
31
+ return ChatCompletionRequest(
32
+ model="test-model",
33
+ messages=[{"role": "user", "content": "hi"}],
34
+ tools=tools,
35
+ tool_choice="auto",
36
+ )
37
+
38
+
39
+def _stream(
40
+ parser: DotsToolParser,
41
+ chunks: Iterable[str],
42
+ request: ChatCompletionRequest,
43
+) -> list[DeltaMessage]:
44
+ previous_text = ""
45
+ messages: list[DeltaMessage] = []
46
+ for chunk in chunks:
47
+ current_text = previous_text + chunk
48
+ delta = parser.extract_tool_calls_streaming(
49
+ previous_text,
50
+ current_text,
51
+ chunk,
52
+ [],
53
+ [],
54
+ [],
55
+ request,
56
+ )
57
+ if delta is not None:
58
+ messages.append(delta)
59
+ previous_text = current_text
60
+ return messages
61
+
62
+
63
+@pytest.fixture
64
+def parser() -> DotsToolParser:
65
+ return DotsToolParser(MagicMock())
66
+
67
+
68
+def test_registered_in_tool_parser_manager() -> None:
69
+ assert ToolParserManager.get_tool_parser("dots") is DotsToolParser
70
+
71
+
72
+def test_non_stream_xml_converts_schema_types_and_resolves_ref(
73
+ parser: DotsToolParser,
74
+) -> None:
75
+ tool = ChatCompletionToolsParam(
76
+ type="function",
77
+ function={
78
+ "name": "set_location",
79
+ "description": "Set location",
80
+ "parameters": {
81
+ "type": "object",
82
+ "properties": {
83
+ "location": {"$ref": "#/$defs/Location"},
84
+ "days": {"type": "integer"},
85
+ "include_weather": {"type": "boolean"},
86
+ },
87
+ "$defs": {
88
+ "Location": {
89
+ "type": "object",
90
+ "properties": {"city": {"type": "string"}},
91
+ }
92
+ },
93
+ },
94
+ },
95
+ )
96
+ request = _request([tool])
97
+ text = (
98
+ "ok<dots_function_call>"
99
+ '<invoke name="set_location">'
100
+ '<parameter name="location">{"city": "Shanghai"}</parameter>'
101
+ '<parameter name="days">3</parameter>'
102
+ '<parameter name="include_weather">true</parameter>'
103
+ "</invoke>"
104
+ "</dots_function_call>"
105
+ )
106
+
107
+ result = parser.extract_tool_calls(text, request)
108
+
109
+ assert result.content == "ok"
110
+ assert result.tools_called
111
+ assert len(result.tool_calls) == 1
112
+ assert result.tool_calls[0].function.name == "set_location"
113
+ assert json.loads(result.tool_calls[0].function.arguments) == {
114
+ "location": {"city": "Shanghai"},
115
+ "days": 3,
116
+ "include_weather": True,
117
+ }
118
+
119
+
120
+def test_non_stream_supports_multiple_invokes_and_json_fallback(
121
+ parser: DotsToolParser,
122
+) -> None:
123
+ tools = [
124
+ _tool("search", {"query": {"type": "string"}}),
125
+ _tool("open", {"id": {"type": "integer"}}),
126
+ ]
127
+ request = _request(tools)
128
+ text = (
129
+ "<dots_function_call>"
130
+ '<invoke name="search"><parameter name="query">chairs</parameter></invoke>'
131
+ '<invoke name="open"><parameter name="id">7</parameter></invoke>'
132
+ "</dots_function_call>"
133
+ '<dots_function_call>{"name":"search","arguments":{"query":"tables"}}'
134
+ "</dots_function_call>"
135
+ )
136
+
137
+ result = parser.extract_tool_calls(text, request)
138
+
139
+ assert [call.function.name for call in result.tool_calls] == [
140
+ "search",
141
+ "open",
142
+ "search",
143
+ ]
144
+ assert [json.loads(call.function.arguments) for call in result.tool_calls] == [
145
+ {"query": "chairs"},
146
+ {"id": 7},
147
+ {"query": "tables"},
148
+ ]
149
+
150
+
151
+def test_non_stream_unknown_tool_is_left_as_content(parser: DotsToolParser) -> None:
152
+ request = _request([_tool("search", {"query": {"type": "string"}})])
153
+ text = (
154
+ "<dots_function_call>"
155
+ '<invoke name="ghost"><parameter name="query">chairs</parameter></invoke>'
156
+ "</dots_function_call>"
157
+ )
158
+
159
+ result = parser.extract_tool_calls(text, request)
160
+
161
+ assert not result.tools_called
162
+ assert result.tool_calls == []
163
+ assert result.content == text
164
+
165
+
166
+def test_streaming_buffers_partial_marker_and_emits_all_complete_calls(
167
+ parser: DotsToolParser,
168
+) -> None:
169
+ request = _request([_tool("search", {"query": {"type": "string"}})])
170
+ chunks = [
171
+ "visible<dots_func",
172
+ (
173
+ "tion_call>"
174
+ '<invoke name="search"><parameter name="query">chairs</parameter></invoke>'
175
+ "</dots_function_call>"
176
+ "<dots_function_call>"
177
+ '<invoke name="search"><parameter name="query">tables</parameter></invoke>'
178
+ "</dots_function_call>"
179
+ ),
180
+ ]
181
+
182
+ messages = _stream(parser, chunks, request)
183
+
184
+ assert "".join(message.content or "" for message in messages) == "visible"
185
+ calls = [call for message in messages for call in message.tool_calls]
186
+ assert [call.index for call in calls] == [0, 1]
187
+ assert [json.loads(call.function.arguments or "") for call in calls] == [
188
+ {"query": "chairs"},
189
+ {"query": "tables"},
190
+ ]
191
+
192
+
193
+def test_streaming_filters_unknown_tools_and_surfaces_content(
194
+ parser: DotsToolParser,
195
+) -> None:
196
+ request = _request([_tool("search", {"query": {"type": "string"}})])
197
+ text = (
198
+ "<dots_function_call>"
199
+ '<invoke name="ghost"><parameter name="query">chairs</parameter></invoke>'
200
+ "</dots_function_call>"
201
+ )
202
+
203
+ messages = _stream(parser, [text], request)
204
+
205
+ assert len(messages) == 1
206
+ assert messages[0].tool_calls == []
207
+ assert "ghost" in (messages[0].content or "")
208
+ assert parser._buffer == ""
209
+
210
+
211
+def test_streaming_malformed_block_does_not_block_later_valid_call(
212
+ parser: DotsToolParser,
213
+) -> None:
214
+ request = _request([_tool("search", {"query": {"type": "string"}})])
215
+ messages = _stream(
216
+ parser,
217
+ [
218
+ "<dots_function_call>garbage</dots_function_call>",
219
+ (
220
+ "<dots_function_call>"
221
+ '<invoke name="search">'
222
+ '<parameter name="query">chairs</parameter>'
223
+ "</invoke></dots_function_call>"
224
+ ),
225
+ ],
226
+ request,
227
+ )
228
+
229
+ assert messages[0].content == "garbage"
230
+ assert messages[0].tool_calls == []
231
+ assert messages[1].tool_calls[0].function.name == "search"
232
+
233
+
234
+def test_streaming_strips_stray_end_marker(parser: DotsToolParser) -> None:
235
+ request = _request([_tool("search", {"query": {"type": "string"}})])
236
+ messages = _stream(parser, ["some text </dots_function_call>"], request)
237
+
238
+ assert len(messages) == 1
239
+ assert messages[0].content == "some text "
240
+ assert messages[0].tool_calls == []
241
+
242
+
243
+def test_streaming_flushes_partial_opening_marker_at_eof(
244
+ parser: DotsToolParser,
245
+) -> None:
246
+ request = _request([_tool("search", {"query": {"type": "string"}})])
247
+
248
+ messages = _stream(parser, ["answer <dots_func"], request)
249
+
250
+ assert messages[0].content == "answer "
251
+ assert parser.flush_pending_normal_text() == "<dots_func"
252
+ assert parser.flush_pending_normal_text() == ""
253
+
254
+
255
+def test_streaming_emits_complete_json_before_end_marker_without_duplication(
256
+ parser: DotsToolParser,
257
+) -> None:
258
+ request = _request([_tool("search", {"query": {"type": "string"}})])
259
+
260
+ messages = _stream(
261
+ parser,
262
+ [
263
+ "<dots_function_call>",
264
+ '{"name":"search","arguments":{"query":"chairs"}}',
265
+ "</dots_function_call>",
266
+ ],
267
+ request,
268
+ )
269
+
270
+ calls = [call for message in messages for call in message.tool_calls]
271
+ assert [call.function.name for call in calls] == ["search", None]
272
+ assert "".join(call.function.arguments or "" for call in calls) == (
273
+ '{"query": "chairs"}'
274
+ )
275
+ assert messages[-1].tool_calls == calls
tests/v1/attention/test_mla_backends.py
+3
@@ -78,6 +78,7 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype(
78
kv_cache_dtype=cache_dtype,
79
head_size=576,
80
non_causal_multi_token_decode=False,
81
+ sliding_window=None,
82
)
83
vllm_config = SimpleNamespace(
84
cache_config=SimpleNamespace(block_size=64), model_config=None
@@ -460,6 +461,7 @@ class MockSparseMLAAttentionLayer:
461
self.qk_nope_head_dim = qk_nope_head_dim
462
self.qk_rope_head_dim = qk_rope_head_dim
463
self.v_head_dim = v_head_dim
464
+ self.q_lora_rank = None
465
self.kv_lora_rank = kv_lora_rank
466
467
# Compute weight matrices in the format expected by forward_impl
@@ -583,6 +585,7 @@ class MockMLAAttentionLayer(MLAAttention):
585
self.qk_nope_head_dim = qk_nope_head_dim
586
self.qk_rope_head_dim = qk_rope_head_dim
587
self.v_head_dim = v_head_dim
588
+ self.q_lora_rank = None
589
self.kv_lora_rank = kv_lora_rank
590
591
# Compute weight matrices from kv_b_proj (like MLAAttention does)
vllm/config/speculative.py
+16
@@ -36,6 +36,7 @@ logger = init_logger(__name__)
36
37
MTPModelTypes = Literal[
38
"deepseek_mtp",
39
+ "dots3_note_mtp",
40
"mimo_mtp",
41
"mimo_v2_mtp",
42
"glm4_moe_mtp",
@@ -336,6 +337,21 @@ class SpeculativeConfig:
337
@staticmethod
338
def hf_config_override(hf_config: PretrainedConfig) -> PretrainedConfig:
339
initial_architecture = hf_config.architectures[0]
340
+ if hf_config.model_type == "dots3_note":
341
+ n_predict = getattr(hf_config, "num_nextn_predict_layers", 1)
342
+ mtp_layer_types = getattr(hf_config, "mtp_layer_types", None)
343
+ if mtp_layer_types is None:
344
+ mtp_layer_types = ["sliding_attention"] * n_predict
345
+ if len(mtp_layer_types) != n_predict:
346
+ raise ValueError(
347
+ "mtp_layer_types must have one entry per MTP layer: "
348
+ f"got {len(mtp_layer_types)} for {n_predict} layers"
349
+ )
350
+ hf_config.layer_types = [*hf_config.layer_types, *mtp_layer_types]
351
+ hf_config.model_type = "dots3_note_mtp"
352
+ hf_config.update(
353
+ {"n_predict": n_predict, "architectures": ["Dots3NoteMTPModel"]}
354
+ )
355
if hf_config.model_type in (
356
"deepseek_v3",
357
"deepseek_v32",
vllm/model_executor/layers/attention/mla_attention.py
+28
-6
@@ -299,6 +299,7 @@ from vllm.v1.kv_cache_interface import (
299
AttentionSpec,
300
KVCacheSpec,
301
MLAAttentionSpec,
302
+ SlidingWindowMLASpec,
303
get_kv_quant_mode,
304
)
305
@@ -406,6 +407,8 @@ class MLAAttention(nn.Module, AttentionLayerBase):
407
indexer: object | None = None,
408
topk_indices_buffer: torch.Tensor | None = None,
409
non_causal_multi_token_decode: bool = False,
410
+ sliding_window: int | None = None,
411
+ prefill_backend_cls: type[MLAPrefillBackend] | None = None,
412
**extra_impl_args,
413
):
414
super().__init__()
@@ -423,6 +426,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
426
self.layer_name = prefix
427
self.indexer = indexer
428
self.non_causal_multi_token_decode = non_causal_multi_token_decode
429
+ self.sliding_window = sliding_window
430
self.num_kv_heads = 1
431
self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
432
@@ -515,7 +519,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
519
scale=self.scale,
520
num_kv_heads=1,
521
alibi_slopes=None,
518
- sliding_window=None,
522
+ sliding_window=sliding_window,
523
kv_cache_dtype=self.kv_cache_dtype,
524
logits_soft_cap=None,
525
attn_type=AttentionType.DECODER,
@@ -551,7 +555,9 @@ class MLAAttention(nn.Module, AttentionLayerBase):
555
self.prefill_backend = None
556
else:
557
try:
554
- prefill_backend_cls = get_mla_prefill_backend(vllm_config)
558
+ prefill_backend_cls = prefill_backend_cls or get_mla_prefill_backend(
559
+ vllm_config
560
+ )
561
except ValueError:
562
if (
563
not self.impl.is_sparse
@@ -1139,13 +1145,21 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1145
kv_cache_dtype = kv_cache_dtype_str_to_dtype(
1146
self.kv_cache_dtype, vllm_config.model_config
1147
)
1142
- return MLAAttentionSpec(
1148
+ common_kwargs = dict(
1149
block_size=vllm_config.cache_config.block_size,
1150
num_kv_heads=1,
1151
head_size=self.head_size,
1152
dtype=kv_cache_dtype,
1153
cache_dtype_str=self.kv_cache_dtype,
1154
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
1155
+ )
1156
+ if self.sliding_window is not None:
1157
+ return SlidingWindowMLASpec(
1158
+ **common_kwargs,
1159
+ sliding_window=self.sliding_window,
1160
+ )
1161
+ return MLAAttentionSpec(
1162
+ **common_kwargs,
1163
non_causal_multi_token_decode=self.non_causal_multi_token_decode,
1164
)
1165
@@ -2056,7 +2070,15 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2070
self.num_heads = get_num_attention_heads_from_layers(
2071
vllm_config, layer_names
2072
) or self.model_config.get_num_attention_heads(parallel_config)
2059
- self.mla_dims = get_mla_dims(self.model_config)
2073
+ # Hybrid MLA models may use different latent dimensions per KV group.
2074
+ layer = self.compilation_config.static_forward_context[layer_names[0]]
2075
+ self.mla_dims = MLADims(
2076
+ q_lora_rank=layer.q_lora_rank,
2077
+ kv_lora_rank=layer.kv_lora_rank,
2078
+ qk_nope_head_dim=layer.qk_nope_head_dim,
2079
+ qk_rope_head_dim=layer.qk_rope_head_dim,
2080
+ v_head_dim=layer.v_head_dim,
2081
+ )
2082
self.aot_schedule = current_platform.is_cuda()
2083
2084
self.kv_cache_spec = kv_cache_spec
@@ -2092,7 +2114,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2114
(
2115
self.chunked_prefill_workspace_size
2116
+ self.chunked_prefill_workspace_size // self.dcp_world_size,
2095
- self.model_config.get_head_size(),
2117
+ self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
2118
),
2119
dtype=torch.bfloat16
2120
if use_packed_fp8_cache
@@ -2109,7 +2131,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):
2131
self.chunked_prefill_workspace = torch.empty(
2132
(
2133
self.chunked_prefill_workspace_size,
2112
- self.model_config.get_head_size(),
2134
+ self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
2135
),
2136
dtype=torch.bfloat16 if use_packed_fp8_cache else self.q_data_type,
2137
device=device,
vllm/model_executor/models/registry.py
+5
@@ -371,6 +371,10 @@ _MULTIMODAL_MODELS = {
371
"DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),
372
"DeepseekOCRForCausalLM": ("deepseek_ocr", "DeepseekOCRForCausalLM"),
373
"DeepseekOCR2ForCausalLM": ("deepseek_ocr2", "DeepseekOCR2ForCausalLM"),
374
+ "Dots3NoteForCausalLM": (
375
+ "vllm.models.dots3_note",
376
+ "Dots3NoteForCausalLM",
377
+ ),
378
"UnlimitedOCRForCausalLM": ("unlimited_ocr", "UnlimitedOCRForCausalLM"),
379
"DotsOCRForCausalLM": ("dots_ocr", "DotsOCRForCausalLM"),
380
"Eagle2_5_VLForConditionalGeneration": (
@@ -645,6 +649,7 @@ _SPECULATIVE_DECODING_MODELS = {
649
"Eagle3DeepseekV3ForCausalLM": ("deepseek_eagle3", "Eagle3DeepseekV2ForCausalLM"),
650
"EagleDeepSeekMTPModel": ("deepseek_eagle", "EagleDeepseekV3ForCausalLM"),
651
"DeepSeekMTPModel": ("deepseek_mtp", "DeepSeekMTP"),
652
+ "Dots3NoteMTPModel": ("vllm.models.dots3_note", "Dots3NoteMTP"),
653
"DeepSeekV4MTPModel": ("vllm.models.deepseek_v4", "DeepSeekV4MTP"),
654
"BailingMoeV3MTPModel": ("bailing_moe_v3_mtp", "BailingMoeV3MTPModel"),
655
"MiniMaxM3MTP": ("vllm.models.minimax_m3", "MiniMaxM3MTP"),
vllm/models/dots3_note/__init__.py
new
+8
@@ -0,0 +1,8 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Dots3Note model entry point."""
4
+
5
+from .nvidia.mtp import Dots3NoteMTP
6
+from .nvidia.multimodal import Dots3NoteForCausalLM
7
+
8
+__all__ = ["Dots3NoteForCausalLM", "Dots3NoteMTP"]
vllm/models/dots3_note/common/__init__.py
new
+3
@@ -0,0 +1,3 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Platform-independent Dots3Note preprocessing."""
vllm/models/dots3_note/common/processor.py
new
+811
@@ -0,0 +1,811 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Multimodal preprocessing for Dots3Note checkpoints."""
4
+
5
+import math
6
+from collections.abc import Mapping, Sequence
7
+from functools import cached_property
8
+from typing import Any, cast
9
+
10
+import numpy as np
11
+import torch
12
+from PIL import Image
13
+from transformers import BatchFeature
14
+
15
+from vllm.config.multimodal import (
16
+ AudioDummyOptions,
17
+ BaseDummyOptions,
18
+ ImageDummyOptions,
19
+ VideoDummyOptions,
20
+)
21
+from vllm.inputs import MultiModalDataDict
22
+from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems
23
+from vllm.multimodal.media import MediaWithBytes
24
+from vllm.multimodal.parse import (
25
+ MultiModalDataItems,
26
+ MultiModalDataParser,
27
+ VideoProcessorItems,
28
+)
29
+from vllm.multimodal.processing import (
30
+ BaseDummyInputsBuilder,
31
+ BaseMultiModalProcessor,
32
+ BaseProcessingInfo,
33
+ PromptReplacement,
34
+ PromptUpdate,
35
+ PromptUpdateDetails,
36
+)
37
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
38
+
39
+from .video import preprocess_dots3_note_video
40
+
41
+IMAGE_START = "<|img|>"
42
+IMAGE_PAD = "<|imgpad|>"
43
+IMAGE_END = "<|endofimg|>"
44
+AUDIO_START = "<|audio_comp_start|>"
45
+AUDIO_PAD = "<|audio_comp_pad|>"
46
+AUDIO_END = "<|audio_comp_end|>"
47
+VIDEO_PLACEHOLDER = "<|video_pad|>"
48
+
49
+_HOP_LENGTH = 160
50
+_DEFAULT_SAMPLE_RATE = 16000
51
+
52
+
53
+def load_note_config_section(
54
+ model: str,
55
+ revision: str | None,
56
+ section: str,
57
+) -> dict[str, Any] | None:
58
+ config = get_hf_file_to_dict("config.json", model, revision)
59
+ value = (config or {}).get(section)
60
+ return value if isinstance(value, dict) else None
61
+
62
+
63
+class Dots3NoteImageProcessor:
64
+ """CPU image preprocessing matching the native Cybertron ViT."""
65
+
66
+ def __init__(
67
+ self,
68
+ config: Mapping[str, Any],
69
+ image_details: Mapping[str, Any] | None = None,
70
+ ) -> None:
71
+ self.min_pixels = int(config["min_pixels"])
72
+ self.max_pixels = int(config["max_pixels"])
73
+ self.patch_size = int(config["patch_size"])
74
+ self.temporal_patch_size = int(config["temporal_patch_size"])
75
+ self.merge_size = int(config["merge_size"])
76
+ self.pre_pixel_shuffle = bool(config.get("pre_pixel_shuffle", True))
77
+ self.image_mean = np.asarray(config["image_mean"], dtype=np.float32)
78
+ self.image_std = np.asarray(config["image_std"], dtype=np.float32)
79
+ self.image_details = dict(image_details or {})
80
+
81
+ @property
82
+ def factor(self) -> int:
83
+ return self.patch_size * self.merge_size
84
+
85
+ @staticmethod
86
+ def _round_by_factor(value: int, factor: int) -> int:
87
+ return round(value / factor) * factor
88
+
89
+ @staticmethod
90
+ def _ceil_by_factor(value: float, factor: int) -> int:
91
+ return math.ceil(value / factor) * factor
92
+
93
+ @staticmethod
94
+ def _floor_by_factor(value: float, factor: int) -> int:
95
+ return math.floor(value / factor) * factor
96
+
97
+ def resized_size(
98
+ self,
99
+ width: int,
100
+ height: int,
101
+ *,
102
+ detail: str = "auto",
103
+ min_pixels: int | None = None,
104
+ max_pixels: int | None = None,
105
+ target_height: int | None = None,
106
+ target_width: int | None = None,
107
+ ) -> tuple[int, int]:
108
+ detail_config = self.image_details.get(detail, {})
109
+ min_pixels = int(
110
+ min_pixels
111
+ if min_pixels is not None
112
+ else detail_config.get("min_pixels", self.min_pixels)
113
+ )
114
+ max_pixels = int(
115
+ max_pixels
116
+ if max_pixels is not None
117
+ else detail_config.get("max_pixels", self.max_pixels)
118
+ )
119
+ height = int(target_height or detail_config.get("target_height") or height)
120
+ width = int(target_width or detail_config.get("target_width") or width)
121
+
122
+ factor = self.factor
123
+ if min(height, width) < factor // 4:
124
+ raise ValueError(
125
+ f"Image height and width must be at least {factor // 4}, "
126
+ f"got {height}x{width}"
127
+ )
128
+ if max(height, width) / min(height, width) > 200:
129
+ raise ValueError("Image aspect ratio must be smaller than 200")
130
+
131
+ resized_h = max(factor, self._round_by_factor(height, factor))
132
+ resized_w = max(factor, self._round_by_factor(width, factor))
133
+ if resized_h * resized_w > max_pixels:
134
+ beta = math.sqrt(height * width / max_pixels)
135
+ resized_h = max(factor, self._floor_by_factor(height / beta, factor))
136
+ resized_w = max(factor, self._floor_by_factor(width / beta, factor))
137
+ elif resized_h * resized_w < min_pixels:
138
+ beta = math.sqrt(min_pixels / (height * width))
139
+ resized_h = self._ceil_by_factor(height * beta, factor)
140
+ resized_w = self._ceil_by_factor(width * beta, factor)
141
+ if resized_h * resized_w > max_pixels:
142
+ beta = math.sqrt(resized_h * resized_w / max_pixels)
143
+ resized_h = max(factor, self._floor_by_factor(resized_h / beta, factor))
144
+ resized_w = max(factor, self._floor_by_factor(resized_w / beta, factor))
145
+ return resized_h, resized_w
146
+
147
+ def preprocess(
148
+ self,
149
+ image: Image.Image,
150
+ *,
151
+ detail: str = "auto",
152
+ min_pixels: int | None = None,
153
+ max_pixels: int | None = None,
154
+ target_height: int | None = None,
155
+ target_width: int | None = None,
156
+ ) -> tuple[torch.Tensor, torch.Tensor]:
157
+ if not isinstance(image, Image.Image):
158
+ raise TypeError(f"Expected a PIL image, got {type(image)}")
159
+ if image.mode == "RGBA":
160
+ background = Image.new("RGB", image.size, (255, 255, 255))
161
+ background.paste(image, mask=image.getchannel("A"))
162
+ image = background
163
+ elif image.mode != "RGB":
164
+ image = image.convert("RGB")
165
+
166
+ resized_h, resized_w = self.resized_size(
167
+ *image.size,
168
+ detail=detail,
169
+ min_pixels=min_pixels,
170
+ max_pixels=max_pixels,
171
+ target_height=target_height,
172
+ target_width=target_width,
173
+ )
174
+ image = image.resize((resized_w, resized_h), Image.Resampling.BICUBIC)
175
+ array = np.asarray(image, dtype=np.float32) / 255.0
176
+ array = (array - self.image_mean) / self.image_std
177
+ patches = array.transpose(2, 0, 1)[None]
178
+ if patches.shape[0] == 1:
179
+ patches = np.tile(patches, (self.temporal_patch_size, 1, 1, 1))
180
+
181
+ channel = patches.shape[1]
182
+ grid_t = patches.shape[0] // self.temporal_patch_size
183
+ grid_h = resized_h // self.patch_size
184
+ grid_w = resized_w // self.patch_size
185
+ if self.pre_pixel_shuffle:
186
+ patches = patches.reshape(
187
+ grid_t,
188
+ self.temporal_patch_size,
189
+ channel,
190
+ grid_h // self.merge_size,
191
+ self.merge_size,
192
+ self.patch_size,
193
+ grid_w // self.merge_size,
194
+ self.merge_size,
195
+ self.patch_size,
196
+ )
197
+ patches = patches.transpose(0, 3, 6, 4, 7, 2, 1, 5, 8)
198
+ else:
199
+ patches = patches.reshape(
200
+ grid_t,
201
+ self.temporal_patch_size,
202
+ channel,
203
+ grid_h,
204
+ self.patch_size,
205
+ grid_w,
206
+ self.patch_size,
207
+ )
208
+ patches = patches.transpose(0, 3, 5, 2, 1, 4, 6)
209
+
210
+ pixel_values = torch.from_numpy(
211
+ patches.reshape(
212
+ grid_t * grid_h * grid_w,
213
+ channel * self.temporal_patch_size * self.patch_size * self.patch_size,
214
+ )
215
+ )
216
+ grid_thw = torch.tensor([grid_t, grid_h, grid_w], dtype=torch.long)
217
+ return pixel_values, grid_thw
218
+
219
+
220
+class Dots3NoteProcessor:
221
+ """Small HF-like processor used by vLLM's multimodal frontend."""
222
+
223
+ def __init__(
224
+ self,
225
+ tokenizer,
226
+ image_processor: Dots3NoteImageProcessor | None,
227
+ *,
228
+ max_model_len: int,
229
+ video_audio_enabled: bool,
230
+ audio_token_stride: int,
231
+ ) -> None:
232
+ self.tokenizer = tokenizer
233
+ self.image_processor = image_processor
234
+ self.max_model_len = max_model_len
235
+ self.video_audio_enabled = video_audio_enabled
236
+ self.audio_token_stride = audio_token_stride
237
+
238
+ def _token_ids(self, text: str) -> list[int]:
239
+ if hasattr(self.tokenizer, "encode"):
240
+ return self.tokenizer.encode(text, add_special_tokens=False)
241
+ token_ids = self.tokenizer(text, add_special_tokens=False)["input_ids"]
242
+ if token_ids and isinstance(token_ids[0], list):
243
+ return token_ids[0]
244
+ return token_ids
245
+
246
+ def _process_video(
247
+ self,
248
+ video,
249
+ *,
250
+ detail: str,
251
+ size_overrides: Mapping[str, int | None],
252
+ seq: int,
253
+ output_reserve: int | None,
254
+ audio_cap: float,
255
+ audio_sample_rate: int,
256
+ k_mode: str,
257
+ max_new_tokens: int,
258
+ question: str,
259
+ ) -> dict[str, torch.Tensor]:
260
+ if self.image_processor is None:
261
+ raise ValueError("This NOTE checkpoint has no vision encoder")
262
+ if not self.video_audio_enabled:
263
+ audio_cap = 0.0
264
+ parts = preprocess_dots3_note_video(
265
+ video,
266
+ tokenizer=self.tokenizer,
267
+ question=question,
268
+ seq=seq,
269
+ output_reserve=output_reserve,
270
+ audio_cap=audio_cap,
271
+ audio_sample_rate=audio_sample_rate,
272
+ k_mode=k_mode,
273
+ max_new_tokens=max_new_tokens,
274
+ )
275
+
276
+ fragments: list[str] = []
277
+ pixel_values: list[torch.Tensor] = []
278
+ image_grids: list[torch.Tensor] = []
279
+ audio_values: list[torch.Tensor] = []
280
+ modalities: list[int] = []
281
+ image_pad_count = 0
282
+ audio_pad_count = 0
283
+
284
+ for part in parts:
285
+ if part.kind == "text":
286
+ fragments.append(cast(str, part.value))
287
+ continue
288
+ if part.kind == "image":
289
+ pixels, grid = self.image_processor.preprocess(
290
+ cast(Image.Image, part.value),
291
+ detail=detail,
292
+ min_pixels=size_overrides.get("min_pixels"),
293
+ max_pixels=size_overrides.get("max_pixels"),
294
+ target_height=size_overrides.get("target_height"),
295
+ target_width=size_overrides.get("target_width"),
296
+ )
297
+ num_tokens = int(grid.prod()) // self.image_processor.merge_size**2
298
+ fragments.append(f"{IMAGE_START}{IMAGE_PAD * num_tokens}{IMAGE_END}")
299
+ pixel_values.append(pixels)
300
+ image_grids.append(grid)
301
+ modalities.append(0)
302
+ image_pad_count += num_tokens
303
+ continue
304
+
305
+ waveform = torch.from_numpy(
306
+ np.ascontiguousarray(part.value, dtype=np.float32)
307
+ ).view(torch.int32)
308
+ num_tokens = math.ceil(waveform.numel() / self.audio_token_stride)
309
+ fragments.append(f"{AUDIO_START}{AUDIO_PAD * num_tokens}{AUDIO_END}")
310
+ audio_values.append(waveform)
311
+ modalities.append(1)
312
+ audio_pad_count += num_tokens
313
+
314
+ if not pixel_values:
315
+ raise ValueError("NOTE video preprocessing produced no frames")
316
+ fragment_ids = self._token_ids("".join(fragments))
317
+ vocab = self.tokenizer.get_vocab()
318
+ embed_mask = torch.isin(
319
+ torch.tensor(fragment_ids, dtype=torch.long),
320
+ torch.tensor([vocab[IMAGE_PAD], vocab[AUDIO_PAD]], dtype=torch.long),
321
+ )
322
+ expected_embeds = image_pad_count + audio_pad_count
323
+ if int(embed_mask.sum()) != expected_embeds:
324
+ raise ValueError(
325
+ "NOTE video placeholder expansion produced an invalid embedding mask: "
326
+ f"expected={expected_embeds}, actual={int(embed_mask.sum())}"
327
+ )
328
+
329
+ empty_audio = torch.empty(0, dtype=torch.int32)
330
+ return {
331
+ "video_pixel_values": torch.cat(pixel_values),
332
+ "video_image_grid_thw": torch.stack(image_grids),
333
+ "video_audio_values": (
334
+ torch.cat(audio_values) if audio_values else empty_audio
335
+ ),
336
+ "video_audio_lengths": torch.tensor(
337
+ [value.numel() for value in audio_values],
338
+ dtype=torch.long,
339
+ ),
340
+ "video_modalities": torch.tensor(modalities, dtype=torch.uint8),
341
+ "video_input_ids": torch.tensor(fragment_ids, dtype=torch.long),
342
+ "video_embed_mask": embed_mask,
343
+ "video_frame_counts": torch.tensor([len(image_grids)], dtype=torch.long),
344
+ "video_audio_counts": torch.tensor([len(audio_values)], dtype=torch.long),
345
+ "video_emission_counts": torch.tensor([len(modalities)], dtype=torch.long),
346
+ "video_prompt_lengths": torch.tensor([len(fragment_ids)], dtype=torch.long),
347
+ "video_patch_counts": torch.tensor(
348
+ [sum(int(grid.prod()) for grid in image_grids)], dtype=torch.long
349
+ ),
350
+ "video_audio_sample_counts": torch.tensor(
351
+ [sum(value.numel() for value in audio_values)], dtype=torch.long
352
+ ),
353
+ }
354
+
355
+ def __call__(self, text: str, **kwargs: object) -> BatchFeature:
356
+ modality_order = [
357
+ key for key in kwargs if key in ("images", "audios", "videos")
358
+ ]
359
+ raw_images = kwargs.pop("images", None)
360
+ raw_audios = kwargs.pop("audios", None)
361
+ raw_videos = kwargs.pop("videos", None)
362
+ images = (
363
+ list(cast(Sequence[Image.Image], raw_images))
364
+ if raw_images is not None
365
+ else []
366
+ )
367
+ audios = (
368
+ list(cast(Sequence[np.ndarray], raw_audios))
369
+ if raw_audios is not None
370
+ else []
371
+ )
372
+ videos = (
373
+ list(cast(Sequence[object], raw_videos)) if raw_videos is not None else []
374
+ )
375
+ if videos and (images or audios):
376
+ raise ValueError(
377
+ "Dots3Note does not support mixing a native video with "
378
+ "separate image/audio inputs"
379
+ )
380
+ if len(videos) > 1:
381
+ raise ValueError("Dots3Note supports one video per request")
382
+ detail = kwargs.pop("image_detail", "auto")
383
+ details = list(detail) if isinstance(detail, (list, tuple)) else None
384
+ size_overrides: dict[str, int | None] = {}
385
+ for key in ("min_pixels", "max_pixels", "target_height", "target_width"):
386
+ if key in kwargs:
387
+ value = kwargs.pop(key)
388
+ size_overrides[key] = None if value is None else int(cast(Any, value))
389
+ video_seq = int(cast(Any, kwargs.pop("seq", self.max_model_len)))
390
+ reserve_value = kwargs.pop("output_reserve", None)
391
+ output_reserve = (
392
+ None if reserve_value is None else int(cast(Any, reserve_value))
393
+ )
394
+ audio_cap = float(cast(Any, kwargs.pop("audio_cap", 1.0)))
395
+ audio_sample_rate = int(cast(Any, kwargs.pop("audio_sr", _DEFAULT_SAMPLE_RATE)))
396
+ k_mode = str(kwargs.pop("k_mode", "eval_ek"))
397
+ max_new_tokens = int(cast(Any, kwargs.pop("max_new_tokens", 0)))
398
+ question = str(kwargs.pop("video_question", text))
399
+
400
+ tokenized = self.tokenizer(text, **kwargs)
401
+ input_ids = tokenized["input_ids"]
402
+ if isinstance(input_ids, torch.Tensor):
403
+ if input_ids.ndim == 1:
404
+ input_ids = input_ids.unsqueeze(0)
405
+ elif input_ids and isinstance(input_ids[0], int):
406
+ input_ids = [input_ids]
407
+ data: dict[str, object] = {"input_ids": input_ids}
408
+
409
+ for modality in modality_order:
410
+ if modality == "images" and images:
411
+ if self.image_processor is None:
412
+ raise ValueError("This NOTE checkpoint has no vision encoder")
413
+ pixel_values = []
414
+ grids = []
415
+ for idx, image in enumerate(images):
416
+ image_detail = details[idx] if details is not None else str(detail)
417
+ pixels, grid = self.image_processor.preprocess(
418
+ image,
419
+ detail=image_detail,
420
+ min_pixels=size_overrides.get("min_pixels"),
421
+ max_pixels=size_overrides.get("max_pixels"),
422
+ target_height=size_overrides.get("target_height"),
423
+ target_width=size_overrides.get("target_width"),
424
+ )
425
+ pixel_values.append(pixels)
426
+ grids.append(grid)
427
+ data["pixel_values"] = torch.cat(pixel_values)
428
+ data["image_grid_thw"] = torch.stack(grids)
429
+ elif modality == "audios" and audios:
430
+ waveforms = [
431
+ torch.from_numpy(
432
+ np.ascontiguousarray(audio, dtype=np.float32)
433
+ ).view(torch.int32)
434
+ for audio in audios
435
+ ]
436
+ data["audio_values"] = torch.cat(waveforms)
437
+ data["audio_lengths"] = torch.tensor(
438
+ [waveform.numel() for waveform in waveforms],
439
+ dtype=torch.long,
440
+ )
441
+ elif modality == "videos" and videos:
442
+ video_detail = details[0] if details is not None else str(detail)
443
+ data.update(
444
+ self._process_video(
445
+ videos[0],
446
+ detail=video_detail,
447
+ size_overrides=size_overrides,
448
+ seq=video_seq,
449
+ output_reserve=output_reserve,
450
+ audio_cap=audio_cap,
451
+ audio_sample_rate=audio_sample_rate,
452
+ k_mode=k_mode,
453
+ max_new_tokens=max_new_tokens,
454
+ question=question,
455
+ )
456
+ )
457
+ return BatchFeature(data=data)
458
+
459
+
460
+class Dots3NoteProcessingInfo(BaseProcessingInfo):
461
+ @cached_property
462
+ def vision_config(self) -> dict[str, Any] | None:
463
+ model_config = self.ctx.model_config
464
+ return load_note_config_section(
465
+ model_config.model,
466
+ model_config.revision,
467
+ "vision_config",
468
+ )
469
+
470
+ @cached_property
471
+ def audio_config(self) -> dict[str, Any] | None:
472
+ model_config = self.ctx.model_config
473
+ return load_note_config_section(
474
+ model_config.model,
475
+ model_config.revision,
476
+ "audio_config",
477
+ )
478
+
479
+ @cached_property
480
+ def image_processor(self) -> Dots3NoteImageProcessor | None:
481
+ if self.vision_config is None:
482
+ return None
483
+ model_config = self.ctx.model_config
484
+ processor_config = get_hf_file_to_dict(
485
+ "preprocessor_config.json",
486
+ model_config.model,
487
+ model_config.revision,
488
+ )
489
+ preprocessor_config = (processor_config or {}).get("vision_config")
490
+ if not isinstance(preprocessor_config, dict):
491
+ raise ValueError(
492
+ "NOTE vision checkpoint is missing preprocessor_config.json"
493
+ )
494
+ image_details = preprocessor_config.get("image_details", {})
495
+ return Dots3NoteImageProcessor(preprocessor_config, image_details)
496
+
497
+ @cached_property
498
+ def processor(self) -> Dots3NoteProcessor:
499
+ mm_config = self.ctx.get_mm_config()
500
+ audio_config = self.audio_config or {}
501
+ audio_token_stride = (
502
+ _HOP_LENGTH
503
+ * (8 if audio_config.get("use_conv2d_stem", True) else 2)
504
+ * int(audio_config.get("merge_factor", 1))
505
+ )
506
+ return Dots3NoteProcessor(
507
+ self.get_tokenizer(),
508
+ self.image_processor,
509
+ max_model_len=self.ctx.model_config.max_model_len,
510
+ video_audio_enabled=(
511
+ self.audio_config is not None
512
+ and (
513
+ mm_config.get_limit_per_prompt("audio") > 0
514
+ or mm_config.get_limit_per_prompt("video") > 0
515
+ )
516
+ ),
517
+ audio_token_stride=audio_token_stride,
518
+ )
519
+
520
+ def get_hf_processor(self, **kwargs: object) -> Dots3NoteProcessor:
521
+ return self.processor
522
+
523
+ def get_data_parser(self) -> MultiModalDataParser:
524
+ sample_rate = int((self.audio_config or {}).get("sampling_rate", 16000))
525
+ return MultiModalDataParser(target_sr=float(sample_rate), target_channels=1)
526
+
527
+ def get_supported_mm_limits(self) -> Mapping[str, int | None]:
528
+ limits: dict[str, int | None] = {}
529
+ if self.vision_config is not None:
530
+ limits["image"] = 512
531
+ limits["video"] = 1
532
+ if self.audio_config is not None:
533
+ limits["audio"] = 128
534
+ return limits
535
+
536
+ def get_max_image_size(self) -> tuple[int, int]:
537
+ image_processor = self.image_processor
538
+ if image_processor is None:
539
+ return 0, 0
540
+ factor = image_processor.factor
541
+ max_tokens = image_processor.max_pixels // (factor * factor)
542
+ for height_factor in range(math.isqrt(max_tokens), 0, -1):
543
+ if max_tokens % height_factor == 0:
544
+ width_factor = max_tokens // height_factor
545
+ if width_factor / height_factor <= 200:
546
+ return factor * width_factor, factor * height_factor
547
+ return factor * max_tokens, factor
548
+
549
+ def get_mm_max_tokens_per_item(
550
+ self,
551
+ seq_len: int,
552
+ mm_counts: Mapping[str, int],
553
+ ) -> Mapping[str, int]:
554
+ result: dict[str, int] = {}
555
+ if self.image_processor is not None:
556
+ factor = self.image_processor.factor
557
+ result["image"] = self.image_processor.max_pixels // (factor * factor)
558
+ if self.audio_config is not None:
559
+ stride = (
560
+ _HOP_LENGTH
561
+ * (8 if self.audio_config.get("use_conv2d_stem", True) else 2)
562
+ * int(self.audio_config.get("merge_factor", 1))
563
+ )
564
+ chunk_samples = int(
565
+ self.audio_config.get("chunk_seconds", 60)
566
+ * self.audio_config.get("sampling_rate", _DEFAULT_SAMPLE_RATE)
567
+ )
568
+ result["audio"] = min(seq_len, math.ceil(chunk_samples / stride))
569
+ if self.image_processor is not None:
570
+ result["video"] = max(1, seq_len - seq_len // 4)
571
+ return result
572
+
573
+
574
+class Dots3NoteDummyInputsBuilder(BaseDummyInputsBuilder[Dots3NoteProcessingInfo]):
575
+ def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
576
+ return (
577
+ f"{IMAGE_START}{IMAGE_PAD}{IMAGE_END}" * mm_counts.get("image", 0)
578
+ + f"{AUDIO_START}{AUDIO_PAD}{AUDIO_END}" * mm_counts.get("audio", 0)
579
+ + VIDEO_PLACEHOLDER * mm_counts.get("video", 0)
580
+ )
581
+
582
+ def get_dummy_mm_data(
583
+ self,
584
+ seq_len: int,
585
+ mm_counts: Mapping[str, int],
586
+ mm_options: Mapping[str, BaseDummyOptions],
587
+ ) -> MultiModalDataDict:
588
+ data: dict[str, Any] = {}
589
+ num_images = mm_counts.get("image", 0)
590
+ if num_images:
591
+ width, height = self.info.get_max_image_size()
592
+ data["image"] = self._get_dummy_images(
593
+ width=width,
594
+ height=height,
595
+ num_images=num_images,
596
+ overrides=cast(ImageDummyOptions | None, mm_options.get("image")),
597
+ )
598
+ num_audios = mm_counts.get("audio", 0)
599
+ if num_audios:
600
+ audio_config = self.info.audio_config or {}
601
+ stride = (
602
+ _HOP_LENGTH
603
+ * (8 if audio_config.get("use_conv2d_stem", True) else 2)
604
+ * int(audio_config.get("merge_factor", 1))
605
+ )
606
+ chunk_samples = int(
607
+ audio_config.get("chunk_seconds", 60)
608
+ * audio_config.get("sampling_rate", _DEFAULT_SAMPLE_RATE)
609
+ )
610
+ data["audio"] = self._get_dummy_audios(
611
+ length=min(chunk_samples, max(1, seq_len) * stride),
612
+ num_audios=num_audios,
613
+ overrides=cast(AudioDummyOptions | None, mm_options.get("audio")),
614
+ )
615
+ num_videos = mm_counts.get("video", 0)
616
+ if num_videos:
617
+ width, height = self.info.get_max_image_size()
618
+ factor = self.info.image_processor.factor # type: ignore[union-attr]
619
+ frame_tokens = max(1, width * height // (factor * factor))
620
+ video_tokens = max(1, seq_len - seq_len // 4)
621
+ num_frames = max(4, math.ceil(video_tokens / frame_tokens))
622
+ data["video"] = self._get_dummy_videos(
623
+ width=width,
624
+ height=height,
625
+ num_frames=num_frames,
626
+ num_videos=num_videos,
627
+ overrides=cast(VideoDummyOptions | None, mm_options.get("video")),
628
+ )
629
+ return data
630
+
631
+
632
+class Dots3NoteMultiModalProcessor(BaseMultiModalProcessor[Dots3NoteProcessingInfo]):
633
+ def _get_hf_mm_data(
634
+ self,
635
+ mm_items: MultiModalDataItems,
636
+ ) -> tuple[Mapping[str, object], Mapping[str, object]]:
637
+ processor_data, passthrough_data = super()._get_hf_mm_data(mm_items)
638
+ if "video" not in mm_items:
639
+ return processor_data, passthrough_data
640
+
641
+ videos = mm_items.get_items("video", VideoProcessorItems)
642
+ raw_videos: list[object] = []
643
+ for index, item in enumerate(videos.data):
644
+ if isinstance(item, MediaWithBytes):
645
+ raw_videos.append(item.original_bytes)
646
+ else:
647
+ raw_videos.append(videos.get(index))
648
+ processor_data = dict(processor_data)
649
+ processor_data["videos"] = raw_videos
650
+ return processor_data, passthrough_data
651
+
652
+ def _call_hf_processor(
653
+ self,
654
+ prompt: str,
655
+ mm_data: Mapping[str, object],
656
+ mm_kwargs: Mapping[str, object],
657
+ tok_kwargs: Mapping[str, object],
658
+ ) -> BatchFeature:
659
+ return super()._call_hf_processor(prompt, mm_data, mm_kwargs, tok_kwargs)
660
+
661
+ def _hf_processor_applies_updates(
662
+ self,
663
+ prompt_text: str,
664
+ mm_items: MultiModalDataItems,
665
+ hf_processor_mm_kwargs: Mapping[str, object],
666
+ tokenization_kwargs: Mapping[str, object],
667
+ ) -> bool:
668
+ return False
669
+
670
+ def _get_mm_fields_config(
671
+ self,
672
+ hf_inputs: BatchFeature,
673
+ hf_processor_mm_kwargs: Mapping[str, object],
674
+ ) -> Mapping[str, MultiModalFieldConfig]:
675
+ fields: dict[str, MultiModalFieldConfig] = {}
676
+ if "pixel_values" in hf_inputs:
677
+ grid_thw = hf_inputs["image_grid_thw"]
678
+ fields["pixel_values"] = MultiModalFieldConfig.flat_from_sizes(
679
+ "image", grid_thw.prod(-1)
680
+ )
681
+ fields["image_grid_thw"] = MultiModalFieldConfig.batched(
682
+ "image", keep_on_cpu=True
683
+ )
684
+ if "audio_values" in hf_inputs:
685
+ fields["audio_values"] = MultiModalFieldConfig.flat_from_sizes(
686
+ "audio", hf_inputs["audio_lengths"]
687
+ )
688
+ fields["audio_lengths"] = MultiModalFieldConfig.batched(
689
+ "audio", keep_on_cpu=True
690
+ )
691
+ if "video_input_ids" in hf_inputs:
692
+ prompt_lengths = hf_inputs["video_prompt_lengths"]
693
+ frame_counts = hf_inputs["video_frame_counts"]
694
+ audio_counts = hf_inputs["video_audio_counts"]
695
+ emission_counts = hf_inputs["video_emission_counts"]
696
+ fields["video_pixel_values"] = MultiModalFieldConfig.flat_from_sizes(
697
+ "video", hf_inputs["video_patch_counts"]
698
+ )
699
+ fields["video_image_grid_thw"] = MultiModalFieldConfig.flat_from_sizes(
700
+ "video", frame_counts, keep_on_cpu=True
701
+ )
702
+ fields["video_audio_values"] = MultiModalFieldConfig.flat_from_sizes(
703
+ "video", hf_inputs["video_audio_sample_counts"]
704
+ )
705
+ fields["video_audio_lengths"] = MultiModalFieldConfig.flat_from_sizes(
706
+ "video", audio_counts, keep_on_cpu=True
707
+ )
708
+ fields["video_modalities"] = MultiModalFieldConfig.flat_from_sizes(
709
+ "video", emission_counts, keep_on_cpu=True
710
+ )
711
+ fields["video_input_ids"] = MultiModalFieldConfig.flat_from_sizes(
712
+ "video", prompt_lengths, keep_on_cpu=True
713
+ )
714
+ fields["video_embed_mask"] = MultiModalFieldConfig.flat_from_sizes(
715
+ "video", prompt_lengths, keep_on_cpu=True
716
+ )
717
+ for key in (
718
+ "video_frame_counts",
719
+ "video_audio_counts",
720
+ "video_emission_counts",
721
+ ):
722
+ fields[key] = MultiModalFieldConfig.batched("video", keep_on_cpu=True)
723
+ return fields
724
+
725
+ def _get_prompt_updates(
726
+ self,
727
+ mm_items: MultiModalDataItems,
728
+ hf_processor_mm_kwargs: Mapping[str, object],
729
+ out_mm_kwargs: MultiModalKwargsItems,
730
+ ) -> Sequence[PromptUpdate]:
731
+ vocab = self.info.get_tokenizer().get_vocab()
732
+ updates: list[PromptUpdate] = []
733
+
734
+ if "image" in out_mm_kwargs:
735
+ image_start_id = vocab[IMAGE_START]
736
+ image_pad_id = vocab[IMAGE_PAD]
737
+ image_end_id = vocab[IMAGE_END]
738
+ merge_size = self.info.image_processor.merge_size # type: ignore[union-attr]
739
+
740
+ def image_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
741
+ grid = out_mm_kwargs["image"][item_idx]["image_grid_thw"].data
742
+ assert isinstance(grid, torch.Tensor)
743
+ num_tokens = int(grid.prod()) // merge_size**2
744
+ return PromptUpdateDetails.select_token_id(
745
+ [image_start_id] + [image_pad_id] * num_tokens + [image_end_id],
746
+ image_pad_id,
747
+ )
748
+
749
+ updates.append(
750
+ PromptReplacement(
751
+ modality="image",
752
+ target=[image_start_id, image_pad_id, image_end_id],
753
+ replacement=image_replacement,
754
+ )
755
+ )
756
+
757
+ if "audio" in out_mm_kwargs:
758
+ audio_start_id = vocab[AUDIO_START]
759
+ audio_pad_id = vocab[AUDIO_PAD]
760
+ audio_end_id = vocab[AUDIO_END]
761
+ config = self.info.audio_config or {}
762
+ stride = (
763
+ _HOP_LENGTH
764
+ * (8 if config.get("use_conv2d_stem", True) else 2)
765
+ * int(config.get("merge_factor", 1))
766
+ )
767
+
768
+ def audio_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
769
+ length = out_mm_kwargs["audio"][item_idx]["audio_lengths"].data
770
+ assert isinstance(length, torch.Tensor)
771
+ num_tokens = math.ceil(int(length.item()) / stride)
772
+ return PromptUpdateDetails.select_token_id(
773
+ [audio_start_id] + [audio_pad_id] * num_tokens + [audio_end_id],
774
+ audio_pad_id,
775
+ )
776
+
777
+ updates.append(
778
+ PromptReplacement(
779
+ modality="audio",
780
+ target=[audio_start_id, audio_pad_id, audio_end_id],
781
+ replacement=audio_replacement,
782
+ )
783
+ )
784
+ if "video" in out_mm_kwargs:
785
+
786
+ def video_replacement(item_idx: int) -> PromptUpdateDetails[list[int]]:
787
+ item = out_mm_kwargs["video"][item_idx]
788
+ input_ids_data = item["video_input_ids"].data
789
+ embed_mask_data = item["video_embed_mask"].data
790
+ assert isinstance(input_ids_data, torch.Tensor)
791
+ assert isinstance(embed_mask_data, torch.Tensor)
792
+ input_ids = input_ids_data.tolist()
793
+ embed_mask = embed_mask_data.bool()
794
+
795
+ def select_video_embeds(tokenizer, full) -> torch.Tensor:
796
+ del tokenizer, full
797
+ return embed_mask
798
+
799
+ return PromptUpdateDetails(
800
+ full=input_ids,
801
+ is_embed=select_video_embeds,
802
+ )
803
+
804
+ updates.append(
805
+ PromptReplacement(
806
+ modality="video",
807
+ target=VIDEO_PLACEHOLDER,
808
+ replacement=video_replacement,
809
+ )
810
+ )
811
+ return updates
vllm/models/dots3_note/common/video.py
new
+497
@@ -0,0 +1,497 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Train-consistent video preprocessing for Dots3Note."""
4
+
5
+from __future__ import annotations
6
+
7
+import hashlib
8
+import io
9
+import math
10
+import os
11
+import random
12
+from dataclasses import dataclass
13
+from typing import Literal
14
+
15
+import numpy as np
16
+from PIL import Image
17
+
18
+_ALIGN = 28
19
+_MIN_FRAMES = 4
20
+_PF_FLOOR = 128
21
+_PF_CEIL = 1024
22
+_FPS_CAP = 1.0
23
+_FPS_MIN = 0.2
24
+_FRAME_OVERHEAD = 15
25
+_BUDGET_OVERHEAD = 2240
26
+_INTERLEAVE_MIN_SECONDS = 1.0
27
+_AUDIO_SAMPLE_RATE = 16000
28
+_AUDIO_SAMPLES_PER_TOKEN = 1280
29
+_AUDIO_CHUNK_SECONDS = 30
30
+
31
+
32
+@dataclass(frozen=True)
33
+class Dots3NoteVideoPart:
34
+ kind: Literal["text", "image", "audio"]
35
+ value: str | Image.Image | np.ndarray
36
+
37
+
38
+def _token_len(tokenizer, text: str) -> int:
39
+ if not text:
40
+ return 0
41
+ if hasattr(tokenizer, "encode"):
42
+ return len(tokenizer.encode(text, add_special_tokens=False))
43
+ return len(tokenizer(text, add_special_tokens=False)["input_ids"])
44
+
45
+
46
+def _compute_target_size(
47
+ orig_h: int,
48
+ orig_w: int,
49
+ min_pixels: int,
50
+ max_pixels: int,
51
+) -> tuple[int, int]:
52
+ height = max(_ALIGN, round(orig_h / _ALIGN) * _ALIGN)
53
+ width = max(_ALIGN, round(orig_w / _ALIGN) * _ALIGN)
54
+ if height * width > max_pixels:
55
+ beta = math.sqrt(orig_h * orig_w / max_pixels)
56
+ height = max(_ALIGN, math.floor(orig_h / beta / _ALIGN) * _ALIGN)
57
+ width = max(_ALIGN, math.floor(orig_w / beta / _ALIGN) * _ALIGN)
58
+ elif height * width < min_pixels:
59
+ beta = math.sqrt(min_pixels / max(1, orig_h * orig_w))
60
+ height = math.ceil(orig_h * beta / _ALIGN) * _ALIGN
61
+ width = math.ceil(orig_w * beta / _ALIGN) * _ALIGN
62
+ if height * width > max_pixels:
63
+ beta = math.sqrt(height * width / max_pixels)
64
+ height = max(_ALIGN, math.floor(height / beta / _ALIGN) * _ALIGN)
65
+ width = max(_ALIGN, math.floor(width / beta / _ALIGN) * _ALIGN)
66
+ return int(height), int(width)
67
+
68
+
69
+def _real_patches_at(
70
+ orig_h: int,
71
+ orig_w: int,
72
+ patch_cap: int,
73
+) -> int:
74
+ target_h, target_w = _compute_target_size(
75
+ orig_h,
76
+ orig_w,
77
+ _PF_FLOOR * _ALIGN * _ALIGN,
78
+ max(_PF_FLOOR, patch_cap) * _ALIGN * _ALIGN,
79
+ )
80
+ return (target_h // _ALIGN) * (target_w // _ALIGN)
81
+
82
+
83
+def _frame_hard_cap(seq_length: int) -> int:
84
+ need = max(1, (seq_length - _BUDGET_OVERHEAD) // (_PF_FLOOR + 15))
85
+ if need <= 1024:
86
+ return 1024
87
+ cap = 1
88
+ while cap < need:
89
+ cap <<= 1
90
+ return cap
91
+
92
+
93
+def _solve_degrade(
94
+ visual_budget: int,
95
+ duration: float,
96
+ orig_h: int,
97
+ orig_w: int,
98
+ orig_fps: float,
99
+ seq_length: int,
100
+) -> tuple[int, int]:
101
+ aligned_h = max(_ALIGN, round(orig_h / _ALIGN) * _ALIGN)
102
+ aligned_w = max(_ALIGN, round(orig_w / _ALIGN) * _ALIGN)
103
+ orig_max_pf = (aligned_h // _ALIGN) * (aligned_w // _ALIGN)
104
+ fps_cap = min(_FPS_CAP, max(orig_fps, 1e-6))
105
+ pf_cap = min(_PF_CEIL, max(orig_max_pf, _PF_FLOOR))
106
+ frame_cap = _frame_hard_cap(seq_length)
107
+
108
+ def usage(scale: float) -> tuple[int, float, int, int]:
109
+ fps = _FPS_MIN + scale * (fps_cap - _FPS_MIN)
110
+ patch_cap = _PF_FLOOR + scale * (pf_cap - _PF_FLOOR)
111
+ num_frames = max(
112
+ _MIN_FRAMES,
113
+ min(int(round(duration * fps)), frame_cap),
114
+ )
115
+ patches = _real_patches_at(
116
+ orig_h,
117
+ orig_w,
118
+ int(round(patch_cap)),
119
+ )
120
+ return (
121
+ num_frames * (patches + _FRAME_OVERHEAD),
122
+ fps,
123
+ int(round(patch_cap)),
124
+ num_frames,
125
+ )
126
+
127
+ if usage(1.0)[0] <= visual_budget:
128
+ _, _, patch_cap, num_frames = usage(1.0)
129
+ return num_frames, patch_cap
130
+
131
+ floor_cost = _real_patches_at(orig_h, orig_w, _PF_FLOOR) + _FRAME_OVERHEAD
132
+ if usage(0.0)[0] > visual_budget:
133
+ return max(
134
+ _MIN_FRAMES,
135
+ min(visual_budget // floor_cost, frame_cap),
136
+ ), _PF_FLOOR
137
+
138
+ low, high = 0.0, 1.0
139
+ for _ in range(50):
140
+ mid = (low + high) / 2
141
+ if usage(mid)[0] <= visual_budget:
142
+ low = mid
143
+ else:
144
+ high = mid
145
+ _, _, patch_cap, num_frames = usage(low)
146
+ return num_frames, patch_cap
147
+
148
+
149
+def _audio_tokens(duration: float, sample_rate: int) -> int:
150
+ if duration <= 0:
151
+ return 0
152
+ total_samples = int(duration * sample_rate)
153
+ chunk_samples = _AUDIO_CHUNK_SECONDS * sample_rate
154
+ num_tokens = 0
155
+ position = 0
156
+ while position < total_samples:
157
+ length = min(chunk_samples, total_samples - position)
158
+ num_tokens += math.ceil(length / _AUDIO_SAMPLES_PER_TOKEN)
159
+ position += chunk_samples
160
+ return num_tokens + 2
161
+
162
+
163
+def _decode_audio(
164
+ video_bytes: bytes,
165
+ sample_rate: int,
166
+) -> tuple[np.ndarray | None, float]:
167
+ from torchcodec.decoders import AudioDecoder
168
+
169
+ try:
170
+ decoder = AudioDecoder(
171
+ video_bytes,
172
+ sample_rate=sample_rate,
173
+ num_channels=1,
174
+ )
175
+ samples = decoder.get_all_samples()
176
+ except Exception:
177
+ return None, 0.0
178
+
179
+ waveform = samples.data
180
+ if waveform is None or waveform.numel() == 0:
181
+ return None, 0.0
182
+ if waveform.ndim == 2:
183
+ waveform = waveform.mean(dim=0) if waveform.shape[0] > 1 else waveform[0]
184
+ array = waveform.cpu().numpy()
185
+ pcm = (np.clip(array, -1.0, 1.0) * 32767.0).astype(np.int16)
186
+ return pcm, float(pcm.shape[0]) / sample_rate
187
+
188
+
189
+def _open_video(video_bytes: bytes):
190
+ from torchcodec.decoders import VideoDecoder
191
+
192
+ try:
193
+ return VideoDecoder(
194
+ video_bytes,
195
+ dimension_order="NHWC",
196
+ num_ffmpeg_threads=1,
197
+ seek_mode="approximate",
198
+ )
199
+ except TypeError:
200
+ return VideoDecoder(
201
+ video_bytes,
202
+ dimension_order="NHWC",
203
+ num_ffmpeg_threads=1,
204
+ )
205
+
206
+
207
+def _jpeg_roundtrip(image: Image.Image, quality: int) -> Image.Image:
208
+ buffer = io.BytesIO()
209
+ image.save(buffer, format="JPEG", quality=quality)
210
+ buffer.seek(0)
211
+ with Image.open(buffer) as decoded:
212
+ return decoded.convert("RGB").copy()
213
+
214
+
215
+def _decode_frames(
216
+ video_bytes: bytes,
217
+ visual_budget: int,
218
+ seq_length: int,
219
+ jpeg_quality: int,
220
+) -> tuple[list[tuple[float, Image.Image]], float]:
221
+ decoder = _open_video(video_bytes)
222
+ metadata = decoder.metadata
223
+ duration = float(metadata.duration_seconds or 0)
224
+ orig_h = int(metadata.height)
225
+ orig_w = int(metadata.width)
226
+ total_frames = int(getattr(metadata, "num_frames", 0) or 0)
227
+ orig_fps = float(getattr(metadata, "average_fps", 0) or 0) or 25.0
228
+ if duration <= 0 or orig_h <= 0 or orig_w <= 0:
229
+ raise ValueError(
230
+ "Invalid video metadata: "
231
+ f"duration={duration}, height={orig_h}, width={orig_w}"
232
+ )
233
+ if total_frames <= 0:
234
+ total_frames = max(1, int(duration * orig_fps))
235
+
236
+ num_frames, patch_cap = _solve_degrade(
237
+ visual_budget,
238
+ duration,
239
+ orig_h,
240
+ orig_w,
241
+ orig_fps,
242
+ seq_length,
243
+ )
244
+ aligned_h = max(_ALIGN, round(orig_h / _ALIGN) * _ALIGN)
245
+ aligned_w = max(_ALIGN, round(orig_w / _ALIGN) * _ALIGN)
246
+ orig_max_pf = (aligned_h // _ALIGN) * (aligned_w // _ALIGN)
247
+ max_pixels = min(patch_cap, orig_max_pf) * _ALIGN * _ALIGN
248
+ target_h, target_w = _compute_target_size(
249
+ orig_h,
250
+ orig_w,
251
+ _PF_FLOOR * _ALIGN * _ALIGN,
252
+ max_pixels,
253
+ )
254
+ num_frames = max(_MIN_FRAMES, min(num_frames, total_frames))
255
+ if num_frames == 1:
256
+ indices = [0]
257
+ else:
258
+ step = (total_frames - 1) / (num_frames - 1)
259
+ indices = [int(round(i * step)) for i in range(num_frames)]
260
+ indices = sorted({max(0, min(index, total_frames - 1)) for index in indices})
261
+
262
+ try:
263
+ decoded = decoder.get_frames_at(indices=indices).data
264
+ except (IndexError, RuntimeError):
265
+ safe = list(indices)
266
+ while safe and safe[-1] > 0:
267
+ safe.pop()
268
+ try:
269
+ decoded = decoder.get_frames_at(indices=safe).data
270
+ indices = safe
271
+ break
272
+ except (IndexError, RuntimeError):
273
+ continue
274
+ else:
275
+ raise
276
+
277
+ frames: list[tuple[float, Image.Image]] = []
278
+ for index, frame in zip(indices, decoded):
279
+ array = frame.cpu().numpy() if hasattr(frame, "cpu") else np.asarray(frame)
280
+ image = Image.fromarray(array)
281
+ if image.size != (target_w, target_h):
282
+ image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
283
+ frames.append(
284
+ (round(index / orig_fps, 3), _jpeg_roundtrip(image, jpeg_quality))
285
+ )
286
+ return frames, duration
287
+
288
+
289
+def _decoded_frames(video) -> tuple[np.ndarray, float]:
290
+ metadata = None
291
+ if isinstance(video, tuple):
292
+ video, metadata = video
293
+ frames = np.asarray(video)
294
+ if frames.ndim != 4 or frames.shape[-1] not in (3, 4):
295
+ raise TypeError(
296
+ "Decoded NOTE video must have shape (frames, height, width, channels)"
297
+ )
298
+ fps = float((metadata or {}).get("fps", 1.0)) if metadata else 1.0
299
+ return frames, max(fps, 1e-6)
300
+
301
+
302
+def _prepare_decoded_frames(
303
+ video,
304
+ visual_budget: int,
305
+ seq_length: int,
306
+ jpeg_quality: int,
307
+) -> tuple[list[tuple[float, Image.Image]], float]:
308
+ frames, fps = _decoded_frames(video)
309
+ duration = len(frames) / fps
310
+ orig_h, orig_w = frames.shape[1:3]
311
+ num_frames, patch_cap = _solve_degrade(
312
+ visual_budget,
313
+ duration,
314
+ orig_h,
315
+ orig_w,
316
+ fps,
317
+ seq_length,
318
+ )
319
+ num_frames = min(max(1, num_frames), len(frames))
320
+ indices = np.linspace(0, len(frames) - 1, num_frames).round().astype(int)
321
+ target_h, target_w = _compute_target_size(
322
+ orig_h,
323
+ orig_w,
324
+ _PF_FLOOR * _ALIGN * _ALIGN,
325
+ patch_cap * _ALIGN * _ALIGN,
326
+ )
327
+ output = []
328
+ for index in sorted(set(indices.tolist())):
329
+ image = Image.fromarray(frames[index]).convert("RGB")
330
+ if image.size != (target_w, target_h):
331
+ image = image.resize((target_w, target_h), Image.Resampling.BICUBIC)
332
+ output.append((round(index / fps, 3), _jpeg_roundtrip(image, jpeg_quality)))
333
+ return output, duration
334
+
335
+
336
+def _format_timestamp(seconds: float) -> str:
337
+ total_centiseconds = int(round(max(seconds, 0.0) * 100))
338
+ hours = total_centiseconds // 360000
339
+ minutes = (total_centiseconds // 6000) % 60
340
+ secs = (total_centiseconds // 100) % 60
341
+ centiseconds = total_centiseconds % 100
342
+ return f"{hours:02d}:{minutes:02d}:{secs:02d}.{centiseconds:02d}"
343
+
344
+
345
+def _group_bounds(
346
+ num_frames: int,
347
+ duration: float,
348
+ mode: str,
349
+ rng: random.Random,
350
+) -> list[int]:
351
+ if num_frames <= 1 or duration <= 0:
352
+ return [0, num_frames]
353
+ max_groups = min(
354
+ num_frames,
355
+ max(1, int(duration // _INTERLEAVE_MIN_SECONDS)),
356
+ )
357
+ if mode == "whole" or max_groups <= 1:
358
+ groups = 1
359
+ elif mode == "eval30":
360
+ groups = round(math.sqrt(max_groups))
361
+ elif mode == "eval_ek":
362
+ groups = round((max_groups - 1) / math.log(max_groups))
363
+ elif mode == "logk":
364
+ groups = round(math.exp(rng.uniform(0.0, math.log(max_groups))))
365
+ else:
366
+ raise ValueError(f"Unsupported video k_mode: {mode}")
367
+ groups = max(1, min(max_groups, groups))
368
+ if groups == 1:
369
+ return [0, num_frames]
370
+ if mode == "logk":
371
+ cuts = sorted(rng.sample(range(1, num_frames), groups - 1))
372
+ else:
373
+ cuts = sorted(
374
+ {
375
+ round(index * num_frames / groups)
376
+ for index in range(1, groups)
377
+ if 0 < round(index * num_frames / groups) < num_frames
378
+ }
379
+ )
380
+ return [0, *cuts, num_frames]
381
+
382
+
383
+def preprocess_dots3_note_video(
384
+ video,
385
+ *,
386
+ tokenizer,
387
+ question: str,
388
+ seq: int,
389
+ output_reserve: int | None = None,
390
+ audio_cap: float = 1.0,
391
+ audio_sample_rate: int = _AUDIO_SAMPLE_RATE,
392
+ k_mode: str = "eval_ek",
393
+ max_new_tokens: int = 0,
394
+) -> list[Dots3NoteVideoPart]:
395
+ """Expand one video into timestamped, interleaved image/audio parts."""
396
+ if seq <= 0:
397
+ raise ValueError(f"seq must be positive, got {seq}")
398
+ configured_reserve = seq // 4 if output_reserve is None else output_reserve
399
+ effective_reserve = max(configured_reserve, max_new_tokens)
400
+ if effective_reserve >= seq:
401
+ raise ValueError(
402
+ "output_reserve/max_new_tokens must leave room for video input"
403
+ )
404
+ if audio_cap < 0:
405
+ raise ValueError(f"audio_cap must be non-negative, got {audio_cap}")
406
+ if audio_sample_rate <= 0:
407
+ raise ValueError(f"audio_sample_rate must be positive, got {audio_sample_rate}")
408
+
409
+ seq_length = seq - effective_reserve
410
+ jpeg_quality = int(os.environ.get("XHS_VIDEO_JPEG_QUALITY", "85"))
411
+ pcm: np.ndarray | None = None
412
+ audio_duration = 0.0
413
+ if isinstance(video, bytes) and audio_cap > 0:
414
+ pcm, audio_duration = _decode_audio(video, audio_sample_rate)
415
+
416
+ audio_token_count = (
417
+ _audio_tokens(audio_duration, audio_sample_rate) if pcm is not None else 0
418
+ )
419
+ estimated_frames = max(1, int(audio_duration * _FPS_CAP))
420
+ max_groups = min(
421
+ estimated_frames,
422
+ max(1, int(audio_duration // _INTERLEAVE_MIN_SECONDS)),
423
+ )
424
+ reserved_audio_tokens = audio_token_count + 3 * max_groups
425
+ min_visual_tokens = _MIN_FRAMES * (_PF_FLOOR + _FRAME_OVERHEAD)
426
+ if (
427
+ audio_token_count > audio_cap * seq_length
428
+ or reserved_audio_tokens + min_visual_tokens + _BUDGET_OVERHEAD > seq_length
429
+ ):
430
+ pcm = None
431
+ audio_duration = 0.0
432
+ reserved_audio_tokens = 0
433
+
434
+ overhead = (
435
+ _token_len(
436
+ tokenizer,
437
+ "<|system|>You are a helpful assistant.<|endofsystem|>\n",
438
+ )
439
+ + 2
440
+ + _token_len(tokenizer, "<video_0>")
441
+ + 64
442
+ )
443
+ visual_budget = max(
444
+ _PF_FLOOR + _FRAME_OVERHEAD,
445
+ seq_length - overhead - reserved_audio_tokens,
446
+ )
447
+ if isinstance(video, bytes):
448
+ frames, video_duration = _decode_frames(
449
+ video,
450
+ visual_budget,
451
+ seq_length,
452
+ jpeg_quality,
453
+ )
454
+ else:
455
+ frames, video_duration = _prepare_decoded_frames(
456
+ video,
457
+ visual_budget,
458
+ seq_length,
459
+ jpeg_quality,
460
+ )
461
+
462
+ if pcm is None:
463
+ parts: list[Dots3NoteVideoPart] = []
464
+ for timestamp, image in frames:
465
+ parts.append(
466
+ Dots3NoteVideoPart("text", f"<{_format_timestamp(timestamp)}>")
467
+ )
468
+ parts.append(Dots3NoteVideoPart("image", image))
469
+ return parts
470
+
471
+ video_id = hashlib.sha1(video).hexdigest()
472
+ record_key = hashlib.sha1(f"{video_id}|{question}".encode()).hexdigest()
473
+ seed_hex = hashlib.sha1(f"42|flatten|{record_key}".encode()).hexdigest()
474
+ rng = random.Random(int(seed_hex[:8], 16))
475
+ bounds = _group_bounds(len(frames), audio_duration, k_mode, rng)
476
+ parts = []
477
+ for group in range(len(bounds) - 1):
478
+ start, end = bounds[group], bounds[group + 1]
479
+ if end <= start:
480
+ continue
481
+ start_time = 0.0 if group == 0 else frames[start][0]
482
+ end_time = audio_duration if group == len(bounds) - 2 else frames[end][0]
483
+ if end_time <= start_time:
484
+ end_time = start_time + audio_duration / max(1, len(bounds) - 1)
485
+ for timestamp, image in frames[start:end]:
486
+ parts.append(
487
+ Dots3NoteVideoPart("text", f"<{_format_timestamp(timestamp)}>")
488
+ )
489
+ parts.append(Dots3NoteVideoPart("image", image))
490
+ sample_start = max(0, int(round(start_time * audio_sample_rate)))
491
+ sample_end = min(len(pcm), int(round(end_time * audio_sample_rate)))
492
+ if sample_end > sample_start:
493
+ segment = np.ascontiguousarray(
494
+ pcm[sample_start:sample_end].astype(np.float32) / 32768.0
495
+ )
496
+ parts.append(Dots3NoteVideoPart("audio", segment))
497
+ return parts
vllm/models/dots3_note/nvidia/__init__.py
new
+3
@@ -0,0 +1,3 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NVIDIA-specific Dots3Note model implementations."""
vllm/models/dots3_note/nvidia/attention.py
new
+807
@@ -0,0 +1,807 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Dots3 NOTE sliding-window MLA attention backends for Hopper.
4
+
5
+Prefill and mixed batches expand the latent cache and use FlashAttention-3
6
+varlen MHA. Decode-only batches use the Triton absorbed-MQA kernel.
7
+"""
8
+
9
+from dataclasses import dataclass
10
+
11
+import torch
12
+
13
+import vllm._custom_ops as ops
14
+from vllm.model_executor.layers.attention.mla_attention import (
15
+ MLACommonDecodeMetadata,
16
+ MLACommonMetadata,
17
+ QueryLenSupport,
18
+)
19
+from vllm.platforms import current_platform
20
+from vllm.triton_utils import tl, triton
21
+from vllm.utils.torch_utils import is_quantized_kv_cache
22
+from vllm.v1.attention.backend import AttentionLayer, CommonAttentionMetadata
23
+from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
24
+ FlashAttnMLASparseBackend,
25
+ FlashAttnMLASparseImpl,
26
+ FlashAttnMLASparseMetadata,
27
+)
28
+from vllm.v1.attention.backends.mla.prefill.base import MLADimensions
29
+from vllm.v1.attention.backends.mla.prefill.flash_attn import (
30
+ FlashAttnPrefillBackend,
31
+)
32
+from vllm.v1.attention.backends.mla.sparse_utils import (
33
+ triton_convert_req_index_to_global_index,
34
+)
35
+from vllm.v1.attention.backends.mla.triton_mla import (
36
+ TritonMLABackend,
37
+ TritonMLAImpl,
38
+ TritonMLAMetadataBuilder,
39
+)
40
+from vllm.v1.worker.workspace import (
41
+ current_workspace_manager,
42
+ is_workspace_manager_initialized,
43
+)
44
+from vllm.vllm_flash_attn.flash_attn_interface import flash_attn_varlen_func
45
+
46
+
47
+@triton.jit
48
+def _gather_swa_kv_kernel(
49
+ kv_cache_ptr,
50
+ block_table_ptr,
51
+ seq_lens_ptr,
52
+ kv_out_ptr,
53
+ valid_out_ptr,
54
+ k_scale_ptr,
55
+ kv_cache_stride_block,
56
+ kv_cache_stride_token,
57
+ kv_cache_stride_dim,
58
+ block_table_stride_req,
59
+ kv_out_stride_req,
60
+ kv_out_stride_token,
61
+ kv_out_stride_dim,
62
+ valid_out_stride_req,
63
+ valid_out_stride_token,
64
+ GATHER_LEN: tl.constexpr,
65
+ KV_DIM: tl.constexpr,
66
+ BLOCK_T: tl.constexpr,
67
+ BLOCK_D: tl.constexpr,
68
+ PAGE_SIZE: tl.constexpr,
69
+):
70
+ req_idx = tl.program_id(0)
71
+ token_block_idx = tl.program_id(1)
72
+ dim_block_idx = tl.program_id(2)
73
+
74
+ seq_len = tl.load(seq_lens_ptr + req_idx).to(tl.int32)
75
+ gather_start = tl.maximum(seq_len - GATHER_LEN, 0)
76
+ token_offsets = token_block_idx * BLOCK_T + tl.arange(0, BLOCK_T)
77
+ logical_tokens = gather_start + token_offsets
78
+ logical_valid = (token_offsets < GATHER_LEN) & (logical_tokens < seq_len)
79
+
80
+ logical_pages = logical_tokens // PAGE_SIZE
81
+ page_offsets = logical_tokens - logical_pages * PAGE_SIZE
82
+ physical_pages = tl.load(
83
+ block_table_ptr + req_idx * block_table_stride_req + logical_pages,
84
+ mask=logical_valid,
85
+ other=-1,
86
+ ).to(tl.int64)
87
+ physical_valid = logical_valid & (physical_pages >= 0)
88
+
89
+ dim_offsets = dim_block_idx * BLOCK_D + tl.arange(0, BLOCK_D)
90
+ values = tl.load(
91
+ kv_cache_ptr
92
+ + physical_pages[:, None] * kv_cache_stride_block
93
+ + page_offsets[:, None] * kv_cache_stride_token
94
+ + dim_offsets[None, :] * kv_cache_stride_dim,
95
+ mask=physical_valid[:, None] & (dim_offsets[None, :] < KV_DIM),
96
+ other=0.0,
97
+ )
98
+ if values.dtype.is_fp8():
99
+ values = values.to(tl.float32) * tl.load(k_scale_ptr)
100
+ tl.store(
101
+ kv_out_ptr
102
+ + req_idx * kv_out_stride_req
103
+ + token_offsets[:, None] * kv_out_stride_token
104
+ + dim_offsets[None, :] * kv_out_stride_dim,
105
+ values,
106
+ mask=(token_offsets[:, None] < GATHER_LEN) & (dim_offsets[None, :] < KV_DIM),
107
+ )
108
+ tl.store(
109
+ valid_out_ptr
110
+ + req_idx * valid_out_stride_req
111
+ + token_offsets * valid_out_stride_token,
112
+ physical_valid,
113
+ mask=(dim_block_idx == 0) & (token_offsets < GATHER_LEN),
114
+ )
115
+
116
+
117
+@triton.jit
118
+def _apply_swa_score_mask_kernel(
119
+ scores_ptr,
120
+ seq_lens_ptr,
121
+ valid_ptr,
122
+ scores_stride_req,
123
+ scores_stride_head,
124
+ scores_stride_query,
125
+ scores_stride_token,
126
+ valid_stride_req,
127
+ valid_stride_token,
128
+ GATHER_LEN: tl.constexpr,
129
+ WINDOW_SIZE: tl.constexpr,
130
+ QUERY_LEN: tl.constexpr,
131
+ BLOCK_T: tl.constexpr,
132
+):
133
+ req_idx = tl.program_id(0)
134
+ head_idx = tl.program_id(1)
135
+ query_idx = tl.program_id(2)
136
+
137
+ token_offsets = tl.arange(0, BLOCK_T)
138
+ seq_len = tl.load(seq_lens_ptr + req_idx).to(tl.int32)
139
+ gather_start = tl.maximum(seq_len - GATHER_LEN, 0)
140
+ kv_positions = gather_start + token_offsets
141
+ query_position = seq_len - QUERY_LEN + query_idx
142
+ page_valid = tl.load(
143
+ valid_ptr + req_idx * valid_stride_req + token_offsets * valid_stride_token,
144
+ mask=token_offsets < GATHER_LEN,
145
+ other=0,
146
+ ).to(tl.int1)
147
+ valid = (
148
+ (token_offsets < GATHER_LEN)
149
+ & page_valid
150
+ & (kv_positions <= query_position)
151
+ & (kv_positions >= query_position - WINDOW_SIZE + 1)
152
+ & (query_position >= 0)
153
+ )
154
+ tl.store(
155
+ scores_ptr
156
+ + req_idx * scores_stride_req
157
+ + head_idx * scores_stride_head
158
+ + query_idx * scores_stride_query
159
+ + token_offsets * scores_stride_token,
160
+ tl.full((BLOCK_T,), -3.4028234663852886e38, tl.float32),
161
+ mask=(token_offsets < GATHER_LEN) & ~valid,
162
+ )
163
+
164
+
165
+@dataclass
166
+class Dots3NoteDecodeMetadata(MLACommonDecodeMetadata):
167
+ query_len: int
168
+
169
+
170
+@dataclass
171
+class _SlidingWindowChunk:
172
+ req_start: int
173
+ req_end: int
174
+ query_start: int
175
+ query_end: int
176
+ cu_seq_lens_q: torch.Tensor
177
+ cu_seq_lens_k: torch.Tensor
178
+ starts: torch.Tensor
179
+ token_to_seq: torch.Tensor
180
+ num_kv_tokens: int
181
+ max_seq_len_q: int
182
+ max_seq_len_k: int
183
+
184
+
185
+@dataclass
186
+class _SlidingWindowMetadata:
187
+ chunks: list[_SlidingWindowChunk]
188
+ workspace: torch.Tensor
189
+
190
+
191
+def _build_sliding_window_metadata(
192
+ *,
193
+ seq_lens_cpu: torch.Tensor,
194
+ query_start_loc_cpu: torch.Tensor,
195
+ sliding_window: int,
196
+ workspace: torch.Tensor,
197
+ workspace_size: int,
198
+ device: torch.device,
199
+) -> _SlidingWindowMetadata:
200
+ """Plan per-request latent-cache gathers for SWA varlen attention."""
201
+ query_lens_cpu = (query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]).to(
202
+ dtype=torch.int32
203
+ )
204
+ seq_lens_cpu = seq_lens_cpu.to(dtype=torch.int32)
205
+ kv_lens_cpu = torch.minimum(seq_lens_cpu, query_lens_cpu + sliding_window - 1)
206
+ starts_cpu = seq_lens_cpu - kv_lens_cpu
207
+
208
+ chunks: list[_SlidingWindowChunk] = []
209
+ req_start = 0
210
+ while req_start < query_lens_cpu.numel():
211
+ req_end = req_start
212
+ num_kv_tokens = 0
213
+ while req_end < query_lens_cpu.numel():
214
+ next_len = int(kv_lens_cpu[req_end].item())
215
+ if num_kv_tokens and num_kv_tokens + next_len > workspace_size:
216
+ break
217
+ if next_len > workspace_size:
218
+ raise ValueError(
219
+ "Dots3 NOTE SWA prefill window exceeds the MLA workspace: "
220
+ f"{next_len} > {workspace_size}"
221
+ )
222
+ num_kv_tokens += next_len
223
+ req_end += 1
224
+
225
+ query_lens = query_lens_cpu[req_start:req_end]
226
+ kv_lens = kv_lens_cpu[req_start:req_end]
227
+ num_reqs = req_end - req_start
228
+ cu_seq_lens_q_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
229
+ cu_seq_lens_k_cpu = torch.zeros(num_reqs + 1, dtype=torch.int32)
230
+ torch.cumsum(query_lens, 0, out=cu_seq_lens_q_cpu[1:])
231
+ torch.cumsum(kv_lens, 0, out=cu_seq_lens_k_cpu[1:])
232
+ token_to_seq_cpu = torch.repeat_interleave(
233
+ torch.arange(num_reqs, dtype=torch.int32), kv_lens
234
+ )
235
+ query_start = int(query_start_loc_cpu[req_start].item())
236
+ query_end = int(query_start_loc_cpu[req_end].item())
237
+ chunks.append(
238
+ _SlidingWindowChunk(
239
+ req_start=req_start,
240
+ req_end=req_end,
241
+ query_start=query_start,
242
+ query_end=query_end,
243
+ cu_seq_lens_q=cu_seq_lens_q_cpu.to(device, non_blocking=True),
244
+ cu_seq_lens_k=cu_seq_lens_k_cpu.to(device, non_blocking=True),
245
+ starts=starts_cpu[req_start:req_end].to(device, non_blocking=True),
246
+ token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),
247
+ num_kv_tokens=num_kv_tokens,
248
+ max_seq_len_q=int(query_lens.max().item()),
249
+ max_seq_len_k=int(kv_lens.max().item()),
250
+ )
251
+ )
252
+ req_start = req_end
253
+
254
+ return _SlidingWindowMetadata(chunks=chunks, workspace=workspace)
255
+
256
+
257
+class Dots3NoteFlashAttnPrefillBackend(FlashAttnPrefillBackend):
258
+ """FA3 varlen prefill for the NOTE SWA MLA dimensions."""
259
+
260
+ @classmethod
261
+ def supports_mla_dimensions(cls, dims: MLADimensions) -> bool:
262
+ return dims == MLADimensions(
263
+ qk_nope_head_dim=192,
264
+ qk_rope_head_dim=64,
265
+ v_head_dim=128,
266
+ )
267
+
268
+ def __init__(self, *args, **kwargs) -> None:
269
+ super().__init__(*args, **kwargs)
270
+ # FA3 on SM90 does not support NOTE's Q/K=256, V=128 combination via
271
+ # the different-head-dimension path. Padding V selects its supported
272
+ # equal-dimension varlen kernel.
273
+ self.requires_v_padding = True
274
+
275
+ def supports_quant_output(self, quant_key) -> bool:
276
+ return False
277
+
278
+ def run_sliding_window(
279
+ self,
280
+ q: torch.Tensor,
281
+ k: torch.Tensor,
282
+ v: torch.Tensor,
283
+ cu_seq_lens_q: torch.Tensor,
284
+ cu_seq_lens_k: torch.Tensor,
285
+ max_seq_len_q: int,
286
+ max_seq_len_k: int,
287
+ sliding_window: int,
288
+ ) -> torch.Tensor:
289
+ output = self._flash_attn_varlen_diff_headdims(
290
+ q=q,
291
+ k=k,
292
+ v=v,
293
+ cu_seqlens_q=cu_seq_lens_q,
294
+ cu_seqlens_k=cu_seq_lens_k,
295
+ max_seqlen_q=max_seq_len_q,
296
+ max_seqlen_k=max_seq_len_k,
297
+ softmax_scale=self.scale,
298
+ causal=True,
299
+ window_size=(sliding_window - 1, 0),
300
+ return_softmax_lse=False,
301
+ )
302
+ assert isinstance(output, torch.Tensor)
303
+ return output
304
+
305
+
306
+class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder):
307
+ """Keep decode on MQA and route prefill/mixed batches through FA3."""
308
+
309
+ query_len_support = QueryLenSupport.UNIFORM
310
+
311
+ def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
312
+ self.sliding_window = kv_cache_spec.sliding_window
313
+ super().__init__(kv_cache_spec, layer_names, vllm_config, device)
314
+
315
+ def _reserve_attn_logits_workspace(self) -> None:
316
+ if not is_workspace_manager_initialized():
317
+ return
318
+ batch = self.vllm_config.scheduler_config.max_num_seqs
319
+ max_query_len = self.reorder_batch_threshold or 1
320
+ gather_len = (self.sliding_window + max_query_len - 1 + 7) // 8 * 8
321
+ current_workspace_manager().get_simultaneous(
322
+ (
323
+ (
324
+ batch,
325
+ gather_len,
326
+ self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim,
327
+ ),
328
+ self.model_config.dtype,
329
+ ),
330
+ ((batch, gather_len), torch.bool),
331
+ )
332
+
333
+ def _build_decode(
334
+ self,
335
+ block_table_tensor: torch.Tensor,
336
+ seq_lens_device: torch.Tensor,
337
+ max_seq_len: int,
338
+ query_start_loc_cpu: torch.Tensor,
339
+ query_start_loc_device: torch.Tensor,
340
+ num_decode_tokens: int,
341
+ dcp_tot_seq_lens_device: torch.Tensor | None,
342
+ ) -> Dots3NoteDecodeMetadata:
343
+ del max_seq_len, query_start_loc_device, num_decode_tokens
344
+ query_len = int(query_start_loc_cpu[1] - query_start_loc_cpu[0])
345
+ return Dots3NoteDecodeMetadata(
346
+ block_table=block_table_tensor,
347
+ seq_lens=seq_lens_device,
348
+ dcp_tot_seq_lens=dcp_tot_seq_lens_device,
349
+ query_len=query_len,
350
+ )
351
+
352
+ def build(
353
+ self,
354
+ common_prefix_len: int,
355
+ common_attn_metadata: CommonAttentionMetadata,
356
+ fast_build: bool = False,
357
+ ) -> MLACommonMetadata:
358
+ metadata = super().build(
359
+ common_prefix_len,
360
+ common_attn_metadata,
361
+ fast_build=fast_build,
362
+ )
363
+ if metadata.prefill is None:
364
+ return metadata
365
+ if self.dcp_world_size > 1:
366
+ raise NotImplementedError(
367
+ "Dots3 NOTE SWA prefill does not support decode context parallelism"
368
+ )
369
+
370
+ reqs_start = metadata.num_decodes
371
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
372
+ assert seq_lens_cpu is not None
373
+ query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu[reqs_start:]
374
+ query_start_loc_cpu = query_start_loc_cpu - query_start_loc_cpu[0]
375
+ sliding_metadata = _build_sliding_window_metadata(
376
+ seq_lens_cpu=seq_lens_cpu[reqs_start:],
377
+ query_start_loc_cpu=query_start_loc_cpu,
378
+ sliding_window=self.sliding_window,
379
+ workspace=self.chunked_prefill_workspace,
380
+ workspace_size=self.chunked_prefill_workspace_size,
381
+ device=self.device,
382
+ )
383
+ metadata.prefill.chunked_context = None
384
+ metadata.prefill.sliding_window = sliding_metadata # type: ignore[attr-defined]
385
+ if metadata.num_decodes > 0 and metadata.num_prefills > 0:
386
+ query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
387
+ for chunk in sliding_metadata.chunks:
388
+ req_start, req_end = chunk.req_start, chunk.req_end
389
+ seq_lens = common_attn_metadata.seq_lens[
390
+ reqs_start + req_start : reqs_start + req_end
391
+ ]
392
+ query_lens = query_lens_cpu[req_start:req_end].to(
393
+ self.device, non_blocking=True
394
+ )
395
+ kv_lens = torch.minimum(seq_lens, query_lens + self.sliding_window - 1)
396
+ chunk.cu_seq_lens_k.zero_()
397
+ torch.cumsum(
398
+ kv_lens,
399
+ dim=0,
400
+ out=chunk.cu_seq_lens_k[1:],
401
+ dtype=torch.int32,
402
+ )
403
+ chunk.starts = seq_lens - kv_lens
404
+ token_ids = torch.arange(
405
+ chunk.num_kv_tokens,
406
+ dtype=torch.int32,
407
+ device=self.device,
408
+ )
409
+ chunk.token_to_seq = torch.searchsorted(
410
+ chunk.cu_seq_lens_k[1:], token_ids, right=True
411
+ ).to(torch.int32)
412
+ chunk.token_to_seq.clamp_max_(req_end - req_start - 1)
413
+ assert metadata.prefill.prefill_backend is not None
414
+ metadata.prefill.prefill_backend.prepare_metadata(metadata.prefill)
415
+ return metadata
416
+
417
+
418
+class Dots3NoteTritonMLABackend(TritonMLABackend):
419
+ """Internal NOTE SWA specialization; not a user-selectable backend."""
420
+
421
+ @classmethod
422
+ def get_supported_head_sizes(cls) -> list[int]:
423
+ return [1088]
424
+
425
+ @classmethod
426
+ def supports_sliding_window(cls) -> bool:
427
+ return True
428
+
429
+ @staticmethod
430
+ def get_impl_cls() -> type["Dots3NoteTritonMLAImpl"]:
431
+ return Dots3NoteTritonMLAImpl
432
+
433
+ @staticmethod
434
+ def get_builder_cls() -> type[Dots3NoteMLAMetadataBuilder]:
435
+ return Dots3NoteMLAMetadataBuilder
436
+
437
+
438
+class Dots3NoteTritonMLAImpl(TritonMLAImpl):
439
+ def __init__(
440
+ self,
441
+ num_heads: int,
442
+ head_size: int,
443
+ scale: float,
444
+ num_kv_heads: int,
445
+ alibi_slopes: list[float] | None,
446
+ sliding_window: int | None,
447
+ kv_cache_dtype: str,
448
+ logits_soft_cap: float | None,
449
+ attn_type: str,
450
+ kv_sharing_target_layer_name: str | None,
451
+ **mla_args,
452
+ ) -> None:
453
+ assert sliding_window is not None
454
+ super().__init__(
455
+ num_heads=num_heads,
456
+ head_size=head_size,
457
+ scale=scale,
458
+ num_kv_heads=num_kv_heads,
459
+ alibi_slopes=alibi_slopes,
460
+ sliding_window=None,
461
+ kv_cache_dtype=kv_cache_dtype,
462
+ logits_soft_cap=logits_soft_cap,
463
+ attn_type=attn_type,
464
+ kv_sharing_target_layer_name=kv_sharing_target_layer_name,
465
+ **mla_args,
466
+ )
467
+ self.sliding_window = sliding_window
468
+
469
+ def _forward_swa_mqa(
470
+ self,
471
+ q: torch.Tensor,
472
+ kv_cache: torch.Tensor,
473
+ decode: Dots3NoteDecodeMetadata,
474
+ k_scale: torch.Tensor,
475
+ ) -> tuple[torch.Tensor, None]:
476
+ query_len = decode.query_len
477
+ num_tokens, num_heads, head_size = q.shape
478
+ assert num_tokens % query_len == 0
479
+ assert head_size == self.head_size
480
+ num_reqs = num_tokens // query_len
481
+ block_table = decode.block_table[:num_reqs]
482
+ seq_lens = decode.seq_lens[:num_reqs]
483
+ gather_len = (self.sliding_window + query_len - 1 + 7) // 8 * 8
484
+
485
+ workspace_specs = (
486
+ ((num_reqs, gather_len, head_size), q.dtype),
487
+ ((num_reqs, gather_len), torch.bool),
488
+ )
489
+ if is_workspace_manager_initialized():
490
+ kv_latent, kv_valid = current_workspace_manager().get_simultaneous(
491
+ *workspace_specs
492
+ )
493
+ else:
494
+ kv_latent = torch.empty(
495
+ workspace_specs[0][0], dtype=q.dtype, device=q.device
496
+ )
497
+ kv_valid = torch.empty(
498
+ workspace_specs[1][0], dtype=torch.bool, device=q.device
499
+ )
500
+
501
+ block_t = 8
502
+ block_d = 128
503
+ _gather_swa_kv_kernel[
504
+ (
505
+ num_reqs,
506
+ triton.cdiv(gather_len, block_t),
507
+ triton.cdiv(head_size, block_d),
508
+ )
509
+ ](
510
+ kv_cache,
511
+ block_table,
512
+ seq_lens,
513
+ kv_latent,
514
+ kv_valid,
515
+ k_scale,
516
+ kv_cache.stride(0),
517
+ kv_cache.stride(1),
518
+ kv_cache.stride(2),
519
+ block_table.stride(0),
520
+ kv_latent.stride(0),
521
+ kv_latent.stride(1),
522
+ kv_latent.stride(2),
523
+ kv_valid.stride(0),
524
+ kv_valid.stride(1),
525
+ GATHER_LEN=gather_len,
526
+ KV_DIM=head_size,
527
+ BLOCK_T=block_t,
528
+ BLOCK_D=block_d,
529
+ PAGE_SIZE=kv_cache.shape[1],
530
+ num_warps=4,
531
+ )
532
+
533
+ q = q.view(num_reqs, query_len, num_heads, head_size)
534
+ scores = torch.bmm(
535
+ q.reshape(num_reqs, query_len * num_heads, head_size),
536
+ kv_latent.transpose(1, 2),
537
+ ).view(num_reqs, query_len, num_heads, gather_len)
538
+ scores = scores.float()
539
+ scores.mul_(self.scale)
540
+ scores_by_head = scores.transpose(1, 2)
541
+ _apply_swa_score_mask_kernel[(num_reqs, num_heads, query_len)](
542
+ scores_by_head,
543
+ seq_lens,
544
+ kv_valid,
545
+ scores_by_head.stride(0),
546
+ scores_by_head.stride(1),
547
+ scores_by_head.stride(2),
548
+ scores_by_head.stride(3),
549
+ kv_valid.stride(0),
550
+ kv_valid.stride(1),
551
+ GATHER_LEN=gather_len,
552
+ WINDOW_SIZE=self.sliding_window,
553
+ QUERY_LEN=query_len,
554
+ BLOCK_T=triton.next_power_of_2(gather_len),
555
+ num_warps=4,
556
+ )
557
+ probs = torch.softmax(scores, dim=-1).to(q.dtype)
558
+ output = torch.bmm(
559
+ probs.reshape(num_reqs, query_len * num_heads, gather_len),
560
+ kv_latent[..., : self.kv_lora_rank],
561
+ )
562
+ return output.view(num_tokens, num_heads, self.kv_lora_rank), None
563
+
564
+ def forward_mha(
565
+ self,
566
+ q: torch.Tensor,
567
+ kv_c_normed: torch.Tensor,
568
+ k_pe: torch.Tensor,
569
+ kv_c_and_k_pe_cache: torch.Tensor,
570
+ attn_metadata: MLACommonMetadata,
571
+ k_scale: torch.Tensor,
572
+ output: torch.Tensor,
573
+ output_scale: torch.Tensor | None = None,
574
+ ) -> None:
575
+ prefill = attn_metadata.prefill
576
+ sliding = getattr(prefill, "sliding_window", None)
577
+ if prefill is None or sliding is None:
578
+ return super().forward_mha(
579
+ q,
580
+ kv_c_normed,
581
+ k_pe,
582
+ kv_c_and_k_pe_cache,
583
+ attn_metadata,
584
+ k_scale,
585
+ output,
586
+ output_scale,
587
+ )
588
+ assert output_scale is None
589
+ assert isinstance(prefill.prefill_backend, Dots3NoteFlashAttnPrefillBackend)
590
+ use_fp8_prefill = prefill.q_data_type == current_platform.fp8_dtype()
591
+ output = output.view(-1, self.num_heads, self.v_head_dim)
592
+
593
+ for chunk in sliding.chunks:
594
+ toks = chunk.num_kv_tokens
595
+ workspace = sliding.workspace
596
+ block_table = prefill.block_table[chunk.req_start : chunk.req_end]
597
+ if is_quantized_kv_cache(self.kv_cache_dtype) and not use_fp8_prefill:
598
+ ops.gather_and_maybe_dequant_cache(
599
+ src_cache=kv_c_and_k_pe_cache,
600
+ dst=workspace,
601
+ block_table=block_table,
602
+ cu_seq_lens=chunk.cu_seq_lens_k,
603
+ token_to_seq=chunk.token_to_seq,
604
+ num_tokens=toks,
605
+ kv_cache_dtype=self.kv_cache_dtype,
606
+ scale=k_scale,
607
+ seq_starts=chunk.starts,
608
+ )
609
+ else:
610
+ ops.cp_gather_cache(
611
+ src_cache=kv_c_and_k_pe_cache,
612
+ dst=workspace,
613
+ block_table=block_table,
614
+ cu_seq_lens=chunk.cu_seq_lens_k,
615
+ batch_size=chunk.req_end - chunk.req_start,
616
+ seq_starts=chunk.starts,
617
+ )
618
+
619
+ # Slicing the latent part from the packed latent+RoPE cache keeps
620
+ # the full cache-entry row stride. FP8 linear kernels require
621
+ # contiguous rows and would otherwise read them at wrong offsets.
622
+ kv_c = workspace[:toks, : self.kv_lora_rank].contiguous()
623
+ weight_dtype = (
624
+ self.kv_b_proj.weight.dtype
625
+ if hasattr(self.kv_b_proj, "weight")
626
+ else self.kv_b_proj.params_dtype
627
+ )
628
+ if (
629
+ use_fp8_prefill or weight_dtype != current_platform.fp8_dtype()
630
+ ) and weight_dtype != torch.uint8:
631
+ kv_c = kv_c.to(weight_dtype)
632
+ k_pe_chunk = workspace[:toks, self.kv_lora_rank :].unsqueeze(1)
633
+ kv_nope = self.kv_b_proj(kv_c)[0].view(
634
+ -1, self.num_heads, self.qk_nope_head_dim + self.v_head_dim
635
+ )
636
+ if use_fp8_prefill:
637
+ kv_nope = kv_nope.to(prefill.q_data_type)
638
+ k_pe_chunk = k_pe_chunk.to(prefill.q_data_type)
639
+ k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)
640
+ k = self._concat_k_nope_k_pe(k_nope, k_pe_chunk)
641
+ chunk_output = prefill.prefill_backend.run_sliding_window(
642
+ q=q[chunk.query_start : chunk.query_end],
643
+ k=k,
644
+ v=v,
645
+ cu_seq_lens_q=chunk.cu_seq_lens_q,
646
+ cu_seq_lens_k=chunk.cu_seq_lens_k,
647
+ max_seq_len_q=chunk.max_seq_len_q,
648
+ max_seq_len_k=chunk.max_seq_len_k,
649
+ sliding_window=self.sliding_window,
650
+ )
651
+ output[chunk.query_start : chunk.query_end].copy_(
652
+ chunk_output[..., : self.v_head_dim]
653
+ )
654
+
655
+ def forward_mqa(
656
+ self,
657
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
658
+ kv_c_and_k_pe_cache: torch.Tensor,
659
+ attn_metadata: MLACommonMetadata,
660
+ layer: AttentionLayer,
661
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
662
+ assert kv_c_and_k_pe_cache.numel() > 0
663
+ assert attn_metadata.decode is not None
664
+ if isinstance(q, tuple):
665
+ q = torch.cat(q, dim=-1)
666
+
667
+ decode = attn_metadata.decode
668
+ assert isinstance(decode, Dots3NoteDecodeMetadata)
669
+ return self._forward_swa_mqa(
670
+ q,
671
+ kv_c_and_k_pe_cache,
672
+ decode,
673
+ layer._k_scale,
674
+ )
675
+
676
+
677
+class Dots3NotePaddedSparseBackend(FlashAttnMLASparseBackend):
678
+ """NOTE DSA backend for cache rows padded to the SWA latent width."""
679
+
680
+ @staticmethod
681
+ def get_name() -> str:
682
+ return "DOTS3_NOTE_PADDED_MLA_SPARSE"
683
+
684
+ @staticmethod
685
+ def get_impl_cls() -> type["Dots3NotePaddedSparseImpl"]:
686
+ return Dots3NotePaddedSparseImpl
687
+
688
+
689
+class Dots3NotePaddedSparseImpl(FlashAttnMLASparseImpl):
690
+ """Read top-k KV directly from uniformly padded cache rows."""
691
+
692
+ def _logical_cache(self, kv_cache: torch.Tensor) -> torch.Tensor:
693
+ assert kv_cache.shape[-1] >= self.head_size
694
+ return kv_cache[..., : self.head_size]
695
+
696
+ def do_kv_cache_update(
697
+ self,
698
+ kv_c_normed: torch.Tensor,
699
+ k_pe: torch.Tensor,
700
+ kv_cache: torch.Tensor,
701
+ slot_mapping: torch.Tensor,
702
+ kv_cache_dtype: str,
703
+ k_scale: torch.Tensor,
704
+ ) -> None:
705
+ super().do_kv_cache_update(
706
+ kv_c_normed,
707
+ k_pe,
708
+ self._logical_cache(kv_cache),
709
+ slot_mapping,
710
+ kv_cache_dtype,
711
+ k_scale,
712
+ )
713
+
714
+ def forward_mha( # type: ignore[override]
715
+ self,
716
+ q: torch.Tensor,
717
+ kv_c_normed: torch.Tensor,
718
+ k_pe: torch.Tensor,
719
+ kv_c_and_k_pe_cache: torch.Tensor,
720
+ attn_metadata: FlashAttnMLASparseMetadata,
721
+ k_scale: torch.Tensor,
722
+ output: torch.Tensor,
723
+ output_scale: torch.Tensor | None = None,
724
+ ) -> None:
725
+ super().forward_mha(
726
+ q,
727
+ kv_c_normed,
728
+ k_pe,
729
+ self._logical_cache(kv_c_and_k_pe_cache),
730
+ attn_metadata,
731
+ k_scale,
732
+ output,
733
+ output_scale,
734
+ )
735
+
736
+ def forward_mqa(
737
+ self,
738
+ q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],
739
+ kv_c_and_k_pe_cache: torch.Tensor,
740
+ attn_metadata: FlashAttnMLASparseMetadata,
741
+ layer: AttentionLayer,
742
+ ) -> tuple[torch.Tensor, torch.Tensor | None]:
743
+ if not isinstance(q, tuple):
744
+ raise NotImplementedError(
745
+ "Dots3NotePaddedSparseImpl expects split (q_nope, q_rope) input."
746
+ )
747
+ q_nope, q_rope = q
748
+ num_actual_toks = q_rope.shape[0]
749
+
750
+ assert self.topk_indices_buffer is not None
751
+ topk_indices = self.topk_indices_buffer[:num_actual_toks]
752
+ topk_indices, valid_counts = triton_convert_req_index_to_global_index(
753
+ attn_metadata.req_id_per_token[:num_actual_toks],
754
+ attn_metadata.block_table,
755
+ topk_indices,
756
+ BLOCK_SIZE=attn_metadata.block_size,
757
+ NUM_TOPK_TOKENS=topk_indices.shape[1],
758
+ return_valid_counts=True,
759
+ )
760
+ physical_head_size = kv_c_and_k_pe_cache.shape[-1]
761
+ block_stride, token_stride, head_stride = kv_c_and_k_pe_cache.stride()
762
+ if (
763
+ token_stride != physical_head_size
764
+ or head_stride != 1
765
+ or block_stride % physical_head_size != 0
766
+ ):
767
+ raise RuntimeError(
768
+ "Dots3 NOTE padded DSA cache must contain contiguous token rows; "
769
+ f"got shape={tuple(kv_c_and_k_pe_cache.shape)} and "
770
+ f"stride={kv_c_and_k_pe_cache.stride()}"
771
+ )
772
+ block_row_stride = block_stride // physical_head_size
773
+ block_indices = torch.div(
774
+ topk_indices,
775
+ attn_metadata.block_size,
776
+ rounding_mode="floor",
777
+ )
778
+ block_indices.clamp_min_(0).mul_(block_row_stride - attn_metadata.block_size)
779
+ topk_indices.add_(block_indices)
780
+
781
+ num_cache_rows = (
782
+ kv_c_and_k_pe_cache.shape[0] - 1
783
+ ) * block_row_stride + attn_metadata.block_size
784
+ kv_cache = kv_c_and_k_pe_cache.as_strided(
785
+ (num_cache_rows, physical_head_size),
786
+ (physical_head_size, 1),
787
+ )
788
+ cu_seqlens_q = torch.arange(
789
+ num_actual_toks + 1,
790
+ dtype=torch.int32,
791
+ device=q_rope.device,
792
+ )
793
+ output = flash_attn_varlen_func(
794
+ q=q_rope,
795
+ k=kv_cache[:, self.kv_lora_rank : self.head_size].unsqueeze(1).unsqueeze(1),
796
+ v=kv_cache[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1),
797
+ q_v=q_nope,
798
+ max_seqlen_q=1,
799
+ cu_seqlens_q=cu_seqlens_q,
800
+ max_seqlen_k=topk_indices.shape[1],
801
+ seqused_k=valid_counts,
802
+ block_table=topk_indices,
803
+ softmax_scale=self.scale,
804
+ causal=True,
805
+ fa_version=3,
806
+ )
807
+ return output, None
vllm/models/dots3_note/nvidia/audio.py
new
+305
@@ -0,0 +1,305 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import json
4
+import math
5
+from functools import lru_cache
6
+
7
+import torch
8
+import torch.nn.functional as F
9
+from torch import nn
10
+from transformers import WhisperConfig
11
+from transformers.audio_utils import mel_filter_bank
12
+
13
+from .audio_encoder import DotsSpeechEncoder
14
+
15
+SAMPLE_RATE = 16000
16
+N_FFT = 400
17
+HOP_LENGTH = 160
18
+DEFAULT_CHUNK_LENGTH_S = 60
19
+DEFAULT_CONV_TEMPORAL_STRIDE = 8
20
+DEFAULT_MERGE_FACTOR = 1
21
+N_SAMPLES = DEFAULT_CHUNK_LENGTH_S * SAMPLE_RATE
22
+STRIDE = HOP_LENGTH * DEFAULT_CONV_TEMPORAL_STRIDE * DEFAULT_MERGE_FACTOR
23
+ENCODER_SEQ_LEN = N_SAMPLES // HOP_LENGTH // 2
24
+
25
+
26
+class Dots3NoteAudioConfig:
27
+ def __init__(self, **kwargs):
28
+ self.encoder_type = kwargs.get("encoder_type", "dots")
29
+ self.whisper_config = kwargs.get("whisper_config", {})
30
+ self.whisper_adapter_in_dim = kwargs.get(
31
+ "whisper_adapter_in_dim", kwargs.get("adapter_in_dim", 1280)
32
+ )
33
+ self.whisper_adapter_out_dim = kwargs.get(
34
+ "whisper_adapter_out_dim", kwargs.get("adapter_out_dim", 2048)
35
+ )
36
+ self.sampling_rate = kwargs.get("sampling_rate", SAMPLE_RATE)
37
+ self.audio_comp_start = kwargs.get("audio_comp_start", "<|audio_comp_start|>")
38
+ self.audio_comp_span = kwargs.get("audio_comp_span", "<|audio_comp_pad|>")
39
+ self.audio_comp_end = kwargs.get("audio_comp_end", "<|audio_comp_end|>")
40
+ self.merge_factor = kwargs.get("merge_factor", DEFAULT_MERGE_FACTOR)
41
+ self.chunk_seconds = kwargs.get("chunk_seconds", DEFAULT_CHUNK_LENGTH_S)
42
+ self.use_conv2d_stem = kwargs.get("use_conv2d_stem", True)
43
+ self.use_rope = kwargs.get("use_rope", True)
44
+ self.use_rms_norm = kwargs.get("use_rms_norm", True)
45
+ self.use_causal = kwargs.get("use_causal", False)
46
+ self.downsample_hidden_size = kwargs.get("downsample_hidden_size", 480)
47
+ self.conv_chunksize = kwargs.get("conv_chunksize", 500)
48
+ self.conv_stem_gradient_checkpointing = kwargs.get(
49
+ "conv_stem_gradient_checkpointing", False
50
+ )
51
+ self.conv_bucket_step = kwargs.get("conv_bucket_step")
52
+ self.conv_bucket_max_elements = kwargs.get("conv_bucket_max_elements")
53
+ self.rope_parameters = kwargs.get(
54
+ "rope_parameters",
55
+ {
56
+ "partial_rotary_factor": 0.5,
57
+ "rope_theta": 10000.0,
58
+ "rope_type": "default",
59
+ },
60
+ )
61
+
62
+ @classmethod
63
+ def from_json(cls, path):
64
+ with open(path) as f:
65
+ return cls(**json.load(f))
66
+
67
+ @property
68
+ def conv_temporal_stride(self) -> int:
69
+ return 8 if self.use_conv2d_stem else 2
70
+
71
+ @property
72
+ def token_stride(self) -> int:
73
+ return HOP_LENGTH * self.conv_temporal_stride * self.merge_factor
74
+
75
+ @property
76
+ def chunk_samples(self) -> int:
77
+ return int(self.chunk_seconds * SAMPLE_RATE)
78
+
79
+ @property
80
+ def chunk_mel_frames(self) -> int:
81
+ return int(self.chunk_seconds * 100)
82
+
83
+
84
+def pad_or_trim(array, length=N_SAMPLES, axis=-1):
85
+ if array.shape[axis] > length:
86
+ array = array.index_select(
87
+ dim=axis, index=torch.arange(length, device=array.device)
88
+ )
89
+ if array.shape[axis] < length:
90
+ pad_widths = [(0, 0)] * array.ndim
91
+ pad_widths[axis] = (0, length - array.shape[axis])
92
+ array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
93
+ return array
94
+
95
+
96
+@lru_cache(maxsize=4)
97
+def _mel_filters(device, n_mels=128):
98
+ filters = mel_filter_bank(
99
+ num_frequency_bins=1 + N_FFT // 2,
100
+ num_mel_filters=n_mels,
101
+ min_frequency=0.0,
102
+ max_frequency=float(SAMPLE_RATE) / 2.0,
103
+ sampling_rate=SAMPLE_RATE,
104
+ norm="slaney",
105
+ mel_scale="slaney",
106
+ )
107
+ return torch.from_numpy(filters).T.contiguous().float().to(device)
108
+
109
+
110
+@lru_cache(maxsize=4)
111
+def _hann_window(device):
112
+ # Generate on CPU (default) then move, to preserve the exact reference
113
+ # window values. Direct on-device generation changes numerics slightly.
114
+ return torch.hann_window(N_FFT).to(device)
115
+
116
+
117
+def log_mel_spectrogram(audio, n_mels=128):
118
+ window = _hann_window(audio.device)
119
+ stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
120
+ magnitudes = stft[..., :-1].abs() ** 2
121
+ filters = _mel_filters(audio.device, n_mels)
122
+ mel_spec = filters @ magnitudes
123
+ log_spec = torch.clamp(mel_spec, min=1e-10).log10()
124
+ log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
125
+ log_spec = (log_spec + 4.0) / 4.0
126
+ return log_spec
127
+
128
+
129
+def compute_audio_token_length(
130
+ num_samples,
131
+ *,
132
+ sample_rate: int = SAMPLE_RATE,
133
+ chunk_seconds: int = DEFAULT_CHUNK_LENGTH_S,
134
+ hop_length: int = HOP_LENGTH,
135
+ conv_temporal_stride: int = DEFAULT_CONV_TEMPORAL_STRIDE,
136
+ merge_factor: int = DEFAULT_MERGE_FACTOR,
137
+):
138
+ stride = hop_length * conv_temporal_stride * merge_factor
139
+ total = 0
140
+ time_step = 0
141
+ while time_step * sample_rate < num_samples:
142
+ chunk_len = min(
143
+ num_samples - time_step * sample_rate, chunk_seconds * sample_rate
144
+ )
145
+ total += math.ceil(chunk_len / stride)
146
+ time_step += chunk_seconds
147
+ return total
148
+
149
+
150
+class DotsEncoderWithMask(nn.Module):
151
+ def __init__(self, config: Dots3NoteAudioConfig):
152
+ super().__init__()
153
+ whisper_config = WhisperConfig(**config.whisper_config)
154
+ whisper_config.use_rope = config.use_rope
155
+ whisper_config.rope_parameters = config.rope_parameters
156
+ whisper_config.use_rms_norm = config.use_rms_norm
157
+ whisper_config.use_causal = config.use_causal
158
+ whisper_config.use_conv2d_stem = config.use_conv2d_stem
159
+ whisper_config.downsample_hidden_size = config.downsample_hidden_size
160
+ whisper_config.conv_chunksize = config.conv_chunksize
161
+ whisper_config.conv_stem_gradient_checkpointing = (
162
+ config.conv_stem_gradient_checkpointing
163
+ )
164
+ whisper_config.conv_bucket_step = config.conv_bucket_step
165
+ whisper_config.conv_bucket_max_elements = config.conv_bucket_max_elements
166
+
167
+ self.speech_encoder = DotsSpeechEncoder(whisper_config)
168
+ self.merge_factor = config.merge_factor
169
+ self.chunk_seconds = config.chunk_seconds
170
+ self.chunk_samples = config.chunk_samples
171
+ self.chunk_mel_frames = config.chunk_mel_frames
172
+ self.conv_temporal_stride = config.conv_temporal_stride
173
+
174
+ @property
175
+ def device(self):
176
+ return next(self.speech_encoder.parameters()).device
177
+
178
+ def _forward_speech_encoder(
179
+ self,
180
+ mel_features: torch.Tensor,
181
+ input_seq_lens: torch.Tensor,
182
+ audio_sample_lens: list[int],
183
+ ) -> torch.Tensor:
184
+ """Run the eager speech encoder without server-side slicing/batching."""
185
+ mel_features = mel_features.to(dtype=torch.bfloat16, device=self.device)
186
+ return self.speech_encoder(
187
+ mel_features,
188
+ return_dict=True,
189
+ input_seq_lens=input_seq_lens,
190
+ audio_sample_lens=audio_sample_lens,
191
+ ).last_hidden_state
192
+
193
+ def encode_waveform(self, audio_waveform: torch.Tensor) -> torch.Tensor:
194
+ segments = []
195
+ time_step = 0
196
+ while time_step * SAMPLE_RATE < audio_waveform.shape[0]:
197
+ segments.append(
198
+ audio_waveform[
199
+ time_step * SAMPLE_RATE : (time_step + self.chunk_seconds)
200
+ * SAMPLE_RATE
201
+ ]
202
+ )
203
+ time_step += self.chunk_seconds
204
+
205
+ mel_features = []
206
+ token_lens = []
207
+ audio_sample_lens = []
208
+ for audio_segment in segments:
209
+ segment_length = audio_segment.shape[0]
210
+ token_len = (segment_length - 1) // (
211
+ HOP_LENGTH * self.conv_temporal_stride * self.merge_factor
212
+ ) + 1
213
+ pad_audio = pad_or_trim(audio_segment.flatten(), length=self.chunk_samples)
214
+ mel = log_mel_spectrogram(pad_audio)
215
+ assert mel.shape[1] == self.chunk_mel_frames
216
+ mel_features.append(mel)
217
+ token_lens.append(token_len)
218
+ audio_sample_lens.append(segment_length)
219
+
220
+ mel_features = torch.stack(mel_features, dim=0)
221
+ # Keep input_seq_lens on CPU: the conv2d bucket path reads it via
222
+ # ``.item()`` and CPU scalars avoid device->host syncs. The encoder's
223
+ # varlen path moves it to the device itself.
224
+ input_seq_lens = torch.tensor(token_lens, dtype=torch.long) * self.merge_factor
225
+ audio_embedding = self._forward_speech_encoder(
226
+ mel_features, input_seq_lens, audio_sample_lens
227
+ )
228
+
229
+ chunk_embeddings = []
230
+ for idx, token_len in enumerate(token_lens):
231
+ chunk_embeddings.append(
232
+ audio_embedding[idx, : token_len * self.merge_factor, :]
233
+ )
234
+ return torch.cat(chunk_embeddings, dim=0).unsqueeze(0)
235
+
236
+
237
+class AudioAdapter(nn.Module):
238
+ def __init__(self, in_dim, out_dim):
239
+ super().__init__()
240
+ self.proj = nn.Sequential(
241
+ nn.LayerNorm(in_dim),
242
+ nn.Linear(in_dim, out_dim),
243
+ nn.GELU(),
244
+ nn.Linear(out_dim, out_dim),
245
+ )
246
+
247
+ def forward(self, x):
248
+ return self.proj(x)
249
+
250
+
251
+class Dots3NoteAudioModel(nn.Module):
252
+ def __init__(self, config: Dots3NoteAudioConfig):
253
+ super().__init__()
254
+ self.config = config
255
+ if config.encoder_type != "dots":
256
+ raise ValueError("Dots3Note only supports encoder_type='dots'")
257
+ self.merge_factor = config.merge_factor
258
+ self.audio_adapter = AudioAdapter(
259
+ config.whisper_adapter_in_dim,
260
+ config.whisper_adapter_out_dim,
261
+ )
262
+
263
+ self.dots_encoder = DotsEncoderWithMask(config)
264
+
265
+ @property
266
+ def device(self):
267
+ return self.dots_encoder.device
268
+
269
+ def _merge_embeddings(self, embedding: torch.Tensor) -> torch.Tensor:
270
+ if self.merge_factor <= 1:
271
+ return embedding
272
+ return embedding.reshape(
273
+ embedding.shape[0],
274
+ embedding.shape[1] // self.merge_factor,
275
+ embedding.shape[2] * self.merge_factor,
276
+ )
277
+
278
+ def _encode_single_audio(self, audio_waveform):
279
+ embedding = self.dots_encoder.encode_waveform(audio_waveform)
280
+ embedding = self._merge_embeddings(embedding)
281
+ embedding = self.audio_adapter(embedding)
282
+ return embedding.squeeze(0)
283
+
284
+ def _split_waveforms(self, audio_inputs, lengths):
285
+ # Convert lengths to CPU ints once so the per-audio slicing below stays
286
+ # on the host and does not trigger a device->host sync per audio.
287
+ if isinstance(audio_inputs, list):
288
+ return audio_inputs
289
+ lengths_list = lengths.tolist()
290
+ waveforms = []
291
+ audio_start = 0
292
+ for length in lengths_list:
293
+ waveforms.append(audio_inputs[audio_start : audio_start + length])
294
+ audio_start += length
295
+ return waveforms
296
+
297
+ def forward(self, audio_inputs, lengths):
298
+ waveforms = self._split_waveforms(audio_inputs, lengths)
299
+ all_embeddings = []
300
+ token_lengths = []
301
+ for waveform in waveforms:
302
+ embedding = self._encode_single_audio(waveform)
303
+ all_embeddings.append(embedding)
304
+ token_lengths.append(embedding.shape[0])
305
+ return torch.cat(all_embeddings, dim=0), token_lengths
vllm/models/dots3_note/nvidia/audio_encoder.py
new
+736
@@ -0,0 +1,736 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+"""Dots-path speech encoder for inference only (single GPU).
5
+
6
+Ported from cybertron_alm ``dots_audio_encoder/modeling_whisper.py``.
7
+Upstream ``WhisperEncoder`` is exposed as :class:`DotsSpeechEncoder`.
8
+"""
9
+
10
+import math
11
+from typing import Any, ClassVar
12
+
13
+import torch
14
+from torch import nn
15
+from transformers.activations import ACT2FN
16
+from transformers.modeling_outputs import BaseModelOutput
17
+from transformers.modeling_utils import PreTrainedModel
18
+from transformers.models.whisper.configuration_whisper import WhisperConfig
19
+from transformers.utils import logging
20
+
21
+from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported
22
+
23
+_FLASH_ATTN_VERSION = 3 if is_fa_version_supported(3) else 2
24
+
25
+logger = logging.get_logger(__name__)
26
+
27
+__all__ = ["DotsSpeechEncoder"]
28
+
29
+
30
+class RMSNorm(nn.Module):
31
+ def __init__(self, dim: int, eps: float = 1e-6):
32
+ super().__init__()
33
+ self.weight = nn.Parameter(torch.ones(dim))
34
+ self.eps = eps
35
+
36
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
37
+ var = x.pow(2).mean(-1, keepdim=True)
38
+ x = x * torch.rsqrt(var + self.eps)
39
+ return self.weight * x
40
+
41
+
42
+def swiglu(x: torch.Tensor) -> torch.Tensor:
43
+ x1, x2 = x.chunk(2, dim=-1)
44
+ return torch.nn.functional.silu(x1) * x2
45
+
46
+
47
+class RotaryEmbedding(nn.Module):
48
+ def __init__(
49
+ self,
50
+ head_dim: int,
51
+ rope_parameters: dict,
52
+ base_seq_len: int = 0,
53
+ ):
54
+ super().__init__()
55
+ self.partial_rotary_factor = float(
56
+ rope_parameters.get("partial_rotary_factor", 1.0)
57
+ )
58
+ self.rope_theta = float(rope_parameters.get("rope_theta", 10000.0))
59
+ self.rope_type = rope_parameters.get("rope_type", "default")
60
+ rotary_dim = int(head_dim * self.partial_rotary_factor)
61
+ # keep even rotary dimension to align cos/sin pairs
62
+ self.rotary_dim = (rotary_dim // 2) * 2
63
+ self.attention_scaling = 1.0
64
+ if self.rope_type != "default":
65
+ logger.warning(
66
+ "RoPE type %s is not implemented in WhisperLv3; using default",
67
+ self.rope_type,
68
+ )
69
+ if self.rotary_dim > 0:
70
+ inv_freq = 1.0 / (
71
+ self.rope_theta
72
+ ** (
73
+ torch.arange(0, self.rotary_dim, 2, dtype=torch.float)
74
+ / max(self.rotary_dim, 1)
75
+ )
76
+ )
77
+ else:
78
+ inv_freq = torch.tensor([])
79
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
80
+ self.base_seq_len = max(0, int(base_seq_len))
81
+ self._cache: (
82
+ tuple[int, torch.dtype, torch.device, torch.Tensor, torch.Tensor] | None
83
+ ) = None
84
+
85
+ @torch.no_grad()
86
+ def get_cos_sin(
87
+ self, position_ids: torch.Tensor, dtype: torch.dtype, device: torch.device
88
+ ) -> tuple[torch.Tensor | None, torch.Tensor | None]:
89
+ if self.rotary_dim == 0:
90
+ return None, None
91
+ seq_len = position_ids.shape[-1]
92
+ if position_ids.shape[0] == 1 and self._cache is not None:
93
+ cached_seq_len, cached_dtype, cached_device, cached_cos, cached_sin = (
94
+ self._cache
95
+ )
96
+ if (
97
+ cached_seq_len >= seq_len
98
+ and cached_dtype == dtype
99
+ and cached_device == device
100
+ ):
101
+ return cached_cos[:, :seq_len, :], cached_sin[:, :seq_len, :]
102
+
103
+ if position_ids.shape[0] == 1:
104
+ cache_seq_len = max(seq_len, self.base_seq_len)
105
+ position_ids = torch.arange(cache_seq_len, device=device)[None, :]
106
+ # [1, D/2, 1]
107
+ if self.inv_freq.dtype != torch.float32:
108
+ inv_freq = 1.0 / (
109
+ self.rope_theta
110
+ ** (
111
+ torch.arange(
112
+ 0, self.rotary_dim, 2, dtype=torch.float, device=device
113
+ )
114
+ / max(self.rotary_dim, 1)
115
+ )
116
+ )
117
+ else:
118
+ inv_freq = self.inv_freq.to(device=device)
119
+ inv_freq_expanded = inv_freq[None, :, None]
120
+ # [B, 1, T]
121
+ position_ids_expanded = position_ids[:, None, :].float()
122
+ # Force float32 since bfloat16 loses precision on long contexts
123
+ # See https://github.com/huggingface/transformers/pull/29285
124
+ device_type = device.type if isinstance(device, torch.device) else "cpu"
125
+ with torch.autocast(device_type=device_type, enabled=False):
126
+ freqs = (
127
+ inv_freq_expanded.float() @ position_ids_expanded.float()
128
+ ).transpose(1, 2)
129
+ emb = torch.cat((freqs, freqs), dim=-1)
130
+ cos = emb.cos() * self.attention_scaling
131
+ sin = emb.sin() * self.attention_scaling
132
+ cos = cos.to(dtype)
133
+ sin = sin.to(dtype)
134
+ if position_ids.shape[0] == 1:
135
+ self._cache = (position_ids.shape[-1], dtype, device, cos, sin)
136
+ return cos[:, :seq_len, :], sin[:, :seq_len, :]
137
+ return cos, sin
138
+
139
+
140
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
141
+ x1 = x[..., : x.shape[-1] // 2]
142
+ x2 = x[..., x.shape[-1] // 2 :]
143
+ return torch.cat((-x2, x1), dim=-1)
144
+
145
+
146
+def apply_rotary_pos_emb(
147
+ q: torch.Tensor,
148
+ k: torch.Tensor,
149
+ cos: torch.Tensor | None,
150
+ sin: torch.Tensor | None,
151
+) -> tuple[torch.Tensor, torch.Tensor]:
152
+ if cos is None or sin is None:
153
+ return q, k
154
+
155
+ if q.dim() == 4:
156
+ # [B, T, H, D]
157
+ rotary_dim = cos.shape[-1]
158
+ cos = cos.unsqueeze(2) # [B, T, 1, rotary_dim]
159
+ sin = sin.unsqueeze(2)
160
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
161
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
162
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
163
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
164
+ return torch.cat([q_embed, q_pass], dim=-1), torch.cat(
165
+ [k_embed, k_pass], dim=-1
166
+ )
167
+
168
+ if q.dim() == 3:
169
+ # [tokens, H, D] (varlen flattened)
170
+ rotary_dim = cos.shape[-1]
171
+ cos = cos.unsqueeze(1) # [tokens, 1, rotary_dim]
172
+ sin = sin.unsqueeze(1)
173
+ q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]
174
+ k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]
175
+ q_embed = (q_rot * cos) + (rotate_half(q_rot) * sin)
176
+ k_embed = (k_rot * cos) + (rotate_half(k_rot) * sin)
177
+ return torch.cat([q_embed, q_pass], dim=-1), torch.cat(
178
+ [k_embed, k_pass], dim=-1
179
+ )
180
+
181
+ return q, k
182
+
183
+
184
+class WhisperPositionalEmbedding(nn.Embedding):
185
+ def __init__(
186
+ self, num_positions: int, embedding_dim: int, padding_idx: int | None = None
187
+ ):
188
+ super().__init__(num_positions, embedding_dim)
189
+
190
+ def forward(self, input_ids, past_key_values_length=0):
191
+ return self.weight[
192
+ past_key_values_length : past_key_values_length + input_ids.shape[1]
193
+ ]
194
+
195
+
196
+class WhisperAttention(nn.Module):
197
+ """Multi-headed attention from 'Attention Is All You Need' paper"""
198
+
199
+ def __init__(
200
+ self,
201
+ embed_dim: int,
202
+ num_heads: int,
203
+ dropout: float = 0.0,
204
+ is_decoder: bool = False,
205
+ bias: bool = True,
206
+ ):
207
+ super().__init__()
208
+ self.embed_dim = embed_dim
209
+ self.num_heads = num_heads
210
+ self.dropout = dropout
211
+ self.head_dim = embed_dim // num_heads
212
+
213
+ if (self.head_dim * num_heads) != self.embed_dim:
214
+ raise ValueError(
215
+ f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
216
+ f" and `num_heads`: {num_heads})."
217
+ )
218
+ self.scaling = self.head_dim**-0.5
219
+ self.is_decoder = is_decoder
220
+
221
+ self.k_proj = nn.Linear(embed_dim, embed_dim, bias=False)
222
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
223
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
224
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
225
+
226
+ def forward_flash_attn(
227
+ self,
228
+ hidden_states,
229
+ cu_seqlens_q=None,
230
+ cu_seqlens_kv=None,
231
+ max_seqlen_q=None,
232
+ max_seqlen_kv=None,
233
+ output_attentions=False,
234
+ rotary_cos=None,
235
+ rotary_sin=None,
236
+ ):
237
+ """Dense eager attention with SGLang FA3 for packed variable-length input."""
238
+ query_states = self.q_proj(hidden_states)
239
+ key_states = self.k_proj(hidden_states)
240
+ value_states = self.v_proj(hidden_states)
241
+ if cu_seqlens_q is None:
242
+ bsz, tgt_len, _ = hidden_states.size()
243
+ query_states = query_states.view(
244
+ bsz, tgt_len, self.num_heads, self.head_dim
245
+ )
246
+ key_states = key_states.view(bsz, tgt_len, self.num_heads, self.head_dim)
247
+ value_states = value_states.view(
248
+ bsz, tgt_len, self.num_heads, self.head_dim
249
+ )
250
+ if rotary_cos is not None and rotary_sin is not None:
251
+ cos = rotary_cos[:, :tgt_len, :]
252
+ sin = rotary_sin[:, :tgt_len, :]
253
+ query_states, key_states = apply_rotary_pos_emb(
254
+ query_states, key_states, cos, sin
255
+ )
256
+ attn_output, attn_probs = self._eager_attention(
257
+ query_states, key_states, value_states, output_attentions
258
+ )
259
+ attn_output = attn_output.view(bsz, tgt_len, self.embed_dim)
260
+ else:
261
+ query_states = query_states.view(-1, self.num_heads, self.head_dim)
262
+ key_states = key_states.view(-1, self.num_heads, self.head_dim)
263
+ value_states = value_states.view(-1, self.num_heads, self.head_dim)
264
+ if rotary_cos is not None and rotary_sin is not None:
265
+ query_states, key_states = apply_rotary_pos_emb(
266
+ query_states, key_states, rotary_cos, rotary_sin
267
+ )
268
+ result = flash_attn_varlen_func(
269
+ query_states,
270
+ key_states,
271
+ value_states,
272
+ cu_seqlens_q=cu_seqlens_q,
273
+ cu_seqlens_k=cu_seqlens_kv,
274
+ max_seqlen_q=max_seqlen_q,
275
+ max_seqlen_k=max_seqlen_kv,
276
+ softmax_scale=self.scaling,
277
+ causal=self.is_decoder,
278
+ return_softmax_lse=output_attentions,
279
+ fa_version=_FLASH_ATTN_VERSION,
280
+ )
281
+ if isinstance(result, tuple):
282
+ attn_output = result[0]
283
+ attn_probs = result[1] if output_attentions else None
284
+ else:
285
+ attn_output = result
286
+ attn_probs = None
287
+ attn_output = attn_output.view(-1, self.embed_dim)
288
+
289
+ return self.out_proj(attn_output), attn_probs
290
+
291
+ def _eager_attention(self, query, key, value, output_attentions):
292
+ query = query.transpose(1, 2)
293
+ key = key.transpose(1, 2)
294
+ value = value.transpose(1, 2)
295
+ scores = torch.matmul(query, key.transpose(-2, -1)) * self.scaling
296
+ if self.is_decoder:
297
+ seq_len = scores.shape[-1]
298
+ causal_mask = torch.ones(
299
+ seq_len, seq_len, dtype=torch.bool, device=scores.device
300
+ ).triu(1)
301
+ scores.masked_fill_(causal_mask, torch.finfo(scores.dtype).min)
302
+ probabilities = torch.softmax(scores, dim=-1, dtype=torch.float32).to(
303
+ query.dtype
304
+ )
305
+ output = torch.matmul(probabilities, value).transpose(1, 2)
306
+ return output, probabilities if output_attentions else None
307
+
308
+
309
+# Copied from transformers.models.mbart.modeling_mbart.MBartEncoderLayer with MBart->Whisper
310
+class WhisperEncoderLayer(nn.Module):
311
+ def __init__(self, config: WhisperConfig):
312
+ super().__init__()
313
+ self.embed_dim = config.d_model
314
+ use_causal = getattr(config, "use_causal", False)
315
+ self.self_attn = WhisperAttention(
316
+ embed_dim=self.embed_dim,
317
+ num_heads=config.encoder_attention_heads,
318
+ dropout=config.attention_dropout,
319
+ is_decoder=use_causal, # enable causal attention when use_causal=True
320
+ )
321
+ self.use_causal = use_causal
322
+ norm_cls = RMSNorm if getattr(config, "use_rms_norm", False) else nn.LayerNorm
323
+ self.self_attn_layer_norm = norm_cls(self.embed_dim)
324
+ self.dropout = config.dropout
325
+ self.activation_fn = (
326
+ ACT2FN[config.activation_function]
327
+ if config.activation_function != "swiglu"
328
+ else swiglu
329
+ )
330
+ self.use_swiglu = config.activation_function == "swiglu"
331
+ self.activation_dropout = config.activation_dropout
332
+ ffn_dim = config.encoder_ffn_dim
333
+ fc1_out = ffn_dim * 2 if self.use_swiglu else ffn_dim
334
+ self.fc1 = nn.Linear(self.embed_dim, fc1_out)
335
+ self.fc2 = nn.Linear(ffn_dim, self.embed_dim)
336
+ self.final_layer_norm = norm_cls(self.embed_dim)
337
+
338
+ def forward(
339
+ self,
340
+ hidden_states: torch.Tensor,
341
+ cu_seqlens_q: torch.Tensor = None,
342
+ cu_seqlens_kv: torch.Tensor = None,
343
+ max_seqlen_q: int | None = None,
344
+ max_seqlen_kv: int | None = None,
345
+ output_attentions: bool = False,
346
+ rotary_cos: torch.Tensor | None = None,
347
+ rotary_sin: torch.Tensor | None = None,
348
+ ) -> torch.Tensor:
349
+ """
350
+ Args:
351
+ hidden_states (`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
352
+ attention_mask (`torch.FloatTensor`): attention mask of size
353
+ `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
354
+ layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size
355
+ `(encoder_attention_heads,)`.
356
+ output_attentions (`bool`, *optional*):
357
+ Whether or not to return the attentions tensors of all attention layers. See `attentions` under
358
+ returned tensors for more detail.
359
+ """
360
+ residual = hidden_states
361
+ hidden_states = self.self_attn_layer_norm(hidden_states)
362
+
363
+ hidden_states, attn_weights = self.self_attn.forward_flash_attn(
364
+ hidden_states=hidden_states,
365
+ cu_seqlens_q=cu_seqlens_q,
366
+ cu_seqlens_kv=cu_seqlens_kv,
367
+ max_seqlen_q=max_seqlen_q,
368
+ max_seqlen_kv=max_seqlen_kv,
369
+ output_attentions=output_attentions,
370
+ rotary_cos=rotary_cos,
371
+ rotary_sin=rotary_sin,
372
+ )
373
+ hidden_states = nn.functional.dropout(
374
+ hidden_states, p=self.dropout, training=self.training
375
+ )
376
+ hidden_states = residual + hidden_states
377
+
378
+ residual = hidden_states
379
+ hidden_states = self.final_layer_norm(hidden_states)
380
+ if self.use_swiglu:
381
+ hidden_states = swiglu(self.fc1(hidden_states))
382
+ else:
383
+ hidden_states = self.activation_fn(self.fc1(hidden_states))
384
+ hidden_states = nn.functional.dropout(
385
+ hidden_states, p=self.activation_dropout, training=self.training
386
+ )
387
+ hidden_states = self.fc2(hidden_states)
388
+ hidden_states = nn.functional.dropout(
389
+ hidden_states, p=self.dropout, training=self.training
390
+ )
391
+ hidden_states = residual + hidden_states
392
+
393
+ if hidden_states.dtype == torch.float16 and (
394
+ torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()
395
+ ):
396
+ clamp_value = torch.finfo(hidden_states.dtype).max - 1000
397
+ hidden_states = torch.clamp(
398
+ hidden_states, min=-clamp_value, max=clamp_value
399
+ )
400
+
401
+ outputs: tuple[Any, ...] = (hidden_states,)
402
+
403
+ if output_attentions:
404
+ outputs += (attn_weights,)
405
+
406
+ return outputs
407
+
408
+
409
+class DotsSpeechPreTrainedModel(PreTrainedModel):
410
+ config_class = WhisperConfig
411
+ base_model_prefix = "model"
412
+ main_input_name = "input_features"
413
+ supports_gradient_checkpointing = False
414
+ _no_split_modules: ClassVar[list[str]] = ["WhisperEncoderLayer"]
415
+
416
+ def _init_weights(self, module):
417
+ std = self.config.init_std
418
+ if isinstance(module, (nn.Linear, nn.Conv1d, nn.Conv2d)):
419
+ module.weight.data.normal_(mean=0.0, std=std)
420
+ if module.bias is not None:
421
+ module.bias.data.zero_()
422
+ elif isinstance(module, nn.Embedding):
423
+ module.weight.data.normal_(mean=0.0, std=std)
424
+ if module.padding_idx is not None:
425
+ module.weight.data[module.padding_idx].zero_()
426
+
427
+
428
+class DotsSpeechEncoder(DotsSpeechPreTrainedModel):
429
+ """
430
+ Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a
431
+ [`WhisperEncoderLayer`].
432
+
433
+ Args:
434
+ config: WhisperConfig
435
+ """
436
+
437
+ def __init__(self, config: WhisperConfig):
438
+ super().__init__(config)
439
+ self.dropout = config.dropout
440
+
441
+ embed_dim = config.d_model
442
+ self.use_latent_input = getattr(config, "use_latent_input", False)
443
+ self.use_causal = getattr(config, "use_causal", False)
444
+ self.use_conv2d_stem = getattr(config, "use_conv2d_stem", False)
445
+ latent_dim = getattr(config, "latent_dim", None)
446
+ if self.use_latent_input and latent_dim is None:
447
+ raise ValueError(
448
+ "DotsSpeechEncoder: use_latent_input=True requires config.latent_dim"
449
+ )
450
+ if self.use_conv2d_stem and self.use_latent_input:
451
+ raise ValueError(
452
+ "DotsSpeechEncoder: use_conv2d_stem and use_latent_input are mutually exclusive"
453
+ )
454
+ self.num_mel_bins = (
455
+ latent_dim if self.use_latent_input and latent_dim is not None else 128
456
+ )
457
+ self.padding_idx = config.pad_token_id
458
+ self.max_source_positions = config.max_source_positions
459
+ self.embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0
460
+
461
+ if self.use_conv2d_stem:
462
+ # Conv2D stem: 3 layers of stride-2 for 8x downsampling
463
+ dhs = getattr(config, "downsample_hidden_size", 480)
464
+ # Causal: keep freq padding, remove time padding (handled by pad(14,0) in forward)
465
+ conv_padding = (1, 0) if self.use_causal else 1
466
+ self.conv2d1 = nn.Conv2d(
467
+ 1, dhs, kernel_size=3, stride=2, padding=conv_padding
468
+ )
469
+ self.conv2d2 = nn.Conv2d(
470
+ dhs, dhs, kernel_size=3, stride=2, padding=conv_padding
471
+ )
472
+ self.conv2d3 = nn.Conv2d(
473
+ dhs, dhs, kernel_size=3, stride=2, padding=conv_padding
474
+ )
475
+ # After 3x stride-2 on freq=128: 128→64→32→16; linear projects dhs*16 → embed_dim
476
+ freq_after = self.num_mel_bins
477
+ for _ in range(3):
478
+ freq_after = (freq_after + 1) // 2
479
+ self.conv_out = nn.Linear(dhs * freq_after, embed_dim, bias=False)
480
+ self.conv1 = None
481
+ self.conv2 = None
482
+ elif self.use_latent_input:
483
+ # Latent path keeps length (stride=1) and applies GLU after conv1
484
+ self.conv1 = nn.Conv1d(
485
+ self.num_mel_bins, embed_dim * 2, kernel_size=3, stride=1, padding=1
486
+ )
487
+ self.conv2 = nn.Conv1d(
488
+ embed_dim, embed_dim, kernel_size=3, stride=1, padding=1
489
+ )
490
+ else:
491
+ self.conv1 = nn.Conv1d(
492
+ self.num_mel_bins, embed_dim, kernel_size=3, padding=1
493
+ )
494
+ self.conv2 = nn.Conv1d(
495
+ embed_dim, embed_dim, kernel_size=3, stride=2, padding=1
496
+ )
497
+
498
+ self.use_rope = getattr(config, "use_rope", False)
499
+ rope_parameters = getattr(config, "rope_parameters", {}) or {}
500
+ self.rope_parameters = rope_parameters
501
+ if self.use_rope:
502
+ head_dim = embed_dim // config.encoder_attention_heads
503
+ self.rotary_embedding = RotaryEmbedding(
504
+ head_dim,
505
+ rope_parameters,
506
+ base_seq_len=self.max_source_positions,
507
+ )
508
+ self.embed_positions = None
509
+ else:
510
+ self.embed_positions = nn.Embedding(self.max_source_positions, embed_dim)
511
+
512
+ self.layers = nn.ModuleList(
513
+ [WhisperEncoderLayer(config) for _ in range(config.encoder_layers)]
514
+ )
515
+ norm_cls = RMSNorm if getattr(config, "use_rms_norm", False) else nn.LayerNorm
516
+ self.layer_norm = norm_cls(config.d_model)
517
+
518
+ self.post_init()
519
+
520
+ def get_input_embeddings(self) -> nn.Module:
521
+ if self.use_conv2d_stem:
522
+ return self.conv2d1
523
+ return self.conv1
524
+
525
+ def set_input_embeddings(self, value: nn.Module):
526
+ self.conv1 = value
527
+
528
+ @staticmethod
529
+ def _temporal_mask(feat, valid_lens):
530
+ """Zero out temporal positions >= valid_lens. feat: [B,C,F,T], valid_lens: [B] tensor."""
531
+ T = feat.shape[-1]
532
+ mask = torch.arange(T, device=feat.device)[None, :] < valid_lens[:, None]
533
+ return mask[:, None, None, :]
534
+
535
+ def _conv2d_stem_one_chunk(self, chunk, chunk_valid_mel_lens=None):
536
+ """Run 3x Conv2d(stride=2) + GELU with per-layer masking."""
537
+ # Causal: left-pad time by 14 = 2 + 4 + 8 (receptive field of the 3-layer
538
+ # stride-2 stack mapped back to input time). Combined with conv padding=(1,0),
539
+ # this makes the whole conv2d stem strictly causal.
540
+ if self.use_causal:
541
+ chunk = nn.functional.pad(chunk, (14, 0))
542
+ if chunk_valid_mel_lens is not None:
543
+ chunk_valid_mel_lens = chunk_valid_mel_lens + 14
544
+ # Step 0: mask mel — silence_mel(-1.5) → 0
545
+ if chunk_valid_mel_lens is not None:
546
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
547
+ chunk = nn.functional.gelu(self.conv2d1(chunk))
548
+ # Step 1: mask conv1 output — gelu(conv(0)+bias) → 0
549
+ if chunk_valid_mel_lens is not None:
550
+ chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
551
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
552
+ chunk = nn.functional.gelu(self.conv2d2(chunk))
553
+ # Step 2: mask conv2 output — gelu(bias) → 0
554
+ if chunk_valid_mel_lens is not None:
555
+ chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
556
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
557
+ chunk = nn.functional.gelu(self.conv2d3(chunk))
558
+ # Step 3: mask conv3 output — gelu(bias) → 0
559
+ if chunk_valid_mel_lens is not None:
560
+ chunk_valid_mel_lens = (chunk_valid_mel_lens + 1) // 2
561
+ chunk = chunk * self._temporal_mask(chunk, chunk_valid_mel_lens)
562
+ return chunk
563
+
564
+ def _forward_conv2d_stem(
565
+ self, input_features, input_seq_lens=None, audio_sample_lens=None
566
+ ):
567
+ """Conv2D stem: 3x Conv2d(stride=2) → 8x downsample. Returns [B, T/8, embed_dim]."""
568
+ x = input_features.unsqueeze(1) # [B, 1, 128, T]
569
+
570
+ if audio_sample_lens is not None:
571
+ hop_length = 160
572
+ if not isinstance(audio_sample_lens, torch.Tensor):
573
+ audio_sample_lens = torch.tensor(audio_sample_lens, device=x.device)
574
+ valid_mel_lens = audio_sample_lens.to(x.device) // hop_length
575
+ else:
576
+ valid_mel_lens = None
577
+ x = self._conv2d_stem_one_chunk(x, valid_mel_lens)
578
+ # [B, dhs, F_out, T/8] → [B, T/8, dhs*freq_after] → [B, T/8, embed_dim]
579
+ batch, channels, frequency, time = x.shape
580
+ x = x.permute(0, 3, 1, 2).reshape(batch, time, channels * frequency)
581
+ inputs_embeds = self.conv_out(x)
582
+
583
+ return inputs_embeds
584
+
585
+ def _forward_conv1d_stem(self, input_features):
586
+ """Standard Conv1d stem: conv1 + conv2(stride=2), 2x downsample. Returns [B, T/2, embed_dim]."""
587
+ assert self.conv1 is not None and self.conv2 is not None
588
+ if self.use_causal:
589
+ x = nn.functional.pad(input_features, (2, 0))
590
+ x = nn.functional.gelu(self.conv1(x))
591
+ x = nn.functional.pad(x, (2, 0))
592
+ x = nn.functional.gelu(self.conv2(x))
593
+ else:
594
+ x = nn.functional.gelu(self.conv1(input_features))
595
+ x = nn.functional.gelu(self.conv2(x))
596
+ return x.permute(0, 2, 1)
597
+
598
+ def _forward_latent_stem(self, input_features):
599
+ """Latent input stem: conv1+GLU + conv2, stride=1, no downsample. Returns [B, T, embed_dim]."""
600
+ assert self.conv1 is not None and self.conv2 is not None
601
+ if self.use_causal:
602
+ x = nn.functional.pad(input_features, (2, 0))
603
+ x = nn.functional.glu(self.conv1(x), dim=1)
604
+ x = nn.functional.pad(x, (2, 0))
605
+ x = nn.functional.gelu(self.conv2(x))
606
+ else:
607
+ x = nn.functional.glu(self.conv1(input_features), dim=1)
608
+ x = nn.functional.gelu(self.conv2(x))
609
+ return x.permute(0, 2, 1)
610
+
611
+ def forward(
612
+ self,
613
+ input_features,
614
+ input_seq_lens=None,
615
+ audio_sample_lens=None,
616
+ output_attentions=None,
617
+ output_hidden_states=None,
618
+ return_dict=None,
619
+ ):
620
+ """Mel `[B, n_mels, T]` → encoder hidden states. Inference-only."""
621
+ output_attentions = (
622
+ output_attentions
623
+ if output_attentions is not None
624
+ else self.config.output_attentions
625
+ )
626
+ output_hidden_states = (
627
+ output_hidden_states
628
+ if output_hidden_states is not None
629
+ else self.config.output_hidden_states
630
+ )
631
+ return_dict = (
632
+ return_dict if return_dict is not None else self.config.use_return_dict
633
+ )
634
+ if self.use_conv2d_stem:
635
+ inputs_embeds = self._forward_conv2d_stem(
636
+ input_features, input_seq_lens, audio_sample_lens
637
+ )
638
+ elif self.use_latent_input:
639
+ inputs_embeds = self._forward_latent_stem(input_features)
640
+ else:
641
+ inputs_embeds = self._forward_conv1d_stem(input_features)
642
+ rotary_cos = None
643
+ rotary_sin = None
644
+ position_ids = torch.arange(
645
+ inputs_embeds.shape[1], device=inputs_embeds.device
646
+ )[None, :]
647
+ if self.use_rope:
648
+ rotary_cos, rotary_sin = self.rotary_embedding.get_cos_sin(
649
+ position_ids, inputs_embeds.dtype, inputs_embeds.device
650
+ )
651
+ hidden_states = inputs_embeds
652
+ else:
653
+ assert self.embed_positions is not None
654
+ embed_pos = self.embed_positions.weight[: inputs_embeds.shape[1]]
655
+ hidden_states = inputs_embeds + embed_pos
656
+
657
+ encoder_states: tuple[torch.Tensor, ...] | None = (
658
+ () if output_hidden_states else None
659
+ )
660
+ all_attentions: tuple[torch.Tensor, ...] | None = (
661
+ () if output_attentions else None
662
+ )
663
+
664
+ if input_seq_lens is not None:
665
+ # Build varlen metadata and pack valid tokens without per-sample cat loops.
666
+ # Read max on whatever device the caller provided: when it is a CPU
667
+ # tensor this avoids a device->host sync (one per forward).
668
+ max_seqlen_q = int(input_seq_lens.max().item())
669
+ input_seq_lens = input_seq_lens.to(
670
+ device=hidden_states.device, dtype=torch.long
671
+ )
672
+ B, S, D = hidden_states.shape
673
+ max_seqlen_kv = max_seqlen_q
674
+ cu_seqlens_q = torch.nn.functional.pad(
675
+ input_seq_lens.cumsum(0, dtype=torch.int32), (1, 0)
676
+ )
677
+ cu_seqlens_kv = cu_seqlens_q
678
+
679
+ token_positions = torch.arange(S, device=hidden_states.device)[None, :]
680
+ valid_token_mask = token_positions < input_seq_lens[:, None]
681
+ hidden_states = hidden_states[valid_token_mask]
682
+ if rotary_cos is not None and rotary_sin is not None:
683
+ packed_positions = token_positions.expand(B, S)[valid_token_mask]
684
+ rotary_cos = rotary_cos.squeeze(0).index_select(0, packed_positions)
685
+ rotary_sin = rotary_sin.squeeze(0).index_select(0, packed_positions)
686
+ else:
687
+ cu_seqlens_q = None
688
+ cu_seqlens_kv = None
689
+ max_seqlen_q = None
690
+ max_seqlen_kv = None
691
+
692
+ for encoder_layer in self.layers:
693
+ if output_hidden_states:
694
+ assert encoder_states is not None
695
+ encoder_states = encoder_states + (hidden_states,)
696
+ layer_outputs = encoder_layer(
697
+ hidden_states,
698
+ cu_seqlens_q,
699
+ cu_seqlens_kv,
700
+ max_seqlen_q,
701
+ max_seqlen_kv,
702
+ output_attentions=output_attentions,
703
+ rotary_cos=rotary_cos,
704
+ rotary_sin=rotary_sin,
705
+ )
706
+ hidden_states = layer_outputs[0]
707
+ if output_attentions:
708
+ assert all_attentions is not None
709
+ all_attentions = all_attentions + (layer_outputs[1],)
710
+
711
+ if input_seq_lens is not None:
712
+ # Recover packed varlen output directly on device.
713
+ recover_positions = torch.arange(max_seqlen_q, device=hidden_states.device)[
714
+ None, :
715
+ ]
716
+ recover_mask = recover_positions < input_seq_lens[:, None]
717
+ recovered = hidden_states.new_zeros(B, max_seqlen_q, D)
718
+ recovered[recover_mask] = hidden_states
719
+ hidden_states = recovered
720
+
721
+ hidden_states = self.layer_norm(hidden_states)
722
+ if output_hidden_states:
723
+ assert encoder_states is not None
724
+ encoder_states = encoder_states + (hidden_states,)
725
+
726
+ if not return_dict:
727
+ return tuple(
728
+ v
729
+ for v in [hidden_states, encoder_states, all_attentions]
730
+ if v is not None
731
+ )
732
+ return BaseModelOutput(
733
+ last_hidden_state=hidden_states,
734
+ hidden_states=encoder_states,
735
+ attentions=all_attentions,
736
+ )
vllm/models/dots3_note/nvidia/model.py
new
+684
@@ -0,0 +1,684 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NVIDIA implementation of the Dots3Note language model."""
4
+
5
+import copy
6
+from collections.abc import Iterable, Iterator
7
+from dataclasses import replace
8
+
9
+import torch
10
+from torch import nn
11
+
12
+import vllm.envs as envs
13
+from vllm.config import VllmConfig
14
+from vllm.distributed import (
15
+ get_pp_group,
16
+ get_tensor_model_parallel_rank,
17
+ get_tensor_model_parallel_world_size,
18
+ tensor_model_parallel_all_gather,
19
+ tensor_model_parallel_all_reduce,
20
+)
21
+from vllm.model_executor.layers.attention import MLAAttention
22
+from vllm.model_executor.layers.layernorm import RMSNorm
23
+from vllm.model_executor.layers.linear import (
24
+ ColumnParallelLinear,
25
+ DCPGroupColumnParallelLinear,
26
+ ReplicatedLinear,
27
+ RowParallelLinear,
28
+)
29
+from vllm.model_executor.layers.quantization import QuantizationConfig
30
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
31
+ GroupShape,
32
+ scaled_dequantize,
33
+)
34
+from vllm.model_executor.layers.rotary_embedding import get_rope
35
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
36
+from vllm.model_executor.models.deepseek_v2 import (
37
+ DeepSeekV2FusedQkvAProjLinear,
38
+ DeepseekV2MLAAttention,
39
+ DeepseekV2MLP,
40
+ DeepseekV2MoE,
41
+)
42
+from vllm.model_executor.models.utils import (
43
+ PPMissingLayer,
44
+ make_empty_intermediate_tensors_factory,
45
+ make_layers,
46
+ sequence_parallel_chunk,
47
+)
48
+from vllm.models.deepseek_v32.nvidia.model import (
49
+ DeepseekV32DecoderLayer,
50
+ DeepseekV32ForCausalLM,
51
+ DeepseekV32Model,
52
+)
53
+from vllm.platforms import current_platform
54
+from vllm.v1.kv_cache_interface import MLAAttentionSpec
55
+
56
+from .attention import (
57
+ Dots3NoteFlashAttnPrefillBackend,
58
+ Dots3NotePaddedSparseBackend,
59
+ Dots3NoteTritonMLABackend,
60
+)
61
+
62
+
63
+def _padded_mlp_size(
64
+ intermediate_size: int,
65
+ quant_config: QuantizationConfig | None,
66
+ is_sequence_parallel: bool = False,
67
+) -> int:
68
+ block_size = getattr(quant_config, "weight_block_size", None)
69
+ if block_size is None or is_sequence_parallel:
70
+ return intermediate_size
71
+ tp_size = get_tensor_model_parallel_world_size()
72
+ blocks = (intermediate_size + block_size[0] - 1) // block_size[0]
73
+ return ((blocks + tp_size - 1) // tp_size) * block_size[0] * tp_size
74
+
75
+
76
+class Dots3NoteMoE(DeepseekV2MoE):
77
+ """DeepSeek MoE with NOTE's block-padded shared expert."""
78
+
79
+ def __init__(
80
+ self,
81
+ config,
82
+ parallel_config,
83
+ quant_config: QuantizationConfig | None = None,
84
+ reduce_results: bool = True,
85
+ prefix: str = "",
86
+ apply_routed_scale_to_output: bool = False,
87
+ ) -> None:
88
+ num_shared_experts = config.n_shared_experts
89
+ routed_config = copy.copy(config)
90
+ object.__setattr__(routed_config, "n_shared_experts", None)
91
+ super().__init__(
92
+ config=routed_config,
93
+ parallel_config=parallel_config,
94
+ quant_config=quant_config,
95
+ reduce_results=False,
96
+ prefix=prefix,
97
+ apply_routed_scale_to_output=apply_routed_scale_to_output,
98
+ )
99
+ self.n_shared_experts = num_shared_experts
100
+ self.reduce_results = reduce_results
101
+ self.shared_experts = DeepseekV2MLP(
102
+ hidden_size=config.hidden_size,
103
+ intermediate_size=_padded_mlp_size(
104
+ config.moe_intermediate_size * num_shared_experts,
105
+ quant_config,
106
+ self.is_sequence_parallel,
107
+ ),
108
+ hidden_act=config.hidden_act,
109
+ quant_config=quant_config,
110
+ is_sequence_parallel=self.is_sequence_parallel,
111
+ reduce_results=False,
112
+ prefix=f"{prefix}.shared_experts",
113
+ )
114
+
115
+ def forward(
116
+ self,
117
+ hidden_states: torch.Tensor,
118
+ already_sequence_parallel: bool = False,
119
+ ) -> torch.Tensor:
120
+ num_tokens = hidden_states.shape[0]
121
+ gather_output = self.is_sequence_parallel and not already_sequence_parallel
122
+ if gather_output:
123
+ hidden_states = sequence_parallel_chunk(hidden_states)
124
+ assert self.shared_experts is not None
125
+ output = super().forward(
126
+ hidden_states, already_sequence_parallel=True
127
+ ) + self.shared_experts(hidden_states)
128
+ if gather_output:
129
+ return tensor_model_parallel_all_gather(output, 0)[:num_tokens]
130
+ if self.reduce_results:
131
+ return tensor_model_parallel_all_reduce(output)
132
+ return output
133
+
134
+
135
+def _forward_note_mla(
136
+ attention,
137
+ positions: torch.Tensor,
138
+ hidden_states: torch.Tensor,
139
+ *,
140
+ g_proj: nn.Module,
141
+ k_rope_only_layernorm: nn.Module,
142
+ attention_gate_type: str,
143
+ q_lora_scale: float,
144
+ kv_lora_scale: float,
145
+ llama_4_scaling: torch.Tensor | None = None,
146
+) -> torch.Tensor:
147
+ qkv_lora = attention.fused_qkv_a_proj(hidden_states)[0]
148
+ q_c, kv_lora = qkv_lora.split(
149
+ [
150
+ attention.q_lora_rank,
151
+ attention.kv_lora_rank + attention.qk_rope_head_dim,
152
+ ],
153
+ dim=-1,
154
+ )
155
+ q_c = attention.q_a_layernorm(q_c) * q_lora_scale
156
+ kv_c, k_pe = kv_lora.split(
157
+ [attention.kv_lora_rank, attention.qk_rope_head_dim], dim=-1
158
+ )
159
+ kv_c_normed = attention.kv_a_layernorm(kv_c) * kv_lora_scale
160
+ k_pe = k_rope_only_layernorm(k_pe).unsqueeze(1)
161
+
162
+ q = attention.q_b_proj(q_c)[0]
163
+ heads = attention.num_heads
164
+ if attention.dcp_q_replicate:
165
+ heads *= attention.q_b_proj.group_size
166
+ q = q.view(-1, heads, attention.qk_head_dim)
167
+ q[..., attention.qk_nope_head_dim :], k_pe = attention.rotary_emb(
168
+ positions, q[..., attention.qk_nope_head_dim :], k_pe
169
+ )
170
+
171
+ if attention.indexer and attention.is_sparse and not attention.skip_topk:
172
+ attention.indexer(hidden_states, q_c, positions, attention.indexer_rope_emb)
173
+ if llama_4_scaling is not None:
174
+ q *= llama_4_scaling
175
+
176
+ q_dcp_replicated = None
177
+ if attention.dcp_q_replicate:
178
+ q_dcp_replicated, q = q, attention.q_b_proj._local_view(q)
179
+ attn_out = attention.mla_attn(
180
+ q,
181
+ kv_c_normed,
182
+ k_pe,
183
+ output_shape=(
184
+ hidden_states.shape[0],
185
+ attention.num_heads * attention.v_head_dim,
186
+ ),
187
+ q_dcp_replicated=q_dcp_replicated,
188
+ )
189
+
190
+ gate = g_proj(hidden_states)[0]
191
+ if attention_gate_type == "headwise":
192
+ if gate.shape[-1] != attention.num_heads:
193
+ rank = get_tensor_model_parallel_rank()
194
+ gate = gate.narrow(-1, rank * attention.num_heads, attention.num_heads)
195
+ attn_out = attn_out.view(-1, attention.num_heads, attention.v_head_dim)
196
+ gate = torch.sigmoid(gate.float()).to(attn_out.dtype)
197
+ attn_out = (attn_out * gate.unsqueeze(-1)).flatten(-2)
198
+ else:
199
+ gate = torch.sigmoid(gate.float()).to(attn_out.dtype)
200
+ attn_out = attn_out * gate
201
+ return attention.o_proj(attn_out)[0]
202
+
203
+
204
+class Dots3NotePaddedMLAAttention(MLAAttention):
205
+ """MLA layer whose physical cache rows match NOTE's SWA rows."""
206
+
207
+ def __init__(self, *args, physical_head_size: int, **kwargs) -> None:
208
+ kwargs["attn_backend"] = Dots3NotePaddedSparseBackend
209
+ super().__init__(*args, **kwargs)
210
+ assert physical_head_size >= self.head_size
211
+ self.physical_head_size = physical_head_size
212
+
213
+ def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MLAAttentionSpec:
214
+ spec = super().get_kv_cache_spec(vllm_config)
215
+ assert isinstance(spec, MLAAttentionSpec)
216
+ return replace(spec, head_size=self.physical_head_size)
217
+
218
+
219
+class Dots3NoteFullAttention(DeepseekV2MLAAttention):
220
+ """NOTE DSA on the existing DeepSeek sparse MLA attention."""
221
+
222
+ def __init__(
223
+ self,
224
+ vllm_config: VllmConfig,
225
+ config,
226
+ prefix: str,
227
+ topk_indices_buffer: torch.Tensor | None = None,
228
+ ) -> None:
229
+ quant_config = vllm_config.quant_config
230
+ local_config = copy.copy(config)
231
+ object.__setattr__(
232
+ local_config,
233
+ "rope_parameters",
234
+ {
235
+ "rope_type": "default",
236
+ "rope_theta": config.rope_parameters["rope_theta"],
237
+ },
238
+ )
239
+ super().__init__(
240
+ vllm_config=vllm_config,
241
+ config=local_config,
242
+ hidden_size=config.hidden_size,
243
+ num_heads=config.num_attention_heads,
244
+ qk_nope_head_dim=config.qk_nope_head_dim,
245
+ qk_rope_head_dim=config.qk_rope_head_dim,
246
+ v_head_dim=config.v_head_dim,
247
+ q_lora_rank=config.q_lora_rank,
248
+ kv_lora_rank=config.kv_lora_rank,
249
+ max_position_embeddings=config.max_position_embeddings,
250
+ cache_config=vllm_config.cache_config,
251
+ quant_config=quant_config,
252
+ prefix=prefix,
253
+ topk_indices_buffer=topk_indices_buffer,
254
+ reduce_results=False,
255
+ )
256
+ wrapper = self._modules.pop("mla_attn")
257
+ inner_attention = wrapper.mla_attn
258
+ vllm_config.compilation_config.static_forward_context.pop(
259
+ inner_attention.layer_name
260
+ )
261
+ prefill_backend_cls = (
262
+ type(inner_attention.prefill_backend)
263
+ if inner_attention.prefill_backend is not None
264
+ else None
265
+ )
266
+ wrapper.mla_attn = Dots3NotePaddedMLAAttention(
267
+ num_heads=inner_attention.num_heads,
268
+ scale=inner_attention.scale,
269
+ q_lora_rank=inner_attention.q_lora_rank,
270
+ kv_lora_rank=inner_attention.kv_lora_rank,
271
+ qk_nope_head_dim=inner_attention.qk_nope_head_dim,
272
+ qk_rope_head_dim=inner_attention.qk_rope_head_dim,
273
+ v_head_dim=inner_attention.v_head_dim,
274
+ kv_b_proj=inner_attention.kv_b_proj,
275
+ dcp_q_replicate=inner_attention.dcp_q_replicate,
276
+ cache_config=vllm_config.cache_config,
277
+ quant_config=quant_config,
278
+ prefix=inner_attention.layer_name,
279
+ use_sparse=True,
280
+ indexer=wrapper.indexer,
281
+ topk_indices_buffer=topk_indices_buffer,
282
+ prefill_backend_cls=prefill_backend_cls,
283
+ physical_head_size=(config.swa_kv_lora_rank + config.swa_qk_rope_head_dim),
284
+ )
285
+ gate_type = config.attention_gate_type
286
+ gate_cls = ReplicatedLinear if gate_type == "headwise" else ColumnParallelLinear
287
+ gate_size = (
288
+ config.num_attention_heads
289
+ if gate_type == "headwise"
290
+ else config.num_attention_heads * config.v_head_dim
291
+ )
292
+ self.g_proj = gate_cls(
293
+ config.hidden_size,
294
+ gate_size,
295
+ bias=False,
296
+ quant_config=quant_config,
297
+ prefix=f"{prefix}.g_proj",
298
+ )
299
+ self.k_rope_only_layernorm = RMSNorm(
300
+ config.qk_rope_head_dim, eps=config.rms_norm_eps
301
+ )
302
+ self.attention_gate_type = gate_type
303
+ self.q_lora_scale = 1.0
304
+ self.kv_lora_scale = 1.0
305
+ if config.apply_mla_qkv_lora_rescale:
306
+ self.q_lora_scale = (config.hidden_size / config.q_lora_rank) ** 0.5
307
+ self.kv_lora_scale = (config.hidden_size / config.kv_lora_rank) ** 0.5
308
+ self.mla_attn = wrapper
309
+
310
+ def forward(
311
+ self,
312
+ positions: torch.Tensor,
313
+ hidden_states: torch.Tensor,
314
+ llama_4_scaling: torch.Tensor | None = None,
315
+ ) -> torch.Tensor:
316
+ return _forward_note_mla(
317
+ self.mla_attn,
318
+ positions,
319
+ hidden_states,
320
+ g_proj=self.g_proj,
321
+ k_rope_only_layernorm=self.k_rope_only_layernorm,
322
+ attention_gate_type=self.attention_gate_type,
323
+ q_lora_scale=self.q_lora_scale,
324
+ kv_lora_scale=self.kv_lora_scale,
325
+ llama_4_scaling=llama_4_scaling,
326
+ )
327
+
328
+
329
+class Dots3NoteSlidingAttention(nn.Module):
330
+ """NOTE SWA constructed directly as dense sliding-window MLA."""
331
+
332
+ def __init__(
333
+ self,
334
+ vllm_config: VllmConfig,
335
+ config,
336
+ prefix: str,
337
+ topk_indices_buffer: torch.Tensor | None = None,
338
+ ) -> None:
339
+ super().__init__()
340
+ del topk_indices_buffer
341
+ num_heads = config.swa_num_attention_heads
342
+ qk_nope_head_dim = config.swa_qk_nope_head_dim
343
+ qk_rope_head_dim = config.swa_qk_rope_head_dim
344
+ v_head_dim = config.swa_v_head_dim
345
+ q_lora_rank = config.swa_q_lora_rank
346
+ kv_lora_rank = config.swa_kv_lora_rank
347
+ qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
348
+ tp_size = get_tensor_model_parallel_world_size()
349
+ assert num_heads % tp_size == 0
350
+ quant_config = vllm_config.quant_config
351
+
352
+ fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
353
+ config.hidden_size,
354
+ [q_lora_rank, kv_lora_rank + qk_rope_head_dim],
355
+ quant_config=quant_config,
356
+ prefix=f"{prefix}.fused_qkv_a_proj",
357
+ )
358
+ qrep_enabled = (
359
+ envs.VLLM_DCP_Q_REPLICATE
360
+ and vllm_config.parallel_config.decode_context_parallel_size > 1
361
+ and vllm_config.parallel_config.prefill_context_parallel_size <= 1
362
+ )
363
+ q_proj_cls = (
364
+ DCPGroupColumnParallelLinear if qrep_enabled else ColumnParallelLinear
365
+ )
366
+ q_a_layernorm = RMSNorm(q_lora_rank, eps=config.rms_norm_eps)
367
+ q_b_proj = q_proj_cls(
368
+ q_lora_rank,
369
+ num_heads * qk_head_dim,
370
+ bias=False,
371
+ quant_config=quant_config,
372
+ prefix=f"{prefix}.q_b_proj",
373
+ )
374
+ kv_a_layernorm = RMSNorm(kv_lora_rank, eps=config.rms_norm_eps)
375
+ kv_b_proj = ColumnParallelLinear(
376
+ kv_lora_rank,
377
+ num_heads * (qk_nope_head_dim + v_head_dim),
378
+ bias=False,
379
+ quant_config=quant_config,
380
+ prefix=f"{prefix}.kv_b_proj",
381
+ )
382
+ o_proj = RowParallelLinear(
383
+ num_heads * v_head_dim,
384
+ config.hidden_size,
385
+ bias=False,
386
+ reduce_results=False,
387
+ quant_config=quant_config,
388
+ prefix=f"{prefix}.o_proj",
389
+ )
390
+ gate_type = config.swa_attention_gate_type
391
+ gate_cls = ReplicatedLinear if gate_type == "headwise" else ColumnParallelLinear
392
+ gate_size = num_heads if gate_type == "headwise" else num_heads * v_head_dim
393
+ g_proj = gate_cls(
394
+ config.hidden_size,
395
+ gate_size,
396
+ bias=False,
397
+ quant_config=quant_config,
398
+ prefix=f"{prefix}.g_proj",
399
+ )
400
+ k_rope_only_layernorm = RMSNorm(qk_rope_head_dim, eps=config.rms_norm_eps)
401
+ rotary_emb = get_rope(
402
+ qk_rope_head_dim,
403
+ max_position=config.max_position_embeddings,
404
+ rope_parameters={
405
+ "rope_type": "default",
406
+ "rope_theta": config.swa_rope_theta,
407
+ },
408
+ is_neox_style=False,
409
+ )
410
+ cache_config = vllm_config.cache_config
411
+ if cache_config.cache_dtype == "fp8_ds_mla":
412
+ cache_config = copy.copy(cache_config)
413
+ cache_config.cache_dtype = "fp8"
414
+ apply_rescale = config.apply_mla_qkv_lora_rescale
415
+ self.hidden_size = config.hidden_size
416
+ self.num_heads = num_heads // tp_size
417
+ self.qk_nope_head_dim = qk_nope_head_dim
418
+ self.qk_rope_head_dim = qk_rope_head_dim
419
+ self.qk_head_dim = qk_head_dim
420
+ self.v_head_dim = v_head_dim
421
+ self.q_lora_rank = q_lora_rank
422
+ self.kv_lora_rank = kv_lora_rank
423
+ self.fused_qkv_a_proj = fused_qkv_a_proj
424
+ self.q_a_layernorm = q_a_layernorm
425
+ self.q_b_proj = q_b_proj
426
+ self.kv_a_layernorm = kv_a_layernorm
427
+ self.kv_b_proj = kv_b_proj
428
+ self.rotary_emb = rotary_emb
429
+ self.o_proj = o_proj
430
+ self.g_proj = g_proj
431
+ self.k_rope_only_layernorm = k_rope_only_layernorm
432
+ self.indexer = None
433
+ self.indexer_rope_emb = None
434
+ self.is_sparse = False
435
+ self.skip_topk = False
436
+ self.dcp_q_replicate = getattr(q_b_proj, "qrep_active", False)
437
+ self.attention_gate_type = gate_type
438
+ self.q_lora_scale = (
439
+ (config.hidden_size / q_lora_rank) ** 0.5 if apply_rescale else 1.0
440
+ )
441
+ self.kv_lora_scale = (
442
+ (config.hidden_size / kv_lora_rank) ** 0.5 if apply_rescale else 1.0
443
+ )
444
+ self.mla_attn = MLAAttention(
445
+ num_heads=self.num_heads,
446
+ scale=qk_head_dim**-0.5,
447
+ qk_nope_head_dim=qk_nope_head_dim,
448
+ qk_rope_head_dim=qk_rope_head_dim,
449
+ v_head_dim=v_head_dim,
450
+ q_lora_rank=q_lora_rank,
451
+ kv_lora_rank=kv_lora_rank,
452
+ kv_b_proj=kv_b_proj,
453
+ dcp_q_replicate=self.dcp_q_replicate,
454
+ cache_config=cache_config,
455
+ quant_config=quant_config,
456
+ prefix=f"{prefix}.attn",
457
+ sliding_window=config.sliding_window_size,
458
+ attn_backend=Dots3NoteTritonMLABackend,
459
+ prefill_backend_cls=Dots3NoteFlashAttnPrefillBackend,
460
+ )
461
+
462
+ def forward(
463
+ self,
464
+ positions: torch.Tensor,
465
+ hidden_states: torch.Tensor,
466
+ llama_4_scaling: torch.Tensor | None = None,
467
+ ) -> torch.Tensor:
468
+ return _forward_note_mla(
469
+ self,
470
+ positions,
471
+ hidden_states,
472
+ g_proj=self.g_proj,
473
+ k_rope_only_layernorm=self.k_rope_only_layernorm,
474
+ attention_gate_type=self.attention_gate_type,
475
+ q_lora_scale=self.q_lora_scale,
476
+ kv_lora_scale=self.kv_lora_scale,
477
+ llama_4_scaling=llama_4_scaling,
478
+ )
479
+
480
+
481
+class Dots3NoteDecoderLayer(DeepseekV32DecoderLayer):
482
+ """DeepSeek-V3.2 decoder orchestration with NOTE-local modules."""
483
+
484
+ def __init__(
485
+ self,
486
+ vllm_config: VllmConfig,
487
+ prefix: str,
488
+ config=None,
489
+ topk_indices_buffer: torch.Tensor | None = None,
490
+ ) -> None:
491
+ nn.Module.__init__(self)
492
+ if config is None:
493
+ config = vllm_config.model_config.hf_config
494
+ quant_config = vllm_config.quant_config
495
+ parallel_config = vllm_config.parallel_config
496
+ self.hidden_size = config.hidden_size
497
+ layer_idx = int(prefix.split(sep=".")[-1])
498
+ self.layer_idx = layer_idx
499
+ self.use_mha = False
500
+ attention_cls = (
501
+ Dots3NoteSlidingAttention
502
+ if config.layer_types[layer_idx] == "sliding_attention"
503
+ else Dots3NoteFullAttention
504
+ )
505
+ self.self_attn = attention_cls(
506
+ vllm_config=vllm_config,
507
+ config=config,
508
+ prefix=f"{prefix}.self_attn",
509
+ topk_indices_buffer=topk_indices_buffer,
510
+ )
511
+
512
+ moe_layer_freq = getattr(config, "moe_layer_freq", 1)
513
+ is_moe = (
514
+ layer_idx < config.num_hidden_layers
515
+ and config.n_routed_experts is not None
516
+ and layer_idx >= config.first_k_dense_replace
517
+ and layer_idx % moe_layer_freq == 0
518
+ )
519
+ if is_moe:
520
+ self.mlp = Dots3NoteMoE(
521
+ config=config,
522
+ parallel_config=parallel_config,
523
+ quant_config=quant_config,
524
+ reduce_results=False,
525
+ prefix=f"{prefix}.mlp",
526
+ apply_routed_scale_to_output=False,
527
+ )
528
+ else:
529
+ self.mlp = DeepseekV2MLP(
530
+ hidden_size=config.hidden_size,
531
+ intermediate_size=_padded_mlp_size(
532
+ config.intermediate_size, quant_config
533
+ ),
534
+ hidden_act=config.hidden_act,
535
+ quant_config=quant_config,
536
+ prefix=f"{prefix}.mlp",
537
+ reduce_results=False,
538
+ )
539
+ self.use_sequence_parallel_moe = (
540
+ parallel_config.use_sequence_parallel_moe
541
+ and parallel_config.pipeline_parallel_size == 1
542
+ and isinstance(self.mlp, DeepseekV2MoE)
543
+ )
544
+ self.tp_size = parallel_config.tensor_parallel_size
545
+ self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
546
+ self.post_attention_layernorm = RMSNorm(
547
+ config.hidden_size, eps=config.rms_norm_eps
548
+ )
549
+ self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
550
+
551
+
552
+class Dots3NoteModel(DeepseekV32Model):
553
+ """DeepSeek-V3.2 runtime shell with NOTE-local decoder layers."""
554
+
555
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
556
+ nn.Module.__init__(self)
557
+ config = vllm_config.model_config.hf_config
558
+ quant_config = vllm_config.quant_config
559
+ self.config = config
560
+ self.device = current_platform.device_type
561
+ self.vocab_size = config.vocab_size
562
+ self.is_v32 = True
563
+ self._weight_block_size = getattr(quant_config, "weight_block_size", None)
564
+ topk_indices_buffer = torch.empty(
565
+ vllm_config.scheduler_config.max_num_batched_tokens,
566
+ config.index_topk,
567
+ dtype=torch.int32,
568
+ device=self.device,
569
+ )
570
+ if get_pp_group().is_first_rank:
571
+ self.embed_tokens = VocabParallelEmbedding(
572
+ config.vocab_size,
573
+ config.hidden_size,
574
+ quant_config=quant_config,
575
+ prefix=f"{prefix}.embed_tokens",
576
+ )
577
+ else:
578
+ self.embed_tokens = PPMissingLayer()
579
+ self.start_layer, self.end_layer, self.layers = make_layers(
580
+ config.num_hidden_layers,
581
+ lambda prefix: Dots3NoteDecoderLayer(
582
+ vllm_config=vllm_config,
583
+ prefix=prefix,
584
+ topk_indices_buffer=topk_indices_buffer,
585
+ ),
586
+ prefix=f"{prefix}.layers",
587
+ )
588
+ if get_pp_group().is_last_rank:
589
+ self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
590
+ else:
591
+ self.norm = PPMissingLayer()
592
+ self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory(
593
+ ["hidden_states", "residual"], config.hidden_size
594
+ )
595
+ self.aux_hidden_state_layers = tuple[int, ...]()
596
+ self.num_redundant_experts = (
597
+ vllm_config.parallel_config.eplb_config.num_redundant_experts
598
+ )
599
+
600
+ def _pad_dense_mlp_weight(
601
+ self, name: str, loaded_weight: torch.Tensor
602
+ ) -> torch.Tensor:
603
+ block_size = self._weight_block_size
604
+ if block_size is None or ".mlp.experts." in name:
605
+ return loaded_weight
606
+ if not any(
607
+ proj_name in name
608
+ for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
609
+ ):
610
+ return loaded_weight
611
+ dim = 1 if ".down_proj." in name else 0
612
+ block_step = 1 if name.endswith("weight_scale_inv") else block_size[0]
613
+ multiple = get_tensor_model_parallel_world_size() * block_step
614
+ pad = (-loaded_weight.shape[dim]) % multiple
615
+ if pad == 0:
616
+ return loaded_weight
617
+ pad_shape = list(loaded_weight.shape)
618
+ pad_shape[dim] = pad
619
+ return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
620
+
621
+ def _adapt_weights(
622
+ self, weights: Iterable[tuple[str, torch.Tensor]]
623
+ ) -> Iterator[tuple[str, torch.Tensor]]:
624
+ pending_indexer_fp8: dict[tuple[str, str], dict[str, torch.Tensor]] = {}
625
+ for name, weight in weights:
626
+ if name.startswith("mtp."):
627
+ continue
628
+ if ".indexer." in name:
629
+ parts = name.split(".")
630
+ if "layers" in parts:
631
+ layer_idx = int(parts[parts.index("layers") + 1])
632
+ if (
633
+ layer_idx < self.config.num_hidden_layers
634
+ and self.config.layer_types[layer_idx] == "sliding_attention"
635
+ ):
636
+ continue
637
+
638
+ projection = next(
639
+ (
640
+ candidate
641
+ for candidate in ("wk", "weights_proj")
642
+ if f".indexer.{candidate}." in name
643
+ ),
644
+ None,
645
+ )
646
+ is_fp8_weight = (
647
+ name.endswith(".weight") and weight.dtype == torch.float8_e4m3fn
648
+ )
649
+ is_scale = "weight_scale" in name
650
+ if projection is not None and (is_fp8_weight or is_scale):
651
+ layer_prefix = name.rsplit(f".{projection}.", 1)[0]
652
+ key = (layer_prefix, projection)
653
+ entry = pending_indexer_fp8.setdefault(key, {})
654
+ entry["weight" if is_fp8_weight else "scale"] = weight
655
+ if "weight" not in entry or "scale" not in entry:
656
+ continue
657
+ weight_fp8 = entry["weight"]
658
+ scale_inv = entry["scale"]
659
+ del pending_indexer_fp8[key]
660
+ group_shape = GroupShape(
661
+ weight_fp8.shape[0] // scale_inv.shape[0],
662
+ weight_fp8.shape[1] // scale_inv.shape[1],
663
+ )
664
+ dequantized = scaled_dequantize(
665
+ weight_fp8,
666
+ scale_inv,
667
+ group_shape=group_shape,
668
+ out_dtype=torch.bfloat16,
669
+ )
670
+ yield f"{layer_prefix}.{projection}.weight", dequantized
671
+ continue
672
+ yield name, self._pad_dense_mlp_weight(name, weight)
673
+ if pending_indexer_fp8:
674
+ missing = ", ".join(
675
+ f"{prefix}.{proj}" for prefix, proj in pending_indexer_fp8
676
+ )
677
+ raise ValueError(f"Incomplete FP8 indexer weights: {missing}")
678
+
679
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
680
+ return super().load_weights(self._adapt_weights(weights))
681
+
682
+
683
+class Dots3NoteLanguageModelForCausalLM(DeepseekV32ForCausalLM):
684
+ model_cls = Dots3NoteModel
vllm/models/dots3_note/nvidia/mtp.py
new
+189
@@ -0,0 +1,189 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NVIDIA multi-token predictor for Dots3Note."""
4
+
5
+from collections.abc import Iterable, Iterator
6
+
7
+import torch
8
+import torch.nn as nn
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.distributed import (
12
+ get_tensor_model_parallel_world_size,
13
+ tensor_model_parallel_all_gather,
14
+ tensor_model_parallel_all_reduce,
15
+)
16
+from vllm.model_executor.layers.layernorm import RMSNorm
17
+from vllm.model_executor.layers.linear import ReplicatedLinear
18
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
+from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
20
+from vllm.model_executor.models.deepseek_mtp import SharedHead
21
+from vllm.model_executor.models.utils import maybe_prefix
22
+from vllm.models.deepseek_v32.common.kernels import fused_eh_norm
23
+from vllm.models.deepseek_v32.nvidia.mtp import (
24
+ DeepseekV32MTP,
25
+ DeepseekV32MultiTokenPredictor,
26
+)
27
+
28
+from .model import Dots3NoteDecoderLayer
29
+
30
+
31
+class Dots3NoteMultiTokenPredictorLayer(nn.Module):
32
+ def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
33
+ super().__init__()
34
+ assert vllm_config.speculative_config is not None
35
+ config = vllm_config.speculative_config.draft_model_config.hf_config
36
+ quant_config = vllm_config.quant_config
37
+
38
+ self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
39
+ self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
40
+ self.eh_proj = ReplicatedLinear(
41
+ config.hidden_size * 2,
42
+ config.hidden_size,
43
+ bias=False,
44
+ quant_config=quant_config,
45
+ prefix=f"{prefix}.eh_proj",
46
+ )
47
+ self.shared_head = SharedHead(
48
+ config=config, prefix=prefix, quant_config=quant_config
49
+ )
50
+ self.mtp_block = Dots3NoteDecoderLayer(
51
+ vllm_config=vllm_config,
52
+ config=config,
53
+ prefix=prefix,
54
+ )
55
+
56
+ def forward(
57
+ self,
58
+ input_ids: torch.Tensor,
59
+ positions: torch.Tensor,
60
+ previous_hidden_states: torch.Tensor,
61
+ inputs_embeds: torch.Tensor | None = None,
62
+ spec_step_index: int = 0,
63
+ ) -> torch.Tensor:
64
+ del input_ids, spec_step_index
65
+ assert inputs_embeds is not None
66
+ eh_input = fused_eh_norm(
67
+ positions,
68
+ inputs_embeds,
69
+ previous_hidden_states,
70
+ self.enorm.weight,
71
+ self.hnorm.weight,
72
+ self.enorm.variance_epsilon,
73
+ )
74
+ hidden_states = self.eh_proj(eh_input)[0]
75
+ hidden_states, residual = self.mtp_block(
76
+ positions=positions, hidden_states=hidden_states, residual=None
77
+ )
78
+ is_sequence_parallel = self.mtp_block.use_sequence_parallel_moe
79
+ if not is_sequence_parallel:
80
+ hidden_states = tensor_model_parallel_all_reduce(hidden_states)
81
+ hidden_states, _ = self.shared_head.norm(hidden_states, residual)
82
+ if is_sequence_parallel:
83
+ hidden_states = tensor_model_parallel_all_gather(hidden_states, 0)
84
+ hidden_states = hidden_states[: positions.shape[0]]
85
+ return hidden_states
86
+
87
+
88
+class Dots3NoteMultiTokenPredictor(DeepseekV32MultiTokenPredictor):
89
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
90
+ nn.Module.__init__(self)
91
+ assert vllm_config.speculative_config is not None
92
+ config = vllm_config.speculative_config.draft_model_config.hf_config
93
+ self.mtp_start_layer_idx = config.num_hidden_layers
94
+ self.num_mtp_layers = config.num_nextn_predict_layers
95
+ self.layers = nn.ModuleDict(
96
+ {
97
+ str(idx): Dots3NoteMultiTokenPredictorLayer(
98
+ vllm_config, f"{prefix}.layers.{idx}"
99
+ )
100
+ for idx in range(
101
+ self.mtp_start_layer_idx,
102
+ self.mtp_start_layer_idx + self.num_mtp_layers,
103
+ )
104
+ }
105
+ )
106
+ self.embed_tokens = VocabParallelEmbedding(
107
+ config.vocab_size,
108
+ config.hidden_size,
109
+ prefix=maybe_prefix(prefix, "embed_tokens"),
110
+ )
111
+ self.register_buffer(
112
+ "max_token_id",
113
+ torch.tensor(config.vocab_size - 1, dtype=torch.int64),
114
+ persistent=False,
115
+ )
116
+ self.logits_processor = LogitsProcessor(config.vocab_size)
117
+
118
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
119
+ input_ids = torch.minimum(input_ids, self.max_token_id)
120
+ return self.embed_tokens(input_ids)
121
+
122
+ def forward(
123
+ self,
124
+ input_ids: torch.Tensor,
125
+ positions: torch.Tensor,
126
+ previous_hidden_states: torch.Tensor,
127
+ inputs_embeds: torch.Tensor | None = None,
128
+ spec_step_idx: int = 0,
129
+ ) -> torch.Tensor:
130
+ if inputs_embeds is None:
131
+ inputs_embeds = self.embed_input_ids(input_ids)
132
+ return super().forward(
133
+ input_ids,
134
+ positions,
135
+ previous_hidden_states,
136
+ inputs_embeds,
137
+ spec_step_idx,
138
+ )
139
+
140
+
141
+class Dots3NoteMTP(DeepseekV32MTP):
142
+ has_own_embed_tokens = True
143
+ has_own_lm_head = False
144
+
145
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
146
+ nn.Module.__init__(self)
147
+ self.config = vllm_config.model_config.hf_config
148
+ self.quant_config = vllm_config.quant_config
149
+ self.model = Dots3NoteMultiTokenPredictor(
150
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
151
+ )
152
+ self.set_moe_parameters()
153
+
154
+ def _pad_dense_mlp_weight(
155
+ self, name: str, loaded_weight: torch.Tensor
156
+ ) -> torch.Tensor:
157
+ block_size = getattr(self.quant_config, "weight_block_size", None)
158
+ if block_size is None or ".mlp.experts." in name:
159
+ return loaded_weight
160
+ if not any(
161
+ proj_name in name
162
+ for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
163
+ ):
164
+ return loaded_weight
165
+ dim = 1 if ".down_proj." in name else 0
166
+ block_step = 1 if name.endswith("weight_scale_inv") else block_size[0]
167
+ multiple = get_tensor_model_parallel_world_size() * block_step
168
+ pad = (-loaded_weight.shape[dim]) % multiple
169
+ if pad == 0:
170
+ return loaded_weight
171
+ pad_shape = list(loaded_weight.shape)
172
+ pad_shape[dim] = pad
173
+ return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
174
+
175
+ def _adapt_weights(
176
+ self, weights: Iterable[tuple[str, torch.Tensor]]
177
+ ) -> Iterator[tuple[str, torch.Tensor]]:
178
+ mtp_layer = self.config.num_hidden_layers
179
+ for name, weight in weights:
180
+ if name.startswith("model.mtp.embed_tokens."):
181
+ name = name.replace(
182
+ "model.mtp.embed_tokens.",
183
+ f"model.layers.{mtp_layer}.embed_tokens.",
184
+ 1,
185
+ )
186
+ yield name, self._pad_dense_mlp_weight(name, weight)
187
+
188
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
189
+ return super().load_weights(self._adapt_weights(weights))
vllm/models/dots3_note/nvidia/multimodal.py
new
+304
@@ -0,0 +1,304 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""vLLM composition layer for Dots3Note image and audio encoders."""
4
+
5
+from collections.abc import Iterable
6
+
7
+import torch
8
+from torch import nn
9
+
10
+from vllm.config import VllmConfig
11
+from vllm.model_executor.models.interfaces import (
12
+ MultiModalEmbeddings,
13
+ SupportsMultiModal,
14
+ SupportsPP,
15
+)
16
+from vllm.model_executor.models.module_mapping import MultiModelKeys
17
+from vllm.model_executor.models.utils import (
18
+ AutoWeightsLoader,
19
+ IntermediateTensors,
20
+ WeightsMapper,
21
+ maybe_prefix,
22
+)
23
+from vllm.multimodal import MULTIMODAL_REGISTRY
24
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
25
+
26
+from ..common.processor import (
27
+ AUDIO_END,
28
+ AUDIO_PAD,
29
+ AUDIO_START,
30
+ IMAGE_END,
31
+ IMAGE_PAD,
32
+ IMAGE_START,
33
+ VIDEO_PLACEHOLDER,
34
+ Dots3NoteDummyInputsBuilder,
35
+ Dots3NoteMultiModalProcessor,
36
+ Dots3NoteProcessingInfo,
37
+ load_note_config_section,
38
+)
39
+from .audio import Dots3NoteAudioConfig, Dots3NoteAudioModel
40
+from .model import Dots3NoteLanguageModelForCausalLM
41
+from .vision import DotsMoEVitConfig, DotsMoEVitModel
42
+
43
+
44
+@MULTIMODAL_REGISTRY.register_processor(
45
+ Dots3NoteMultiModalProcessor,
46
+ info=Dots3NoteProcessingInfo,
47
+ dummy_inputs=Dots3NoteDummyInputsBuilder,
48
+)
49
+class Dots3NoteForCausalLM(nn.Module, SupportsMultiModal, SupportsPP):
50
+ """Dots3Note model with optional image and audio towers."""
51
+
52
+ supports_encoder_tp_data = True
53
+
54
+ hf_to_vllm_mapper = WeightsMapper(
55
+ orig_to_new_prefix={
56
+ "model.": "language_model.model.",
57
+ "lm_head.": "language_model.lm_head.",
58
+ "mtp.": "language_model.mtp.",
59
+ "vision_encoder.": "visual.",
60
+ "audio_encoder.": "audio_tower.",
61
+ }
62
+ )
63
+
64
+ @classmethod
65
+ def get_placeholder_str(cls, modality: str, i: int) -> str | None:
66
+ if modality.startswith("image"):
67
+ return f"{IMAGE_START}{IMAGE_PAD}{IMAGE_END}"
68
+ if modality.startswith("audio"):
69
+ return f"{AUDIO_START}{AUDIO_PAD}{AUDIO_END}"
70
+ if modality.startswith("video"):
71
+ return VIDEO_PLACEHOLDER
72
+ raise ValueError(f"Unsupported modality: {modality}")
73
+
74
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
75
+ super().__init__()
76
+ model_config = vllm_config.model_config
77
+ self.config = model_config.hf_config
78
+ self.quant_config = vllm_config.quant_config
79
+ self.multimodal_config = model_config.multimodal_config
80
+ assert self.multimodal_config is not None
81
+
82
+ added_tokens = get_hf_file_to_dict(
83
+ "added_tokens.json",
84
+ model_config.model,
85
+ model_config.revision,
86
+ )
87
+ if added_tokens is None or IMAGE_PAD not in added_tokens:
88
+ raise ValueError("NOTE tokenizer is missing the image padding token")
89
+ self.config.image_token_index = int(added_tokens[IMAGE_PAD])
90
+
91
+ vision_config_dict = load_note_config_section(
92
+ model_config.model,
93
+ model_config.revision,
94
+ "vision_config",
95
+ )
96
+ audio_config_dict = load_note_config_section(
97
+ model_config.model,
98
+ model_config.revision,
99
+ "audio_config",
100
+ )
101
+ video_enabled = (
102
+ vision_config_dict is not None
103
+ and self.multimodal_config.get_limit_per_prompt("video") > 0
104
+ )
105
+ image_enabled = vision_config_dict is not None and (
106
+ self.multimodal_config.get_limit_per_prompt("image") > 0 or video_enabled
107
+ )
108
+ audio_enabled = audio_config_dict is not None and (
109
+ self.multimodal_config.get_limit_per_prompt("audio") > 0 or video_enabled
110
+ )
111
+
112
+ with self._mark_tower_model(vllm_config, {"image", "audio", "video"}):
113
+ self.visual: DotsMoEVitModel | None
114
+ if image_enabled:
115
+ assert vision_config_dict is not None
116
+ self.visual = DotsMoEVitModel(DotsMoEVitConfig(**vision_config_dict))
117
+ else:
118
+ self.visual = None
119
+ self.audio_tower: Dots3NoteAudioModel | None
120
+ if audio_enabled:
121
+ assert audio_config_dict is not None
122
+ self.audio_tower = Dots3NoteAudioModel(
123
+ Dots3NoteAudioConfig(**audio_config_dict)
124
+ )
125
+ else:
126
+ self.audio_tower = None
127
+ # The native encoder service casts each complete tower before
128
+ # loading its checkpoint. This also converts explicitly-created
129
+ # floating buffers (for example RoPE tables and router state),
130
+ # which merely constructing under vLLM's default dtype does not.
131
+ if self.visual is not None:
132
+ self.visual.to(dtype=model_config.dtype)
133
+ if self.audio_tower is not None:
134
+ self.audio_tower.to(dtype=model_config.dtype)
135
+ with self._mark_language_model(vllm_config):
136
+ self.language_model = Dots3NoteLanguageModelForCausalLM(
137
+ vllm_config=vllm_config,
138
+ prefix=maybe_prefix(prefix, "language_model"),
139
+ )
140
+ self.make_empty_intermediate_tensors = (
141
+ self.language_model.make_empty_intermediate_tensors
142
+ )
143
+
144
+ def _process_image_input(
145
+ self,
146
+ pixel_values: torch.Tensor,
147
+ image_grid_thw: torch.Tensor,
148
+ ) -> tuple[torch.Tensor, ...]:
149
+ if self.visual is None:
150
+ return ()
151
+ image_embeds = self.visual(pixel_values, image_grid_thw)
152
+ merge_size = self.visual.spatial_merge_size
153
+ sizes = (image_grid_thw.prod(-1) // merge_size**2).tolist()
154
+ return image_embeds.split(sizes)
155
+
156
+ def _process_audio_input(
157
+ self,
158
+ audio_values: torch.Tensor,
159
+ audio_lengths: torch.Tensor,
160
+ ) -> tuple[torch.Tensor, ...]:
161
+ if self.audio_tower is None:
162
+ return ()
163
+ if audio_values.dtype != torch.int32:
164
+ raise TypeError(
165
+ "NOTE audio values must carry float32 waveform bits as int32, "
166
+ f"got {audio_values.dtype}"
167
+ )
168
+ waveforms = audio_values.contiguous().view(torch.float32)
169
+ audio_embeds, item_lengths = self.audio_tower(waveforms, audio_lengths)
170
+ return audio_embeds.split(item_lengths)
171
+
172
+ def _process_video_input(
173
+ self,
174
+ pixel_values: torch.Tensor,
175
+ image_grid_thw: torch.Tensor,
176
+ audio_values: torch.Tensor,
177
+ audio_lengths: torch.Tensor,
178
+ modalities: torch.Tensor,
179
+ frame_counts: torch.Tensor,
180
+ audio_counts: torch.Tensor,
181
+ emission_counts: torch.Tensor,
182
+ ) -> tuple[torch.Tensor, ...]:
183
+ image_embeds = self._process_image_input(pixel_values, image_grid_thw)
184
+ audio_embeds = (
185
+ self._process_audio_input(audio_values, audio_lengths)
186
+ if audio_lengths.numel()
187
+ else ()
188
+ )
189
+ orders = modalities.tolist()
190
+ frame_counts_list = frame_counts.tolist()
191
+ audio_counts_list = audio_counts.tolist()
192
+ emission_counts_list = emission_counts.tolist()
193
+ outputs: list[torch.Tensor] = []
194
+ image_idx = audio_idx = order_idx = 0
195
+ for num_frames, num_audios, num_emissions in zip(
196
+ frame_counts_list,
197
+ audio_counts_list,
198
+ emission_counts_list,
199
+ ):
200
+ video_order = orders[order_idx : order_idx + num_emissions]
201
+ video_parts: list[torch.Tensor] = []
202
+ video_image_start = image_idx
203
+ video_audio_start = audio_idx
204
+ for modality in video_order:
205
+ if modality == 0:
206
+ video_parts.append(image_embeds[image_idx])
207
+ image_idx += 1
208
+ elif modality == 1:
209
+ if audio_idx >= len(audio_embeds):
210
+ raise ValueError("NOTE video audio tower output is missing")
211
+ video_parts.append(audio_embeds[audio_idx])
212
+ audio_idx += 1
213
+ else:
214
+ raise ValueError(f"Unknown NOTE video modality id: {modality}")
215
+ if image_idx - video_image_start != num_frames:
216
+ raise ValueError("NOTE video frame order/count mismatch")
217
+ if audio_idx - video_audio_start != num_audios:
218
+ raise ValueError("NOTE video audio order/count mismatch")
219
+ outputs.append(torch.cat(video_parts))
220
+ order_idx += num_emissions
221
+ if image_idx != len(image_embeds) or audio_idx != len(audio_embeds):
222
+ raise ValueError("NOTE video encoder outputs were not fully consumed")
223
+ return tuple(outputs)
224
+
225
+ def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
226
+ multimodal_embeddings: list[torch.Tensor] = []
227
+ handled: set[str] = set()
228
+ for input_key in kwargs:
229
+ if input_key == "pixel_values" and "image" not in handled:
230
+ pixel_values = kwargs.get("pixel_values")
231
+ image_grid_thw = kwargs.get("image_grid_thw")
232
+ if isinstance(pixel_values, torch.Tensor) and isinstance(
233
+ image_grid_thw, torch.Tensor
234
+ ):
235
+ multimodal_embeddings.extend(
236
+ self._process_image_input(pixel_values, image_grid_thw)
237
+ )
238
+ handled.add("image")
239
+ elif input_key == "audio_values" and "audio" not in handled:
240
+ audio_values = kwargs.get("audio_values")
241
+ audio_lengths = kwargs.get("audio_lengths")
242
+ if isinstance(audio_values, torch.Tensor) and isinstance(
243
+ audio_lengths, torch.Tensor
244
+ ):
245
+ multimodal_embeddings.extend(
246
+ self._process_audio_input(audio_values, audio_lengths)
247
+ )
248
+ handled.add("audio")
249
+ elif input_key == "video_pixel_values" and "video" not in handled:
250
+ video_inputs = (
251
+ kwargs.get("video_pixel_values"),
252
+ kwargs.get("video_image_grid_thw"),
253
+ kwargs.get("video_audio_values"),
254
+ kwargs.get("video_audio_lengths"),
255
+ kwargs.get("video_modalities"),
256
+ kwargs.get("video_frame_counts"),
257
+ kwargs.get("video_audio_counts"),
258
+ kwargs.get("video_emission_counts"),
259
+ )
260
+ if all(isinstance(value, torch.Tensor) for value in video_inputs):
261
+ multimodal_embeddings.extend(
262
+ self._process_video_input(*video_inputs) # type: ignore[arg-type]
263
+ )
264
+ handled.add("video")
265
+ return tuple(multimodal_embeddings)
266
+
267
+ def forward(
268
+ self,
269
+ input_ids: torch.Tensor | None,
270
+ positions: torch.Tensor,
271
+ intermediate_tensors: IntermediateTensors | None = None,
272
+ inputs_embeds: torch.Tensor | None = None,
273
+ **kwargs: object,
274
+ ) -> torch.Tensor | IntermediateTensors:
275
+ return self.language_model(
276
+ input_ids,
277
+ positions,
278
+ intermediate_tensors,
279
+ inputs_embeds,
280
+ )
281
+
282
+ def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
283
+ return self.language_model.compute_logits(hidden_states)
284
+
285
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
286
+ skip_prefixes = []
287
+ if self.visual is None:
288
+ skip_prefixes.append("visual.")
289
+ if self.audio_tower is None:
290
+ skip_prefixes.append("audio_tower.")
291
+ return AutoWeightsLoader(self, skip_prefixes=skip_prefixes).load_weights(
292
+ weights,
293
+ mapper=self.hf_to_vllm_mapper,
294
+ )
295
+
296
+ def process_weights_after_loading(self) -> None:
297
+ if self.visual is not None:
298
+ self.visual.process_weights_after_loading()
299
+
300
+ def get_mm_mapping(self) -> MultiModelKeys:
301
+ return MultiModelKeys.from_string_field(
302
+ language_model="language_model",
303
+ tower_model=["visual", "audio_tower"],
304
+ )
vllm/models/dots3_note/nvidia/vision.py
new
+677
@@ -0,0 +1,677 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+import math
5
+from typing import Any
6
+
7
+import torch
8
+import torch.nn.functional as F
9
+from torch import nn
10
+from torch.nn import LayerNorm
11
+from transformers.configuration_utils import PretrainedConfig
12
+from transformers.modeling_utils import PreTrainedModel
13
+
14
+from vllm.third_party.deep_gemm import per_block_cast_to_fp8
15
+
16
+from .vision_attention import (
17
+ VisionRotaryEmbedding,
18
+ apply_vision_attention_residual,
19
+ attn_uses_seqlens,
20
+ build_vision_attention,
21
+ prepare_seqlens_for_attention,
22
+ resolve_attn_implementation,
23
+)
24
+from .vision_moe import note_vision_fused_moe_fp8
25
+
26
+
27
+class DotsMoEVitConfig(PretrainedConfig):
28
+ model_type: str = "dots_moe_vit"
29
+
30
+ def __init__(
31
+ self,
32
+ embed_dim: int = 1536,
33
+ hidden_size: int = 2048,
34
+ intermediate_size: int = 4224,
35
+ moe_intermediate_size: int = 2112,
36
+ num_hidden_layers: int = 42,
37
+ num_attention_heads: int = 24,
38
+ num_channels: int = 3,
39
+ patch_size: int = 14,
40
+ spatial_merge_size: int = 2,
41
+ temporal_patch_size: int = 1,
42
+ rms_norm_eps: float = 1e-5,
43
+ use_bias: bool = False,
44
+ use_qk_norm: bool = True,
45
+ attn_implementation="flash_attention_3", # "eager","eager_v2","sdpa","flash_attention_2","flash_attention_3"
46
+ initializer_range=0.02,
47
+ is_causal=False,
48
+ post_norm=True,
49
+ gradient_checkpointing=False,
50
+ pyramid_num_routed: list[int] | None = None,
51
+ capacity_factor: float = 2.0,
52
+ router_scoring_func: str = "sigmoid",
53
+ router_scale: float = 1.0,
54
+ adapter_in_dim: int = 1536,
55
+ adapter_out_dim: int = 2048,
56
+ adapter_merge_size: int = 2,
57
+ # ``pixel_shuffle_mlp`` reshapes 2x2 spatial neighbours via NHWC
58
+ # pixel-shuffle, then runs layer norm and a two-layer MLP.
59
+ # ``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.
60
+ adapter_type: str = "pixel_shuffle_mlp",
61
+ # If True the preprocessor already emits patches in 2x2-grouped order (qwen ``merge_size=2`` flatten path), so RoPE positions must also be regrouped
62
+ # to match the in-block layout. Older checkpoints were trained with
63
+ # ``pre_pixel_shuffle=False``, so the default keeps row-major RoPE.
64
+ pre_pixel_shuffle: bool = False,
65
+ # Keep the native encoder's post-load compile behavior. Compiling in
66
+ # ``__init__`` would rename checkpoint keys through ``_orig_mod``.
67
+ enable_torch_compile: bool = True,
68
+ # If True, use FP8 MoE implementation
69
+ enable_fp8_moe: bool = True,
70
+ **kwargs: Any,
71
+ ):
72
+ super().__init__(**kwargs)
73
+ self.embed_dim = embed_dim
74
+ self.hidden_size = hidden_size
75
+ self.intermediate_size = intermediate_size
76
+ self.moe_intermediate_size = moe_intermediate_size
77
+ self.num_hidden_layers = num_hidden_layers
78
+ self.num_attention_heads = num_attention_heads
79
+ self.num_channels = num_channels
80
+ self.patch_size = patch_size
81
+ self.spatial_merge_size = spatial_merge_size
82
+ self.temporal_patch_size = temporal_patch_size
83
+ self.rms_norm_eps = rms_norm_eps
84
+ self.use_bias = use_bias
85
+ self.use_qk_norm = use_qk_norm
86
+ self.attn_implementation = attn_implementation
87
+ self.initializer_range = initializer_range
88
+ self.is_causal = is_causal
89
+ self.post_norm = post_norm
90
+ self.gradient_checkpointing = gradient_checkpointing
91
+ self.pyramid_num_routed = pyramid_num_routed or []
92
+ self.capacity_factor = capacity_factor
93
+ self.router_scoring_func = router_scoring_func
94
+ self.router_scale = router_scale
95
+ self.adapter_in_dim = adapter_in_dim
96
+ self.adapter_out_dim = adapter_out_dim
97
+ self.adapter_merge_size = adapter_merge_size
98
+ if adapter_type not in ("pixel_shuffle_mlp", "patch_merger"):
99
+ raise ValueError(
100
+ f"adapter_type must be 'pixel_shuffle_mlp' or 'patch_merger', got {adapter_type!r}"
101
+ )
102
+ self.adapter_type = adapter_type
103
+ self.pre_pixel_shuffle = pre_pixel_shuffle
104
+ self.enable_torch_compile = enable_torch_compile
105
+ self.enable_fp8_moe = enable_fp8_moe
106
+
107
+
108
+class RMSNorm(nn.Module):
109
+ def __init__(self, dim: int, eps: float = 1e-6):
110
+ super().__init__()
111
+ self.weight = nn.Parameter(torch.ones(dim))
112
+ self.eps = eps
113
+
114
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
115
+ output = self._norm(x.float()).type_as(x)
116
+ return output * self.weight
117
+
118
+ def extra_repr(self) -> str:
119
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
120
+
121
+ def _norm(self, x: torch.Tensor) -> torch.Tensor:
122
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
123
+
124
+
125
+# ---- FFN modules ----
126
+
127
+
128
+class DotsSwiGLUFFN(nn.Module):
129
+ def __init__(self, in_features, hidden_features, bias=False):
130
+ super().__init__()
131
+ self.fc1 = nn.Linear(in_features, hidden_features, bias=bias)
132
+ self.fc2 = nn.Linear(hidden_features, in_features, bias=bias)
133
+ self.fc3 = nn.Linear(in_features, hidden_features, bias=bias)
134
+
135
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
136
+ return self.fc2(F.silu(self.fc1(x)) * self.fc3(x))
137
+
138
+
139
+class MoESwiGLUFFN(nn.Module):
140
+ """MoE FFN with per-expert SwiGLU experts, sigmoid/softmax gating, top-k routing."""
141
+
142
+ def __init__(self, config: DotsMoEVitConfig, layer_number: int):
143
+ super().__init__()
144
+ self.config = config
145
+ self.layer_number = layer_number
146
+ self.hidden_size = config.embed_dim
147
+ self.num_routed = config.pyramid_num_routed[layer_number]
148
+ self.capacity_factor = config.capacity_factor
149
+ self.router_scoring_func = config.router_scoring_func
150
+ self.router_scale = config.router_scale
151
+
152
+ self.register_buffer(
153
+ "router_bias", torch.zeros(self.num_routed, dtype=torch.float32)
154
+ )
155
+
156
+ self.experts = nn.ModuleList(
157
+ [
158
+ DotsSwiGLUFFN(
159
+ self.hidden_size, config.moe_intermediate_size, bias=config.use_bias
160
+ )
161
+ for _ in range(self.num_routed)
162
+ ]
163
+ )
164
+
165
+ self.gate_weight = nn.Parameter(
166
+ torch.empty((self.num_routed, self.hidden_size), dtype=torch.float32)
167
+ )
168
+ nn.init.kaiming_uniform_(self.gate_weight, a=math.sqrt(5))
169
+
170
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
171
+ # Mirror cybertron's ``AIMv2MoEGEMMSwiGLUFFN.forward``: keep ``gating_prob`` /
172
+ # router-bias add in fp32 and run topk in fp32. The legacy ``.type_as(x_flat)``
173
+ # path on bf16 made expert routing diverge whenever two routed-expert scores
174
+ # tied within bf16 precision, which is the dominant source of numerical drift
175
+ # observed against latest cybertron checkpoints.
176
+ epsilon = 1e-9
177
+ x_flat = x.contiguous()
178
+ num_tokens = x_flat.shape[0]
179
+
180
+ gate_logits = F.linear(x_flat.float(), self.gate_weight.float())
181
+
182
+ if self.router_scoring_func == "sigmoid":
183
+ gating_prob = torch.sigmoid(gate_logits)
184
+ else:
185
+ gating_prob = torch.softmax(gate_logits, dim=-1, dtype=torch.float32)
186
+
187
+ aggregated_output = torch.zeros_like(x_flat)
188
+ aggregated_gate = torch.zeros(num_tokens, dtype=x.dtype, device=x.device)
189
+
190
+ topk = min(int(self.capacity_factor), self.num_routed)
191
+
192
+ gating_with_bias = gating_prob + self.router_bias.to(torch.float32).unsqueeze(0)
193
+ _, topk_indices = torch.topk(gating_with_bias, k=topk, dim=-1, sorted=False)
194
+
195
+ routed_weights = gating_prob.gather(1, topk_indices)
196
+ if self.router_scoring_func == "sigmoid" and topk > 1:
197
+ routed_weights = routed_weights / (
198
+ routed_weights.sum(dim=-1, keepdim=True) + epsilon
199
+ )
200
+ routed_weights = (routed_weights * self.router_scale).to(x_flat.dtype)
201
+
202
+ for expert_idx in range(self.num_routed):
203
+ selected_mask = topk_indices == expert_idx
204
+ if selected_mask.sum() == 0:
205
+ continue
206
+ n_idx, top = torch.where(selected_mask)
207
+ # Fancy indexing can yield non-contiguous rows; cuBLAS bf16 GEMM may then fail
208
+ # with ``CUBLAS_STATUS_INVALID_VALUE`` inside ``F.linear``.
209
+ x_selected = x_flat[n_idx].contiguous()
210
+ expert_output = self.experts[expert_idx](x_selected)
211
+ contrib = expert_output * routed_weights[n_idx, top].unsqueeze(-1)
212
+ aggregated_output[n_idx] = aggregated_output[n_idx] + contrib
213
+ aggregated_gate[n_idx] = aggregated_gate[n_idx] + routed_weights[n_idx, top]
214
+
215
+ aggregated_output = aggregated_output / (
216
+ aggregated_gate.unsqueeze(-1) + epsilon
217
+ )
218
+ return aggregated_output
219
+
220
+
221
+def _ceil_to_multiple(value: int, multiple: int) -> int:
222
+ return ((value + multiple - 1) // multiple) * multiple
223
+
224
+
225
+def _per_block_cast_to_fp8_padded(
226
+ weight: torch.Tensor,
227
+ block_size: int = 128,
228
+) -> tuple[torch.Tensor, torch.Tensor]:
229
+ rows, columns = weight.shape
230
+ padded = weight.new_zeros(
231
+ _ceil_to_multiple(rows, block_size),
232
+ _ceil_to_multiple(columns, block_size),
233
+ )
234
+ padded[:rows, :columns] = weight
235
+ return per_block_cast_to_fp8(
236
+ padded.contiguous(),
237
+ use_ue8m0=False,
238
+ gran_k=block_size,
239
+ )
240
+
241
+
242
+class MoESwiGLUFFNFP8(MoESwiGLUFFN):
243
+ """NOTE vision MoE using the checkpoint's local block-FP8 semantics."""
244
+
245
+ @torch.no_grad()
246
+ def process_weights_after_loading(self) -> None:
247
+ if hasattr(self, "_fused_w13_fp8"):
248
+ return
249
+
250
+ w13_weights = []
251
+ w13_scales = []
252
+ w2_weights = []
253
+ w2_scales = []
254
+ for expert in self.experts:
255
+ w1, s1 = _per_block_cast_to_fp8_padded(expert.fc1.weight)
256
+ w3, s3 = _per_block_cast_to_fp8_padded(expert.fc3.weight)
257
+ w2, s2 = _per_block_cast_to_fp8_padded(expert.fc2.weight)
258
+ w13_weights.append(torch.cat((w1, w3), dim=0))
259
+ w13_scales.append(torch.cat((s1, s3), dim=0))
260
+ w2_weights.append(w2)
261
+ w2_scales.append(s2)
262
+
263
+ self.register_buffer(
264
+ "_fused_w13_fp8",
265
+ torch.stack(w13_weights).contiguous(),
266
+ persistent=False,
267
+ )
268
+ self.register_buffer(
269
+ "_fused_w13_scale",
270
+ torch.stack(w13_scales).contiguous(),
271
+ persistent=False,
272
+ )
273
+ self.register_buffer(
274
+ "_fused_w2_fp8",
275
+ torch.stack(w2_weights).contiguous(),
276
+ persistent=False,
277
+ )
278
+ self.register_buffer(
279
+ "_fused_w2_scale",
280
+ torch.stack(w2_scales).contiguous(),
281
+ persistent=False,
282
+ )
283
+ del self.experts
284
+
285
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
286
+ if not hasattr(self, "_fused_w13_fp8"):
287
+ raise RuntimeError("NOTE vision FP8 weights were not initialized")
288
+
289
+ gate_logits = F.linear(x.float(), self.gate_weight.float())
290
+ if self.router_scoring_func == "sigmoid":
291
+ scores = torch.sigmoid(gate_logits)
292
+ else:
293
+ scores = torch.softmax(gate_logits, dim=-1, dtype=torch.float32)
294
+
295
+ topk = min(int(self.capacity_factor), self.num_routed)
296
+ biased_scores = scores + self.router_bias.float().unsqueeze(0)
297
+ topk_ids = torch.topk(biased_scores, k=topk, dim=-1, sorted=False).indices
298
+ topk_weights = scores.gather(1, topk_ids)
299
+ if self.router_scoring_func == "sigmoid" and topk > 1:
300
+ topk_weights = topk_weights / (
301
+ topk_weights.sum(dim=-1, keepdim=True) + 1e-9
302
+ )
303
+ topk_weights = topk_weights * self.router_scale
304
+
305
+ output = note_vision_fused_moe_fp8(
306
+ x.contiguous(),
307
+ self._fused_w13_fp8,
308
+ self._fused_w2_fp8,
309
+ topk_weights,
310
+ topk_ids.to(torch.int32),
311
+ self._fused_w13_scale,
312
+ self._fused_w2_scale,
313
+ )
314
+ denominator = topk_weights.sum(dim=-1, keepdim=True).clamp_min(1e-9)
315
+ return (output / denominator).type_as(x)
316
+
317
+
318
+# ---- PatchEmbed ----
319
+
320
+
321
+class DotsPatchEmbed(nn.Module):
322
+ def __init__(self, config: DotsMoEVitConfig):
323
+ super().__init__()
324
+ self.num_channels = config.num_channels
325
+ self.patch_size = config.patch_size
326
+ self.temporal_patch_size = config.temporal_patch_size
327
+ self.embed_dim = config.embed_dim
328
+ self.proj = nn.Conv2d(
329
+ config.num_channels,
330
+ config.embed_dim,
331
+ kernel_size=(config.patch_size, config.patch_size),
332
+ stride=(config.patch_size, config.patch_size),
333
+ )
334
+ self.norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
335
+
336
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
337
+ x = x.view(
338
+ -1,
339
+ self.num_channels,
340
+ self.temporal_patch_size,
341
+ self.patch_size,
342
+ self.patch_size,
343
+ )[:, :, 0]
344
+ x = self.proj(x).view(-1, self.embed_dim)
345
+ x = self.norm(x)
346
+ return x
347
+
348
+
349
+# ---- Block ----
350
+
351
+
352
+class MoEVisionBlock(nn.Module):
353
+ def __init__(self, config: DotsMoEVitConfig, layer_number: int):
354
+ super().__init__()
355
+ attn_impl = resolve_attn_implementation(
356
+ config.attn_implementation, eager_fallback="eager_v2"
357
+ )
358
+ self.attn = build_vision_attention(attn_impl, config, eager_fallback="eager_v2")
359
+ self._attn_uses_seqlens = attn_uses_seqlens(attn_impl)
360
+ self.norm_1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
361
+ self.norm_2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
362
+
363
+ is_moe = (
364
+ config.pyramid_num_routed
365
+ and layer_number < len(config.pyramid_num_routed)
366
+ and config.pyramid_num_routed[layer_number] > 0
367
+ )
368
+ if is_moe:
369
+ mlp_cls = MoESwiGLUFFNFP8 if config.enable_fp8_moe else MoESwiGLUFFN
370
+ self.mlp = mlp_cls(config, layer_number)
371
+ else:
372
+ self.mlp = DotsSwiGLUFFN(
373
+ config.embed_dim, config.intermediate_size, bias=config.use_bias
374
+ )
375
+
376
+ def forward(
377
+ self,
378
+ hidden_states,
379
+ cu_seqlens,
380
+ rotary_pos_emb,
381
+ max_seqlen: int,
382
+ seqlens: list[int] | None = None,
383
+ ) -> torch.Tensor:
384
+ hidden_states = apply_vision_attention_residual(
385
+ self.attn,
386
+ self.norm_1,
387
+ hidden_states,
388
+ cu_seqlens,
389
+ max_seqlen,
390
+ rotary_pos_emb,
391
+ seqlens=seqlens,
392
+ uses_seqlens=self._attn_uses_seqlens,
393
+ )
394
+ hidden_states = hidden_states + self.mlp(self.norm_2(hidden_states))
395
+ return hidden_states
396
+
397
+
398
+# ---- Adapter (pixel_shuffle + MLP) ----
399
+
400
+
401
+def _pixel_shuffle(x, scale_factor=0.5):
402
+ if x.size(1) % 2 == 1:
403
+ x = torch.cat([x[:, :1], x], dim=1)
404
+ if x.size(2) % 2 == 1:
405
+ x = torch.cat([x[:, :, :1], x], dim=2)
406
+ n, h, w, c = x.size()
407
+ x = x.reshape(n, h, int(w * scale_factor), int(c / scale_factor))
408
+ x = x.permute(0, 2, 1, 3).contiguous()
409
+ x = x.reshape(
410
+ n,
411
+ int(w * scale_factor),
412
+ int(h * scale_factor),
413
+ int(c / (scale_factor * scale_factor)),
414
+ )
415
+ x = x.permute(0, 2, 1, 3).contiguous()
416
+ return x
417
+
418
+
419
+class PixelShuffleAdapter(nn.Module):
420
+ """Legacy adapter: NHWC pixel-shuffle spatial merge + LayerNorm + 2-layer MLP.
421
+
422
+ Mirrors ``cybertron`` ``FCAdapter(pool_kind='pixel_shuffle', proj_kind='mlp2x_ln_gelu')``.
423
+ State-dict keys: ``proj.0`` (LayerNorm of in_dim*merge**2), ``proj.1`` / ``proj.3`` (Linear).
424
+ """
425
+
426
+ def __init__(self, config: DotsMoEVitConfig):
427
+ super().__init__()
428
+ in_dim = config.adapter_in_dim
429
+ out_dim = config.adapter_out_dim
430
+ merge_size = config.adapter_merge_size
431
+ merged_dim = in_dim * merge_size**2
432
+ self.proj = nn.Sequential(
433
+ LayerNorm(merged_dim),
434
+ nn.Linear(merged_dim, out_dim),
435
+ nn.GELU(),
436
+ nn.Linear(out_dim, out_dim),
437
+ )
438
+
439
+ def forward(
440
+ self,
441
+ patch_embed: torch.Tensor,
442
+ grid_thw: torch.Tensor,
443
+ ) -> torch.Tensor:
444
+ assert patch_embed.dim() == 2 and grid_thw is not None
445
+ image_features = []
446
+ token_index = 0
447
+ for i in range(grid_thw.shape[0]):
448
+ grid_t, grid_h, grid_w = grid_thw[i]
449
+ images_token_length = grid_t * grid_h * grid_w
450
+ _pe = patch_embed[token_index : token_index + images_token_length]
451
+ token_index += images_token_length
452
+ if grid_t == 1:
453
+ _pe = _pe.reshape(int(grid_h), int(grid_w), -1).unsqueeze(0)
454
+ else:
455
+ _pe = _pe.reshape(int(grid_t), int(grid_h), int(grid_w), -1)
456
+ _pe = _pixel_shuffle(_pe, scale_factor=0.5)
457
+ _pe = _pe.squeeze(0) if grid_t == 1 else _pe.reshape(-1, _pe.shape[-1])
458
+ image_features.append(_pe.reshape(-1, _pe.shape[-1]))
459
+ out = torch.cat(image_features, dim=0)
460
+ out = self.proj(out)
461
+ return out
462
+
463
+
464
+class PatchMergerAdapter(nn.Module):
465
+ """Cybertron ``PatchMerger`` (``pool_kind='patch_merger', proj_kind='identity'``).
466
+
467
+ Assumes the encoder output is already laid out in ``merge_size``x``merge_size`` groups
468
+ (qwen ``pre_pixel_shuffle`` preprocessor + RoPE grouped accordingly), so merging is a
469
+ simple ``view(-1, merge**2 * in_dim)`` of consecutive tokens. State-dict layout matches
470
+ cybertron's ``PatchMerger`` (``ln_q`` over the per-token dim, ``mlp.0`` / ``mlp.2`` Linear).
471
+ """
472
+
473
+ def __init__(self, config: DotsMoEVitConfig):
474
+ super().__init__()
475
+ in_dim = config.adapter_in_dim
476
+ out_dim = config.adapter_out_dim
477
+ merge_size = config.adapter_merge_size
478
+ merged_dim = in_dim * merge_size**2
479
+ self.merge_size = merge_size
480
+ self.merged_dim = merged_dim
481
+ self.ln_q = LayerNorm(in_dim, eps=1e-6)
482
+ self.mlp = nn.Sequential(
483
+ nn.Linear(merged_dim, merged_dim),
484
+ nn.GELU(),
485
+ nn.Linear(merged_dim, out_dim),
486
+ )
487
+
488
+ def forward(
489
+ self,
490
+ patch_embed: torch.Tensor,
491
+ grid_thw: torch.Tensor,
492
+ ) -> torch.Tensor:
493
+ assert patch_embed.dim() == 2 and grid_thw is not None
494
+ x = self.ln_q(patch_embed)
495
+ x = x.reshape(-1, self.merged_dim)
496
+ return self.mlp(x)
497
+
498
+
499
+_ADAPTER_CLASSES: dict[str, type[nn.Module]] = {
500
+ "pixel_shuffle_mlp": PixelShuffleAdapter,
501
+ "patch_merger": PatchMergerAdapter,
502
+}
503
+
504
+
505
+# ---- Full Model ----
506
+
507
+
508
+class DotsMoEVitModel(PreTrainedModel):
509
+ config_class = DotsMoEVitConfig
510
+
511
+ def __init__(self, config: DotsMoEVitConfig) -> None:
512
+ super().__init__(config)
513
+ self.config = config
514
+ self.spatial_merge_size = config.spatial_merge_size
515
+
516
+ self.patch_embed = DotsPatchEmbed(config)
517
+
518
+ head_dim = config.embed_dim // config.num_attention_heads
519
+ self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2, cache_seq_len=100000)
520
+
521
+ self.blocks = nn.ModuleList(
522
+ [MoEVisionBlock(config, i) for i in range(config.num_hidden_layers)]
523
+ )
524
+
525
+ if config.post_norm:
526
+ self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
527
+
528
+ adapter_cls = _ADAPTER_CLASSES.get(config.adapter_type)
529
+ if adapter_cls is None:
530
+ raise ValueError(f"Unknown adapter_type {config.adapter_type!r}")
531
+ self.adapter = adapter_cls(config)
532
+
533
+ self.gradient_checkpointing = False
534
+ self._gradient_checkpointing_func = torch.utils.checkpoint.checkpoint
535
+
536
+ @property
537
+ def dtype(self) -> torch.dtype:
538
+ mlp = self.blocks[0].mlp
539
+ if hasattr(mlp, "fc13"):
540
+ return mlp.fc13.weight.dtype
541
+ if hasattr(mlp, "fc1"):
542
+ return mlp.fc1.weight.dtype
543
+ expert = mlp.experts[0]
544
+ if hasattr(expert, "fc13"):
545
+ return expert.fc13.weight.dtype
546
+ return expert.fc1.weight.dtype
547
+
548
+ @property
549
+ def device(self) -> torch.device:
550
+ return self.patch_embed.proj.weight.device
551
+
552
+ def process_weights_after_loading(self) -> None:
553
+ for block in self.blocks:
554
+ if isinstance(block.mlp, MoESwiGLUFFNFP8):
555
+ block.mlp.process_weights_after_loading()
556
+ if self.config.enable_torch_compile:
557
+ self.compile_block_modules()
558
+
559
+ def compile_block_modules(self, **compile_kwargs: Any) -> None:
560
+ for block in self.blocks:
561
+ block.attn = torch.compile(block.attn, **compile_kwargs)
562
+ block.norm_1 = torch.compile(block.norm_1, **compile_kwargs)
563
+ block.norm_2 = torch.compile(block.norm_2, **compile_kwargs)
564
+
565
+ def get_pos_ids_by_grid(self, grid_thw):
566
+ # Mirrors ``cybertron`` ``AIMv2NativeModel.rot_pos_emb``: when ``pre_pixel_shuffle``
567
+ # is set, RoPE positions follow the qwen ``merge_size`` grouped layout (default 2x2);
568
+ # otherwise positions are flat row-major regardless of ``spatial_merge_size``.
569
+ if self.config.pre_pixel_shuffle:
570
+ rope_merge_size = (
571
+ self.spatial_merge_size if self.spatial_merge_size > 1 else 2
572
+ )
573
+ else:
574
+ rope_merge_size = 1
575
+ pos_ids = []
576
+ for t, h, w in grid_thw:
577
+ hpos_ids = (
578
+ torch.arange(h, device=grid_thw.device).unsqueeze(1).expand(-1, w)
579
+ )
580
+ hpos_ids = hpos_ids.reshape(
581
+ h // rope_merge_size,
582
+ rope_merge_size,
583
+ w // rope_merge_size,
584
+ rope_merge_size,
585
+ )
586
+ hpos_ids = hpos_ids.permute(0, 2, 1, 3).flatten()
587
+
588
+ wpos_ids = (
589
+ torch.arange(w, device=grid_thw.device).unsqueeze(0).expand(h, -1)
590
+ )
591
+ wpos_ids = wpos_ids.reshape(
592
+ h // rope_merge_size,
593
+ rope_merge_size,
594
+ w // rope_merge_size,
595
+ rope_merge_size,
596
+ )
597
+ wpos_ids = wpos_ids.permute(0, 2, 1, 3).flatten()
598
+ pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))
599
+ return pos_ids
600
+
601
+ def rot_pos_emb(self, grid_thw):
602
+ pos_ids = self.get_pos_ids_by_grid(grid_thw)
603
+ pos_ids = torch.cat(pos_ids, dim=0)
604
+ max_grid_size = grid_thw[:, 1:].max()
605
+ rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)
606
+ rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)
607
+ return rotary_pos_emb
608
+
609
+ def _build_cu_seqlens_from_grid(self, grid_thw: torch.Tensor):
610
+ cu_seqlens = torch.repeat_interleave(
611
+ grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]
612
+ ).cumsum(
613
+ dim=0,
614
+ dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,
615
+ )
616
+ cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)
617
+ # Same as (cu_seqlens[1:] - cu_seqlens[:-1]).max(); computed here to avoid D2H inside attention.
618
+ max_seqlen = int((grid_thw[:, 1] * grid_thw[:, 2]).max().item())
619
+ return cu_seqlens, max_seqlen
620
+
621
+ def _build_single_temporal_cu_seqlens_from_grid(self, grid_thw: torch.Tensor):
622
+ seq_lens = grid_thw[:, 1] * grid_thw[:, 2]
623
+ cu_seqlens = torch.empty(
624
+ (grid_thw.shape[0] + 1,), device=grid_thw.device, dtype=torch.int32
625
+ )
626
+ cu_seqlens[0] = 0
627
+ torch.cumsum(seq_lens, dim=0, dtype=torch.int32, out=cu_seqlens[1:])
628
+ max_seqlen = int(seq_lens.max().item())
629
+ return cu_seqlens, max_seqlen
630
+
631
+ def forward(
632
+ self, hidden_states: torch.Tensor, grid_thw: torch.Tensor, bf16=True
633
+ ) -> torch.Tensor:
634
+ grid_thw = grid_thw.to(device=self.device)
635
+ if bf16:
636
+ hidden_states = hidden_states.to(device=self.device, dtype=self.dtype)
637
+ hidden_states = self.patch_embed(hidden_states)
638
+
639
+ rotary_pos_emb = self.rot_pos_emb(grid_thw)
640
+
641
+ if grid_thw[:, 0].sum().item() == grid_thw.shape[0]:
642
+ cu_seqlens, max_seqlen = self._build_single_temporal_cu_seqlens_from_grid(
643
+ grid_thw
644
+ )
645
+ else:
646
+ cu_seqlens, max_seqlen = self._build_cu_seqlens_from_grid(grid_thw)
647
+
648
+ seqlens = prepare_seqlens_for_attention(
649
+ self.config.attn_implementation,
650
+ cu_seqlens,
651
+ eager_fallback="eager_v2",
652
+ )
653
+
654
+ for blk in self.blocks:
655
+ if self.gradient_checkpointing and self.training:
656
+ hidden_states = self._gradient_checkpointing_func(
657
+ blk.__call__,
658
+ hidden_states,
659
+ cu_seqlens,
660
+ rotary_pos_emb,
661
+ max_seqlen,
662
+ seqlens,
663
+ )
664
+ else:
665
+ hidden_states = blk(
666
+ hidden_states,
667
+ cu_seqlens,
668
+ rotary_pos_emb,
669
+ max_seqlen,
670
+ seqlens=seqlens,
671
+ )
672
+
673
+ if self.config.post_norm:
674
+ hidden_states = self.post_trunk_norm(hidden_states)
675
+
676
+ hidden_states = self.adapter(hidden_states, grid_thw)
677
+ return hidden_states
vllm/models/dots3_note/nvidia/vision_attention.py
new
+477
@@ -0,0 +1,477 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+# ruff: noqa: E501
4
+"""Shared vision attention stack for Dots dense / MoE ViT encoders.
5
+
6
+Exports every attention backend (eager, eager_v2, sdpa, flash_attention_2,
7
+flash_attention_3), RoPE helpers, and block wiring utilities.
8
+"""
9
+
10
+from __future__ import annotations
11
+
12
+import math
13
+from dataclasses import dataclass
14
+from typing import Any
15
+
16
+import torch
17
+import torch.nn.functional as F
18
+from torch import nn
19
+
20
+from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported
21
+
22
+_flash_attn_available = is_fa_version_supported(2)
23
+_flash_attn3_available = is_fa_version_supported(3)
24
+_flash_attn3_varlen_func = flash_attn_varlen_func
25
+_flash_attn3_use_positional_args = False
26
+
27
+
28
+# ---------------------------------------------------------------------------
29
+# RoPE
30
+# ---------------------------------------------------------------------------
31
+
32
+
33
+def rotate_half(x: torch.Tensor) -> torch.Tensor:
34
+ x1 = x[..., : x.shape[-1] // 2]
35
+ x2 = x[..., x.shape[-1] // 2 :]
36
+ return torch.cat((-x2, x1), dim=-1)
37
+
38
+
39
+def apply_rotary_pos_emb_vision(
40
+ tensor: torch.Tensor, freqs: torch.Tensor
41
+) -> torch.Tensor:
42
+ orig_dtype = tensor.dtype
43
+ tensor = tensor.float()
44
+ cos = freqs.cos()
45
+ sin = freqs.sin()
46
+ cos = cos.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
47
+ sin = sin.unsqueeze(1).repeat(1, 1, 2).unsqueeze(0).float()
48
+ output = (tensor * cos) + (rotate_half(tensor) * sin)
49
+ return output.to(orig_dtype)
50
+
51
+
52
+class VisionRotaryEmbedding(nn.Module):
53
+ """2D vision RoPE frequency table.
54
+
55
+ When ``cache_seq_len`` is set, precomputes and reuses frequencies up to that
56
+ length (MoE default). Otherwise frequencies are computed on each forward
57
+ (dense default).
58
+ """
59
+
60
+ def __init__(
61
+ self,
62
+ dim: int,
63
+ theta: float = 10000.0,
64
+ cache_seq_len: int | None = None,
65
+ ) -> None:
66
+ super().__init__()
67
+ inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim))
68
+ self.register_buffer("inv_freq", inv_freq, persistent=False)
69
+ self._cache_seq_len = cache_seq_len
70
+ self.freqs_cache: torch.Tensor
71
+ if cache_seq_len is not None:
72
+ self.register_buffer(
73
+ "freqs_cache",
74
+ self._compute_freqs(cache_seq_len),
75
+ persistent=False,
76
+ )
77
+
78
+ def _compute_freqs(self, seqlen: int) -> torch.Tensor:
79
+ seq = torch.arange(
80
+ seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype
81
+ )
82
+ return torch.outer(seq, self.inv_freq)
83
+
84
+ def forward(self, seqlen: int) -> torch.Tensor:
85
+ if self._cache_seq_len is None:
86
+ return self._compute_freqs(seqlen)
87
+ if seqlen > self.freqs_cache.shape[0]:
88
+ self.freqs_cache = self._compute_freqs(seqlen)
89
+ return self.freqs_cache[:seqlen]
90
+
91
+
92
+# ---------------------------------------------------------------------------
93
+# Attention modules
94
+# ---------------------------------------------------------------------------
95
+
96
+
97
+class _RMSNorm(nn.Module):
98
+ """Q/K norm inside attention; matches Dots ViT RMSNorm (fp32 reduce, then cast back)."""
99
+
100
+ def __init__(self, dim: int, eps: float = 1e-6):
101
+ super().__init__()
102
+ self.weight = nn.Parameter(torch.ones(dim))
103
+ self.eps = eps
104
+
105
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
106
+ output = self._norm(x.float()).type_as(x)
107
+ return output * self.weight
108
+
109
+ def _norm(self, x: torch.Tensor) -> torch.Tensor:
110
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
111
+
112
+
113
+@dataclass(frozen=True)
114
+class VisionAttentionParams:
115
+ dim: int
116
+ num_heads: int
117
+ bias: bool = True
118
+ is_causal: bool = False
119
+ rms_norm_eps: float = 1e-5
120
+ use_qk_norm: bool = False
121
+
122
+
123
+def vision_attention_params_from_config(config: Any) -> VisionAttentionParams:
124
+ return VisionAttentionParams(
125
+ dim=config.embed_dim,
126
+ num_heads=config.num_attention_heads,
127
+ bias=getattr(config, "use_bias", True),
128
+ is_causal=config.is_causal,
129
+ rms_norm_eps=config.rms_norm_eps,
130
+ use_qk_norm=getattr(config, "use_qk_norm", False),
131
+ )
132
+
133
+
134
+class _VisionAttentionBase(nn.Module):
135
+ """QKV projection, optional Q/K norm, and output projection."""
136
+
137
+ def __init__(self, params: VisionAttentionParams) -> None:
138
+ super().__init__()
139
+ self.num_heads = params.num_heads
140
+ self.head_dim = params.dim // params.num_heads
141
+ self.is_causal = params.is_causal
142
+ self.use_qk_norm = params.use_qk_norm
143
+ self.qkv = nn.Linear(params.dim, params.dim * 3, bias=params.bias)
144
+ self.proj = nn.Linear(params.dim, params.dim, bias=params.bias)
145
+ if self.use_qk_norm:
146
+ self.q_norm = _RMSNorm(self.head_dim, eps=params.rms_norm_eps)
147
+ self.k_norm = _RMSNorm(self.head_dim, eps=params.rms_norm_eps)
148
+
149
+ def _qkv_with_rope(
150
+ self,
151
+ hidden_states: torch.Tensor,
152
+ rotary_pos_emb: torch.Tensor | None,
153
+ ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
154
+ seq_length = hidden_states.shape[0]
155
+ q, k, v = (
156
+ self.qkv(hidden_states)
157
+ .reshape(seq_length, 3, self.num_heads, -1)
158
+ .permute(1, 0, 2, 3)
159
+ .unbind(0)
160
+ )
161
+ if self.use_qk_norm:
162
+ q = self.q_norm(q)
163
+ k = self.k_norm(k)
164
+ q = apply_rotary_pos_emb_vision(q.unsqueeze(0), rotary_pos_emb).squeeze(0)
165
+ k = apply_rotary_pos_emb_vision(k.unsqueeze(0), rotary_pos_emb).squeeze(0)
166
+ return q, k, v
167
+
168
+
169
+class VisionAttention(_VisionAttentionBase):
170
+ """Eager attention with a dense block-diagonal mask (cu_seqlens boundaries)."""
171
+
172
+ def forward(
173
+ self,
174
+ hidden_states: torch.Tensor,
175
+ cu_seqlens: torch.Tensor,
176
+ max_seqlen: int,
177
+ rotary_pos_emb: torch.Tensor | None = None,
178
+ seqlens: list[int] | None = None,
179
+ ) -> torch.Tensor:
180
+ seq_length = hidden_states.shape[0]
181
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
182
+
183
+ attention_mask = torch.full(
184
+ [1, seq_length, seq_length],
185
+ torch.finfo(q.dtype).min,
186
+ device=q.device,
187
+ dtype=q.dtype,
188
+ )
189
+ for i in range(1, len(cu_seqlens)):
190
+ attention_mask[
191
+ ...,
192
+ cu_seqlens[i - 1] : cu_seqlens[i],
193
+ cu_seqlens[i - 1] : cu_seqlens[i],
194
+ ] = 0
195
+
196
+ q = q.transpose(0, 1)
197
+ k = k.transpose(0, 1)
198
+ v = v.transpose(0, 1)
199
+ attn_weights = torch.matmul(q, k.transpose(1, 2)) / math.sqrt(self.head_dim)
200
+ attn_weights = attn_weights + attention_mask
201
+ attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(q.dtype)
202
+ attn_output = torch.matmul(attn_weights, v)
203
+ attn_output = attn_output.transpose(0, 1).reshape(seq_length, -1)
204
+ return self.proj(attn_output)
205
+
206
+
207
+class VisionAttentionV2(_VisionAttentionBase):
208
+ """Eager attention per varlen segment (lower peak memory than :class:`VisionAttention`)."""
209
+
210
+ def forward(
211
+ self,
212
+ hidden_states: torch.Tensor,
213
+ cu_seqlens: torch.Tensor,
214
+ max_seqlen: int,
215
+ rotary_pos_emb: torch.Tensor | None = None,
216
+ seqlens: list[int] | None = None,
217
+ ) -> torch.Tensor:
218
+ seq_length = hidden_states.shape[0]
219
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
220
+
221
+ if seqlens is None:
222
+ seqlens = torch.diff(cu_seqlens).tolist()
223
+
224
+ outputs = []
225
+ for q_i, k_i, v_i in zip(
226
+ torch.split(q, seqlens, 0),
227
+ torch.split(k, seqlens, 0),
228
+ torch.split(v, seqlens, 0),
229
+ ):
230
+ q_i = q_i.transpose(0, 1)
231
+ k_i = k_i.transpose(0, 1)
232
+ v_i = v_i.transpose(0, 1)
233
+ out = torch.matmul(q_i, k_i.transpose(1, 2)) / math.sqrt(self.head_dim)
234
+ out = F.softmax(out, dim=-1, dtype=torch.float32).to(q.dtype)
235
+ out = torch.matmul(out, v_i).transpose(0, 1)
236
+ outputs.append(out)
237
+
238
+ attn_output = torch.concat(outputs, dim=0).reshape(seq_length, -1)
239
+ return self.proj(attn_output)
240
+
241
+
242
+# MoE checkpoints historically named this implementation ``VisionEagerAttention``.
243
+VisionEagerAttention = VisionAttentionV2
244
+
245
+
246
+class VisionFlashAttention2(_VisionAttentionBase):
247
+ def forward(
248
+ self,
249
+ hidden_states: torch.Tensor,
250
+ cu_seqlens: torch.Tensor,
251
+ max_seqlen: int,
252
+ rotary_pos_emb: torch.Tensor | None = None,
253
+ seqlens: list[int] | None = None,
254
+ ) -> torch.Tensor:
255
+ seq_length = hidden_states.shape[0]
256
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
257
+ attn_output = flash_attn_varlen_func(
258
+ q,
259
+ k,
260
+ v,
261
+ cu_seqlens_q=cu_seqlens,
262
+ cu_seqlens_k=cu_seqlens,
263
+ max_seqlen_q=max_seqlen,
264
+ max_seqlen_k=max_seqlen,
265
+ causal=self.is_causal,
266
+ fa_version=2,
267
+ ).reshape(seq_length, -1)
268
+ return self.proj(attn_output)
269
+
270
+
271
+class VisionFlashAttention3(_VisionAttentionBase):
272
+ def forward(
273
+ self,
274
+ hidden_states: torch.Tensor,
275
+ cu_seqlens: torch.Tensor,
276
+ max_seqlen: int,
277
+ rotary_pos_emb: torch.Tensor | None = None,
278
+ seqlens: list[int] | None = None,
279
+ ) -> torch.Tensor:
280
+ seq_length = hidden_states.shape[0]
281
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
282
+ if _flash_attn3_use_positional_args:
283
+ attn_output = _flash_attn3_varlen_func(
284
+ q,
285
+ k,
286
+ v,
287
+ cu_seqlens,
288
+ cu_seqlens,
289
+ max_seqlen,
290
+ max_seqlen,
291
+ causal=self.is_causal,
292
+ )
293
+ else:
294
+ result = _flash_attn3_varlen_func(
295
+ q,
296
+ k,
297
+ v,
298
+ cu_seqlens_q=cu_seqlens,
299
+ cu_seqlens_k=cu_seqlens,
300
+ max_seqlen_q=max_seqlen,
301
+ max_seqlen_k=max_seqlen,
302
+ causal=self.is_causal,
303
+ fa_version=3,
304
+ )
305
+ attn_output = result[0] if isinstance(result, tuple) else result
306
+ attn_output = attn_output.reshape(seq_length, -1)
307
+ return self.proj(attn_output)
308
+
309
+
310
+class VisionSdpaAttention(_VisionAttentionBase):
311
+ def forward(
312
+ self,
313
+ hidden_states: torch.Tensor,
314
+ cu_seqlens: torch.Tensor,
315
+ max_seqlen: int,
316
+ rotary_pos_emb: torch.Tensor | None = None,
317
+ seqlens: list[int] | None = None,
318
+ ) -> torch.Tensor:
319
+ seq_length = hidden_states.shape[0]
320
+ q, k, v = self._qkv_with_rope(hidden_states, rotary_pos_emb)
321
+
322
+ attention_mask = torch.zeros(
323
+ [1, seq_length, seq_length], device=q.device, dtype=torch.bool
324
+ )
325
+ for i in range(1, len(cu_seqlens)):
326
+ attention_mask[
327
+ ...,
328
+ cu_seqlens[i - 1] : cu_seqlens[i],
329
+ cu_seqlens[i - 1] : cu_seqlens[i],
330
+ ] = True
331
+
332
+ q = q.transpose(0, 1).unsqueeze(0)
333
+ k = k.transpose(0, 1).unsqueeze(0)
334
+ v = v.transpose(0, 1).unsqueeze(0)
335
+
336
+ if attention_mask.stride(-1) != 1:
337
+ attention_mask = torch.empty_like(
338
+ attention_mask, memory_format=torch.contiguous_format
339
+ ).copy_(attention_mask)
340
+
341
+ from torch.nn.attention import SDPBackend, sdpa_kernel
342
+
343
+ with sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION):
344
+ attn_output = F.scaled_dot_product_attention(
345
+ q, k, v, attention_mask, dropout_p=0.0
346
+ )
347
+
348
+ attn_output = attn_output.squeeze(0).transpose(0, 1).reshape(seq_length, -1)
349
+ return self.proj(attn_output)
350
+
351
+
352
+DOTS_VISION_ATTENTION_CLASSES: dict[str, type[nn.Module]] = {
353
+ "eager": VisionAttention,
354
+ "eager_v2": VisionAttentionV2,
355
+ "flash_attention_2": VisionFlashAttention2,
356
+ "flash_attention_3": VisionFlashAttention3,
357
+ "sdpa": VisionSdpaAttention,
358
+}
359
+
360
+
361
+def resolve_attn_implementation(
362
+ attn_implementation: str,
363
+ *,
364
+ eager_fallback: str = "eager",
365
+) -> str:
366
+ """Apply FA3 → FA2 → eager fallback when backends are missing.
367
+
368
+ ``eager_fallback`` selects which eager backend is used when flash-attn is
369
+ unavailable. MoE historically used per-segment eager (``eager_v2``); dense
370
+ uses full-mask ``eager``.
371
+ """
372
+ if attn_implementation == "flash_attention_3" and not _flash_attn3_available:
373
+ print(
374
+ "flash attention 3 not available! fallback to flash attention 2 implementation"
375
+ )
376
+ attn_implementation = "flash_attention_2"
377
+ if attn_implementation == "flash_attention_2" and not _flash_attn_available:
378
+ print("flash attention 2 not available! fallback to eager implementation")
379
+ attn_implementation = eager_fallback
380
+ return attn_implementation
381
+
382
+
383
+def build_vision_attention(
384
+ attn_implementation: str,
385
+ config: Any,
386
+ *,
387
+ dim: int | None = None,
388
+ num_heads: int | None = None,
389
+ bias: bool | None = None,
390
+ eager_fallback: str = "eager",
391
+) -> nn.Module:
392
+ """Instantiate a vision attention module from ``attn_implementation`` and config."""
393
+ attn_implementation = resolve_attn_implementation(
394
+ attn_implementation, eager_fallback=eager_fallback
395
+ )
396
+ if attn_implementation not in DOTS_VISION_ATTENTION_CLASSES:
397
+ raise ValueError(
398
+ f"Unknown attn_implementation {attn_implementation!r}; "
399
+ f"expected one of {sorted(DOTS_VISION_ATTENTION_CLASSES)}"
400
+ )
401
+ if dim is None:
402
+ params = vision_attention_params_from_config(config)
403
+ else:
404
+ params = VisionAttentionParams(
405
+ dim=dim,
406
+ num_heads=num_heads
407
+ if num_heads is not None
408
+ else config.num_attention_heads,
409
+ bias=bias if bias is not None else getattr(config, "use_bias", True),
410
+ is_causal=config.is_causal,
411
+ rms_norm_eps=config.rms_norm_eps,
412
+ use_qk_norm=getattr(config, "use_qk_norm", False),
413
+ )
414
+ return DOTS_VISION_ATTENTION_CLASSES[attn_implementation](params)
415
+
416
+
417
+def attn_uses_seqlens(attn_implementation: str) -> bool:
418
+ return attn_implementation == "eager_v2"
419
+
420
+
421
+def prepare_seqlens_for_attention(
422
+ attn_implementation: str,
423
+ cu_seqlens: torch.Tensor,
424
+ *,
425
+ eager_fallback: str = "eager",
426
+) -> list[int] | None:
427
+ """Return per-segment lengths when the resolved backend needs ``seqlens``."""
428
+ resolved = resolve_attn_implementation(
429
+ attn_implementation, eager_fallback=eager_fallback
430
+ )
431
+ if attn_uses_seqlens(resolved):
432
+ return torch.diff(cu_seqlens).tolist()
433
+ return None
434
+
435
+
436
+def apply_vision_attention_residual(
437
+ attn: nn.Module,
438
+ norm: nn.Module,
439
+ hidden_states: torch.Tensor,
440
+ cu_seqlens: torch.Tensor,
441
+ max_seqlen: int,
442
+ rotary_pos_emb: torch.Tensor,
443
+ *,
444
+ seqlens: list[int] | None = None,
445
+ uses_seqlens: bool = False,
446
+) -> torch.Tensor:
447
+ """Pre-norm residual attention used by dense / MoE vision blocks."""
448
+ attn_kwargs: dict[str, Any] = {"rotary_pos_emb": rotary_pos_emb}
449
+ if uses_seqlens:
450
+ attn_kwargs["seqlens"] = seqlens
451
+ return hidden_states + attn(
452
+ norm(hidden_states),
453
+ cu_seqlens,
454
+ max_seqlen,
455
+ **attn_kwargs,
456
+ )
457
+
458
+
459
+__all__ = [
460
+ "DOTS_VISION_ATTENTION_CLASSES",
461
+ "VisionAttention",
462
+ "VisionAttentionParams",
463
+ "VisionAttentionV2",
464
+ "VisionEagerAttention",
465
+ "VisionFlashAttention2",
466
+ "VisionFlashAttention3",
467
+ "VisionRotaryEmbedding",
468
+ "VisionSdpaAttention",
469
+ "apply_rotary_pos_emb_vision",
470
+ "apply_vision_attention_residual",
471
+ "attn_uses_seqlens",
472
+ "build_vision_attention",
473
+ "prepare_seqlens_for_attention",
474
+ "resolve_attn_implementation",
475
+ "rotate_half",
476
+ "vision_attention_params_from_config",
477
+]
vllm/models/dots3_note/nvidia/vision_moe.py
new
+149
@@ -0,0 +1,149 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""NOTE vision MoE execution matching the native encoder's FP8 semantics."""
4
+
5
+import torch
6
+
7
+from vllm import _custom_ops as ops
8
+from vllm.model_executor.layers.fused_moe.activation import (
9
+ MoEActivation,
10
+ apply_moe_activation,
11
+)
12
+from vllm.model_executor.layers.fused_moe.fused_moe import (
13
+ _prepare_expert_assignment,
14
+ dispatch_fused_moe_kernel,
15
+ try_get_optimal_moe_config,
16
+)
17
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18
+ per_token_group_quant_fp8,
19
+)
20
+from vllm.triton_utils import tl
21
+
22
+_BLOCK_SHAPE = [128, 128]
23
+
24
+
25
+def note_vision_fused_moe_fp8(
26
+ hidden_states: torch.Tensor,
27
+ w13: torch.Tensor,
28
+ w2: torch.Tensor,
29
+ topk_weights: torch.Tensor,
30
+ topk_ids: torch.Tensor,
31
+ w13_scale: torch.Tensor,
32
+ w2_scale: torch.Tensor,
33
+) -> torch.Tensor:
34
+ """Run top-k SwiGLU MoE with the encoder's block-FP8 formula.
35
+
36
+ Unlike vLLM's generic functional MoE path, the native NOTE encoder keeps
37
+ dynamic activation scales as FP32 instead of rounding them to E8M0. Weight
38
+ scales are block 128x128 and activations are quantized per token/group-128.
39
+ """
40
+ if hidden_states.dtype != torch.bfloat16:
41
+ raise TypeError(
42
+ "NOTE vision FP8 MoE requires bfloat16 activations, got "
43
+ f"{hidden_states.dtype}"
44
+ )
45
+
46
+ num_tokens = hidden_states.shape[0]
47
+ num_experts, intermediate_size, _ = w13.shape
48
+ output_size = w2.shape[1]
49
+ topk = topk_ids.shape[1]
50
+ config = try_get_optimal_moe_config(
51
+ w13.shape,
52
+ w2.shape,
53
+ topk,
54
+ "fp8_w8a8",
55
+ num_tokens,
56
+ block_shape=_BLOCK_SHAPE,
57
+ )
58
+
59
+ cache = torch.empty(
60
+ num_tokens * topk * max(intermediate_size, output_size),
61
+ dtype=hidden_states.dtype,
62
+ device=hidden_states.device,
63
+ )
64
+ first_gemm = cache[: num_tokens * topk * intermediate_size].view(
65
+ num_tokens, topk, intermediate_size
66
+ )
67
+ second_gemm = cache[: num_tokens * topk * output_size].view(
68
+ num_tokens, topk, output_size
69
+ )
70
+ activated_size = intermediate_size // 2
71
+ activated = torch.empty(
72
+ (num_tokens * topk, activated_size),
73
+ dtype=hidden_states.dtype,
74
+ device=hidden_states.device,
75
+ )
76
+
77
+ quantized_input, input_scale = per_token_group_quant_fp8(
78
+ hidden_states,
79
+ _BLOCK_SHAPE[1],
80
+ use_ue8m0=False,
81
+ )
82
+ sorted_token_ids, expert_ids, num_tokens_post_padded = _prepare_expert_assignment(
83
+ topk_ids,
84
+ config,
85
+ num_tokens,
86
+ topk,
87
+ num_experts,
88
+ None,
89
+ block_shape=_BLOCK_SHAPE,
90
+ )
91
+ dispatch_fused_moe_kernel(
92
+ quantized_input,
93
+ w13,
94
+ first_gemm,
95
+ input_scale,
96
+ w13_scale,
97
+ None,
98
+ topk_weights,
99
+ sorted_token_ids,
100
+ expert_ids,
101
+ num_tokens_post_padded,
102
+ False,
103
+ topk,
104
+ config,
105
+ tl.bfloat16,
106
+ True,
107
+ False,
108
+ False,
109
+ False,
110
+ False,
111
+ _BLOCK_SHAPE,
112
+ )
113
+
114
+ apply_moe_activation(
115
+ MoEActivation.SILU,
116
+ activated,
117
+ first_gemm.view(-1, intermediate_size),
118
+ )
119
+ quantized_activated, activated_scale = per_token_group_quant_fp8(
120
+ activated,
121
+ _BLOCK_SHAPE[1],
122
+ use_ue8m0=False,
123
+ )
124
+ dispatch_fused_moe_kernel(
125
+ quantized_activated,
126
+ w2,
127
+ second_gemm,
128
+ activated_scale,
129
+ w2_scale,
130
+ None,
131
+ topk_weights,
132
+ sorted_token_ids,
133
+ expert_ids,
134
+ num_tokens_post_padded,
135
+ True,
136
+ 1,
137
+ config,
138
+ tl.bfloat16,
139
+ True,
140
+ False,
141
+ False,
142
+ False,
143
+ False,
144
+ _BLOCK_SHAPE,
145
+ )
146
+
147
+ output = torch.empty_like(hidden_states)
148
+ ops.moe_sum(second_gemm, output)
149
+ return output
vllm/tool_parsers/__init__.py
+4
@@ -22,6 +22,10 @@ Example:
22
23
24
_TOOL_PARSERS_TO_REGISTER = {
25
+ "dots": (
26
+ "dots_tool_parser",
27
+ "DotsToolParser",
28
+ ),
29
"deepseek_v3": ( # name
30
"deepseekv3_tool_parser", # filename
31
"DeepSeekV3ToolParser", # class_name
vllm/tool_parsers/dots_tool_parser.py
new
+434
@@ -0,0 +1,434 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import json
5
+from collections.abc import Sequence
6
+from typing import Any
7
+
8
+import regex as re
9
+
10
+from vllm.entrypoints.chat_utils import make_tool_call_id
11
+from vllm.entrypoints.openai.chat_completion.protocol import (
12
+ ChatCompletionRequest,
13
+ ChatCompletionToolsParam,
14
+)
15
+from vllm.entrypoints.openai.engine.protocol import (
16
+ DeltaFunctionCall,
17
+ DeltaMessage,
18
+ DeltaToolCall,
19
+ ExtractedToolCallInformation,
20
+ FunctionCall,
21
+ ToolCall,
22
+)
23
+from vllm.logger import init_logger
24
+from vllm.tokenizers import TokenizerLike
25
+from vllm.tool_parsers.abstract_tool_parser import Tool, ToolParser
26
+from vllm.tool_parsers.utils import is_complete_json, partial_tag_overlap
27
+
28
+logger = init_logger(__name__)
29
+
30
+
31
+class DotsToolParser(ToolParser):
32
+ """Parse Dots tool calls in their XML wrapper format.
33
+
34
+ The canonical body contains one or more ``invoke`` elements::
35
+
36
+ <dots_function_call>
37
+ <invoke name="search">
38
+ <parameter name="query">weather in Shanghai</parameter>
39
+ </invoke>
40
+ </dots_function_call>
41
+
42
+ A JSON object with ``name`` and ``arguments`` is also accepted as a
43
+ fallback. Multiple wrapper blocks and multiple invokes per block are
44
+ supported.
45
+ """
46
+
47
+ supports_required_and_named = False
48
+
49
+ tool_call_start_token = "<dots_function_call>"
50
+ tool_call_end_token = "</dots_function_call>"
51
+
52
+ _block_regex = re.compile(
53
+ rf"{re.escape(tool_call_start_token)}\s*(.*?)\s*"
54
+ rf"{re.escape(tool_call_end_token)}",
55
+ re.DOTALL,
56
+ )
57
+ _invoke_regex = re.compile(
58
+ r"<invoke\s+name\s*=\s*(?P<name>[^>]+)>(?P<body>.*?)</invoke>",
59
+ re.DOTALL,
60
+ )
61
+ _parameter_regex = re.compile(
62
+ r"<parameter\s+name\s*=\s*(?P<name>[^>]+)>(?P<value>.*?)</parameter>",
63
+ re.DOTALL,
64
+ )
65
+
66
+ def __init__(
67
+ self,
68
+ tokenizer: TokenizerLike,
69
+ tools: list[Tool] | None = None,
70
+ ) -> None:
71
+ super().__init__(tokenizer, tools)
72
+ self._buffer = ""
73
+
74
+ @staticmethod
75
+ def _extract_name(value: str) -> str:
76
+ value = value.strip()
77
+ if len(value) >= 2 and value[0] == value[-1] and value[0] in {'"', "'"}:
78
+ return value[1:-1]
79
+ return value
80
+
81
+ @staticmethod
82
+ def _convert_param_value(value: str, param_type: Any) -> Any:
83
+ if value.lower() == "null":
84
+ return None
85
+
86
+ if isinstance(param_type, list):
87
+ param_type = next((item for item in param_type if item != "null"), "string")
88
+ if not isinstance(param_type, str):
89
+ param_type = str(param_type)
90
+ param_type = param_type.lower()
91
+
92
+ if param_type in {"string", "str", "text"}:
93
+ return value
94
+ if param_type in {"integer", "int"}:
95
+ try:
96
+ return int(value)
97
+ except (TypeError, ValueError):
98
+ return value
99
+ if param_type in {"number", "float"}:
100
+ try:
101
+ number = float(value)
102
+ return int(number) if number.is_integer() else number
103
+ except (TypeError, ValueError):
104
+ return value
105
+ if param_type in {"boolean", "bool"}:
106
+ return value.lower() in {"true", "1"}
107
+
108
+ try:
109
+ return json.loads(value)
110
+ except (json.JSONDecodeError, TypeError, ValueError):
111
+ return value
112
+
113
+ def _resolve_param_type(
114
+ self,
115
+ schema: Any,
116
+ defs: dict[str, Any],
117
+ depth: int = 0,
118
+ ) -> Any | None:
119
+ if not isinstance(schema, dict) or depth > 10:
120
+ return None
121
+ if "type" in schema:
122
+ return schema["type"]
123
+
124
+ ref = schema.get("$ref")
125
+ if isinstance(ref, str) and ref.startswith("#/$defs/"):
126
+ return self._resolve_param_type(
127
+ defs.get(ref.rsplit("/", 1)[-1]), defs, depth + 1
128
+ )
129
+
130
+ for keyword in ("anyOf", "oneOf", "allOf"):
131
+ alternatives = schema.get(keyword)
132
+ if not isinstance(alternatives, list):
133
+ continue
134
+ for alternative in alternatives:
135
+ if isinstance(alternative, dict) and alternative.get("type") == "null":
136
+ continue
137
+ resolved = self._resolve_param_type(alternative, defs, depth + 1)
138
+ if resolved is not None:
139
+ return resolved
140
+ return None
141
+
142
+ @staticmethod
143
+ def _tool_schema(
144
+ name: str,
145
+ tools: list[ChatCompletionToolsParam] | None,
146
+ ) -> tuple[dict[str, Any], dict[str, Any]]:
147
+ for tool in tools or []:
148
+ if tool.function.name != name:
149
+ continue
150
+ schema = tool.function.parameters
151
+ if not isinstance(schema, dict):
152
+ break
153
+ properties = schema.get("properties", {})
154
+ defs = schema.get("$defs", {})
155
+ return (
156
+ properties if isinstance(properties, dict) else {},
157
+ defs if isinstance(defs, dict) else {},
158
+ )
159
+ return {}, {}
160
+
161
+ def _parse_xml_invoke(
162
+ self,
163
+ match: re.Match[str],
164
+ tools: list[ChatCompletionToolsParam] | None,
165
+ ) -> dict[str, Any]:
166
+ name = self._extract_name(match.group("name"))
167
+ properties, defs = self._tool_schema(name, tools)
168
+ arguments: dict[str, Any] = {}
169
+
170
+ for parameter in self._parameter_regex.finditer(match.group("body")):
171
+ param_name = self._extract_name(parameter.group("name"))
172
+ value = parameter.group("value").strip()
173
+ param_type: Any = "string"
174
+ if param_name in properties:
175
+ param_type = (
176
+ self._resolve_param_type(properties[param_name], defs) or "string"
177
+ )
178
+ arguments[param_name] = self._convert_param_value(value, param_type)
179
+
180
+ return {"name": name, "arguments": arguments}
181
+
182
+ def _parse_block(
183
+ self,
184
+ content: str,
185
+ tools: list[ChatCompletionToolsParam] | None,
186
+ ) -> list[dict[str, Any]]:
187
+ content = content.strip()
188
+ if content.startswith("<invoke"):
189
+ return [
190
+ self._parse_xml_invoke(match, tools)
191
+ for match in self._invoke_regex.finditer(content)
192
+ ]
193
+
194
+ parsed = json.loads(content)
195
+ if not isinstance(parsed, dict):
196
+ raise TypeError("Dots JSON tool call must be an object")
197
+ return [parsed]
198
+
199
+ @staticmethod
200
+ def _known_tool_names(
201
+ tools: list[ChatCompletionToolsParam] | None,
202
+ ) -> set[str]:
203
+ return {tool.function.name for tool in tools or []}
204
+
205
+ def _validated_call(
206
+ self,
207
+ parsed: dict[str, Any],
208
+ tools: list[ChatCompletionToolsParam] | None,
209
+ ) -> tuple[str, dict[str, Any]] | None:
210
+ name = parsed.get("name")
211
+ if not isinstance(name, str) or name not in self._known_tool_names(tools):
212
+ return None
213
+ arguments = parsed.get("arguments", parsed.get("parameters", {})) or {}
214
+ if not isinstance(arguments, dict):
215
+ return None
216
+ return name, arguments
217
+
218
+ def extract_tool_calls(
219
+ self,
220
+ model_output: str,
221
+ request: ChatCompletionRequest,
222
+ ) -> ExtractedToolCallInformation:
223
+ marker_index = model_output.find(self.tool_call_start_token)
224
+ if marker_index == -1:
225
+ return ExtractedToolCallInformation(
226
+ tools_called=False,
227
+ tool_calls=[],
228
+ content=model_output,
229
+ )
230
+
231
+ tool_calls: list[ToolCall] = []
232
+ for block in self._block_regex.finditer(model_output):
233
+ try:
234
+ for parsed in self._parse_block(block.group(1), request.tools):
235
+ validated = self._validated_call(parsed, request.tools)
236
+ if validated is None:
237
+ continue
238
+ name, arguments = validated
239
+ tool_calls.append(
240
+ ToolCall(
241
+ function=FunctionCall(
242
+ name=name,
243
+ arguments=json.dumps(arguments, ensure_ascii=False),
244
+ )
245
+ )
246
+ )
247
+ except (json.JSONDecodeError, TypeError, ValueError) as exc:
248
+ logger.warning("Failed to parse Dots tool call: %s", exc)
249
+
250
+ normal_text = model_output[:marker_index].strip()
251
+ if not tool_calls:
252
+ return ExtractedToolCallInformation(
253
+ tools_called=False,
254
+ tool_calls=[],
255
+ content=model_output,
256
+ )
257
+ return ExtractedToolCallInformation(
258
+ tools_called=True,
259
+ tool_calls=tool_calls,
260
+ content=normal_text or None,
261
+ )
262
+
263
+ def _append_complete_stream_call(
264
+ self,
265
+ name: str,
266
+ arguments: dict[str, Any],
267
+ tool_calls: list[DeltaToolCall],
268
+ ) -> None:
269
+ self.current_tool_id += 1
270
+ serialized = json.dumps(arguments, ensure_ascii=False)
271
+ self.prev_tool_call_arr.append({"name": name, "arguments": arguments})
272
+ self.streamed_args_for_tool.append(serialized)
273
+ tool_calls.append(
274
+ DeltaToolCall(
275
+ index=self.current_tool_id,
276
+ id=make_tool_call_id(),
277
+ type="function",
278
+ function=DeltaFunctionCall(name=name, arguments=serialized),
279
+ )
280
+ )
281
+
282
+ def _stream_complete_json_body(
283
+ self,
284
+ tools: list[ChatCompletionToolsParam] | None,
285
+ tool_calls: list[DeltaToolCall],
286
+ ) -> None:
287
+ content = self._buffer[len(self.tool_call_start_token) :].strip()
288
+ if not content or not is_complete_json(content):
289
+ return
290
+
291
+ try:
292
+ parsed = json.loads(content)
293
+ except (json.JSONDecodeError, TypeError, ValueError):
294
+ return
295
+ if not isinstance(parsed, dict):
296
+ return
297
+ validated = self._validated_call(parsed, tools)
298
+ if validated is None:
299
+ return
300
+
301
+ name, arguments = validated
302
+ serialized = json.dumps(arguments, ensure_ascii=False)
303
+ if not self.current_tool_name_sent:
304
+ self.current_tool_id += 1
305
+ tool_calls.append(
306
+ DeltaToolCall(
307
+ index=self.current_tool_id,
308
+ id=make_tool_call_id(),
309
+ type="function",
310
+ function=DeltaFunctionCall(name=name, arguments=""),
311
+ )
312
+ )
313
+ self.prev_tool_call_arr.append({"name": name, "arguments": arguments})
314
+ self.streamed_args_for_tool.append("")
315
+ self.current_tool_name_sent = True
316
+
317
+ streamed = self.streamed_args_for_tool[self.current_tool_id]
318
+ if serialized.startswith(streamed):
319
+ argument_diff = serialized[len(streamed) :]
320
+ if argument_diff:
321
+ tool_calls.append(
322
+ DeltaToolCall(
323
+ index=self.current_tool_id,
324
+ function=DeltaFunctionCall(arguments=argument_diff),
325
+ )
326
+ )
327
+ self.streamed_args_for_tool[self.current_tool_id] += argument_diff
328
+
329
+ def extract_tool_calls_streaming(
330
+ self,
331
+ previous_text: str,
332
+ current_text: str,
333
+ delta_text: str,
334
+ previous_token_ids: Sequence[int],
335
+ current_token_ids: Sequence[int],
336
+ delta_token_ids: Sequence[int],
337
+ request: ChatCompletionRequest,
338
+ ) -> DeltaMessage | None:
339
+ del current_text, previous_token_ids, current_token_ids, delta_token_ids
340
+ if not previous_text:
341
+ self._buffer = ""
342
+ self.prev_tool_call_arr = []
343
+ self.current_tool_id = -1
344
+ self.current_tool_name_sent = False
345
+ self.streamed_args_for_tool = []
346
+
347
+ self._buffer += delta_text
348
+ normal_parts: list[str] = []
349
+ tool_calls: list[DeltaToolCall] = []
350
+
351
+ while self._buffer:
352
+ marker_index = self._buffer.find(self.tool_call_start_token)
353
+ if marker_index == -1:
354
+ partial_len = partial_tag_overlap(
355
+ self._buffer, self.tool_call_start_token
356
+ )
357
+ if partial_len:
358
+ normal_parts.append(self._buffer[:-partial_len])
359
+ self._buffer = self._buffer[-partial_len:]
360
+ else:
361
+ normal_parts.append(self._buffer)
362
+ self._buffer = ""
363
+ normal_parts = [
364
+ part.replace(self.tool_call_end_token, "") for part in normal_parts
365
+ ]
366
+ break
367
+
368
+ if marker_index > 0:
369
+ normal_parts.append(self._buffer[:marker_index])
370
+ self._buffer = self._buffer[marker_index:]
371
+
372
+ end_index = self._buffer.find(
373
+ self.tool_call_end_token, len(self.tool_call_start_token)
374
+ )
375
+ if end_index == -1:
376
+ self._stream_complete_json_body(request.tools, tool_calls)
377
+ break
378
+
379
+ content = self._buffer[len(self.tool_call_start_token) : end_index]
380
+ self._buffer = self._buffer[end_index + len(self.tool_call_end_token) :]
381
+ try:
382
+ parsed_calls = self._parse_block(content, request.tools)
383
+ if not parsed_calls:
384
+ raise ValueError("Dots tool-call block contains no invoke")
385
+
386
+ block_had_streamed_call = self.current_tool_name_sent
387
+ valid_calls = [
388
+ validated
389
+ for parsed in parsed_calls
390
+ if (validated := self._validated_call(parsed, request.tools))
391
+ is not None
392
+ ]
393
+ if self.current_tool_name_sent and valid_calls:
394
+ name, arguments = valid_calls.pop(0)
395
+ serialized = json.dumps(arguments, ensure_ascii=False)
396
+ streamed = self.streamed_args_for_tool[self.current_tool_id]
397
+ if serialized.startswith(streamed):
398
+ remaining = serialized[len(streamed) :]
399
+ if remaining:
400
+ tool_calls.append(
401
+ DeltaToolCall(
402
+ index=self.current_tool_id,
403
+ function=DeltaFunctionCall(arguments=remaining),
404
+ )
405
+ )
406
+ self.prev_tool_call_arr[self.current_tool_id] = {
407
+ "name": name,
408
+ "arguments": arguments,
409
+ }
410
+ self.streamed_args_for_tool[self.current_tool_id] = serialized
411
+
412
+ for name, arguments in valid_calls:
413
+ self._append_complete_stream_call(name, arguments, tool_calls)
414
+
415
+ if not valid_calls and not block_had_streamed_call:
416
+ normal_parts.append(content.strip())
417
+ except (json.JSONDecodeError, TypeError, ValueError) as exc:
418
+ logger.warning("Failed to parse streamed Dots tool call: %s", exc)
419
+ normal_parts.append(content.strip())
420
+
421
+ self.current_tool_name_sent = False
422
+
423
+ content_delta = "".join(normal_parts)
424
+ if not content_delta and not tool_calls:
425
+ return None
426
+ return DeltaMessage(content=content_delta or None, tool_calls=tool_calls)
427
+
428
+ def flush_pending_normal_text(self) -> str:
429
+ """Return a partial opening marker as text when generation ends."""
430
+ if not self._buffer or self.tool_call_start_token in self._buffer:
431
+ return ""
432
+ normal_text = self._buffer.replace(self.tool_call_end_token, "")
433
+ self._buffer = ""
434
+ return normal_text
vllm/transformers_utils/config.py
+1
@@ -87,6 +87,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
87
deepseek_vl_v2="DeepseekVLV2Config",
88
deepseek_v32="DeepseekV3Config",
89
deepseek_v4="DeepseekV4Config",
90
+ dots3_note="Dots3NoteConfig",
91
k3_dspark="K3DSparkConfig",
92
flex_olmo="FlexOlmoConfig",
93
fireredlid="FireRedLIDConfig",
vllm/transformers_utils/configs/__init__.py
+2
@@ -35,6 +35,7 @@ _CLASS_TO_MODULE: dict[str, str] = {
35
"DiffusionGemmaTextConfig": "vllm.transformers_utils.configs.diffusion_gemma",
36
"DeepseekVLV2Config": "vllm.transformers_utils.configs.deepseek_vl2",
37
"DeepseekV4Config": "vllm.transformers_utils.configs.deepseek_v4",
38
+ "Dots3NoteConfig": "vllm.transformers_utils.configs.dots3_note",
39
"K3DSparkConfig": "vllm.transformers_utils.configs.k3_dspark",
40
"DotsOCRConfig": "vllm.transformers_utils.configs.dotsocr",
41
"EAGLEConfig": "vllm.transformers_utils.configs.eagle",
@@ -127,6 +128,7 @@ __all__ = [
128
"DeepseekVLV2Config",
129
"DeepseekV3Config",
130
"DeepseekV4Config",
131
+ "Dots3NoteConfig",
132
"K3DSparkConfig",
133
"DotsOCRConfig",
134
"EAGLEConfig",
vllm/transformers_utils/configs/dots3_note.py
new
+25
@@ -0,0 +1,25 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from transformers import DeepseekV3Config
5
+
6
+
7
+class Dots3NoteConfig(DeepseekV3Config):
8
+ """Configuration for the Dots3Note hybrid DSA/SWA model."""
9
+
10
+ model_type = "dots3_note"
11
+
12
+ def __init__(self, **kwargs):
13
+ # These fields are absent from existing Note config.json files. Do
14
+ # not inherit DeepSeek-V3's 8-group/4-group router defaults: Note was
15
+ # trained with an ungrouped (1/1) noaux_tc router, as defined by the
16
+ # reference Dots3Config in SGLang haimain. A different grouping
17
+ # changes the selected experts at every MoE layer.
18
+ kwargs.setdefault("n_group", 1)
19
+ kwargs.setdefault("topk_group", 1)
20
+ # Dots3 projects the indexer RoPE coordinates in adjacent (GPT-J)
21
+ # pairs. DeepSeek-V3.2 defaults to the split-half NeoX layout when
22
+ # this flag is absent, which rotates different learned coordinates.
23
+ kwargs.setdefault("indexer_rope_interleave", True)
24
+ kwargs.setdefault("num_nextn_predict_layers", 1)
25
+ super().__init__(**kwargs)
vllm/transformers_utils/model_arch_config_convertor.py
+1
@@ -317,6 +317,7 @@ class ModelArchConfigConvertorBase:
317
"deepseek_v3",
318
"deepseek_v32",
319
"deepseek_v4",
320
+ "dots3_note",
321
"deepseek_mtp",
322
"k3_dspark",
323
"glm_moe_dsa",