[Model] Add native Dots3 NOTE multimodal support (#51255)

Signed-off-by: KurodaKanbei <mistergalahad@gmail.com>
Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai>
Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>

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)

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",

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

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<dots_function_call>"
+        '<invoke name="set_location">'
+        '<parameter name="location">{"city": "Shanghai"}</parameter>'
+        '<parameter name="days">3</parameter>'
+        '<parameter name="include_weather">true</parameter>'
+        "</invoke>"
+        "</dots_function_call>"
+    )
+
+    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 = (
+        "<dots_function_call>"
+        '<invoke name="search"><parameter name="query">chairs</parameter></invoke>'
+        '<invoke name="open"><parameter name="id">7</parameter></invoke>'
+        "</dots_function_call>"
+        '<dots_function_call>{"name":"search","arguments":{"query":"tables"}}'
+        "</dots_function_call>"
+    )
+
+    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 = (
+        "<dots_function_call>"
+        '<invoke name="ghost"><parameter name="query">chairs</parameter></invoke>'
+        "</dots_function_call>"
+    )
+
+    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<dots_func",
+        (
+            "tion_call>"
+            '<invoke name="search"><parameter name="query">chairs</parameter></invoke>'
+            "</dots_function_call>"
+            "<dots_function_call>"
+            '<invoke name="search"><parameter name="query">tables</parameter></invoke>'
+            "</dots_function_call>"
+        ),
+    ]
+
+    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 = (
+        "<dots_function_call>"
+        '<invoke name="ghost"><parameter name="query">chairs</parameter></invoke>'
+        "</dots_function_call>"
+    )
+
+    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,
+        [
+            "<dots_function_call>garbage</dots_function_call>",
+            (
+                "<dots_function_call>"
+                '<invoke name="search">'
+                '<parameter name="query">chairs</parameter>'
+                "</invoke></dots_function_call>"
+            ),
+        ],
+        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 </dots_function_call>"], 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 <dots_func"], request)
+
+    assert messages[0].content == "answer "
+    assert parser.flush_pending_normal_text() == "<dots_func"
+    assert parser.flush_pending_normal_text() == ""
+
+
+def test_streaming_emits_complete_json_before_end_marker_without_duplication(
+    parser: DotsToolParser,
+) -> None:
+    request = _request([_tool("search", {"query": {"type": "string"}})])
+
+    messages = _stream(
+        parser,
+        [
+            "<dots_function_call>",
+            '{"name":"search","arguments":{"query":"chairs"}}',
+            "</dots_function_call>",
+        ],
+        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

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)

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",

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,

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"),

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"]

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."""

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

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, "<video_0>")
+        + 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

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."""

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

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

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,
+        )

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

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))

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"],
+        )

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

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",
+]

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

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

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::
+
+        <dots_function_call>
+        <invoke name="search">
+        <parameter name="query">weather in Shanghai</parameter>
+        </invoke>
+        </dots_function_call>
+
+    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 = "<dots_function_call>"
+    tool_call_end_token = "</dots_function_call>"
+
+    _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"<invoke\s+name\s*=\s*(?P<name>[^>]+)>(?P<body>.*?)</invoke>",
+        re.DOTALL,
+    )
+    _parameter_regex = re.compile(
+        r"<parameter\s+name\s*=\s*(?P<name>[^>]+)>(?P<value>.*?)</parameter>",
+        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("<invoke"):
+            return [
+                self._parse_xml_invoke(match, tools)
+                for match in self._invoke_regex.finditer(content)
+            ]
+
+        parsed = json.loads(content)
+        if not isinstance(parsed, dict):
+            raise TypeError("Dots JSON tool call must be an object")
+        return [parsed]
+
+    @staticmethod
+    def _known_tool_names(
+        tools: list[ChatCompletionToolsParam] | None,
+    ) -> 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

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",

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",

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)

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",