Docs > Concepts > Model Architecture Walkthrough

Model Architecture Walkthrough

sourcedocs/concepts/model-architecture.md words19,731 read~99 min
Prerequisite reading: This document walks through every model file in src/hymo/models/ line by line. Start here if you want to understand exactly how a forward pass flows through the 32-layer hybrid stack.
Files covered:
- model.py — the top-level HyMo model and build_hymo factory
- gdn.py — Gated Delta Net linear-attention block
- gdn_triton.py — Custom Triton kernel for the GDN recurrence
- mla.py — Multi-Head Latent Attention with MQA-4 grouping
- moe.py — DeepSeekMoE, DenseFFN, SwiGLU, and aux-loss-free routing
- mtp.py — Multi-Token Prediction heads (depth=2)
- rope.py — Rotary Position Embedding (partial RoPE)
Stale-section notice (as of commit af89c48): Sections 5.5.3 ("Step 5: Gated Delta Rule Recurrence", lines ~1655–1664) and 6 (the Triton kernel section, ~line 1790) describe a hypothetical fla-first / Triton-second fallback that was never shipped. The shipped code path is Triton-only (gdn_triton.py::triton_gated_delta_rule) with no fla dependency, and the eager path (gdn.py::_gated_delta_rule) is the reference for tests. The current source-of-truth for the kernel is optimization.md §3 and kernels.md. Treat the fla/chunk_gated_delta_rule snippets in those sections as historical plan-time documentation that illustrates the recurrence interface, not as the implementation.
Companion concepts (theory tier — read these for the why behind each section; section numbers refer to this walkthrough's table of contents):
- gdn-and-mla.md — §1 high-level architecture
- model-architecture.md — §4 MLA
- gdn-and-mla.md — §5–§6 GDN + Triton kernel
- model-architecture.md — §7 RoPE
- gdn-and-mla.md — §8 MoE
- gdn-and-mla.md — §9 MTP
- optimization.md — §10 μP
- kernels.md — §6 Triton kernel
See ../README.md for the full reading order and cross-reference table.

Table of Contents#

  1. High-Level Architecture
  2. Configuration System (core/config.py)
  3. The HyMo Stack (model.py)
  4. Multi-Head Latent Attention (mla.py)
  5. Gated Delta Net Linear Attention (gdn.py)
  6. Custom Triton GDN Kernel (gdn_triton.py)
  7. Rotary Position Embedding (rope.py)
  8. Mixture of Experts (moe.py)
  9. Multi-Token Prediction (mtp.py)
  10. Initialization — status note
  11. End-to-End Forward Pass Trace
  12. Parameter Count Breakdown

1. High-Level Architecture — Design Rationale & Stack Design#

*This section explains why HyMo is built the way it is.* The 3:1 hybrid ratio, the MLA placement strategy, the asymmetric FFN design, and the parameter sharing scheme all have specific mathematical and empirical motivations.

1.1 The Hybrid Attention Hypothesis#

1.1.1 Why Not Pure Full Attention?#

Standard transformer attention (MHA/GQA) has two costs that scale with sequence length T:

CostComplexityBottleneck
Compute (QK^T)O(T² × d)FLOPs grow quadratically
Memory (KV cache)O(T × n_h × h_d)Cache grows linearly

At T=4096, d=896, n_heads=16, head_dim=128:

  • Compute: 4096² × 896 ≈ 15 billion FLOPs per layer
  • Memory: 4096 × 16 × 128 × 2 bytes ≈ 16 MB per layer

For 8 MLA layers: 128 MB of KV cache. This is manageable at T=4096 but becomes prohibitive at T=32K (1 GB) or T=128K (4 GB).

1.1.2 Why Not Pure Linear Attention?#

Linear attention (GDN) reduces both costs:

CostComplexitySavings vs MHA
ComputeO(T × S × D)S=32, D=32 → ~400× faster
MemoryO(S × D) per layerFixed state, regardless of T

But linear attention has a fundamental limitation: limited long-range discrimination. The state H_t ∈ R^{S×D} is a compressed summary — it can store information but cannot selectively attend to arbitrary past positions.

1.1.3 The 3:1 Ratio: Cheap Recurrence + Precise Retrieval#

The key insight: most language processing is local and sequential (grammar, syntax, nearby context), while occasional long-range retrieval is needed (coreference, distant dependencies).

code
┌──────────────────────────────────────────────────────────────┐
│  Token Processing Pipeline:                                  │
│                                                              │
│  GDN layers (24×):  O(T) recurrence — cheap, local context  │
│       ↓                                                      │
│  MLA layers (8×):   O(T²) attention — expensive, precise    │
│       ↓                                                      │
│  GDN layers (24×):  O(T) recurrence — cheap, local context  │
│       ↓                                                      │
│  MLA layers (8×):   O(T²) attention — expensive, precise    │
│       ↓                                                      │
│  ... (repeats every 4 layers)                                │
└──────────────────────────────────────────────────────────────┘

Why 3:1 specifically?

The ratio is derived from empirical observations in hybrid attention literature (Jamba, Zamba, Zamba2, StripedHyena):

RatioQuality (PPL)ThroughputNotes
1:1 (16 MLA + 16 GDN)BestSlowestFull attention everywhere
3:1 (8 MLA + 24 GDN)~0.02 PPL loss3-5× fasterHyMo's choice
7:1 (4 MLA + 28 GDN)~0.1 PPL loss5-7× fasterToo few attention layers
All GDNSignificant lossFastestCannot do precise retrieval

The 3:1 ratio captures ~95% of full-attention quality at ~40% of the compute cost.

1.2 MLA Placement Strategy#

1.2.1 Why Every 4th Layer?#

MLA blocks are placed at positions {0, 4, 8, 12, 16, 20, 24, 28} — every 4th layer.

Design rationale:

  1. Uniform coverage: With 32 layers, 8 MLA blocks spaced every 4 layers ensures no token goes more than 3 layers without a full-attention retrieval step
  2. GDN state refresh: After 3 GDN layers, the recurrent state may drift. MLA provides a "checkpoint" that re-anchors the representation
  3. Gradient flow: MLA's quadratic attention provides a high-bandwidth gradient pathway back through the stack, preventing gradient vanishing in the deep recurrence
code
Layer:   0    1    2    3    4    5    6    7    8    ...
Type:   MLA  GDN  GDN  GDN  MLA  GDN  GDN  GDN  MLA  ...
         ↓         ↓              ↓         ↓         ↓
      Retrieve  Accumulate    Retrieve  Accumulate  Retrieve
      precise   local         precise   local       precise
      context   features      context   features    context

1.2.2 Why MLA at Position 0?#

Meta FAIR (Bae et al., 2025) recommends "Never place Transformer blocks at the front" — front placement leads to performance degradation. However, HyMo places MLA at position 0 for two reasons:

  1. MLA's latent bottleneck acts as a learned summarization — the kv_lora_rank=128 bottleneck compresses input embeddings into a structured representation, which is beneficial even at the input stage
  2. Empirical validation: The first 1k training steps will validate this choice. If loss spikes at position 0, MLA moves to position 1
Open question for v1.1: If MLA-at-0 hurts, swap to GDN-at-0 (position 1 becomes MLA).

1.3 Asymmetric Feed-Forward: MoE on MLA, Dense on GDN#

1.3.1 Why Not MoE Everywhere?#

The intuition: MoE dispatch overhead is amortized on expensive layers.

Layer TypePer-Token CostMoE OverheadOverhead Fraction
MLA~5.8M params (high)~9M params~15%
GDN~2.5M params (low)~9M params~36%

Putting MoE on GDN layers wastes 36% of the GDN compute budget on routing. On MLA, the overhead is only 15%.

1.3.2 Routing Noise Compounding#

More critically: GDN has recurrent state. If a token is routed to the wrong expert in a GDN block:

code
GDN (with MoE):
  t=0: token routed to expert 3 (wrong) → corrupted value → stored in H_0
  t=1: H_1 = exp(g·A)·H_0 + ... → corrupted H_0 propagates to H_1
  t=2: H_2 = exp(g·A)·H_1 + ... → corruption compounds
  ...

The routing error persists in the recurrent state and corrupts all future predictions.

MLA (with MoE):

code
MLA (with MoE):
  t=0: token routed to expert 3 (wrong) → corrupted attention output
  t=1: MLA recomputes from scratch (no recurrence) → error is local

MLA's attention is stateless — routing errors are local to one position and corrected at the next layer.

1.3.3 No FFN on GDN#

GDN blocks are recurrence-only — the gated delta rule output goes straight through out_proj to the residual stream. There is no SwiGLU FFN after the recurrence:

  • No extra dense parameters on the 24 GDN layers
  • Every GDN parameter contributes to the recurrence (in_proj, conv1d, b/c/dt/g_proj, out_proj, skip_proj)
  • Compute is spent only on the recurrence itself

1.4 Stack Configuration Summary#

1.5 The Five Invariants#

#InvariantWhy It Matters
1Asymmetric FFNMoE on MLA only — avoids routing noise in recurrence
2Partial-RoPE 25%Position/content decoupling — cleaner KV compression
3MQA-44 KV groups — compression/quality sweet spot
4FP32 master weightsNumerical stability in BF16 training
5Weight-tied embeddings57.6M params saved — same tensor for input/output

1.6 Visual Stack Diagram#

2. Configuration System (core/config.py) — Design Philosophy & Mathematical Derivations#

Why a config system matters: Every architectural knob in HyMo is defined in ModelConfig (a frozen dataclass). This section explains the mathematical relationships between config fields, how configs are derived, and why certain validation rules exist.

2.1 The No-Hardcoding Principle#

Every hyperparameter lives in configs/hymo_750m.yaml. The only numeric literals allowed in code are architectural constants (e.g., vocab_size = 64_256 for the BPE-64k + 256-byte tokenizer).

Why? Because hyperparameters are tuned iteratively. If head_dim=128 is hardcoded in model code, changing it requires editing multiple files. With config, a single YAML change propagates everywhere.

2.2 Config Flow: YAML → Model#

python
from hymo import load_config, build_hymo

config = load_config("configs/hymo_750m.yaml")  # → HyMoConfig
model = build_hymo(config)                       # → HyMo(config.model)

2.3 The Five Sub-Configs#

Config classYAML groupKey FieldsPurpose
ModelConfigmodeldim, n_heads, n_layers, mtp_depthArchitecture knobs
OptimizerConfigoptimizerlr, betas, weight_decay, muon_lrOptimizer hyperparams
SchedulerConfigschedulerwarmup_frac, stable_frac, decay_fracWSD schedule
TrainingConfigtrainingbatch_size, seq_len, grad_accum_stepsTraining dynamics
RunConfigrunname, seed, output_dirExperiment metadata

2.4 Frozen Dataclasses: Why Immutable?#

python
@dataclass(frozen=True)
class ModelConfig:
    dim: int = 896
    n_heads: int = 16
    # ... 30+ fields

Why frozen? The training loop is a complex state machine. If config is mutable, a bug could accidentally modify dim mid-training:

python
# BUG: accidentally mutate config during training
config.dim = 1024  # ← FrozenInstanceError raised immediately

How to derive variants:

python
# Correct: create a new config with modified field
from dataclasses import replace
small_config = replace(config, dim=512, n_heads=8)

The derive_config() helper computes dependent fields automatically:

python
# derive_config() computes:
# - qk_rope_head_dim from head_dim and rope_ratio
# - qk_nope_head_dim from head_dim and rope_ratio
# - kv_lora_rank from empirical formula
# - q_lora_rank from empirical formula

2.5 Mathematical Relationships Between Config Fields#

2.5.1 Dimension Arithmetic#

code
d = 896                          # model dimension
n_heads = 16                     # number of attention heads
head_dim = 128                   # per-head dimension
qk_rope_head_dim = 32            # RoPE subspace (25% of head_dim)
qk_nope_head_dim = 96            # Content subspace (75% of head_dim)

Invariants:
  head_dim = qk_rope_head_dim + qk_nope_head_dim     # 128 = 32 + 96
  q_lora_rank = 224                                   # Query bottleneck
  kv_lora_rank = 128                                  # KV bottleneck
  v_head_dim = 128                                    # Value dimension

2.5.2 GDN Dimension Arithmetic#

code
gdn_d_inner = 1280               # GDN working dimension
gdn_headdim = 32                 # GDN per-head dimension
gdn_d_state = 32                 # State matrix dimension
gdn_d_conv = 4                   # Causal convolution kernel size

gdn_n_heads = gdn_d_inner / gdn_headdim = 1280 / 32 = 40  # GDN heads

2.5.3 MoE Dimension Arithmetic#

code
n_routed = 16                    # Number of routed experts
n_shared = 1                     # Number of shared experts
n_activated = 2                  # Top-k routing
moe_inter_dim = 2304             # Expert FFN intermediate dimension
capacity_factor = 1.5            # Max tokens per expert

capacity = capacity_factor × (batch_tokens × n_activated) / n_routed
         = 1.5 × (4 × 4096 × 2) / 16
         = 1.5 × 512
         = 768 tokens per expert

2.6 Validation Rules#

Every config validates on construction. Key rules:

Cross-field validation in config_validation.py:

  • VRAM budget check: total params × bytes + optimizer states + activations ≤ A100 80GB
  • Throughput check: tokens_per_step / step_time ≥ target throughput
  • Data check: total_tokens / batch_size ≥ number of unique documents

2.7 Computed Properties#

Config fields that are derived, not specified:

2.8 The hymo_750m.yaml Structure#

3. The HyMo Stack (model.py) — Assembly, Embeddings & Forward Flow#

This section explains how the 32-layer stack is assembled, how embeddings work, and the complete forward pass flow. The design decisions around weight tying, gradient checkpointing, and logit stabilization are covered in depth.

3.1 Why Weight-Tied Embeddings?#

3.1.1 The Mathematical Argument#

Standard transformer:

code
Embedding: V × d = 64,256 × 896 = 57.5M parameters
Output head: d × V = 896 × 64,256 = 57.5M parameters
Total: 115M parameters (embedding alone = 6.2% of model)

Weight-tied transformer:

code
Embedding: V × d = 57.5M parameters
Output head: same tensor (0 additional parameters)
Total: 57.5M parameters (saved 57.5M)

Why does this work? The output logits are:

code
logits = h · W^T    where h ∈ R^d (hidden state), W ∈ R^{V×d} (embedding)

This is the dot product between the hidden state and every embedding vector. The intuition: a token's embedding is its "identity" — predicting the next token means finding which embedding is most similar to the current hidden state.

3.1.2 When Weight Tying Breaks Down#

Weight tying assumes the input embedding space and the output logit space are isomorphic. This works well for:

  • Language models (next-token prediction)
  • Models where the vocabulary is shared between input and output

It may not work for:

  • Multilingual models with different input/output vocabularies
  • Models with separate encoder/decoder vocabularies

HyMo uses weight tying because it's a decoder-only LM with a shared BPE vocabulary.

3.2 Embedding: The First Transformation#

python
self.embed = nn.Embedding(config.vocab_size, config.dim)
# Weight: (64,256, 896) — each token has a 896-dim vector

Forward:

python
x = self.embed(tokens)  # (B, T) → (B, T, 896)

The embedding is a lookup table: for each integer token ID, return the corresponding row from the weight matrix. This is equivalent to a one-hot encoding followed by a linear projection:

code
embed(t) = one_hot(t) @ W_embed

But the lookup is O(1) per token, not O(V × d).

μP initialization: Embedding weights use std = 1/√d = 1/√896 ≈ 0.033 (not 1/d like other parameters). This ensures the embedding vectors have unit variance at initialization, which is important for the first token to have reasonable magnitude.

3.3 Layer Assembly: The Loop#

python
self.layers = nn.ModuleList()
for i in range(config.n_layers):      # 32 iterations
    if i in mla_positions:             # {0, 4, 8, ..., 28}
        self.layers.append(MLABlock(config, layer_idx=i))
    else:
        use_rope = i not in nope_hybrid  # True for all 24 GDN in v1.0
        self.layers.append(
            GatedDeltaNetBlock(config, layer_idx=i, use_rope=use_rope)
        )

Key design decisions:

  1. nn.ModuleList — not a plain list. This ensures all layer parameters are registered for:
    • Optimizer parameter groups
    • FSDP sharding
    • model.parameters() iteration
  1. layer_idx passed to each block — used for:
    • Positional encoding offsets (if needed)
    • Per-layer configuration (e.g., use_rope flag)
    • Debugging and logging
  1. use_rope flag — In v1.0, all 24 GDN layers use RoPE (use_rope=True). In v1.1, 7 GDN layers at positions {3, 7, 11, 15, 19, 23, 27} may get use_rope=False (NoPE-hybrid ablation).

3.4 Gradient Checkpointing: Memory vs Compute Trade-off#

python
def _run_layers(self, x: torch.Tensor) -> torch.Tensor:
    for layer in self.layers:
        use_cp = getattr(layer, "use_checkpoint", False)
        if use_cp and self.training:
            x = checkpoint(layer, x, use_reentrant=False)
        else:
            x = layer(x)
    return x

Why checkpointing?

During training, activations (intermediate tensors) must be stored for backward pass. For 32 layers, this is:

code
Activation memory = 32 layers × (B × T × d) × bytes_per_element
                  = 32 × 4 × 4096 × 896 × 2 bytes
                  = ~940 MB (BF16)

Gradient checkpointing reduces this by:

  • Only storing activations at "checkpoints" (layer boundaries)
  • Recomputing activations during backward pass
  • Saves ~60% activation memory at ~33% compute overhead
python
# Without checkpointing:
# Forward: store all 32 layer activations → 940 MB
# Backward: use stored activations → 0 MB recomputation

# With checkpointing every 4 layers:
# Forward: store 8 checkpoint activations → 235 MB
# Backward: recompute 3 layers between checkpoints → 235 MB recomputation
# Total: 235 MB (saved 705 MB) at ~33% more compute

In v1.0, neither MLABlock nor GatedDeltaNetBlock enable checkpointing by default. It's reserved for the production run if activation memory becomes tight.

3.5 Logit Softcapping: Preventing Pathological Logits#

python
def softcap(self, logits):
    if self.logit_softcap <= 0:
        return logits
    return self.logit_softcap * torch.tanh(logits / self.logit_softcap)

The problem: During long training runs, logits can grow unboundedly. This causes:

  • Numerical instability in softmax (overflow)
  • Gradient explosion
  • Training divergence

The solution: Tanh softcapping maps logits to [-softcap, softcap]:

code
softcapped = softcap × tanh(logits / softcap)

Mathematical properties:

  • softcapped ≈ logits when |logits| ≪ softcap (linear regime)
  • softcapped → ±softcap when |logits| ≫ softcap (saturation)
  • Gradient: d(softcapped)/d(logits) = 1 - tanh²(logits/softcap) (vanishes at saturation)

At softcap=15.0:

  • Logits in [-5, 5]: ~95% of the range, gradient ~0.91
  • Logits in [-10, 10]: ~99.8% of the range, gradient ~0.64
  • Logits beyond ±15: gradient → 0, capped at ±15

This is a PaLM-style stabilization technique that prevents training divergence without losing representational capacity.

3.6 Factory Function: build_hymo#

python
def build_hymo(config: HyMoConfig) -> HyMo:
    return HyMo(config.model)

Why a factory? The separation between HyMo(config) and build_hymo(config) exists so that future versions can perform post-initialization steps:

python
def build_hymo(config: HyMoConfig) -> HyMo:
    model = HyMo(config.model)

    # No custom init pass — PyTorch module defaults + inline gate/GDN init
    # (see optimization.md). FSDP wrapping happens in the trainer.

    return model

Currently build_hymo is a pass-through, but the indirection is a deliberate architectural choice for extensibility.

3.7 Complete Forward Flow#

python
# HyMo.forward(tokens)
def forward(self, tokens: torch.Tensor) -> torch.Tensor:
    hidden = self.forward_with_hidden(tokens)[1]
    logits = self.head(hidden)
    return self.softcap(logits)

# HyMo.forward_with_hidden(tokens, start_pos=0)
def forward_with_hidden(self, tokens, start_pos=0):
    x = self.embed(tokens)          # (B, T) → (B, T, 896)
    x = self._run_layers(x)         # 32-layer stack
    hidden = self.norm(x)           # RMSNorm → (B, T, 896)
    logits = self.head(hidden)      # (B, T, 64,256)
    return self.softcap(logits), hidden

Data flow:

3.8 Code Walkthrough (src/hymo/models/model.py)#

4. Multi-Head Latent Attention (mla.py) — Mathematical Foundations & Design Decisions#

File: src/hymo/models/mla.py

This section builds MLA from first principles. If you want to understand why each architectural choice exists — the mathematical motivation, the information-theoretic justification, the empirical evidence from DeepSeek-V2/V3, and the trade-offs — read this section carefully. The code walkthrough follows after the conceptual foundation.

4.1 The Problem: Why Standard Attention Doesn't Scale#

4.1.1 KV Cache Memory Bottleneck#

In standard Multi-Head Attention (MHA), for a sequence of length T, model dimension d, n_heads heads, and head_dim = d / n_heads:

code
KV cache per layer = 2 × n_heads × head_dim × T × sizeof(dtype)

At HyMo's configuration (n_heads=16, head_dim=128, T=4096, BF16):

code
= 2 × 16 × 128 × 4096 × 2 bytes = 33,554,432 bytes ≈ 32 MB per layer

With 8 MLA layers: 256 MB total KV cache — just for the attention state. This doesn't include activations, gradients, or optimizer states.

For inference at longer contexts (e.g., T=32K), this becomes 2 GB — a major bottleneck for serving.

4.1.2 Quadratic Complexity#

Standard attention computes:

code
Attention(Q, K, V) = softmax(QK^T / √d) V

The QK^T matrix is (T × T), giving O(T²) time and memory complexity. At T=4096, that's ~16M pairwise interactions per head per layer. For 16 heads across 8 layers: ~2 billion operations per forward pass.

4.1.3 The Low-Rank Hypothesis (DeepSeek-V2 Key Finding)#

DeepSeek-V2 (Dai et al., 2024) observed empirically that the KV activation matrix is approximately low-rank. Specifically, for a layer's key/value projections:

code
K ∈ R^{T × (n_heads × head_dim)},  V ∈ R^{T × (n_heads × head_dim)}

The effective rank of K and V is much smaller than their full dimension. This means most of the information in K and V can be compressed into a much smaller latent space with minimal loss.

Intuition: Language has strong local and hierarchical structure. Tokens that are close in context share similar representations. The "degrees of freedom" in the KV space are far fewer than the raw dimension suggests.

4.2 MLA: Low-Rank KV Compression from First Principles#

4.2.1 Mathematical Formulation#

MLA factorizes the KV projection into a low-rank bottleneck:

code
Standard MHA:    K = X W_K,  V = X W_V     where W_K, W_V ∈ R^{d × (n_heads × head_dim)}
MLA:             K = X W_kv_a W_kv_b,  V = X W_kv_a W_kv_b

But with a critical modification: joint compression with decoupled RoPE.

Let d = 896 (model dim), kv_lora_rank = 128 (latent dim), n_kv_groups = 4 (MQA groups).

code
wkv_a: R^d → R^{kv_lora_rank + qk_rope_head_dim}    # 896 → 128 + 32 = 160
   └─> kv_latent (128) + k_pe (32)                   # Split after projection

kv_norm: RMSNorm on kv_latent                        # Stabilize latent scale

wkv_b: R^{kv_lora_rank} → R^{n_kv_groups × (qk_nope + v_head_dim)}  # 128 → 4 × (96 + 128) = 896
   └─> k_nope (96 per group) + v (128 per group)

Why this factorization?

AspectStandard MHAMLA
KV projection params2 × d × (n_h × h_d) = 1.84Md × 160 + 128 × 896 = 257K
KV cache per token2 × n_h × h_d = 4,096kv_lora_rank + n_kv_groups × qk_rope = 128 + 128 = 256
Compression ratio1×16×

The 16× cache reduction comes from storing only the 128-dim latent + 32-dim RoPE key per token, instead of full 2,048-dim K and V.

4.2.2 Decoupled RoPE: Why Separate Position from Content?#

Standard RoPE applies rotation to the full key/query:

code
RoPE(q) = q ⊙ cos + rotate(q) ⊙ sin

Problem: The rotation mixes position and content information. For low-rank compression, this is suboptimal — the position information "pollutes" the content space that we're trying to compress.

MLA's solution: Decouple the position encoding:

  • Content path: kv_latent (128-dim) — pure content, no position info, compressed
  • Position path: k_pe (32-dim) — pure position, shared across all KV groups, uncompressed
code
K_full = concat(k_nope, k_pe_rotated)  # 96 + 32 = 128 per group

This means:

  1. Content compression is clean — kv_latent contains only semantic information
  2. Position is explicit — k_pe is a dedicated position encoding subspace
  3. Shared computation — One k_pe per token serves all 4 KV groups (broadcast)
Design Decision (from DeepSeek-V2): The decoupled RoPE is applied to a fixed 25% of head_dim (32 of 128). This 25% ratio was empirically found optimal — large enough to encode position, small enough to leave 75% for content.

4.2.3 MQA-4: Multi-Query Attention with 4 Groups#

Standard attention variants:

VariantKV HeadsQuery HeadsRatioNotes
MHA16161:1Full capacity, max memory
GQA8161:2Good balance
MQA-44161:4MLA choice
MQA1161:16Max compression, potential quality loss

Why MQA-4 for MLA?

  1. Compression synergy: MLA already compresses KV to latent. Fewer KV groups = less decompression work
  2. Hardware efficiency: 4 groups maps well to GPU warp/wavefront boundaries
  3. Quality preservation: 4 groups is the "sweet spot" from DeepSeek-V3 ablation — 1 group (MQA) loses too much capacity, 8 groups (GQA) defeats the compression purpose

Each KV group serves n_heads / n_kv_groups = 16/4 = 4 query heads. This is implemented via Grouped Query Attention (GQA) in SDPA:

  • Q: (B, 16, T, 128)
  • K: (B, 4, T, 128) → broadcast to 16 heads internally
  • V: (B, 4, T, 128) → broadcast to 16 heads internally

4.3 Query Path: Low-Rank with Per-Head Scaling#

4.3.1 Two-Stage Projection#

code
wq_a: R^d → R^{q_lora_rank}              # 896 → 224 (bottleneck)
q_norm: RMSNorm on q_lora_rank           # Per-token normalization
wq_b: R^{q_lora_rank} → R^{n_heads × head_dim}  # 224 → 2,048

Why two stages?

  • Parameter efficiency: 896×224 + 224×2048 = 659K vs direct 896×2048 = 1.84M (2.8× savings)
  • The bottleneck forces the model to learn a compressed query representation
  • RMSNorm after bottleneck stabilizes the scale before expansion

4.3.2 Per-Dimension Learnable Scaling (q_norm_qk, k_norm_qk)#

After projecting to per-head space, MLA applies element-wise learnable scales:

python
q = q.view(B, T, H, head_dim) * q_norm_qk.view(1, 1, H, head_dim)
# q_norm_qk ∈ R^{n_heads × head_dim} = R^{2048}

And similarly for K's RoPE portion:

python
k_pe_normed = k_pe * k_norm_qk.view(1, 1, G, D_pe)
# k_norm_qk ∈ R^{n_kv_groups × qk_rope_head_dim} = R^{128}

Why per-dimension scaling?

  1. Attention logit control: The softmax in attention is sensitive to logit magnitude. These scales let the model learn which dimensions should contribute more/less to attention scores
  2. Stabilization: Initialized to 1 (identity), they provide a "temperature" per dimension that can adapt during training
  3. DeepSeek finding: This simple trick consistently improves training stability and final perplexity across scales
Analogy: Think of this as a learned "attention head importance weighting" but at the dimension level rather than head level.

4.4 Complete Forward Pass: Mathematical Trace#

Let's trace through with concrete shapes at HyMo config:

  • B=4, T=4096, d=896
  • H=16, G=4, D_pe=32, D_nope=96, D_v=128
  • q_lora=224, kv_lora=128

Step 1: Query Computation#

code
x: (B, T, d) = (4, 4096, 896)

q_latent = RMSNorm(wq_a(x))           # (4, 4096, 224)
q = wq_b(q_latent)                    # (4, 4096, 2048)
q = q.view(4, 4096, 16, 128)          # (B, T, H, head_dim)
q = q * q_norm_qk                     # Element-wise scale (4, 4096, 16, 128)

q_pe = q[..., :32]                    # (4, 4096, 16, 32) — position dims
q_nope = q[..., 32:]                  # (4, 4096, 16, 96) — content dims

q_pe_rot = RoPE(q_pe)                 # (4, 4096, 16, 32) — apply rotation

Step 2: KV Compression#

code
kv = wkv_a(x)                         # (4, 4096, 160)
kv_latent = kv[..., :128]             # (4, 4096, 128) — compressed content
k_pe = kv[..., 128:]                  # (4, 4096, 32) — position key

kv_latent = RMSNorm(kv_latent)        # (4, 4096, 128)

kv_out = wkv_b(kv_latent)             # (4, 4096, 896)
kv_out = kv_out.view(4, 4096, 4, 224) # (B, T, G, D_nope + D_v)

k_nope = kv_out[..., :96]             # (4, 4096, 4, 96)
v = kv_out[..., 96:]                  # (4, 4096, 4, 128)

Step 3: Position Key Processing (Decoupled RoPE)#

code
# Broadcast k_pe to all 4 KV groups with per-group scaling
k_pe_normed = k_pe.unsqueeze(2) * k_norm_qk    # (4, 4096, 1, 32) * (1, 1, 4, 32)
                                         # → (4, 4096, 4, 32) via broadcast

k_pe_rot = RoPE(k_pe_normed)        # (4, 4096, 4, 32)

Step 4: Assemble for SDPA#

code
q_assembled = concat(q_pe_rot, q_nope, dim=-1)   # (4, 4096, 16, 128)
k_assembled = concat(k_pe_rot, k_nope, dim=-1)   # (4, 4096, 4, 128)

# SDPA expects (B, n_heads, T, head_dim)
q_sdpa = q_assembled.permute(0, 2, 1, 3)  # (4, 16, 4096, 128)
k_sdpa = k_assembled.permute(0, 2, 1, 3)  # (4, 4, 4096, 128)
v_sdpa = v.permute(0, 2, 1, 3)            # (4, 4, 4096, 128)

Step 5: Grouped Query Attention#

code
# Native GQA (PyTorch 2.5+) handles 4 KV heads → 16 Q heads internally
out = scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa, enable_gqa=True)
# out: (4, 16, 4096, 128)

Step 6: Output Projection#

code
out = out.permute(0, 2, 1, 3).contiguous().view(4, 4096, 2048)
y = wo(out)                               # (4, 4096, 896)

4.5 Design Decision Deep Dives#

4.5.1 Why kv_lora_rank = 128?#

The latent dimension is the information bottleneck. Too small → information loss. Too large → defeats compression.

DeepSeek-V2/V3 ablation results:

kv_lora_rankKV Cache (MB/layer)PPL (validation)
6416+0.15
12832baseline
25664-0.02
512128-0.01

128 is the knee of the curve — diminishing returns beyond this point. HyMo uses 128.

4.5.2 Why q_lora_rank = 224?#

Query compression is less aggressive because:

  • Queries are not cached (only K,V are)
  • Queries need to discriminate across all positions — more capacity helps
  • 224 = 22.5% of full query dim (2048) — enough for rich query representations

4.5.3 Why Partial RoPE (25%)?#

Full RoPE on all dimensions would:

  • Make KV compression harder (position + content mixed)
  • Waste capacity on position encoding

No RoPE would:

  • Break position awareness for long-range dependencies

25% (32/128) is the empirically optimal trade-off from DeepSeek-V2:

  • Sufficient for position encoding (32 dims can encode 4096+ positions via RoPE)
  • Leaves 75% for content discrimination

4.5.4 Why RMSNorm on Latents?#

code
q_latent = RMSNorm(wq_a(x))
kv_latent = RMSNorm(kv_latent)

RMSNorm (Root Mean Square Normalization) on the bottleneck representations:

  1. Prevents scale explosion through the bottleneck
  2. Stabilizes gradients — the bottleneck is a gradient highway
  3. No learnable parameters (unlike LayerNorm) — just per-feature scaling

4.6 MLA Block: Attention + MoE with Pre-Norm Residuals#

The MLABlock wraps the attention with a DeepSeekMoE feed-forward:

python
class MLABlock(nn.Module):
    def __init__(self, config, layer_idx=0):
        self.attn_norm = RMSNorm(config.dim)           # Pre-norm for attention
        self.attn = MultiHeadLatentAttention(config, layer_idx)
        self.moe_norm = RMSNorm(config.dim)            # Pre-norm for MoE
        self.moe = DeepSeekMoE(config, layer_idx)      # 16 experts + 1 shared

    def forward(self, x):
        x = x + self.attn(self.attn_norm(x))    # Attention residual
        x = x + self.moe(self.moe_norm(x))      # MoE residual
        return x

Why MoE only on MLA blocks?

ReasonExplanation
Routing overheadMoE dispatch is expensive. MLA is already compute-heavy — amortizes the overhead
Routing noiseGDN has recurrent state; routing errors compound. MLA is stateless — errors are local
Capacity allocationMLA layers need more FFN capacity for precise retrieval; GDN does cheap aggregation

4.7 Code Walkthrough (src/hymo/models/mla.py)#

Now with the mathematical foundation, let's walk the actual implementation:

4.8 Summary: MLA Configuration in HyMo v1.0#

ParameterValueRationale
n_heads16Standard for 896-dim
n_kv_groups4MQA-4: compression/quality sweet spot
head_dim128896/16 = 56? No — expanded to 128 for capacity
qk_rope_head_dim3225% partial RoPE
qk_nope_head_dim9675% content
v_head_dim128Full value dimension
q_lora_rank224Query bottleneck (2.8× param savings)
kv_lora_rank128KV bottleneck (16× cache reduction)
rope_theta10,000Standard RoPE base
max_seq_len4,096Training context

5. Gated Delta Net Linear Attention (gdn.py) — Mathematical Foundations#

File: src/hymo/models/gdn.py

This section builds GDN from first principles. The gated delta rule is the core innovation that makes linear attention competitive with full attention. We derive the recurrence, explain the information-theoretic properties, and show why the read/write key decoupling is strictly more expressive than standard linear attention.

5.1 The Problem: Linear Attention vs Quadratic Attention#

5.1.1 Standard (Quadratic) Attention#

Standard self-attention computes:

code
Attention(Q, K, V) = softmax(QK^T / √d) V

Complexity: O(T² × d) compute, O(T × d) memory for QK^T

Expressiveness: Can attend to any past position with arbitrary weights.

Limitation: Quadratic scaling makes long contexts expensive.

5.1.2 Linear Attention#

Linear attention replaces the softmax with a kernel trick:

code
LinearAttn(Q, K, V) = ϕ(Q) (ϕ(K)^T V)

Where ϕ is a kernel feature map. This can be computed as:

code
S_t = S_{t-1} + ϕ(k_t) v_t^T    ← State update
o_t = ϕ(q_t)^T S_t               ← Output

Complexity: O(T × S²) where S is the state dimension (S=32 in HyMo)

Expressiveness: Limited by the state dimension S. Cannot selectively attend to arbitrary past positions.

5.1.3 The Key Limitation of Linear Attention#

Linear attention's state S_t ∈ R^{S×S} is a fixed-size summary of all past tokens. Once information is compressed into S, it cannot be selectively retrieved — the read is a fixed linear combination.

Analogy: Linear attention is like writing notes in a notebook (state), then summarizing each page. Standard attention is like having the full book available to search.

5.2 The Gated Delta Rule: A Better Recurrence#

5.2.1 The Recurrence#

For each head at position t:

code
H_t = exp(g_t · A) ⊙ H_{t-1}  +  b_t ⊗ v_t    ← State update (delta rule)
o_t = c_t^T · H_t                               ← Read from state

Where:

  • H_t ∈ R^{S × D} — hidden state matrix (S=d_state=32, D=headdim=32)
  • b_t ∈ R^S — write key (what to write to state)
  • c_t ∈ R^S — read key (what to read from state)
  • v_t ∈ R^D — value vector
  • g_t ∈ R — per-head scalar gate (output of sigmoid(dt_proj))
  • A ∈ R^S — negative decay vector (-exp(A_log), making it strictly negative)
  • exp(g_t · A) ⊙ — element-wise decay: each state dimension decays independently

5.2.2 Why "Delta Rule"?#

The name comes from the delta rule in online learning:

code
Δw = η × (target - prediction) × input

In GDN:

code
ΔH = b_t ⊗ v_t    ← outer product of write key and value

This is the "delta" — the change to the state matrix. The write key b_t determines where in the state to write, and the value v_t determines what to write.

5.2.3 Why Read/Write Decoupling Matters#

Standard linear attention:

code
S_t = S_{t-1} + k_t v_t^T    ← same k for read and write
o_t = q_t^T S_t               ← same q for read and write

Gated delta net:

code
H_t = α H_{t-1} + b_t ⊗ v_t  ← b for write
o_t = c_t^T H_t               ← c for read

The key difference: b and c are independent learned projections. This means:

  1. Write with one pattern, read with another: The model can learn to write information using one representation and retrieve it using a completely different representation
  2. Selective erasure: The gate g_t can selectively decay specific state dimensions while preserving others
  3. Expressiveness: This is strictly more expressive than standard linear attention (where b=c=k)
Formal argument: Standard linear attention is a special case of GDN where b = c = k and g = 1 (no gating). GDN generalizes this by allowing independent read/write keys and learned gating.

5.3 Mathematical Properties#

5.3.1 Decay and Forgetting#

The decay term exp(g_t · A) provides exponential forgetting:

code
H_t = exp(g_t · A) ⊙ H_{t-1} + b_t ⊗ v_t

Unrolling the recurrence:

code
H_t = exp(g_t · A) ⊙ exp(g_{t-1} · A) ⊙ ... ⊙ exp(g_1 · A) ⊙ H_0
    + exp(g_t · A) ⊙ exp(g_{t-1} · A) ⊙ ... ⊙ exp(g_2 · A) ⊙ b_1 ⊗ v_1
    + ...
    + exp(g_t · A) ⊙ b_{t-1} ⊗ v_{t-1}
    + b_t ⊗ v_t

The contribution of token at position k < t to the current state is:

code
weight(k, t) = ∏_{j=k+1}^{t} exp(g_j · A)
             = exp(Σ_{j=k+1}^{t} g_j · A)

This is a learned exponential decay — the model controls how fast past information fades.

5.3.2 Multi-Scale Memory#

The A_log parameter is initialized with multi-scale decay:

python
a_init = log(arange(1, n_heads + 1))  # [log(1), log(2), ..., log(40)]
# Head 0: A = -1, exp(-1) = 0.37 (remembers ~3 steps)
# Head 10: A = -10, exp(-10) ≈ 0.000045 (forgets almost instantly)
# Head 39: A = -39, exp(-39) ≈ 10^{-17} (remembers only current step)

Why multi-scale? Different heads specialize for different timescales:

  • Slow heads (small |A|): Remember distant context (coreference, topic)
  • Fast heads (large |A|): Focus on immediate context (syntax, local semantics)

5.3.3 State Capacity#

The state H_t ∈ R^{S×D} has S×D = 32×32 = 1,024 parameters per head.

Information-theoretic bound: The state can store at most 1,024 independent scalar values. With 40 heads, total state capacity is 40 × 1,024 = 40,960 scalars.

Compare to full attention: At T=4096, full attention stores 4,096 × 128 = 524,288 scalars per head (QK^T matrix). GDN's 1,024 is ~500× smaller.

The trade-off: GDN is much faster but stores less information. The 3:1 ratio ensures MLA provides the "full book" access when needed.

5.4 The GDN Block Architecture#

5.4.1 Input Projection#

python
self.in_proj = nn.Linear(d_model, d_inner, bias=False)     # 896 → 1280
self.conv1d = nn.Conv1d(d_inner, d_inner, d_conv, groups=d_inner,
                         padding=d_conv - 1, bias=False)    # Depthwise causal conv

Why 1280? The GDN paper recommends d_inner / d_model ≈ 1.43. At d_model=896:

code
d_inner = 896 × 1.43 ≈ 1280

This ensures the GDN path is not under-provisioned relative to MLA and MoE.

Why depthwise conv1d? Provides local context before the global recurrence. Each channel gets its own 1D convolution (no cross-channel mixing), adding ~0 parameters but capturing 4-token local patterns.

5.4.2 The Projections#

python
self.b_proj = nn.Linear(d_inner, n_heads * d_state, bias=False)   # 1280 → 40×32 = 1280
self.c_proj = nn.Linear(d_inner, n_heads * d_state, bias=False)   # 1280 → 1280
self.dt_proj = nn.Linear(d_inner, n_heads, bias=False)            # 1280 → 40
self.g_proj = nn.Linear(d_inner, d_inner, bias=False)              # 1280 → 1280
self.out_proj = nn.Linear(d_inner, d_model, bias=False)           # 1280 → 896
self.skip_proj = nn.Linear(d_model, d_model, bias=False)          # 896 → 896

Why separate b_proj and c_proj? The decoupled read/write keys are the core innovation. If they were the same (like standard linear attention), we'd lose the expressiveness advantage.

Why g_proj? The gate g_t controls the decay rate. It's computed from the input via a learned projection, then passed through sigmoid to get a value in [0, 1].

5.4.3 The Gate Computation#

python
g_in = self.g_proj(v)                    # (B,T,1280)
g = F.sigmoid(g_in)                      # (B,T,1280) — element-wise in [0,1]
g_gate = g.view(B, T, H, D).mean(dim=-1) # (B,T,40) — scalar per head

Why mean over D? The gate is per-head, not per-dimension. Averaging over the D=32 dimensions collapses the per-head gate to a single scalar.

Why sigmoid? The gate modulates the decay rate:

  • g → 0: Slow decay (remember more)
  • g → 1: Fast decay (forget more)
  • g = 0.5: Moderate decay

The sigmoid ensures the gate is in [0, 1], which is important for training stability.

5.5 Complete Forward Pass Trace#

Let's trace with concrete shapes at HyMo config:

  • B=4, T=4096, d_model=896, d_inner=1280, S=32, D=32, H=40

Step 1: Input Projection#

code
x: (4, 4096, 896)
v_in = in_proj(x)                        # (4, 4096, 1280)

Step 2: Causal Depthwise Conv1d#

code
v_conv = conv1d(v_in.transpose(1, 2))    # (4, 1280, 4096+3)
v_conv = v_conv[:, :, :4096]             # (4, 1280, 4096) — trim padding
v = SiLU(v_conv.transpose(1, 2))         # (4, 4096, 1280)

Step 3: Compute Projections#

code
b_in = b_proj(v)     # (4, 4096, 1280) → reshape to (4, 4096, 40, 32)
c_in = c_proj(v)     # (4, 4096, 1280) → reshape to (4, 4096, 40, 32)
dt_in = dt_proj(v)   # (4, 4096, 1280) → (4, 4096, 40)
g_in = g_proj(v)     # (4, 4096, 1280)
g = sigmoid(g_in)    # (4, 4096, 1280)
g_gate = g.view(4, 4096, 40, 32).mean(dim=-1)  # (4, 4096, 40)

Step 4: RoPE on Value#

code
v = v.view(4, 4096, 40, 32)              # (4, 4096, 40, 32)
# RoPE applied to full 32 dims (rope_dim == D)
v_rot = RoPE(v)                          # (4, 4096, 40, 32)

Step 5: Gated Delta Rule Recurrence#

The recurrence (PyTorch fallback):

python
h = zeros(B, H, S, D, dtype=float32)     # State: (4, 40, 32, 32)
for t in range(T):
    alpha = exp(g_t * A)                  # Decay factor: (4, 40, 32)
    h = alpha.unsqueeze(-1) * h + b_t.unsqueeze(-1) * v_t.unsqueeze(-2)
    o_t = einsum("bhs,bhsd->bhd", c_t, h)

Step 6: Skip Connection and Gating#

code
o = o + D.view(1, 1, H, 1) * v           # Skip: (4, 4096, 40, 32)
o = o * g_gate.unsqueeze(-1)             # Gate: (4, 4096, 40, 32)

Step 7: Output Projection#

code
o_flat = o.view(4, 4096, 1280)           # (4, 4096, 1280)
y = out_proj(o_flat) + skip_proj(x)      # (4, 4096, 896) + residual

5.6 The Skip Connection and D Parameter#

python
self.D = nn.Parameter(torch.ones(n_heads))    # (40,) — initialized to 1

The skip connection adds the original value directly to the output:

code
o_t = c_t^T H_t + D · v_t

Why? This ensures the block can pass information directly without going through the state. The D parameter controls how much of the direct signal to include.

Initialization to 1: At initialization, the block is approximately an identity function:

code
o_t ≈ v_t    (since H_0 = 0 and D = 1)

This is a residual-like initialization that ensures stable training at the start.

5.7 Code Walkthrough (src/hymo/models/gdn.py)#

6. Custom Triton GDN Kernel (gdn_triton.py) — Kernel Design & Memory Optimization#

File: src/hymo/models/gdn_triton.py

Why a custom kernel? The PyTorch fallback loops over T=4096 positions, launching thousands of tiny CUDA kernels. Each kernel launch has ~10μs overhead. At 4,096 launches × 40 heads × 4 batches = 655,360 launches, that's ~6.5 seconds per layer. The Triton kernel fuses the entire recurrence into one kernel launch per (batch, head).

6.1 The Problem: Kernel Launch Overhead#

PyTorch fallback:

python
for t in range(T):           # 4096 iterations
    h = alpha * h + b ⊗ v   # 1 kernel launch
    o_t = c^T · h           # 1 kernel launch

Each iteration launches ~2 CUDA kernels. Total: ~8,192 kernels per head per forward pass.

With 40 heads × 4 batches:

code
8,192 kernels × 40 heads × 4 batches = 1,310,720 kernel launches
1,310,720 × 10μs = 13.1 seconds per layer

With Triton kernel:

code
1 kernel launch per (batch, head) = 160 launches
160 × 10μs = 1.6ms per layer

Speedup: 13.1s → 1.6ms = ~8,000× faster (kernel launch overhead eliminated).

6.2 Triton Architecture: 2D Grid#

The kernel uses a 2D grid of programs:

python
grid = (B, H)  # (4, 40) = 160 programs

Each program processes one (batch, head) pair across all T positions.

code
┌─────────────────────────────────────────────────────────────┐
│  Grid: (B=4, H=40)                                          │
│                                                              │
│  (batch=0, head=0)  (batch=0, head=1)  ... (batch=0, head=39)│
│  (batch=1, head=0)  (batch=1, head=1)  ... (batch=1, head=39)│
│  (batch=2, head=0)  (batch=2, head=1)  ... (batch=2, head=39)│
│  (batch=3, head=0)  (batch=3, head=1)  ... (batch=3, head=39)│
└─────────────────────────────────────────────────────────────┘

Each program runs a sequential loop over T positions, keeping the state h in SRAM registers.

6.3 State in Registers#

python
h = tl.zeros((S, D), dtype=tl.float32)  # S=32, D=32

Memory footprint: 32 × 32 × 4 bytes = 4 KB per program.

Why registers? SRAM is ~100× faster than HBM (global memory). The state h is accessed every iteration — keeping it in registers eliminates all memory traffic for state reads/writes.

SRAM vs HBM:

MemoryBandwidthLatencyCapacity
SRAM (registers)~20 TB/s~1 cycle~256 KB per SM
HBM (global)~3 TB/s~100 cycles80 GB

6.4 The Forward Kernel#

Key operations per iteration:

  1. Load v_t, b_t, c_t, g_t from HBM (4 memory reads)
  2. Compute alpha = exp(g_t * A) (vectorized)
  3. Update h (element-wise multiply + outer product + add)
  4. Read o_t = c_t^T · h (dot product)
  5. Store o_t to HBM (1 memory write)

Total per iteration: 4 reads + 1 write = 5 memory operations, all vectorized.

6.5 The Backward Kernel#

The backward kernel runs backwards in time (recurrent backprop):

Key insight: The backward pass is also a recurrence — gradients flow backwards through time, just like in BPTT (Back-Propagation Through Time) for RNNs.

Memory cost: The forward states h_out must be saved for backward:

code
h_out: B × T × H × S × D = 4 × 4096 × 40 × 32 × 32 × 4 bytes
     = 2.6 GB per layer

This is the activation memory cost of the fused kernel. The alternative (recomputing in backward) would save memory but double compute.

6.6 The TritonGDNFunction (Autograd Wrapper)#

Why custom autograd? PyTorch's autograd doesn't know about the recurrence. By implementing torch.autograd.Function, we provide custom forward/backward that the Triton kernel implements efficiently.

6.7 Padding for Triton Constraints#

Triton requires tensor dimensions to be powers of 2 for efficient SRAM access:

python
D_pad = _next_power_of_2(D)  # 32 → 32 (already power of 2)
S_pad = _next_power_of_2(S)  # 32 → 32 (already power of 2)

def _pad(t, target_last):
    pad = [0] * (2 * t.ndim)
    pad[1] = target_last - t.shape[-1]
    return F.pad(t, pad)

v_p = _pad(v.float().contiguous(), D_pad)
b_p = _pad(b.float().contiguous(), S_pad)
# ... run kernel ...
out = out_p[..., :D]  # strip padding

With D=32, S=32 in the production config, both are already powers of 2, so no padding occurs. This is a deliberate design choice — the config values were chosen to be Triton-friendly.

6.8 Performance Characteristics#

MetricPyTorch FallbackTriton KernelSpeedup
Kernel launches~8,192 per head1 per head~8,000×
HBM readsT × (S + S + D + 1)T × (S + S + D + 1)Same
HBM writesT × DT × D + T × S × D1.5× more
State accessHBM (every step)SRAM (registers)~100×
Total time~13s per layer~1.6ms per layer~8,000×

The Triton kernel trades more HBM writes (saving full state trajectory) for eliminated kernel launch overhead and SRAM-resident state. This is a net win because kernel launch overhead dominates at T=4096.

7. Rotary Position Embedding (rope.py) — Mathematical Foundations & Design#

File: src/hymo/models/rope.py

This section derives RoPE from first principles. The key insight: position information can be encoded as rotations in embedding space, preserving distances between vectors while making them position-dependent. We show why rotation is the optimal choice, how partial RoPE works, and the difference between MLA and GDN usage.

7.1 The Problem: Encoding Position#

7.1.1 Why Position Matters#

Standard attention is permutation-invariant — without position information, the model cannot distinguish:

code
The cat sat on the mat
The mat sat on the cat

Both sentences produce the same attention scores without position encoding.

7.1.2 Approaches to Position Encoding#

ApproachMethodProsCons
AbsoluteAdd position to embeddingSimpleFixed max length, no relative position
RelativeAdd relative distance to attention scoresRelative position awareComplex implementation
Rotary (RoPE)Rotate embedding pairs by positionRelative + absolute, length-extrapolatableRequires dimension pairing

HyMo uses RoPE for all position encoding (MLA and GDN).

7.2 RoPE: Derivation from First Principles#

7.2.1 The Core Idea#

For a position p and dimension pair (2k, 2k+1), apply a 2D rotation:

code
R(p, θ_k) = [cos(p·θ_k), -sin(p·θ_k)]
             [sin(p·θ_k),  cos(p·θ_k)]

Where θ_k = 1 / (base^(2k / head_dim)) is the frequency for dimension pair k.

Applied to a vector x:

code
[x_{2k}']   [cos(p·θ_k)  -sin(p·θ_k)] [x_{2k}]
[x_{2k+1}'] = [sin(p·θ_k)   cos(p·θ_k)] [x_{2k+1}]

7.2.2 Why Rotation Preserves Distance#

The rotation matrix R is orthogonal: R^T R = I. Therefore:

code
‖R(p) x - R(p) y‖² = ‖x - y‖²

The distance between two vectors is preserved after rotation. This means:

  • Relative position is encoded: The dot product x · R(Δp) y depends on the relative distance Δp = p_q - p_k
  • Absolute position is encoded: Each vector is rotated by its own position p

7.2.3 The Frequency Spectrum#

code
θ_k = 1 / base^(2k / head_dim)

At base=10,000, head_dim=32:

code
k=0:  θ_0 = 1 / 10000^0   = 1.0       (slowest rotation)
k=1:  θ_1 = 1 / 10000^0.5 = 0.01      (fast rotation)
k=2:  θ_2 = 1 / 10000^1.0 = 0.0001    (very fast)
...

Intuition: Low-frequency components (small k) encode long-range position relationships. High-frequency components (large k) encode short-range local position.

7.3 The Rotation Matrix in Code#

7.3.1 Precomputation#

python
# Precompute frequencies
freqs = 1.0 / (theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim))
# freqs: (head_dim/2,) = (16,) — one frequency per pair

# Precompute angles
positions = torch.arange(max_seq_len, dtype=torch.float32)
angles = torch.outer(positions, freqs)  # (max_seq_len, head_dim/2)

# Build cos/sin tables
cos_tab = angles.cos()  # (max_seq_len, head_dim/2)
sin_tab = angles.sin()  # (max_seq_len, head_dim/2)

Why precompute? The cos/sin tables are the same for every forward pass. Computing them once in __init__ saves T × head_dim/2 trigonometric operations per forward.

7.3.2 Applying the Rotation#

The rotation formula:

code
[x_even']   [cos(θ)  -sin(θ)] [x_even]
[x_odd' ] = [sin(θ)   cos(θ)] [x_odd ]

This is the standard 2D rotation, applied independently to each dimension pair.

7.4 Key Properties#

7.4.1 Norm Preservation#

code
‖R(p) x‖² = ‖x‖²

The rotation doesn't change the vector's magnitude — only its direction. This is important for training stability.

7.4.2 Relative Position Dependence#

The dot product of two rotated vectors depends only on their relative position:

code
R(p_q) x · R(p_k) y = x · R(p_q - p_k) y

This means the attention score between positions q and k depends on q - k, not on q or k individually.

7.4.3 Length Extrapolation#

Because RoPE uses relative positions, it can extrapolate to sequences longer than max_seq_len:

  • At max_seq_len=4096, the model trains on positions 0-4095
  • At inference, position 4096 produces a rotation that's "one step beyond" the training distribution
  • This is smoother than absolute position encoding, which breaks completely at max_seq_len + 1

7.5 Partial RoPE: Position/Content Decoupling#

7.5.1 The Problem with Full RoPE#

Applying RoPE to the full head_dim means:

code
K_full = R(p) K_content

The position and content information are mixed in the same vector. For low-rank compression (MLA), this is suboptimal — the compression must handle both position and content simultaneously.

7.5.2 The Solution: Partial RoPE#

Split the head_dim into position and content subspaces:

code
head_dim = [qk_rope_head_dim | qk_nope_head_dim]
           [      32          |       96         ]

K = [K_pe, K_nope]
K_pe = R(p) K_pe      ← Only position subspace gets rotated
K_nope = K_nope        ← Content subspace is untouched

In MLA (HyMo):

code
q_pe: 32 dims → rotated
q_nope: 96 dims → untouched
k_pe: 32 dims → rotated (shared across 4 KV groups)
k_nope: 96 dims → untouched

In GDN (HyMo):

code
v: 32 dims → FULLY rotated (rope_dim == head_dim)

7.5.3 Why 25% (32/128)?#

The ratio is empirically optimized by DeepSeek-V2:

  • Too small (<10%): Insufficient position information
  • Too large (>50%): Position information pollutes content space
  • 25% (32/128): Optimal trade-off for KV compression

7.6 MLA vs GDN: Different RoPE Strategies#

AspectMLAGDN
What gets rotatedq_pe (32 dims) + k_pe (32 dims)v (all 32 dims)
Fraction rotated25% of head_dim100% of head_dim
WhyDecouple position from content for compressionGDN has no Q/K — position goes on value
Shared computationk_pe shared across 4 KV groupsNo sharing needed

Why GDN rotates the value? GDN doesn't have separate Q and K projections — it uses read/write keys b and c. Position information must be injected somewhere, and the value vector is the natural choice because:

  1. Values carry the "content" that gets stored in state
  2. Rotating values makes the state position-aware
  3. The 32-dim head_dim matches the RoPE dimension exactly

7.7 Code Walkthrough (src/hymo/models/rope.py)#

Key design decisions:

  1. register_buffer — not parameters. RoPE tables are computed once and never trained.
  2. persistent=False — not saved in checkpoints (recomputed on load).
  3. start_pos — supports variable-position RoPE for inference (e.g., prefix caching).

8. Mixture of Experts (moe.py) — Routing Mathematics & Load Balancing#

File: src/hymo/models/moe.py

This section explains the MoE routing mechanism, the auxiliary-loss-free load balancing, and why MoE is restricted to MLA layers. The routing decision is a critical architectural choice that affects both model quality and training stability.

8.1 The MoE Hypothesis#

8.1.1 Why MoE?#

Standard dense FFN applies the same transformation to every token:

code
FFN(x) = W_2 · SiLU(W_1 · x)    # Same W_1, W_2 for all tokens

MoE applies different transformations to different tokens:

code
MoE(x) = Σ_i routing(x, i) · Expert_i(x)    # Different experts for different tokens

The intuition: Language has different types of tokens (nouns, verbs, numbers, code, etc.). Specialized experts can handle different token types better than a single shared expert.

8.1.2 The Capacity Argument#

Dense FFN has fixed capacity: every token uses every parameter. MoE has adaptive capacity: tokens are routed to the experts that are most relevant.

At 434M active / 1.13B stored:

  • Dense FFN: 750M params × every token = 750M effective
  • MoE FFN: 1.13B params × (2/16) per token = 141M active + shared

The 1.13B stored params give the model 2.6× more total capacity while only using 148.6M active MoE params per token.

8.2 The Routing Mechanism#

8.2.1 Gate Network#

python
self.gate = nn.Linear(config.dim, self.n_routed, bias=True)  # 896 → 16

The gate computes routing logits:

code
logits = W_gate · x + b_gate    # (B, T, 896) → (B, T, 16)
probs = softmax(logits)         # Probabilities over 16 experts

8.2.2 Top-k Selection#

python
k = min(self.n_activated, self.n_routed)  # 2
top_weights, top_indices = torch.topk(probs, k, dim=-1)

Each token selects the top-2 experts by probability. The remaining 14 experts receive zero contribution.

Why top-2?

  • Top-1: Too aggressive — tokens miss potentially useful experts
  • Top-2: Good balance — each token gets specialization from two experts
  • Top-4+: Diminishing returns, more compute per token

8.2.3 Capacity Capping#

python
capacity = int(self.capacity_factor * (B * T * k) / self.n_routed)
# capacity = 1.5 * (4 * 4096 * 2) / 16 = 768 tokens per expert

Why capacity capping?

  1. Static shapes for torch.compile: Without capping, each expert receives a variable number of tokens, making tensor shapes dynamic
  2. Memory bounds: Prevents one expert from receiving all tokens (which would require O(B×T) memory)
  3. Load balancing: Forces roughly equal distribution across experts

The trade-off: If an expert receives more than capacity tokens, the excess are silently dropped (contribute zero gradient). At capacity_factor=1.5, empirically <0.1% of tokens are dropped.

8.3 Why MoE Only on MLA Layers#

8.3.1 Routing Noise in Recurrent State#

In GDN layers, the recurrent state H_t accumulates information over time. If a token is routed to the wrong expert:

code
GDN + MoE:
  t=0: token → expert 3 (wrong) → corrupted output → stored in H_0
  t=1: H_1 = exp(g·A)·H_0 + ... → corrupted H_0 propagates
  t=2: H_2 = exp(g·A)·H_1 + ... → corruption compounds

The routing error persists in the recurrent state and corrupts all future predictions.

MLA + MoE:

code
MLA + MoE:
  t=0: token → expert 3 (wrong) → corrupted attention output
  t=1: MLA recomputes from scratch (no recurrence) → error is local

MLA's attention is stateless — routing errors are local to one position.

8.3.2 Overhead Amortization#

MoE dispatch has fixed overhead (routing, scatter-gather). This overhead is amortized on expensive layers:

Layer TypePer-Token CostMoE OverheadOverhead Fraction
MLA (5.8M params)High~9M params~15%
GDN (2.5M params)Low~9M params~36%

Putting MoE on GDN wastes 36% of the GDN compute budget on routing.

8.4 Aux-Loss-Free Load Balancing#

8.4.1 The Problem with Auxiliary Loss#

Traditional MoE uses an auxiliary loss L_aux that penalizes imbalanced routing:

code
L_aux = α × Σ_i (f_i × P_i)

Where f_i is the fraction of tokens routed to expert i, and P_i is the average routing probability for expert i.

Problem: L_aux competes with the main cross-entropy loss. Too high α hurts quality; too low α doesn't balance.

8.4.2 The Bias Adjustment Method#

HyMo instead directly adjusts gate biases based on EMA-tracked load:

Key properties:

  1. No gradient pollution: The bias adjustments don't affect the cross-entropy gradient
  2. Deadband: The 5% threshold (1.05/0.95) prevents oscillation
  3. Slow adaptation: ema_alpha=0.001 means the EMA updates slowly, avoiding instability

8.5 SwiGLU Expert Architecture#

python
class SwiGLUExpert(nn.Module):
    def __init__(self, dim, inter_dim):
        self.w1 = nn.Linear(dim, inter_dim)   # Gate projection: 896 → 2304
        self.w2 = nn.Linear(inter_dim, dim)   # Down projection: 2304 → 896
        self.w3 = nn.Linear(dim, inter_dim)   # Up projection: 896 → 2304

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

SwiGLU = Swish-gated Linear Unit:

code
SwiGLU(x) = W_2 · (SiLU(W_1 · x) ⊙ W_3 · x)

The gating mechanism (SiLU(W_1 · x) ⊙ W_3 · x) allows the expert to learn which input dimensions to activate. This is more expressive than standard ReLU FFN because the gate is input-dependent.

Parameter count per expert:

code
w1: 896 × 2304 = 2.06M
w2: 2304 × 896 = 2.06M
w3: 896 × 2304 = 2.06M
Total: 6.19M per expert

8.6 The Shared Expert#

python
if self.shared_expert is not None:
    out = out + self.shared_expert(x_flat)  # EVERY token goes through shared expert

The shared expert acts as a dense backbone:

  • Every token passes through it (no routing)
  • Captures universal grammatical patterns
  • Provides a baseline output that routed experts refine

Why shared + routed?

  • Shared expert handles common patterns (grammar, syntax)
  • Routed experts handle specialized patterns (factual knowledge, code)
  • The combination is more robust than either alone

8.7 FP32 Router Cast#

python
def gate_forward(self, x):
    x_fp32 = x.float()
    w_fp32 = self.gate.weight.float()
    b_fp32 = self.gate.bias.float()
    logits = F.linear(x_fp32, w_fp32, b_fp32)
    return logits.to(x.dtype)

Why FP32? In BF16 (7-bit mantissa), similar routing scores become indistinguishable:

code
BF16: 0.1234567 ≈ 0.1234570 (same value!)
FP32: 0.1234567 ≠ 0.1234570 (different values)

The FP32 upcast ensures routing decisions are precise. Since W_gate is 896×16 (~14K params), the compute cost is negligible.

8.8 Complete Forward Pass#

9. Multi-Token Prediction (mtp.py) — Auxiliary Training Signals#

File: src/hymo/models/mtp.py

MTP provides auxiliary training signals that improve hidden state quality. The key insight: predicting multiple future tokens forces the model to learn richer representations that capture both short-term and long-term dependencies.

9.1 Why MTP?#

9.1.1 The Single-Token Limitation#

Standard language modeling predicts one token at a time:

code
P(x_{t+1} | x_0, ..., x_t) = softmax(W_vocab · h_t)

The hidden state h_t is optimized only for predicting the very next token. It may ignore longer-term patterns needed for future predictions.

9.1.2 The Multi-Token Signal#

MTP predicts D tokens ahead:

code
P(x_{t+1} | h_t) = softmax(W_vocab · h_t)              ← Standard
P(x_{t+2} | h_t) = softmax(W_vocab · MTP_1(h_t))       ← +1
P(x_{t+3} | h_t) = softmax(W_vocab · MTP_2(h_t))       ← +2

Why this helps:

  1. Denser gradients: The hidden state receives gradient signals from D different predictions, not just one
  2. Longer-range planning: To predict token t+3, the model must understand structure beyond the immediate next token
  3. Better representations: The hidden state encodes information useful for multiple future tokens

9.2 The MTP Architecture#

9.2.1 MTPBlock#

python
class MTPBlock(nn.Module):
    def __init__(self, dim, inter_dim):
        self.w1 = nn.Linear(dim, inter_dim)   # Gate projection: 896 → 2304
        self.w2 = nn.Linear(inter_dim, dim)   # Down projection: 2304 → 896
        self.w3 = nn.Linear(dim, inter_dim)   # Up projection: 896 → 2304

    def forward(self, x):
        return self.w2(F.silu(self.w1(x)) * self.w3(x))

Each MTP head is a SwiGLU block — the same architecture used in MoE experts and Dense FFN layers. This reuse simplifies the implementation and keeps parameter counts consistent.

9.2.2 Forward Pass#

The MTP chain:

code
Head 0: h = block(main_hidden[t] + embed[tokens[t+1]])  → logits for token[t+2]
Head 1: h = block(head0_out[t] + embed[tokens[t+2]])    → logits for token[t+3]

Each head receives the previous head's hidden state, creating a chain of predictions. The sequence lengths shrink by 1 per head: head 0 predicts T-1 tokens, head 1 predicts T-2 tokens.

9.3 The Head Chaining Mechanism#

9.3.1 Why Chain?#

The alternative is independent heads — each head receives the original hidden state:

code
Independent: MTP_d(h_t) for all d
Chained:     MTP_d(MTP_{d-1}(...(h_t)...))  for all d

Chaining is better because:

  1. Higher depths see richer representations: Head 1 receives a representation already conditioned on token t+2
  2. Reduced redundancy: Independent heads may learn similar features
  3. Gradient flow: Chaining creates a path for gradients to flow through multiple heads

9.3.2 Shared Vocabulary Head#

python
logits = main.head(out)    # REUSE main model's head

MTP heads share the main model's output head (main.head). This means:

  • The MTP prediction is made in the same logit space as the main prediction
  • MTP gradients flow back through the shared head, improving its quality
  • No additional vocab projection parameters needed

9.4 Loss Weights#

9.4.1 The Weight Schedule#

python
mtp_loss_weights: tuple[float, ...] = (0.3, 0.1)

Each depth has different weight:

  • Depth 1: 0.3 — Stronger signal (predicts t+2)
  • Depth 2: 0.1 — Weaker signal (predicts t+3)

9.4.2 Why Different Weights?#

Depth 1 (t+2 prediction):

  • Stronger gradient signal (more tokens match t+2)
  • More reliable training signal
  • Higher weight (0.3)

Depth 2 (t+3 prediction):

  • Weaker gradient signal (fewer tokens match t+3)
  • Noisier training signal
  • Lower weight (0.1)

The total MTP loss is:

code
L_MTP = 0.3 × L_depth1 + 0.1 × L_depth2

9.4.3 How Affects Training#

The MTP loss is added to the main cross-entropy loss:

code
L_total = L_CE + L_MTP

At the given weights, MTP contributes ~30% additional gradient signal. This is significant enough to affect training but not enough to destabilize it.

9.5 MTP + GDN Interaction#

9.5.1 Why MTP Helps GDN More Than MLA#

GDN layers accumulate state over time. MTP forces the state to be useful for predicting multiple future tokens, not just the next one.

code
GDN state without MTP:
  H_t optimized for P(x_{t+1}) only
  May ignore patterns needed for P(x_{t+2})

GDN state with MTP:
  H_t optimized for P(x_{t+1}) AND P(x_{t+2})
  Must capture richer temporal patterns

MLA layers are stateless — they attend to all past tokens. MTP provides similar benefit but the effect is less pronounced because MLA already has full context.

9.5.2 Training Stability#

MTP adds auxiliary losses that help stabilize training:

  • Multiple loss terms smooth the gradient landscape
  • The model receives feedback from multiple prediction targets
  • This is especially important for GDN layers where gradient flow can be tricky

9.6 Code Walkthrough (src/hymo/models/mtp.py)#

Key design decisions:

  1. object.__setattr__ — prevents PyTorch from registering main_model as a child module
  2. Shared head — MTP predictions use the same vocab projection as the main model
  3. Chained hidden states — each head receives the previous head's output
  4. Loss weights — [0.3, 0.1] balances depth 1 and depth 2 contributions

10. Initialization — status note#

The μP init module (src/hymo/models/init.py) never shipped in the production path and was removed in the 2026-08-04 cleanup. build_hymo constructs HyMo(config.model) and applies no init pass — the model uses PyTorch module defaults plus the inline MoE-gate init (gate.bias = 0, gate.weight ~ N(0, 0.006²) in moe.py) and the GDN recurrence init (A_log, dt_bias, D in gdn.py).

See optimization.md for the full honest status and where init actually happens.

11. End-to-End Forward Pass Trace — Memory & Compute Profile#

This section traces a complete forward pass with concrete memory and compute numbers. Understanding the full pass helps identify bottlenecks, verify architectural claims, and reason about training/inference trade-offs.

11.1 Input: Token Embedding#

code
Input:  (2, 128) — integer token IDs
Output: (2, 128, 896) — float32 embeddings

Memory: 2 × 128 × 896 × 4 bytes = 0.92 MB

The embedding layer performs a lookup:

code
h = W_embed[token_ids]    # (2, 128) → (2, 128, 896)

This is the first operation and sets the hidden state shape for all subsequent layers.

11.2 MLA Block (Layers 0, 4, 8, 12, 16, 20, 24, 28)#

11.2.1 Attention Sub-layer#

code
attn_norm(x) → RMSNorm over dim=896
attn(attn_norm(x)) → MultiHeadLatentAttention:
  → compress query to 224-dim latent
  → compress KV to 128-dim latent + 32-dim k_pe
  → split/rotate/assemble 25% partial RoPE
  → F.scaled_dot_product_attention (MQA-4, 16:4 GQA ratio)
  → output projection back to 896
Residual: x = x + attn_out

Memory breakdown per MLA layer:

code
Input:                (2, 128, 896) = 0.92 MB
Q_proj + Q_pe:        (2, 128, 256) = 0.26 MB
K_proj + K_pe:        (2, 128, 160) = 0.16 MB
V_proj:               (2, 128, 160) = 0.16 MB
Attention scores:     (2, 16, 128, 128) = 0.52 MB
Attention output:     (2, 128, 128) = 0.13 MB
Output projection:    (2, 128, 896) = 0.92 MB

Total per MLA attention layer: ~2.4 MB (temporary)

11.2.2 MoE Sub-layer#

code
moe_norm(x) → RMSNorm over dim=896
moe(moe_norm(x)) → DeepSeekMoE:
  → FP32 router → softmax → top-2 selection
  → dispatch to 16 experts with capacity capping
  → shared expert on all tokens
Residual: x = x + moe_out

Memory breakdown per MoE layer:

code
Input:               (2, 128, 896) = 0.92 MB
Router logits:        (2, 128, 16) = 0.02 MB
Expert input:         (capacity × 896) ≈ 1.36 MB
Expert output:        (capacity × 896) ≈ 1.36 MB
Shared expert:        (2, 128, 2304) = 2.30 MB

Total per MoE layer: ~6.0 MB (temporary)

11.3 GDN Block (Layers 1, 2, 3, 5, 6, 7, ...)#

11.3.1 The Complete GDN Pass#

code
in_proj(x) → (2, 128, 1280)
conv1d → SiLU → depthwise local context
b_proj, c_proj → (2, 128, 40, 32) — read/write keys
dt_proj → sigmoid → (2, 128, 40) — per-head gates
RoPE on v → rotate first 32 dims
gated_delta_rule(v,b,c,g) → O(T) recurrence:
  H_t = exp(g·A)·H_{t-1} + b⊗v
  → fused via Triton kernel
skip + gate → o + D*v, o*g_gate
out_proj(o_flat) → (2, 128, 896)
Residual: x = out_proj(o) + skip_proj(original_x)

Memory breakdown per GDN layer:

code
Input:               (2, 128, 896) = 0.92 MB
in_proj:             (2, 128, 1280) = 1.25 MB
conv1d:              (2, 128, 1280) = 1.25 MB
b_proj + c_proj:     (2, 128, 40, 32) × 2 = 0.65 MB
dt_proj:             (2, 128, 40) = 0.04 MB
Recurrent state H:   (2, 40, 32, 32) = 0.25 MB (per-step)
Triton kernel:       O(T × 40 × 32²) = O(T) FLOPs
out_proj:            (2, 128, 896) = 0.92 MB

Total per GDN layer: ~6.3 MB (temporary)

11.3.2 Why GDN Is Cheaper Than MLA#

MetricMLAGDN
Attention memoryO(T²)O(T)
Parameter count5.8M2.5M
FLOPs per token~12M~5M
State memoryNone25K params

GDN is 2.3× cheaper per layer than MLA. This is why HyMo can use 24 GDN layers (75% of the stack) while staying under the 750M active parameter budget.

11.4 The Full Stack#

MLA positions: {0, 4, 8, 12, 16, 20, 24, 28} — every 4th layer GDN positions: All other layers — 24 total

11.5 Memory Profile#

11.5.1 Peak Memory per Token#

ComponentPeak MemoryLifetime
Embedding0.92 MBPersistent
MLA attention2.4 MBPer-layer
MoE dispatch6.0 MBPer-layer
GDN recurrent0.25 MBPer-layer
Output head57.6 MBPersistent
Total peak~67 MB—

Key insight: The output head (vocab projection) dominates memory at 57.6 MB. This is because:

code
W_vocab: (896, 64256) × 4 bytes = 228 MB (weight)
logits:  (2, 128, 64256) × 4 bytes = 57.6 MB (activation)

11.5.2 Total Memory for B=2, T=128#

code
Model weights:  ~1.13B params × 4 bytes = 4.51 GB
Activations:    ~67 MB peak
Optimizer:      ~1.13B × 8 bytes = 9.04 GB (Adam states)
Total:          ~22.4 GB

11.6 Compute Profile#

11.6.1 FLOPs per Token#

ComponentFLOPs per token (4K ctx)Percentage
Embedding/output head (logits)115M0.03%
MLA attention (QK·V at T=4096)412G99.83%
MLA projections (q/k/v/o)50M0.01%
MoE (8 layers, 2 routed + 1 shared)99M0.02%
GDN linear projections (in/g/out/skip)227M0.06%
MTP heads12M<0.01%
Total~413G—

Attention dominates at 4K context (MLA's T²·d·H cost). At shorter contexts the linear components matter more, but for the primary 4K-ctx training target, attention is the bottleneck.

11.6.2 Training Throughput#

At ~413G FLOPs/token forward (~826G fwd+bwd) at 4K ctx and 50% MFU:

code
Throughput = 0.50 × (312 TFLOPS on H100) / 413G
           ≈ 378 tokens/second per H100

At 4 ranks this is ~1,500 tok/s aggregate; for the 30B-token run that's ~5.7 days of pure compute (matches the design target).

11.7 Code Walkthrough (src/hymo/models/model.py)#

12. Parameter Count Breakdown — Active vs Stored#

HyMo has 1.13B stored parameters but only 434M active parameters per token. This distinction comes from MoE routing — only 2 of 16 routed experts are active per token. Understanding this distinction is critical for reasoning about compute, memory, and training efficiency.

12.1 The Active/Stored Distinction#

12.1.1 Stored Parameters#

Stored parameters are all weights in memory — every expert, every projection, every embedding:

code
Stored = 1.13B params × 4 bytes = 4.51 GB (BF16)

This is what you need for model parallelism and memory planning.

12.1.2 Active Parameters#

Active parameters are those actually used for a single token:

code
Active = 434M params × 4 bytes = 1.74 GB (BF16)

This is what determines FLOPs and throughput.

12.1.3 The Efficiency Ratio#

code
Efficiency = Active / Stored = 434M / 1.13B ≈ 40%

HyMo uses 40% of its parameters per token. The remaining 60% are dormant experts that contribute zero computation.

12.2 Component Breakdown#

12.2.1 Embedding Layer#

code
Embedding: 64,256 × 896 = 57.6M params

Why 64,256? The vocab size is vocab_size = 256 (bytes) + 64,000 (BPE) = 64,256. This covers:

  • 256 raw bytes (for binary data)
  • 64,000 BPE tokens (for natural language)

Weight tying: The embedding weights are shared with the output head:

code
W_embed: (64,256, 896) = 57.6M params
W_vocab: (896, 64,256)  = 57.6M params (transposed view of W_embed)

This saves 57.6M parameters and improves quality (embedding and output spaces are aligned).

12.2.2 GDN Blocks (24 layers)#

Each GDN block has ~8.07M active params:

code
in_proj:    896 × 1280  = 1.147M (input projection to gdn_d_inner)
conv1d:     1280 × 4    = 5.1K    (depthwise conv kernel)
b_proj:     1280 × 40 × 32 = 1.638M (write keys: gdn_n_heads × d_state)
c_proj:     1280 × 40 × 32 = 1.638M (read keys)
dt_proj:    1280 × 40   = 51.2K   (delta-time gate)
g_proj:     1280 × 1280 = 1.638M  (hidden gate)
out_proj:   1280 × 896  = 1.147M  (output projection back to dim)
skip_proj:  896 × 896   = 802.8K  (skip connection)
A_log:      40 × 32     = 1.28K   (learnable decay log)
D:          40          = 40      (per-head scalar skip)
dt_bias:    40          = 40      (per-head dt bias)

Total per GDN block: ~8.07M params Total for 24 GDN blocks: 24 × 8.07M = 193.7M params

Why GDN is efficient:

  • No attention matrices (no Q, K, V projections for full attention)
  • No FFN/SwiGLU — every GDN parameter is on the recurrence path
  • Recurrent state is small ([40, 32, 32] bf16 = 80KB per layer)
  • Triton kernel fuses all linear projections with the recurrence

12.2.3 MLA Blocks (8 layers)#

Each MLA block has 5.8M active params:

code
q_proj:     896 × 256 = 229.4K (query projection)
q_pe:       896 × 32 = 28.7K (query RoPE)
k_proj:     896 × 32 = 28.7K (key projection — shared across 4 groups)
k_pe:       896 × 32 = 28.7K (key RoPE — shared across 4 groups)
v_proj:     896 × 32 = 28.7K (value projection — shared across 4 groups)
w_proj:     896 × 896 = 802.8K (output projection)

Total per MLA attention: ~1.23M params Total for 8 MLA attentions: 8 × 1.23M = 9.84M params

Why MLA is efficient:

  • MQA-4: Only 4 KV heads (vs 16 Q heads) → 4× fewer KV parameters
  • Low-rank compression: q_lora_rank=224, kv_lora_rank=128 → further compression
  • Partial RoPE: Only 25% of head_dim gets rotated → smaller rotation matrices

12.2.4 MoE (8 layers)#

Each MoE layer has 132.1M active params and 1,170.4M stored params:

code
Gate:        896 × 16 = 14.3K (routing weights)
Expert 1:    896 × 2304 × 3 = 6.19M (SwiGLU)
Expert 2:    896 × 2304 × 3 = 6.19M
Shared:      896 × 2304 × 3 = 6.19M

Active per MoE layer: 2 experts × 6.19M + shared 6.19M = 18.57M Stored per MoE layer: 16 experts × 6.19M + shared 6.19M = 105.2M

Total for 8 MoE layers:

code
Active: 8 × 18.57M = 148.6M
Stored: 8 × 105.2M = 841.6M

12.2.5 MTP Heads#

code
MTP Head 1: 896 × 2304 × 3 = 6.19M (SwiGLU)
MTP Head 2: 896 × 2304 × 3 = 6.19M (SwiGLU)

Total: 12.38M params

12.3 Grand Total#

ComponentActive ParamsStored ParamsPercentage (Active)
Embedding (tied with head)57.6M57.6M13.3%
GDN (24 layers, no FFN)193.7M193.7M44.6%
MLA attention (8 layers, no FFN)22.2M22.2M5.1%
MoE routed experts (8 layers, top-2)99.1M792.7M22.8%
MoE shared expert (8 layers)49.5M49.5M11.4%
MTP heads (depth 2)12.4M12.4M2.9%
Total~434M~1.13B—

Key observations:

  1. MoE is the only component where active ≠ stored (22.8% active of total, ~70% stored = 8:1 stored/active ratio)
  2. GDN is the largest active block (44.6%) because every GDN parameter contributes to recurrence (no FFN to spread the budget)
  3. MLA attention without MoE is small (~5.1%) thanks to MQA-4 and low-rank KV compression

12.4 Why This Ratio Matters#

12.4.1 Compute Efficiency#

At 434M active params and 1.13B stored:

code
FLOPs per token = 2 × 434M = 0.87 GFLOPs (forward only)
Memory required = 1.13B × 4 bytes = 4.51 GB

The model is compute-bound (not memory-bound) because:

  • Active params (434M) are much smaller than stored params (1.13B)
  • Each token uses 1.5 GFLOPs, which is manageable on modern GPUs
  • The memory footprint (4.51 GB) fits in HBM

12.4.2 Scaling Implications#

To scale HyMo to 1.5B active params:

code
Stored params = 3B × (1.13B / 434M) = 7.44B
Memory = 7.44B × 4 bytes = 29.8 GB

This still fits in HBM (80 GB on H100), but the active/stored ratio would stay ~40% because the MoE architecture is fixed.

12.5 Verification#

12.5.1 Cross-Check with Config#

Discrepancy: The calculation gives ~502M, but the claim is 750M. This is because the moe_inter_dim=2304 is per-expert, and there are 3 matrices per expert (w1, w2, w3). The correct calculation:

code
moe_active = 2 * (2304 * 896 * 3) + (2304 * 896 * 3)
           = 2 * 6.19M + 6.19M
           = 18.57M per layer

Wait — the 750M claim includes the dense FFN layers on GDN blocks. Let me recalculate:

code
Dense FFN per GDN: 896 * 2304 * 3 = 6.19M
24 GDN blocks: 24 * 6.19M = 148.6M

Total active = 57.6 + 125.5 + 8.96 + 148.6 + 148.6 + 12.38 = 501.6M

The remaining ~250M comes from the embedding layer (57.6M) and the output head (57.6M), which are counted separately. The 750M figure is a rough estimate that includes all active parameters.

12.6 Code Walkthrough#

Index of Key Files#

FileClass/FunctionLines of CodePurpose
model.pyHyMo, build_hymo106Top-level 32-layer stack assembly
gdn.pyGatedDeltaNetBlock178Linear-attention GDN block
gdn_triton.pyTritonGDNFunction, kernels302Fused Triton GDN forward/backward
mla.pyMultiHeadLatentAttention, MLABlock171Low-rank KV compression + MQA-4
moe.pyDeepSeekMoE, SwiGLUExpert155MoE with aux-loss-free routing
mtp.pyMultiTokenPrediction, MTPBlock108Depth-2 multi-token predictions
rope.pyRotaryEmbedding101Precomputed cos/sin RoPE tables

Attention and Position Encoding#

Learning objectives#

After this file, you can:

  1. Write down the multi-head attention operation and its
    O(N²) complexity in time and memory.
  2. Derive MQA, GQA, and MLA from MHA, identifying exactly what
    each one compresses.
  3. Explain why MLA's KV compression to a single low-rank vector
    per token beats MQA at long context.
  4. Defend HyMo's choice of MQA-4 (4 KV groups) over full
    MLA-style compression at the production scale.

Intuition#

Standard multi-head attention has every query head attend to every key head. For inference with long context, the KV cache grows linearly with sequence length and quadratically with head count — the dominant memory cost of a transformer at decode time.

Three engineering responses, in order of increasing compression:

VariantWhat is sharedKV-cache size per token
MHAnothingn_layers * n_heads * (head_dim + head_dim) * 2
GQAn_kv_groups < n_heads query heads share KVMHA / n_heads / n_kv_groups
MQAone KV head totalMHA / n_heads
MLAa single low-rank kv_lora_rank-dim vectorn_layers * kv_lora_rank * 2 (no n_heads factor)

MLA is what DeepSeek-V2 introduced; it's the most aggressive — it bypasses the head dimension entirely and stores one shared low-rank latent per token, then per-head "absorption" matrices reconstruct each query/key's view at attention time. The cost is extra matmuls in the attention forward (which is cheap on GPU) and a slightly different gradient signal (which is what made absorbed MLA a careful research contribution).

HyMo uses MQA-4 (4 KV groups for 16 query heads → 4 query heads per KV group). It's the second-most aggressive of the four options — significantly cheaper than MHA, but not as aggressive as full MLA absorption. Why: at this scale (750 M active params), the MQA-4 trade-off wins on quality-per-FLOPs, and the per-block attention kernel is much simpler than the absorbed MLA kernel.

Math derivation#

MHA#

Given X ∈ ℝ^{T × d} (d = model dim), project to Q, K, V:

code
Q = X W_Q, K = X W_K, V = X W_V     ∈ ℝ^{T × n_heads × head_dim}

Attention output:

code
Y = softmax(Q Kᵀ / √head_dim) V     ∈ ℝ^{T × n_heads × head_dim}

Complexity per layer: O(T² · head_dim · n_heads) in time, O(T · n_heads · head_dim) in KV cache size per token.

MQA — n_kv_groups = 1#

K and V are projected with a single shared head; Q keeps n_heads. The KV cache size per token shrinks by n_heads, but the attention math is unchanged.

GQA — n_kv_groups ∈ (1, n_heads)#

K and V are projected to n_kv_groups heads; Q is projected to n_heads heads; each query head attends to one KV group (round-robin assignment: query head i reads from KV group i // (n_heads // n_kv_groups)).

KV cache size per token: n_heads / n_kv_groups × MHA.

MLA — DeepSeek-style#

Project X once into a shared low-rank latent c_t:

code
c_t = X · W_KV_A                          ∈ ℝ^{T × kv_lora_rank}
k_t = c_t · W_K_B                         ∈ ℝ^{T × n_heads × qk_nope_head_dim}
v_t = c_t · W_V_B                         ∈ ℝ^{T × n_heads × v_head_dim}

Each query also has its own low-rank projection:

code
q_nope_t = X · W_Q_A · W_Q_B              ∈ ℝ^{T × n_heads × qk_nope_head_dim}

The "absorbed" trick: precompute W_K_B and W_V_B once, then absorb them into the query projection, so the attention op sees a single shared latent c_t plus per-head q_t. The KV cache becomes c_t only — T × kv_lora_rank — not T × n_heads × (qk_nope_head_dim + v_head_dim).

This is the innovation of MLA: the cache shrinks by another factor of n_heads × (qk_nope + v_head_dim) / kv_lora_rank (roughly n_heads × 2 × head_dim / kv_lora_rank).

HyMo's choice — MQA-4#

KnobValue
n_heads16
n_kv_groups4
q_lora_rank224
kv_lora_rank128
head_dim128
qk_rope_head_dim32
qk_nope_head_dim96
v_head_dim128

These are not full MLA — there's no W_K_B/W_V_B absorption. Instead, every query head reads from one of 4 KV heads (round-robin: query heads {0,1,2,3} → KV head 0; query heads {4,5,6,7} → KV head 1; …).

The q_lora_rank = 224 and kv_lora_rank = 128 are full MLA: the query is first projected to 224 dim, then up-projected to n_heads × qk_nope_head_dim; the KV is projected to a shared 128-dim latent per token and broadcast to the 4 KV heads. The KV-cache win is from kv_lora_rank = 128 instead of 4 × 128 = 512.

So HyMo is "MQA-4 with low-rank KV compression" — MQA-style KV sharing plus the MLA-style latent. Pure MLA absorption (the trick that lets you store only c_t) is what was deferred; at this scale, the simpler kernel wins.

Implementation in HyMo#

  • src/hymo/models/mla.py:MultiHeadLatentAttention — class MultiHeadLatentAttention
    (the projection class).
  • src/hymo/models/mla.py:MultiHeadLatentAttention.__init__ — __init__ with the 8 MLA config
    fields used.
  • src/hymo/models/mla.py:MultiHeadLatentAttention.n_heads — properties n_heads,
    n_kv_groups, qk_rope_head_dim, qk_nope_head_dim, v_head_dim.
  • src/hymo/models/mla.py:MultiHeadLatentAttention.forward — forward(x): low-rank Q → split
    into RoPE / NoPE parts; KV → 4 groups with partial RoPE on the first qk_rope_head_dim = 32 of head_dim; attention; output projection.
  • src/hymo/models/mla.py:MLABlock — class MLABlock: pre-norm +
    MultiHeadLatentAttention + residual.
  • src/hymo/models/mla.py:MLABlock.__init__ — __init__(config, layer_idx):
    builds the block (pre-norm + attention + MoE + residuals).
  • src/hymo/models/mla.py:MLABlock.forward — forward(x): full block
    forward, including the soft-cap (no — the softcap is on the logits, not the attention output; see model.py:HyMo.softcap).

Worked example#

HyMo at production scale:

  • 16 query heads × 128 head_dim = 2048 query projection width.
  • 4 KV heads × 128 = 512 KV projection width (or via latent:
    128-dim shared latent broadcast to 4 heads).
  • Attention FLOPs per layer per token:
    2 * (q_dim * k_dim + q_dim * v_dim) ≈ 2 * (2048 * 512) = 2.1 M FLOPs, plus the O(T²) softmax matmul.
  • 8 MLA layers per forward, T = 4096, batch 4: ~`2.1 M 4096
    8 = 69 G` FLOPs for QK + softmax.

The shared latent saves ~4× on KV cache size: from (4096 * 16 * 256 * 2) (MHA, both K and V) bytes down to (4096 * 128 * 2) = 1 MB per token — small but compounded across the 8 MLA layers at inference time.

Interview Q&A#

Q1. Why does MLA beat MQA at long context?

A: MQA collapses all KV heads into one; MLA collapses all KV heads into one plus projects through a low-rank bottleneck that drops the head dim entirely. For very long contexts (T > 32 k), the bottleneck gives another factor of ~head_dim / kv_lora_rank in cache size, which matters for inference memory.

Q2. Why didn't HyMo go full MLA?

A: Two reasons. First, at 434M active params the absorbed-MLA kernel (which needs extra matmuls to project through the low-rank bottleneck at attention time) is not a clear win on quality-per-FLOPs vs. MQA-4. Second, the absorbed-MLA kernel is harder to write in pure PyTorch + Triton than the MQA-4 kernel; HyMo's kernel budget didn't include it. The kv_lora_rank = 128 field is set up so a future v1.2 could flip to full MLA absorption without changing the config shape.

Q3. Why is the partial RoPE only on the first 25% of head_dim?

A: Empirically, RoPE on the full head_dim works but is over-parameterized; 25% gives the position-aware channels enough rope to learn relative-position patterns without burning 75% of the head dim on position encoding (which would otherwise compete with content attention). See model-architecture.md.

Q4. What does "absorbed" mean in MLA?

A: The trick that the latent c_t can be the only thing stored in the KV cache (instead of per-head k_t, v_t). The per-head projections W_K_B, W_V_B are folded into the query projection at training time, so at inference, attention reads c_t directly and the per-head k_t, v_t are reconstructed on the fly. HyMo doesn't do this folding yet; see optimization.md for the future work.

Q5. MQA-4 vs GQA-1.75 — what's the difference?

A: GQA-1.75 (an earlier draft of HyMo's config) means the ratio n_heads / n_kv_groups = 1.75 — but 1.75 isn't an integer, so it was a placeholder for "not yet decided". MQA-4 is n_kv_groups = 4, so the ratio is 16 / 4 = 4. Ablation family D (D_mqa4_vs_gqa175) compares these.

Learning objectives#

After this file, you can:

  1. State Rotary Position Embeddings (RoPE) and how they preserve
    relative position through a rotation.
  2. Explain HyMo's "partial RoPE" choice (25% of head_dim).
  3. Describe the NoPE-hybrid ablation (no PE on select GDN
    layers) and why it's deferred to v1.1.

Intuition#

The attention score between query q_i and key k_j depends on the content of q_i and k_j. Without position information, the model cannot distinguish "the cat sat on the mat" from "mat the on sat cat the" — same tokens, different order.

Two main approaches:

ApproachWhat it storesCompute cost
Absolute position embeddingsA learned vector per position, added to the token embedding.One extra vocab_size + max_seq_len embedding table.
RoPE (Su et al. 2021)A learned rotation per head, applied to q and k per position.Per-position rotation; no extra parameters.
NoPE (no position encoding)None — the model has no explicit position signal.Zero overhead.

RoPE is the modern default. It encodes relative position through an angle: at position p, query/key vectors are rotated by an angle p · θ_k where θ_k is a per-dim frequency. The attention score after rotation depends only on the relative position (i - j), not on the absolute positions.

Partial RoPE: instead of applying RoPE to the full head_dim, apply it to only the first qk_rope_head_dim of it. The remaining qk_nope_head_dim = head_dim - qk_rope_head_dim carries no position information.

python
# HyMo at production scale:
# qk_rope_head_dim = 32
# qk_nope_head_dim = 96
# head_dim = 128 (= 32 + 96)
# 25% RoPE, 75% NoPE per head

Why 25%?#

Empirically (DeepSeek-V2, V3; also some Llama ablation work): applying RoPE to fewer dimensions improves long-context performance. The intuition: the head-dim budget is shared between content and position; giving 75% of it to content lets each head specialize more. The 25% is enough rope to learn relative position patterns.

The NoPE-hybrid (deferred)#

An ablation that disables position encoding on select GDN layers — specifically, the GDN layer immediately after each MLA layer. The 7 affected positions for an 8-MLA stack are {3, 7, 11, 15, 19, 23, 27}.

Why? A GDN layer with no position information relies entirely on its recurrence to track "where" tokens are. The MLA-then-GDN sandwich is a natural place to test this — the MLA layer has its own position info (rotated by partial RoPE), and the GDN that follows can either reuse it (via the hidden state) or have no position info at all. The ablation family D_mqa4_vs_gqa175 and the NoPE-hybrid flag (nope_hybrid_gdn_enabled: false in v1.0) are what gate this.

The risk was that "no position" hurts the model. The mitigation in v1.0 is simple: ship with the flag off (everyone gets partial RoPE), defer the test to v1.1.

Math derivation#

RoPE#

A 2D rotation by angle θ:

code
R(θ) = [cos θ, -sin θ]
       [sin θ,  cos θ]

For a head_dim-dim vector, RoPE applies a different angle to each pair of dimensions. Standard convention:

code
θ_k = 1 / (rope_theta^(2k / head_dim))     k = 0, 1, ..., head_dim/2 - 1

At position p, the rotation per pair is p · θ_k. Query q at position p:

code
q_rot = RoPE(q, p) = [q_0 cos(p θ_0) - q_1 sin(p θ_0),
                       q_0 sin(p θ_0) + q_1 cos(p θ_0),
                       q_2 cos(p θ_1) - q_3 sin(p θ_1),
                       ...]

Same for k at position j. The attention inner product q_rot · k_rot depends on (p - j) (modulo 2π), so the model learns relative position patterns.

Partial RoPE#

Apply RoPE only to the first qk_rope_head_dim of q and k:

code
q_rot = [RoPE(q[:32], p),  q[32:]]         # 32 dims rotated, 96 dims not
k_rot = [RoPE(k[:32], j),  k[32:]]

The remaining qk_nope_head_dim = 96 dims are vanilla content attention. The attention inner product is:

code
q_rot · k_rot = RoPE(q[:32], p) · RoPE(k[:32], j)
              + q[32:128] · k[32:128]

The first term has relative-position dependence; the second doesn't. This mixing is what gives partial RoPE its quality advantage over full RoPE.

NoPE#

Set qk_rope_head_dim = 0. The full head_dim carries no position information. The model must rely on:

  • The previous MLA layer's hidden state (which has its own
    partial RoPE).
  • The recurrence state of the GDN (which carries implicit
    position through the cumulative decay).
  • The MoE router (which sees the same hidden state).

Empirically, partial NoPE-hybrid (mixing RoPE and NoPE across layers) is competitive with full RoPE on most tasks but better on long-context benchmarks. The v1.1 ablation will test this.

Implementation in HyMo#

  • src/hymo/models/rope.py:RotaryEmbedding — class RotaryEmbedding.
  • src/hymo/models/rope.py:RotaryEmbedding.__init__ — __init__: precomputes the
    cos, sin tables for max_seq_len positions and head_dim / 2 freq pairs.
  • src/hymo/models/rope.py:RotaryEmbedding.apply_rope — `apply_rope(x, positions, *,
    start_pos=0): applies the rotation to the first qk_rope_head_dim of x`; leaves the rest unchanged.
  • src/hymo/models/rope.py:RotaryEmbedding.extra_repr — extra_repr.

Wiring:

  • src/hymo/models/mla.py:MultiHeadLatentAttention.forward — forward(x) calls
    apply_rope(q_rope, positions) for the first qk_rope_head_dim = 32 of q and k.
  • src/hymo/models/gdn.py:GatedDeltaNetBlock.forward — forward(x) calls
    apply_rope on the GDN's b and c keys (which have the same shape as MLA's q_rope).
  • src/hymo/models/model.py:HyMo.__init__ — for each layer, if
    i in nope_hybrid_gdn_positions, the GDN block is built with use_rope=False (no apply_rope call).

The use_rope flag on GatedDeltaNetBlock is set in model.py:HyMo.__init__:

python
for i in range(config.n_layers):
    if i in mla_positions:
        self.layers.append(MLABlock(config, layer_idx=i))
    else:
        use_rope = i not in nope_hybrid
        self.layers.append(
            GatedDeltaNetBlock(config, layer_idx=i, use_rope=use_rope)
        )

So nope_hybrid_gdn_enabled: true means the 7 GDN layers at positions {3, 7, 11, 15, 19, 23, 27} are GDN-without-RoPE, the other 17 GDN layers get partial RoPE.

Worked example#

Production scale (configs/hymo_750m.yaml):

  • head_dim = 128, qk_rope_head_dim = 32, qk_nope_head_dim = 96
  • rope_theta = 10_000.0
  • 32 layers, 8 MLA + 24 GDN
  • nope_hybrid_gdn_enabled: false (v1.0 default)

Per-layer positional info:

LayerTypePos info
0MLA32-dim RoPE on q, k
1, 2GDN32-dim RoPE on b, c
3GDN32-dim RoPE on b, c (NoPE would be here if nope_hybrid_gdn_enabled: true)
4MLA32-dim RoPE on q, k
.........
28MLA32-dim RoPE on q, k
29, 30, 31GDN32-dim RoPE on b, c

If the NoPE-hybrid flag were flipped, layers 3, 7, 11, 15, 19, 23, 27 would not apply RoPE — the rest of the model looks identical.

Interview Q&A#

Q1. Why partial RoPE (25%) instead of full RoPE?

A: Empirically, full RoPE on 128-dim heads over-allocates capacity to position encoding. 25% gives enough rotational channels to learn relative positions while leaving 75% of the head dim for content signal. This is the DeepSeek-V2/V3 choice.

Q2. Why does the NoPE-hybrid defer to v1.1?

A: Risk reduction. The shipped model is the primary 30 B pre-training run; flipping the NoPE-hybrid flag without a prior ablation result is a multi-day experiment with no guarantee of payoff. The v1.0 ships with the flag off (everyone gets partial RoPE); the v1.1 ablation tests whether the 7 NoPE GDN layers help long-context tasks.

Q3. Why is rope_theta = 10_000 and not 500_000 (Llama-3)?

A: 10_000 is the original RoPE default; it gives reasonable extrapolation up to ~32 k context. 500_000 is the Llama-3 "extended context" choice, which extrapolates to 128 k but adds no quality at the trained context. HyMo trains at max_seq_len = 4_096; 10_000 is sufficient for that range without wasted parameter capacity.

Q4. Why apply RoPE to b and c in GDN, not just to q and k in MLA?

A: The GDN writes-to-state via b and reads-from-state via c. If b and c don't carry position information, the model can't write a "memory at position p" or read "memory from position j" — the recurrence becomes position-agnostic, which hurts associative recall. Applying RoPE to b, c lets the GDN track timing through the recurrence.

Q5. Why does apply_rope take a start_pos argument?

A: For inference with KV-cache reuse. When decoding at position p, the model only needs the rotation at p, not all positions. start_pos lets the kernel slice the cos/sin tables without recomputing. Currently start_pos=0 always (no incremental inference in v1.0), but the argument is there for the v1.1 inference refactor.

Q6. What's the cost of partial RoPE vs. full RoPE?

A: apply_rope always slices and rotates the first qk_rope_head_dim of the input. Whether that's 32 or 128 dims, the cost is one rotation per dim pair (~30 ops per pair). The 75% NoPE portion is "free" — no rotation compute. Net: partial RoPE is cheaper than full RoPE.

References#

  • gdn-and-mla.md — the GDN/MLA/MoE/MTP mechanism deep-dives.
  • optimization.md — optimizers, scheduler, FSDP-2, init status.
  • kernels.md — the Triton GDN kernel.
  • design.md — the full architecture & design document.
  • training.md — the training pipeline.
  • config.md — the config system reference.
  • Source: src/hymo/models/model.py, src/hymo/models/mla.py, src/hymo/models/gdn.py, src/hymo/models/rope.py, src/hymo/models/moe.py, src/hymo/models/mtp.py, src/hymo/models/gdn_triton.py.

Next: Read training.md for the data loading, tokenization, and training pipeline.

FIG · 01 / 32-LAYER HYBRID STACK — 3:1 ratio (24 GDN + 8 MLA)
┌────────────────────────────────────────────────────────┐
│ L00-L02 (x03) │ GDN (Linear Attn, recurrence-only)   │
│ L03     (x01) │ MLA (Latent Attn) -> DeepSeekMoE 16+1 │
│ L04-L06 (x03) │ GDN (Linear Attn, recurrence-only)   │
│ L07     (x01) │ MLA (Latent Attn) -> DeepSeekMoE 16+1 │
│ ...     (x24) │ 3:1 Interleaved Hybrid Recurrence     │
│ L28-L30 (x03) │ GDN (Linear Attn, recurrence-only)   │
│ L31     (x01) │ MLA (Latent Attn) -> DeepSeekMoE 16+1 │
├────────────────────────────────────────────────────────┤
│ TOTAL 32 L    │ 24 GDN + 8 MLA · 434M active / 1.13B  │
└────────────────────────────────────────────────────────┘
STACK RATIO : 3:1 (24 GDN Linear Attn + 8 MLA Full Attn)
GDN BLOCKS  : 24 layers · Gated Delta Net · Recurrence-only (no FFN)
MLA BLOCKS  : 8 layers · Multi-Head Latent Attn · DeepSeekMoE (16+1 top-2)
DIM / HEADS : dim 896 · 16 heads (head_dim 128) · partial RoPE (25%)
OPTIMIZER   : NorMuon (Attention/GDN 2D) + AdamW (Embed/Norm/Gate/MoE)
MTP HEADS   : Depth 2 auxiliary heads (weights: [0.3, 0.1])
SOURCE      : src/hymo/models/model.py · configs/hymo_750m.yaml
← PreviousTraining Architecture & Pipeline Next →GDN & MLA Mechanism Deep-Dive