From 3fbd79af64721578d6639bac2d35871e73932fb9 Mon Sep 17 00:00:00 2001 From: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Date: Wed, 11 Mar 2026 23:16:02 -0700 Subject: [PATCH 1/3] [None][feat] Add AD custom model for DeepSeek-V3.2 family Add AutoDeploy custom model for the DeepSeek-V3.2 model family (deepseek-ai/DeepSeek-V3.2, DeepSeek-V3.2-Speciale, nvidia/DeepSeek-V3.2-NVFP4). The core forward path (MLA + MoE + RMSNorm + YaRN RoPE) is identical to DeepSeek-V3. V3.2 adds: - Indexer module (sparse attention selection) in every layer - included for weight loading but not used in forward (no-op for seq <= index_topk=2048) - MTP (Multi-Token Prediction) layers - extra decoder layers included for weight loading but not used for main logits Config subclasses DeepseekV3Config from transformers, adding V3.2-specific fields (index_head_dim, index_n_heads, index_topk, num_nextn_predict_layers). Also fixes FP8 tensor indexing in the shared MLA RoPE de-interleaving hook (mla_rope_utils.py) which is needed for FP8-quantized checkpoints. Files: - modeling_deepseek_v32.py: Custom model implementation - __init__.py: Registration - mla_rope_utils.py: FP8 support for de-interleaving - test_deepseek_v32_modeling.py: Hierarchical tests + HF V3 equivalence Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> --- .../auto_deploy/models/custom/__init__.py | 2 + .../models/custom/mla_rope_utils.py | 15 +- .../models/custom/modeling_deepseek_v32.py | 755 ++++++++++++++ .../models/test_deepseek_v32_modeling.py | 944 ++++++++++++++++++ 4 files changed, 1712 insertions(+), 4 deletions(-) create mode 100644 tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py create mode 100644 tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py index 93e56cb1bfb5..67c36e6cf6ea 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py @@ -1,6 +1,7 @@ from .modeling_decilm import DeciLMForCausalLM from .modeling_deepseek import DeepSeekV3ForCausalLM from .modeling_deepseek_v2 import DeepSeekV2ForCausalLM +from .modeling_deepseek_v32 import DeepSeekV32ForCausalLM from .modeling_glm4_moe_lite import Glm4MoeLiteForCausalLM from .modeling_granite_moe_hybrid import GraniteMoeHybridForCausalLM from .modeling_hunyuan_dense_v1 import HunYuanDenseV1ForCausalLM @@ -16,6 +17,7 @@ "DeciLMForCausalLM", "DeepSeekV2ForCausalLM", "DeepSeekV3ForCausalLM", + "DeepSeekV32ForCausalLM", "Glm4MoeLiteForCausalLM", "HunYuanDenseV1ForCausalLM", "GraniteMoeHybridForCausalLM", diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/mla_rope_utils.py b/tensorrt_llm/_torch/auto_deploy/models/custom/mla_rope_utils.py index 55e78588309b..87e3551399c9 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/mla_rope_utils.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/mla_rope_utils.py @@ -3,7 +3,7 @@ """Shared MLA RoPE utilities for auto_deploy custom models. Contains helper functions for RoPE weight de-interleaving, -used by DeepSeek V3 and GLM4 MoE Lite model implementations. +used by DeepSeek V3, V3.2, and GLM4 MoE Lite model implementations. """ from typing import Dict @@ -11,6 +11,13 @@ import torch +def _permute_with_fp8_support(tensor: torch.Tensor, perm: torch.Tensor, dim: int) -> torch.Tensor: + """Index-select along a dimension, casting through float32 for FP8 tensors.""" + if tensor.dtype in (torch.float8_e4m3fn, torch.float8_e5m2): + return tensor.to(torch.float32).index_select(dim, perm).to(tensor.dtype) + return tensor.index_select(dim, perm) + + def _rope_deinterleave_load_hook( state_dict: Dict[str, torch.Tensor], prefix: str, @@ -43,7 +50,7 @@ def _rope_deinterleave_load_hook( w = w.view(num_heads, qk_head_dim, -1) w_nope = w[:, :qk_nope_head_dim, :] w_rope = w[:, qk_nope_head_dim:, :] - w_rope = w_rope[:, perm, :] + w_rope = _permute_with_fp8_support(w_rope, perm, dim=1) w = torch.cat([w_nope, w_rope], dim=1) state_dict[q_key] = w.view(-1, w.shape[-1]) @@ -53,7 +60,7 @@ def _rope_deinterleave_load_hook( w = state_dict[kv_key] w_kv = w[:kv_lora_rank, :] w_pe = w[kv_lora_rank:, :] - w_pe = w_pe[perm, :] + w_pe = _permute_with_fp8_support(w_pe, perm, dim=0) state_dict[kv_key] = torch.cat([w_kv, w_pe], dim=0) # --- kv_a_proj_with_mqa.bias (if present) --- @@ -62,5 +69,5 @@ def _rope_deinterleave_load_hook( b = state_dict[kv_bias_key] b_kv = b[:kv_lora_rank] b_pe = b[kv_lora_rank:] - b_pe = b_pe[perm] + b_pe = _permute_with_fp8_support(b_pe, perm, dim=0) state_dict[kv_bias_key] = torch.cat([b_kv, b_pe]) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py new file mode 100644 index 000000000000..c10da8171e06 --- /dev/null +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py @@ -0,0 +1,755 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +"""Slimmed down PyTorch DeepSeekV32 model implementation for auto_deploy export. + +Source: +https://huggingface.co/deepseek-ai/DeepSeek-V3.2 + +This implementation differs from the original in the following ways: +* Simplified for prefill-only inference (no KV caching) +* Uses auto_deploy custom ops for export compatibility +* Removed flash attention variants (uses torch_mla custom op) +* Removed gradient checkpointing and training code paths +* Removed attention dropout (inference only) +* Indexer module included for weight loading but not used in forward (no-op for seq <= index_topk) +* MTP (Multi-Token Prediction) layers included for weight loading but not used for main logits + +This allows us to have a clean export-ready implementation with auto_deploy custom ops. +""" + +import math +from dataclasses import dataclass +from functools import partial +from typing import Optional, Tuple + +import torch +import torch.nn.functional as F +from torch import nn +from transformers import AutoConfig +from transformers.activations import ACT2FN +from transformers.generation import GenerationMixin +from transformers.modeling_utils import PreTrainedModel +from transformers.models.deepseek_v3.configuration_deepseek_v3 import ( + DeepseekV3Config as _BaseDeepseekV3Config, +) +from transformers.utils import ModelOutput + +from tensorrt_llm._torch.auto_deploy.models.custom import mla_rope_utils +from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory +from tensorrt_llm._torch.utils import ActivationType + + +class DeepseekV32Config(_BaseDeepseekV3Config): + """Configuration class for DeepSeek-V3.2 models. + + Extends DeepseekV3Config with V3.2-specific fields: Indexer (sparse attention + selection) and MTP (multi-token prediction) layers. + """ + + model_type = "deepseek_v32" + + def __init__( + self, + index_head_dim=128, + index_n_heads=64, + index_topk=2048, + num_nextn_predict_layers=1, + ep_size=1, + **kwargs, + ): + super().__init__(**kwargs) + self.index_head_dim = index_head_dim + self.index_n_heads = index_n_heads + self.index_topk = index_topk + self.num_nextn_predict_layers = num_nextn_predict_layers + self.ep_size = ep_size + + +AutoConfig.register("deepseek_v32", DeepseekV32Config, exist_ok=True) + + +class DeepSeekV32RMSNorm(nn.Module): + """RMS Normalization for DeepSeekV32.""" + + def __init__(self, hidden_size: int, eps: float = 1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return torch.ops.auto_deploy.torch_rmsnorm( + hidden_states, self.weight, self.variance_epsilon + ).to(hidden_states.dtype) + + +class DeepSeekV32RotaryEmbedding(nn.Module): + """Rotary Position Embedding for DeepSeekV32. + + Simplified version that precomputes and caches cos/sin values. + Returns full cached values (not sliced by seq_len) to enable export. + """ + + def __init__(self, dim: int, max_position_embeddings: int = 2048, base: float = 10000.0): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float() / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + # Build cos/sin cache + self._set_cos_sin_cache(max_position_embeddings) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + t = torch.arange(seq_len, dtype=self.inv_freq.dtype) + freqs = torch.outer(t, self.inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("_ad_cos_cached", emb.cos(), persistent=False) + self.register_buffer("_ad_sin_cached", emb.sin(), persistent=False) + + def forward( + self, x: torch.Tensor, seq_len: Optional[int] = None + ) -> Tuple[torch.Tensor, torch.Tensor]: + return ( + self._ad_cos_cached.to(dtype=x.dtype, device=x.device), + self._ad_sin_cached.to(dtype=x.dtype, device=x.device), + ) + + +class DeepSeekV32YarnRotaryEmbedding(DeepSeekV32RotaryEmbedding): + """YaRN-extended rotary embedding for DeepSeekV32.""" + + def __init__( + self, + dim: int, + max_position_embeddings: int = 2048, + base: float = 10000.0, + scaling_factor: float = 1.0, + original_max_position_embeddings: int = 4096, + beta_fast: int = 32, + beta_slow: int = 1, + mscale: float = 1.0, + mscale_all_dim: float = 0.0, + ): + self.scaling_factor = scaling_factor + self.original_max_position_embeddings = original_max_position_embeddings + self.beta_fast = beta_fast + self.beta_slow = beta_slow + self.mscale = mscale + self.mscale_all_dim = mscale_all_dim + super().__init__(dim, max_position_embeddings, base) + + def _set_cos_sin_cache(self, seq_len: int): + self.max_seq_len_cached = seq_len + dim = self.dim + + freq_extra = 1.0 / (self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim)) + freq_inter = 1.0 / ( + self.scaling_factor * self.base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + + low, high = self._yarn_find_correction_range( + self.beta_fast, + self.beta_slow, + dim, + self.base, + self.original_max_position_embeddings, + ) + inv_freq_mask = 1.0 - self._yarn_linear_ramp_mask(low, high, dim // 2) + inv_freq = freq_inter * (1 - inv_freq_mask) + freq_extra * inv_freq_mask + self.register_buffer("inv_freq", inv_freq, persistent=False) + + t = torch.arange(seq_len, dtype=torch.float32) + freqs = torch.outer(t, inv_freq) + + _mscale = float( + self._yarn_get_mscale(self.scaling_factor, self.mscale) + / self._yarn_get_mscale(self.scaling_factor, self.mscale_all_dim) + ) + + emb = torch.cat((freqs, freqs), dim=-1) + self.register_buffer("_ad_cos_cached", (emb.cos() * _mscale), persistent=False) + self.register_buffer("_ad_sin_cached", (emb.sin() * _mscale), persistent=False) + + @staticmethod + def _yarn_find_correction_dim( + num_rotations: float, dim: int, base: float = 10000, max_position_embeddings: int = 2048 + ) -> float: + return (dim * math.log(max_position_embeddings / (num_rotations * 2 * math.pi))) / ( + 2 * math.log(base) + ) + + def _yarn_find_correction_range( + self, low_rot: int, high_rot: int, dim: int, base: float, max_position_embeddings: int + ) -> Tuple[int, int]: + low = math.floor( + self._yarn_find_correction_dim(low_rot, dim, base, max_position_embeddings) + ) + high = math.ceil( + self._yarn_find_correction_dim(high_rot, dim, base, max_position_embeddings) + ) + return max(low, 0), min(high, dim - 1) + + @staticmethod + def _yarn_get_mscale(scale: float = 1.0, mscale: float = 1.0) -> float: + if scale <= 1: + return 1.0 + return 0.1 * mscale * math.log(scale) + 1.0 + + @staticmethod + def _yarn_linear_ramp_mask(min_val: float, max_val: float, dim: int) -> torch.Tensor: + if min_val == max_val: + max_val += 0.001 + linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (max_val - min_val) + return torch.clamp(linear_func, 0, 1) + + +class DeepSeekV32MLP(nn.Module): + """MLP layer for DeepSeekV32 (SwiGLU activation).""" + + def __init__( + self, config, hidden_size: Optional[int] = None, intermediate_size: Optional[int] = None + ): + super().__init__() + self.config = config + self.hidden_size = hidden_size or config.hidden_size + self.intermediate_size = intermediate_size or config.intermediate_size + + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + self.act_fn = ACT2FN[config.hidden_act] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class DeepSeekV32MoEGate(nn.Module): + """MoE Gating for DeepSeekV32 with noaux_tc top-k selection.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.top_k = config.num_experts_per_tok + self.n_routed_experts = config.n_routed_experts + self.routed_scaling_factor = config.routed_scaling_factor + self.n_group = config.n_group + self.topk_group = config.topk_group + + self.weight = nn.Parameter( + torch.empty((self.n_routed_experts, config.hidden_size), dtype=torch.float32) + ) + self.register_buffer( + "e_score_correction_bias", + torch.zeros(self.n_routed_experts, dtype=torch.float32), + ) + self.reset_parameters() + + def reset_parameters(self) -> None: + torch.nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) + + def forward(self, hidden_states: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + bsz, seq_len, hidden_dim = hidden_states.shape + hidden_states_flat = hidden_states.view(-1, hidden_dim) + + # Router logits via plain linear (CPU-compatible) + router_logits = F.linear(hidden_states_flat.float(), self.weight.float()) + + # Sigmoid scoring + scores = router_logits.sigmoid() + original_scores = scores + + # Apply bias correction + scores = scores + self.e_score_correction_bias + + # Group-limited top-k selection + if self.n_group > 1: + scores = scores.view(-1, self.n_group, self.n_routed_experts // self.n_group) + group_scores = scores.topk(2, dim=-1)[0].sum(dim=-1) + group_indices = group_scores.topk(self.topk_group, dim=-1)[1] + mask = torch.ones_like(group_scores, dtype=torch.bool).scatter_(1, group_indices, False) + scores = scores.masked_fill_(mask.unsqueeze(-1), float("-inf")).flatten(1) + + topk_indices = scores.topk(self.top_k, dim=-1)[1] + topk_weights = original_scores.gather(1, topk_indices) + + # Normalize and scale + topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True) + topk_weights = topk_weights * self.routed_scaling_factor + + return topk_indices, topk_weights + + +class DeepSeekV32MoE(nn.Module): + """Mixture of Experts layer for DeepSeekV32.""" + + def __init__(self, config): + super().__init__() + self.config = config + self.num_experts_per_tok = config.num_experts_per_tok + + self.experts = nn.ModuleList( + [ + DeepSeekV32MLP(config, intermediate_size=config.moe_intermediate_size) + for _ in range(config.n_routed_experts) + ] + ) + + self.gate = DeepSeekV32MoEGate(config) + + if config.n_shared_experts is not None: + intermediate_size = config.moe_intermediate_size * config.n_shared_experts + self.shared_experts = DeepSeekV32MLP(config, intermediate_size=intermediate_size) + else: + self.shared_experts = None + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + identity = hidden_states + orig_shape = hidden_states.shape + + selected_experts, routing_weights = self.gate(hidden_states) + + final_hidden_states = torch.ops.auto_deploy.torch_moe( + hidden_states.view(-1, hidden_states.shape[-1]), + selected_experts, + routing_weights, + w1_weight=[expert.gate_proj.weight for expert in self.experts], + w2_weight=[expert.down_proj.weight for expert in self.experts], + w3_weight=[expert.up_proj.weight for expert in self.experts], + is_gated_mlp=True, + act_fn=int(ActivationType.Silu), + ) + + final_hidden_states = final_hidden_states.view(*orig_shape) + + if self.shared_experts is not None: + final_hidden_states = final_hidden_states + self.shared_experts(identity) + + return final_hidden_states.to(hidden_states.dtype) + + +class DeepSeekV32Indexer(nn.Module): + """Indexer module for sparse attention selection in DeepSeekV32. + + This module exists for weight loading compatibility. It is NOT used in the + forward pass — the torch_mla canonical op handles causal attention without + additional masking. For sequences <= index_topk (2048), the Indexer is + effectively a no-op in the original model. + """ + + def __init__(self, config): + super().__init__() + self.wq_b = nn.Linear( + config.q_lora_rank, config.index_n_heads * config.index_head_dim, bias=False + ) + self.wk = nn.Linear(config.hidden_size, config.index_head_dim, bias=False) + self.k_norm = nn.LayerNorm(config.index_head_dim) + self.weights_proj = nn.Linear(config.hidden_size, config.index_n_heads, bias=False) + + +class DeepSeekV32Attention(nn.Module): + """Multi-head Latent Attention (MLA) for DeepSeekV32. + + Uses compressed KV representation with latent projections. + """ + + def __init__(self, config, layer_idx: Optional[int] = None): + super().__init__() + self.config = config + self.layer_idx = layer_idx + + self.hidden_size = config.hidden_size + self.num_heads = config.num_attention_heads + self.q_lora_rank = config.q_lora_rank + self.kv_lora_rank = config.kv_lora_rank + self.qk_nope_head_dim = config.qk_nope_head_dim + self.qk_rope_head_dim = config.qk_rope_head_dim + self.v_head_dim = config.v_head_dim + self.q_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim + + self.max_position_embeddings = config.max_position_embeddings + self.rope_theta = config.rope_theta + + # Q projection (with optional LoRA) + if self.q_lora_rank is None: + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.q_head_dim, bias=False) + else: + self.q_a_proj = nn.Linear( + self.hidden_size, config.q_lora_rank, bias=config.attention_bias + ) + self.q_a_layernorm = DeepSeekV32RMSNorm(config.q_lora_rank) + self.q_b_proj = nn.Linear( + config.q_lora_rank, self.num_heads * self.q_head_dim, bias=False + ) + + # KV projection with MQA + self.kv_a_proj_with_mqa = nn.Linear( + self.hidden_size, + self.kv_lora_rank + self.qk_rope_head_dim, + bias=config.attention_bias, + ) + self.kv_a_layernorm = DeepSeekV32RMSNorm(self.kv_lora_rank) + self.kv_b_proj = nn.Linear( + self.kv_lora_rank, + self.num_heads * (self.qk_nope_head_dim + self.v_head_dim), + bias=False, + ) + + # Output projection + self.o_proj = nn.Linear( + self.num_heads * self.v_head_dim, self.hidden_size, bias=config.attention_bias + ) + + # Indexer (for weight loading — not used in forward) + self.indexer = DeepSeekV32Indexer(config) + + # Initialize rotary embedding + self._init_rope() + + # Softmax scale + self.softmax_scale = self.q_head_dim ** (-0.5) + if config.rope_scaling is not None: + mscale_all_dim = config.rope_scaling.get("mscale_all_dim", 0) + scaling_factor = config.rope_scaling["factor"] + if mscale_all_dim: + mscale = DeepSeekV32YarnRotaryEmbedding._yarn_get_mscale( + scaling_factor, mscale_all_dim + ) + self.softmax_scale = self.softmax_scale * mscale * mscale + + def _init_rope(self): + if self.config.rope_scaling is None: + self.rotary_emb = DeepSeekV32RotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + else: + scaling_type = self.config.rope_scaling["type"] + scaling_factor = self.config.rope_scaling["factor"] + + if scaling_type == "yarn": + kwargs = { + key: self.config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in self.config.rope_scaling + } + self.rotary_emb = DeepSeekV32YarnRotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + scaling_factor=scaling_factor, + base=self.rope_theta, + **kwargs, + ) + else: + self.rotary_emb = DeepSeekV32RotaryEmbedding( + self.qk_rope_head_dim, + max_position_embeddings=self.max_position_embeddings, + base=self.rope_theta, + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + bsz, q_len, _ = hidden_states.size() + + # Q projection + if self.q_lora_rank is None: + q = self.q_proj(hidden_states) + else: + q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states))) + + # Shape: [B, S, N, q_head_dim] (BSND layout) + q = q.view(bsz, q_len, self.num_heads, self.q_head_dim) + q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + + # KV projection - keep compressed form + kv_a_output = self.kv_a_proj_with_mqa(hidden_states) + compressed_kv, k_pe = torch.split( + kv_a_output, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + + compressed_kv = self.kv_a_layernorm(compressed_kv) + + # k_pe: [B, S, 1, qk_rope_head_dim] (BSND layout, shared across heads) + k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) + + kv_seq_len = q_len + + cos, sin = self.rotary_emb(hidden_states, seq_len=kv_seq_len) + cos = cos[position_ids] + sin = sin[position_ids] + + # Apply RoPE using custom op (weights pre-permuted to NeoX format at load time) + q_pe_rotated, kpe = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin( + q_pe, + k_pe, + cos, + sin, + 2, # unsqueeze_dim=2 for BSND layout + ) + + # Call MLA with compressed KV + attn_output = torch.ops.auto_deploy.torch_mla( + q_nope, + q_pe_rotated, + compressed_kv, + kpe, + self.kv_b_proj.weight, + True, # is_causal + self.softmax_scale, + "bsnd", + ) + + # Output: [B, S, N, v_head_dim] -> [B, S, N * v_head_dim] + attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim) + attn_output = self.o_proj(attn_output) + + return attn_output + + +class DeepSeekV32DecoderLayer(nn.Module): + """Transformer decoder layer for DeepSeekV32.""" + + def __init__(self, config, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + # Attention + self.self_attn = DeepSeekV32Attention(config, layer_idx=layer_idx) + + # MLP or MoE + use_moe = ( + config.n_routed_experts is not None + and layer_idx >= config.first_k_dense_replace + and layer_idx % config.moe_layer_freq == 0 + ) + if use_moe: + self.mlp = DeepSeekV32MoE(config) + else: + self.mlp = DeepSeekV32MLP(config) + + # Layer norms + self.input_layernorm = DeepSeekV32RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = DeepSeekV32RMSNorm( + config.hidden_size, eps=config.rms_norm_eps + ) + + def forward( + self, + hidden_states: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + # Self attention + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = residual + hidden_states + + # MLP/MoE + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class DeepSeekV32MTPSharedHead(nn.Module): + """Shared prediction head for MTP layers (weight loading only).""" + + def __init__(self, config): + super().__init__() + self.norm = DeepSeekV32RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + +class DeepSeekV32MTPDecoderLayer(DeepSeekV32DecoderLayer): + """MTP decoder layer with additional modules for multi-token prediction. + + Extends the standard decoder layer with embed_tokens, eh_proj, enorm, hnorm, + and shared_head. These are included for weight loading compatibility. + In the forward pass, this layer operates as a standard decoder layer. + """ + + def __init__(self, config, layer_idx: int): + super().__init__(config, layer_idx=layer_idx) + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size) + self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False) + self.enorm = DeepSeekV32RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.hnorm = DeepSeekV32RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.shared_head = DeepSeekV32MTPSharedHead(config) + + +@dataclass +class DeepSeekV32Output(ModelOutput): + """Output for DeepSeekV32Model.""" + + last_hidden_state: Optional[torch.FloatTensor] = None + + +@dataclass +class DeepSeekV32CausalLMOutput(ModelOutput): + """Output for DeepSeekV32ForCausalLM.""" + + logits: Optional[torch.FloatTensor] = None + + +class DeepSeekV32PreTrainedModel(PreTrainedModel): + """Base class for DeepSeekV32 models.""" + + config_class = DeepseekV32Config + base_model_prefix = "model" + _no_split_modules = ["DeepSeekV32DecoderLayer", "DeepSeekV32MTPDecoderLayer"] + supports_gradient_checkpointing = False + + def _init_weights(self, module): + std = self.config.initializer_range + if isinstance(module, nn.Linear): + 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 DeepSeekV32Model(DeepSeekV32PreTrainedModel): + """DeepSeekV32 transformer decoder model.""" + + def __init__(self, config): + super().__init__(config) + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + + # Main decoder layers + layers = [ + DeepSeekV32DecoderLayer(config, layer_idx=idx) + for idx in range(config.num_hidden_layers) + ] + # MTP layers (for weight loading — not used in main forward path) + num_mtp = getattr(config, "num_nextn_predict_layers", 0) or 0 + for i in range(num_mtp): + layers.append( + DeepSeekV32MTPDecoderLayer(config, layer_idx=config.num_hidden_layers + i) + ) + self.layers = nn.ModuleList(layers) + + self.norm = DeepSeekV32RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + self.post_init() + + def get_input_embeddings(self): + return self.embed_tokens + + def set_input_embeddings(self, value): + self.embed_tokens = value + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> DeepSeekV32Output: + assert position_ids is not None, "position_ids is required for DeepSeekV32Model" + + if input_ids is not None and inputs_embeds is not None: + raise ValueError("Cannot specify both input_ids and inputs_embeds") + elif input_ids is None and inputs_embeds is None: + raise ValueError("Must specify either input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + + # Only process main decoder layers (skip MTP layers) + for idx in range(self.config.num_hidden_layers): + hidden_states = self.layers[idx](hidden_states, position_ids) + + hidden_states = self.norm(hidden_states) + + return DeepSeekV32Output(last_hidden_state=hidden_states) + + +class DeepSeekV32ForCausalLM(DeepSeekV32PreTrainedModel, GenerationMixin): + """DeepSeekV32 model with language modeling head.""" + + _tied_weights_keys = ["lm_head.weight"] + + def __init__(self, config): + super().__init__(config) + self.model = DeepSeekV32Model(config) + self.vocab_size = config.vocab_size + self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) + + # Pre-permute RoPE weight rows from interleaved to NeoX format at load time + self._register_load_state_dict_pre_hook( + partial( + mla_rope_utils._rope_deinterleave_load_hook, + qk_rope_head_dim=config.qk_rope_head_dim, + qk_nope_head_dim=config.qk_nope_head_dim, + num_heads=config.num_attention_heads, + kv_lora_rank=config.kv_lora_rank, + num_layers=config.num_hidden_layers + + (getattr(config, "num_nextn_predict_layers", 0) or 0), + ) + ) + + self.post_init() + + def get_input_embeddings(self): + return self.model.embed_tokens + + def set_input_embeddings(self, value): + self.model.embed_tokens = value + + def get_output_embeddings(self): + return self.lm_head + + def set_output_embeddings(self, new_embeddings): + self.lm_head = new_embeddings + + def get_decoder(self): + return self.model + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + position_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.FloatTensor] = None, + **kwargs, + ) -> DeepSeekV32CausalLMOutput: + assert position_ids is not None, "position_ids is required for DeepSeekV32ForCausalLM" + + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + logits = self.lm_head(hidden_states).float() + + return DeepSeekV32CausalLMOutput(logits=logits) + + +# Register with AutoModelForCausalLMFactory +AutoModelForCausalLMFactory.register_custom_model_cls("DeepseekV32Config", DeepSeekV32ForCausalLM) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py new file mode 100644 index 000000000000..3b5ee5447172 --- /dev/null +++ b/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py @@ -0,0 +1,944 @@ +"""Tests for DeepSeekV32 custom model implementation for auto_deploy export. + +Covers the DeepSeek-V3.2 model family (model_type="deepseek_v32"): +- deepseek-ai/DeepSeek-V3.2 +- deepseek-ai/DeepSeek-V3.2-Speciale +- nvidia/DeepSeek-V3.2-NVFP4 + +The core forward path (MLA + MoE + RMSNorm + RoPE) is identical to DeepSeek-V3. +V3.2 adds an Indexer module (for sparse attention) and MTP layers (multi-token +prediction), both included for weight loading but not used in the prefill forward path. +Numerical equivalence is tested against the HF DeepSeek-V3 reference since the +forward math is identical. +""" + +import pytest +import torch +from torch.export import Dim + +from tensorrt_llm._torch.auto_deploy.export import torch_export_to_gm +from tensorrt_llm._torch.auto_deploy.models.custom.modeling_deepseek_v32 import ( + DeepSeekV32Attention, + DeepseekV32Config, + DeepSeekV32DecoderLayer, + DeepSeekV32ForCausalLM, + DeepSeekV32Indexer, + DeepSeekV32MLP, + DeepSeekV32MoE, + DeepSeekV32MTPDecoderLayer, + DeepSeekV32MTPSharedHead, + DeepSeekV32RMSNorm, + DeepSeekV32RotaryEmbedding, + DeepSeekV32YarnRotaryEmbedding, +) +from tensorrt_llm._torch.auto_deploy.utils._graph import move_to_device + +_BATCH_AND_SEQUENCE_TEST_CASES = ((2, 6), (1, 8)) + + +@pytest.fixture(scope="function", autouse=True) +def set_seed(): + torch.manual_seed(42) + + +# ============================================================================= +# Config helpers +# ============================================================================= + + +def _create_config(**kwargs): + """Create a small DeepSeekV32 config for testing.""" + defaults = dict( + vocab_size=1000, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=512, + rms_norm_eps=1e-6, + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=32, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + first_k_dense_replace=1, + moe_layer_freq=1, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + pad_token_id=0, + index_head_dim=16, + index_n_heads=4, + index_topk=64, + num_nextn_predict_layers=1, + ) + defaults.update(kwargs) + return DeepseekV32Config(**defaults) + + +def _create_config_no_mtp(**kwargs): + """Create config without MTP layers.""" + return _create_config(num_nextn_predict_layers=0, **kwargs) + + +# ============================================================================= +# RMSNorm Tests +# ============================================================================= + + +class TestDeepSeekV32RMSNorm: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def test_forward_shape(self): + hidden_size = 64 + norm = DeepSeekV32RMSNorm(hidden_size).to(self.device, self.dtype) + x = torch.randn(2, 4, hidden_size, dtype=self.dtype, device=self.device) + output = norm(x) + assert output.shape == x.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_output_normalized(self): + hidden_size = 64 + norm = DeepSeekV32RMSNorm(hidden_size).to(self.device, torch.float32) + x = torch.randn(2, 4, hidden_size, dtype=torch.float32, device=self.device) + output = norm(x) + rms = torch.sqrt((output**2).mean(-1)) + assert torch.allclose(rms, torch.ones_like(rms), atol=0.1) + + +# ============================================================================= +# Rotary Embedding Tests +# ============================================================================= + + +class TestDeepSeekV32RotaryEmbedding: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def test_base_rope_shape(self): + dim = 8 + max_pos = 512 + rope = DeepSeekV32RotaryEmbedding(dim, max_pos).to(self.device) + x = torch.randn(2, 4, 4, dim, dtype=self.dtype, device=self.device) + cos, sin = rope(x) + assert cos.shape == (max_pos, dim) + assert sin.shape == (max_pos, dim) + + def test_yarn_rope_shape(self): + dim = 8 + max_pos = 512 + rope = DeepSeekV32YarnRotaryEmbedding( + dim, + max_pos, + scaling_factor=2.0, + original_max_position_embeddings=256, + ).to(self.device) + x = torch.randn(2, 4, 4, dim, dtype=self.dtype, device=self.device) + cos, sin = rope(x) + assert cos.shape == (max_pos, dim) + assert sin.shape == (max_pos, dim) + + +# ============================================================================= +# MLP Tests +# ============================================================================= + + +class TestDeepSeekV32MLP: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def test_forward_shape(self): + config = _create_config() + mlp = DeepSeekV32MLP(config).to(self.device, self.dtype) + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = mlp(x) + assert output.shape == x.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + +# ============================================================================= +# MoE Tests +# ============================================================================= + + +class TestDeepSeekV32MoE: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def _create_moe(self, config): + moe = DeepSeekV32MoE(config).to(self.device, self.dtype) + moe.gate.weight = torch.nn.Parameter(torch.randn_like(moe.gate.weight)) + return moe + + def test_forward_shape(self): + config = _create_config() + moe = self._create_moe(config) + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = moe(x) + assert output.shape == x.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_with_shared_experts(self): + config = _create_config(n_shared_experts=2) + moe = self._create_moe(config) + assert moe.shared_experts is not None + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = moe(x) + assert output.shape == x.shape + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_without_shared_experts(self): + config = _create_config(n_shared_experts=None) + moe = self._create_moe(config) + assert moe.shared_experts is None + x = torch.randn(2, 4, config.hidden_size, dtype=self.dtype, device=self.device) + output = moe(x) + assert output.shape == x.shape + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_expert_structure(self): + config = _create_config() + moe = DeepSeekV32MoE(config) + assert isinstance(moe.experts, torch.nn.ModuleList) + assert len(moe.experts) == config.n_routed_experts + for i, expert in enumerate(moe.experts): + assert hasattr(expert, "gate_proj"), f"Expert {i} missing gate_proj" + assert hasattr(expert, "up_proj"), f"Expert {i} missing up_proj" + assert hasattr(expert, "down_proj"), f"Expert {i} missing down_proj" + + +# ============================================================================= +# Attention Tests +# ============================================================================= + + +class TestDeepSeekV32Attention: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def test_forward_shape_with_qlora(self): + config = _create_config() + assert config.q_lora_rank is not None + attn = DeepSeekV32Attention(config, layer_idx=0).to(self.device, self.dtype) + + B, S = 2, 4 + hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + + output = attn(hidden_states, position_ids) + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_forward_shape_no_qlora(self): + config = _create_config(q_lora_rank=None) + attn = DeepSeekV32Attention(config, layer_idx=0).to(self.device, self.dtype) + + B, S = 2, 4 + hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + + output = attn(hidden_states, position_ids) + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_indexer_present(self): + config = _create_config() + attn = DeepSeekV32Attention(config, layer_idx=0) + assert isinstance(attn.indexer, DeepSeekV32Indexer) + # Verify indexer has expected sub-modules + assert hasattr(attn.indexer, "wq_b") + assert hasattr(attn.indexer, "wk") + assert hasattr(attn.indexer, "k_norm") + assert hasattr(attn.indexer, "weights_proj") + + def test_different_batch_sizes(self): + config = _create_config() + attn = DeepSeekV32Attention(config, layer_idx=0).to(self.device, self.dtype) + for B in [1, 2, 4]: + S = 4 + hidden_states = torch.randn( + B, S, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + output = attn(hidden_states, position_ids) + assert output.shape == (B, S, config.hidden_size) + + def test_different_sequence_lengths(self): + config = _create_config() + attn = DeepSeekV32Attention(config, layer_idx=0).to(self.device, self.dtype) + for S in [1, 4, 16]: + B = 2 + hidden_states = torch.randn( + B, S, config.hidden_size, dtype=self.dtype, device=self.device + ) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + output = attn(hidden_states, position_ids) + assert output.shape == (B, S, config.hidden_size) + + +# ============================================================================= +# Decoder Layer Tests +# ============================================================================= + + +class TestDeepSeekV32DecoderLayer: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def test_dense_layer(self): + config = _create_config() + layer = DeepSeekV32DecoderLayer(config, layer_idx=0).to(self.device, self.dtype) + assert isinstance(layer.mlp, DeepSeekV32MLP) + + B, S = 2, 4 + hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + output = layer(hidden_states, position_ids) + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + def test_moe_layer(self): + config = _create_config() + layer = DeepSeekV32DecoderLayer(config, layer_idx=1).to(self.device, self.dtype) + assert isinstance(layer.mlp, DeepSeekV32MoE) + layer.mlp.gate.weight = torch.nn.Parameter(torch.randn_like(layer.mlp.gate.weight)) + + B, S = 2, 4 + hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + output = layer(hidden_states, position_ids) + assert output.shape == hidden_states.shape + assert torch.isfinite(output).all() + assert not torch.allclose(output, torch.zeros_like(output)) + + +# ============================================================================= +# MTP Layer Tests +# ============================================================================= + + +class TestDeepSeekV32MTPLayer: + def test_mtp_layer_structure(self): + config = _create_config() + mtp_layer = DeepSeekV32MTPDecoderLayer(config, layer_idx=config.num_hidden_layers) + + # MTP-specific modules + assert hasattr(mtp_layer, "embed_tokens") + assert hasattr(mtp_layer, "eh_proj") + assert hasattr(mtp_layer, "enorm") + assert hasattr(mtp_layer, "hnorm") + assert hasattr(mtp_layer, "shared_head") + assert isinstance(mtp_layer.shared_head, DeepSeekV32MTPSharedHead) + assert hasattr(mtp_layer.shared_head, "head") + assert hasattr(mtp_layer.shared_head, "norm") + + # Standard decoder layer modules + assert hasattr(mtp_layer, "self_attn") + assert hasattr(mtp_layer, "mlp") + assert hasattr(mtp_layer, "input_layernorm") + assert hasattr(mtp_layer, "post_attention_layernorm") + + def test_mtp_layer_weight_keys(self): + """Verify MTP layer state dict keys match checkpoint expectations.""" + config = _create_config() + mtp_layer = DeepSeekV32MTPDecoderLayer(config, layer_idx=config.num_hidden_layers) + state_dict = mtp_layer.state_dict() + + expected_keys = [ + "embed_tokens.weight", + "eh_proj.weight", + "enorm.weight", + "hnorm.weight", + "shared_head.head.weight", + "shared_head.norm.weight", + "self_attn.indexer.wq_b.weight", + "self_attn.indexer.wk.weight", + "self_attn.indexer.k_norm.weight", + "self_attn.indexer.k_norm.bias", + "self_attn.indexer.weights_proj.weight", + ] + for key in expected_keys: + assert key in state_dict, f"Missing key: {key}" + + +# ============================================================================= +# Full Model Tests +# ============================================================================= + + +class TestDeepSeekV32ForCausalLM: + @pytest.fixture(autouse=True) + def setup_method(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + self.dtype = torch.bfloat16 + + def test_forward(self): + config = _create_config() + model = DeepSeekV32ForCausalLM(config).to(self.device, self.dtype) + + B, S = 2, 4 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=self.device) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + + output = model(input_ids=input_ids, position_ids=position_ids) + assert output.logits.shape == (B, S, config.vocab_size) + assert torch.isfinite(output.logits).all() + assert not torch.allclose(output.logits, torch.zeros_like(output.logits)) + + def test_output_dtype(self): + config = _create_config() + model = DeepSeekV32ForCausalLM(config).to(self.device, self.dtype) + + B, S = 2, 4 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=self.device) + position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + + output = model(input_ids=input_ids, position_ids=position_ids) + assert output.logits.dtype == torch.float32 + assert torch.isfinite(output.logits).all() + + def test_position_ids_required(self): + config = _create_config() + model = DeepSeekV32ForCausalLM(config).to(self.device, self.dtype) + + B, S = 2, 4 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=self.device) + + with pytest.raises(AssertionError, match="position_ids is required"): + model(input_ids=input_ids) + + def test_layer_types(self): + config = _create_config() + model = DeepSeekV32ForCausalLM(config) + # Layer 0 should be dense + assert type(model.model.layers[0].mlp).__name__ == "DeepSeekV32MLP" + # Layer 1+ should be MoE + for i in range(1, config.num_hidden_layers): + assert type(model.model.layers[i].mlp).__name__ == "DeepSeekV32MoE" + + def test_mtp_layers_present(self): + config = _create_config() + model = DeepSeekV32ForCausalLM(config) + total_layers = config.num_hidden_layers + config.num_nextn_predict_layers + assert len(model.model.layers) == total_layers + # Last layer should be MTP + assert isinstance(model.model.layers[-1], DeepSeekV32MTPDecoderLayer) + + def test_no_mtp_layers(self): + config = _create_config_no_mtp() + model = DeepSeekV32ForCausalLM(config) + assert len(model.model.layers) == config.num_hidden_layers + + def test_indexer_in_all_layers(self): + config = _create_config() + model = DeepSeekV32ForCausalLM(config) + for i, layer in enumerate(model.model.layers): + assert hasattr(layer.self_attn, "indexer"), f"Layer {i} missing indexer" + assert isinstance(layer.self_attn.indexer, DeepSeekV32Indexer) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_full_model(B, S, dtype): + """Test full model produces valid output.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + config = _create_config() + model = DeepSeekV32ForCausalLM(config) + model.to(device=device, dtype=dtype) + model.eval() + + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + output = model(input_ids=input_ids, position_ids=position_ids) + + assert output.logits.shape == (B, S, config.vocab_size) + assert not torch.isnan(output.logits).any() + assert not torch.isinf(output.logits).any() + assert not torch.allclose(output.logits, torch.zeros_like(output.logits)) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_moe_layer(B, S, dtype): + """Test MoE layer produces valid output.""" + device = "cuda" if torch.cuda.is_available() else "cpu" + config = _create_config() + + moe = DeepSeekV32MoE(config).to(device=device, dtype=dtype) + moe.gate.weight = torch.nn.Parameter(torch.randn_like(moe.gate.weight)) + moe.eval() + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + output = moe(x) + assert output.shape == x.shape + assert not torch.isnan(output).any() + assert not torch.isinf(output).any() + assert not torch.allclose(output, torch.zeros_like(output)) + + +# ============================================================================= +# Export Test +# ============================================================================= + + +def test_deepseek_v32_model_can_be_exported(): + """Test that the custom model can be exported with torch_export_to_gm.""" + device = "cuda" + dtype = torch.bfloat16 + config = _create_config_no_mtp() + + model = DeepSeekV32ForCausalLM(config) + model.to(device=device, dtype=dtype) + model.eval() + + B, S = 2, 8 + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + batch_size_dynamic = Dim.DYNAMIC + seq_len_dynamic = Dim.DYNAMIC + dynamic_shapes = ( + {0: batch_size_dynamic, 1: seq_len_dynamic}, + {0: batch_size_dynamic, 1: seq_len_dynamic}, + ) + + gm = torch_export_to_gm( + model, + args=tuple(), + kwargs={"input_ids": input_ids, "position_ids": position_ids}, + dynamic_shapes=dynamic_shapes, + ) + + move_to_device(gm, device) + + with torch.inference_mode(): + out_gm = gm(input_ids=input_ids, position_ids=position_ids) + + assert "logits" in out_gm + logits = out_gm["logits"] + assert logits.shape == (B, S, config.vocab_size) + assert torch.isfinite(logits).all() + + # Test with different input shape (dynamic shapes) + B2, S2 = 1, 4 + input_ids2 = torch.randint(0, config.vocab_size, (B2, S2), device=device) + position_ids2 = torch.arange(S2, device=device).unsqueeze(0).expand(B2, -1) + + with torch.inference_mode(): + out_gm2 = gm(input_ids=input_ids2, position_ids=position_ids2) + + logits2 = out_gm2["logits"] + assert logits2.shape == (B2, S2, config.vocab_size) + assert torch.isfinite(logits2).all() + + +# ============================================================================= +# Registration Test +# ============================================================================= + + +def test_deepseek_v32_model_registration(): + """Test that DeepSeekV32ForCausalLM is registered with the factory.""" + from tensorrt_llm._torch.auto_deploy.models.hf import AutoModelForCausalLMFactory + + assert "DeepseekV32Config" in AutoModelForCausalLMFactory._custom_model_mapping + assert ( + AutoModelForCausalLMFactory._custom_model_mapping["DeepseekV32Config"] + == DeepSeekV32ForCausalLM + ) + + +def test_deepseek_v32_config_registration(): + """Test that DeepseekV32Config is registered with AutoConfig.""" + from transformers import AutoConfig + + config = AutoConfig.for_model("deepseek_v32") + assert config.__class__.__name__ == "DeepseekV32Config" + + +# ============================================================================= +# Numerical Equivalence Tests (against HF DeepSeek-V3 reference) +# ============================================================================= + + +def _get_hf_v3_model_class(): + """Get the HF DeepseekV3ForCausalLM class.""" + try: + from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3ForCausalLM as HFDeepseekV3ForCausalLM, + ) + + return HFDeepseekV3ForCausalLM + except ImportError: + return None + + +def _get_hf_v3_config_class(): + """Get the HF DeepseekV3Config class.""" + try: + from transformers.models.deepseek_v3.configuration_deepseek_v3 import ( + DeepseekV3Config as HFDeepseekV3Config, + ) + + return HFDeepseekV3Config + except ImportError: + return None + + +def _get_hf_v3_mlp_class(): + """Get the HF DeepseekV3MLP class.""" + try: + from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3MLP as HFDeepseekV3MLP, + ) + + return HFDeepseekV3MLP + except ImportError: + return None + + +def _get_hf_v3_moe_class(): + """Get the HF DeepseekV3MoE class.""" + try: + from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3MoE as HFDeepseekV3MoE, + ) + + return HFDeepseekV3MoE + except ImportError: + return None + + +def _get_hf_v3_rmsnorm_class(): + """Get the HF DeepseekV3RMSNorm class.""" + try: + from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3RMSNorm as HFDeepseekV3RMSNorm, + ) + + return HFDeepseekV3RMSNorm + except ImportError: + return None + + +def _create_hf_v3_config(): + """Create HF V3 config matching our small test config.""" + HFConfig = _get_hf_v3_config_class() + if HFConfig is None: + return None + + return HFConfig( + vocab_size=1000, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=3, + num_attention_heads=4, + num_key_value_heads=4, + hidden_act="silu", + max_position_embeddings=512, + rms_norm_eps=1e-6, + q_lora_rank=32, + kv_lora_rank=32, + qk_nope_head_dim=8, + qk_rope_head_dim=8, + v_head_dim=16, + n_routed_experts=4, + n_shared_experts=1, + num_experts_per_tok=2, + moe_intermediate_size=32, + n_group=1, + topk_group=1, + routed_scaling_factor=1.0, + topk_method="noaux_tc", + scoring_func="sigmoid", + norm_topk_prob=True, + first_k_dense_replace=1, + moe_layer_freq=1, + rope_theta=10000.0, + rope_scaling=None, + attention_bias=False, + attention_dropout=0.0, + pad_token_id=0, + ) + + +def _filter_hf_state_dict(hf_sd): + """Filter out HF-specific buffers/keys not present in custom model.""" + return {k: v for k, v in hf_sd.items() if "ep_rank" not in k and "experts_per_rank" not in k} + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_rmsnorm_numerical_equivalence(B, S, dtype): + """Test RMSNorm produces numerically equivalent output to HF V3.""" + HFRMSNorm = _get_hf_v3_rmsnorm_class() + if HFRMSNorm is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + hidden_size = 64 + + hf_norm = HFRMSNorm(hidden_size, eps=1e-6).to(device=device, dtype=dtype) + hf_norm.eval() + + custom_norm = DeepSeekV32RMSNorm(hidden_size, eps=1e-6).to(device=device, dtype=dtype) + custom_norm.load_state_dict(hf_norm.state_dict()) + custom_norm.eval() + + x = torch.randn(B, S, hidden_size, device=device, dtype=dtype) + + hf_out = hf_norm(x) + custom_out = custom_norm(x) + + torch.testing.assert_close(custom_out, hf_out, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_mlp_numerical_equivalence(B, S, dtype): + """Test MLP produces numerically equivalent output to HF V3.""" + HFMLP = _get_hf_v3_mlp_class() + if HFMLP is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + config = _create_config_no_mtp() + hf_config = _create_hf_v3_config() + + hf_mlp = HFMLP(hf_config).to(device=device, dtype=dtype) + hf_mlp.eval() + + custom_mlp = DeepSeekV32MLP(config).to(device=device, dtype=dtype) + custom_mlp.load_state_dict(hf_mlp.state_dict()) + custom_mlp.eval() + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + + hf_out = hf_mlp(x) + custom_out = custom_mlp(x) + + torch.testing.assert_close(custom_out, hf_out, rtol=1e-3, atol=1e-3) + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_moe_numerical_equivalence(B, S, dtype): + """Test MoE produces numerically equivalent output to HF V3.""" + HFMoE = _get_hf_v3_moe_class() + if HFMoE is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + config = _create_config_no_mtp() + hf_config = _create_hf_v3_config() + + hf_moe = HFMoE(hf_config).to(device=device, dtype=dtype) + hf_moe.eval() + hf_moe.gate.weight = torch.nn.Parameter(torch.randn_like(hf_moe.gate.weight)) + + custom_moe = DeepSeekV32MoE(config).to(device=device, dtype=dtype) + hf_sd = _filter_hf_state_dict(hf_moe.state_dict()) + custom_moe.load_state_dict(hf_sd) + custom_moe.eval() + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + + hf_out = hf_moe(x) + if isinstance(hf_out, tuple): + hf_out = hf_out[0] + custom_out = custom_moe(x) + + from _model_test_utils import assert_rmse_close + + assert_rmse_close(custom_out, hf_out, rmse_ratio_tol=0.02, msg="MoE: ") + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_attention_numerical_equivalence(B, S, dtype): + """Test attention (MLA) produces numerically equivalent output to HF V3. + + Uses the full model to load weights (via the model-level de-interleaving hook), + then compares attention layer outputs directly. + """ + HFModel = _get_hf_v3_model_class() + if HFModel is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + config = _create_config_no_mtp() + hf_config = _create_hf_v3_config() + + hf_model = HFModel(hf_config).to(device=device, dtype=dtype) + hf_model.eval() + hf_attn = hf_model.model.layers[0].self_attn + hf_rope = hf_model.model.rotary_emb + + custom_model = DeepSeekV32ForCausalLM(config).to(device=device, dtype=dtype) + hf_sd = _filter_hf_state_dict(hf_model.state_dict()) + # Custom model has extra indexer keys, load with strict=False + custom_model.load_state_dict(hf_sd, strict=False) + custom_model.eval() + custom_attn = custom_model.model.layers[0].self_attn + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + pos_emb = hf_rope(x, position_ids) + hf_out, _ = hf_attn(x, position_embeddings=pos_emb, position_ids=position_ids) + custom_out = custom_attn(x, position_ids) + + from _model_test_utils import assert_rmse_close + + assert_rmse_close(custom_out, hf_out, rmse_ratio_tol=0.10, msg="Attention: ") + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_decoder_layer_numerical_equivalence(B, S, dtype): + """Test decoder layer (dense, layer_idx=0) produces numerically equivalent output to HF V3.""" + HFModel = _get_hf_v3_model_class() + if HFModel is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + config = _create_config_no_mtp() + hf_config = _create_hf_v3_config() + + hf_model = HFModel(hf_config).to(device=device, dtype=dtype) + hf_model.eval() + hf_layer = hf_model.model.layers[0] + hf_rope = hf_model.model.rotary_emb + + custom_model = DeepSeekV32ForCausalLM(config).to(device=device, dtype=dtype) + hf_sd = _filter_hf_state_dict(hf_model.state_dict()) + custom_model.load_state_dict(hf_sd, strict=False) + custom_model.eval() + custom_layer = custom_model.model.layers[0] + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + pos_emb = hf_rope(x, position_ids) + hf_out = hf_layer(x, position_ids=position_ids, position_embeddings=pos_emb) + custom_out = custom_layer(x, position_ids) + + from _model_test_utils import assert_rmse_close + + assert_rmse_close(custom_out, hf_out, rmse_ratio_tol=0.05, msg="Decoder layer (dense): ") + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_moe_decoder_layer_numerical_equivalence(B, S, dtype): + """Test MoE decoder layer (layer_idx=1) produces numerically equivalent output to HF V3.""" + HFModel = _get_hf_v3_model_class() + if HFModel is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + config = _create_config_no_mtp() + hf_config = _create_hf_v3_config() + + hf_model = HFModel(hf_config).to(device=device, dtype=dtype) + hf_model.eval() + for module in hf_model.modules(): + if hasattr(module, "gate") and hasattr(module.gate, "weight"): + module.gate.weight = torch.nn.Parameter(torch.randn_like(module.gate.weight)) + + hf_layer = hf_model.model.layers[1] + hf_rope = hf_model.model.rotary_emb + + custom_model = DeepSeekV32ForCausalLM(config).to(device=device, dtype=dtype) + hf_sd = _filter_hf_state_dict(hf_model.state_dict()) + custom_model.load_state_dict(hf_sd, strict=False) + custom_model.eval() + custom_layer = custom_model.model.layers[1] + + x = torch.randn(B, S, config.hidden_size, device=device, dtype=dtype) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + pos_emb = hf_rope(x, position_ids) + hf_out = hf_layer(x, position_ids=position_ids, position_embeddings=pos_emb) + custom_out = custom_layer(x, position_ids) + + from _model_test_utils import assert_rmse_close + + assert_rmse_close(custom_out, hf_out, rmse_ratio_tol=0.05, msg="Decoder layer (MoE): ") + + +@pytest.mark.parametrize("B,S", _BATCH_AND_SEQUENCE_TEST_CASES) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +@torch.no_grad() +def test_deepseek_v32_full_model_numerical_equivalence(B, S, dtype): + """Test full model produces numerically equivalent output to HF V3.""" + HFModel = _get_hf_v3_model_class() + if HFModel is None: + pytest.skip("transformers doesn't have deepseek_v3 modeling") + + device = "cuda" + config = _create_config_no_mtp() + hf_config = _create_hf_v3_config() + + hf_model = HFModel(hf_config).to(device=device, dtype=dtype) + hf_model.eval() + for module in hf_model.modules(): + if hasattr(module, "gate") and hasattr(module.gate, "weight"): + module.gate.weight = torch.nn.Parameter(torch.randn_like(module.gate.weight)) + + custom_model = DeepSeekV32ForCausalLM(config).to(device=device, dtype=dtype) + hf_sd = _filter_hf_state_dict(hf_model.state_dict()) + custom_model.load_state_dict(hf_sd, strict=False) + custom_model.eval() + + input_ids = torch.randint(0, config.vocab_size, (B, S), device=device) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + + hf_out = hf_model(input_ids=input_ids, position_ids=position_ids) + custom_out = custom_model(input_ids=input_ids, position_ids=position_ids) + + from _model_test_utils import assert_rmse_close + + assert_rmse_close( + custom_out.logits.float(), + hf_out.logits.float(), + rmse_ratio_tol=0.05, + msg="Full model: ", + ) From 8eb28070239b5e05a288af9f8033020559596a80 Mon Sep 17 00:00:00 2001 From: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Date: Thu, 12 Mar 2026 12:01:30 -0700 Subject: [PATCH 2/3] [None][fix] Add attention DP config for DeepSeek-V3.2 registry entries Add deepseek_v2_ep.yaml to V3.2 registry entries to enable attention data parallelism. MLA attention is replicated, experts use EP. This avoids FP8 scale sharding issues with the q_b_proj LoRA projection. Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> --- examples/auto_deploy/model_registry/models.yaml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/examples/auto_deploy/model_registry/models.yaml b/examples/auto_deploy/model_registry/models.yaml index 91b2194507ee..f566da760645 100644 --- a/examples/auto_deploy/model_registry/models.yaml +++ b/examples/auto_deploy/model_registry/models.yaml @@ -241,11 +241,11 @@ models: yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml'] # --- DeepSeek V3.2 (Dec 2025) --- - name: deepseek-ai/DeepSeek-V3.2 - yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml'] + yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml', 'deepseek_v2_ep.yaml'] - name: deepseek-ai/DeepSeek-V3.2-Speciale - yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml'] + yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml', 'deepseek_v2_ep.yaml'] - name: nvidia/DeepSeek-V3.2-NVFP4 - yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml'] + yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'num_hidden_layers_5.yaml', 'deepseek_v2_ep.yaml'] # --- GLM-4.6 (Sep 2025) --- - name: zai-org/GLM-4.6 yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml'] From dc331428e54e533f1bc5da35efb6bc18b378533f Mon Sep 17 00:00:00 2001 From: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> Date: Thu, 12 Mar 2026 12:07:24 -0700 Subject: [PATCH 3/3] [None][fix] Move RoPE to model level, respond to review feedback - Move rotary_emb from per-attention to DeepSeekV32Model level - Compute cos/sin once in model forward and pass pre-sliced to layers - Attention and decoder layer forward now take (cos, sin) instead of position_ids - Update all tests for new forward signature Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com> --- .../models/custom/modeling_deepseek_v32.py | 102 +++++++++--------- .../models/test_deepseek_v32_modeling.py | 46 +++++--- 2 files changed, 83 insertions(+), 65 deletions(-) diff --git a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py index c10da8171e06..46efa170e227 100644 --- a/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py +++ b/tensorrt_llm/_torch/auto_deploy/models/custom/modeling_deepseek_v32.py @@ -405,9 +405,6 @@ def __init__(self, config, layer_idx: Optional[int] = None): # Indexer (for weight loading — not used in forward) self.indexer = DeepSeekV32Indexer(config) - # Initialize rotary embedding - self._init_rope() - # Softmax scale self.softmax_scale = self.q_head_dim ** (-0.5) if config.rope_scaling is not None: @@ -419,47 +416,11 @@ def __init__(self, config, layer_idx: Optional[int] = None): ) self.softmax_scale = self.softmax_scale * mscale * mscale - def _init_rope(self): - if self.config.rope_scaling is None: - self.rotary_emb = DeepSeekV32RotaryEmbedding( - self.qk_rope_head_dim, - max_position_embeddings=self.max_position_embeddings, - base=self.rope_theta, - ) - else: - scaling_type = self.config.rope_scaling["type"] - scaling_factor = self.config.rope_scaling["factor"] - - if scaling_type == "yarn": - kwargs = { - key: self.config.rope_scaling[key] - for key in [ - "original_max_position_embeddings", - "beta_fast", - "beta_slow", - "mscale", - "mscale_all_dim", - ] - if key in self.config.rope_scaling - } - self.rotary_emb = DeepSeekV32YarnRotaryEmbedding( - self.qk_rope_head_dim, - max_position_embeddings=self.max_position_embeddings, - scaling_factor=scaling_factor, - base=self.rope_theta, - **kwargs, - ) - else: - self.rotary_emb = DeepSeekV32RotaryEmbedding( - self.qk_rope_head_dim, - max_position_embeddings=self.max_position_embeddings, - base=self.rope_theta, - ) - def forward( self, hidden_states: torch.Tensor, - position_ids: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, ) -> torch.Tensor: bsz, q_len, _ = hidden_states.size() @@ -484,13 +445,8 @@ def forward( # k_pe: [B, S, 1, qk_rope_head_dim] (BSND layout, shared across heads) k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim) - kv_seq_len = q_len - - cos, sin = self.rotary_emb(hidden_states, seq_len=kv_seq_len) - cos = cos[position_ids] - sin = sin[position_ids] - # Apply RoPE using custom op (weights pre-permuted to NeoX format at load time) + # cos/sin are pre-sliced by position_ids at the model level q_pe_rotated, kpe = torch.ops.auto_deploy.torch_rope_with_explicit_cos_sin( q_pe, k_pe, @@ -549,12 +505,13 @@ def __init__(self, config, layer_idx: int): def forward( self, hidden_states: torch.Tensor, - position_ids: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, ) -> torch.Tensor: # Self attention residual = hidden_states hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, position_ids) + hidden_states = self.self_attn(hidden_states, cos, sin) hidden_states = residual + hidden_states # MLP/MoE @@ -652,8 +609,48 @@ def __init__(self, config): self.norm = DeepSeekV32RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + # RoPE (computed once, shared across all layers) + self._init_rope(config) + self.post_init() + def _init_rope(self, config): + if config.rope_scaling is None: + self.rotary_emb = DeepSeekV32RotaryEmbedding( + config.qk_rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + ) + else: + scaling_type = config.rope_scaling["type"] + scaling_factor = config.rope_scaling["factor"] + + if scaling_type == "yarn": + kwargs = { + key: config.rope_scaling[key] + for key in [ + "original_max_position_embeddings", + "beta_fast", + "beta_slow", + "mscale", + "mscale_all_dim", + ] + if key in config.rope_scaling + } + self.rotary_emb = DeepSeekV32YarnRotaryEmbedding( + config.qk_rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + scaling_factor=scaling_factor, + base=config.rope_theta, + **kwargs, + ) + else: + self.rotary_emb = DeepSeekV32RotaryEmbedding( + config.qk_rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + ) + def get_input_embeddings(self): return self.embed_tokens @@ -679,9 +676,14 @@ def forward( hidden_states = inputs_embeds + # Compute RoPE cos/sin once for all layers + cos, sin = self.rotary_emb(hidden_states) + cos = cos[position_ids] # [B, S, head_dim] + sin = sin[position_ids] # [B, S, head_dim] + # Only process main decoder layers (skip MTP layers) for idx in range(self.config.num_hidden_layers): - hidden_states = self.layers[idx](hidden_states, position_ids) + hidden_states = self.layers[idx](hidden_states, cos, sin) hidden_states = self.norm(hidden_states) diff --git a/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py b/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py index 3b5ee5447172..9a6656c21c31 100644 --- a/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py +++ b/tests/unittest/auto_deploy/singlegpu/models/test_deepseek_v32_modeling.py @@ -90,6 +90,19 @@ def _create_config_no_mtp(**kwargs): return _create_config(num_nextn_predict_layers=0, **kwargs) +def _make_cos_sin(config, B, S, device, dtype): + """Create pre-sliced cos/sin for testing attention/layer directly.""" + rope = DeepSeekV32RotaryEmbedding( + config.qk_rope_head_dim, + max_position_embeddings=config.max_position_embeddings, + base=config.rope_theta, + ).to(device) + x_dummy = torch.randn(1, 1, 1, config.qk_rope_head_dim, dtype=dtype, device=device) + cos, sin = rope(x_dummy) + position_ids = torch.arange(S, device=device).unsqueeze(0).expand(B, -1) + return cos[position_ids].to(dtype), sin[position_ids].to(dtype) + + # ============================================================================= # RMSNorm Tests # ============================================================================= @@ -247,9 +260,9 @@ def test_forward_shape_with_qlora(self): B, S = 2, 4 hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) - position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + cos, sin = _make_cos_sin(config, B, S, self.device, self.dtype) - output = attn(hidden_states, position_ids) + output = attn(hidden_states, cos, sin) assert output.shape == hidden_states.shape assert torch.isfinite(output).all() assert not torch.allclose(output, torch.zeros_like(output)) @@ -260,9 +273,9 @@ def test_forward_shape_no_qlora(self): B, S = 2, 4 hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) - position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) + cos, sin = _make_cos_sin(config, B, S, self.device, self.dtype) - output = attn(hidden_states, position_ids) + output = attn(hidden_states, cos, sin) assert output.shape == hidden_states.shape assert torch.isfinite(output).all() assert not torch.allclose(output, torch.zeros_like(output)) @@ -285,8 +298,8 @@ def test_different_batch_sizes(self): hidden_states = torch.randn( B, S, config.hidden_size, dtype=self.dtype, device=self.device ) - position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) - output = attn(hidden_states, position_ids) + cos, sin = _make_cos_sin(config, B, S, self.device, self.dtype) + output = attn(hidden_states, cos, sin) assert output.shape == (B, S, config.hidden_size) def test_different_sequence_lengths(self): @@ -297,8 +310,8 @@ def test_different_sequence_lengths(self): hidden_states = torch.randn( B, S, config.hidden_size, dtype=self.dtype, device=self.device ) - position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) - output = attn(hidden_states, position_ids) + cos, sin = _make_cos_sin(config, B, S, self.device, self.dtype) + output = attn(hidden_states, cos, sin) assert output.shape == (B, S, config.hidden_size) @@ -320,8 +333,8 @@ def test_dense_layer(self): B, S = 2, 4 hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) - position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) - output = layer(hidden_states, position_ids) + cos, sin = _make_cos_sin(config, B, S, self.device, self.dtype) + output = layer(hidden_states, cos, sin) assert output.shape == hidden_states.shape assert torch.isfinite(output).all() assert not torch.allclose(output, torch.zeros_like(output)) @@ -334,8 +347,8 @@ def test_moe_layer(self): B, S = 2, 4 hidden_states = torch.randn(B, S, config.hidden_size, dtype=self.dtype, device=self.device) - position_ids = torch.arange(S, device=self.device).unsqueeze(0).expand(B, -1) - output = layer(hidden_states, position_ids) + cos, sin = _make_cos_sin(config, B, S, self.device, self.dtype) + output = layer(hidden_states, cos, sin) assert output.shape == hidden_states.shape assert torch.isfinite(output).all() assert not torch.allclose(output, torch.zeros_like(output)) @@ -821,7 +834,8 @@ def test_deepseek_v32_attention_numerical_equivalence(B, S, dtype): pos_emb = hf_rope(x, position_ids) hf_out, _ = hf_attn(x, position_embeddings=pos_emb, position_ids=position_ids) - custom_out = custom_attn(x, position_ids) + cos, sin = _make_cos_sin(config, B, S, device, dtype) + custom_out = custom_attn(x, cos, sin) from _model_test_utils import assert_rmse_close @@ -857,7 +871,8 @@ def test_deepseek_v32_decoder_layer_numerical_equivalence(B, S, dtype): pos_emb = hf_rope(x, position_ids) hf_out = hf_layer(x, position_ids=position_ids, position_embeddings=pos_emb) - custom_out = custom_layer(x, position_ids) + cos, sin = _make_cos_sin(config, B, S, device, dtype) + custom_out = custom_layer(x, cos, sin) from _model_test_utils import assert_rmse_close @@ -897,7 +912,8 @@ def test_deepseek_v32_moe_decoder_layer_numerical_equivalence(B, S, dtype): pos_emb = hf_rope(x, position_ids) hf_out = hf_layer(x, position_ids=position_ids, position_embeddings=pos_emb) - custom_out = custom_layer(x, position_ids) + cos, sin = _make_cos_sin(config, B, S, device, dtype) + custom_out = custom_layer(x, cos, sin) from _model_test_utils import assert_rmse_close