# 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.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, 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 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.""" 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 self.use_sequence_parallel = 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 = False 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.use_sequence_parallel = False 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.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( 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 _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, _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 ) 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