mxfp4 kernels: replace the compute sink with real math (milestone 2)

[rhCat]

Milestone: mxfp4 kernels — replace the compute sink with real math

The structural milestone (see #1) is done: pointer map, trunk streaming, 3-phase cache, and the memory gate are validated on the 304B checkpoint (gate forecast 12.5 GB vs measured 11.92 GB; 0 dropped experts over 1,290 fetches). This milestone replaces the skeleton's checksum sink with real mxfp4 math, pipelined so reads and compute overlap.

Why now — the measured evidence

• Threads sweep (4 → 20): 7.53 → 7.25 s/token. Flat. The wall is the synchronous loop + single-threaded compute sink, not read parallelism.
• GB/token = 4.06, of which ~3.2 GB is expert bytes at FP8. The pool is 137 GB.
• Replay on the real trace: compulsory ceiling 4.11%; LRU hits 0.00% at 2 GB — policy lesson already banked.

Consequence: two independent levers, both in this milestone:

1. mxfp4 halves the expert bytes → ~68 GB pool, ~1.6 GB/token expert reads (the read-speed lever, invariant: 4-bit where streamed).
2. Real threaded kernels + read/compute overlap → kills the sink wall (the threads sweep proves there is headroom).

Scope

1. Quantize pass (tools/convert-ds4f.py quantize)

• Input: packed pool.bin (F8_E4M3 + per-tensor scales, copied as-is in the current convert) + tensor metadata (shapes from the safetensors index).
• Output: pool-mxfp4.bin (E2M1, 4-bit, per-32-element block scale, MX-standard; scale = 8-bit E8M0 per block) + pool-mxfp4.json (block layout + per-expert error stats) + manifest update.
• Deterministic; block-aligned; refuse on shape/dtype surprises (no guessing — same rule as the converter).
• Quality gate: report per-expert max-abs and relative error vs the FP8 originals in the manifest. Numbers, not vibes.
• Trunk (5.26 GB dense) stays at native precision — precision asymmetry invariant: full where resident, 4-bit where streamed.

2. Kernels (src/kernels.c + include/ds4f/kernels.h)

• Scalar-correct mxfp4 first: unpack 32×4-bit block → fp32 × block scale → accumulate fp32. Portable C99.
• Matvec path (M=1 decode): hidden/latent dims from config; per-expert w1/w2/w3 with their scales.
• Router matvec: the resident high-precision router matrix (in trunk, layers.N.ffn.gate*) — small (hidden × 256), needed to produce real top-k.
• SIMD after scalar is proven: AVX2 (x86) and NEON (arm64) paths, verified bit-identical against scalar on fixtures.
• Threaded compute: per-expert matvecs across the topk batch.

3. Pipeline (src/main.c / trunk.c pattern reuse)

• The trunk already has the async-reader pattern (ring + condvar). Extend it to experts: fetch layer L+1's top-k experts while computing layer L.
• Decode loop becomes: bind trunk L (pinned or ring) → route (resident router) → overlapped expert fetch/compute → advance.
• Multi-threaded checksum sink stays available via --sink for comparison runs.

4. Validation ladder (tests/)

• Fixture: hand-built mxfp4 block with known values → decode → fp32 (bit-exact expected).
• Quantize round-trip on synthetic FP8 tensors: error bounds asserted (not just reported).
• Engine e2e with quantized synthetic pool: same request trace as FP8 run (routing independent of cache AND of dtype — invariant 2).
• SIMD vs scalar equivalence on fixtures.

5. Acer validation run (paste back into this issue)

• quantize the real pool (~68 GB out), then rerun the laptop-preset sweep.
• Expected: GB/token ~2.1 (expert bytes halved), sec/token meaningfully lower (threaded + overlapped), RSS similar, 0 dropped.
• If sec/token still flat, the wall moved — report it, don't tune around it.

Order of work

1. quantize tool + manifest shapes (shapes already available in the converter's index)
2. scalar mxfp4 decode + matvec + router (fixture-gated)
3. pipeline restructure (async expert fetch)
4. SIMD (AVX2/NEON) + equivalence gate
5. real-pool quantize + acer sweep

Acceptance

• quantize deterministic, block-aligned, error stats in manifest
• scalar kernel bit-exact on fixture; SIMD == scalar on fixtures
• --pool pool-mxfp4.bin runs end-to-end with real routing
• acer: GB/token ≈ half of FP8 run, 0 dropped, report pasted here

Open decisions (flag in comments)

• Expert tensor shapes (w1/w2/w3 dims) — resolve from the safetensors index during quantize; post the resolved shapes here before locking the kernel dimensions.
• Block size 32 (MX default) vs 64 — 32 unless the error stats say otherwise.
• Router placement: keep in trunk (resident when pinned) — confirm the gate tensor name from the index.

[rhCat]

Step 1 (quantize tool) is implemented, tested (13/13), and pushed.

What it does: reads the HF repo directly (not the 137 GB pool), resolves
per-expert weight layout + F32 scale scheme (per-tensor or per-row;
anything else = refuse with a report), dequantizes F8_E4M3, requantizes
to MX E2M1 with per-32-block E8M0 scales, writes pool-mxfp4.bin +
pool-mxfp4.json (per-tensor offsets/shapes/error stats) + config.json.
numpy-accelerated; pure-python fallback. self-test: 300 round-trips,
deterministic, relative error bounded.

Next on the acer box:

cd ds4f-disk && git pull
python3 tools/convert-ds4f.py quantize /home/rhmaomao/hosting/deepseek_v4_flash_0731 \
    --out ~/ds4f-mxfp4 --dry-run

Paste the dry-run output. It resolves the real scale scheme (tensor
vs row per tensor) and the w1/w2/w3 shapes — the open decision in this
issue. If anything refuses (unexpected dtype, odd scale counts), that
report tells us exactly what to handle next.

Then the real run (needs numpy — the box's hf-venv has it; expect ~68 GB
out, progress every 2000 experts):

python3 tools/convert-ds4f.py quantize /home/rhmaomao/hosting/deepseek_v4_flash_0731 \
    --out ~/ds4f-mxfp4

./ds4f ~/ds4f-mxfp4 --trunk ~/ds4f-model/trunk.bin \
    --offsets ~/ds4f-model/trunk.offsets --pool ~/ds4f-mxfp4/pool-mxfp4.bin \
    --preset laptop --gen 5

The engine runs the mxfp4 pool today in sink mode — expected: expert
bytes/token roughly halved (the read lever) even before the kernels land.

[rhCat]

The dry-run did its job: the real checkpoint's routed experts are I8 with F32 scales (int8 weight-only quantization), not F8_E4M3 — the checkpoint mixes dtypes (attention F8_E4M3, experts I8). Fixed and pushed: the quantizer now accepts F8_E4M3 or I8 per tensor (I8 decoded as signed int8 x scale), with dtype carried through the layout. git pull and re-run the dry-run:

python3 tools/convert-ds4f.py quantize /home/rhmaomao/hosting/deepseek_v4_flash_0731 --out ~/ds4f-mxfp4 --dry-run

Expect: 11008 experts x 3 tensors, dtype I8, scheme resolved per tensor (tensor vs row — the last open shape decision), and a real size estimate (~68 GB). Paste it. If the scheme resolves cleanly, run the real quantize.

[rhCat]

Second assumption caught by the dry-run: the scales are F8_E8M0 (MX exponent-only), not F32 — the checkpoint is already in microscaling form (I8 values + E8M0 scales). Fixed and pushed: the quantizer accepts F32 or F8_E8M0 scales and resolves the granularity by element count (per-tensor / per-row / per-32-block), with dequant handled for all combinations. git pull and re-run:

python3 tools/convert-ds4f.py quantize /home/rhmaomao/hosting/deepseek_v4_flash_0731 --out ~/ds4f-mxfp4 --dry-run

Expect now: 11008 experts x 3 tensors, dtype I8, scale F8_E8M0, and a resolved scheme (likely block32) + a real size estimate (~68 GB). Paste it. If it resolves, the real quantize is the next step.

[rhCat]

The dry-run resolved the scheme: one E8M0 scale per 16 elements (262,144 scales for 4,194,304-element tensors), and confirmed the shapes — w1 [2048,2048], w2 [4096,1024], w3 [2048,2048]. Cross-check: 12,582,912 I8 bytes + 786,432 scale bytes = 13,369,344, exactly the converted expert_nbytes. The layout is fully mapped. block16 support is pushed (13/13).

git pull, then re-run the dry-run:

/home/rhmaomao/hosting/deepseek_v4_flash_0731/hf-venv/bin/python tools/convert-ds4f.py quantize \
    /home/rhmaomao/hosting/deepseek_v4_flash_0731 --out ~/ds4f-mxfp4 --dry-run

Expect: 11008 experts x 3 tensors, dtype I8, scale F8_E8M0, scheme block16, and the real size estimate (~68 GB). If clean, run the real quantize (same command without --dry-run). The mxfp4 output keeps 32-element blocks — the input's block16 granularity is preserved exactly through dequant/requant (dequantize per-16, requantize per-32; no approximation beyond E2M1 rounding).

[rhCat]

Dry-run clean — format fully mapped:

• 11008 experts x 3 tensors, all I8 + F8_E8M0 block16 scales
• w1/w2/w3 shapes [2048,2048] / [4096,1024] / [2048,2048]
• per-expert mxfp4 slot: 3 x (2 MiB values + 128 KiB block scales) = 6.37 MiB
• estimated pool: 68.53 GB (137.06 -> 68.53, the half-bytes lever)

Run the real quantize (needs numpy; hf-venv python — progress every 2000 experts, ~68 GB out):

cd ds4f-disk && git pull
/home/rhmaomao/hosting/deepseek_v4_flash_0731/hf-venv/bin/python tools/convert-ds4f.py quantize \
    /home/rhmaomao/hosting/deepseek_v4_flash_0731 --out ~/ds4f-mxfp4

Then the structural run on the halved pool (config comes from the quantize out dir; trunk/offsets from the earlier convert):

./ds4f ~/ds4f-mxfp4 --trunk ~/ds4f-model/trunk.bin \
    --offsets ~/ds4f-model/trunk.offsets --pool ~/ds4f-mxfp4/pool-mxfp4.bin \
    --preset laptop --gen 5

Expect: GB read per token roughly halved on the expert side (~3.2 -> ~1.6 GB/token), error summary in pool-mxfp4.json quantifying the E2M1 rounding, 0 dropped. Paste the quantize summary + run report.

[rhCat]

Real quantize complete:

• pool-mxfp4.bin: 68.53 GB (11008 experts x 6.38 MB) — exactly the dry-run estimate (137.06 -> 68.53)
• error summary: tensor max abs 64, global rms 0.789745
• config.json: expert_nbytes = 6,684,672

Reading the error numbers honestly:

• global rms 0.79 on int8-range values is small — typical per-element error is well under 1% of the value range. This is the number that matters for quality.
• max abs 64 is the E2M1 worst case at the top magnitudes: 4-bit mantissa-less values round to the nearest of {0, 0.5, 1, 1.5, 2, 3, 4, 6} x 2^k, so the worst element in a large-scale block lands up to ~half a scale step off (up to ~33% relative at magnitude boundaries). Inherent to the format, matches the self-test bound, and now quantified on real weights rather than assumed.
• If RMS on real downstream output ever matters more than bytes, the knob is block size (16 instead of 32) or a wider mantissa — but 0.79 rms says mxfp4 is doing its job.

Structural run on the halved pool:

./ds4f ~/ds4f-mxfp4 --trunk ~/ds4f-model/trunk.bin \
    --offsets ~/ds4f-model/trunk.offsets --pool ~/ds4f-mxfp4/pool-mxfp4.bin \
    --preset laptop --gen 5

Expect: expert-side GB/token ~1.6 (was ~3.2), 0 dropped, sec/token similar (still sink-bound until kernels land). Paste the run report.

[rhCat]

Halved-pool structural run — the read lever, measured:

MEMORY PLAN  12.2 GB vs 125.2 GB available   -> passes
cache: 1196 slots (8 GB; slots doubled as slot size halved)
5 tokens in 35.5 s (7.09 s/token)
GB read per token: 2.04   (trunk 0.38 + experts 1.57)   [was 3.76, experts 3.2]
cache: 1290 requests, 58 hits (4.5%), 0 dropped          [was 2.6%]
PEAK RSS 11.92 GB

The number that matters: expert bytes/token 3.2 -> 1.57 GB (49% cut), matching the pool halving (137.06 -> 68.53 GB) almost exactly. The 4-bit-where-streamed invariant pays on real weights.

The other observations confirm the plan rather than surprise:

• sec/token 7.91 -> 7.09 (~10%): still sink-bound, as flagged — the real time win waits for the kernels (steps 2-3).
• hit rate 2.6% -> 4.5%: the cache now holds 2x the experts per GB, so the same random-routing stream hits more often. Still routing-bound; real routing is what makes this number meaningful.
• RSS 11.92 GB unchanged: cache arena is 8 GB in both runs; the saving is in bytes read, not bytes resident.

Milestone 2, step 1 (quantize tool + real-pool conversion) is closed with evidence:

• format fully mapped (I8 + F8_E8M0 block16 -> E2M1 + E8M0 block32)
• pool 137.06 -> 68.53 GB, deterministic, error quantified (rms 0.79)
• engine runs the mxfp4 pool today, 0 dropped, expert reads halved

Remaining in the milestone: step 2 scalar mxfp4 kernels + router matvec (fixture-gated), step 3 read/compute overlap pipeline, step 4 SIMD equivalence.

[rhCat]

Step 2 done: scalar mxfp4 kernels, fixture-gated (14/14 tests).

• src/kernels.c + include/ds4f/kernels.h — portable C99 reference kernels for the exact pool format quantize writes:
  • ds4f_mxfp4_decode — 2 values/byte, even index low nibble, E8M0 scale per 16/32-block
  • ds4f_mxfp4_matvec — row-major W x activation (decode into scratch, dot)
  • ds4f_router_scores — resident fp32 router matrix dot + bias
• tests/test_kernels.c gates: E8M0 value decode, hand-built decode vector (bit-exact), encode->decode round trip at block16 AND block32, matvec vs naive dequant+dot on odd shapes (7x40), router vs naive dot. All exact except FP addition-order tolerance on the router bias.

The kernels are the scalar-correct reference the SIMD paths must verify bit-identical against (step 4).

Next: step 3 — pipeline (async expert fetch overlapping compute + real router wiring). That needs one converter addition first: a trunk tensor-layout manifest (trunk.json) so the engine can find the real router matrix (layers.N.ffn.gate*) and attention tensors inside the packed trunk bytes. Small converter change; re-convert of the 5.26 GB trunk only (fast).

[rhCat]

Step 3 done: real compute path, fixture-gated (15/15 tests).

The engine now runs REAL math when the layouts are given:

• convert emits trunk.json (per-layer tensor layout); quantize emits pool-mxfp4.json (per-expert tensor layout) — the engine loads both
• per layer: real router matvec on the resident fp32 gate matrix → top-k (deterministic tie-break) → the mxfp4 expert chain (w1→w2→w3) computed out of the cache slots → down/up projections → residual
• --dump-state FILE writes the final hidden state; refuses if not finite; two runs produce byte-identical dumps (determinism gate)
• report gains moe: N matvecs, M decoded elements
• json.h now parses arrays of objects (two-pass reserve+copy like objects) — required for the layout files
• fixture realism fixes: valid bounded fp32, sane E8M0 exponents, small I8 weights (the random-blob fixture overflowed to inf once real math ran on it)

Acer runbook (needs a trunk re-convert for trunk.json — 5.26 GB, fast):

cd ds4f-disk && git pull && make
python3 tools/convert-ds4f.py convert /home/rhmaomao/hosting/deepseek_v4_flash_0731 --out ~/ds4f-model

./ds4f ~/ds4f-mxfp4 --trunk ~/ds4f-model/trunk.bin \
    --offsets ~/ds4f-model/trunk.offsets --pool ~/ds4f-mxfp4/pool-mxfp4.bin \
    --layout-trunk ~/ds4f-model/trunk.json \
    --layout-pool ~/ds4f-mxfp4/pool-mxfp4.json \
    --preset laptop --gen 5 --dump-state ~/ds4f-model/state.bin

Expect: router: real matvec on N/43 layers (N depends on the real gate tensor's dtype — if the gate is fp32, N=43; if not, those layers fall back to the hash router), real matvec counts, 0 dropped. This is the reference path the GPU backend must match.

Note on the true overlap: with real routing, layer L+1's expert selection depends on layer L's computed state, so the only way to prefetch ahead is the draft model (DSpark) — that stays on the roadmap as the real pipelining win.

[rhCat]

Segfault fixed, pushed (git pull && make on the acer, then re-run the same command).

Root cause: the pool layout loader computed slot_off in an expression
that could wrap before the int64 cast, and never validated the resulting
rel_v/rel_s against the slot size. At fixture scale (KB offsets) nothing
wrapped; at the real pool (~68 GB) a wrapped rel_v sent the matvec
reading out of bounds → segfault right after the memory plan.

Fix (src/moe.c):

• slot_off computed in int64 throughout
• every tensor's rel_v/rel_s validated against slot bounds at load;
  out-of-bounds = clean refusal naming the offending layer/expert,
  not a crash
• the scratch-overflow guard in the expert chain is now loud (error +
  abort) instead of a silent skip

Gate 15/15, ASan clean on the layout path. The artifacts are untouched —
no re-quantize, no re-convert.

[rhCat]

Segfault fixed — the engine now runs real math end-to-end on the 304B checkpoint.

Root cause (ASan-confirmed on the acer): json.h's array branch stored
e->child = &d->pool[base] before filling the elements. Nested arrays
(the "shape": [4, 32] lists inside every tensor object) recursed into
j_value → j_reserve → realloc → pool moved → the stored child
pointer dangled. Heap-use-after-free, write at j_value:196, freed at
j_reserve:82. The object branch already used index-based access for
exactly this reason; the array branch didn't. macOS realloc didn't move
the pool at my test sizes (204 KB), Linux did at 210 KB — which is why
it was invisible locally and fatal on the acer. Fixed by taking the
entry by index and setting child/nchild after the fill loop.

Also fixed: Makefile now lists headers as deps — make was serving
stale binaries after json.h edits ("Nothing to be done").

First real-math run (this run):

trunk layout: 43 layer entries, 1300 tensors total — ok
pool layout: 43 x 256, max_rc=4194304 — ok
router: real matvec on 40/43 layers, others hash fallback
moe: token 0 layers 0,1,2 ... (ran to completion)
REFUSE: final state not finite, not dumping

Two follow-ons:

1. 40/43 gates — 3 layers have no fp32 gate tensor; they hash-fell-back.
   Need the trunk.json dtype for those 3 layers' gate entries.
2. Non-finite state — the numerics guard (rightly) refused to dump.
   The skeleton's residual chain has no normalization, so fp32 state can
   blow up over 5 chained tokens. Need to decide: clamp/renorm per layer,
   or treat the finite-check as the real-model quality gate it is.

[rhCat]

43/43 real routing confirmed — the BF16 gate fix landed clean. Remaining issue: the skeleton chains 43 residual layers with no normalization, so state magnitudes compound and the finite-guard rightly refuses the dump.

To wire the model's real RMSNorm, need the per-layer norm tensor names:

python3 -c "
import json
d = json.load(open('/home/rhmaomao/ds4f-model/trunk.json'))
import collections
c = collections.Counter(t['n'].split('.')[-1] for l in d['layers'] for t in l['tensors'])
for k, v in sorted(c.items()):
    print(f'{v:5d}  {k}')
"

That histogram shows whether each layer has norm.weight / norm1 / norm2 (block norms) vs only the MLA q_norm/kv_norm (per-head, not what we need). With the names confirmed, the moe step gains a real RMSNorm (BF16 weight decode, rms/eps) applied per layer — deterministic, finite, and identical on both backends. That's the faithful fix vs. a synthetic clamp.