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