61 transformer blocks › 58 MoE blocks › Sparse latent attention › Key–value compression
Key–value latent RMSNorm
×learnedself.kv_norm(kv)
Rescale each position’s features using their root mean square, then apply learned gain.
- in
- [B batch, 1 positions, 512 kv latent]
- out
- [B batch, 1 positions, 512 kv latent]
- parameters
- 512 · × 58 MoE blocks = 29,696
Normalized in float32, multiplied by the gain in float32, then cast back to the input type (model.py 296–301).
Parameters
512 parameters per instance.
| Component | Calculation | Parameters |
|---|---|---|
| Learned gain | 512 | |
| Per-instance total | 512 | |
Each normalization instance owns its learned gain vector.
Across 58 independent MoE blocks: parameters.
Derived from the configured feature width; checkpoint tensors have not been independently counted.
Codemodel.py 280–568
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.dim = dim
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
def forward(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None):
"""
Forward pass for RMSNorm.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Normalized tensor with the same shape as input.
"""
dtype = x.dtype
if residual is None:
x = x.float()
var = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(var + self.eps)
return (self.weight * x).to(dtype)
else:
x = residual = x.float() + residual.float()
var = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(var + self.eps)
return (self.weight * x).to(dtype), residual.to(dtype) self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
self.v_head_dim = args.v_head_dim
self.wq_a = Linear(self.dim, self.q_lora_rank)
self.q_norm = RMSNorm(self.q_lora_rank)
self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)
self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)
self.kv_norm = RMSNorm(self.kv_lora_rank)
self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))
self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)
self.softmax_scale = self.qk_head_dim ** -0.5
self.scale_fmt = args.scale_fmt
if args.max_seq_len > args.original_seq_len:
mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
self.softmax_scale = self.softmax_scale * mscale * mscale
self.indexer = Indexer(args) kv = self.wkv_a(x)
kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
kv = self.kv_norm(kv)
k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)Sources3
Direct evidence
RMSNorm owns one gain per feature and computes in float32; with a residual argument it first adds the residual, fusing the residual addition into the next normalization.
rmsnorm · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 280–306
Read code excerpt
def __init__(self, dim: int, eps: float = 1e-6):
super().__init__()
self.dim = dim
self.eps = eps
self.weight = nn.Parameter(torch.ones(dim, dtype=torch.float32))
def forward(self, x: torch.Tensor, residual: Optional[torch.Tensor] = None):
"""
Forward pass for RMSNorm.
Args:
x (torch.Tensor): Input tensor.
Returns:
torch.Tensor: Normalized tensor with the same shape as input.
"""
dtype = x.dtype
if residual is None:
x = x.float()
var = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(var + self.eps)
return (self.weight * x).to(dtype)
else:
x = residual = x.float() + residual.float()
var = x.pow(2).mean(-1, keepdim=True)
x = x * torch.rsqrt(var + self.eps)
return (self.weight * x).to(dtype), residual.to(dtype)Direct evidence
MLA owns a query down-projection, query RMSNorm and up-projection, a joint key–value down-projection that also emits the shared rotary key, a latent RMSNorm, a key–value up-projection and an output projection.
mla-init · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 523–539
Read code excerpt
self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim
self.v_head_dim = args.v_head_dim
self.wq_a = Linear(self.dim, self.q_lora_rank)
self.q_norm = RMSNorm(self.q_lora_rank)
self.wq_b = ColumnParallelLinear(self.q_lora_rank, self.n_heads * self.qk_head_dim)
self.wkv_a = Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim)
self.kv_norm = RMSNorm(self.kv_lora_rank)
self.wkv_b = ColumnParallelLinear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim))
self.wo = RowParallelLinear(self.n_heads * self.v_head_dim, self.dim)
self.softmax_scale = self.qk_head_dim ** -0.5
self.scale_fmt = args.scale_fmt
if args.max_seq_len > args.original_seq_len:
mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0
self.softmax_scale = self.softmax_scale * mscale * mscale
self.indexer = Indexer(args)Direct evidence
The joint down-projection yields the key–value latent and one rotary key shared by all heads; the latent is normalized and the key rotated.
mla-kv · deepseek-ai/DeepSeek-V3.2-Exp / inference/model.py
Pinned revision 87e509a2e5a1 ↗ · lines 565–568
Read code excerpt
kv = self.wkv_a(x)
kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)
kv = self.kv_norm(kv)
k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis)