iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Performance] Optimize Dots3 NOTE runtime (#53517)

Browse Source 
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>
This commit is contained in:
范裕达
2026-08-31 00:05:18 -07:00
committed by GitHub
co-authored by Jee Jee Li Jee Jee Li
parent 7292ee2791
commit da0b2d8b17
6 changed files with 94 additions and 85 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/config/vllm.py
+2
View File
@@ -78,6 +78,8 @@ DEFAULT_BREAKABLE_CUDAGRAPH_ARCHITECTURES = frozenset(
"DeepseekV32ForCausalLM",
"DeepseekV4ForCausalLM",
"DeepSeekV4MTPModel",
"Dots3NoteForCausalLM",
"Dots3NoteMTPModel",
"GlmMoeDsaForCausalLM",
"HYV4ForCausalLM",
"HYV4MTPModel",
vllm/models/dots3_note/nvidia/attention.py
+7 -1
View File
@@ -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):
vllm/models/dots3_note/nvidia/model.py
+25 -22
View File
@@ -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
vllm/models/dots3_note/nvidia/mtp.py
+37 -35
View File
@@ -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))
vllm/models/dots3_note/nvidia/vision.py
+7 -21
View File
@@ -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:
vllm/models/dots3_note/nvidia/vision_attention.py
+16 -6
View File
@@ -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",