MiniCPM-SALA — Complete Blueprint

Every weight, every kernel, every memory access • From token ID to "42"

9.48B params • 18.95GB bf16

8× minicpm4 (softmax GQA 32:2)

24× lightning (Simple GLA 32:32)

hidden=4096 • heads=32 • d=128

vocab=73,448 • max_pos=524,288

Architecture

Model Overview

MiniCPM-SALA is a hybrid: 8 layers use softmax attention (with sparse patterns for long context via infllmv2 CUDA kernels), 24 layers use Simple Gated Linear Attention (constant-memory). The serving framework is SGLang with a custom MiniCPMHybridReqToTokenPool that manages KV cache + compressed K1/K2 tables + GLA recurrent states.

Parameters

9,477,203,968

18.95 GB BFloat16

Hidden

4,096

32 heads × 128 dim

Context

524,288

Sparse attn > 8192

Residual Scale

0.24749

1.4/√32 (µP)

32-Layer Map

minicpm4 Layers 0, 9, 16, 17, 22, 29, 30, 31 — softmax GQA, no RoPE, o_gate, KV cache + sparse for seq>8192 lightning 24 other layers — Simple GLA, RoPE, QK norm, z_proj gate, o_norm, fixed 1MB state/layer

minicpm4 vs lightning: Every Difference

	minicpm4 (8 layers)	lightning-attn (24 layers)
Attention	Softmax GQA via FlashAttention3 + infllmv2	Simple GLA via fla kernels (chunk + recurrent)
Q / KV heads	32 / 2 (16× GQA)	32 / 32
RoPE	❌ attn_use_rope=false	✅ θ=10,000
QK Norm	❌	✅ RMSNorm(128)
Output Gate	o_gate: σ(h@W) ⊙ out, W=[4096×4096]	z_proj: σ(h@W) ⊙ out, W=[4096×4096]
Output Norm	❌	✅ RMSNorm(4096) o_norm
Decode State	KV cache: 1KB/token/layer (grows)	Fixed state: [32×128×128] = 1MB/layer
Long ctx	seq≤8192: dense (full attention). seq>8192: same attention math, but block mask from infllmv2 skips unimportant 64-token blocks	Recurrent — O(1) per token always, no mask needed
Params/layer	253,763,584 (484 MB)	285,225,216 (544 MB)
Backend class	`FlashAttentionBackend`	`SimpleGLAAttnBackend`
Kernel (decode)	`flash_attn_with_kvcache` (FA3)	`fused_recurrent_simple_gla` (fla)
Kernel (prefill)	`flash_attn_with_kvcache` + sparse: `infllmv2_attn_stage1`	`chunk_simple_gla` (fla, chunk_size=64)

Deep Understanding

Why Not Just Use Linear Attention Everywhere?

Linear attention IS just reordering the parentheses: (QK)V = Q(KV). Matrix multiply is associative. So why keep softmax at all? Because softmax does something that linear attention fundamentally cannot.

The Core Difference: Sharp vs Blurry

Softmax Attention — SHARP retrieval

// "What is the capital of France?"
// Q scans ALL tokens individually:

scores = Q @ K^T = [0.1, 0.2, 0.1, ..., 8.9, ..., 0.3]
// ^^^
// "capital of France is Paris"

// softmax turns small differences into HUGE differences:
probs = softmax(scores)
= [0.00, 0.00, 0.00, ..., 0.97, ..., 0.00]
// ^^^^
// 97% weight on the ONE relevant token!

output ≈ V[that one token] // precise retrieval → "Paris" ✓

exp() is the key. softmax(x)ᵢ = exp(xᵢ) / Σexp(xⱼ). The exponential makes large scores dominate exponentially — it's a "winner-take-all" filter. Even a small score gap becomes a huge probability gap.

Linear Attention — BLURRY summary

// Same question, but linear attention doesn't scan tokens individually.
// It queries a compressed state matrix S [128×128]:

output = Q @ S // S is a blend of ALL past tokens

// S = α^n·k₁ᵀ⊗v₁ + α^(n-1)·k₂ᵀ⊗v₂ + ... + k_nᵀ⊗v_n
// ^^^^^^^^^^^^ ^^^^^^^^^
// heavily decayed (old) full weight (recent)

// "Paris" and "Tokyo" and everything else is mixed in S
// Can't say "give me EXACTLY what token #47,832 said"
// Output is a weighted blend → might get "Paris", might get "Tokyo"

No exp(), no winner-take-all. Linear attention retrieves a weighted mix from the state. Like a blurry photo — you see the general picture but can't read fine print.

Why Linear Attention Loses Information: Two Problems

Problem 1: Exponential Decay Kills Old Information

// Unroll the recurrence to see what S actually contains:

After token 1: S = α·S₀ + k₁ᵀ⊗v₁
After token 2: S = α²·S₀ + α¹·k₁ᵀ⊗v₁ + α⁰·k₂ᵀ⊗v₂
After token 3: S = α³·S₀ + α²·k₁ᵀ⊗v₁ + α¹·k₂ᵀ⊗v₂ + α⁰·k₃ᵀ⊗v₃
...
After token n: S = αⁿ·S₀ + Σᵢ αⁿ⁻ⁱ · kᵢᵀ⊗vᵢ

// S is a WEIGHTED SUM of all past outer products.
// Recent tokens have weight ≈ 1, old tokens have weight → 0.
// Nothing is "overwritten" — it's BLENDED with exponential decay.

Concrete decay over distance:

Tokens ago	Head 0 (α=0.40)	Head 15 (α=0.80)	Head 31 (α=0.999)
1	0.40 (40%)	0.80 (80%)	0.999 (99.9%)
10	0.000105	0.107	0.990
100	≈ 0	2×10⁻¹⁰	0.905
1,000	≈ 0	≈ 0	0.368
10,000	≈ 0	≈ 0	0.0000454
100,000	≈ 0	≈ 0	≈ 0

Head 0 forgets everything beyond ~5 tokens. Head 31 remembers ~1000 tokens well. No head can remember 100K+ tokens. The information isn't "deleted" — it's multiplied by a number so small it's below float precision.

Problem 2: Fixed Capacity — Even Without Decay

Even if α = 1.0 (no decay at all), there's a deeper problem:

// With α = 1.0, S is just a sum of outer products:
S = k₁ᵀ⊗v₁ + k₂ᵀ⊗v₂ + ... + k₁₀₀₀₀₀ᵀ⊗v₁₀₀₀₀₀

// S is [128 × 128] = 16,384 values
// Each kᵀ⊗v is a RANK-1 matrix (outer product of two 128-vectors)
// Maximum rank of S = 128 (can't exceed matrix dimension)

// After ~128 linearly independent tokens: S is "full rank"
// Any new kᵀ⊗v INTERFERES with existing information
// 100,000 directions crammed into 128-dimensional space → collision

🎨 Paint Bucket Analogy

S is a bucket of paint (128×128 = fixed capacity).
Each token adds a drop of colored paint (kᵀ⊗v).

After 128 unique colors: bucket is "full"
After 100,000 drops: everything blends into brown mud
You can't extract "that blue drop from 50,000 steps ago"

Decay (α < 1) makes it worse: old drops also evaporate

📦 KV Cache Contrast (Softmax)

KV cache stores each token SEPARATELY:

K_cache[0] = k₁ ← intact, independent
K_cache[1] = k₂ ← intact, independent
...
K_cache[99999] = k₁₀₀₀₀₀

Query computes score per token individually.
Softmax picks the best → EXACT retrieval
No blending. No decay. No rank limit.
Just costs more memory.

Why MiniCPM-SALA Uses Both — The Strategic Placement

The 8 minicpm4 layers are placed at specific positions to act as "anchor points" where the model can precisely retrieve individual tokens from the full KV cache. The 24 lightning layers in between do fast, cheap feature combination that doesn't need per-token precision.

0

1

2

3-8

9

10-15

16

17

18-21

22

23-28

29

30

31

Layer 0: First layer — precise token retrieval for raw input features
Layers 1-8: Cheap feature combination, pattern matching
Layer 9: "Anchor" — re-ground representations to specific tokens
Layers 10-15: More feature extraction
Layers 16-17: Double anchor at middle — critical for multi-step reasoning
Layers 18-21: High-level reasoning
Layer 22: Re-ground again before final processing
Layers 23-28: Final feature refinement
Layers 29-31: 3× anchors at end — precise token selection for final answer

The Memory Tradeoff — Why This Matters

Architecture	KV Memory at 128K	Concurrent Reqs (32 GB GPU)
Pure softmax (32 layers)	32 × 128K × 1 KB = 4 GB	~3 requests
Hybrid (8 softmax + 24 GLA)	8 × 128K × 1 KB + 24 MB = 1.05 GB	~12 requests
Pure linear (32 layers)	32 × 1 MB = 32 MB	~400+ requests

Pure linear would be 125× cheaper but can't do precise retrieval → bad quality on long-range tasks. Pure softmax has perfect retrieval but 4× more memory → fewer concurrent users. Hybrid gets ~4× the concurrency of pure softmax with minimal quality loss — the softmax "anchors" compensate for linear layers' blurriness.

Token Prediction

"The answer to the universe is" → 42

We follow the last token "is" during decode (batch=1). All prior tokens are already encoded in KV caches (8 minicpm4 layers) and recurrent states (24 lightning layers).

0 Token ID for "is"

// scalar int64, e.g. 374 ∈ [0, 73447]

1 Embedding Lookup + µP Scale

token_id

scalar

→

embed_tokens

73,448 × 4,096

601 MB

→

row

1 × 4,096

scale_emb

h₀

1 × 4,096

h₀ = embed_tokens[374] × 12 // 300.84M params, NOT tied to lm_head

HBM Read 1 row (8KB) from 601MB table — random access

h = [1, 4096] bf16

REPEAT × 32 LAYERS (i = 0…31)

ATTENTION RESIDUAL BLOCK

residual = h

A RMSNorm (input_layernorm)

// RMSNorm = Root Mean Square Layer Normalization
// WHY: Normalizes activations so they don't explode/vanish across 32 layers.
// Unlike LayerNorm, it does NOT subtract the mean — only divides by RMS.
// Simpler and faster (no mean computation), works just as well in practice.

// Input: h = [4096] vector (one value per hidden dimension)

// Step 1: Compute Root Mean Square
rms = √( (h[0]² + h[1]² + ... + h[4095]²) / 4096 + 1e-6 )
// └── sum of squares ──────────────┘ └──── ε prevents division by zero
// = √( mean(h²) + ε )

// Step 2: Normalize (make unit-scale)
h_norm = h / rms // [4096] / scalar → [4096], each element divided by same rms

// Step 3: Scale by learned weight (element-wise)
h = h_norm ⊙ weight // [4096] ⊙ [4096] → [4096]
// weight is a learned parameter [4096] — allows each dimension to have its own scale
// 8 KB in bf16, stays in L2 cache

B1 minicpm4 Attention 8 layers

Softmax GQA 32:2 via FlashAttentionBackend. Sparse only changes which tokens are attended.

══════════════════════════════════════════
STEP 1: QKV Projection (always the same)
══════════════════════════════════════════
q = h @ W_q^T // [1,4096]×[4096,4096] → [1,4096] → reshape [32 heads, 128 dim] 32MB
k = h @ W_k^T // [1,4096]×[4096,256] → [1,256] → reshape [2 heads, 128 dim] 2MB
v = h @ W_v^T // [1,4096]×[4096,256] → [1,256] → reshape [2 heads, 128 dim] 2MB
// NO RoPE (attn_use_rope=false)

══════════════════════════════════════════
STEP 2: Append new k,v to KV cache (always the same)
══════════════════════════════════════════
k_buffer[layer_id][loc] = k // [2 heads, 128 dim]
v_buffer[layer_id][loc] = v // [2 heads, 128 dim]

══════════════════════════════════════════
STEP 3: ★ THE ATTENTION — this is where sparse differs ★
══════════════════════════════════════════

PATH A: seq_len ≤ 8,192 → Dense (standard FlashAttention3)

// Q attends to ALL tokens in KV cache — nothing special
out = flash_attn_with_kvcache(
  q, k_cache, v_cache,
  page_table, cache_seqlens,
  scale=0.08839,
  blockmask=None // ← no mask, attend to everything
)
// Standard: scores = (Q @ K_all^T) × 0.08839
// attn = softmax(scores) @ V_all

PATH B: seq_len > 8,192 → Sparse (infllmv2 block mask)

// ❶ Compress all cached keys into two summary levels
K1 = mean_pool(K_cache, stride=16, window=32) // ~seq/16 compressed keys
K2 = mean_pool(K_cache, stride=64, window=128) // ~seq/64 compressed keys

// ❷ Cheap approximate scoring against compressed keys
scores = infllmv2_attn_stage1(Q, K1, K2) // CUDA kernel

// ❸ Pool scores into 64-token blocks, pick top 64
block_scores = max_pooling_1d_varlen(scores, block_size=64)
topk_idx = block_scores.topk(64).indices // 64 best blocks

// ❹ Convert to uint64 bitmask for the CUDA kernel
block_mask = topk_to_uint64(topk_idx, max_seqlen_k, 64)

// ❺ SAME attention kernel, but with block mask!
out = infllm_cuda.varlen_fwd(
q, k_cache, v_cache, ...,
blockmask=block_mask // ← ONLY DIFFERENCE: skips masked blocks
)
// Same softmax math, but only on ~6,208 tokens
// instead of all 524,288 → ~84× fewer blocks!

Attend to: init 64 tokens + local 2,048 + topk 64×64 = ≤ 6,208 tokens instead of full seq_len

══════════════════════════════════════════
STEP 4: Output Gate + Projection (always the same)
══════════════════════════════════════════

// WHY a gate? The attention output may contain information from ALL
// attended tokens, but not all of it is useful for this layer's job.
// The gate learns to SUPPRESS irrelevant dimensions and AMPLIFY useful ones.
// Think of it as: "attention found these things, but which ones do I actually need?"
// Without the gate, every dimension of attention output passes through equally.
// With the gate, the model can selectively zero out or keep each dimension.

// σ(x) = Sigmoid function = 1 / (1 + exp(-x))
// Maps any value → range (0, 1)
// σ(-∞) → 0 (fully closed, block this dimension)
// σ(0) → 0.5 (half open)
// σ(+∞) → 1 (fully open, pass this dimension through)

// Step 4a: Compute gate values from ORIGINAL h (not attention output!)
gate_logits = h @ W_gate^T // [1,4096] × [4096,4096] → [1,4096] 32MB weight
gate = σ(gate_logits) // sigmoid([4096]) → [4096], each value in (0,1)

// Step 4b: Apply gate — element-wise multiply
out = attn_out ⊙ gate // [4096] ⊙ [4096] → [4096]
// If gate[i] ≈ 0 → out[i] ≈ 0 (suppress this dimension)
// If gate[i] ≈ 1 → out[i] ≈ attn_out[i] (keep this dimension)

// Step 4c: Project back to hidden size
out = out @ W_o^T // [1,4096] × [4096,4096] → [1,4096] 32MB weight

HBM reads: 36 + 32 + 32 = 100 MB weights + KV cache. Sparse adds K1/K2 compression + topk overhead.

B2 lightning-attn (Simple GLA) 24 layers

Gated Linear Attention via SimpleGLAAttnBackend → fused_recurrent_simple_gla

── QKV Projection ──
q,k,v = h @ W_q,W_k,W_v // each [4096×4096] → 3×32MB
// All produce [32h, 128] (no GQA grouping)

── QK RMSNorm (per-head) ──
q = RMSNorm(q.reshape(32,128), w=[128])
k = RMSNorm(k.reshape(32,128), w=[128])

── RoPE (θ=10000) ──
q,k = RoPE(q.float(), k.float(), pos).bf16()

── ★ Simple GLA Recurrence ──
// Load full state from MambaPool.temporal[layer_idx][mamba_indices]
// S: [32 heads, 128, 128] = 524,288 values = 1 MB bf16
// g_gamma: [32] floats — precomputed ALiBi slopes, one per head:
// g_gamma[0] = -0.917 (head 0, slowest decay, longest memory)
// g_gamma[1] = -0.841
// ...
// g_gamma[31] = -0.0001 (head 31, fastest decay, shortest memory)

for each head_idx = 0..31 (all 32 heads in parallel on GPU):

// Sizes for ONE head:
S_h // [128 × 128] = 16,384 values = 32 KB — this head's state matrix
q_h // [1 × 128] = 128 values — this head's query vector
k_h // [1 × 128] = 128 values — this head's key vector
v_h // [1 × 128] = 128 values — this head's value vector

// α is a SCALAR — g_gamma[head_idx] is one float from a [32] array
α = exp(g_gamma[head_idx]) // exp(one float) = one float
// e.g. head 0: α = exp(-0.917) ≈ 0.40 (forgets 60% each step)
// e.g. head 31: α = exp(-0.0001) ≈ 1.00 (remembers almost everything)

// ❶ Decay: scale down old memory (scalar × matrix)
S_h = α · S_h // scalar × [128,128] → [128,128]

// ❷ Update: add new information via outer product
S_h = S_h + k_h^T ⊗ v_h // [128,1] × [1,128] = [128,128] outer product, added to state

// ❸ Query: read from state
o_h = q_h @ S_h × 0.08839 // [1,128] @ [128,128] × scale = [1,128] output

// Concat all 32 heads: [32 × 128] = [4096] output
// Save updated S back to MambaPool.temporal

One head's recurrence (head_idx=0..31, shown for one):

scalar
exp(g_γ[idx])

S_h^old

128×128

k^T

128×1

⊗

1×128

S_h^new

128×128

→

1×128

S_h

128×128

o_h

1×128

            32 heads × [128×128] state = 524,288 values = 1 MB bf16 per layer • 24 layers = 24 MB per request • Fixed size regardless of sequence length
          

Why Different α Per Head? — Multi-Scale Memory

No head is "more important" — each head specializes in a different time scale, like 32 zoom levels on the past. The model needs all scales simultaneously.

// Each head sees the past at a different "resolution":

Head 0: α=0.40 memory ~5 tokens → "what was the LAST word?"
Head 5: α=0.65 memory ~20 tokens → "what was this sentence about?"
Head 15: α=0.80 memory ~200 tokens → "what was this paragraph about?"
Head 25: α=0.95 memory ~2000 tokens → "what was this section about?"
Head 31: α=0.999 memory ~5000 tokens → "what was the general topic?"

// For input: "The cat sat on the mat. The dog ran in the park. She went to the store..."

Head 0 sees: "...store" // only the very last word
Head 10 sees: "...went to the store" // recent phrase
Head 20 sees: "...dog ran...store" // blurry recent paragraph
Head 31 sees: "cat...dog...store" // vague gist of everything

Why not make ALL heads slow-decay?

If every head has α≈1.0, you'd blend 5000+ tokens equally. Can't tell what the last word was — drowned in the average. Fast-decay heads stay sharp on recent tokens. Slow-decay heads remember the big picture. You need both.

Why fixed slopes, not learned?

Some models (GDN) learn per-token gating. Simple GLA uses fixed ALiBi slopes — hardcoded at init, never change. ALiBi paper showed fixed geometric slopes work well empirically. Saves parameters + compute. The output projection W_o [4096×4096] learns to mix heads — weighting head 0 heavily for one task, head 31 for another.

── Output Norm + Gate + Projection ──

// First normalize the output (lightning layers have this extra norm)
o = RMSNorm(o, w=[4096]) // o_norm — same formula as input_layernorm

// Gate (same concept as minicpm4's o_gate, just named z_proj here)
// WHY: The recurrent state accumulates info from ALL past tokens.
// The gate lets the model learn which dimensions of the recurrent
// output are relevant NOW vs which are stale/irrelevant.
z = σ(h @ W_z^T) // [1,4096]×[4096,4096] → σ → [4096] each in (0,1) 32MB
o = o ⊙ z // element-wise: z[i]≈0 suppresses, z[i]≈1 passes through

out = o @ W_o^T // [1,4096]×[4096,4096] → [1,4096] 32MB

HBM reads: 96 + 32 + 32 = 160 MB weights + 2MB state R/W

h = residual + attn_output × 0.24749 // scale_depth/√layers = 1.4/√32

MLP RESIDUAL BLOCK (all 32 layers identical)

residual = h

C RMSNorm (post_attention_layernorm)

h = RMSNorm(h, w=[4096], ε=1e-6)

D SwiGLU MLP 201.33M params/layer

// WHY SwiGLU? A standard MLP is just: out = W_down @ ReLU(W_up @ h)
// SwiGLU adds a GATING mechanism: two parallel projections — one makes a
// "gate" that controls how much of the other projection passes through.
// This gives the model more expressiveness: it can learn to selectively
// activate or suppress different features in the intermediate space.
// SwiGLU consistently outperforms ReLU/GELU MLPs in practice (PaLM, LLaMA, etc.)

═══ Step 1: Fused gate + up projection (single GEMM) ═══
[gate, up] = h @ [W_gate | W_up]^T
// [1, 4096] × [4096, 32768] → [1, 32768]
// Then split in half:
// gate = first 16384 values (controls what passes through)
// up = last 16384 values (the actual features to gate)
// W_gate: [16384 × 4096] = 128 MB W_up: [16384 × 4096] = 128 MB
// Fused as one [32768 × 4096] matrix = 256 MB total

═══ Step 2: SiLU activation on gate, then element-wise multiply ═══

// SiLU (Sigmoid Linear Unit) = also called "Swish"
// Formula: SiLU(x) = x × σ(x) = x × (1 / (1 + exp(-x)))
//
// What it does for different values:
// SiLU(-5) = -5 × σ(-5) = -5 × 0.007 = -0.034 ≈ 0 (kills negative)
// SiLU(-1) = -1 × σ(-1) = -1 × 0.269 = -0.269 (slightly negative OK)
// SiLU(0) = 0 × σ(0) = 0 × 0.5 = 0
// SiLU(1) = 1 × σ(1) = 1 × 0.731 = 0.731 (passes most)
// SiLU(5) = 5 × σ(5) = 5 × 0.993 = 4.97 ≈ x (passes through)
//
// Unlike ReLU (which hard-zeroes negatives), SiLU is smooth
// and allows small negative values — better gradients during training.

x = SiLU(gate) ⊙ up
// SiLU([16384]) ⊙ [16384] → [16384]
// gate decides HOW MUCH of each "up" feature passes through:
// If gate[i] is large positive → SiLU ≈ gate[i] → up[i] passes fully
// If gate[i] is large negative → SiLU ≈ 0 → up[i] is suppressed
// This is the "GLU" (Gated Linear Unit) part of SwiGLU

═══ Step 3: Down projection (compress back to hidden size) ═══
out = x @ W_down^T
// [1, 16384] × [16384, 4096] → [1, 4096]
// W_down: [4096 × 16384] = 128 MB
// Compresses 16384-dim intermediate back to 4096-dim hidden

HBM 256 MB (gate+up) + 128 MB (down) = 384 MB per layer — 67.9% of all layer params. Deeply memory-bound at batch=1.

h = residual + mlp_output × 0.24749

h = [1, 4096] after layer 31

3 Final Norm → 4 µP Scale → 5 LM Head → 6 Sample

h = RMSNorm(h, w=[4096]) // model.norm
h = h ÷ 16 // scale_width = hidden/dim_model_base = 4096/256
logits = h @ lm_head.weight^T // [1,4096]×[4096,73448] → [1,73448] 601MB
probs = softmax(logits) // [73448] → probs["42"] is highest
next_token = sample(probs) → "42" ✓

HBM lm_head: 601 MB — single biggest per-token op. NOT tied to embed_tokens.

Deep Dive

Sparse Attention in minicpm4 Layers (seq > 8,192)

Key insight: Sparse attention does NOT change the model architecture. The QKV projections, output gate, and output projection are 100% identical in dense and sparse mode. The only difference is that for long sequences, a block mask is computed beforehand and passed to the attention CUDA kernel, telling it to skip unimportant 64-token blocks. Same softmax math, fewer blocks.

Sparse Config & infllmv2 Pipeline

dense_len

8,192

Dense below this

window

2,048

Local sliding window

topk blocks

Global selected

block_size

Tokens per block

K1 kernel

32/16

size / stride

K2 kernel

128/64

4× coarser

init_blocks

Always first 64 tokens

Block mask dim

infllmv2 N_BLOCK_DIM

DENSE (seq ≤ 8,192)

Q attends to ALL tokens:

all seq_len tokens (up to 8,192)

blockmask = None

flash_attn_with_kvcache → standard FA3

SPARSE (seq > 8,192)

Q attends to SELECTED blocks only:

init

local window

init block topk=64 selected skipped local 2048

blockmask = uint64 bitmask from infllmv2

infllm_cuda.varlen_fwd → same math, skips masked blocks

⚠ The Cliff: Dense → Sparse Transition

There is a hard cutoff at dense_len = 8,192. One token past the threshold, attention drops from 100% of tokens to ~75%. At long contexts, sparse attention covers only a tiny fraction — but selects the most important blocks.

seq_len	Mode	Tokens Attended	% of Sequence	Why
8,192	Dense	8,192	100%	Below threshold — full attention
8,193	Sparse	≤ 6,208	75.8%	← Cliff! 24% drop in one token
16,384	Sparse	≤ 6,208	37.9%	topk picks 64 of ~256 blocks
32,768	Sparse	≤ 6,208	18.9%	topk picks 64 of ~512 blocks
131,072	Sparse	≤ 6,208	4.7%	topk picks 64 of ~2,048 blocks
524,288	Sparse	≤ 6,208	1.2%	topk picks 64 of ~8,192 blocks

Tokens attended (bar) vs total sequence (full width):

8,192

8,192 (100%)

8,193

← cliff!

16K

32K

128K

524K

// The ≤ 6,208 is constant regardless of seq_len:
attended = init_blocks × block_size + topk × block_size + window
= 1 × 64 + 64 × 64 + 2048
= 64 + 4096 + 2048 = 6,208

// BUT at seq=8,193 there are only 128 total blocks (8193/64).
// topk=64 picks 64 out of 128 — still attending to ~half via topk.
// Plus window + init → still quite dense near the threshold.
// The real sparsity benefit is at 128K+ where 64 out of 2,000+ blocks.

// Also note: topk blocks may OVERLAP with local window near the end,
// so actual unique tokens attended can be < 6,208.

How the block mask is computed (Steps ❶→❹ only run for seq > 8,192):

Step 1: Compress Keys (Triton kernel: `compress_k_complete_kernel_new`)

// K1: mean-pool 32 consecutive K vectors with stride 16
K1[i] = mean(K[i*16 : i*16+32]) // → ~seq_len/16 compressed keys [num_kv_heads, 128]

// K2: mean-pool 128 K vectors with stride 64 (4× coarser)
K2[i] = mean(K[i*64 : i*64+128]) // → ~seq_len/64 compressed keys

// Stored in MiniCPMHybridReqToTokenPool.req_to_sparse_k1_token / k2_token
// Triton kernel handles history + new chunks in single fused pass

Step 2: Score & Select TopK (`infllmv2_attn_stage1` → `max_pooling_1d_varlen`)

// infllmv2_attn_stage1 calls infllm_cuda.varlen_fwd_stage1
// Q heads reshaped for GQA: [total_q, 32h, 128] → [total_q*16, 2h, 128]
// This makes each of 2 KV heads see its 16 Q heads as separate tokens

scores = infllmv2_attn_stage1(Q, K1, K2, ...) // CUDA kernel
// Returns: [num_kv_heads, total_q, max_seqlen_k1]

// Max-pool scores into block-level scores:
block_scores = max_pooling_1d_varlen(scores,
block_size=64, stride=16,
local_blocks=32, init_blocks=1) // infllmv2 CUDA kernel

// Select top-64 blocks globally:
topk_idx = block_scores.topk(64).indices.sort() // [kv_heads, total_q, 64]

Step 3: Sparse Attention via Block Mask (`infllmv2_attn_varlen_func`)

// Convert topk_idx → uint64 block mask for CUDA kernel
fwd_blockmask_uint64 = topk_to_uint64(topk_idx, max_seqlen_k, 64) // N_BLOCK_DIM=64

// Run modified FlashAttention with block mask
out = infllm_cuda.varlen_fwd(q, k, v, ..., fwd_blockmask_uint64)

// Effective attention pattern per query:
init_blocks: first 64 tokens
local_window: last 2,048 tokens
topk_blocks: 64 × 64 = 4,096 selected tokens
// Total ≤ 6,208 instead of 524,288 → ~84× reduction!

Decode Shortcut (split_stage1 mode)

// For decode (batch=1, seq_len=1), stage1 uses simpler BMM path:
q_reshape = q.reshape(bs, 1, q_head, head_dim).transpose(...) // GQA reshape
k_reshape = k.reshape(bs, k1_len//bs, kv_head, head_dim).transpose(-2,-1)
score = torch.bmm(q_reshape, k_reshape).mul_(1/√128)
score = softmax(score)
score = score.reshape(kv_head, bs, group_size, ...).sum(dim=2) // aggregate groups
// Then max_pooling_1d_varlen + topk as above

Memory Architecture

SGLang Memory Pool Stack

MiniCPM-SALA uses a custom memory pool hierarchy that manages three separate caches: KV cache for minicpm4 softmax layers, recurrent states for lightning GLA layers, and compressed key tables for sparse attention.

Memory Pool Architecture — Two Separate Systems

The model needs two completely different cache systems because minicpm4 and lightning layers store state differently. These are two independent pools, managed by two separate objects. A request gets an index in each.

POOL 1: MiniCPMHybridReqToTokenPool

Manages: which token is stored where + GLA recurrent states + compressed key indices. One instance for the whole server.

req_to_token

Shape & Size:

[max_reqs, 524,288] int32

= max_reqs × 524,288 × 4 bytes

What it stores:

// Each row = one request. Each value = slot index in KV cache.
// Example: request 3 has 5 tokens stored at these KV cache slots:
req_to_token[3] = [1042, 1043, 1099, 1100, 1101, 0, 0, ...]
// ↑ slot in k_buffer/v_buffer where token 0's KV lives

req_to_sparse_k1_token

[max_reqs, 32,766] int32

(524,288 − 32) / 16 + 1 = 32,766 slots

// Slot indices for compressed K1 keys
// K1 = mean of 32 K vectors, stride 16
// One entry per compressed chunk

req_to_sparse_k2_token

[max_reqs, 8,190] int32

(524,288 − 128) / 64 + 1 = 8,190 slots

// Slot indices for compressed K2 keys
// K2 = mean of 128 K vectors, stride 64
// 4× coarser than K1

req_index_to_mamba_index_mapping

[max_reqs] int32

// Maps request index → GLA state slot index in MambaPool
// Example: request 3 has GLA state at slot 7:
mapping[3] = 7 // → mamba_pool.temporal[:, 7, :, :, :] is this request's states

MambaPool (mamba_pool)

Stores the fixed-size recurrent states for all 24 lightning layers. Indexed by the mapping above.

mamba_cache.temporal — THE GLA STATES

[24 layers, pool_size+1, 32 heads, 128, 128] bf16

layer 0

pool×
32×128
×128

layer 1

…

layer 23

            Per request, per layer:

              32 heads × 128 × 128 = 524,288 values

              × 2 bytes (bf16) = 1 MB

            Per request, all 24 layers:

              24 × 1 MB = 24 MB

            Total pool (e.g. pool=100 requests):

              24 × 100 × 1 MB = 2.4 GB

mamba_cache.conv = [] — empty, GLA has no convolution states

mamba_map — layer ID translation

// Global layer IDs (0-31) don't match MambaPool's internal indices (0-23)
// mamba_map translates: which slot in temporal[layer_dim] stores which layer

mamba_map = {
  1→0, 2→1, 3→2, 4→3, 5→4, 6→5, 7→6, 8→7, // layers 1-8 → indices 0-7
  10→8, 11→9, ..., 15→13, // layers 10-15 → indices 8-13
  18→14, ..., 21→17, // layers 18-21 → indices 14-17
  23→18, ..., 28→23 // layers 23-28 → indices 18-23
}
// Notice: layers 0, 9, 16, 17, 22, 29, 30, 31 are MISSING — those are minicpm4!

POOL 2: HybridLinearKVPool (token_to_kv_pool)

Manages: actual K and V tensors for the 8 minicpm4 (softmax) layers only. Lightning layers don't use this at all.

MHATokenToKVPool (full_kv_pool) — actual KV storage

k_buffer — 8 separate tensors

[pool_size+1, 2 heads, 128 dim] bf16

            Per token per layer: 2 × 128 × 2B = 512 bytes

            Per token all 8 layers: 4 KB

L0

L9

L16

L17

L22

L29

L30

L31

v_buffer — 8 separate tensors

[pool_size+1, 2 heads, 128 dim] bf16

            Per token per layer: 2 × 128 × 2B = 512 bytes

            Per token all 8 layers: 4 KB

L0

L9

L16

L17

L22

L29

L30

L31

full_attention_layer_id_mapping — layer ID translation

// k_buffer and v_buffer only have 8 entries (one per minicpm4 layer)
// But layer IDs in the model are 0-31. This maps global → internal:

{ 0→0, 9→1, 16→2, 17→3, 22→4, 29→5, 30→6, 31→7 }

// So when layer 22 calls get_kv_buffer(22), it maps to internal index 4:
// k = k_buffer[4][page_table_indices] — NOT k_buffer[22]

Per-Request Memory Cost Summary

Component	Size Formula	At 4K tokens	At 128K tokens	Growth
KV cache (8 minicpm4 layers)	seq_len × 8 layers × 2 heads × 128 dim × 2B × 2(K+V)	32 MB	1 GB	O(n) — grows with seq!
GLA states (24 lightning layers)	24 × 32 × 128 × 128 × 2B	24 MB	24 MB	O(1) — constant!
Sparse K1 table	min(seq/16, 32766) × 4B	1 KB	32 KB	O(n) but tiny (indices only)
Sparse K2 table	min(seq/64, 8190) × 4B	256 B	8 KB	O(n) but tiny
Total per request		~56 MB	~1.02 GB	KV cache dominates at long ctx

      Key insight: Lightning layers (75% of the model) contribute only 24 MB constant to memory regardless of context length.
      At 128K context, the 8 minicpm4 layers consume 1 GB — this is the bottleneck for concurrent requests.
    

How a decode step reads from both pools:

// For request #3, layer 22 (minicpm4 → softmax attention):
req_pool_idx = 3
page_table = req_to_token[3] // → [1042, 1043, 1099, 1100, 1101, ...]
internal_id = full_attention_layer_id_mapping[22] // → 4
K = k_buffer[4][page_table] // → gather K vectors from KV pool slots
V = v_buffer[4][page_table] // → gather V vectors from KV pool slots

// For request #3, layer 23 (lightning → GLA recurrence):
mamba_idx = req_index_to_mamba_index_mapping[3] // → 7
internal_layer = mamba_map[23] // → 18
S = mamba_pool.temporal[18][7] // → [32 heads, 128, 128] = this request's state for layer 23
// ... run GLA recurrence, update S in-place ...
mamba_pool.temporal[18][7] = S_updated // write back

How State is Accessed During Decode

Layer Type	Backend Call	State Read From	State Written To
minicpm4	`FlashAttentionBackend.forward_decode()`	`full_kv_pool.k_buffer[mapped_id][page_table]`	New k,v appended via `set_kv_buffer()`
lightning	`SimpleGLAAttnBackend.forward()`	`mamba_pool.temporal[mamba_map[layer_id]][mamba_indices]`	Updated state written back to same location
minicpm4 (sparse)	`compressed_attention()`	`req_to_sparse_k1_token` + `req_to_sparse_k2_token`	New compressed keys via `compress_k_core_new()`

GPU & Memory

RTX PRO GPU — Chip Architecture & Memory Hierarchy

Visual diagram of the GPU showing where MiniCPM-SALA's data physically lives during decode. Data flows from slow/large (GDDR7) → fast/small (Registers).

GPU Chip Diagram — Data Placement During Decode

GDDR7 — 32 GB @ ~1,792 GB/s

SLOWEST • LARGEST

GDDR7 MEMORY CHIPS

8 GB

8 GB

8 GB

8 GB

STORED IN GDDR7:

Model Weights 18.95 GB

KV Cache (8 layers) ~32MB-1GB

GLA States (24 layers) 24 MB/req

Sparse K1/K2 tables

embed_tokens 601 MB

lm_head 601 MB

↕ Memory Bus

~1,792 GB/s bandwidth

BOTTLENECK for batch=1 decode ↕

GPU DIE

L2 Cache — ~96 MB @ ~12 TB/s

h vector [4096] 8 KB

RMSNorm weights 8 KB ea.

MLP intermediate 64 KB

ALiBi slopes 128 B

RoPE sin/cos ~2 KB

Shared across all SMs. Holds frequently-accessed small data. 6.7× faster than GDDR7.

~170 Streaming Multiprocessors (SMs), grouped into ~12 GPCs:

SM (Streaming Multiprocessor) — detailed view

Registers — 256 KB per SM @ INSTANT

FASTEST memory in the GPU. Each thread has its own.

GEMM operands

SiLU(x)=x·σ(x)

sigmoid gate

RoPE rotation

0.08839 0.24749

Shared Memory / L1 — 128 KB per SM @ ~19 TB/s

Programmer-controlled scratchpad. Shared by all threads in a block.

GEMM tiles 128×128

FlashAttn softmax accum

GLA state [128×128] 32KB

QK norm weights 256B

Tensor Cores

4th gen
BF16 matrix multiply
All GEMMs run here

CUDA Cores

128 FP32 cores
RMSNorm, softmax
SiLU, sigmoid, RoPE

Special Func

SFU units
exp(), rsqrt()
for softmax, norms

Warp Schedulers ×4

Each manages 32 threads. Issue instructions to tensor/CUDA cores every cycle.

SM 1

REG 256K
L1 128K

SM 2

REG 256K
L1 128K

SM 3

REG 256K
L1 128K

SM 4

REG 256K
L1 128K

SM 5

REG 256K
L1 128K

…

~170 SMs
total

GigaThread Engine
Distributes kernel blocks to SMs

Memory Controllers
Drive the GDDR7 bus

PCIe Gen5 ×16
CPU ↔ GPU data transfer

Video Decode
(unused for LLM)

Memory Hierarchy — Speed vs Capacity Tradeoff:

Registers

~∞ TB/s (instant) — 256 KB/SM × ~170 = ~43 MB total

L1/SHMEM

~19 TB/s — 128 KB/SM × ~170 = ~21.8 MB total

L2 Cache

~12 TB/s — 96 MB shared across all SMs

GDDR7

~1.8 TB/s — 32 GB off-chip

      At batch=1 decode, every GEMM is a [1,N]×[N,M] vector-matrix multiply. The weight matrix (N×M) must stream from GDDR7 at 1.8 TB/s.
      The GPU has 419 TFLOPS of compute but can only feed it ~1.8 TB/s of data → arithmetic intensity is ~0.01 → 99%+ of time waiting for memory.
    

What Lives Where — Detailed

🔴 GDDR7 / HBM — 32 GB @ ~1,792 GB/s

All model weights	18.95 GB	Every GEMM streams from here
KV cache (8 minicpm4 layers)	8 KB/token × seq_len	k_buffer + v_buffer in MHATokenToKVPool
GLA states (24 lightning layers)	24 MB/request	mamba_pool.temporal — read+write per layer
Compressed K1/K2 tables	variable	req_to_sparse_k1_token / k2_token
Embedding + LM head	1.2 GB	Two separate 601MB tables

🟡 L2 Cache — ~96 MB @ ~12 TB/s

Hidden state h [4096] bf16	8 KB — stays hot between ops
RMSNorm weights (all ~66 norms)	8 KB each — frequently accessed
MLP intermediate [32768]	64 KB — short-lived
ALiBi slope table g_gamma [32]	128 bytes — always cached
RoPE sin/cos tables	~2 KB — precomputed

🟢 Shared Memory / L1 — 128 KB per SM

GEMM weight tiles (tensor cores)	128×128 or 64×256 tiles from HBM
FlashAttention partial max/sum	Online softmax accumulators
GLA per-head state during update	[128×128] = 32KB fits in SHMEM

🔵 Registers — 256 KB per SM

Tensor core operands	Current GEMM fragment
SiLU, sigmoid, RoPE	All element-wise in registers
Constants: 0.08839, 0.24749, 12, 16	Scaling factors

Per-Token Decode Bandwidth (Batch=1, 4K ctx)

Operation	HBM Read	Est. Time	Source
minicpm4 layer (×8)
QKV proj	36 MB	0.020 ms	HBM
Softmax attn (4K ctx)	4 MB KV	0.003 ms	HBM
o_gate + o_proj	64 MB	0.036 ms	HBM
MLP (gate+up+down)	384 MB	0.214 ms	HBM
Total	~488 MB	~0.28 ms
lightning layer (×24)
QKV proj	96 MB	0.054 ms	HBM
GLA recurrence (state R/W)	2 MB	0.001 ms	HBM
z_proj + o_proj	64 MB	0.036 ms	HBM
MLP (gate+up+down)	384 MB	0.214 ms	HBM
Total	~546 MB	~0.31 ms
Embed + LM head + norms	~602 MB	~0.34 ms	HBM

8×0.28 + 24×0.31 + 0.34 ≈ 10.0 ms/token → ~100 tok/s // theoretical max, batch=1, RTX PRO

SOAR Competition

Optimization Landscape

Key bottlenecks and optimization opportunities based on the codebase analysis. The contest scores 40% on max-concurrent=1 (latency), 30% on max-concurrent=8, and 30% on unlimited concurrency.

Where Time is Spent (Batch=1 Decode)

MLP 67.9%

QKV 11.3%

Gate+O 9.9%

LM head 6.3%

<5%

At batch=1, weight streaming dominates. MLP alone reads 384MB × 32 layers = 12.29 GB per token.

Optimization Vectors

Technique	Impact	Key Files
W4A16 Quantization (GPTQ/Marlin)	~2× weight size reduction → ~2× decode speedup. MLP and proj weights go from 18.95GB → ~9.5GB.	gptq.py, marlin CUDA kernels
FP8 KV Cache	Halves KV cache memory → more concurrent requests → better throughput. Only helps minicpm4 layers (8 of 32).	kv_cache.py, server_args `--kv-cache-dtype fp8_e5m2`
Kernel Fusion (MLP)	Fuse gate_up+SiLU+down or overlap weight loads. MLP is 67.9% of time.	minicpm_sala.py MiniCPMMLP
GLA Kernel Optimization	fused_recurrent_simple_gla runs 24× per token. BV=8 (min) — room to tune tile sizes.	fla/fused_recurrent.py, hybrid_linear_attn_backend.py
Sparse Attention Kernel	infllmv2_attn_stage1 + max_pooling_1d_varlen are custom CUDA — profile for bottlenecks on long ctx.	infllmv2_sparse_attention.py, minicpm_sparse_utils.py
Chunked Prefill Tuning	Default `--chunked-prefill-size 32768`. Tune for concurrency profile. GLA uses chunk_size=64 internally.	server_args, chunk.py (FLA_CHUNK_SIZE)
CUDA Graph	Reduces kernel launch overhead for decode. Already supported but `--disable-cuda-graph` may be set for stability.	FlashAttentionBackend CUDA graph methods
Speculative Decoding (EAGLE3)	Multi-token prediction. Main challenge: lightning layers' recurrent state doesn't support tree verification natively.	hybrid_linear_attn_backend.py (is_target_verify handling)

Input Distribution (Speed Benchmark)

Input Length

0-4K 25%

4-16K

16-32K

32-128K 35%

128-160K

50% of inputs exceed 32K → sparse attention is critical

Output Length

0-512 35%

512-2K 25%

2-4K

4-16K

16-32K

30% outputs > 4K → decode throughput matters heavily

All Numbers

Complete Reference

All Constants

Scaling Factor	Value	Used In
scale_emb	12	After embedding lookup
scale_depth/√layers	0.24749	Every residual connection
scale_width	16	Before lm_head (4096/256)
attn scale	0.08839	1/√128, both attn types
rope_theta	10,000	Lightning layers only
rms_norm_eps	1e-6	All RMSNorm

Architecture	Value
hidden_size	4,096
intermediate_size	16,384
num_hidden_layers	32
num_attention_heads	32
num_key_value_heads (minicpm4)	2
lightning_nkv	32
head_dim	128
vocab_size	73,448
max_position_embeddings	524,288
tie_word_embeddings	false

Complete Parameter Count (sorted by %)

Component	Total Params	%	BF16
mlp.gate_proj (×32)	2,147,483,648	22.66%	4.0 GB
mlp.up_proj (×32)	2,147,483,648	22.66%	4.0 GB
mlp.down_proj (×32)	2,147,483,648	22.66%	4.0 GB
self_attn.q_proj (×32)	536,870,912	5.66%	1.0 GB
self_attn.o_proj (×32)	536,870,912	5.66%	1.0 GB
self_attn.k_proj (×32)	411,041,792	4.34%	0.77 GB
self_attn.v_proj (×32)	411,041,792	4.34%	0.77 GB
self_attn.z_proj (×24 lightning)	402,653,184	4.25%	0.75 GB
embed_tokens	300,843,008	3.17%	601 MB
lm_head	300,843,008	3.17%	601 MB
self_attn.o_gate (×8 minicpm4)	134,217,728	1.42%	256 MB
All norms	~370K	~0.00%	~740 KB
TOTAL	9,477,203,968	100%	18.95 GB

  MiniCPM-SALA Blueprint • Built from config.json + minicpm_sala.py + hybrid_linear_attn_backend.py + flashattention_backend.py + infllmv2_sparse_attention.py + minicpm_sparse_utils.py + memory_pool.py + fused_recurrent.py + chunk.py + forward_batch_info.py