The Algorithm Layer: From Attention and KV Cache to FFN and MoE

The algorithm layer answers not “which GPU runs this?” but “what result should this operation produce mathematically?” In zLLM it begins with device-independent specifications, dataflow, and CPU f32 references. The CPU oracle provides a readable and testable correctness anchor; Metal, CUDA, ROCm, Vulkan, and NPU kernels are accelerated implementations of the same semantics.

This distinction matters. A fast kernel may change layouts, fuse operations, lower precision, or submit work in parallel, but it must not silently change the causal mask, RoPE, Top-K routing, KV write position, or numerical definition. New algorithms first establish shapes, boundaries, and results in CPU/reference code, then device implementations are compared against that oracle. Passing an operator oracle establishes local semantic agreement; complete prefill/decode execution still has to validate output and performance.

1. From a token to the next token

A Transformer layer can be compared to one structured step of thought, but the analogy is not a biological claim. “Angle, shape, vision, semantics, history, and relationships” are not six predefined slots. After training, such features are distributed across directions in a high-dimensional representation, and an individual coordinate rarely has one stable human interpretation.

The complete path is:

1. Tokenization. The tokenizer maps text or multimodal placeholders to token IDs. A token may be a word, subword, byte fragment, or special symbol; it need not equal one written word.
2. Vectorization. An embedding table maps each discrete ID to a hidden_size-dimensional vector. Positional encoding then lets attention distinguish order and relative position.
3. Layer-by-layer transformation. Each layer uses attention to collect relevant information from current and historical tokens, then an FFN or MoE transforms each token representation. Normalization and residual connections run throughout the block.
4. Output. Final normalization and the LM head project hidden state to vocabulary logits. Sampling or greedy selection chooses a token, and decode feeds it into the next round until EOS or a length limit is reached.

What expansion, matrix multiplication, merging, and reduction mean

For one hidden row x ∈ R^d, a linear layer is fundamentally y = xW. Matrix multiplication forms weighted combinations of the old coordinates in a new coordinate system. It may project d dimensions into a wider d_ff space, form Q/K/V, or map concatenated head outputs back to d dimensions.

• Expansion does not create facts from nothing. It supplies a wider intermediate workspace in which gates and nonlinearities can express more feature combinations.
• Multiple heads project the same hidden state into several subspaces. Heads may learn locality, semantics, position, reference, or visual relations, but no rule says “head 3 is shape.”
• Merging usually means a weighted sum of values followed by concatenation across heads. It aggregates information rather than copying the source sentence.
• Reduction uses an output projection or FFN down projection to return to hidden_size, so the result can join the residual stream and enter the next layer. It is learned projection, not truncating the last few coordinates.

A layer's “thinking” is therefore better described as alternating cross-token retrieval in attention with per-token feature transformation in the FFN. Depth provides successive revisions of the representation; width provides the workspace available to each revision.

2. Attention: choosing what information to read now

Classic scaled dot-product attention is:

Attention(Q,K,V) = softmax(QKᵀ / √dₖ + mask)V

Q asks “what am I looking for now?”, K is the index under which each piece of history can be found, and V is the content retrieved after a match. A causal mask prevents token N from seeing the future. “Attention Is All You Need” replaced the recurrent backbone common in earlier sequence models with global self-attention, but modern models now use several compute/storage trade-offs.

These techniques are not all forms of “gated grouping.” They differ in where compression occurs:

Design	Core idea	Primarily reduces	zLLM mapping
MHA, multi-head attention	Every Q head has its own K/V head	Baseline; expressive but KV-heavy	General multi-head geometry and block/reference semantics
MQA	All Q heads share one K/V pair	KV cache and K/V bandwidth	The num_kv_heads = 1 limit of GQA
GQA, grouped-query attention	A group of Q heads shares one K/V head	KV cache, K/V projection, bandwidth	attention/gqa.rs, including full/sliding windows and hybrid layers
MLA / Gated MLA	Compress KV into a low-rank latent and reconstruct it when needed; optionally gate output	KV representation and projection cost	attention/mla.rs, composed by GLM-5.2, DeepSeek-V3, Kimi-K3, and other runtimes
Sliding / block attention	Read only a recent window or explicitly visible blocks	Long-context attention compute	gqa::CausalWindow, attention/block.rs
DSA / MSA	A learned indexer or block index selects Top-K tokens/blocks per query	Long-context QK and AV work	attention/dsa.rs, attention/msa.rs
Compressed sparse attention	Keep a recent window and pool older history before selection or full reading	Distant-history storage and compute	attention/compressed_sparse.rs
Gated DeltaNet / KDA	Recursively summarize history in short-convolution and fixed-size recurrent state	Avoid a full KV cache that grows with context	attention/gated_delta_net.rs, attention/kda.rs

Gating, grouping, low rank, and sparse selection are four different mechanisms. A gate controls how much information passes; GQA shares KV among query heads; MLA compresses representation; DSA/MSA let the current query focus on only part of history. They resemble selective human attention: a question need not examine all memories at equal strength. Engineering still has to ask whether selection itself scans all history—sparse output does not automatically imply sparse computation.

3. KV cache: a reusable representation of history

Autoregressive generation adds only one token per round. Recomputing K and V for every earlier token each time would repeat a large amount of matrix multiplication. A KV cache stores the historical K/V representation—or equivalent state—already produced by each layer. Decode only computes the new token, queries the saved history, and appends a new record.

It is useful to compare this with an index and contents in working memory, but it is not a lossless copy of the original conversation:

• GQA stores fewer KV heads; MHA/GQA capacity still generally grows linearly with token count.
• MLA stores a normalized low-rank latent plus a RoPE component and reconstructs the required representation when read.
• zLLM's MLA cache supports F16 and also defines a layout with per-group INT8 latent values while retaining RoPE in F16. Quantization saves capacity at the cost of measurable error.
• Sliding windows retain only the effective window; DSA/MSA also maintain indexing information.
• Gated DeltaNet and KDA retain fixed-size recurrent and short-convolution state. Their update semantics differ from full-attention KV and must not be forced into one storage abstraction.

kv_cache/mod.rs owns device-independent semantics such as logical-layer-to-slot mapping, GQA/MLA shapes, formats, strides, capacity, and valid length. The backend owns physical buffer allocation, placement, quantization kernels, and synchronization. KV stays with the device that computes its layer; when multi-node execution partitions consecutive complete layers, per-layer KV does not travel across the network.

4. FFN: feature processing inside each layer

Attention answers “what should I retrieve from other tokens?” The FFN answers “how should this token transform the information it now has?” A mainstream gated MLP can be summarized as:

y = W_down(act(xW_gate) ⊙ (xW_up))

Gate and up projections expand hidden state into a wider intermediate space; activation and elementwise multiplication make a nonlinear selection; down projection returns to hidden size. zLLM's moe/dense_mlp.rs preserves this device-independent dataflow. Its current activation specifications cover SiLU, clamped SiLU, SiTU, OpenAI-style SwiGLU, and GELU-Tanh, while the backend capability implements the matrix operations.

Dense FFN and MoE

In a dense FFN, the complete gate/up/down weight set participates for every token. It is casually described as “all experts working in a layer,” but in source terms a dense MLP is one complete network block, not a collection of pre-existing experts that were all selected.

MoE instead provides many expert FFNs. A router scores them and activates only the Top-K routed experts for each token; shared experts may always run. This resembles the idea that different brain regions become active for different tasks: total parameter capacity can be large while active computation per token remains smaller. The analogy explains sparse activation only; it does not mean Transformer experts map to fixed brain regions or possess named personalities.

zLLM divides FFN/MoE responsibility as follows:

Module	Responsibility
moe/dense_mlp.rs	Dense gated-MLP specification, activations, and gate/up → activation → down dataflow
moe/routing.rs	Routing scores, Top-K, grouped assignments, and active-expert accounting
moe/topk_moe.rs	Routed/shared expert composition, post-routing scaling, and output accumulation
moe/latent_moe.rs	Route from original hidden state while routed experts consume a low-dimensional latent; shared MLP still reads original hidden
moe/prefill.rs	Grouped expert execution and merging for multi-token prefill
moe/expert_predictor.rs	Predict later experts from real route history for asynchronous prefetch; owns neither weights nor I/O

MoE primarily reduces computation per token; it does not automatically reduce total model weights. High-throughput, low-latency deployments generally keep many or all expert weights resident in memory/VRAM, so memory usage remains close to that of the complete MoE model. The router saves FLOPs but does not remove unselected experts from the checkpoint.

When capacity is tight and performance requirements are lower, cold experts can reside in CPU memory or on SSD and be streamed after routing. zLLM accordingly places expert source, prefetch, and prediction feedback at backend/runtime boundaries. This introduces I/O latency, prefetch hit rate, concurrent working-set, and jitter problems. Streaming trades residency for acceptable latency only when loading genuinely overlaps current-layer compute and prediction hits often enough.

5. A map of zLLM's algorithm modules

Directory/module	What it defines	What it does not own
attention/mod.rs	Device-independent entry point for attention-family specifications	Device buffers or command submission
attention/rope.rs	Positional rotation, layouts, and reference behavior	Model tokenization
attention/gqa.rs, mla.rs	GQA/MLA geometry, windows, projection relationships, and f32 references	Platform-specific fused kernels
attention/dsa.rs, msa.rs	Token/block scoring, causal Top-K, and sparse-selection semantics	A claim that every selection path is inherently O(K)
attention/compressed_sparse.rs	Recent window, compressed history, visibility, and selection plans	Physical placement of compression buffers
attention/gated_delta_net.rs, kda.rs	Recurrent/conv state shapes, update semantics, and references	Pretending this state is ordinary KV cache
attention/hybrid.rs	Combined full-attention and linear/recurrent layer state	Device-placement policy
attention/hyper_connection.rs, attn_res.rs	Multi-residual-stream and attention-residual semantics	HTTP sessions or cross-node transport
moe/*	Dense FFN, Top-K/latent MoE, routing, shared experts, and prefetch feedback semantics	SSD/VRAM allocation and DMA
kv_cache/*	Logical layouts, formats, capacity, valid length, and persistence boundaries	Platform allocation, synchronization, and kernels
runtime/prefill.rs	Generic chunk, batch, stage, and complete-layer loops	A specific model's layer order
runtime/generation.rs	Token generation, EOS, and position-advance lifecycle	Attention/FFN mathematics
runtime/<model>/	Compose embedding, per-layer attention + FFN, normalization, LM head from model specifications	Duplicating backend algorithms

Production execution has only three complete task forms: NewPrefill creates a session and its cache, AppendPrefill appends a token span after existing history, and DecodeRound consumes one token, runs every layer, appends state, and produces the next token. The algorithm layer defines the correct order, runtime composes a complete model, backend manages resources and submission, and kernels accelerate local computation.

6. How the CPU oracle protects correctness

The CPU oracle is valuable not because “the CPU can just about run the model,” but because it turns complex device optimization into testable problems:

1. Use small but real shapes to verify post-tokenization positions, masks, RoPE, matrix dimensions, and cache append behavior.
2. Compare deterministic CPU f32 references with device kernels. Check shape, finite values, absolute/relative error, and discrete results such as sparse indices and MoE expert IDs.
3. Compare hidden state, routing, and logits layer by layer to find where error first grows.
4. Run prefill/decode regressions on fixed token sequences, including cache boundaries, appended prefixes, and long context.
5. Finally run the complete model task. Successful compilation, a passing operator oracle, fixed-token agreement, and acceptable on-device performance are four distinct states.

Tolerance must follow dtype, quantization method, and error propagation; one generous global threshold is not correctness. Top-K routing, causal selection, cache length, and commit position often require exact equality. Evidence-based atol/rtol belongs on F16/BF16/quantized matrix outputs. Even a kernel that passes its local oracle may hurt end-to-end performance through extra conversion, synchronization, or unsuitable shapes, so complete prefill/decode throughput, peak memory, and stability still need validation.

That is the role of zLLM's algorithm layer: not another device implementation, but the mathematical contract shared by every implementation. CPU references make the contract executable and regression-testable; attention/KV define how history is read, FFN/MoE define how the current representation is processed, and runtime composes them into one complete, verifiable generation.