[Performance] Optimize Dots3 NOTE runtime (#53517)
Signed-off-by: KurodaKanbei <mistergalahad@gmail.com> Signed-off-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <jeejeelee@inferact.ai> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
范裕达 committed
Aug 31, 2026 at 15:05 UTC
da0b2d8b17c9203bdd5eeb7a9af575f75a050c09
6 files changed
+94
-85
vllm/config/vllm.py
+2
@@ -78,6 +78,8 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
78
"DeepseekV32ForCausalLM",
79
"DeepseekV4ForCausalLM",
80
"DeepSeekV4MTPModel",
81
+ "Dots3NoteForCausalLM",
82
+ "Dots3NoteMTPModel",
83
"GlmMoeDsaForCausalLM",
84
"HYV4ForCausalLM",
85
"HYV4MTPModel",
vllm/models/dots3_note/nvidia/attention.py
+7
-1
@@ -7,6 +7,7 @@ varlen MHA. Decode-only batches use the Triton absorbed-MQA kernel.
7
"""
8
9
from dataclasses import dataclass
10
+from typing import ClassVar
11
12
import torch
13
@@ -19,7 +20,11 @@ from vllm.model_executor.layers.attention.mla_attention import (
20
from vllm.platforms import current_platform
21
from vllm.triton_utils import tl, triton
22
from vllm.utils.torch_utils import is_quantized_kv_cache
22
-from vllm.v1.attention.backend import AttentionLayer, CommonAttentionMetadata
23
+from vllm.v1.attention.backend import (
24
+ AttentionCGSupport,
25
+ AttentionLayer,
26
+ CommonAttentionMetadata,
27
+)
28
from vllm.v1.attention.backends.mla.flashattn_mla_sparse import (
29
FlashAttnMLASparseBackend,
30
FlashAttnMLASparseImpl,
@@ -307,6 +312,7 @@ class Dots3NoteFlashAttnPrefillBackend(FlashAttnPrefillBackend):
312
class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder):
313
"""Keep decode on MQA and route prefill/mixed batches through FA3."""
314
315
+ _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH
316
query_len_support = QueryLenSupport.UNIFORM
317
318
def __init__(self, kv_cache_spec, layer_names, vllm_config, device):
vllm/models/dots3_note/nvidia/model.py
+25
-22
@@ -19,6 +19,7 @@ from vllm.distributed import (
19
tensor_model_parallel_all_reduce,
20
)
21
from vllm.model_executor.layers.attention import MLAAttention
22
+from vllm.model_executor.layers.fused_embed_norm import has_full_vocab_on_rank
23
from vllm.model_executor.layers.layernorm import RMSNorm
24
from vllm.model_executor.layers.linear import (
25
ColumnParallelLinear,
@@ -73,6 +74,28 @@ def _padded_mlp_size(
74
return ((blocks + tp_size - 1) // tp_size) * block_size[0] * tp_size
75
76
77
+def _pad_dense_mlp_weight(
78
+ name: str,
79
+ loaded_weight: torch.Tensor,
80
+ weight_block_size: list[int] | None,
81
+) -> torch.Tensor:
82
+ if weight_block_size is None or ".mlp.experts." in name:
83
+ return loaded_weight
84
+ if not any(
85
+ proj_name in name for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
86
+ ):
87
+ return loaded_weight
88
+ dim = 1 if ".down_proj." in name else 0
89
+ block_step = 1 if name.endswith("weight_scale_inv") else weight_block_size[0]
90
+ multiple = get_tensor_model_parallel_world_size() * block_step
91
+ pad = (-loaded_weight.shape[dim]) % multiple
92
+ if pad == 0:
93
+ return loaded_weight
94
+ pad_shape = list(loaded_weight.shape)
95
+ pad_shape[dim] = pad
96
+ return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
97
+
98
+
99
class Dots3NoteMoE(DeepseekV2MoE):
100
"""DeepSeek MoE with NOTE's block-padded shared expert."""
101
@@ -574,6 +597,7 @@ class Dots3NoteModel(DeepseekV32Model):
597
)
598
else:
599
self.embed_tokens = PPMissingLayer()
600
+ self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens)
601
self.start_layer, self.end_layer, self.layers = make_layers(
602
config.num_hidden_layers,
603
lambda prefix: Dots3NoteDecoderLayer(
@@ -595,27 +619,6 @@ class Dots3NoteModel(DeepseekV32Model):
619
vllm_config.parallel_config.eplb_config.num_redundant_experts
620
)
621
598
- def _pad_dense_mlp_weight(
599
- self, name: str, loaded_weight: torch.Tensor
600
- ) -> torch.Tensor:
601
- block_size = self._weight_block_size
602
- if block_size is None or ".mlp.experts." in name:
603
- return loaded_weight
604
- if not any(
605
- proj_name in name
606
- for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
607
- ):
608
- return loaded_weight
609
- dim = 1 if ".down_proj." in name else 0
610
- block_step = 1 if name.endswith("weight_scale_inv") else block_size[0]
611
- multiple = get_tensor_model_parallel_world_size() * block_step
612
- pad = (-loaded_weight.shape[dim]) % multiple
613
- if pad == 0:
614
- return loaded_weight
615
- pad_shape = list(loaded_weight.shape)
616
- pad_shape[dim] = pad
617
- return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
618
-
622
def _adapt_weights(
623
self, weights: Iterable[tuple[str, torch.Tensor]]
624
) -> Iterator[tuple[str, torch.Tensor]]:
@@ -667,7 +670,7 @@ class Dots3NoteModel(DeepseekV32Model):
670
)
671
yield f"{layer_prefix}.{projection}.weight", dequantized
672
continue
670
- yield name, self._pad_dense_mlp_weight(name, weight)
673
+ yield name, _pad_dense_mlp_weight(name, weight, self._weight_block_size)
674
if pending_indexer_fp8:
675
missing = ", ".join(
676
f"{prefix}.{proj}" for prefix, proj in pending_indexer_fp8
vllm/models/dots3_note/nvidia/mtp.py
+37
-35
@@ -9,10 +9,13 @@ import torch.nn as nn
9
10
from vllm.config import VllmConfig
11
from vllm.distributed import (
12
- get_tensor_model_parallel_world_size,
12
tensor_model_parallel_all_gather,
13
tensor_model_parallel_all_reduce,
14
)
15
+from vllm.model_executor.layers.fused_embed_norm import (
16
+ fused_embed_eh_norm,
17
+ has_full_vocab_on_rank,
18
+)
19
from vllm.model_executor.layers.layernorm import RMSNorm
20
from vllm.model_executor.layers.linear import ReplicatedLinear
21
from vllm.model_executor.layers.logits_processor import LogitsProcessor
@@ -25,7 +28,7 @@ from vllm.models.deepseek_v32.nvidia.mtp import (
28
DeepseekV32MultiTokenPredictor,
29
)
30
28
-from .model import Dots3NoteDecoderLayer
31
+from .model import Dots3NoteDecoderLayer, _pad_dense_mlp_weight
32
33
34
class Dots3NoteMultiTokenPredictorLayer(nn.Module):
@@ -59,18 +62,30 @@ class Dots3NoteMultiTokenPredictorLayer(nn.Module):
62
positions: torch.Tensor,
63
previous_hidden_states: torch.Tensor,
64
inputs_embeds: torch.Tensor | None = None,
65
+ embed_table: torch.Tensor | None = None,
66
spec_step_index: int = 0,
67
) -> torch.Tensor:
64
- del input_ids, spec_step_index
65
- assert inputs_embeds is not None
66
- eh_input = fused_eh_norm(
67
- positions,
68
- inputs_embeds,
69
- previous_hidden_states,
70
- self.enorm.weight,
71
- self.hnorm.weight,
72
- self.enorm.variance_epsilon,
73
- )
68
+ del spec_step_index
69
+ if embed_table is not None:
70
+ eh_input = fused_embed_eh_norm(
71
+ positions,
72
+ input_ids,
73
+ embed_table,
74
+ previous_hidden_states,
75
+ self.enorm.weight,
76
+ self.hnorm.weight,
77
+ self.enorm.variance_epsilon,
78
+ )
79
+ else:
80
+ assert inputs_embeds is not None
81
+ eh_input = fused_eh_norm(
82
+ positions,
83
+ inputs_embeds,
84
+ previous_hidden_states,
85
+ self.enorm.weight,
86
+ self.hnorm.weight,
87
+ self.enorm.variance_epsilon,
88
+ )
89
hidden_states = self.eh_proj(eh_input)[0]
90
hidden_states, residual = self.mtp_block(
91
positions=positions, hidden_states=hidden_states, residual=None
@@ -108,6 +123,7 @@ class Dots3NoteMultiTokenPredictor(DeepseekV32MultiTokenPredictor):
123
config.hidden_size,
124
prefix=maybe_prefix(prefix, "embed_tokens"),
125
)
126
+ self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens)
127
self.register_buffer(
128
"max_token_id",
129
torch.tensor(config.vocab_size - 1, dtype=torch.int64),
@@ -128,7 +144,7 @@ class Dots3NoteMultiTokenPredictor(DeepseekV32MultiTokenPredictor):
144
spec_step_idx: int = 0,
145
) -> torch.Tensor:
146
if inputs_embeds is None:
131
- inputs_embeds = self.embed_input_ids(input_ids)
147
+ input_ids = torch.minimum(input_ids, self.max_token_id)
148
return super().forward(
149
input_ids,
150
positions,
@@ -151,27 +167,6 @@ class Dots3NoteMTP(DeepseekV32MTP):
167
)
168
self.set_moe_parameters()
169
154
- def _pad_dense_mlp_weight(
155
- self, name: str, loaded_weight: torch.Tensor
156
- ) -> torch.Tensor:
157
- block_size = getattr(self.quant_config, "weight_block_size", None)
158
- if block_size is None or ".mlp.experts." in name:
159
- return loaded_weight
160
- if not any(
161
- proj_name in name
162
- for proj_name in (".gate_proj.", ".up_proj.", ".down_proj.")
163
- ):
164
- return loaded_weight
165
- dim = 1 if ".down_proj." in name else 0
166
- block_step = 1 if name.endswith("weight_scale_inv") else block_size[0]
167
- multiple = get_tensor_model_parallel_world_size() * block_step
168
- pad = (-loaded_weight.shape[dim]) % multiple
169
- if pad == 0:
170
- return loaded_weight
171
- pad_shape = list(loaded_weight.shape)
172
- pad_shape[dim] = pad
173
- return torch.cat([loaded_weight, loaded_weight.new_zeros(pad_shape)], dim=dim)
174
-
170
def _adapt_weights(
171
self, weights: Iterable[tuple[str, torch.Tensor]]
172
) -> Iterator[tuple[str, torch.Tensor]]:
@@ -183,7 +178,14 @@ class Dots3NoteMTP(DeepseekV32MTP):
178
f"model.layers.{mtp_layer}.embed_tokens.",
179
1,
180
)
186
- yield name, self._pad_dense_mlp_weight(name, weight)
181
+ yield (
182
+ name,
183
+ _pad_dense_mlp_weight(
184
+ name,
185
+ weight,
186
+ getattr(self.quant_config, "weight_block_size", None),
187
+ ),
188
+ )
189
190
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
191
return super().load_weights(self._adapt_weights(weights))
vllm/models/dots3_note/nvidia/vision.py
+7
-21
@@ -15,6 +15,7 @@ from transformers.modeling_utils import PreTrainedModel
15
from vllm.third_party.deep_gemm import per_block_cast_to_fp8
16
17
from .vision_attention import (
18
+ VisionRMSNorm,
19
VisionRotaryEmbedding,
20
VisionRotaryPositionEmbedding,
21
apply_vision_attention_residual,
@@ -107,23 +108,6 @@ class DotsMoEVitConfig(PretrainedConfig):
108
self.enable_fp8_moe = enable_fp8_moe
109
110
110
-class RMSNorm(nn.Module):
111
- def __init__(self, dim: int, eps: float = 1e-6):
112
- super().__init__()
113
- self.weight = nn.Parameter(torch.ones(dim))
114
- self.eps = eps
115
-
116
- def forward(self, x: torch.Tensor) -> torch.Tensor:
117
- output = self._norm(x.float()).type_as(x)
118
- return output * self.weight
119
-
120
- def extra_repr(self) -> str:
121
- return f"{tuple(self.weight.shape)}, eps={self.eps}"
122
-
123
- def _norm(self, x: torch.Tensor) -> torch.Tensor:
124
- return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
125
-
126
-
111
# ---- FFN modules ----
112
113
@@ -333,7 +317,7 @@ class DotsPatchEmbed(nn.Module):
317
kernel_size=(config.patch_size, config.patch_size),
318
stride=(config.patch_size, config.patch_size),
319
)
336
- self.norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
320
+ self.norm = VisionRMSNorm(config.embed_dim, eps=config.rms_norm_eps)
321
322
def forward(self, x: torch.Tensor) -> torch.Tensor:
323
x = x.view(
@@ -359,8 +343,8 @@ class MoEVisionBlock(nn.Module):
343
)
344
self.attn = build_vision_attention(attn_impl, config, eager_fallback="eager_v2")
345
self._attn_uses_seqlens = attn_uses_seqlens(attn_impl)
362
- self.norm_1 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
363
- self.norm_2 = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
346
+ self.norm_1 = VisionRMSNorm(config.embed_dim, eps=config.rms_norm_eps)
347
+ self.norm_2 = VisionRMSNorm(config.embed_dim, eps=config.rms_norm_eps)
348
349
is_moe = (
350
config.pyramid_num_routed
@@ -534,7 +518,9 @@ class DotsMoEVitModel(PreTrainedModel):
518
)
519
520
if config.post_norm:
537
- self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps)
521
+ self.post_trunk_norm = VisionRMSNorm(
522
+ config.embed_dim, eps=config.rms_norm_eps
523
+ )
524
525
adapter_cls = _ADAPTER_CLASSES.get(config.adapter_type)
526
if adapter_cls is None:
vllm/models/dots3_note/nvidia/vision_attention.py
+16
-6
@@ -17,8 +17,11 @@ import torch
17
import torch.nn.functional as F
18
from torch import nn
19
20
+from vllm.logger import init_logger
21
from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported
22
23
+logger = init_logger(__name__)
24
+
25
_flash_attn_available = is_fa_version_supported(2)
26
_flash_attn3_available = is_fa_version_supported(3)
27
_flash_attn3_varlen_func = flash_attn_varlen_func
@@ -107,8 +110,8 @@ class VisionRotaryEmbedding(nn.Module):
110
# ---------------------------------------------------------------------------
111
112
110
-class _RMSNorm(nn.Module):
111
- """Q/K norm inside attention; matches Dots ViT RMSNorm (fp32 reduce, then cast back)."""
113
+class VisionRMSNorm(nn.Module):
114
+ """Dots ViT RMSNorm with an fp32 reduction and input-dtype output."""
115
116
def __init__(self, dim: int, eps: float = 1e-6):
117
super().__init__()
@@ -119,6 +122,9 @@ class _RMSNorm(nn.Module):
122
output = self._norm(x.float()).type_as(x)
123
return output * self.weight
124
125
+ def extra_repr(self) -> str:
126
+ return f"{tuple(self.weight.shape)}, eps={self.eps}"
127
+
128
def _norm(self, x: torch.Tensor) -> torch.Tensor:
129
return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
130
@@ -156,8 +162,8 @@ class _VisionAttentionBase(nn.Module):
162
self.qkv = nn.Linear(params.dim, params.dim * 3, bias=params.bias)
163
self.proj = nn.Linear(params.dim, params.dim, bias=params.bias)
164
if self.use_qk_norm:
159
- self.q_norm = _RMSNorm(self.head_dim, eps=params.rms_norm_eps)
160
- self.k_norm = _RMSNorm(self.head_dim, eps=params.rms_norm_eps)
165
+ self.q_norm = VisionRMSNorm(self.head_dim, eps=params.rms_norm_eps)
166
+ self.k_norm = VisionRMSNorm(self.head_dim, eps=params.rms_norm_eps)
167
168
def _qkv_with_rope(
169
self,
@@ -385,12 +391,15 @@ def resolve_attn_implementation(
391
uses full-mask ``eager``.
392
"""
393
if attn_implementation == "flash_attention_3" and not _flash_attn3_available:
388
- print(
394
+ logger.warning_once(
395
"flash attention 3 not available! fallback to flash attention 2 implementation"
396
)
397
attn_implementation = "flash_attention_2"
398
if attn_implementation == "flash_attention_2" and not _flash_attn_available:
393
- print("flash attention 2 not available! fallback to eager implementation")
399
+ logger.warning_once(
400
+ "flash attention 2 not available! fallback to %s implementation",
401
+ eager_fallback,
402
+ )
403
attn_implementation = eager_fallback
404
return attn_implementation
405
@@ -481,6 +490,7 @@ __all__ = [
490
"VisionFlashAttention3",
491
"VisionRotaryEmbedding",
492
"VisionRotaryPositionEmbedding",
493
+ "VisionRMSNorm",
494
"VisionSdpaAttention",
495
"apply_rotary_pos_emb_vision",
496
"apply_vision_attention_residual",