Module deepseek_v41_ref

Module deepseek_v41_ref 

Source
Expand description

DeepSeek-V4.1 Flash tiny-graph reference harness — stage S2 of the single-Spark Q2_K bring-up (#1059, under #970).

Design artifact, not a production path. Nothing here runs on a GPU and nothing here is wired into a forward pass. Its job is to give the V4.1 graph a checkable truth BEFORE it is written, so engram, the shared-attention runtime and the reworked indexer get built once against DeepSeek’s own reference instead of twice.

§Where the truth comes from

ds41_tiny_golden.json is produced by gen_ds41_tiny_golden.py (kept beside it), which runs DeepSeek’s OWN inference/model.py (deepseek-ai/DeepSeek-V4.1-Flash @ dba1be0a) on a tiny synthetic model that keeps every V4.1 structural feature: hc_mult=4 hyper-connections with the delayed mixes, engram on two layers with DeepSeek’s hash (engram.py), shared compressed attention with two kv/index sources and a candidate prefilter whose consumers share its ratio, sqrt-softplus routing with correction bias, route_scale, swiglu clamp and a shared expert. No equation is re-derived on the Python side. The five tilelang kernels the reference calls unconditionally are replaced by ds41_ref_shims.py, pure torch written from the tilelang source, line-referenced.

§Determinism, and why only outputs are committed

Weights and inputs are RNG-free and order-free: fixed_value(name, i, scale, offset) here reproduces the generator bit for bit (fnv1a64 salt, splitmix64, 24-bit uniform, f64 arithmetic, f32 result, bf16 RNE where the parameter is stored bf16). So the golden carries only OUTPUTS: each capture as {shape, n, stride, ck, data} where data is a prime-strided sample and ck is the fp64 index-weighted checksum over the whole tensor. The tests below prove the regeneration matches before any graph code exists.

§Tier 2: the engram hash is pinned (proven 2026-09-15)

gen_ds41_engram_hash_check.py recomputes, from the real V4.1 tokenizer through engram.py, the compressed token map (129,280 -> 99,092), the 48 bucket primes, the 48 slot offsets and the 8 hash multipliers, and they are IDENTICAL to the deepseek41.engram.* metadata the Q2_K GGUF ships. The sum of each layer’s primes equals its table’s row count exactly (384,006,168 / 384,016,682), so engram_num_embeddings is derivable from the file. An Atlas engram can therefore be built from GGUF metadata alone, with nothing reverse-engineered.

Modules§

attn
DeepSeek-V4.1 attention, ratio-0 layers (Attention.forward with compress_ratio == 0): low-rank Q with RMSNorm, a single latent KV row per token (wkv + kv_norm), RoPE on the last rope_head_dim of both, the KV row quantised to fp8 in place (block 32, ue8m0 scale), a sliding-window ring cache, sparse_attn with the attention sink, the inverse rotation on the output, then the grouped low-rank output projection (wo_a block-diagonal over groups, wo_b). CPU reference against the golden.
compress
DeepSeek-V4.1 compressed attention (compress_ratio > 0): the Compressor (softmax pooling of ratio tokens into one latent, or a plain projection at ratio 1), the Indexer (fp4 query heads against one shared fp4 key per compressed position, rectified scores combined by weights_proj), select_candidate_blocks (level one of the two-level top-k), the SharedAttentionRuntime slots layers hand down the stack, YaRN RoPE with compress_rope_theta, the two fp4 quantisers, and the full Attention.forward for any ratio. CPU reference against the golden; the ratio-0 pieces are shared with attn.rs.
engram
DeepSeek-V4.1 engram, CPU reference, checked sub-op by sub-op against the golden.
hc
DeepSeek-V4.1 block scaffold: hyper-connections (hc_mixes with the Sinkhorn split, hc_pre, hc_post) and RMSNorm, CPU reference against the golden.
model
DeepSeek-V4.1 full tiny-graph forward (Block.forward, Transformer.forward): embed, expand to hc_mult copies, the one-hot initial pre-mix, engram before layers 1 and 4, six blocks with the delayed mixes, collapse, final RMSNorm, f32 logits. Every component is the CPU reference already checked on its own; this file wires them in the reference’s order and carries every per-layer cache and the shared slots across prefill and decode, so the golden’s sampled captures for layers 1 to 5 and the logits become checkable.
moe
DeepSeek-V4.1 MoE (Gate, Expert, MoE in the reference), CPU reference against the golden: sqrt-softplus routing with a correction bias that picks experts but does not scale them, top-k renormalisation, route_scale, the SwiGLU clamp, and one shared expert every token goes through.

Structs§

Golden
GoldenTensor
One committed capture.
WeightMeta
One parameter’s regeneration rule, from weights_meta.

Constants§

GOLDEN_JSON
The tracked golden. Regenerate with the script beside it; the fixture is RNG-free.

Functions§

assert_close
Elementwise max-abs comparison that names the worst index and both values.
checksum
fp64 index-weighted checksum sum(v[i] * (i + 1)), the generator’s ck. Summed sequentially here; torch reduces in a different order, so compare with a relative tolerance.
fixed_int
The generator’s integer filler: splitmix64(...) % modulus.
fixed_raw
Raw stream value for element i of the tensor named name.
fixed_value
The generator’s filler: u = (z >> 40) / 2^24, f32((2u - 1) * scale + offset) in f64.
fnv1a64
FNV-1a 64 over the UTF-8 bytes of s; the per-tensor salt.
regen_param
Regenerate a bf16- or f32-stored parameter exactly as the generator initialised it: the rule (scale, offset) and the storage dtype both come from weights_meta, so no test hard-codes a rule the generator might change. fp8/e8m0 tables have their own path in engram.
splitmix64
splitmix64 finaliser, exactly as gen_ds41_tiny_golden.py::splitmix64.
to_bf16_rne
Round an f32 to bf16 precision with round-to-nearest-even, as torch’s bf16 cast does.