From da0b2d8b17c9203bdd5eeb7a9af575f75a050c09 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E8=8C=83=E8=A3=95=E8=BE=BE?= Date: Mon, 31 Aug 2026 15:05:18 +0800 Subject: [PATCH] [Performance] Optimize Dots3 NOTE runtime (#53517) Signed-off-by: KurodaKanbei Signed-off-by: Jee Jee Li Co-authored-by: Jee Jee Li Co-authored-by: Jee Jee Li --- vllm/config/vllm.py | 2 + vllm/models/dots3_note/nvidia/attention.py | 8 ++- vllm/models/dots3_note/nvidia/model.py | 47 ++++++------ vllm/models/dots3_note/nvidia/mtp.py | 72 ++++++++++--------- vllm/models/dots3_note/nvidia/vision.py | 28 ++------ .../dots3_note/nvidia/vision_attention.py | 22 ++++-- 6 files changed, 94 insertions(+), 85 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 37304ec188..aadaccc82d 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -78,6 +78,8 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset( "DeepseekV32ForCausalLM", "DeepseekV4ForCausalLM", "DeepSeekV4MTPModel", + "Dots3NoteForCausalLM", + "Dots3NoteMTPModel", "GlmMoeDsaForCausalLM", "HYV4ForCausalLM", "HYV4MTPModel", diff --git a/vllm/models/dots3_note/nvidia/attention.py b/vllm/models/dots3_note/nvidia/attention.py index af17094d4d..cbd9ad7f4d 100644 --- a/vllm/models/dots3_note/nvidia/attention.py +++ b/vllm/models/dots3_note/nvidia/attention.py @@ -7,6 +7,7 @@ varlen MHA. Decode-only batches use the Triton absorbed-MQA kernel. """ from dataclasses import dataclass +from typing import ClassVar import torch @@ -19,7 +20,11 @@ from vllm.model_executor.layers.attention.mla_attention import ( 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.backend import ( + AttentionCGSupport, + AttentionLayer, + CommonAttentionMetadata, +) from vllm.v1.attention.backends.mla.flashattn_mla_sparse import ( FlashAttnMLASparseBackend, FlashAttnMLASparseImpl, @@ -307,6 +312,7 @@ class Dots3NoteFlashAttnPrefillBackend(FlashAttnPrefillBackend): class Dots3NoteMLAMetadataBuilder(TritonMLAMetadataBuilder): """Keep decode on MQA and route prefill/mixed batches through FA3.""" + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH query_len_support = QueryLenSupport.UNIFORM def __init__(self, kv_cache_spec, layer_names, vllm_config, device): diff --git a/vllm/models/dots3_note/nvidia/model.py b/vllm/models/dots3_note/nvidia/model.py index e713911e61..15128132be 100644 --- a/vllm/models/dots3_note/nvidia/model.py +++ b/vllm/models/dots3_note/nvidia/model.py @@ -19,6 +19,7 @@ from vllm.distributed import ( tensor_model_parallel_all_reduce, ) from vllm.model_executor.layers.attention import MLAAttention +from vllm.model_executor.layers.fused_embed_norm import has_full_vocab_on_rank from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, @@ -73,6 +74,28 @@ def _padded_mlp_size( return ((blocks + tp_size - 1) // tp_size) * block_size[0] * tp_size +def _pad_dense_mlp_weight( + name: str, + loaded_weight: torch.Tensor, + weight_block_size: list[int] | None, +) -> torch.Tensor: + if weight_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 weight_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) + + class Dots3NoteMoE(DeepseekV2MoE): """DeepSeek MoE with NOTE's block-padded shared expert.""" @@ -574,6 +597,7 @@ class Dots3NoteModel(DeepseekV32Model): ) else: self.embed_tokens = PPMissingLayer() + self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens) self.start_layer, self.end_layer, self.layers = make_layers( config.num_hidden_layers, lambda prefix: Dots3NoteDecoderLayer( @@ -595,27 +619,6 @@ class Dots3NoteModel(DeepseekV32Model): 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]]: @@ -667,7 +670,7 @@ class Dots3NoteModel(DeepseekV32Model): ) yield f"{layer_prefix}.{projection}.weight", dequantized continue - yield name, self._pad_dense_mlp_weight(name, weight) + yield name, _pad_dense_mlp_weight(name, weight, self._weight_block_size) if pending_indexer_fp8: missing = ", ".join( f"{prefix}.{proj}" for prefix, proj in pending_indexer_fp8 diff --git a/vllm/models/dots3_note/nvidia/mtp.py b/vllm/models/dots3_note/nvidia/mtp.py index d2e31dfe8e..dcc42b4e21 100644 --- a/vllm/models/dots3_note/nvidia/mtp.py +++ b/vllm/models/dots3_note/nvidia/mtp.py @@ -9,10 +9,13 @@ 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.fused_embed_norm import ( + fused_embed_eh_norm, + has_full_vocab_on_rank, +) 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 @@ -25,7 +28,7 @@ from vllm.models.deepseek_v32.nvidia.mtp import ( DeepseekV32MultiTokenPredictor, ) -from .model import Dots3NoteDecoderLayer +from .model import Dots3NoteDecoderLayer, _pad_dense_mlp_weight class Dots3NoteMultiTokenPredictorLayer(nn.Module): @@ -59,18 +62,30 @@ class Dots3NoteMultiTokenPredictorLayer(nn.Module): positions: torch.Tensor, previous_hidden_states: torch.Tensor, inputs_embeds: torch.Tensor | None = None, + embed_table: 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, - ) + del spec_step_index + if embed_table is not None: + eh_input = fused_embed_eh_norm( + positions, + input_ids, + embed_table, + previous_hidden_states, + self.enorm.weight, + self.hnorm.weight, + self.enorm.variance_epsilon, + ) + else: + 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 @@ -108,6 +123,7 @@ class Dots3NoteMultiTokenPredictor(DeepseekV32MultiTokenPredictor): config.hidden_size, prefix=maybe_prefix(prefix, "embed_tokens"), ) + self.replicated_embed = has_full_vocab_on_rank(self.embed_tokens) self.register_buffer( "max_token_id", torch.tensor(config.vocab_size - 1, dtype=torch.int64), @@ -128,7 +144,7 @@ class Dots3NoteMultiTokenPredictor(DeepseekV32MultiTokenPredictor): spec_step_idx: int = 0, ) -> torch.Tensor: if inputs_embeds is None: - inputs_embeds = self.embed_input_ids(input_ids) + input_ids = torch.minimum(input_ids, self.max_token_id) return super().forward( input_ids, positions, @@ -151,27 +167,6 @@ class Dots3NoteMTP(DeepseekV32MTP): ) 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]]: @@ -183,7 +178,14 @@ class Dots3NoteMTP(DeepseekV32MTP): f"model.layers.{mtp_layer}.embed_tokens.", 1, ) - yield name, self._pad_dense_mlp_weight(name, weight) + yield ( + name, + _pad_dense_mlp_weight( + name, + weight, + getattr(self.quant_config, "weight_block_size", None), + ), + ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return super().load_weights(self._adapt_weights(weights)) diff --git a/vllm/models/dots3_note/nvidia/vision.py b/vllm/models/dots3_note/nvidia/vision.py index 70ea31cb6f..4aba14d3c6 100644 --- a/vllm/models/dots3_note/nvidia/vision.py +++ b/vllm/models/dots3_note/nvidia/vision.py @@ -15,6 +15,7 @@ from transformers.modeling_utils import PreTrainedModel from vllm.third_party.deep_gemm import per_block_cast_to_fp8 from .vision_attention import ( + VisionRMSNorm, VisionRotaryEmbedding, VisionRotaryPositionEmbedding, apply_vision_attention_residual, @@ -107,23 +108,6 @@ class DotsMoEVitConfig(PretrainedConfig): 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 ---- @@ -333,7 +317,7 @@ class DotsPatchEmbed(nn.Module): 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) + self.norm = VisionRMSNorm(config.embed_dim, eps=config.rms_norm_eps) def forward(self, x: torch.Tensor) -> torch.Tensor: x = x.view( @@ -359,8 +343,8 @@ class MoEVisionBlock(nn.Module): ) 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) + self.norm_1 = VisionRMSNorm(config.embed_dim, eps=config.rms_norm_eps) + self.norm_2 = VisionRMSNorm(config.embed_dim, eps=config.rms_norm_eps) is_moe = ( config.pyramid_num_routed @@ -534,7 +518,9 @@ class DotsMoEVitModel(PreTrainedModel): ) if config.post_norm: - self.post_trunk_norm = RMSNorm(config.embed_dim, eps=config.rms_norm_eps) + self.post_trunk_norm = VisionRMSNorm( + config.embed_dim, eps=config.rms_norm_eps + ) adapter_cls = _ADAPTER_CLASSES.get(config.adapter_type) if adapter_cls is None: diff --git a/vllm/models/dots3_note/nvidia/vision_attention.py b/vllm/models/dots3_note/nvidia/vision_attention.py index 4d5821094b..bc81297c59 100644 --- a/vllm/models/dots3_note/nvidia/vision_attention.py +++ b/vllm/models/dots3_note/nvidia/vision_attention.py @@ -17,8 +17,11 @@ import torch import torch.nn.functional as F from torch import nn +from vllm.logger import init_logger from vllm.vllm_flash_attn import flash_attn_varlen_func, is_fa_version_supported +logger = init_logger(__name__) + _flash_attn_available = is_fa_version_supported(2) _flash_attn3_available = is_fa_version_supported(3) _flash_attn3_varlen_func = flash_attn_varlen_func @@ -107,8 +110,8 @@ class VisionRotaryEmbedding(nn.Module): # --------------------------------------------------------------------------- -class _RMSNorm(nn.Module): - """Q/K norm inside attention; matches Dots ViT RMSNorm (fp32 reduce, then cast back).""" +class VisionRMSNorm(nn.Module): + """Dots ViT RMSNorm with an fp32 reduction and input-dtype output.""" def __init__(self, dim: int, eps: float = 1e-6): super().__init__() @@ -119,6 +122,9 @@ class _RMSNorm(nn.Module): 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) @@ -156,8 +162,8 @@ class _VisionAttentionBase(nn.Module): 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) + self.q_norm = VisionRMSNorm(self.head_dim, eps=params.rms_norm_eps) + self.k_norm = VisionRMSNorm(self.head_dim, eps=params.rms_norm_eps) def _qkv_with_rope( self, @@ -385,12 +391,15 @@ def resolve_attn_implementation( uses full-mask ``eager``. """ if attn_implementation == "flash_attention_3" and not _flash_attn3_available: - print( + logger.warning_once( "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") + logger.warning_once( + "flash attention 2 not available! fallback to %s implementation", + eager_fallback, + ) attn_implementation = eager_fallback return attn_implementation @@ -481,6 +490,7 @@ __all__ = [ "VisionFlashAttention3", "VisionRotaryEmbedding", "VisionRotaryPositionEmbedding", + "VisionRMSNorm", "VisionSdpaAttention", "apply_rotary_pos_emb_vision", "apply_vision_attention_residual",