Multi-Head Latent Attention
Decoupled RoPE and non-RoPE query/key dimensions, Q and KV low-rank projections, and YaRN position scaling.
STARPOWER MODEL ARCHITECTURE
arXiv-WVY-base
A compact DeepSeek-V3-style causal language model pretrained on arXiv titles and abstracts. The notebook preserves the complete architecture pattern—MLA, YaRN, routed and shared experts, dense-to-MoE layers, and sequential Multi-Token Prediction—while reducing widths and counts for Kaggle-scale training.
HIDDEN384
MAIN LAYERS6
MTP LAYERS1
ATTN HEADS6
ROUTED EXPERTS8
CONTEXT4,096
01 / SYSTEM
Full architecture, reduced scale.
DeepSeekMoE
Eight routed experts, one shared expert, sigmoid grouped top-k routing, correction bias, and two experts selected per token.
Sequential MTP
A full appended Transformer/MoE prediction layer with enorm, hnorm, eh_proj, shared embeddings, and a shared output head.
02 / CODE
Build the model.
These blocks come directly from the supplied training notebook. Every block can be copied independently.
Configuration / Python
config = DeepseekV3Config(
vocab_size=VOCAB_SIZE,
hidden_size=384,
intermediate_size=1024,
moe_intermediate_size=384,
num_hidden_layers=6,
num_nextn_predict_layers=1,
mtp_loss_weight=0.3,
num_attention_heads=6,
num_key_value_heads=6,
n_shared_experts=1,
n_routed_experts=8,
ep_size=1,
routed_scaling_factor=1.0,
kv_lora_rank=96,
q_lora_rank=128,
qk_rope_head_dim=16,
qk_nope_head_dim=48,
v_head_dim=64,
topk_method="noaux_tc",
n_group=4,
topk_group=2,
num_experts_per_tok=2,
moe_layer_freq=1,
first_k_dense_replace=2,
norm_topk_prob=True,
scoring_func="sigmoid",
hidden_act="silu",
max_position_embeddings=4096,
initializer_range=0.006,
rms_norm_eps=1e-6,
use_cache=False,
tie_word_embeddings=False,
rope_theta=10000.0,
rope_scaling={
"type": "yarn",
"factor": 4.0,
"original_max_position_embeddings": 1024,
"beta_fast": 32,
"beta_slow": 1,
"mscale": 1.0,
"mscale_all_dim": 1.0,
},
attention_bias=False,
attention_dropout=0.0,
)
config._attn_implementation = "eager"Multi-Head Latent Attention / Python
class DeepseekV3Attention(nn.Module):
def __init__(self, config, layer_idx=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_rope_head_dim = config.qk_rope_head_dim
self.qk_nope_head_dim = config.qk_nope_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.q_a_proj = nn.Linear(
self.hidden_size, config.q_lora_rank,
bias=config.attention_bias,
)
self.q_a_layernorm = DeepseekV3RMSNorm(config.q_lora_rank)
self.q_b_proj = nn.Linear(
config.q_lora_rank,
self.num_heads * self.q_head_dim,
bias=False,
)
self.kv_a_proj_with_mqa = nn.Linear(
self.hidden_size,
config.kv_lora_rank + config.qk_rope_head_dim,
bias=config.attention_bias,
)
self.kv_a_layernorm = DeepseekV3RMSNorm(config.kv_lora_rank)
self.kv_b_proj = nn.Linear(
config.kv_lora_rank,
self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
bias=False,
)Grouped Top-K MoE Routing / Python
def forward(self, hidden_states):
batch_size, sequence_length, hidden_size = hidden_states.shape
hidden_states = hidden_states.view(-1, hidden_size)
logits = F.linear(
hidden_states.float(),
self.weight.float(),
None,
)
scores = logits.sigmoid()
scores_for_choice = (
scores.view(batch_size * sequence_length, -1)
+ self.e_score_correction_bias.unsqueeze(0)
)
group_scores = (
scores_for_choice
.view(batch_size * sequence_length, self.n_group, -1)
.topk(2, dim=-1)[0]
.sum(dim=-1)
)
group_idx = torch.topk(
group_scores,
k=self.topk_group,
dim=-1,
sorted=False,
)[1]
group_mask = torch.zeros_like(group_scores)
group_mask.scatter_(1, group_idx, 1)
score_mask = (
group_mask.unsqueeze(-1)
.expand(
batch_size * sequence_length,
self.n_group,
self.n_routed_experts // self.n_group,
)
.reshape(batch_size * sequence_length, -1)
)
routed_scores = scores_for_choice.masked_fill(
~score_mask.bool(), float("-inf")
)
topk_idx = torch.topk(
routed_scores,
k=self.top_k,
dim=-1,
sorted=False,
)[1]
topk_weight = scores.gather(1, topk_idx)
topk_weight /= topk_weight.sum(dim=-1, keepdim=True) + 1e-20
return topk_idx, topk_weight * self.routed_scaling_factorSequential Multi-Token Prediction / Python
class DeepseekV3MTPDecoderLayer(DeepseekV3DecoderLayer):
def __init__(self, config, layer_idx):
super().__init__(config, layer_idx)
self.embed_tokens = nn.Embedding(
config.vocab_size,
config.hidden_size,
config.pad_token_id,
)
self.enorm = DeepseekV3RMSNorm(
config.hidden_size,
eps=config.rms_norm_eps,
)
self.hnorm = DeepseekV3RMSNorm(
config.hidden_size,
eps=config.rms_norm_eps,
)
self.eh_proj = nn.Linear(
2 * config.hidden_size,
config.hidden_size,
bias=False,
)
self.shared_head = DeepseekV3MTPSharedHead(config)
def fuse_inputs(
self,
previous_hidden_states,
future_token_embeddings,
):
hidden = self.hnorm(previous_hidden_states)
future = self.enorm(future_token_embeddings)
return self.eh_proj(torch.cat([hidden, future], dim=-1))Instantiate + Verify / Python
model = DeepseekV3ForCausalLM(config).to(DEVICE)
model.config.use_cache = False
total_params = sum(parameter.numel() for parameter in model.parameters())
trainable_params = sum(
parameter.numel()
for parameter in model.parameters()
if parameter.requires_grad
)
assert len(model.model.layers) == (
config.num_hidden_layers
+ config.num_nextn_predict_layers
)
mtp_layer = model.model.layers[config.num_hidden_layers]
# The MTP path physically shares both parameter matrices.
assert mtp_layer.embed_tokens.weight is model.model.embed_tokens.weight
assert mtp_layer.shared_head.head.weight is model.lm_head.weight
print(f"Total parameters: {total_params:,}")
print(f"Trainable parameters: {trainable_params:,}")03 / SOURCE
Complete files.
The full 2,028-line architecture implementation, configuration class, and original Kaggle notebook are available directly.