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.forwardwithcompress_ratio == 0): low-rank Q with RMSNorm, a single latent KV row per token (wkv+kv_norm), RoPE on the lastrope_head_dimof both, the KV row quantised to fp8 in place (block 32, ue8m0 scale), a sliding-window ring cache,sparse_attnwith the attention sink, the inverse rotation on the output, then the grouped low-rank output projection (wo_ablock-diagonal over groups,wo_b). CPU reference against the golden. - compress
- DeepSeek-V4.1 compressed attention (
compress_ratio > 0): theCompressor(softmax pooling ofratiotokens into one latent, or a plain projection at ratio 1), theIndexer(fp4 query heads against one shared fp4 key per compressed position, rectified scores combined byweights_proj),select_candidate_blocks(level one of the two-level top-k), theSharedAttentionRuntimeslots layers hand down the stack, YaRN RoPE withcompress_rope_theta, the two fp4 quantisers, and the fullAttention.forwardfor any ratio. CPU reference against the golden; the ratio-0 pieces are shared withattn.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_mixeswith 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 tohc_multcopies, 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,MoEin 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
- Golden
Tensor - One committed capture.
- Weight
Meta - 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’sck. 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
iof the tensor namedname. - 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 inengram. - 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.