diff --git a/.gitignore b/.gitignore index 43de3ffcb..a4bd8f87d 100644 --- a/.gitignore +++ b/.gitignore @@ -32,3 +32,4 @@ k3parts.bin # containers a fuzzer kept for reproduction fuzz-case-*/ sweep +kernel_kl diff --git a/CLAUDE.md b/CLAUDE.md index 7232a6c5b..6faf2b1be 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -98,6 +98,10 @@ WASTE_CHUNK=1 ./test_forward ... # chunked prefill inste ./test_image /tmp && ./test_state MODEL && ./test_tokenizer MODEL "text" ./test_k3parts out.bin && uv run --with torch python tools/k3parts_ref.py out.bin +# two trunk kernels over a long prompt, every position: KL, argmax, routes, +# perplexity on the real text. Run 0 against 0 first — it must be all zero. +./kernel_kl MODEL ids.txt 256 0 2,3 512 + # why two paths disagree: identical routes, a tie, or a real divergence WASTE_DUMP_ROUTE=a.route WASTE_DUMP_SCORES=a.scores ./test_forward M IDS a.bin 0 WASTE_BACKEND=cpu WASTE_DUMP_ROUTE=b.route ./test_forward M IDS b.bin 0 @@ -127,6 +131,9 @@ fast group rather than waking the whole pool — 4 MB, measured; see docs/LEARNED.md §67), `WASTE_Q8=0` (dequantize the trunk to f32 at load, any width — 8x the RAM on a 4-bit trunk, so it is out of reach on K3), `WASTE_I8MM=1`, +`WASTE_TRUNK_KERNEL` (the 4-bit trunk matvec: 0 f32, the exact reference; +1 SDOT; 2 i8mm, which a Qwen load selects when this is unset; 3 SMLAL — +LEARNED §77), `WASTE_TOK_PLAIN=1`, `WASTE_VIS_STAGE`, `WASTE_DUMP_LATENT/HIDDEN`, `WASTE_DUMP_DSA` (the sparse-attention selection: which pools won and on what scores, so two implementations can be diffed on the decision rather @@ -152,6 +159,12 @@ int8 lookup table raises the bar that far. Profiling a decode step: `WASTE_PROFILE=1 WASTE_CACHE_MB=17735 ./test_forward MODEL ids out.bin 5`. +`WASTE_PROFILE=decode` leaves the prompt steps out — they are the ones that +find the cache empty, so on a short run they are most of the expert I/O. A +Qwen container prints its own phase tree (HyperConnection, PLE, GDN, QSA, +router, shared expert) with ms/step; the gap between `wall` and `accounted` +is what no phase covers, and more than a few percent means a phase is +missing. ## Architecture @@ -237,8 +250,10 @@ RSS actually stays inside the ceiling. Residency also decides *scheduling*: `moe_layer` runs one task per routed expert when the layer's experts are already cached and one per row range when they are not, because holding K records before doing any arithmetic is -a barrier against the read-ahead. `WASTE_XPAR=0/1` forces it; the default -asks the cache. The two paths are **bit-identical** and `tests/run.sh` +a barrier against the read-ahead. `qwen_moe_layer` asks per expert instead: +the resident ones run first as tasks while the misses read, then the misses +(LEARNED §82). `WASTE_XPAR=0/1` forces it; the default +asks the cache. The paths are **bit-identical** and `tests/run.sh` asserts it — an automatic choice that changed the numbers would make results depend on how warm the cache happened to be. diff --git a/Makefile b/Makefile index 773ad7075..0b97ea88b 100644 --- a/Makefile +++ b/Makefile @@ -127,7 +127,8 @@ CFLAGS += -MMD -MP SRC := src/model.c src/kda.c src/backend.c src/ecache.c src/version.c \ src/tokenizer.c src/waste.c src/vq.c src/vision.c src/image.c \ - src/crc32.c src/memory.c + src/crc32.c src/memory.c src/qwen_gdn.c src/qwen_qsa.c \ + src/qwen_hc.c src/qwen_ple.c src/qwen_moe.c # Match what backend.c tests for. Linux/aarch64 reports "aarch64", which # does not contain "arm" — the old findstring left kda_neon.c out of the # build while backend.c still emitted the call to it, so the link failed @@ -258,9 +259,9 @@ waste$(EXE): cli/main.o libwaste.a # the two failures tests/run.sh was written to catch, so a binary that # `test` builds and `clean` forgets defeats the check meant to notice it. TESTNAMES := test_kda test_container test_forward test_tokenizer test_k3parts \ - test_state test_vision test_vision_glm test_vision_ds41 \ - test_image test_memory \ - test_cpus test_lock sweep + test_qwenparts test_state test_vision test_vision_glm \ + test_vision_ds41 test_image test_memory test_cpus test_lock sweep \ + kernel_kl test_qsa_pick test_qsa_attn TESTBINS := $(addsuffix $(EXE),$(TESTNAMES)) test: $(TESTBINS) @@ -284,10 +285,21 @@ test_forward$(EXE): tests/test_forward.o libwaste.a # library the checks do and must never drift from it. sweep$(EXE): tests/sweep.o libwaste.a $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) + +kernel_kl$(EXE): tests/kernel_kl.o libwaste.a + $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) + +test_qsa_pick$(EXE): tests/test_qsa_pick.o libwaste.a + $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) + +test_qsa_attn$(EXE): tests/test_qsa_attn.o libwaste.a + $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) test_tokenizer$(EXE): tests/test_tokenizer.o libwaste.a $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) test_k3parts$(EXE): tests/test_k3parts.o libwaste.a $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) +test_qwenparts$(EXE): tests/test_qwenparts.o libwaste.a + $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) test_image$(EXE): tests/test_image.o libwaste.a $(CC) $(CFLAGS) -o $@ $^ $(LDLIBS) diff --git a/docs/LEARNED.md b/docs/LEARNED.md index 41087fb39..941f91cd5 100644 --- a/docs/LEARNED.md +++ b/docs/LEARNED.md @@ -5413,7 +5413,1087 @@ is the other half. **A test that only compares a thing to itself is not a weak oracle, it is not an oracle**, and the three protocols this repo renders now each have one. -## 74. A feasibility gate does not need the download (2026-09-15) +## 74. Qwen3.8-Flash-Next: what was new, and what only looked new (2026-09-04) + +Four architectural pieces this engine had never run — Gated DeltaNet, +Qwen Sparse Attention, HyperConnection, and a per-layer n-gram embedding — +and the useful finding is how little of the *engine* they touched. The +container format did not change: a packed `[E, 2I, H]` gate_up beside an +`[E, H, I]` down splits into exactly the WEXP records everything else +writes, one expert per 4 KiB-aligned record, so routing still costs one +`pread` and the expert cache, the read-ahead and the expert-parallel path +were reused unmodified. What is genuinely Qwen's is five self-contained +kernel files and a loader branch. + +**The 80 GiB trunk that is 2.6 GB resident.** The n-gram tables are 16 +heads of ~20 M rows, 78 of the trunk file's 80.51 GiB. Held resident they +would exceed the whole RAM budget on any machine this targets; read a row +per head per token they cost one Q8G row each. The same exclusion the +embedding table has always had, for the same reason, and it is what makes +a 176.94 B model open with a 3.11 GB floor. The converter has the mirror +problem: a head is ~12 GiB as f32, so it is written 64 Ki rows at a time, +relying on Q8G grouping along the last dimension to make the batches +reconstruct what quantizing the head whole would have produced. + +**8 GiB of expert cache, and nothing above it.** Measured over 48 greedy +tokens: 4 GiB gives 3.20 tok/s at an 8% hit rate, 8 GiB gives 4.97 at 64%, +16 GiB gives 4.92 at 88%. The collapse below one working set is §3's rule +again — below a multiple the hit rate is zero, not low. The flat top is +the more useful half: **24 points of hit rate bought nothing**, because at +64% the remaining reads already overlap the arithmetic. A cache sized to +the machine rather than to the knee spends RAM for no tokens. + +**A tokenizer difference that no vocabulary test could see.** Qwen's +pre-tokenization pattern is `\p{N}` where Kimi's and GLM's are +`\p{N}{1,3}`: every digit is its own piece. `tools/hf_tokenizer.py` was +right to refuse the pattern rather than approximate it, and the engine now +carries `tokenizer_digit_run` the way it already carried +`tokenizer_han_split`. The trap is in the checking, not the fixing — +**Qwen's vocabulary contains no multi-digit token at all**, so on this +checkpoint the two settings produce identical ids and every parity test +passes either way. "202" has no merge to reach. The flag is therefore +tested on a synthetic vocabulary that does hold `20`, where the pre-token +boundary is directly visible. A parity test against the release would have +green-lit the wrong pattern for the next member of the family. + +**An unexplained 4e-3.** The container-native oracle and the engine read +the same trunk dequantized to the same f32, so their difference should be +summation order — around 1e-6. It is 4e-3 relative per layer. Routed +expert ids and weights match exactly at every layer and the argmax matches, +so nothing observable is wrong, and end-to-end generation is coherent on +the real checkpoint. It is recorded here undiagnosed rather than absorbed +into a tolerance: 4e-3 is close to bf16's epsilon and nothing in that path +should be rounding to bf16. The suite gates at the measured value so the +number cannot grow while nobody is looking, which is the least a check can +do about a thing it does not understand. + +## 75. Where a Qwen decode step goes, and a cache curve that climbs to 17 GiB (2026-09-11) + +`WASTE_PROFILE` stopped at Qwen's door. The forward pass had timers only +in the expert-parallel branch of its MoE, so §74 could say how fast a token +was and not where it went. It now times HyperConnection, PLE, GDN with its +recurrence as a sub-phase, QSA with block selection and attention as a +sub-phase, the router, the shared expert and the head, and both routed +expert paths rather than one. The new phases take slots after the old ones +— `tests/sweep.c` reads slots by number — and `test_forward` prints Qwen's +as a tree with ms/step and a `wall` line. `WASTE_PROFILE=decode` leaves +out the prompt steps, which are the ones that find the cache empty. + +The timers change nothing they time: final logits and every generated token +are byte-identical with profiling on and off, on both MoE paths of the +synthetic Qwen fixture. Their cost is inside run-to-run noise — 7.03 against +7.01 tok/s at a 16 GiB cache and 6.35 against 6.33 at 8 GiB, profile on +first — and the phases account for 99.9–100% of wall time on the real +container, so the tree is not missing a branch. + +Everything below: the pinned checkpoint, an M4 Pro with 8 performance and +4 efficiency cores, 48 GiB, the container on the internal SSD; an 18-token +prompt, 200 greedy decode tokens, `WASTE_THREADS=8`, `test_forward`, one +process per arm. + +**Where 141.5 ms goes**, at a 16 GiB cache and 7.03 tok/s: + +| phase | ms/step | +|---|---:| +| MoE, all of it | 67.5 | +| ├ routed expert arithmetic | 52.2 | +| ├ routed expert I/O | 6.6 | +| ├ shared expert | 6.4 | +| └ router | 2.1 | +| GDN | 35.7 | +| └ recurrence | 5.4 | +| HyperConnection | 18.0 | +| QSA | 13.3 | +| └ selection and attention | 4.0 | +| lm_head | 6.5 | +| PLE | 0.6 | + +Expert I/O is 5% of that step, and 11% (17.3 ms) of the same step at 8 GiB. +At either size Qwen on this machine is bound by arithmetic, not by the disk. + +Two rows are not what their names suggest. Only 5.4 of GDN's 35.7 ms is the +recurrence; the rest is five projections, a short conv and a gated norm. +And the trunk matvecs, spread across every row, are 154,600 calls in 200 +steps — 500 GB at 37.4 GB/s overall, but 17.1 GB/s below 1 MB, 31.0 from 1 +to 8 MB, 39.0 from 8 to 32 MB, and 98.3 for the head alone above that. The +small calls are slow per byte, and `WASTE_WIDE_MIN` is not why: with eight +threads on eight performance cores the fast group is the whole pool. What +they have in common is a dispatch each. GDN's four input projections read +the same vector, as do HyperConnection's down and inject projections, and +they are four and two dispatches where one would do. + +**The cache curve**, profile off; hit rates and bytes are the whole run, +prompt included: + +| expert cache | tok/s | hit rate | evictions | read | peak RSS | +|---:|---:|---:|---:|---:|---:| +| 8 GiB | 6.33 | 80.3% | 16,013 | 35.61 GB | — | +| 12 GiB | 6.86 | 88.0% | 5,637 | 21.72 GB | 15.70 GB | +| 16 GiB | 7.01 | 90.3% | 914 | 17.58 GB | 20.01 GB | +| 20 GiB | 7.06 | 90.4% | 0 | 17.33 GB | 21.46 GB | +| 24 GiB | 7.07 | 90.4% | 0 | 17.33 GB | 21.48 GB | + +`WASTE_DUMP_ROUTE` over the same run gives the number the curve bends at: +10,049 distinct records, 17.33 GiB. At 8 GiB, 10,603 of the 20,652 misses +were re-reads of records the cache had evicted; at 16 GiB, 144 of 10,193. +Past the distinct set nothing is left to miss but first use, which is why +20 and 24 GiB read exactly the same 17.33 GB. + +That set belongs to the session, not the architecture. New records per +token: 188 across the prompt, 74 over decode tokens 0–49, 29 over 50–99, +and 15 over both 100–149 and 150–199 — still about 28 MB a token when the +run stopped. The pre-rewrite branch +(`archive/qwen38-flash-next-pre-rewrite-20260904`) measured 95.39% at 8 GiB +over 216 tokens of a different prompt that touched 4,783 records; the same +cache size is 95% or 80% depending on what is being written. On 48 GiB a +20 GiB cache peaked at 21.5 GB resident with no swap, so this machine holds +the whole distinct set of a run this long, and a longer one will keep +asking for more. + +**§74's flat top does not survive a longer run.** §74 measured 4.97 tok/s +at 8 GiB and 4.92 at 16 over 48 tokens and concluded that above 8 GiB a +better hit rate buys nothing. Over 200 tokens the same step is worth 11% +(6.33 to 7.01). The runs differ in length, prompt and thread count, and +this entry does not isolate which of those flattened §74's curve. + +**The first PLE number was a cold page cache.** 3.54 ms/step on the first +run after the container came back from network storage, 0.58 on every run +after it. `trunk.bin` is opened without `F_NOCACHE`, unlike the expert +banks, so on-disk n-gram rows go through the page cache. + +## 76. One dispatch per Qwen layer, when the cache already holds it (2026-09-11) + +The pre-rewrite branch had two measured Qwen CPU defaults that the rewrite +dropped: the expert-parallel path forced on at a batch of four (+12.8%), +and a "one-join layer job" that held all ten routed records and ran them in +one pool dispatch (+10.9% on top). Both were measured at a 95.39% hit rate, +and neither ports as it was. + +The layer job needs no port at all. Its CPU half holds all K records, +builds the gate and up tables once and hands all K to one +`waste_parallel_for` — which is this branch's expert-parallel loop with a +batch of K. Its struct was a seam for a Metal backend; the CPU speed was +the batch. So both changes could be measured from the environment before +writing any code, on §75's protocol, 16 GiB arms alternated: + +| arm | 16 GiB tok/s | 8 GiB tok/s | +|---|---:|---:| +| cache decides, batch 4 (the default) | 7.00, 6.94, 7.10, 7.01 | 6.35 | +| forced on, batch 4 | 6.85, 6.87 | — | +| forced on, batch 10 | 7.34, 7.32 | 5.97 | +| cache decides, batch 10 | 7.60, 7.48 | 6.52 | + +**Forcing the path is §44's barrier, measured on a third model.** Forced +at batch 10 and 16 GiB, expert arithmetic fell from 52 to 34 ms a step and +expert I/O rose from 6.3 to 19.3 ms, for 5%. At 8 GiB, where a fifth of the +records miss, the same arm's I/O was 45.6 ms against 17.3 and the step lost +6%. Forced at batch 4 it lost at 16 GiB as well. The archive's gains were +real at 95% and are not a property of the model: at a lower hit rate the +hold waits on reads that the row split would have overlapped. + +Letting the cache decide and then taking all ten keeps the arithmetic +without the wait, because a layer only gets there with every record +already held. It is now the Qwen default: `qwen_moe_layer` uses a batch of +K when the cache chose the path, and a forced `WASTE_XPAR=1` or an explicit +`WASTE_XPAR_BATCH` keeps the batch it was given. On that build, against +`WASTE_XPAR_BATCH=4` in the same binary: + +| ms/step unless stated | 16 GiB | 16 GiB | 8 GiB | +|---|---:|---:|---:| +| tok/s, batch 4 → K | 7.08 → 7.65 | 7.02 → 7.50 | 6.34 → 6.52 | +| MoE | 67.3 → 57.4 | 67.6 → 58.5 | 82.6 → 79.3 | +| ├ expert arithmetic | 51.9 → 42.0 | 52.2 → 43.0 | 56.2 → 52.9 | +| └ expert I/O | 6.6 → 6.6 | 6.5 → 6.4 | 17.3 → 17.2 | + ++7.4% at 16 GiB and +2.8% at 8 GiB, the same bytes read, and no other phase +moved by more than 0.6 ms. The gain is smaller at 8 GiB because fewer +layers find all ten resident, and the path those layers take is unchanged. +Kimi's and GLM's `moe_layer` keep a batch of four: the reasoning carries +over, but nobody has measured it there. + +All 27 real-container runs in §75 and here generated the same 200 tokens. + +**Two harness mistakes, neither of which changed a conclusion.** The +synthetic Qwen fixture opens with no expert cache under `test_forward`'s +defaults, and the expert-parallel path needs four slots per routed expert. +So every "both paths" comparison made on it — §75's profiling check and +the first version of this entry's suite check — compared the row split with +itself. With `WASTE_CACHE_MB=1`, 256 slots and the whole bank, the row +split, forced batches of 4 and 64 and the default give identical logits and +generated tokens with profiling on and off, and the forced arms record no +LUT apply at all, which is how we know they took the parallel path. The +suite check now sets that cache. And the token comparisons read the second +whitespace field of `[ 0] 248068`, which below step 100 is the step index, +so they compared 100 tokens rather than 200. Re-read from the saved logs +with the index stripped, all 200 agree in every run. + +## 77. Qwen's trunk through i8mm: 29% for a difference the text cannot see (2026-09-13) + +§75 left GDN at 35.7 ms a step and HyperConnection at 18.0. Both looked +like places for NEON loops — HyperConnection alone runs a sigmoid over +10,240 values 96 times a token. They are not. `WASTE_PROFILE` now splits +every phase into how much of it was trunk matvec (`PROF_START` notes the +matvec total and `PROF_END` charges the difference) and prints matvec time +by tensor role, a tensor's name with its layer number taken out, because +the size buckets could not tell GDN's `out_proj` from QSA's `o_proj`. At +f32, 7.33 tok/s: + +| phase | ms/step | of which matvec | +|---|---:|---:| +| HyperConnection | 19.1 | 16.1 | +| GDN (recurrence 5.3) | 35.4 | 27.8 | +| QSA | 13.0 | 8.6 | +| shared expert | 6.5 | 6.2 | +| lm_head | 6.3 | 6.3 | + +HyperConnection's loops are 3 ms; the rest of both phases is 4-bit +projections. By tensor, at f32 and 7.57 tok/s: `in_proj_qkv` 11.9 ms at +39.8 GB/s, `in_proj_z` 7.5 at 37.6, GDN `out_proj` 7.5 at 37.8, +HyperConnection's down projections 9.1 at 17.1 and its up projections 6.0 +at 31.3, the shared expert's three 5.9 at 18–21, the router 1.8 at 17.3, +and GDN's 48-row `in_proj_a`/`in_proj_b` 0.8 at 5.3. The large shapes run +at about 5 GB/s a core on eight cores — compute-bound, not bandwidth-bound — +and the small ones below that. + +**The kernel.** `WASTE_TRUNK_KERNEL` already had three alternatives to +the f32 path. `sweep trunk=0,2,3,1`, one process, two repeats, 200 tokens +teacher-forced against f32, 16 GiB cache and eight threads: + +| kernel | tok/s | KL vs f32 | top-10 | argmax | +|---|---:|---:|---:|---:| +| f32 | 7.45, 7.42 | — (second repeat: 0) | | | +| i8mm | 9.54, 9.51 | 1.9e-4 | 99.1% | 199/200 | +| SMLAL | 9.02, 9.02 | 1.3e-4 | 99.1% | 199/200 | +| SDOT | 9.73, 9.91 | 6.1e-3 | 95.0% | 197/200 | + +f32 scoring exactly zero against itself on the second repeat is the check +that `waste_model_reset` clears Qwen's state between arms. i8mm takes +`in_proj_qkv` from 39.8 to 98.9 GB/s and GDN from 35.3 to 19.7 ms. Over +512 positions i8mm's KL is 2.1e-4 (511/512) and SMLAL's 1.4e-4 (510/512). +SDOT is out at thirty times the KL for 2% more speed. + +**That KL measured the easy positions.** `sweep` scores only generated +tokens, where the model is predicting its own confident output. On a +document it did not write, i8mm's KL is forty times higher, and a question +about 200 generated tokens says nothing about whether a sparse-attention +selection that only starts choosing past 2,048 tokens starts choosing +differently. `tests/kernel_kl.c` loads the container once per kernel, +steps every copy through the same tokens with the trunk kernel switched +between them, and scores each position as it goes: KL, argmax, top-10, +the logits' relative L2, how many of each layer's ten routed experts +agree, and each kernel's next-token perplexity on the real text — the one +column that says whether a copy further from f32 is any worse. Against +itself every column is exactly zero. + +The prompt was 5,918 tokens: `docs/QWEN.md`, `src/qwen_qsa.c`, +`src/qwen_gdn.c` and a question about both, tokenized a file at a time; +then 256 tokens generated from f32's greedy choices. + +| over the prompt | i8mm | SMLAL | +|---|---:|---:| +| perplexity (f32 3.712) | 3.698 (−0.40%) | 3.714 (+0.03%) | +| KL | 9.2e-3 | 7.0e-3 | +| argmax | 5,711/5,918 (96.5%) | 5,734/5,918 (96.9%) | +| top-10 | 93.2% | 94.2% | +| routed experts agreeing | 97.65% | 97.93% | +| layer-positions with any expert different | 20.3% | 18.1% | +| generated: KL, argmax | 1.5e-3, 251/256 | 1.2e-3, 251/256 | + +Nothing grows. By 512-position window i8mm's KL is 2.9e-2 over the first +window, where the text is prose the model finds hardest, and falls to +2–5e-3 across the source files; the window straddling 2,048 is 1.4e-2 +against 1.0e-2 before it, with perplexity 0.8% *lower*. Expert agreement is +97.1–98.1% in every window, and perplexity per window moves between −2.3% +and +0.8% with no trend. One expert in forty goes elsewhere and one argmax +in thirty differs, and neither shows up in how well the model predicts the +text. + +**Free-running, it is a coin that lands both ways.** Answering the +5,918-token prompt, greedy for 400 tokens, f32 walked through QSA's +selection correctly and i8mm stated the indexer's key width as 256 (it is +128) and spent its tokens re-reading the table. Five short prompts at both +kernels, greedy to 600 tokens, raw completion — a train catch-up (6:00 pm), +reversing a linked list in C, why the sky is blue, ordering four people by +age, a summary in exactly three bullets — gave the same correct answers on +the first four. On the fifth it was f32 that deliberated until the token +limit, one bullet started, and i8mm that finished in 479 tokens with three. +Greedy decoding parts at the first near-tie and after that the two are +writing different answers; one sample each way is what that looks like. + +**Speed is short-context speed.** At 200 decode tokens the default build +measures 9.91 tok/s against 7.70 with `WASTE_TRUNK_KERNEL=0` in the same +binary, 29%; the five prompts gained 26–28%. Decoding at a 6K context it +was 5.19 against 4.67, 11% — the attention the context adds is not trunk +matvec, and those two runs hit 88% and 93% of the cache on different text. + +So it is the Qwen default: a Qwen load selects i8mm when +`WASTE_TRUNK_KERNEL` is not set, and the variable pins either kernel. The +kernel is process-wide, so a process that loads Qwen and then another +architecture keeps i8mm for both. The suite needs no change: the +container-native oracle runs with `WASTE_Q8=0`, which dequantizes the +trunk at load, and Qwen's chunked prefill is its decode step in a loop. +K3 is the reason this was measured rather than assumed: SDOT measured KL +0.289 there, teacher-forced, because its recurrence carried the error +forward (the note above `WASTE_TRUNK_KERNEL` in `model.c`), and nothing +here says the next model is Qwen. + +**Chunk size, a smaller and exact change.** A matvec's rows went to the +pool at no fewer than 64 a chunk, and the pool rounds a chunk up to a +multiple of that floor: on eight threads the 320-row HyperConnection down +projection split into five chunks, the shared expert's 640 rows into five. +Chunks are now about 16 KB of weights, a power of two of 2 to 64 rows, so +i8mm's two-row tiles land on the same rows — tokens identical at both +kernels. HyperConnection down went 17.1 → 19.4 GB/s at f32 and 29.5 → 31.6 +at i8mm, the shared expert 25.1 → 28.8. The first version also split the +48-row projections and took `in_proj_a` from 7.5 to 2.7 GB/s — 63 KB is +cheaper on the calling thread than dispatched — so anything under 256 KB +stays there. Net: i8mm 9.59 → 9.82 and 9.79 tok/s, f32 inside its noise. +It changes chunking on every model and was measured on this one. +HyperConnection down is still 32 GB/s against `in_proj_qkv`'s 99, with a +10,240-wide activation quantized serially before each call as the suspect, +unmeasured. + +`sweep`'s route columns read 0% on Qwen: `qwen_moe_layer` does not write +the capture it compares. `kernel_kl` reads routes from +`waste_model_step`'s own argument instead. + +## 78. HyperConnection's down projection was waiting for the pool to wake (2026-09-13) + +§77 left HyperConnection's 320×10,240 down projection at 31.6 GB/s under +i8mm, a third of `in_proj_qkv`'s 99. Two explanations fit that: the +activation quantizer, scalar and on the calling thread, has four times the +input to chew on at that shape; or the kernel is slower there. The profile +now charges each tensor the time spent quantizing its activation and +prints the kernel's speed with that taken out, and one thread against +eight separates the kernel from the dispatch: + +| tensor | kernel GB/s, 1 thread | 8 threads | scaling | quantizing, ms/step | +|---|---:|---:|---:|---:| +| `in_proj_qkv` 13.1 MB | 16.5 | 81.1 | 4.9× | 0.14 | +| GDN `out_proj` 7.9 MB | 16.1 | 70.8 | 4.4× | 0.33 | +| HyperConnection up 2.0 MB | 16.6 | 47.3 | 2.9× | 0.03 | +| HyperConnection down 1.6 MB | 16.1 | 37.3 | 2.3× | 0.75 | +| shared expert gate 0.8 MB | 16.3 | 27.6 | 1.7× | 0.19 | + +The kernel is 16 GB/s a core at every shape. What falls off is how much +eight threads get out of a small call, and HyperConnection down's +quantization is a quarter of its time but not the rest of it. + +**A pool worker parks after 8 µs.** `WASTE_SPIN`'s default is 20,000 +iterations of an atomic load and `yield`, and the loop alone measures 6–8 +µs on this machine (50,000: 19; 100,000: 32–47; 200,000: 60–76). The number +was chosen in iterations against a 54 µs wake measured on another model +(§67). In front of the down projection sit a combine over 4×2,560, an +RMSNorm over 10,240 and the quantization — timed alone, 2, 9 and 4 µs (to +the 1 µs the clock resolves; about 15 µs inside the engine) — so that +matvec always found the pool asleep. After the up projection, a sigmoid +over 10,240 through `expf` (11 µs) did the same to the next block's first +projection. + +Spinning longer confirms it and is not the fix. Profiled, one run each: +0 → 8.95 tok/s, 20,000 → 9.68, 100,000 → 10.81 with HyperConnection down +at 89.1 GB/s, 400,000 → 10.43. Unprofiled, two runs each, with user+system +CPU seconds per token as the only energy measure available without root: + +| `WASTE_SPIN` | tok/s | CPU s/token | +|---:|---|---:| +| 20,000 | 9.78, 9.90 | 0.49 | +| 50,000 | 10.10, 10.30 | 0.52 | +| 100,000 | 10.10, 10.42 | 0.56 | +| 200,000 | 10.33, 10.82 | 0.58 | + +At 200,000, 7.5% more speed costs 17% more CPU — fewer tokens per +CPU-second, which is the trade §67 bounded the spin to avoid. The profiled +sweep's 11.7% was the profiler's own doing: a lock and two clock reads on +every matvec lengthen exactly the gaps being measured, so it inflates +anything that keeps workers awake. Judge those unprofiled, with repeats. + +**The change is to stop leaving the gaps.** HyperConnection's RMSNorm now +runs one stream per task and its sigmoid in ranges, on the fast group, and +`quant_act4_mm` quantizes one weight group per task once an input has 32 +or more (4,096 activations; GDN's 2,560-wide inputs stay serial). Each +element goes through the same function in the same order, so the bytes +are the serial loops' — tokens identical in every run. Profiled, +HyperConnection went 12.0 → 9.5 ms a step, its down projection 43.5 → 76 +GB/s and its up 47 → 88. Unprofiled, against a build of the previous +commit in the same process sequence: + +| | default spin, three runs | CPU s/token | spin 200,000 | +|---|---|---:|---:| +| before | 9.80, 9.84, 9.91 | 0.487 | 10.64 | +| after | 10.07, 10.00, 10.07 | 0.501 | 11.19 | + ++2.0%, every run of the new build above every run of the old, for 3% more +CPU a token. Smaller than the profile said, and cheaper than the spin that +bought 3.7% for 7%. + +**What is left is not HyperConnection.** The new build is still 11% faster +at spin 200,000, so other serial stretches still put the pool to sleep. The +largest is GDN's recurrence, 5.3 ms a step on the calling thread — 147 µs a +layer, sitting between `in_proj_qkv` and `out_proj` — then QSA's selection +and attention at 4 ms. + +One harness note, for whoever measures this next: the session scratchpad +was emptied twice mid-session and took reference logs with it. Every +comparison above ran in one command against a build of `HEAD` made with +`git archive`, so none of it depends on a file surviving between runs. + +## 79. GDN's recurrence, one value head per task (2026-09-14) + +§78 ended on the largest serial stretch left in a Qwen decode step: GDN's +recurrence, 5.3 ms a step on the calling thread — 147 µs in each of 36 +layers, between `in_proj_qkv` and `out_proj`. Its 48 value heads share +nothing they write. A head reads its own rows of `v`, the decay and `beta` +and of `S`, plus the QK head it is repeated from, and writes only its own +rows of `S` and the output; the one shared buffer was a `Dv`-float scratch. + +So `qwen_gdn.c` stays the kernel file, now with +`waste_qwen_gdn_step_heads(h0, h1, ...)` for a range of value heads, and +`waste_qwen_gdn_step` is that range over all of them — the reference check +in `tests/test_qwenparts.c` still calls the whole step and still passes. +`qwen_gdn_layer` hands the heads to the fast group, each task with its own +scratch on its stack. Same code per head in the same order, so the state +and output are the serial loop's bit for bit. + +Against a build of the previous commit, unprofiled, 200 decode tokens, +16 GiB cache, eight threads: + +| | three runs, tok/s | mean | CPU s/token | +|---|---|---:|---:| +| before | 9.96, 9.86, 9.93 | 9.92 | 0.505 | +| after | 10.31, 10.21, 10.15 | 10.22 | 0.509 | + ++3.1%, every run above every run before, tokens identical in all seven +runs including the profiled one, and CPU per token within 1% — this one is +parallel work, not a worker spinning. Profiled, the recurrence went from +5.3 to 1.64 ms a step and GDN from 19.8 to 15.7, which is the same 3.7 ms. + +What it did not do is the other half of the reason given for it. GDN's +`out_proj`, which followed the serial recurrence and so was expected to be +paying for a parked pool, measured about 88 GB/s against 86 before. That +projection is 7.9 MB and was already long enough to hide a wake; the gaps +§78 found were costly in front of HyperConnection's 1.6 MB matvecs, not +in front of every matvec. + +The largest serial stretch left is QSA's block selection and attention, +4 ms a step across 12 layers. + +## 80. QSA's attention was a third of a long-context step, on one core (2026-09-14) + +Every profile until now ran at a context of about 220 tokens, and at that +length QSA's selection and attention looked like the 4 ms §79 ended on. +Three of its four parts grow with the context rather than the token, so +that number could not say what a long conversation costs. The profile now +splits it into the RoPE table the block scores rotate by, block pooling +and top-k, the BF16-to-F32 gather of the selected K/V, and the attention +itself, and was run at both lengths — the same 18-token prompt with 200 +decode tokens, and `docs/QWEN.md` (2,801 tokens) with 32: + +| ms/step | ~220-token context | ~2,830-token context | +|---|---:|---:| +| RoPE table | 0.13 | 4.3 | +| block pooling and top-k | 0.13 | 5.2 | +| K/V gather | 0.30 | 5.3 | +| attention | 3.4 | 58.1 | +| QSA, all of it | 8.3 | 77.7 | +| the step | 96.7 (10.34 tok/s) | 170 (5.88 tok/s) | + +Nothing else in the step moved with the context — MoE, GDN and +HyperConnection cost the same at both lengths. At 2,830 tokens QSA was 46% +of the step and its attention alone a third: 24 query heads, each over +every selected token at dimension 256, one after another on the calling +thread. Its cost stops rising only when the selection fills its 2,048-token +budget, so a long context sits near that figure. + +The heads are independent. A head reads its own query row and its KV head's +keys and values and writes its own row of the output; the one thing they +shared was a buffer of scores the width of the selection. `qwen_qsa.c` +gains `waste_qwen_qsa_attn_heads(h0, h1, ...)`, `waste_qwen_qsa_attn` is +that over every head (so `test_qwenparts` still checks the whole thing +against the reference), and `qwen_qsa_layer` runs one head per task with +its own row of scores. `qsa_scr` is therefore `n_heads` rows of the maximum +selection — about 200 KB on this model — and `waste_plan_memory` counts +the same, so the floor still describes what the load allocates. + +Against a build of the previous commit, unprofiled, 16 GiB cache, eight +threads: + +| | before | after | +|---|---:|---:| +| ~220-token context, decode, two runs | 10.33, 10.24 | 10.65, 10.36 | +| 2,801-token context, decode | 5.95 | 8.58 (+44%) | +| 2,801-token context, reading the prompt | 7.25 | 9.39 (+30%) | + +The first-position logits are byte-identical and every generated token +the same in all six runs. The long-context rows are one run a build; the +gap is fifteen times the run-to-run spread measured so far. Profiled at the +short context, attention went from 3.4 to 0.74 ms a step. + +What is left grows with the context and has not been touched: at 2,830 +tokens, 15 ms a step between the RoPE table (every past position's row +recomputed every token, though a row never changes), block pooling (every +complete block re-pooled, though a full block never changes, then a top-k +that rescans every block once per block kept) and the gather (the whole +selection converted on one core). + +## 81. The rest of QSA: work redone every token, and an argmax per block (2026-09-14) + +§80 left 15 ms a step of QSA at 2,830 tokens that grows with the context. +All three parts are now bit-identical rewrites; none adds state. + +**The RoPE table.** Every token, every QSA layer rewrote the cos/sin rows +for all T positions. A row is a function of its position alone, so once +written it is right for every later token, every layer and any session a +reset or restore produces. The model counts the rows it has filled +(`qsa_cs_n`) and writes only those past it: 4.3 ms a step → 0.00. + +**Block selection.** `waste_qwen_qsa_select` is now `score_blocks`, which +pools, rotates and scores a range of blocks — each block writes only its +own pooled row and its own score, so `qwen_qsa_layer` scores them on the +pool once there are 32 — followed by `pick`. The pick was a pass over every +block for each block kept, about 360,000 comparisons a layer at this +length; it is now a heapsort in the order that argmax took: the higher +score first, and a tie to the earlier block, over exactly the scores it +could ever have taken (above -1e30, which leaves out NaN). The order is the +point — attention sums the selected tokens in it, so the same set in a +different sequence would move the bits. `tests/test_qsa_pick.c` holds the +old loop verbatim and compares it with the new one over 4,000 cases built +to break an ordering: ties everywhere, NaN, -1e30 and -inf scores, and a +budget below, at and above the block count. Pooling was kept per token +rather than cached: a cache of pooled blocks would have been one more +thing a reset, a restore and a rewound position all had to invalidate, and +scoring them at once already took the part to 1.3 ms from 5.2. + +**The gather.** Each selected index writes its own rows of the F32 K/V +and its own slot of the selection, so the BF16 conversion goes in ranges: +5.3 → 1.2 ms. + +Against a build of §80's commit, unprofiled: + +| | before | after | +|---|---:|---:| +| ~220-token context, decode, two runs | 10.51, 10.53 | 10.74, 10.66 | +| 2,801-token context, decode | 8.53 | 9.47 (+11%) | +| 2,801-token context, reading the prompt | 9.37 | 9.83 (+5%) | + +First-position logits byte-identical and every token the same in all seven +runs. Across §80 and §81, decode at 2,801 tokens went from 5.95 to 9.47 +tok/s and QSA at that length from 77.7 ms a step to 15.0, of which the +attention — on the pool now — is 8.3. + +The step at either length is now mostly MoE: 57 ms of it at 2,801 tokens, +and its expert arithmetic the largest single part. + +## 82. A layer missing one expert ran all ten as rows (2026-09-14) + +§81 left MoE the largest part of a Qwen step. On §75's protocol (18-token +prompt, 200 decode tokens, 16 GiB cache, eight threads) it was 55.7 ms of +it, and 42.8 of those the routed experts' arithmetic. The profile hid where: +its LUT apply row is timed only on the row split, so the expert-parallel +layers showed up as a remainder of about 19 ms with no row of its own. + +**Counted per layer**, over that run with the prompt included: + +| path | layers | ms per layer | +|---|---:|---:| +| expert-parallel, all ten resident | 5,741 | 0.69 | +| row split, anything missing | 4,723 | 1.71 | + +§76's rule sent a layer down the row split if any of its ten records was +absent, and 2,343 of those 4,723 layers were missing exactly one. For one +read, nine resident experts gave up the single dispatch and took thirty, +over rows too short to fill the pool. + +**The thread split was a smaller thing than it looked.** A batch of ten on +eight threads went to `waste_parallel_for`, which cuts n into equal ranges: +five ranges of two, three threads with no expert. `waste_parallel_for_each` +now gives each item its own range and lets every participant take the next +one. On its own it measured within noise — 10.41 and 10.65 tok/s before, +10.43 and 10.50 after, expert arithmetic 43.75 → 42.95 ms — because it +only touches the layers that were already fast. It stays, since the staged +path below runs through it. + +**The staged path.** When the cache decides, `qwen_moe_layer` asks it +about each expert rather than the layer. The residents are held and run +first, one task per expert: they need no read, so holding them is no +barrier. The hint issued the reads for the rest before the first hold, so +those run underneath. Then the misses are held and run: as rows when there +are fewer than four, as tasks when there are more, because one expert on +one thread is slower than its rows on eight. The shared expert, which needs +no record either, is computed between the two stages. A forced `WASTE_XPAR` +or an explicit `WASTE_XPAR_BATCH` keeps §76's fixed batches. + +Each expert writes its own slice, and the sum still runs in route order +afterwards, so the order the experts are computed in does not reach the +bits. First-position logits were byte-identical and every generated token +the same in all ten real-container runs below. The suite's schedule check +gains a cold-cache arm: the fixture preloads its whole bank, so the default +arm never met a miss and only ever ran the first stage; with +`WASTE_PRELOAD=0`, five of its layers take the second. + +The threshold of four, as single runs on an instrumented build before the +shared expert moved: 11.57 tok/s at four, 11.16 with every miss a task, +11.33 with every miss as rows. + +Against a build of §81's commit, unprofiled: + +| | before | after | +|---|---:|---:| +| 16 GiB, decode, two runs | 10.37, 10.28 | 11.30, 11.14 (+8.7%) | +| 8 GiB, decode | 8.39 | 9.08 (+8.2%) | +| 2,801-token context, decode | 9.41 | 10.03 (+6.6%) | +| 2,801-token context, reading the prompt | 9.83 | 10.55 (+7.3%) | + +Hit rates and bytes read are unchanged at either cache size (90.2% and +17.72 GB; 79.8% and 36.4 GB). Profiled, decode only, 16 GiB: + +| ms/step | before | after | +|---|---:|---:| +| MoE | 55.7 | 49.0 | +| ├ expert arithmetic | 42.8 | 34.7 | +| │ ├ LUT build | 5.6 | 3.0 | +| │ └ LUT apply (row split only) | 18.6 | 2.6 | +| ├ expert I/O | 7.1 | 8.2 | +| ├ shared expert | 3.9 | 4.1 | +| └ router | 1.6 | 1.7 | + +**Expert I/O rose, and that is the next thing.** On the row split a read +landed under the arithmetic of the experts routed ahead of it; now the +residents finish first and the misses are waited for. Timed on the +instrumented build, over decode only: the second stage waited 9.4 ms a +step and computed for 3.8. A record is 1.72 MB and a read took 0.88 ms, +against about 0.6 ms of first-stage arithmetic. More readers did not buy +it back: every read got slower and so did the arithmetic beside it. + +| readers / depth | tok/s | ms per read | +|---|---:|---:| +| 2 / 2 (the default) | 11.29 | 0.885 | +| 4 / 4 | 10.91 | 1.239 | +| 8 / 8 | 10.73 | 1.445 | +| 4 / 10 | 10.94 | 1.228 | + +What would help is starting those reads earlier than the layer's own +router. Kimi's `moe_layer` already does, through `predict_next_moe` +(§34), and Qwen's does not. + +## 83. Qwen's router lookahead: the cheap half of the next layer's mix (2026-09-14) + +§82 ended on expert I/O: 8.2 ms of a step spent waiting for the reads of +experts no layer had asked for until its own router ran. Kimi starts them +a layer early (§34, §35); Qwen did not. Its default `WASTE_LOOKAHEAD` of 6 +now applies to Qwen as well, with a predictor of its own. + +**Which input to give layer L+1's router**, measured before any of it read +a byte: per layer transition over §75's 200 decode tokens at a 16 GiB +cache, against L+1's real routing and the cache's residency at that +moment. 0.71 of L+1's ten experts missed per transition. + +| predictor, top 6 | misses it would have started | wasted reads a layer | +|---|---:|---:| +| (a) layer L's MoE input — Kimi's | 32.1% | 0.43 | +| (b) the same plus L's MoE output | 32.7% | 0.42 | +| (c) L+1's MLP HyperConnection mix, on the streams after L | 43.0% | 0.09 | +| (d) the streams normalized with (c)'s weights, averaged, no gate | 42.7% | 0.23 | +| (e) the raw streams averaged | 38.0% | 0.45 | + +Kimi's predictor is weak here for a reason Kimi does not have: the MoE +input is one mix of four streams, and the next router will see a different +mix of different ones. (c) is nearly L+1's real MoE input, missing only +L+1's attention, and wider it is better still — 75% of misses at top 10 for +0.43 wasted. It is also a down projection over 10,240 values and an up +projection back, per layer: 10.71–11.07 tok/s against 11.32–11.48 without +the lookahead. (d) keeps its norms and drops the gate. + +A sixth, (f), predicted layer L from its own attention input, which is +already computed and so costs only the router projection. It chose worse: +93.8% hit rate and 21.05 GB read, against (d)'s 94.6% and 19.20 GB. + +**Width**, with (d), 16 GiB unless stated: + +| width | tok/s | hit rate | read | +|---|---|---:|---:| +| off | 11.34, 11.41 | 90.2% | 17.72 GB | +| 3 | 11.46, 11.45 | 92.3% | 18.08 GB | +| 4 | 11.09, 11.58 | 93.1% | 18.34 GB | +| 6 | 11.00, 11.68 | 94.6% | 19.20 GB | +| 8 | 11.11, 11.68 | 95.9% | 20.69 GB | +| 10 | 11.23, 11.66 | 97.0% | 22.84 GB | +| 12 | 11.32, 11.47 | 97.6% | 26.19 GB | +| 16 | 11.03, 11.27 | 98.3% | 35.15 GB | +| 8 GiB: off, 6, 10 | 9.09, 9.56, 9.07 | 79.8%, 88.3%, 92.7% | 36.4, 42.3, 55.3 GB | + +The first pass of this table dipped from width 4 to 8 and the second did +not; this machine's run-to-run spread was ±3% all afternoon. Six is where +the 8 GiB row peaks and the bytes have not yet started to climb. + +**Asked earlier, it waits the same.** Issuing (d)'s guess straight after +layer L routes instead of after L's MoE gives the reads a whole MoE more to +land in. Its hit rate was 94.3% against 94.6%, and expert I/O was 4.13 ms +a step against 4.19 at 16 GiB, 16.47 against 16.27 at 8. What is still +waited for is the misses no top-6 guess contains, not guesses that land +late — so the guess stays where the buffers it needs are already dead. + +**Against a build of §82's commit**, unprofiled: + +| | before | after | +|---|---|---| +| 16 GiB, decode, three runs | 11.40, 11.01, 11.00 | 11.86, 11.01, 11.67 | +| 8 GiB, decode, two runs | 9.01, 9.08 | 9.74, 9.86 (+8.2%) | +| 2,801-token context, decode | 10.09 | 10.06 | +| 2,801-token context, reading the prompt | 10.50 | 10.79 (+2.8%) | +| read: 16 GiB / 8 GiB / 2,801-token run | 17.7 / 36.4 / 179 GB | 19.2 / 42.2 / 230 GB | + +First-position logits byte-identical and every token the same in all +eleven runs; the suite's Qwen schedule check gains a cold arm with the +lookahead off, and on the fixture the default's demand misses fall from +7 to 3 with the same logits. + +Profiled, decode only, 16 GiB: expert I/O 8.25 → 4.58 ms a step, the +lookahead itself 1.88 ms (a row of its own now, inside MoE), MoE 49.5 → +46.5. + +**It costs bytes, and that is a choice rather than a finding.** The +2,801-token run read 28% more, 29,000 more records at 1.72 MB — about the +0.23 wasted reads a layer the table above predicted, over 2,833 steps, most +of them spent reading the prompt. On this machine's internal SSD reads are +not the budget and the lookahead was not slower in any configuration +measured, so it is on by default for Qwen. This file judges K3's changes on +bytes per token, and by that measure this one is a cost. A confidence +cutoff — prefetch a guess only when its score clears the rest by a margin — +might keep the useful reads and drop some of the wasted ones. It is not +measured. + +## 84. The expert kernel was not waiting on memory, it was waiting on a thread with two (2026-09-14) + +After §83, 16 GiB and eight threads, the step was 85.6 ms and 34.7 of it +the routed experts' arithmetic. Across thread counts, 64 decode tokens each: + +| ms/step | 1 thread | 2 | 4 | 8 | 12 | 1 → 8 | +|---|---:|---:|---:|---:|---:|---:| +| expert arithmetic | 152.5 | 87.9 | 52.0 | 34.7 | 32.4 | 4.40x | +| lm_head | 48.2 | 24.3 | 12.4 | 6.4 | 7.3 | 7.55x | +| GDN | 72.0 | 39.9 | 23.2 | 16.2 | 19.8 | 4.44x | +| HyperConnection | 26.2 | 14.7 | 10.1 | 9.6 | 10.5 | 2.72x | +| tok/s | 2.98 | 5.24 | 8.49 | 11.09 | 10.46 | | + +The trunk kernel runs 15 GB/s on one core at every call size. On eight its +large calls reach 92–100 GB/s, a third of this machine's 273 GB/s, and its +calls under 1 MB 25–28: those are dispatch, not arithmetic. Twelve threads +lose, as §47 found on other models — the efficiency cores are stragglers. + +The expert kernel falls behind at two threads already (1.7x against lm_head's +2.0x), which looked like memory: a gather is a load, an address and a load, +and its tables sit in a cache the performance cores share. **It is not.** +One engine thread, 32 decode tokens, with six other cores running each of +five loads, two passes in opposite orders: + +| load | expert arithmetic | lm_head | +|---|---|---| +| none | 149.9, 154.1 | 46.6, 48.2 | +| spin | 157.4, 158.9 | 49.3, 50.2 | +| memcpy, 64 MB buffers | 157.4, 157.3 | 49.5, 49.5 | +| address-dependent reads over 1 GB | 159.5, 159.8 | 49.4, 49.4 | +| the same over 2 MB each | 157.6, 157.9 | 49.3, 49.3 | + +Everything lost 4–5% to any load at all, spin included — the cluster +sharing power — and memory traffic added at most 1% on top. So the table +was not quantized; `WASTE_VQ8`'s case rests on its kernel being faster, not +on memory being the wall, and this entry did not test it. + +**What it was.** §82 gave every routed expert one task. Ten experts of +equal size on eight threads is two threads with two experts and a barrier +waiting for them: ten experts of work in two experts of wall time, 5x at +best, and 4.4x measured. `experts_staged` cuts the work into equal pieces +instead, in the three stages an expert depends on — every expert's gate and +up rows 128 at a time, then each expert's activation and down table, then +every expert's down rows 128 at a time. Three dispatches a layer, each +piece writing only its own rows through the same `vq_rows` and +`lutb_range` the per-expert task called, so the logits are unchanged. It +runs both of §82's stages, and replaces the split between rows and tasks +for the misses: every threshold of that split was slower. + +Against a build of §83's commit, unprofiled, logits byte-identical, every +token the same and the bytes read unchanged in all eleven runs: + +| | before | after | +|---|---|---| +| 16 GiB, decode, three runs | 11.87, 11.41, 11.51 | 12.23, 11.95, 11.98 (+3.9%) | +| 8 GiB, decode, two runs | 9.80, 9.75 | 10.19, 10.27 (+4.6%) | +| 2,801-token context, decode | 10.03 | 10.56 (+5.3%) | +| 2,801-token context, reading the prompt | 10.47 | 11.12 (+6.2%) | +| expert arithmetic, profiled, ms/step | 34.3 | 26.8 | + +In one binary with a switch, an hour earlier, the same change measured +11.19 and 11.14 against 12.17 and 11.92. The machine drifted by that much +between the two sessions, which is why the table above is the one kept. + +Measured and not adopted, each within noise of the plain version at 16 GiB: + +| variant | tok/s | +|---|---| +| rows per piece 64 / 128 / 256 / 512 | 12.18, 11.83 / 12.14, 12.14 / 12.03, 12.34 / 11.94, 12.18 | +| stage 2 in 16-vector pieces, the activation serial | 12.14, 11.99, 12.22 | +| gate and up tables built in one dispatch | 11.86, 12.07, 12.12 | +| both | 11.90, 12.09, 12.14 | +| none of them | 12.25, 11.90, 12.02 | + +The suite cannot see the splitting: the synthetic Qwen fixture's matrices +are 16 and 32 rows, under one piece. The eleven real-container runs above +are what checks it. + +**The rest of this measurement, for the next entry.** The SSD holding the +container is the internal one, 97% full: uncontended, one reader gets 3.1 +GB/s at 0.59 ms a 1.77 MB record, two get 4.4 GB/s at 0.84 ms, four get +3.9 GB/s at 1.88 ms. Beside eight spinning threads two readers take 1.08 ms +a record, beside eight memcpy threads 1.15 — the engine's 0.88 ms is this +drive plus the arithmetic beside it. Two readers is the drive's best. + +## 85. A bigger cache is a long-context fix, and one quantization per vector (2026-09-15) + +Two of the three places §84 left to look. + +**The cache.** 16 GiB has been every Qwen measurement's protocol since §76, +not a recommendation. On the same build, one process per run, peak RSS from +`/usr/bin/time -l`, swap unused before and after: + +| expert cache | 200 tokens, tok/s | read | peak RSS | +|---|---|---:|---:| +| 16 GiB | 11.91, 11.47 | 19.20 GB | 20.0 GB | +| 20 GiB | 12.60, 11.58 | 18.34 GB | 22.5 GB | +| 24 GiB | 12.06, 11.33 | 18.34 GB | 22.6 GB | +| 28 GiB | 11.76, 11.94 | 18.34 GB | 22.6 GB | + +| expert cache | 2,801-token prompt | decode | hit rate | read | peak RSS | +|---|---:|---:|---:|---:|---:| +| 16 GiB | 11.20 tok/s | 10.70 | 95.3% | 229.6 GB | 20.1 GB | +| 24 GiB | 11.88 | 11.19 | 98.5% | 72.2 GB | 28.7 GB | + +A 200-token session evicts nothing from 20 GiB up, and what it still misses +is first use: nothing a cache can hold. The second pass of that table ran +beside two compiles and says nothing about speed. A long prompt is +different — 69% fewer bytes and about 5% on both phases, single runs. + +None of it needs a change. `waste run` with no `--budget` already takes +35.18 of this machine's 48 GB, 32.31 of it expert cache. + +**One quantization per vector.** The trunk's calls under 1 MB ran at 25.9 +GB/s against the same kernel's 15 GB/s on one core (§84): per call, an +activation quantization and a dispatch of their own. Qwen reads one vector +three and four times over. GDN projects its input through `in_proj_qkv`, +`_z`, `_a` and `_b`; QSA through q, k, v and the indexer; the MoE through +the router and the shared expert's gate, up and gate scalar. + +`matvec_t_batch` quantizes the vector once — i8mm's planes are a function +of the vector and the group size alone — and cuts every tensor's rows at +the multiples `mv_chunk` would have, so i8mm's two-row tiles stay where +they were, then hands every piece to one dispatch. Each row is the same +kernel call on the same bytes, so the logits do not move; anything the +shared planes do not fit goes through `matvec_t` as before. GDN's conv +reads only the first projection and now follows all four. The shared +expert keeps its gate and up outputs in `m->ff` until it runs; the serial +loop is the one path that writes there, and it has the shared expert +redo them. + +Against a build of §84's commit, unprofiled: + +| | before | after | +|---|---|---| +| 16 GiB, decode, three runs | 12.02, 12.03, 12.00 | 13.07, 12.28, 12.41 (+4.7%) | +| 2,801-token context, decode | 10.56 | 10.88 (+3.0%) | +| 2,801-token context, reading the prompt | 11.14 | 11.69 (+4.9%) | + +Logits byte-identical and every token the same in all eight runs. Profiled, +decode only, 16 GiB: + +| | before | after | +|---|---:|---:| +| trunk calls under 1 MB | 25.9 GB/s | 38.9 GB/s | +| QSA k, v, indexer projections | 26–31 GB/s | 97 GB/s | +| GDN, ms/step | 16.6 | 15.1 | +| QSA, ms/step | 5.5 | 4.9 | +| router and shared expert, ms/step | 5.6 | 4.1 | + +The router row now carries the shared expert's gate and up projections; +the profile splits a batch's time among its tensors by bytes. + +**For the CLI, not measured on it.** `--threads 0` is one thread per logical +CPU, twelve here, and §84 measured twelve 6% slower than eight on Qwen: +the efficiency cores are stragglers. A default that counted performance +cores would be worth measuring on every model before it is one. + +## 86. The pool parks 300 times a token, and closing gaps is worth 2% (2026-09-16) + +§84 left seven milliseconds of the expert stages above their eight-thread +ideal and HyperConnection's non-matvec work at 4.3 ms. Timed inside, with +clock reads around each piece, 200 decode tokens at 16 GiB: + +| HyperConnection, per mix | us | ms/step | +|---|---:|---:| +| RMSNorm, four streams on the pool | 21.6 | 2.09 | +| down projection, its quantization included | 40.2 | 3.90 | +| up projection | 28.0 | 2.72 | +| sigmoid over 10,240, on the pool | 10.1 | 0.98 | +| weighted sum over streams | 4.4 | 0.42 | +| inject projection | 5.0 | 0.49 | +| combine | 3.6 | 0.35 | + +| expert stages, per call | us | ms/step | +|---|---:|---:| +| holds — the misses' reads | 103.6 | 6.66 | +| gate and up tables | 40.0 | 2.57 | +| stage 1, gate and up rows | 214.8 | 13.79 | +| stage 2, activation and down table | 38.7 | 2.48 | +| stage 3, down rows | 127.5 | 8.19 | + +A norm over 10,240 floats is about 10 us of arithmetic; it took 21.6. With +`WASTE_SPIN=200000` in the same instrumented build it took 8.1, the down +projection 21.0, the sigmoid 4.8 and the expert tables 19.4. What the pieces +cost is mostly waking a parked pool. + +**How often.** A probe on `waste__pool_run` timed the calling thread's +serial stretch before every dispatch, keyed by the function dispatched. +Per token, the dispatches that followed a stretch longer than a worker's +8 us spin: + +| next dispatch | per token | after a gap | serial ms | +|---|---:|---:|---:| +| GDN, QSA and router batches (§85) | 96 | 96 | 1.47 | +| HyperConnection norm | 97 | 65 | 1.53 | +| GDN recurrence | 36 | 36 | 1.66 | +| i8mm matvecs | 339 | 55 | 1.41 | +| activation quantization | 48 | 46 | 0.62 | +| VQ tables | 96 | 28 | 0.54 | +| stage 1 (after the misses' reads) | 64 | 14 | 5.98 | + +About 330 wakes a token, and 7.7 ms of serial time outside the expert +reads. `waste_find` was a suspect — a linear `strcmp` over every tensor — +and is not: one lookup averaged 0.1 us. + +**A wake is the scheduler, not the primitive.** Seven workers at the fast +group's quality of service, parked, woken together after 2 ms idle, 500 +times, until the first and the last acknowledged: + +| primitive | first, median | last, median | last, p90 | last, p99 | +|---|---:|---:|---:|---:| +| condvar broadcast under a mutex | 18.6 us | 37.2 | 67.6 | 184 | +| `os_sync_wake_by_address_all` | 18.6 | 37.5 | 66.2 | 178 | + +No cheaper wake to swap in, and a barrier waits for the last one. + +**Spinning through them is §78's trade again.** Unprofiled, same build, +user+system CPU for the whole run: + +| `WASTE_SPIN` | tok/s | CPU s | +|---:|---|---:| +| 20,000 | 12.89, 12.63 | 100.6, 103.6 | +| 50,000 | 13.26, 13.32 | 108.7, 107.5 | +| 100,000 | 13.33, 13.61 | 113.4, 110.0 | +| 200,000 | 13.50, 13.68 | 114.9, 112.3 | + ++6.5% for +12% CPU. The instrumented build said +10%: clock reads lengthen +the gaps being measured, as §78 found of the profiler. + +**Closing gaps instead**, each the same function over the same elements in +the same order, so the logits do not move: + +- *HyperConnection's norm task* now also does the combine that finishes the + previous block, when the mix is the MLP's, and quantizes the stream's + weight groups for the down projection, which reads the planes through + `matvec_t_prequant`. `prequant_ok` says when a tensor can: i8mm, four + bits, and the caller's piece a whole number of groups. +- *Its gate* is one job of pieces: the sigmoid and the sum over streams, a + 256-wide hidden range at a time, and the inject projection's four rows, + each the `dotf` `matvec` would have taken. A quantized inject falls back + to `matvec_t`. +- *GDN's short conv* runs a range of channels per task, 46 us of SiLU a + layer that had been serial in front of the recurrence. +- *GDN's gated RMSNorm* and out_proj's quantization moved into the + recurrence's per-head tasks. + +HyperConnection is three dispatches where it was six. In the instrumented +build the down projection went 42.4 → 27.3 us and the sum and sigmoid 16.0 +→ 11.4; the recurrence's serial gap 1.70 → 0.20 ms a token and the parked +quantizations 47 → 11. Folding the inject and the combine in cut the +batches' parked count 96 → 82, which is how we know most of their gap is +not those two. + +Three runs a side could not see any of it — the arms landed within 0.2 +tok/s of each other in both directions. Against a build of §85's commit, +six alternated pairs: + +| | tok/s | mean | CPU s | +|---|---|---:|---:| +| before | 13.04, 13.24, 13.03, 12.96, 13.09, 13.15 | 13.09 | 101.3 | +| after | 13.73, 13.26, 13.17, 13.32, 13.29, 13.32 | 13.35 | 100.8 | + ++2.0%, faster in all six pairs, at the same CPU. The 2,801-token context +did not move (prompt 12.13 → 12.22, decode 11.27 → 11.23, one run each). +Logits byte-identical and every token the same in every run, and also +against the same commit with `WASTE_TRUNK_KERNEL=0` and `=1`, where +nothing can take prequantized planes and every fallback runs. Profiled: +HyperConnection 10.1 → 9.3 ms a step, GDN 14.7 → 13.8. + +Two percent here was not obviously worth the code. It went in because the +gaps are a property of this machine's wake latency and eight cores, and a +machine with more cores or a slower scheduler pays more for each one. + +## 87. QSA's attention, four scores at a time — and the product that must not fuse (2026-09-16) + +At a 2,801-token context QSA's attention was 8.3 ms of a 87 ms step, on +the pool since §80 and scalar inside: per selected token a 256-wide dot +against the query, then a 256-wide accumulation of that token's values. + +The dot was not short of arithmetic, it was short of independence — each +element's multiply-add waits for the one before it, ~4 cycles deep, +whatever else the core could issue. So four selected tokens are scored in +one pass now, each with its own accumulator summing its own dimensions in +its own order: §41's trick on the VQ gather, for the same reason. The +value accumulation is the other way round — every output dimension sums +the tokens in order — so its lanes run along the dimension, four vectors +at a time. Both leave every element's sequence where it was. + +Against a build of §86's commit, 16 GiB, three pairs: + +| | before | after | +|---|---|---| +| 2,801-token context, decode | 11.35, 11.35, 11.47 | 11.68, 11.81, 11.72 (+3.0%) | +| 2,801-token context, reading the prompt | 11.89, 12.13, 12.15 | 12.32, 12.51, 12.35 (+2.6%) | +| attention, ms/step | 8.34 | 5.36 | +| QSA, ms/step | 14.59 | 11.78 | + +A short context selects a handful of tokens and measured 0.5% slower in +three pairs of three, so the four-at-a-time pass is taken from 32 +selections up; below that the plain loop runs and the short prompt is back +to level. + +**What nearly shipped instead.** The first two versions were not +bit-identical, and neither was wrong about the order of anything. The loop +they replaced compiles — at -O2, no fast-math — to four *products* in one +vector and a scalar chain of adds: the products are independent, the sum +order is not, so clang vectorizes the multiplies and leaves the additions +alone. Each product is therefore rounded on its own. Write the same +arithmetic as `s += q * k` in four accumulators and clang's SLP pass packs +them into a vector too, but emits a multiply and an add where the original +contracted to one fused multiply-add; write it as `fmaf` and every product +fuses. Both round differently from the original — by one ulp, on a dot of +256 terms, over 2,048 tokens and 24 heads. The logits moved in the last +bits and the text diverged a few hundred tokens in. + +The fix is to say exactly what the original does: a product, then a sum, +as two statements, which C's contraction rules leave alone. `-fno-vectorize +-fno-slp-vectorize` also fixes it, and is not a fix — it is a build flag +this file cannot rely on. + +**A bit-identical kernel needs a bit-identical test.** `tests/run.sh` ran +the whole suite green through both wrong versions: its Qwen checks compare +paths of *one* build against each other, and both paths had the same +kernel. `tests/test_qsa_attn.c` holds the old loops verbatim and compares +them with the new ones over 40 random cases with out-of-range selections +mixed in, in one translation unit, where a compiler that transforms one +and not the other is exactly what is being looked for. It is the same +shape as §81's `test_qsa_pick`, and it was written after the fact rather +than before, which is the part to do differently next time. +## 88. A feasibility gate does not need the download (2026-09-15) DeepSeek-V4.1-Flash is 510 GB in 48 shards. The gate that had to run before any of it was fetched — does 3-bit VQ survive experts that are *already* fp4? @@ -5457,7 +6537,7 @@ bytes.** Two of the eight gates in this file spent their cost on a download or a conversion that the measurement itself did not need. Range requests against a published index is a way to not do that again. -## 75. Twenty-one strings is not a tokenizer corpus, and it never was (2026-09-15) +## 89. Twenty-one strings is not a tokenizer corpus, and it never was (2026-09-15) `tools/tokdiff.py` opens with a comment saying "twelve short ASCII strings is not a tokenizer corpus". It then lists twenty-one strings and, until @@ -5519,7 +6599,7 @@ which is what it did before and is why "22914/24021 identical" passed. That last part is §73 again, in the one file that had already written the warning down. -## 76. The oracle was wrong, and only the shape of the error said so (2026-09-15) +## 90. The oracle was wrong, and only the shape of the error said so (2026-09-15) DeepSeek-V4.1's forward pass came up 0.7% off against `tools/ds41_ref.py` on the first run — same argmax, plausible logits, a number that could have @@ -5569,10 +6649,10 @@ Three things this is evidence for: - **The test corpus has to reach the mechanism.** Four tokens against a four-slot window never wraps the ring, never fills a compressed cache and never gives the candidate filter two blocks to choose between. Twelve - does. §75 was the same lesson about a tokenizer corpus, three commits + does. §89 was the same lesson about a tokenizer corpus, three commits earlier, and it did not transfer on its own. -## 77. A speculative batch of five reads 3.45 tokens' worth (2026-09-15) +## 91. A speculative batch of five reads 3.45 tokens' worth (2026-09-15) Speculative decoding verifies K draft tokens in one backbone pass, and on a GPU that is nearly free — the pass is compute-bound and the K tokens ride @@ -5623,7 +6703,7 @@ and it is the one number left. What is settled is everything else: the threshold it has to clear, and that a batched CSA2/mHC forward path is the price of finding out. -## 78. Four releases of a JSON reader that did not decode JSON (2026-09-15) +## 92. Four releases of a JSON reader that did not decode JSON (2026-09-15) DeepSeek-V4.1 loaded, ran, matched its oracle to 0.0025% on the real container — and could not tokenize `<｜User｜>`. Markup mode returned @@ -5675,14 +6755,14 @@ the same corrupted string read from the same file. A rendering artifact at the end of a correct generation looks like a cosmetic bug in the printer. It was the tokenizer's security boundary, seen from the other side. -## 79. Three checks that were wrong about a correct engine (2026-09-15) +## 93. Three checks that were wrong about a correct engine (2026-09-15) The DeepSeek-V4.1 container converted, matched its oracle to 0.0025% and answered "The capital of France is" with " Paris." — and the suite said 4 failures. Every one of them was the check, not the engine. - **`tests/run.sh` tested the tokenizer with `grep -q identical`**, and the - string `"22914/24021 identical"` contains that word. §75. + string `"22914/24021 identical"` contains that word. §89. - **`verify_container.py` kept a second copy of how a checkpoint names its experts**, an inline probe for `mlp/gate_proj` falling back to `block_sparse_moe/w1`, while `convert.py` had the same fact in @@ -5696,7 +6776,7 @@ failures. Every one of them was the check, not the engine. rate is zero and not low. It answered 284 misses → 286 on one run and fewer on the next: a verdict decided by noise. -Plus §78, the escaped `specials.json`, which was a real defect — so the +Plus §92, the escaped `specials.json`, which was a real defect — so the board read 4 failures over 1 bug. **A suite is only exercised by a model it has not seen.** These three sat diff --git a/docs/QWEN.md b/docs/QWEN.md new file mode 100644 index 000000000..0a68e90f5 --- /dev/null +++ b/docs/QWEN.md @@ -0,0 +1,198 @@ +# Qwen3.8-Flash-Next + +Text inference for `qwen4_exp` / `qwen4_exp_text`. Four architectural +pieces that no other model here has, and a container format that does not +change to hold them. + +Pinned for every measurement below: + + Qwen/Qwen3.8-Flash-Next + revision de4b8e4d43b917e7706784d8bb445c9af86a3540 + +131 shards, 360,013,002,208 bytes on the hub, index SHA-256 +`99e815241ef03325536b0aaa4441deea45174c17fae31e10f0bb456410c590de`. +Fetch the revision, not `main` — this repository's `main` has moved since, +and a container built from a different one is not the container these +numbers describe. + +## What is different about it + +| | Kimi / GLM | Qwen3.8-Flash-Next | +|---|---|---| +| recurrent layers | Kimi Delta Attention | **Gated DeltaNet** — per-head scalar decay, 16 QK heads repeated onto 48 V heads | +| attention layers | MLA over a compressed latent | **Qwen Sparse Attention** — an indexer scores 4-key mean-pooled blocks, keeps the best 512 plus the token tail, and attention runs over the *original* K/V | +| residual | one stream (four on GLM's mHC) | **HyperConnection** — four streams mixed through a rank-320 bottleneck and recombined through a per-branch gate | +| extra embedding | none | **PLE** — an n-gram lookup at one layer, 16 hashed head tables of ~20 M rows each | +| router | sigmoid plus a learned bias | **softmax over all 512, then top-k** | +| experts on disk | one tensor per expert per matrix | **two packed tensors per layer** | + +Layer mix is 36 `linear_attention` and 12 `full_attention`, stated by +`layer_types` and read from it. A container that does not say is refused: +"every layer is GDN" is a plausible-looking default that would be wrong in +every token and visible in none. + +Shape: hidden 2560, 48 layers, vocab 248320, 512 experts with top-10 +routed plus one gated shared expert of width 640, GDN 16×128 QK and 48×128 +V, QSA 24 query and 2 KV heads at dim 256 with 64 rotary dims, indexer MQA +at dim 128 with a 2048 budget, `hc_count` 4 / `hc_lowrank` 320, n-gram +size 3 over 16 heads, `ple_layer_ids: [2]`. + +`ple_layer_ids` is 1-based and the PLE tensors live on +`model.language_model.layers.1`. Two spellings of the same layer; the +container carries the 0-based one. + +## Converting + +```bash +tools/fetch_weights.sh --repo Qwen/Qwen3.8-Flash-Next \ + --revision de4b8e4d43b917e7706784d8bb445c9af86a3540 \ + --dest /path/to/raw + +uv run --with torch --no-project python tools/convert.py \ + --src /path/to/raw --out /path/to/qwen38-flash-next.waste +``` + +No Qwen-specific flags. Three things happen that do not happen for any +other family: + +* **Packed experts.** `experts.gate_up_proj` [E, 2I, H] and + `experts.down_proj` [E, H, I] hold a whole layer's routed experts. They + are split into ordinary WEXP records — one expert per 4 KiB-aligned + record, format v0, nothing about the read path changes. The layout is + validated against its invariant rather than against Flash-Next's + dimensions, so a larger family member with the same packing converts on + the same code. + +* **PLE.** 128 source shards become 16 Q8G head tensors on the trunk. A + head is ~12 GiB as f32, so it is quantized 64 Ki rows at a time with one + source shard resident; Q8G groups along the last dimension, which is + what makes the concatenation of the batches equal to quantizing the head + whole. The i64 offset and vocabulary-size tables go into the manifest as + integers — they are primes near 2×10⁷ and do not survive a float. + +* **`--jobs` defaults to 1.** A worker holds a whole layer's packed pair, + 3.3 GiB of BF16 before the f32 it dequantizes into. Three at once is + what turns a conversion into a swap storm. An explicit `--jobs` wins. + +`--reclaim` works, with the n-gram shards as a consumer of their own: +`build_ple` runs after the trunk pass, so a shard holding both an expert +and an n-gram slice is not released until the 16 heads are on the trunk. +The vision tower and the MTP layer have no consumer at all here and are +released first. + +`tools/verify_container.py` reads packed sources and checks PLE rows +against their source shard, three rows per head, without dequantizing a +head. + +## Running + +```bash +./waste run /path/to/qwen38-flash-next.waste "The capital of France is" \ + --budget 8G +``` + +Measured on an Apple silicon laptop, 12 logical CPUs (8 performance, 4 +efficiency), 48 GiB of RAM, container on the internal SSD, at the commit +that ships this file: + +| | | +|---|---:| +| container | 123.9 GiB | +| ├ `trunk.bin` | 80.51 GiB | +| ├ expert banks | 42.5 GiB | +| resident trunk | 2.60 GB | +| memory floor | 3.11 GB | +| one token's expert working set | 888.7 MB | +| parameters | 176.94 B total, 57.87 B active per token | + +The gap between an 80.51 GiB trunk file and a 2.60 GB resident trunk is +the point of the PLE design: the 16 n-gram heads are 78 GiB of the file +and are read one row per head per token, never held. `waste_plan_memory` +excludes them from the resident set for the same reason it excludes the +embedding table. + +Throughput, 48 tokens greedy from the same prompt: + +| expert cache | tok/s | hit rate | +|---:|---:|---:| +| 4 GiB | 3.20 | 8% | +| 8 GiB | 4.97 | 64% | +| 16 GiB | 4.92 | 88% | + +**8 GiB is the knee and there is nothing above it.** Below one token's +working set the hit rate collapses rather than degrades — the 4 GiB row is +that, not a gentle slope — and above 8 GiB a better hit rate buys no time +at all: the reads it saves were already overlapping the arithmetic. That +is the whole reason the default budget is not "as much as the machine +has". + +Threads, at 8 GiB: 5.04 tok/s at the default (one per logical CPU), 5.38 +at `--threads 8`, 5.14 at 6, 5.09 at 12. Eight — the performance-core +count on this machine — is worth about 7%, which is real but is a property +of this machine and not of the architecture, so it is left to `--threads` +rather than compiled in as a default. LEARNED §47 has the same finding +inverting between two other models. + +Routed experts go through the existing expert-parallel path and its +existing per-layer decision: one task per expert when the records are +already resident, one task per row range when holding a batch would +barrier the read-ahead. `WASTE_XPAR=0/1` still forces it. + +The throughput above was measured with the f32 trunk kernel, before the +two changes that moved it most; LEARNED §75–77 have the numbers since, and +the table stands as what that commit measured. + +### Trunk kernel + +A Qwen load runs the 4-bit trunk through **i8mm** unless +`WASTE_TRUNK_KERNEL` is set — 0 pins the exact f32 arithmetic, 2 i8mm, +3 SMLAL. Against f32 on this machine, 16 GiB cache and eight threads: 9.91 +tok/s against 7.70 over 200 decode tokens, and 11% at a 6K-token context. +It is not the exact arithmetic. Over a 5,918-token prompt of real text +i8mm's next-token perplexity is 3.698 against f32's 3.712, top-1 +agreement 96.5%, 97.65% of routed experts the same, and nothing grows past +QSA's 2,048-token selection budget; `tests/kernel_kl.c` is the harness and +LEARNED §77 the measurement. The kernel is process-wide: a process that +loads Qwen and then another architecture keeps i8mm for both. + +## Correctness + +* **Components.** `tests/test_qwenparts.c` dumps every intermediate of the + five kernels and `tools/qwenparts_ref.py` recomputes them in PyTorch + from the published equations — at the official geometry as well as at + toy sizes. Largest disagreement across the dump: 2.4e-7 absolute. + +* **Whole forward pass.** `tools/qwen_container_ref.py` implements the + same forward pass in PyTorch reading the same container, so the + comparison is against an independent decode of identical weights. On the + synthetic fixture the routed expert ids and weights match **exactly** at + every layer and the logits argmax matches. The worst hidden-state + difference is 2.4e-7 absolute and the final-logit difference is 1.9e-6. + +* **Tokenizer.** Qwen's pre-tokenization pattern splits every digit into + its own piece where Kimi's and GLM's take up to three. The engine is + told which through `tokenizer_digit_run`; `tools/hf_tokenizer.py` + refuses a pattern it does not recognise rather than approximating one. + Against the release's own tokenizer the C encoder matches on every + string in `tests/test_qwen_tok.py`, numbers included. The rank file this + path writes is byte-identical to one built from `vocab.json` through the + GPT-2 byte map. + +* **End to end.** On the pinned checkpoint at an 8 GiB budget, "The + capital of France is" completes to "**Paris**". + +## Not supported + +Text only. Not converted, not executed, and refused rather than half-done: + +* the **vision tower** (`model.visual.*`, 333 tensors) and everything that + configures it — no `vision.json` is written, so a container cannot be + handed an image; +* **video and audio**; +* the **MTP layer** (`mtp.*`, 31 tensors) and speculative decoding; +* **native Qwen serving** — `serve/` has no Qwen chat format, so the + OpenAI-compatible server falls back to `chatfmt.py`, which refuses + tools, thinking and images by name rather than dropping them. + +The chat template is carried into the container as metadata, as every +other release's is. Nothing reads it. diff --git a/src/model.c b/src/model.c index 5412e237d..55967a6a7 100644 --- a/src/model.c +++ b/src/model.c @@ -28,6 +28,11 @@ #include "platform.h" #include "threads.h" #include "kda.h" +#include "qwen_gdn.h" +#include "qwen_hc.h" +#include "qwen_moe.h" +#include "qwen_ple.h" +#include "qwen_qsa.h" #include "simd.h" #include "waste_backend.h" #include "waste_metal.h" @@ -37,8 +42,13 @@ /* ---- lightweight phase profiling (WASTE_PROFILE=1) --------------------- */ #include -double waste_prof[16]; -uint64_t waste_prof_n[16]; +double waste_prof[32]; +uint64_t waste_prof_n[32]; +/* How much of each phase was trunk matvec: P_TMV's total as it moved while + * the phase ran. A phase that is mostly projections and a phase that is + * mostly the loops between them look the same in waste_prof, and they are + * not fixed the same way. */ +double waste_prof_tmv[32]; uint64_t waste_tmv_bytes; int *waste_route_cap; int waste_route_n, waste_route_cap_n; /* WASTE_TRUNK_CHECK=1: run the f32 reference beside whichever quantized @@ -53,8 +63,26 @@ unsigned long long waste_tcheck_n; /* matvec_t by call size: [<1MB, <8MB, <32MB, rest] */ double waste_tmv_t[4]; uint64_t waste_tmv_b[4], waste_tmv_c[4]; +/* Slots are read by number in tests/test_forward.c and tests/sweep.c, so a + * phase is appended, never inserted. Qwen reuses the roles it shares with + * Kimi — P_KDA is its recurrent layer (GDN), P_MLA its attention layer + * (QSA), P_KDAK the recurrence inside the first, P_ROUTE the whole MoE — + * and gets slots of its own for the pieces nothing else has. */ enum { P_LUTB, P_KDA, P_MLA, P_ROUTE, P_EDEQ, P_EMM, P_HEAD, P_LUTA, P_MM, - P_TMV, P_KDAK }; + P_TMV, P_KDAK, + P_QHC, /* HyperConnection mixes and combines, final mixer too */ + P_QPLE, /* n-gram embedding: row reads, projections, conv */ + P_QSHX, /* shared expert and its gate */ + P_QSAK, /* QSA block selection, K/V gather, attention */ + P_QRTR, /* router projection and top-k */ + /* Inside P_QSAK. Three of the four grow with the context rather + * than the token, so which one matters depends on how long the + * context is, and a short-prompt profile cannot say. */ + P_QSAR, /* the RoPE cos/sin table the block scores rotate by */ + P_QSAS, /* block pooling, scoring and top-k */ + P_QSAG, /* selected K/V, BF16 to F32 */ + P_QSAA, /* attention over the selection */ + P_QLAH }; /* router lookahead: the next layer's guess, inside P_ROUTE */ static int prof_on = -1; static pthread_mutex_t prof_mu = PTHREAD_MUTEX_INITIALIZER; static double pnow(void) @@ -62,9 +90,11 @@ static double pnow(void) struct timespec t; clock_gettime(CLOCK_MONOTONIC, &t); return t.tv_sec + t.tv_nsec / 1e9; } -#define PROF_START(b) double _t##b = prof_on ? pnow() : 0 +#define PROF_START(b) double _t##b = prof_on ? pnow() : 0, \ + _m##b = prof_on ? waste_prof[P_TMV] : 0 #define PROF_END(b) do { if (prof_on) { pthread_mutex_lock(&prof_mu); \ waste_prof[b] += pnow() - _t##b; waste_prof_n[b]++; \ + waste_prof_tmv[b] += waste_prof[P_TMV] - _m##b; \ pthread_mutex_unlock(&prof_mu); } } while (0) static char *slurp(const char *path, size_t *len) @@ -184,6 +214,7 @@ static int sdot_on = 0; /* 1 = also quantize activations (SDOT path) */ * costs in accuracy. */ enum { TK_F32 = 0, TK_SDOT = 1, TK_I8MM = 2, TK_SMLAL = 3 }; static int trunk_kern = TK_F32; /* WASTE_TRUNK_KERNEL */ +static int trunk_kern_env = 0; /* set explicitly; waste_model_load */ static int sdot4_sg = 32; /* TK_SDOT only: activations per int8 scale */ static int i8mm_on = 0; /* SMMLA batched matmul; costs activation int8 */ static const char *dump_route = NULL; /* WASTE_DUMP_ROUTE, see moe_layer */ @@ -231,6 +262,7 @@ static inline void pf_wide(int bit, int n, int min_chunk, waste_range_fn fn, else waste_parallel_for(n, min_chunk, fn, arg); } static int xpar_batch = 4; /* WASTE_XPAR_BATCH, see moe_layer */ +static int xpar_batch_set = 0; /* given explicitly; qwen_moe_layer */ static pthread_once_t model_opts_once = PTHREAD_ONCE_INIT; static void model_opts_init(void) @@ -271,6 +303,7 @@ static void model_opts_init(void) * is also why the same kernel measures KL 0.0013 on Kimi-Linear's 27 * layers. i8mm buys 43x the accuracy for 83% of the speed. */ e = getenv("WASTE_TRUNK_KERNEL"); + trunk_kern_env = e != NULL; trunk_kern = e ? atoi(e) : TK_F32; if (trunk_kern < 0 || trunk_kern > TK_SMLAL) trunk_kern = TK_F32; if ((trunk_kern == TK_SDOT || trunk_kern == TK_I8MM) && @@ -340,6 +373,7 @@ static void model_opts_init(void) /* Experts held — and so barriered — at a time. Small keeps the reads * overlapping the arithmetic; large gives the pool more to chew on. */ { const char *e2 = getenv("WASTE_XPAR_BATCH"); + xpar_batch_set = e2 != NULL; xpar_batch = e2 ? atoi(e2) : 4; if (xpar_batch < 1) xpar_batch = 1; if (xpar_batch > WASTE_PF_MAX) xpar_batch = WASTE_PF_MAX; } @@ -616,10 +650,10 @@ static void mvq4_rows_smlal(int b, int e, void *p) * Inside the same guard as its only caller, or -Wunused-function fires on * every build that cannot reach it. */ #if defined(__ARM_NEON) || defined(__aarch64__) -static void quant_act4_mm(const float *x, int n, int g, int8_t *q, float *sc) +static void quant_act4_mm_group(const float *x, int n, int g, int k, + int8_t *q, float *sc) { - const int ng = (n + g - 1) / g; - for (int k = 0; k < ng; k++) { + { const int beg = k * g, end = (beg + g < n) ? beg + g : n; float amax = 0; for (int i = beg; i < end; i++) { @@ -648,6 +682,28 @@ static void quant_act4_mm(const float *x, int n, int g, int8_t *q, float *sc) } } +typedef struct { const float *x; int n, g; int8_t *q; float *sc; } qa4_arg; + +static void quant_act4_mm_range(int b, int e, void *p) +{ + const qa4_arg *a = (const qa4_arg *)p; + for (int k = b; k < e; k++) quant_act4_mm_group(a->x, a->n, a->g, k, a->q, a->sc); +} + +/* Every group reads its own activations and writes its own planes and + * scale, so the groups go to the pool as they are and the bytes are the + * serial loop's. Serially this sat on the calling thread in front of every + * i8mm matvec — about 15 us before Qwen's 10,240-wide HyperConnection down + * projection, twice a pool worker's 8 us spin, so that matvec started by + * waking the pool. Below 32 groups the dispatch is not worth it. */ +static void quant_act4_mm(const float *x, int n, int g, int8_t *q, float *sc) +{ + const int ng = (n + g - 1) / g; + qa4_arg a = { x, n, g, q, sc }; + if (ng >= 32) waste_parallel_for_fast(ng, 4, quant_act4_mm_range, &a); + else quant_act4_mm_range(0, ng, &a); +} + #endif /* Activations for mvq4_rows_smlal: int16, one amax per weight group, @@ -785,32 +841,113 @@ static void matvec_t_inner(waste_model *m, float *y, const waste_tensor *t, * Q4G read once per token — the single largest byte term in a decode step * (docs/LEARNED.md §59). It has no bucket of its own in the profile, which * is why "kda 29%" was being read as if it were all recurrence. */ +waste_tmv_role waste_tmv_roles[WASTE_TMV_ROLES]; +int waste_tmv_nroles; +/* The quantization inside the matvec call being timed. Only matvec_t's + * calling thread writes it, and only under WASTE_PROFILE. */ +static double tmv_quant_dt; + +/* "model.layers.17.linear_attn.in_proj_z.weight" -> "linear_attn.in_proj_z" */ +static void tmv_role_name(const char *name, char *dst, size_t cap) +{ + const char *p = strstr(name, "layers."); + if (p) { + p += 7; + while (*p >= '0' && *p <= '9') p++; + if (*p == '.') p++; + } else { + p = name; + if (!strncmp(p, "model.", 6)) p += 6; + } + snprintf(dst, cap, "%s", p); + const size_t n = strlen(dst); + if (n > 7 && !strcmp(dst + n - 7, ".weight")) dst[n - 7] = 0; +} + +/* One call's worth of profile: the P_TMV total, its size bucket, and the + * tensor's role row. Caller holds prof_mu. */ +static void tmv_account(const waste_tensor *t, int out, int in, double dt, double dtq) +{ + const uint64_t nb = t ? (uint64_t)out * t->rowbytes : 0; + const int bk = nb < (1u<<20) ? 0 : nb < (8u<<20) ? 1 : nb < (32u<<20) ? 2 : 3; + waste_prof[P_TMV] += dt; waste_prof_n[P_TMV]++; + waste_tmv_bytes += nb; + waste_tmv_t[bk] += dt; waste_tmv_b[bk] += nb; waste_tmv_c[bk]++; + if (t) { + /* The slot is the profiler's own cache on a struct the model owns, + * so writing through the const is writing to what was calloc'd. */ + waste_tensor *tw = (waste_tensor *)t; + int si = tw->prof_slot - 1; + if (si < 0) { + char role[96]; + tmv_role_name(t->name, role, sizeof role); + for (si = 0; si < waste_tmv_nroles; si++) + if (!strcmp(waste_tmv_roles[si].role, role)) break; + if (si == waste_tmv_nroles && si < WASTE_TMV_ROLES) { + waste_tmv_role *r = &waste_tmv_roles[si]; + memset(r, 0, sizeof *r); + snprintf(r->role, sizeof r->role, "%s", role); + r->out = out; r->in = in; r->bits = t->q ? t->bits : 32; + waste_tmv_nroles++; + } + if (si < WASTE_TMV_ROLES) tw->prof_slot = si + 1; + } + if (si >= 0 && si < WASTE_TMV_ROLES) { + waste_tmv_role *r = &waste_tmv_roles[si]; + r->calls++; + r->bytes += t->q ? nb : (uint64_t)out * (uint64_t)in * sizeof(float); + r->t += dt; + r->tq += dtq; + } + } +} + static void matvec_t(waste_model *m, float *y, const waste_tensor *t, const float *x, int out, int in) { if (!prof_on) { matvec_t_inner(m, y, t, x, out, in); return; } + tmv_quant_dt = 0; const double t0 = pnow(); matvec_t_inner(m, y, t, x, out, in); const double dt = pnow() - t0; - const uint64_t nb = t ? (uint64_t)out * t->rowbytes : 0; - const int bk = nb < (1u<<20) ? 0 : nb < (8u<<20) ? 1 : nb < (32u<<20) ? 2 : 3; pthread_mutex_lock(&prof_mu); - waste_prof[P_TMV] += dt; waste_prof_n[P_TMV]++; - waste_tmv_bytes += nb; - waste_tmv_t[bk] += dt; waste_tmv_b[bk] += nb; waste_tmv_c[bk]++; + tmv_account(t, out, in, dt, tmv_quant_dt); pthread_mutex_unlock(&prof_mu); } +/* Rows per matvec chunk. A floor of 64 rows was right for wide calls and + * starved narrow ones, because the pool rounds a chunk up to a multiple of + * its floor: on eight threads Qwen's 320-row HyperConnection down + * projection split into five chunks, the shared expert's 640 rows into + * five, and GDN's 48-row in_proj_a never left the calling thread. Sized by + * bytes instead, about 16 KB of weights a chunk, and kept a power of two + * of at least two — so every chunk but the last has an even row count and + * i8mm's two-row tiles fall on exactly the rows they did before. + * + * Below 256 KB of weights the whole call stays on the calling thread, as + * those 48 rows always had: split up, in_proj_a measured 2.7 GB/s against + * 7.5 left alone, the dispatch costing more than the work it shared. */ +static int mv_chunk(int out, size_t rowbytes) +{ + if ((size_t)out * rowbytes < ((size_t)256 << 10)) return out > 0 ? out : 1; + const size_t per = rowbytes ? ((size_t)16 << 10) / rowbytes : 64; + int c = 2; + while (c < 64 && (size_t)c * 2 <= per) c *= 2; + return c; +} + static void matvec_t_inner(waste_model *m, float *y, const waste_tensor *t, const float *x, int out, int in) { if (!t || (!t->q && !t->data)) { memset(y, 0, (size_t)out * sizeof(float)); return; } if (!t->q) { matvec(y, t->data, x, out, in); return; } const int g = t->group, ng = (in + g - 1) / g; + const int mc = mv_chunk(out, t->rowbytes); if (trunk_kern != TK_F32 && t->bits == 4 && (g & 31) == 0) { mvq4_arg a = { y, (const uint8_t *)t->q, t->qs, m->xq, m->xs, in, ng, g, sdot4_sg, g / sdot4_sg, t->rowbytes }; waste_range_fn fn = NULL; + const double tq0 = prof_on ? pnow() : 0; if (trunk_kern == TK_SDOT && g % sdot4_sg == 0) { quant_act4(x, in, g, sdot4_sg, m->xq, m->xs); fn = mvq4_rows_sdot; @@ -823,8 +960,9 @@ static void matvec_t_inner(waste_model *m, float *y, const waste_tensor *t, quant_act4_16(x, in, g, m->xq, m->xs); fn = mvq4_rows_smlal; } + if (prof_on) tmv_quant_dt = pnow() - tq0; if (fn) { - waste_parallel_for_work(out, 64, fn, &a, + waste_parallel_for_work(out, mc, fn, &a, (size_t)out * t->rowbytes); if (trunk_check) { float *ref = (float *)malloc((size_t)out * sizeof(float)); @@ -852,15 +990,144 @@ static void matvec_t_inner(waste_model *m, float *y, const waste_tensor *t, if (sdot_on && t->bits == 8) { quant_act(x, in, g, m->xq, m->xs); mvq_arg a = { y, t->q, t->qs, m->xq, m->xs, in, ng, g, 8, (size_t)ng * g }; - waste_parallel_for_work(out, 64, mvq_rows, &a, + waste_parallel_for_work(out, mc, mvq_rows, &a, (size_t)out * t->rowbytes); } else { mvq_arg a = { y, t->q, t->qs, NULL, x, in, ng, g, t->bits, t->rowbytes }; - run_rows(out, 64, waste_k.mvq_rows_f32, &a, + run_rows(out, mc, waste_k.mvq_rows_f32, &a, (size_t)out * t->rowbytes); } } +/* Several projections of the same vector, as one. + * + * i8mm's activation planes are a function of the vector and the group size + * alone, so projections that read the same input can share them: quantized + * once, and every tensor's rows cut at the multiples mv_chunk would have cut + * them at — even boundaries, so i8mm's two-row tiles fall where they did — + * and handed to the pool as one job. Each row is the same kernel call on the + * same bytes as matvec_t's, so the outputs are its bit for bit. Qwen reads + * the same vector three and four times over: GDN's four input projections, + * QSA's four, the router beside the shared expert's gate and up — and each + * of those was a quantization and a dispatch of its own, half of them too + * small to wake the pool for (LEARNED §85). + * + * Anything the shared planes do not fit — another kernel, another group, a + * float tensor, the trunk check — goes through matvec_t as before. */ +typedef struct { float *y; const waste_tensor *t; int out; } mvb_item; + +#if defined(__ARM_NEON) || defined(__aarch64__) +enum { MVB_MAX = 8 }; +typedef struct { + mvq4_arg a[MVB_MAX]; + int base[MVB_MAX + 1], mc[MVB_MAX], out[MVB_MAX], n; +} mvb_arg; + +static void mvb_pieces(int b, int e, void *p) +{ + const mvb_arg *a = (const mvb_arg *)p; + int i = 0; + for (int k = b; k < e; k++) { + while (k >= a->base[i + 1]) i++; + const int r0 = (k - a->base[i]) * a->mc[i]; + const int r1 = r0 + a->mc[i] < a->out[i] ? r0 + a->mc[i] : a->out[i]; + waste_mvq4_rows_i8mm(r0, r1, (void *)&a->a[i]); + } +} +#endif + +static void matvec_t_batch(waste_model *m, const float *x, int in, + const mvb_item *it, int n) +{ +#if defined(__ARM_NEON) || defined(__aarch64__) + const int shared = trunk_kern == TK_I8MM && !trunk_check && n > 1 && + n <= MVB_MAX && it[0].t && it[0].t->q; + const int g = shared ? it[0].t->group : 0; + mvb_arg a; + a.n = 0; + a.base[0] = 0; + int left[MVB_MAX], nleft = 0; + size_t bytes = 0; + for (int i = 0; i < n; i++) { + const waste_tensor *t = it[i].t; + if (shared && t && t->q && t->bits == 4 && t->group == g && (g & 31) == 0) { + const int k = a.n++; + a.a[k] = (mvq4_arg){ it[i].y, (const uint8_t *)t->q, t->qs, m->xq, m->xs, + in, (in + g - 1) / g, g, sdot4_sg, g / sdot4_sg, + t->rowbytes }; + a.mc[k] = mv_chunk(it[i].out, t->rowbytes); + a.out[k] = it[i].out; + a.base[k + 1] = a.base[k] + (it[i].out + a.mc[k] - 1) / a.mc[k]; + bytes += (size_t)it[i].out * t->rowbytes; + } else { + left[nleft++] = i; + } + } + if (a.n >= 2) { + const double t0 = prof_on ? pnow() : 0; + quant_act4_mm(x, in, g, m->xq, m->xs); + const double tq = prof_on ? pnow() - t0 : 0; + waste_parallel_for_work(a.base[a.n], 1, mvb_pieces, &a, bytes); + if (prof_on) { + const double dt = pnow() - t0; + pthread_mutex_lock(&prof_mu); + for (int i = 0; i < n; i++) { + const waste_tensor *t = it[i].t; + if (!(t && t->q && t->bits == 4 && t->group == g && (g & 31) == 0)) continue; + const double share = bytes ? (double)it[i].out * t->rowbytes / bytes : 0; + tmv_account(t, it[i].out, in, dt * share, tq * share); + } + pthread_mutex_unlock(&prof_mu); + } + for (int j = 0; j < nleft; j++) + matvec_t(m, it[left[j]].y, it[left[j]].t, x, it[left[j]].out, in); + return; + } +#endif + for (int i = 0; i < n; i++) matvec_t(m, it[i].y, it[i].t, x, it[i].out, in); +} + +/* The same kernel over planes the caller has already filled. + * + * i8mm quantizes each weight group of the input on its own, so a caller + * that computes the input a piece at a time — a HyperConnection stream, a + * GDN value head — can quantize each piece's groups in the task that wrote + * it, instead of leaving a serial stretch and a second dispatch in front of + * the projection. `prequant_ok` says whether `t` reads planes laid out that + * way: i8mm, four bits, a group the kernel takes, and `span` — the size of + * the caller's pieces — a whole number of groups. */ +static int prequant_ok(const waste_tensor *t, int span) +{ +#if defined(__ARM_NEON) || defined(__aarch64__) + return t && t->q && trunk_kern == TK_I8MM && !trunk_check && t->bits == 4 && + t->group > 0 && (t->group & 31) == 0 && span % t->group == 0; +#else + (void)t; (void)span; + return 0; +#endif +} + +static void matvec_t_prequant(waste_model *m, float *y, const waste_tensor *t, + int out, int in) +{ +#if defined(__ARM_NEON) || defined(__aarch64__) + const double t0 = prof_on ? pnow() : 0; + const int g = t->group; + mvq4_arg a = { y, (const uint8_t *)t->q, t->qs, m->xq, m->xs, + in, (in + g - 1) / g, g, sdot4_sg, g / sdot4_sg, t->rowbytes }; + waste_parallel_for_work(out, mv_chunk(out, t->rowbytes), waste_mvq4_rows_i8mm, &a, + (size_t)out * t->rowbytes); + if (prof_on) { + const double dt = pnow() - t0; + pthread_mutex_lock(&prof_mu); + tmv_account(t, out, in, dt, 0.0); + pthread_mutex_unlock(&prof_mu); + } +#else + (void)m; (void)y; (void)t; (void)out; (void)in; +#endif +} + /* Dequantize one row of a trunk tensor into dst[cols]. * * matvec_t fuses this with the dot product, which is right when every row @@ -1115,7 +1382,8 @@ static int load_trunk(waste_model *m, const char *dir, const js_doc *d, int trun * row is read per token. Keeping 1.11 GB resident to touch 7 KB * of it is a bad trade against the expert cache, so leave it on * disk and pread the row. */ - if (strstr(t->name, "embed_tokens.weight")) { + if (strstr(t->name, "embed_tokens.weight") || + strstr(t->name, "ngram_head.")) { t->on_disk = 1; t->file_off = off; t->file_scale_off = soff; @@ -1206,12 +1474,87 @@ static int bad_tensor(const char *name) do { const char *rn_ = (name); if (!tensor_data_ok(m, rn_, (n))) \ return bad_tensor(rn_); } while (0) +static int validate_qwen_tensors(waste_model *m) +{ + const waste_config *c = &m->cfg; + const int hid = c->hidden, hc = c->hc_count, lr = c->hc_lowrank; + const int H = hc * hid; + REQUIRE_MATRIX(tname("%smodel.embed_tokens.weight", c->prefix), c->vocab, hid); + REQUIRE_MATRIX(tname("%slm_head.weight", c->prefix), c->vocab, hid); + REQUIRE_VECTOR(tname("%smodel.hyper_connection_mixer.hc_norm.weight", c->prefix), H); + REQUIRE_MATRIX(tname("%smodel.hyper_connection_mixer.input_mix_weight_down.weight", c->prefix), lr, H); + REQUIRE_MATRIX(tname("%smodel.hyper_connection_mixer.input_mix_weight_up.weight", c->prefix), H, lr); + + const int Hk = c->gdn_k_heads, Hv = c->gdn_v_heads; + const int Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int qkv = 2 * Hk * Dk + Hv * Dv; + const int qd = c->n_heads * c->qsa_head_dim; + const int kvd = c->qsa_n_kv * c->qsa_head_dim; + const int idxd = (c->idx_n_heads + c->idx_kv_heads) * c->idx_head_dim; + const int shared = c->shared_inter ? c->shared_inter : c->moe_inter; + + for (int L = 0; L < c->n_layers; L++) { + const char *side[2] = { "attn_hyper_connection", "mlp_hyper_connection" }; + for (int s = 0; s < 2; s++) { + REQUIRE_VECTOR(tname("%smodel.layers.%d.%s.hc_norm.weight", c->prefix, L, side[s]), H); + REQUIRE_MATRIX(tname("%smodel.layers.%d.%s.block_inject_weight.weight", c->prefix, L, side[s]), hc, H); + REQUIRE_MATRIX(tname("%smodel.layers.%d.%s.input_mix_weight_down.weight", c->prefix, L, side[s]), lr, H); + REQUIRE_MATRIX(tname("%smodel.layers.%d.%s.input_mix_weight_up.weight", c->prefix, L, side[s]), H, lr); + } + if (!c->qwen_full[L]) { + REQUIRE_DATA(tname("%smodel.layers.%d.linear_attn.A_log", c->prefix, L), Hv); + REQUIRE_DATA(tname("%smodel.layers.%d.linear_attn.dt_bias", c->prefix, L), Hv); + REQUIRE_DATA(tname("%smodel.layers.%d.linear_attn.conv1d.weight", c->prefix, L), + (size_t)qkv * c->conv_k); + REQUIRE_MATRIX(tname("%smodel.layers.%d.linear_attn.in_proj_qkv.weight", c->prefix, L), qkv, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.linear_attn.in_proj_z.weight", c->prefix, L), Hv * Dv, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.linear_attn.in_proj_a.weight", c->prefix, L), Hv, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.linear_attn.in_proj_b.weight", c->prefix, L), Hv, hid); + REQUIRE_VECTOR(tname("%smodel.layers.%d.linear_attn.norm.weight", c->prefix, L), Dv); + REQUIRE_MATRIX(tname("%smodel.layers.%d.linear_attn.out_proj.weight", c->prefix, L), hid, Hv * Dv); + } else { + REQUIRE_MATRIX(tname("%smodel.layers.%d.self_attn.q_proj.weight", c->prefix, L), qd * 2, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.self_attn.k_proj.weight", c->prefix, L), kvd, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.self_attn.v_proj.weight", c->prefix, L), kvd, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.self_attn.o_proj.weight", c->prefix, L), hid, qd); + REQUIRE_VECTOR(tname("%smodel.layers.%d.self_attn.q_norm.weight", c->prefix, L), c->qsa_head_dim); + REQUIRE_VECTOR(tname("%smodel.layers.%d.self_attn.k_norm.weight", c->prefix, L), c->qsa_head_dim); + REQUIRE_MATRIX(tname("%smodel.layers.%d.self_attn.indexer.index_qk_proj.weight", c->prefix, L), idxd, hid); + REQUIRE_VECTOR(tname("%smodel.layers.%d.self_attn.indexer.q_layernorm.weight", c->prefix, L), c->idx_head_dim); + REQUIRE_VECTOR(tname("%smodel.layers.%d.self_attn.indexer.k_layernorm.weight", c->prefix, L), c->idx_head_dim); + } + REQUIRE_MATRIX(tname("%smodel.layers.%d.mlp.gate.weight", c->prefix, L), c->n_experts, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.mlp.shared_expert.gate_proj.weight", c->prefix, L), shared, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.mlp.shared_expert.up_proj.weight", c->prefix, L), shared, hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.mlp.shared_expert.down_proj.weight", c->prefix, L), hid, shared); + REQUIRE_MATRIX(tname("%smodel.layers.%d.mlp.shared_expert_gate.weight", c->prefix, L), 1, hid); + if (L == c->ple_layer) { + REQUIRE_MATRIX(tname("%smodel.layers.%d.ple.key_proj.weight", c->prefix, L), H, c->ple_embed ? c->ple_embed : hid); + REQUIRE_MATRIX(tname("%smodel.layers.%d.ple.value_proj.weight", c->prefix, L), hid, c->ple_embed ? c->ple_embed : hid); + REQUIRE_VECTOR(tname("%smodel.layers.%d.ple.norm_key.weight", c->prefix, L), H); + REQUIRE_VECTOR(tname("%smodel.layers.%d.ple.norm_query.weight", c->prefix, L), H); + REQUIRE_VECTOR(tname("%smodel.layers.%d.ple.norm_conv.weight", c->prefix, L), H); + REQUIRE_DATA(tname("%smodel.layers.%d.ple.conv1d.weight", c->prefix, L), + (size_t)H * c->ple_conv_k); + for (int h = 0; h < WASTE_QWEN_PLE_HEADS; h++) { + const int rows = c->ple_sz[h] > 0 ? (int)c->ple_sz[h] : 1; + const int width = (c->ple_embed && c->heads_per_ngram) + ? c->ple_embed / ((c->ngram_size - 1) * c->heads_per_ngram) : 8; + REQUIRE_MATRIX(tname("%smodel.layers.%d.ple.ple_embedding.ngram_head.%d.weight", + c->prefix, L, h), rows, width); + } + } + } + return 1; +} + /* Validate every tensor shape the text forward pass indexes. Kernel calls * receive dimensions from config rather than from the tensor, so merely * checking that a name exists is not enough: a shorter, correctly named * tensor is an out-of-bounds read. */ static int validate_text_tensors(waste_model *m) { + if (m->cfg.arch_qwen) return validate_qwen_tensors(m); const waste_config *c = &m->cfg; const int hid = c->hidden; REQUIRE_MATRIX(tname("%smodel.embed_tokens.weight", c->prefix), c->vocab, hid); @@ -1484,6 +1827,13 @@ static int cfg_sane(const waste_config *c) if (c->kda_heads < 0 || c->kda_heads > (1 << 16)) return 0; if (c->kda_dim < 0 || c->kda_dim > (1 << 16)) return 0; if (c->conv_k < 0 || c->conv_k > 64) return 0; + /* A PLE conv kernel of 0 allocates a zero-length ring while the + * weight check and the step walk as if it were 4, so a container + * the loader should have refused reads past the allocation. Only + * a container that declares a PLE layer ever reaches the ring; + * the field stays 0 elsewhere and is meaningless there. */ + if (c->ple_layer >= 0 && (c->ple_conv_k < 1 || c->ple_conv_k > 64)) + return 0; if (c->kv_lora < 0 || c->kv_lora > (1 << 20) || c->q_lora < 0 || c->q_lora > (1 << 20)) return 0; if (c->qk_nope < 0 || c->qk_nope > (1 << 20) || @@ -1499,7 +1849,7 @@ static int cfg_sane(const waste_config *c) if (c->kda_heads < 1 || c->kda_dim < 1 || c->conv_k < 1) return 0; if ((int64_t)c->kda_heads * c->kda_dim > INT_MAX) return 0; } - if (n_kda < c->n_layers && !c->ds41) { + if (!c->arch_qwen && !c->ds41 && n_kda < c->n_layers) { const int64_t qd = (int64_t)c->qk_nope + c->qk_rope; if (c->kv_lora < 1 || qd < 1 || c->v_head < 1) return 0; if ((int64_t)c->n_heads * qd > INT_MAX || @@ -1534,6 +1884,40 @@ static int cfg_sane(const waste_config *c) } else if (c->index_heads || c->index_dim) { return 0; } + /* Qwen states its shapes in its own keys, and every one of them sizes + * an allocation or indexes a loop below. A container that omits one is + * refused here rather than opened and read out of bounds. */ + if (c->arch_qwen) { + if (c->qwen_n_layer_types != c->n_layers) return 0; + if (c->gdn_k_heads < 1 || c->gdn_v_heads < 1 || + c->gdn_k_dim < 1 || c->gdn_v_dim < 1) return 0; + if (c->gdn_v_heads % c->gdn_k_heads != 0) return 0; + if (c->hc_count < 1 || c->hc_count > 16 || + c->hc_lowrank < 1 || c->hc_lowrank > (1 << 16)) return 0; + if (c->qsa_head_dim < 1 || c->qsa_n_kv < 1) return 0; + if (c->idx_n_heads < 1 || c->idx_head_dim < 1 || + c->idx_compress < 1 || c->idx_budget < 1) return 0; + if (c->idx_budget % c->idx_compress != 0) return 0; + if (c->idx_kv_heads != 1) return 0; + if (c->n_heads % c->qsa_n_kv != 0) return 0; + if (c->ngram_size < 1 || c->ngram_size > 8) return 0; + if (c->rotary_dim < 0 || c->rotary_dim > 256) return 0; + if (c->rotary_dim / 2 > WASTE_MAX_ROPE_HALF) return 0; + if ((int64_t)c->hc_count * c->hidden > INT_MAX) return 0; + if ((int64_t)c->gdn_v_heads * c->gdn_k_dim * c->gdn_v_dim > INT_MAX) + return 0; + /* The QSA query buffers are n_heads * head_dim and the indexer's + * work is idx_n_heads * idx_head_dim; both are computed as int + * before they reach a size_t, so bound the products, not just the + * factors. */ + if ((int64_t)c->n_heads * c->qsa_head_dim > INT_MAX / 4) return 0; + if ((int64_t)c->idx_n_heads * c->idx_head_dim > INT_MAX / 4) return 0; + if ((int64_t)c->idx_budget + c->idx_compress > INT_MAX / 4) return 0; + if (c->ple_layer >= 0) { + for (int h = 0; h < WASTE_QWEN_PLE_HEADS; h++) + if (c->ple_sz[h] <= 0) return 0; + } + } if (c->n_experts && c->moe_inter < 1) return 0; if ((!c->n_experts || c->first_dense) && c->dense_inter < 1) return 0; if ((int64_t)c->moe_inter * (c->n_shared ? c->n_shared : 1) > INT_MAX) @@ -1793,6 +2177,8 @@ static void cfg_from_json(waste_config *c, const js_doc *d, int cfg) c->hidden = (int)js_int(d, js_get(d, cfg, "hidden_size"), 0); c->n_experts = (int)js_int(d, js_get(d, cfg, "num_experts"), 0); c->top_k = (int)js_int(d, js_get(d, cfg, "num_experts_per_token"), 0); + if (!c->top_k) + c->top_k = (int)js_int(d, js_get(d, cfg, "num_experts_per_tok"), 0); c->moe_inter = (int)js_int(d, js_get(d, cfg, "moe_intermediate_size"), 0); c->dense_inter = (int)js_int(d, js_get(d, cfg, "intermediate_size"), 0); c->n_shared = (int)js_int(d, js_get(d, cfg, "num_shared_experts"), 0); @@ -1859,6 +2245,7 @@ static void cfg_from_json(waste_config *c, const js_doc *d, int cfg) c->index_dim = (int)js_int(d, js_get(d, cfg, "index_head_dim"), 0); c->index_tail = js_get(d, cfg, "index_kpool_always_select_tail") >= 0; c->tok_han_split = js_bool(d, js_get(d, cfg, "tokenizer_han_split"), 1); + c->tok_digit_run = (int)js_int(d, js_get(d, cfg, "tokenizer_digit_run"), 3); c->tok_pattern = (int)js_int(d, js_get(d, cfg, "tokenizer_pattern"), 0); int lac = js_get(d, cfg, "linear_attn_config"); @@ -1873,6 +2260,96 @@ static void cfg_from_json(waste_config *c, const js_doc *d, int cfg) int v = (int)js_int(d, js_at(d, kl, i), -1) - 1; /* list is 1-based */ if (v >= 0 && v < 128) c->kda_layer[v] = 1; } + + c->arch_qwen = 0; + c->ple_layer = -1; + { + char mt[40]; + js_str(d, js_get(d, cfg, "model_type"), mt, sizeof mt); + if (strcmp(mt, "qwen4_exp_text") == 0 || + strstr(c->arch, "Qwen4Exp") != NULL) + c->arch_qwen = 1; + } + if (!c->arch_qwen) return; + + /* Qwen is not Kimi: do not fill kda_layer from a missing linear_attn_config. */ + memset(c->kda_layer, 0, sizeof c->kda_layer); + c->qsa_n_kv = (int)js_int(d, js_get(d, cfg, "num_key_value_heads"), 0); + c->qsa_head_dim = (int)js_int(d, js_get(d, cfg, "head_dim"), 0); + c->gdn_k_heads = (int)js_int(d, js_get(d, cfg, "linear_num_key_heads"), 0); + c->gdn_v_heads = (int)js_int(d, js_get(d, cfg, "linear_num_value_heads"), 0); + c->gdn_k_dim = (int)js_int(d, js_get(d, cfg, "linear_key_head_dim"), 0); + c->gdn_v_dim = (int)js_int(d, js_get(d, cfg, "linear_value_head_dim"), 0); + c->conv_k = (int)js_int(d, js_get(d, cfg, "linear_conv_kernel_dim"), 4); + c->hc_count = (int)js_int(d, js_get(d, cfg, "hc_count"), 0); + c->hc_lowrank = (int)js_int(d, js_get(d, cfg, "hc_lowrank"), 0); + c->idx_n_heads = (int)js_int(d, js_get(d, cfg, "indexer_n_heads"), 4); + c->idx_kv_heads = (int)js_int(d, js_get(d, cfg, "indexer_kv_heads"), 1); + c->idx_head_dim = (int)js_int(d, js_get(d, cfg, "indexer_head_dim"), 128); + c->idx_budget = (int)js_int(d, js_get(d, cfg, "indexer_budget"), 2048); + c->idx_compress = (int)js_int(d, js_get(d, cfg, "indexer_compress_ratio"), 4); + c->ngram_size = (int)js_int(d, js_get(d, cfg, "ngram_size"), 3); + c->heads_per_ngram = (int)js_int(d, js_get(d, cfg, "heads_per_ngram"), 8); + c->ple_embed = (int)js_int(d, js_get(d, cfg, "ple_embed_dim"), 0); + c->ple_conv_k = (int)js_int(d, js_get(d, cfg, "ple_conv_kernel_size"), 4); + c->shared_inter = (int)js_int(d, js_get(d, cfg, "shared_expert_intermediate_size"), + c->moe_inter); + /* Not read from the config: Qwen4ExpTextTopKRouter renormalizes + * unconditionally, so a container that happens to omit the key must + * still renormalize. */ + c->renorm = 1; + { + const int lt = js_get(d, cfg, "layer_types"); + memset(c->qwen_full, 0, sizeof c->qwen_full); + /* Kept so cfg_sane can insist on one entry per layer. All-zero is + * a valid-looking answer that means "every layer is GDN", and a + * container whose `layer_types` is missing or short would attend + * with a recurrence on layers that need sparse attention — wrong + * everywhere and diagnosable nowhere. */ + c->qwen_n_layer_types = js_size(d, lt); + for (int i = 0; i < js_size(d, lt) && i < WASTE_MAX_LAYERS; i++) { + char kind[32]; + js_str(d, js_at(d, lt, i), kind, sizeof kind); + c->qwen_full[i] = (strcmp(kind, "full_attention") == 0); + } + } + { + const int ids = js_get(d, cfg, "ple_layer_ids"); + if (js_size(d, ids) > 0) { + const int one = (int)js_int(d, js_at(d, ids, 0), 0); + c->ple_layer = one > 0 ? one - 1 : -1; + } + } + { + const int off = js_get(d, cfg, "ple_head_offsets"); + const int sz = js_get(d, cfg, "ple_head_vocab_sizes"); + const int mul = js_get(d, cfg, "ple_layer_multipliers"); + for (int h = 0; h < WASTE_QWEN_PLE_HEADS; h++) { + c->ple_off[h] = js_int(d, js_at(d, off, h), 0); + c->ple_sz[h] = js_int(d, js_at(d, sz, h), 0); + } + for (int i = 0; i < 8; i++) + c->ple_mult[i] = js_int(d, js_at(d, mul, i), 0); + } + { + const int rp = js_get(d, cfg, "rope_parameters"); + const double pf = js_num(d, js_get(d, rp, "partial_rotary_factor"), + js_num(d, js_get(d, cfg, "partial_rotary_factor"), 0.25)); + const int hd = c->qsa_head_dim ? c->qsa_head_dim : 256; + c->rotary_dim = (int)(hd * pf); + const int sec = js_get(d, rp, "mrope_section"); + c->mrope_section[0] = (int)js_int(d, js_at(d, sec, 0), 11); + c->mrope_section[1] = (int)js_int(d, js_at(d, sec, 1), 11); + c->mrope_section[2] = (int)js_int(d, js_at(d, sec, 2), 10); + const double base = js_num(d, js_get(d, rp, "rope_theta"), + js_num(d, js_get(d, cfg, "rope_theta"), 10000000.0)); + const int half = c->rotary_dim / 2; + if (half > 0 && half <= WASTE_MAX_ROPE_HALF) { + for (int j = 0; j < half; j++) + c->rope_inv_freq[j] = (float)(1.0 / pow(base, (double)(2 * j) / c->rotary_dim)); + c->rope_err[0] = 0; + } + } } /* Defined below, next to record_check; the cache needs it at load. */ @@ -2176,6 +2653,9 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, snprintf(m->cfg.prefix, sizeof m->cfg.prefix, "language_model."); } cfg_from_json(&m->cfg, &d, cfg); + /* Qwen containers are accepted once the kernels and planner exist. + * Kimi still never sees Qwen tensors: arch_qwen selects a distinct + * forward, not KDA/MLA/AttnRes. */ } if (!cfg_sane(&m->cfg)) { fprintf(stderr, "waste: manifest config is out of range " @@ -2185,6 +2665,16 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, js_free(&d); free(src); return -2; /* -> WASTE_E_FORMAT */ } + /* Qwen's 4-bit trunk goes through i8mm unless WASTE_TRUNK_KERNEL says + * otherwise. It is not the exact arithmetic, and the error is not the + * K3 kind that a recurrence carries forward: against f32 over 5,918 + * tokens of real text, perplexity 3.712 against 3.698, no growth past + * QSA's 2,048-token selection budget, for 7.57 -> 9.59 tok/s + * (LEARNED §77). The kernel is one setting for the whole process, so a + * process that loads Qwen and then another architecture keeps i8mm for + * both; the variable pins it either way. */ + if (m->cfg.arch_qwen && !trunk_kern_env) + waste_model_set_sdot4(TK_I8MM, sdot4_sg); /* rope_init leaves no table for a shape it does not implement. Running * anyway would apply no rotation, which is not a degraded result but an * unordered one, so refuse instead. */ @@ -2606,22 +3096,81 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, * closes: kv_cap tokens make exactly kv_cap/kpool of them, and rounding * up costs one vector and removes a bound to get wrong. */ m->pool_cap = c->index_kpool ? kv_cap / c->index_kpool + 1 : 0; - if (c->ds41 && csa2_alloc(m, kv_cap) < 0) return -1; - for (int L = 0; L < c->n_layers; L++) { - if (c->ds41) continue; /* csa2_alloc did this */ - if (c->kda_layer[L]) { - m->S[L] = (float *)calloc((size_t)H * D * D, sizeof(float)); - m->conv[L] = (float *)calloc((size_t)3 * C * (c->conv_k - 1), sizeof(float)); - } else { - m->has_mla = 1; /* this is what makes kv_cap a real bound */ - m->latcache[L] = (float *)calloc( - (size_t)kv_cap * (c->kv_lora + c->qk_rope), sizeof(float)); - if (c->index_topk) { - m->idxpool[L] = (float *)calloc( - (size_t)m->pool_cap * c->index_dim, sizeof(float)); - m->idxbuf[L] = (float *)calloc( - (size_t)c->index_kpool * 2 * c->index_dim, sizeof(float)); - if (!m->idxpool[L] || !m->idxbuf[L]) return -1; + if (c->arch_qwen) { + const int Hv = c->gdn_v_heads, Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int Hk = c->gdn_k_heads; + const int qkv = 2 * Hk * Dk + Hv * Dv; + const int nkv = c->qsa_n_kv, hd = c->qsa_head_dim; + const int idim = c->idx_head_dim; + const int compress = c->idx_compress > 0 ? c->idx_compress : 4; + const int nblk = compress > 0 ? (kv_cap + compress - 1) / compress : 0; + const int max_sel = c->idx_budget + compress; + const int rot = c->rotary_dim > 0 ? c->rotary_dim : 1; + for (int L = 0; L < c->n_layers; L++) { + if (!c->qwen_full[L]) { + m->S[L] = (float *)calloc((size_t)Hv * Dk * Dv, sizeof(float)); + m->conv[L] = (float *)calloc((size_t)qkv * (c->conv_k > 0 ? c->conv_k - 1 : 0), + sizeof(float)); + } else { + m->has_qsa = 1; + m->qsa_k[L] = (uint16_t *)calloc((size_t)kv_cap * nkv * hd, 2); + m->qsa_v[L] = (uint16_t *)calloc((size_t)kv_cap * nkv * hd, 2); + m->qsa_rawk[L] = (float *)calloc((size_t)kv_cap * idim, sizeof(float)); + } + } + m->hcx = (float *)calloc((size_t)c->hc_count * c->hidden, sizeof(float)); + { + const int R = (c->ple_conv_k > 1 && c->ngram_size > 0) + ? (c->ple_conv_k - 1) * c->ngram_size : 0; + m->ple_ring = (float *)calloc((size_t)c->hc_count * c->hidden * (R > 0 ? R : 1), + sizeof(float)); + } + { + const int pe = c->ple_embed ? c->ple_embed : c->hidden; + m->ple_emb = (float *)calloc((size_t)(pe > 0 ? pe : 1), sizeof(float)); + } + m->gdn_g = (float *)calloc((size_t)(Hv > 0 ? Hv : 1), sizeof(float)); + { + const int qd = c->n_heads * hd; + m->qsa_q = (float *)calloc((size_t)(qd > 0 ? qd : 1), sizeof(float)); + m->qsa_gate = (float *)calloc((size_t)(qd > 0 ? qd : 1), sizeof(float)); + m->qsa_attn = (float *)calloc((size_t)(qd > 0 ? qd : 1), sizeof(float)); + const size_t compact = (size_t)(max_sel > 0 ? max_sel : 1) * (size_t)nkv * hd; + m->qsa_kf = (float *)calloc(compact > 0 ? compact : 1, sizeof(float)); + m->qsa_vf = (float *)calloc(compact > 0 ? compact : 1, sizeof(float)); + /* One row of scores per query head, so the heads can attend at + * once (qwen_qsa_layer). src/waste.c plans the same size. */ + m->qsa_scr = (float *)calloc((size_t)(max_sel > 0 ? max_sel : 1) * + (size_t)(c->n_heads > 0 ? c->n_heads : 1), + sizeof(float)); + m->qsa_sel = (int *)calloc((size_t)(max_sel > 0 ? max_sel : 1), sizeof(int)); + const size_t work = (size_t)nblk * idim + (size_t)nblk + (size_t)idim; + m->qsa_work = (float *)calloc(work > 0 ? work : 1, sizeof(float)); + m->qsa_taken = (int *)calloc((size_t)(nblk > 0 ? nblk : 1), sizeof(int)); + m->qsa_cs = (float *)calloc((size_t)2 * kv_cap * rot, sizeof(float)); + } + m->moe_prob = (float *)calloc((size_t)(c->n_experts > 0 ? c->n_experts : 1), + sizeof(float)); + m->moe_used = (uint8_t *)calloc((size_t)(c->n_experts > 0 ? c->n_experts : 1), 1); + for (int i = 0; i < 8; i++) m->ple_prev[i] = c->eos_token_id; + } else if (c->ds41) { + if (csa2_alloc(m, kv_cap) < 0) return -1; + } else { + for (int L = 0; L < c->n_layers; L++) { + if (c->kda_layer[L]) { + m->S[L] = (float *)calloc((size_t)H * D * D, sizeof(float)); + m->conv[L] = (float *)calloc((size_t)3 * C * (c->conv_k - 1), sizeof(float)); + } else { + m->has_mla = 1; /* this is what makes kv_cap a real bound */ + m->latcache[L] = (float *)calloc( + (size_t)kv_cap * (c->kv_lora + c->qk_rope), sizeof(float)); + if (c->index_topk) { + m->idxpool[L] = (float *)calloc( + (size_t)m->pool_cap * c->index_dim, sizeof(float)); + m->idxbuf[L] = (float *)calloc( + (size_t)c->index_kpool * 2 * c->index_dim, sizeof(float)); + if (!m->idxpool[L] || !m->idxbuf[L]) return -1; + } } } } @@ -2636,7 +3185,16 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, c->index_heads, sizeof(float)); if (!m->idxsel || !m->idxscore || !m->idxrank || !m->idxq) return -1; } - const int big = c->hidden > C ? c->hidden : C; + int big = c->hidden > C ? c->hidden : C; + if (c->arch_qwen) { + const int qkv = 2 * c->gdn_k_heads * c->gdn_k_dim + + c->gdn_v_heads * c->gdn_v_dim; + const int hcH = c->hc_count * c->hidden; + const int qsa = c->n_heads * c->qsa_head_dim * 2; + if (qkv > big) big = qkv; + if (hcH > big) big = hcH; + if (qsa > big) big = qsa; + } /* mHC keeps hc_mult residual streams instead of one. Every other user * of m->x reads stream 0, which is where the single-stream models put * the whole thing, so the multiplier is confined to this allocation @@ -2654,16 +3212,19 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, } m->h = (float *)calloc((size_t)c->hidden, sizeof(float)); m->tmp = (float *)calloc((size_t)8 * big + 8 * c->moe_inter + 8 * c->dense_inter - + (size_t)4 * c->n_heads * (c->v_head + c->qk_nope + c->qk_rope) + + (size_t)4 * c->n_heads * (c->v_head + c->qk_nope + c->qk_rope + + c->qsa_head_dim) + (size_t)2 * (c->q_lora ? c->q_lora : 1) + 256, sizeof(float)); /* Sized for every user of the buffer, not just the one it is named * after — see WASTE_ATT_ROUTER_OFF in model.h. */ { - size_t need = (size_t)kv_cap * (size_t)c->n_heads; /* MLA scores */ + size_t need = (size_t)kv_cap * (size_t)c->n_heads; /* MLA/QSA scores */ const size_t kda = (size_t)c->kda_heads * (size_t)c->kda_dim; + const size_t gdn = (size_t)c->gdn_v_heads * (size_t)c->gdn_k_dim; const size_t route = WASTE_ATT_ROUTER_OFF + 2u * (size_t)c->n_experts; if (kda > need) need = kda; + if (gdn > need) need = gdn; if (route > need) need = route; m->att = (float *)calloc(need + 1024, sizeof(float)); } @@ -2726,8 +3287,11 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, * matvec in the model: 16384 against the dense FFN's 12288. Sized * from the FFN alone this buffer is 4096 activations short of what * hc_collapse quantizes into it, and the 1024 bytes of slack below - * hide that at test scale and not at model scale. */ - const int64_t hcw = (int64_t)(c->hc_mult ? c->hc_mult : 1) * c->hidden; + * hide that at test scale and not at model scale. Qwen's + * HyperConnection mix reads the same shape under its own key. */ + const int64_t hcw = (int64_t)(c->hc_mult ? c->hc_mult : + c->arch_qwen && c->hc_count ? c->hc_count : 1) + * c->hidden; if (hcw > nmax) nmax = (int)hcw; /* Two bytes per activation: the i8mm path writes two int8 planes * and the SMLAL path writes int16, both over the padded group @@ -2810,13 +3374,28 @@ int waste_model_load(waste_model *m, const char *dir, int kv_cap, return -1; if (m->index_bits == 6 && (!m->lut8 || !m->lut8_scale)) return -1; for (int L = 0; L < c->n_layers; L++) { - /* CSA2 caches a window rather than a latent, and only four of its - * forty layers cache anything more — csa2_alloc has already - * checked its own. */ - if (c->ds41) { if (!m->winkv[L]) return -1; } - else if (c->kda_layer[L]) { if (!m->S[L] || !m->conv[L]) return -1; } - else if (!m->latcache[L]) return -1; + if (c->arch_qwen) { + if (!c->qwen_full[L]) { + if (!m->S[L] || !m->conv[L]) return -1; + } else if (!m->qsa_k[L] || !m->qsa_v[L] || !m->qsa_rawk[L]) { + return -1; + } + } else if (c->ds41) { + /* CSA2 caches a window rather than a latent, and only four of + * its forty layers cache anything more; csa2_alloc checked it. */ + if (!m->winkv[L]) return -1; + } else if (c->kda_layer[L]) { + if (!m->S[L] || !m->conv[L]) return -1; + } else if (!m->latcache[L]) { + return -1; + } } + if (c->arch_qwen && (!m->hcx || !m->ple_ring || !m->ple_emb || !m->gdn_g || + !m->qsa_q || !m->qsa_gate || !m->qsa_attn || + !m->qsa_kf || !m->qsa_vf || !m->qsa_scr || !m->qsa_work || + !m->qsa_cs || !m->qsa_sel || !m->qsa_taken || + !m->moe_prob || !m->moe_used)) + return -1; if (c->attn_res_block && !m->blockres) return -1; /* Last, so a load that fails leaves no thread reading a model nobody * owns — every return above this line is a failure. */ @@ -2856,10 +3435,19 @@ void waste_model_free(waste_model *m) for (int L = 0; L < 128; L++) { free(m->S[L]); free(m->conv[L]); free(m->latcache[L]); free(m->idxpool[L]); free(m->idxbuf[L]); + free(m->qsa_k[L]); free(m->qsa_v[L]); free(m->qsa_rawk[L]); free(m->winkv[L]); free(m->ckvc[L]); free(m->ikey[L]); free(m->cpool[L]); for (int s = 0; s < WASTE_MAX_SHARDS; s++) if (m->bank[L].fd[s] >= 0) close(m->bank[L].fd[s]); } + free(m->hcx); + free(m->ple_ring); + free(m->ple_emb); + free(m->gdn_g); + free(m->qsa_q); free(m->qsa_gate); free(m->qsa_attn); + free(m->qsa_kf); free(m->qsa_vf); free(m->qsa_scr); free(m->qsa_work); + free(m->qsa_cs); free(m->qsa_sel); free(m->qsa_taken); + free(m->moe_prob); free(m->moe_used); free(m->x); free(m->h); free(m->tmp); free(m->att); free(m->logits); free(m->ff); free(m->e_gate); free(m->e_up); free(m->e_down); waste_dio_free(m->lut); free(m->lut8); free(m->lut8_scale); @@ -3171,7 +3759,7 @@ void waste_model_clear_read_error(waste_model *m) * position and is not bounded here. */ int waste_model_ctx_max(const waste_model *m) { - return m->has_mla ? m->kv_cap : 0; + return (m->has_mla || m->has_qsa) ? m->kv_cap : 0; } int waste_model_ctx_full(const waste_model *m) { return m->ctx_full; } @@ -3585,6 +4173,7 @@ typedef struct { const float *lut_gate, *lut_up; const int8_t *q_gate, *q_up; const float *qs_gate, *qs_up; + const int *jmap; /* task t is expert jmap[t]; NULL: j_off+t */ } xpar_arg; static void moe_expert_range(int b, int e, void *p) @@ -3595,7 +4184,7 @@ static void moe_expert_range(int b, int e, void *p) const int inter = a->inter, lat = a->lat; for (int t = b; t < e; t++) { - const int j = a->j_off + t; + const int j = a->jmap ? a->jmap[t] : a->j_off + t; const uint8_t *rec = a->recs[j]; const waste_expert_hdr *h = (const waste_expert_hdr *)rec; const uint16_t *sc = (const uint16_t *)(rec + h->chan_corr_off); @@ -3624,6 +4213,104 @@ static void moe_expert_range(int b, int e, void *p) } } +/* ---- routed experts in row ranges, a stage at a time --------------------- + * + * One task per expert leaves the barrier waiting on whichever thread got two: + * ten equal experts on eight threads is two experts of wall time for ten of + * work, 5x at best, and Qwen's measured 4.4x on eight cores. It is not the + * memory — on one thread the kernel lost no more to six cores of random + * reads over 1 GB than to six cores spinning (LEARNED §84). So the work is + * cut into equal pieces instead, in the three stages an expert's arithmetic + * depends on: + * + * 1. every expert's gate and up rows, VQ_TILE * VQ_SUPER rows a task + * 2. every expert's activation and down table, one task each + * 3. every expert's down rows, in the same ranges + * + * Each piece writes only its own rows, through the kernels vq_apply_serial + * and vq_matvec_serial call, so the result is moe_expert_range's bit for + * bit. VQ3R with float tables only: that is what Qwen's experts use, and a + * VQ4P or WASTE_VQ8 layer keeps the per-expert tasks. */ +enum { XS_ROWS = VQ_TILE * VQ_SUPER }; + +typedef struct { + waste_model *m; + const uint8_t **recs; + const int *list; /* piece t's expert is list[t / per] */ + int inter, hid, n_gu, n_dn; /* row ranges per matrix */ + const float *lut_gate, *lut_up; +} xstage_arg; + +static void xstage_gate_up(int b, int e, void *p) +{ + const xstage_arg *a = (const xstage_arg *)p; + waste_model *m = a->m; + const int inter = a->inter, per = 2 * a->n_gu; + for (int k = b; k < e; k++) { + const int j = a->list[k / per], mat = (k % per) / a->n_gu; + const int r0 = ((k % per) % a->n_gu) * XS_ROWS; + const int r1 = r0 + XS_ROWS < inter ? r0 + XS_ROWS : inter; + const uint8_t *rec = a->recs[j]; + const waste_expert_hdr *h = (const waste_expert_hdr *)rec; + const uint16_t *sc = (const uint16_t *)(rec + h->chan_corr_off); + vq_arg va = { (mat ? m->xub : m->xga) + (size_t)j * inter, + rec + (mat ? h->up_off : h->gate_off), sc + mat * inter, + mat ? a->lut_up : a->lut_gate, + a->hid / m->vec_dim, m->stages, m->cb_entries }; + vq_rows(r0, r1, &va); + } +} + +static void xstage_down_lut(int b, int e, void *p) +{ + const xstage_arg *a = (const xstage_arg *)p; + waste_model *m = a->m; + const int inter = a->inter; + for (int t = b; t < e; t++) { + const int j = a->list[t]; + const waste_expert_hdr *h = (const waste_expert_hdr *)a->recs[j]; + float *ga = m->xga + (size_t)j * inter; + waste_act_pair_range(&m->cfg, ga, m->xub + (size_t)j * inter, inter); + lutb_arg la = { m->xlut + (size_t)j * m->xlut_sz, m->codebooksT, ga, + h->codebook_id + 2 * m->stages, m->stages, + m->cb_entries, m->vec_dim }; + waste_k.lutb_range(0, inter / m->vec_dim, &la); + } +} + +static void xstage_down(int b, int e, void *p) +{ + const xstage_arg *a = (const xstage_arg *)p; + waste_model *m = a->m; + const int inter = a->inter, hid = a->hid; + for (int k = b; k < e; k++) { + const int j = a->list[k / a->n_dn]; + const int r0 = (k % a->n_dn) * XS_ROWS; + const int r1 = r0 + XS_ROWS < hid ? r0 + XS_ROWS : hid; + const uint8_t *rec = a->recs[j]; + const waste_expert_hdr *h = (const waste_expert_hdr *)rec; + const uint16_t *sc = (const uint16_t *)(rec + h->chan_corr_off); + vq_arg va = { m->xacc + (size_t)j * hid, rec + h->down_off, sc + 2 * inter, + m->xlut + (size_t)j * m->xlut_sz, + inter / m->vec_dim, m->stages, m->cb_entries }; + vq_rows(r0, r1, &va); + } +} + +/* n experts, named by list, into their xacc slices. */ +static void experts_staged(waste_model *m, const uint8_t **recs, const int *list, + int n, int inter, int hid, + const float *lut_gate, const float *lut_up) +{ + xstage_arg a = { m, recs, list, inter, hid, + (inter + XS_ROWS - 1) / XS_ROWS, (hid + XS_ROWS - 1) / XS_ROWS, + lut_gate, lut_up }; + const int w = g_pool.nthreads; + waste_parallel_for_each(n * 2 * a.n_gu, xstage_gate_up, &a, w); + waste_parallel_for_each(n, xstage_down_lut, &a, w); + waste_parallel_for_each(n * a.n_dn, xstage_down, &a, w); +} + /* ---- layers ------------------------------------------------------------ */ /* Log-space decay gate, in place over [H][D]. @@ -5134,7 +5821,7 @@ static void moe_layer(waste_model *m, int L, const float *in, float *out, int *r lut_done = 1; } xpar_arg pa = { m, c, recs, w, j0, inter, lat, lut_gate, lut_up, - q_gate, q_up, qs_gate, qs_up }; + q_gate, q_up, qs_gate, qs_up, NULL }; waste_parallel_for(j1 - j0, 1, moe_expert_range, &pa); PROF_END(P_EMM); waste_ecache_release(&m->cache); @@ -5456,6 +6143,30 @@ static void state_fill(const waste_model *m, waste_state_hdr *h, int pos) const waste_config *c = &m->cfg; memset(h, 0, sizeof *h); h->magic = WASTE_MAGIC_KDASTATE; + h->n_layers = c->n_layers; + h->hidden = c->hidden; + h->pos = pos; + if (c->arch_qwen) { + /* Version 2: GDN/QSA/HC/PLE state. The fields keep their Kimi + * names and carry Qwen's shapes, so the struct stays one size and + * every shape this file's length depends on is still compared. + * `hc_mult` and `index_dim` are free to reuse here because Qwen + * has neither mHC nor the DSA indexer — the version guards the + * two readings apart. */ + h->version = 2; + h->kda_heads = c->gdn_v_heads; + h->kda_dim = c->gdn_k_dim; + h->conv_k = c->conv_k; + h->n_heads = c->qsa_n_kv; + h->qk_nope = c->hc_count; + h->qk_rope = c->qsa_head_dim; + h->v_head = c->idx_head_dim; + h->attn_res_block = c->idx_compress; + h->hc_mult = c->gdn_v_dim; + h->index_dim = c->gdn_k_heads; + h->n_blockres = 0; + return; + } h->version = 1; h->n_layers = c->n_layers; h->hidden = c->hidden; h->kda_heads = c->kda_heads; h->kda_dim = c->kda_dim; h->conv_k = c->conv_k; @@ -5468,6 +6179,42 @@ static void state_fill(const waste_model *m, waste_state_hdr *h, int pos) h->pos = pos; h->n_blockres = m->n_blockres; } +static int qwen_state_hdr_ok(const waste_config *c, const waste_state_hdr *h) +{ + return h->version == 2 && h->n_layers == c->n_layers && + h->hidden == c->hidden && h->kda_heads == c->gdn_v_heads && + h->kda_dim == c->gdn_k_dim && h->conv_k == c->conv_k && + h->n_heads == c->qsa_n_kv && h->qk_nope == c->hc_count && + h->qk_rope == c->qsa_head_dim && h->v_head == c->idx_head_dim && + h->attn_res_block == c->idx_compress && + h->hc_mult == c->gdn_v_dim && + h->index_dim == c->gdn_k_heads && h->n_blockres == 0; +} + +static uint64_t qwen_layer_state_bytes(const waste_config *c, int L, int T) +{ + if (c->qwen_full[L]) { + const int Hkv = c->qsa_n_kv, D = c->qsa_head_dim, Dk = c->idx_head_dim; + return 12ULL + (uint64_t)T * (uint64_t)Dk * 4ULL + + (uint64_t)T * (uint64_t)Hkv * (uint64_t)D * 4ULL; + } + const int Hv = c->gdn_v_heads, Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int Hk = c->gdn_k_heads; + const int qkv = 2 * Hk * Dk + Hv * Dv; + const int ck = c->conv_k > 0 ? c->conv_k - 1 : 0; + return (uint64_t)Hv * (uint64_t)Dk * (uint64_t)Dv * 4ULL + + (uint64_t)qkv * (uint64_t)ck * 4ULL; +} + +static uint64_t qwen_state_tail_bytes(const waste_config *c) +{ + const int R = (c->ple_conv_k > 1 && c->ngram_size > 0) + ? (c->ple_conv_k - 1) * c->ngram_size : 0; + return (uint64_t)c->hc_count * (uint64_t)c->hidden * 4ULL + + (uint64_t)c->hc_count * (uint64_t)c->hidden * + (uint64_t)(R > 0 ? R : 1) * 4ULL + 8ULL * 4ULL; +} + /* Every buffer a session accumulates into, back to the state of a fresh * open. Lived in waste.c reaching into the model's fields; it is here so * there is one copy, and so a measurement harness that drives the model @@ -5476,17 +6223,42 @@ void waste_model_reset(waste_model *m) { const waste_config *c = &m->cfg; for (int L = 0; L < c->n_layers; L++) { - if (m->S[L]) - memset(m->S[L], 0, (size_t)c->kda_heads * c->kda_dim * c->kda_dim * sizeof(float)); - if (m->conv[L]) - memset(m->conv[L], 0, - (size_t)3 * c->kda_heads * c->kda_dim * (c->conv_k - 1) * sizeof(float)); + if (c->arch_qwen) { + const int Hv = c->gdn_v_heads, Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int qkv = 2 * c->gdn_k_heads * Dk + Hv * Dv; + if (m->S[L]) + memset(m->S[L], 0, (size_t)Hv * Dk * Dv * sizeof(float)); + if (m->conv[L]) + memset(m->conv[L], 0, + (size_t)qkv * (c->conv_k > 0 ? c->conv_k - 1 : 0) * sizeof(float)); + m->n_qsa_blk[L] = 0; + m->n_qsa_tail[L] = 0; + if (m->qsa_rawk[L]) + memset(m->qsa_rawk[L], 0, + (size_t)m->kv_cap * c->idx_head_dim * sizeof(float)); + } else { + if (m->S[L]) + memset(m->S[L], 0, (size_t)c->kda_heads * c->kda_dim * c->kda_dim * sizeof(float)); + if (m->conv[L]) + memset(m->conv[L], 0, + (size_t)3 * c->kda_heads * c->kda_dim * (c->conv_k - 1) * sizeof(float)); + } m->n_kv[L] = 0; } m->n_blockres = 0; if (m->x) memset(m->x, 0, (size_t)(c->hc_mult ? c->hc_mult : 1) * c->hidden * sizeof(float)); for (int L = 0; L < c->n_layers; L++) m->n_pool[L] = 0; + if (c->arch_qwen) { + if (m->hcx) + memset(m->hcx, 0, (size_t)c->hc_count * c->hidden * sizeof(float)); + if (m->ple_ring) { + const int R = (c->ple_conv_k > 1 && c->ngram_size > 0) + ? (c->ple_conv_k - 1) * c->ngram_size : 0; + memset(m->ple_ring, 0, (size_t)c->hc_count * c->hidden * (R > 0 ? R : 1) * sizeof(float)); + } + for (int i = 0; i < 8; i++) m->ple_prev[i] = c->eos_token_id; + } if (m->blockres && c->attn_res_block) { const int nb = c->n_layers / c->attn_res_block + 2; memset(m->blockres, 0, (size_t)nb * c->hidden * sizeof(float)); @@ -5657,6 +6429,55 @@ int waste_model_state_save(const waste_model *m, const char *path, int pos) state_fill(m, &h, pos); int rc = fwrite(&h, sizeof h, 1, f) == 1 ? 0 : -1; + if (c->arch_qwen) { + for (int L = 0; L < c->n_layers && !rc; L++) { + if (c->qwen_full[L]) { + const int32_t T = m->n_kv[L]; + const int Hkv = c->qsa_n_kv, D = c->qsa_head_dim, Dk = c->idx_head_dim; + const int32_t blk = m->n_qsa_blk[L], tail = m->n_qsa_tail[L]; + if (fwrite(&T, sizeof T, 1, f) != 1 || + fwrite(&blk, sizeof blk, 1, f) != 1 || + fwrite(&tail, sizeof tail, 1, f) != 1) { + rc = -1; + break; + } + if (T > 0 && m->qsa_rawk[L] && + fwrite(m->qsa_rawk[L], sizeof(float), + (size_t)T * (size_t)Dk, f) != (size_t)T * (size_t)Dk) + rc = -1; + const size_t kvbf = (size_t)T * (size_t)Hkv * (size_t)D; + if (!rc && kvbf && m->qsa_k[L] && + fwrite(m->qsa_k[L], 2, kvbf, f) != kvbf) + rc = -1; + if (!rc && kvbf && m->qsa_v[L] && + fwrite(m->qsa_v[L], 2, kvbf, f) != kvbf) + rc = -1; + } else { + const int Hv = c->gdn_v_heads, Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int Hk = c->gdn_k_heads; + const int qkv = 2 * Hk * Dk + Hv * Dv; + const size_t sn = (size_t)Hv * (size_t)Dk * (size_t)Dv; + const size_t cn = (size_t)qkv * + (size_t)(c->conv_k > 0 ? c->conv_k - 1 : 0); + if (fwrite(m->S[L], sizeof(float), sn, f) != sn) rc = -1; + if (!rc && cn && fwrite(m->conv[L], sizeof(float), cn, f) != cn) rc = -1; + } + } + const int hcH = c->hc_count * c->hidden; + if (!rc && hcH && m->hcx && + fwrite(m->hcx, sizeof(float), (size_t)hcH, f) != (size_t)hcH) + rc = -1; + { + const int R = (c->ple_conv_k > 1 && c->ngram_size > 0) + ? (c->ple_conv_k - 1) * c->ngram_size : 0; + const size_t pr = (size_t)c->hc_count * (size_t)c->hidden * + (size_t)(R > 0 ? R : 1); + if (!rc && pr && m->ple_ring && + fwrite(m->ple_ring, sizeof(float), pr, f) != pr) + rc = -1; + } + if (!rc && fwrite(m->ple_prev, sizeof(int), 8, f) != 8) rc = -1; + } else { const int H = c->kda_heads, D = c->kda_dim, C = H * D; for (int L = 0; L < c->n_layers && !rc; L++) { if (c->ds41) { @@ -5716,6 +6537,8 @@ int waste_model_state_save(const waste_model *m, const char *path, int pos) const size_t xn = (size_t)(c->hc_mult ? c->hc_mult : 1) * c->hidden; if (!rc && fwrite(m->x, sizeof(float), xn, f) != xn) rc = -1; } + } /* end of the non-Qwen state: Qwen's residual is m->hcx, and it + * has neither blockres nor a widened m->x. */ if (!rc && waste_sync_file(f)) rc = -1; if (fclose(f)) rc = -1; if (!rc && waste_replace_file(tmp, path)) rc = -1; @@ -5732,10 +6555,10 @@ int waste_model_state_load(waste_model *m, const char *path, int *pos) waste_state_hdr h, want; state_fill(m, &want, 0); if (fread(&h, sizeof h, 1, f) != 1) { fclose(f); return -1; } - /* Every shape must match; pos and n_blockres are payload, but they also - * bound array indices and therefore need validation before the first - * byte of live state is replaced. */ - if (h.magic != want.magic || h.version != want.version || + if (h.magic != want.magic) { fclose(f); return -2; } + if (c->arch_qwen) { + if (!qwen_state_hdr_ok(c, &h)) { fclose(f); return -2; } + } else if (h.version != want.version || h.n_layers != want.n_layers || h.hidden != want.hidden || h.kda_heads != want.kda_heads || h.kda_dim != want.kda_dim || h.conv_k != want.conv_k || h.n_heads != want.n_heads || @@ -5745,26 +6568,50 @@ int waste_model_state_load(waste_model *m, const char *path, int *pos) h.head_dim != want.head_dim || h.window != want.window || h.engram_n != want.engram_n) { fclose(f); - return -2; /* state does not belong to this model */ + return -2; } const int H = c->kda_heads, D = c->kda_dim, C = H * D; const int nb_max = c->attn_res_block ? c->n_layers / c->attn_res_block + 2 : 0; if (h.pos < 0 || h.pos == INT32_MAX || - (m->has_mla && h.pos > m->kv_cap) || - h.n_blockres < 0 || h.n_blockres > nb_max) { + (c->arch_qwen ? h.pos > m->kv_cap : + (m->has_mla && h.pos > m->kv_cap)) || + (!c->arch_qwen && (h.n_blockres < 0 || h.n_blockres > nb_max))) { fclose(f); return -2; } - /* First walk the complete payload without touching the model. This - * rejects truncated files and bad per-layer KV counts up front, so the - * ordinary failure paths preserve the current conversation. */ const int64_t fsize_i = waste_file_size(fileno(f)); uint64_t off = sizeof h; if (fsize_i < 0) { fclose(f); return -1; } const uint64_t fsize = (uint64_t)fsize_i; + if (c->arch_qwen) { + for (int L = 0; L < c->n_layers; L++) { + uint64_t bytes = 0; + if (c->qwen_full[L]) { + int32_t T = 0, blk = 0, tail = 0; + if (off > fsize || fsize - off < 12 || + waste_pread(fileno(f), &T, 4, (int64_t)off) != 4 || + waste_pread(fileno(f), &blk, 4, (int64_t)(off + 4)) != 4 || + waste_pread(fileno(f), &tail, 4, (int64_t)(off + 8)) != 4) { + fclose(f); return -2; + } + if (T < 0 || T > m->kv_cap || T != h.pos) { + fclose(f); return -2; + } + bytes = qwen_layer_state_bytes(c, L, T); + } else { + bytes = qwen_layer_state_bytes(c, L, 0); + } + if (off > fsize || bytes > fsize - off) { fclose(f); return -2; } + off += bytes; + } + const uint64_t tail = qwen_state_tail_bytes(c); + if (off > fsize || tail > fsize - off || off + tail != fsize) { + fclose(f); return -2; + } + } else { for (int L = 0; L < c->n_layers; L++) { uint64_t bytes = 0; if (c->ds41) { @@ -5833,9 +6680,60 @@ int waste_model_state_load(waste_model *m, const char *path, int *pos) fclose(f); return -2; } } + } if (fseek(f, (long)sizeof h, SEEK_SET)) { fclose(f); return -1; } int rc = 0; + if (c->arch_qwen) { + for (int L = 0; L < c->n_layers && !rc; L++) { + if (c->qwen_full[L]) { + int32_t T = 0, blk = 0, tail = 0; + const int Hkv = c->qsa_n_kv, Dq = c->qsa_head_dim, Dk = c->idx_head_dim; + if (fread(&T, sizeof T, 1, f) != 1 || + fread(&blk, sizeof blk, 1, f) != 1 || + fread(&tail, sizeof tail, 1, f) != 1) { + rc = -1; break; + } + if (T > 0 && m->qsa_rawk[L] && + fread(m->qsa_rawk[L], sizeof(float), + (size_t)T * (size_t)Dk, f) != (size_t)T * (size_t)Dk) + rc = -1; + const size_t kvbf = (size_t)T * (size_t)Hkv * (size_t)Dq; + if (!rc && kvbf && m->qsa_k[L] && + fread(m->qsa_k[L], 2, kvbf, f) != kvbf) + rc = -1; + if (!rc && kvbf && m->qsa_v[L] && + fread(m->qsa_v[L], 2, kvbf, f) != kvbf) + rc = -1; + m->n_kv[L] = T; + m->n_qsa_blk[L] = blk; + m->n_qsa_tail[L] = tail; + } else { + const int Hv = c->gdn_v_heads, Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int Hk = c->gdn_k_heads; + const int qkv = 2 * Hk * Dk + Hv * Dv; + const size_t sn = (size_t)Hv * (size_t)Dk * (size_t)Dv; + const size_t cn = (size_t)qkv * + (size_t)(c->conv_k > 0 ? c->conv_k - 1 : 0); + if (fread(m->S[L], sizeof(float), sn, f) != sn) rc = -1; + if (!rc && cn && fread(m->conv[L], sizeof(float), cn, f) != cn) rc = -1; + } + } + const int hcH = c->hc_count * c->hidden; + if (!rc && hcH && m->hcx && + fread(m->hcx, sizeof(float), (size_t)hcH, f) != (size_t)hcH) + rc = -1; + { + const int R = (c->ple_conv_k > 1 && c->ngram_size > 0) + ? (c->ple_conv_k - 1) * c->ngram_size : 0; + const size_t pr = (size_t)c->hc_count * (size_t)c->hidden * + (size_t)(R > 0 ? R : 1); + if (!rc && pr && m->ple_ring && + fread(m->ple_ring, sizeof(float), pr, f) != pr) + rc = -1; + } + if (!rc && fread(m->ple_prev, sizeof(int), 8, f) != 8) rc = -1; + } else { for (int L = 0; L < c->n_layers && !rc; L++) { if (c->ds41) { const size_t wn = (size_t)c->window * c->head_dim; @@ -5889,6 +6787,8 @@ int waste_model_state_load(waste_model *m, const char *path, int *pos) const size_t xn = (size_t)(c->hc_mult ? c->hc_mult : 1) * c->hidden; if (!rc && fread(m->x, sizeof(float), xn, f) != xn) rc = -1; } + } /* end of the non-Qwen state: Qwen's residual is m->hcx, and it + * has neither blockres nor a widened m->x. */ fclose(f); if (!rc && pos) *pos = h.pos; /* -3 means the file changed or the device failed after the successful @@ -6380,6 +7280,14 @@ const float *waste_model_prefill(waste_model *m, const int *tokens, int n, const waste_config *c = &m->cfg; const int hid = c->hidden; if (n <= 0) return m->logits; + if (c->arch_qwen) { + const float *lg = NULL; + for (int t = 0; t < n; t++) { + lg = waste_model_step(m, tokens[t], pos0 + t, NULL); + if (!lg) return NULL; + } + return lg; + } if (n == 1) return waste_model_step(m, tokens[0], pos0, NULL); /* The chunked path carries one residual per token and one dense * attention per layer. mHC's parallel streams and the DSA indexer's @@ -6534,8 +7442,1097 @@ const float *waste_model_prefill(waste_model *m, const int *tokens, int n, return m->read_error ? NULL : m->logits; } +/* ---- Qwen3.8-Flash-Next forward (not KDA/MLA/AttnRes) ---------------- */ + +static uint16_t f32_to_bf16(float x) +{ + union { float f; uint32_t u; } a; + a.f = x; + const uint32_t u = a.u; + return (uint16_t)((u + 0x7fffu + ((u >> 16) & 1u)) >> 16); +} + +static float bf16_to_f32(uint16_t b) +{ + union { float f; uint32_t u; } a; + a.u = (uint32_t)b << 16; + return a.f; +} + +static void qwen_row(waste_model *m, const waste_tensor *t, long row, float *dst) +{ + const int cols = t->shape[t->ndim - 1]; + if (!t->on_disk && t->data) { + memcpy(dst, t->data + (size_t)row * (size_t)cols, (size_t)cols * sizeof(float)); + return; + } + if (!t->on_disk && t->q) { + waste_deq_row(t, row, cols, dst); + return; + } + const int g = t->group, ng = (cols + g - 1) / g; + const int8_t *q; const uint16_t *sc; + trunk_row(m, t, row, &q, &sc); + for (int k = 0; k < ng; k++) { + const float sv = f16_to_f32(sc[k]); + for (int i = 0; i < g && k * g + i < cols; i++) { + int v; + if (t->bits == 4) { + const uint8_t byte = ((const uint8_t *)q)[(k * g + i) / 2]; + v = (i & 1) ? (byte >> 4) - 8 : (byte & 0x0F) - 8; + } else { + v = q[k * g + i]; + } + dst[k * g + i] = (float)v * sv; + } + } +} + +static void qwen_rope_cs(const waste_config *c, int pos, float *cos, float *sin) +{ + const int half = c->rotary_dim / 2; + float ft[WASTE_MAX_ROPE_HALF], fh[WASTE_MAX_ROPE_HALF]; + float fw[WASTE_MAX_ROPE_HALF], freqs[WASTE_MAX_ROPE_HALF]; + if (half <= 0 || half > WASTE_MAX_ROPE_HALF) return; + for (int j = 0; j < half; j++) { + const float a = (float)pos * c->rope_inv_freq[j]; + ft[j] = fh[j] = fw[j] = a; + } + waste_qwen_mrope_interleave(ft, fh, fw, c->mrope_section, half, freqs); + for (int j = 0; j < half; j++) { + const float cj = cosf(freqs[j]), sj = sinf(freqs[j]); + cos[j] = cj; sin[j] = sj; + cos[j + half] = cj; sin[j + half] = sj; + } +} + +/* One stream of a HyperConnection mix's front half: the combine that + * finishes the previous block when there is one, the RMSNorm, and — when + * the down projection can take them — the i8mm planes of the stream's + * groups. Each is the same function over the same elements as the serial + * loop it replaced, one stream at a time. */ +typedef struct { + float *o, *x; + const float *w; + int group; + float eps; + const float *block, *inj; /* combine into x first, unless NULL */ + int qg, n; /* quantize o in groups of qg, if > 0 */ + int8_t *q; + float *sc; +} hcn_arg; + +static void hc_norm_range(int b, int e, void *p) +{ + const hcn_arg *a = (const hcn_arg *)p; + for (int s = b; s < e; s++) { + const size_t off = (size_t)s * (size_t)a->group; + if (a->block) + waste_qwen_hc_combine(a->x + off, a->block, a->inj + s, 1, a->group, + a->x + off); + waste_qwen_rmsnorm(a->o + off, a->x + off, a->w + off, + a->group, a->group, a->eps); +#if defined(__ARM_NEON) || defined(__aarch64__) + if (a->qg) { + const int per = a->group / a->qg; + for (int k = s * per; k < (s + 1) * per; k++) + quant_act4_mm_group(a->o, a->n, a->qg, k, a->q, a->sc); + } +#endif + } +} + +/* The back half, as pieces of one job: the sigmoid of every stream's gate + * and the weighted sum over streams, a hidden range at a time, and the + * inject projection's rows. The sigmoid is a scalar expf per element and + * the sum runs over streams in order for each element, as the serial loop + * did; an inject row is the dotf matvec() would have taken. */ +enum { HC_SPAN = 256 }; +typedef struct { + float *gate, *mixed, *yi; + const float *normed, *W; /* W: float inject rows, or NULL */ + int hc, hid, n_mix; +} hcg_arg; + +static void hc_gate_mix_piece(int b, int e, void *p) +{ + const hcg_arg *a = (const hcg_arg *)p; + const int hc = a->hc, hid = a->hid; + for (int k = b; k < e; k++) { + if (k < a->n_mix) { + const int d0 = k * HC_SPAN, d1 = d0 + HC_SPAN < hid ? d0 + HC_SPAN : hid; + for (int st = 0; st < hc; st++) { + float *v = a->gate + (size_t)st * hid; + for (int d = d0; d < d1; d++) v[d] = 1.0f / (1.0f + expf(-v[d])); + } + for (int d = d0; d < d1; d++) { + float s = 0.0f; + for (int st = 0; st < hc; st++) + s += a->gate[st * hid + d] * a->normed[st * hid + d]; + a->mixed[d] = s / (float)hc; + } + } else { + const int o = k - a->n_mix, H = hc * hid; + a->yi[o] = dotf(a->W + (size_t)o * H, a->normed, H); + } + } +} + +/* `cblock`, when given, is the block the streams have not taken in yet: + * the combine with the previous mix's `cinj` weights happens here, stream by + * stream in the same tasks as the norm, rather than as a serial pass in front + * of it. `cinj` may be `inj_w` itself — it is copied before anything writes. + * + * The mix used to be six dispatches and four serial stretches, and each + * stretch long enough for the pool to park before the next dispatch: the + * combine (3.5 us), the down projection's quantization, the sum over streams + * (4.4 us) and the inject projection (5 us). It is now three dispatches, the + * norm and the gate each carrying the work that sat between them and the + * matvecs (LEARNED §86). */ +static void qwen_hc_mix_t(waste_model *m, float *hyper, + const float *cblock, const float *cinj, + const waste_tensor *nw, const waste_tensor *down, + const waste_tensor *up, const waste_tensor *inject, + int use_inj, float *mixed, float *inj_w) +{ + const waste_config *c = &m->cfg; + const int hc = c->hc_count, hid = c->hidden, rank = c->hc_lowrank; + const int H = hc * hid; + float inj_prev[16]; + if (cblock) memcpy(inj_prev, cinj, (size_t)hc * sizeof(float)); + if (!nw || !nw->data || !down || !up) { + if (cblock) waste_qwen_hc_combine(hyper, cblock, inj_prev, hc, hid, hyper); + memset(mixed, 0, (size_t)hid * sizeof(float)); + if (inj_w) memset(inj_w, 0, (size_t)hc * sizeof(float)); + return; + } + float *normed = m->tmp; + float *lo = normed + H; + float *gate = lo + rank; + const int pq = prequant_ok(down, hid); + { + hcn_arg na = { normed, hyper, nw->data, hid, c->eps, + cblock, inj_prev, pq ? down->group : 0, H, m->xq, m->xs }; + waste_parallel_for_fast(hc, 1, hc_norm_range, &na); + } + if (pq) matvec_t_prequant(m, lo, down, rank, H); + else matvec_t(m, lo, down, normed, rank, H); + for (int i = 0; i < rank; i++) lo[i] = silu(lo[i] / (float)hc); + matvec_t(m, gate, up, lo, H, rank); + + const int want_inj = use_inj && inject && inj_w && hc <= 16; + const int inj_rows = want_inj && !inject->q && inject->data; + float tmpi[16]; + { + const int n_mix = (hid + HC_SPAN - 1) / HC_SPAN; + hcg_arg ga = { gate, mixed, tmpi, normed, inj_rows ? inject->data : NULL, + hc, hid, n_mix }; + const double t0 = prof_on && inj_rows ? pnow() : 0; + waste_parallel_for_each(n_mix + (inj_rows ? hc : 0), hc_gate_mix_piece, &ga, + waste_pool_fast()); + if (prof_on && inj_rows) { + /* The inject rows share a job with the mix; the matvec table + * gets their share of it by piece count, which is an estimate. */ + pthread_mutex_lock(&prof_mu); + tmv_account(inject, hc, H, (pnow() - t0) * hc / (n_mix + hc), 0.0); + pthread_mutex_unlock(&prof_mu); + } + } + if (want_inj) { + if (!inj_rows) matvec_t(m, tmpi, inject, normed, hc, H); + for (int b = 0; b < hc; b++) + inj_w[b] = 2.0f / (1.0f + expf(-tmpi[b] / (float)hc)); + } +} + +static void qwen_dilated_conv_step(int C, int KS, int dil, const float *w, + float *ring, const float *x, float *y) +{ + const int R = (KS - 1) * dil; + for (int c = 0; c < C; c++) { + const float *wc = w + (size_t)c * KS; + float *rc = ring + (size_t)c * R; + float acc = x[c] * wc[KS - 1]; + for (int k = 0; k < KS - 1; k++) + acc += rc[k * dil] * wc[k]; + for (int j = 0; j + 1 < R; j++) rc[j] = rc[j + 1]; + if (R > 0) rc[R - 1] = x[c]; + y[c] = silu(acc); + } +} + +static void qwen_ple_inject(waste_model *m, int token) +{ + const waste_config *c = &m->cfg; + const int L = c->ple_layer; + if (L < 0) return; + const int hid = c->hidden, hc = c->hc_count, H = hc * hid; + const int pe = c->ple_embed ? c->ple_embed : hid; + const int ngram = c->ngram_size > 0 ? c->ngram_size : 3; + const int heads = (ngram - 1) * (c->heads_per_ngram ? c->heads_per_ngram : 8); + const int ctxn = ngram - 1; + int ids[8]; + for (int i = 0; i < ctxn && i < 8; i++) ids[i] = m->ple_prev[i]; + ids[ctxn] = token; + const int n = ctxn + 1; + int local[WASTE_QWEN_PLE_HEADS]; + waste_qwen_ple_row_ids(ids, n, ctxn, c->eos_token_id, ngram, + c->heads_per_ngram ? c->heads_per_ngram : 8, + c->ple_mult, c->ple_sz, local); + float *emb = m->ple_emb; + if (!emb) return; + memset(emb, 0, (size_t)pe * sizeof(float)); + int off = 0; + for (int h = 0; h < heads && h < WASTE_QWEN_PLE_HEADS; h++) { + if (c->ple_sz[h] <= 0) return; + const waste_tensor *ht = waste_find(m, tname( + "%smodel.layers.%d.ple.ple_embedding.ngram_head.%d.weight", + c->prefix, L, h)); + if (!ht) continue; + const int width = ht->shape[ht->ndim - 1]; + if (off + width > pe) break; + qwen_row(m, ht, local[h], emb + off); + m->ple_reads++; + off += width; + } + float *key = m->tmp, *val = key + H, *qnorm = val + hid; + matvec_t(m, key, waste_find(m, tname("%smodel.layers.%d.ple.key_proj.weight", + c->prefix, L)), emb, H, pe); + matvec_t(m, val, waste_find(m, tname("%smodel.layers.%d.ple.value_proj.weight", + c->prefix, L)), emb, hid, pe); + const waste_tensor *tnk = waste_find(m, tname("%smodel.layers.%d.ple.norm_key.weight", + c->prefix, L)); + const waste_tensor *tnq = waste_find(m, tname("%smodel.layers.%d.ple.norm_query.weight", + c->prefix, L)); + const waste_tensor *tnc = waste_find(m, tname("%smodel.layers.%d.ple.norm_conv.weight", + c->prefix, L)); + if (!tnk || !tnk->data || !tnq || !tnq->data || !tnc || !tnc->data) return; + const float *nk = tnk->data, *nq = tnq->data, *nc = tnc->data; + waste_qwen_rmsnorm(key, key, nk, H, hid, c->eps); + waste_qwen_rmsnorm(qnorm, m->hcx, nq, H, hid, c->eps); + float *gated = qnorm + H; + const float inv = 1.0f / sqrtf((float)hid); + for (int b = 0; b < hc; b++) { + float g = 0.0f; + for (int d = 0; d < hid; d++) + g += key[b * hid + d] * qnorm[b * hid + d]; + g *= inv; + const float mag = sqrtf(fabsf(g) < 1e-6f ? 1e-6f : fabsf(g)); + g = copysignf(mag, g); + const float sg = 1.0f / (1.0f + expf(-g)); + for (int d = 0; d < hid; d++) gated[b * hid + d] = sg * val[d]; + } + float *gnorm = gated + H; + waste_qwen_rmsnorm(gnorm, gated, nc, H, hid, c->eps); + const waste_tensor *cw = waste_find(m, tname("%smodel.layers.%d.ple.conv1d.weight", + c->prefix, L)); + const int KS = c->ple_conv_k; + float *conv_y = gnorm + H; + if (cw && cw->data) + qwen_dilated_conv_step(H, KS, ngram, cw->data, m->ple_ring, gnorm, conv_y); + else + memcpy(conv_y, gnorm, (size_t)H * sizeof(float)); + for (int i = 0; i < H; i++) m->hcx[i] += gated[i] + conv_y[i]; + for (int i = 0; i < ctxn - 1 && i < 7; i++) m->ple_prev[i] = m->ple_prev[i + 1]; + if (ctxn > 0) m->ple_prev[ctxn - 1] = token; +} + +typedef struct { + int Hk, Hv, Dk, Dv; + const float *q, *k, *v, *g_log, *beta; + float *S, *o; +} gdn_arg; + +enum { GDN_SCRATCH = 1024 }; + +static void gdn_heads_range(int b, int e, void *p) +{ + const gdn_arg *a = (const gdn_arg *)p; + float u[GDN_SCRATCH]; + waste_qwen_gdn_step_heads(b, e, a->Hk, a->Hv, a->Dk, a->Dv, a->q, a->k, a->v, + a->g_log, a->beta, a->S, a->o, u); +} + +/* GDN's short conv, a range of channels at a time. Each channel reads its + * own ring and input and writes its own output, through the same kernel the + * whole-layer call used. On the calling thread it was 46 us a layer of + * SiLU — the pool asleep by the time the recurrence reached it. */ +typedef struct { int KS; const float *w; float *ring; const float *x; float *y; } gconv_arg; + +static void gdn_conv_range(int b, int e, void *p) +{ + const gconv_arg *a = (const gconv_arg *)p; + const int R = a->KS - 1; + waste_k.short_conv_step(e - b, a->KS, a->w + (size_t)b * a->KS, NULL, + a->ring + (size_t)b * R, a->x + b, a->y + b); +} + +/* A value head from the recurrence to the output projection's input: its + * state update, its gated RMSNorm, and — when out_proj can take them — the + * i8mm planes of its groups. The norm used to be a serial loop over the heads + * after the recurrence's dispatch, and out_proj's quantization a dispatch + * after that, which found the pool parked 47 times a token. */ +typedef struct { + gdn_arg r; + const float *z, *nw; + float eps; + float *normed; + int qg, n; /* quantize normed in groups, if > 0 */ + int8_t *q; + float *sc; +} gdnf_arg; + +static void gdn_heads_out_range(int b, int e, void *p) +{ + const gdnf_arg *a = (const gdnf_arg *)p; + const int Dv = a->r.Dv; + float u[GDN_SCRATCH]; + waste_qwen_gdn_step_heads(b, e, a->r.Hk, a->r.Hv, a->r.Dk, Dv, a->r.q, a->r.k, + a->r.v, a->r.g_log, a->r.beta, a->r.S, a->r.o, u); + for (int h = b; h < e; h++) + waste_k.rmsnorm_gated(Dv, a->r.o + (size_t)h * Dv, a->z + (size_t)h * Dv, + a->nw, a->eps, a->normed + (size_t)h * Dv); +#if defined(__ARM_NEON) || defined(__aarch64__) + if (a->qg) { + const int per = Dv / a->qg; + for (int k = b * per; k < e * per; k++) + quant_act4_mm_group(a->normed, a->n, a->qg, k, a->q, a->sc); + } +#endif +} + +static void qwen_gdn_layer(waste_model *m, int L, const float *in, float *out) +{ + const waste_config *c = &m->cfg; + const int hid = c->hidden, Hk = c->gdn_k_heads, Hv = c->gdn_v_heads; + const int Dk = c->gdn_k_dim, Dv = c->gdn_v_dim; + const int qkv = 2 * Hk * Dk + Hv * Dv; + memset(out, 0, (size_t)hid * sizeof(float)); + if (!m->gdn_g) return; + float *mixed = m->tmp; + float *conv_y = mixed + qkv; + float *z = conv_y + qkv; + float *a = z + Hv * Dv; + float *b = a + Hv; + float *core = b + Hv; + { + /* Four projections of one vector, as one job. The conv reads only + * the first, so it follows all four. */ + const mvb_item proj[4] = { + { mixed, waste_find(m, tname("%smodel.layers.%d.linear_attn.in_proj_qkv.weight", + c->prefix, L)), qkv }, + { z, waste_find(m, tname("%smodel.layers.%d.linear_attn.in_proj_z.weight", + c->prefix, L)), Hv * Dv }, + { a, waste_find(m, tname("%smodel.layers.%d.linear_attn.in_proj_a.weight", + c->prefix, L)), Hv }, + { b, waste_find(m, tname("%smodel.layers.%d.linear_attn.in_proj_b.weight", + c->prefix, L)), Hv }, + }; + matvec_t_batch(m, in, hid, proj, 4); + } + const waste_tensor *cw = waste_find(m, tname("%smodel.layers.%d.linear_attn.conv1d.weight", + c->prefix, L)); + if (cw && cw->data) { + gconv_arg ca = { c->conv_k, cw->data, m->conv[L], mixed, conv_y }; + waste_parallel_for_fast(qkv, 512, gdn_conv_range, &ca); + } else { + memcpy(conv_y, mixed, (size_t)qkv * sizeof(float)); + } + for (int h = 0; h < Hv; h++) b[h] = 1.0f / (1.0f + expf(-b[h])); + const waste_tensor *tA = waste_find(m, tname("%smodel.layers.%d.linear_attn.A_log", + c->prefix, L)); + const waste_tensor *tdt = waste_find(m, tname("%smodel.layers.%d.linear_attn.dt_bias", + c->prefix, L)); + if (!tA || !tA->data || !tdt || !tdt->data) return; + PROF_START(P_KDAK); + waste_qwen_gdn_decay(a, tA->data, tdt->data, Hv, m->gdn_g); + const float *q = conv_y; + const float *k = conv_y + Hk * Dk; + const float *v = conv_y + 2 * Hk * Dk; + /* One task per value head's range. On the calling thread this was 5.3 ms + * of a step — 147 us a layer, eighteen times the 8 us a pool worker waits + * before parking — between in_proj_qkv and out_proj, so it both ran on + * one core and put the pool to sleep for the projection after it. The + * heads share nothing but the QK rows they read (see qwen_gdn.h), and + * each runs the same code in the same order, so the state and output are + * the serial loop's bit for bit. */ + const waste_tensor *tnw = waste_find(m, tname("%smodel.layers.%d.linear_attn.norm.weight", + c->prefix, L)); + const waste_tensor *top = waste_find(m, tname("%smodel.layers.%d.linear_attn.out_proj.weight", + c->prefix, L)); + if (Dv <= GDN_SCRATCH && tnw && tnw->data) { + /* `normed` is `mixed`: in_proj_qkv's output, which nothing reads + * after the conv, so the heads can write it while they run. */ + const int pq = prequant_ok(top, Dv); + gdnf_arg fa = { { Hk, Hv, Dk, Dv, q, k, v, m->gdn_g, b, m->S[L], core }, + z, tnw->data, c->eps, mixed, + pq ? top->group : 0, Hv * Dv, m->xq, m->xs }; + waste_parallel_for_fast(Hv, 1, gdn_heads_out_range, &fa); + PROF_END(P_KDAK); + if (pq) matvec_t_prequant(m, out, top, hid, Hv * Dv); + else matvec_t(m, out, top, mixed, hid, Hv * Dv); + return; + } + if (Dv <= GDN_SCRATCH) { + gdn_arg ga = { Hk, Hv, Dk, Dv, q, k, v, m->gdn_g, b, m->S[L], core }; + waste_parallel_for_fast(Hv, 1, gdn_heads_range, &ga); + } else { + waste_qwen_gdn_step(Hk, Hv, Dk, Dv, q, k, v, m->gdn_g, b, m->S[L], core, m->att); + } + PROF_END(P_KDAK); + if (!tnw || !tnw->data) return; + float *normed = mixed; + for (int h = 0; h < Hv; h++) + waste_k.rmsnorm_gated(Dv, core + (size_t)h * Dv, z + (size_t)h * Dv, + tnw->data, c->eps, normed + (size_t)h * Dv); + matvec_t(m, out, top, normed, hid, Hv * Dv); +} + +typedef struct { + const float *q, *k, *v; + const int *sel; + float *out, *scr; + int Hq, D, Hkv, n_sel; + float scale; +} qsaa_arg; + +typedef struct { + const float *q_idx, *raw_k, *cos, *sin, *k_ln_w; + float *pooled, *scores; + int Hq, Dk, rot, compress; + float eps; +} qsas_arg; + +static void qsa_score_range(int b, int e, void *p) +{ + const qsas_arg *a = (const qsas_arg *)p; + waste_qwen_qsa_score_blocks(b, e, a->q_idx, a->Hq, a->Dk, a->raw_k, a->cos, a->sin, + a->rot, a->k_ln_w, a->eps, a->compress, a->pooled, + a->scores); +} + +typedef struct { + const uint16_t *kq, *vq; + float *kf, *vf; + int *sel; + int Hkv, D, kvd, T; +} qsag_arg; + +static void qsa_gather_range(int b, int e, void *p) +{ + const qsag_arg *a = (const qsag_arg *)p; + for (int i = b; i < e; i++) { + const int t = a->sel[i]; + a->sel[i] = i; + const int in = t >= 0 && t < a->T; + const uint16_t *kb = in ? a->kq + (size_t)t * a->Hkv * a->D : NULL; + const uint16_t *vb = in ? a->vq + (size_t)t * a->Hkv * a->D : NULL; + for (int j = 0; j < a->kvd; j++) { + a->kf[(size_t)i * a->kvd + j] = in ? bf16_to_f32(kb[j]) : 0.0f; + a->vf[(size_t)i * a->kvd + j] = in ? bf16_to_f32(vb[j]) : 0.0f; + } + } +} + +/* Each head gets the row of scr at h * n_sel, which qsa_scr is sized for. */ +static void qsa_attn_range(int b, int e, void *p) +{ + const qsaa_arg *a = (const qsaa_arg *)p; + for (int h = b; h < e; h++) + waste_qwen_qsa_attn_heads(h, h + 1, a->q, a->Hq, a->D, a->k, a->v, a->Hkv, + a->n_sel, a->sel, a->n_sel, a->scale, a->out, + a->scr + (size_t)h * (size_t)a->n_sel); +} + +static void qwen_qsa_layer(waste_model *m, int L, const float *in, float *out, int pos) +{ + const waste_config *c = &m->cfg; + const int hid = c->hidden, Hq = c->n_heads, Hkv = c->qsa_n_kv, D = c->qsa_head_dim; + const int Dk = c->idx_head_dim, compress = c->idx_compress > 0 ? c->idx_compress : 4; + const int qd = Hq * D, kvd = Hkv * D; + const int idxd = (c->idx_n_heads + c->idx_kv_heads) * Dk; + const int rot = c->rotary_dim; + memset(out, 0, (size_t)hid * sizeof(float)); + if (!m->qsa_q || !m->qsa_gate || !m->qsa_attn || !m->qsa_sel || + !m->qsa_kf || !m->qsa_vf || !m->qsa_rawk[L]) + return; + float *qgate = m->tmp; + float *k = qgate + qd * 2; + float *v = k + kvd; + float *idx = v + kvd; + { + const mvb_item proj[4] = { + { qgate, waste_find(m, tname("%smodel.layers.%d.self_attn.q_proj.weight", + c->prefix, L)), qd * 2 }, + { k, waste_find(m, tname("%smodel.layers.%d.self_attn.k_proj.weight", + c->prefix, L)), kvd }, + { v, waste_find(m, tname("%smodel.layers.%d.self_attn.v_proj.weight", + c->prefix, L)), kvd }, + { idx, waste_find(m, tname("%smodel.layers.%d.self_attn.indexer.index_qk_proj.weight", + c->prefix, L)), idxd }, + }; + matvec_t_batch(m, in, hid, proj, 4); + } + float *q = m->qsa_q, *gate = m->qsa_gate; + for (int h = 0; h < Hq; h++) { + memcpy(q + (size_t)h * D, qgate + (size_t)h * 2 * D, (size_t)D * sizeof(float)); + memcpy(gate + (size_t)h * D, qgate + (size_t)h * 2 * D + D, (size_t)D * sizeof(float)); + } + const waste_tensor *tqn = waste_find(m, tname("%smodel.layers.%d.self_attn.q_norm.weight", + c->prefix, L)); + const waste_tensor *tkn = waste_find(m, tname("%smodel.layers.%d.self_attn.k_norm.weight", + c->prefix, L)); + if (!tqn || !tqn->data || !tkn || !tkn->data) return; + for (int h = 0; h < Hq; h++) + waste_qwen_rmsnorm(q + (size_t)h * D, q + (size_t)h * D, tqn->data, D, D, c->eps); + for (int h = 0; h < Hkv; h++) + waste_qwen_rmsnorm(k + (size_t)h * D, k + (size_t)h * D, tkn->data, D, D, c->eps); + float cos[256], sin[256]; + qwen_rope_cs(c, pos, cos, sin); + for (int h = 0; h < Hq; h++) + if (waste_qwen_rope_apply(q + (size_t)h * D, D, cos, sin, rot) != 0) return; + for (int h = 0; h < Hkv; h++) + if (waste_qwen_rope_apply(k + (size_t)h * D, D, cos, sin, rot) != 0) return; + + if (pos >= 0 && pos < m->kv_cap) { + uint16_t *kb = m->qsa_k[L] + (size_t)pos * Hkv * D; + uint16_t *vb = m->qsa_v[L] + (size_t)pos * Hkv * D; + for (int i = 0; i < kvd; i++) { + kb[i] = f32_to_bf16(k[i]); + vb[i] = f32_to_bf16(v[i]); + } + m->n_kv[L] = pos + 1; + } + const int iq = c->idx_n_heads * Dk; + float *q_idx = idx; + float *raw_k = idx + iq; + const waste_tensor *tqln = waste_find(m, tname( + "%smodel.layers.%d.self_attn.indexer.q_layernorm.weight", c->prefix, L)); + const waste_tensor *tkln = waste_find(m, tname( + "%smodel.layers.%d.self_attn.indexer.k_layernorm.weight", c->prefix, L)); + if (!tqln || !tqln->data || !tkln || !tkln->data) return; + for (int h = 0; h < c->idx_n_heads; h++) { + waste_qwen_rmsnorm(q_idx + (size_t)h * Dk, q_idx + (size_t)h * Dk, + tqln->data, Dk, Dk, c->eps); + if (waste_qwen_rope_apply(q_idx + (size_t)h * Dk, Dk, cos, sin, rot) != 0) + return; + } + if (pos >= 0 && pos < m->kv_cap && m->qsa_rawk[L]) + memcpy(m->qsa_rawk[L] + (size_t)pos * Dk, raw_k, (size_t)Dk * sizeof(float)); + const int T = m->n_kv[L]; + m->n_qsa_blk[L] = compress > 0 ? T / compress : 0; + m->n_qsa_tail[L] = compress > 0 ? T % compress : 0; + + PROF_START(P_QSAK); + const int block_topk = c->idx_budget / compress; + float *full_cos = m->qsa_cs; + float *full_sin = m->qsa_cs + (size_t)m->kv_cap * (rot > 0 ? rot : 1); + PROF_START(P_QSAR); + /* A row of the table is a function of its position and nothing else, so + * once written it is right for every later token, every layer and every + * session — a reset or a restore leaves it valid. Rewriting all T rows + * every token in every QSA layer was 4.3 ms a step at 2,830 tokens and + * growing with the context; only the rows past the last one filled are + * ever new. */ + if (full_cos && rot > 0) { + for (int t = m->qsa_cs_n; t < T; t++) + qwen_rope_cs(c, t, full_cos + (size_t)t * rot, full_sin + (size_t)t * rot); + if (T > m->qsa_cs_n) m->qsa_cs_n = T; + } + PROF_END(P_QSAR); + PROF_START(P_QSAS); + /* waste_qwen_qsa_select, taken apart so the blocks can be scored at once: + * each writes its own pooled row and its own score (qwen_qsa.h). The pick + * is a sort in the order the argmax chose, not a pass per block kept, and + * together they were 5.2 ms a step at 2,830 tokens. */ + int nsel = 0; + if (T >= 1) { + const int qp = pos < 0 ? 0 : (pos >= T ? T - 1 : pos); + const int n_complete = (qp + 1) / compress; + const int n_tail = qp + 1 - n_complete * compress; + const float *scores = m->qsa_work + (size_t)n_complete * Dk; + if (n_complete > 0) { + qsas_arg sa = { q_idx, m->qsa_rawk[L], full_cos, full_sin, tkln->data, + m->qsa_work, m->qsa_work + (size_t)n_complete * Dk, + c->idx_n_heads, Dk, rot, compress, c->eps }; + if (n_complete >= 32) + waste_parallel_for_fast(n_complete, 4, qsa_score_range, &sa); + else + qsa_score_range(0, n_complete, &sa); + } + nsel = waste_qwen_qsa_pick(n_complete > 0 ? scores : NULL, n_complete, block_topk, + compress, n_tail, m->qsa_sel, m->qsa_taken); + } + PROF_END(P_QSAS); + PROF_START(P_QSAG); + float *kf = m->qsa_kf, *vf = m->qsa_vf, *attn = m->qsa_attn, *scr = m->qsa_scr; + /* Each selected index writes its own rows of kf and vf and its own slot of + * sel, so the conversion goes in ranges: 5.3 ms a step at 2,830 tokens on + * one core. */ + { + qsag_arg ga = { m->qsa_k[L], m->qsa_v[L], kf, vf, m->qsa_sel, Hkv, D, kvd, T }; + waste_parallel_for_fast(nsel, 16, qsa_gather_range, &ga); + } + PROF_END(P_QSAG); + PROF_START(P_QSAA); + /* One task per query head. Serial this was 3.4 ms of a step at a 220 + * token context and 58 ms at 2,800 — a third of the step, on one core, + * and it stops growing only when the selection fills its budget. A head + * reads its own query and its KV head's rows and writes its own row of + * attn (qwen_qsa.h), with its own row of scores, so the output is the + * serial loop's bit for bit. */ + if (nsel > 0) { + qsaa_arg aa = { q, kf, vf, m->qsa_sel, attn, scr, Hq, D, Hkv, nsel, + 1.0f / sqrtf((float)D) }; + waste_parallel_for_fast(Hq, 1, qsa_attn_range, &aa); + } else { + memset(attn, 0, (size_t)qd * sizeof(float)); + } + PROF_END(P_QSAA); + PROF_END(P_QSAK); + for (int i = 0; i < qd; i++) + attn[i] *= 1.0f / (1.0f + expf(-gate[i])); + matvec_t(m, out, waste_find(m, tname("%smodel.layers.%d.self_attn.o_proj.weight", + c->prefix, L)), attn, hid, qd); +} + +/* One routed expert through the row-parallel kernels, into `acc`. + * + * The serial loop's body: the path WASTE_XPAR=0 forces, and the one a layer + * falls back to when a record does not read. */ +static void qwen_expert_rows(waste_model *m, const uint8_t *rec, int inter, + int hid, const float *lut_gate, + const float *lut_up, float *ga, float *ub, + float *acc, float *lut_down) +{ + const waste_expert_hdr *h = (const waste_expert_hdr *)rec; + const uint16_t *corr = (const uint16_t *)(rec + h->chan_corr_off); + vq_apply(m, ga, rec + h->gate_off, corr, inter, hid, lut_gate, NULL, NULL); + vq_apply(m, ub, rec + h->up_off, corr + inter, inter, hid, lut_up, NULL, NULL); + for (int i = 0; i < inter; i++) ga[i] = silu(ga[i]) * ub[i]; + vq_matvec(m, acc, rec + h->down_off, corr + 2 * inter, ga, hid, inter, + h->codebook_id + 2 * m->stages, lut_down, NULL, NULL); +} + +/* The shared expert and its gate, into `acc`; returns the gate. The add + * into the layer output is the caller's, and stays after the routed sum, + * so running this early changes when it is computed and not what is + * summed. + * + * `pre` says the router's job already projected `in` through the gate and + * up matrices into m->ff and the gate scalar into `sg_raw` (see + * qwen_moe_layer); what is left is the activation, the down projection and + * the sigmoid. Without it, all of it is done here, as ffn does it. */ +static float qwen_shared_expert(waste_model *m, int L, const float *in, float *acc, + int pre, float sg_raw) +{ + const waste_config *c = &m->cfg; + const int hid = c->hidden; + PROF_START(P_QSHX); + const int shared = c->shared_inter ? c->shared_inter : c->moe_inter; + if (pre) { + waste_act_pair_range(c, m->ff, m->ff + shared, shared); + matvec_t(m, acc, waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert.down_proj.weight", + c->prefix, L)), m->ff, hid, shared); + } else { + ffn(m, + waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert.gate_proj.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert.up_proj.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert.down_proj.weight", c->prefix, L)), + in, acc, shared, hid, 1.0f, 0); + matvec_t(m, &sg_raw, waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert_gate.weight", + c->prefix, L)), in, 1, hid); + } + const float sg = 1.0f / (1.0f + expf(-sg_raw)); + PROF_END(P_QSHX); + return sg; +} + +static void qwen_moe_layer(waste_model *m, int L, const float *in, float *out, int *routed) +{ + const waste_config *c = &m->cfg; + const int E = c->n_experts, K = c->top_k, hid = c->hidden, inter = c->moe_inter; + float *sc = m->att + WASTE_ATT_ROUTER_OFF; + int idx[64]; + float w[64]; + const int shared_in = c->shared_inter ? c->shared_inter : inter; + float sg_raw = 0.0f; + int shared_pre = 1; + PROF_START(P_QRTR); + { + /* The router and the shared expert's gate, up and gate scalar all + * read `in`: one job. The shared expert keeps its outputs in m->ff + * until it runs, which nothing on the routed paths below touches + * except the serial loop — and that clears shared_pre. */ + const mvb_item proj[4] = { + { sc, waste_find(m, tname("%smodel.layers.%d.mlp.gate.weight", c->prefix, L)), E }, + { m->ff, waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert.gate_proj.weight", + c->prefix, L)), shared_in }, + { m->ff + shared_in, waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert.up_proj.weight", + c->prefix, L)), shared_in }, + { &sg_raw, waste_find(m, tname("%smodel.layers.%d.mlp.shared_expert_gate.weight", + c->prefix, L)), 1 }, + }; + matvec_t_batch(m, in, hid, proj, 4); + } + const int route_rc = waste_qwen_moe_route(sc, E, K, c->renorm, idx, w, + m->moe_prob, m->moe_used); + PROF_END(P_QRTR); + if (route_rc != 0) { + memset(out, 0, (size_t)hid * sizeof(float)); + return; + } + if (routed) for (int j = 0; j < K; j++) routed[j] = idx[j]; + if (dump_route) { + FILE *df = fopen(dump_route, (dump_pos0 || L) ? "a" : "wb"); + if (df) { + fprintf(df, "%d %d", dump_pos0, L); + for (int j = 0; j < K; j++) fprintf(df, " %d", idx[j]); + for (int j = 0; j < K; j++) fprintf(df, " %.6g", w[j]); + for (int j = 0; j < K; j++) fprintf(df, " -1"); + fputc('\n', df); + fclose(df); + } + } + waste_ecache_hint(&m->cache, L, idx, K); + memset(out, 0, (size_t)hid * sizeof(float)); + float *ga = m->ff, *ub = ga + inter, *acc = m->e_gate; + const int lut_sz = (hid / m->vec_dim) * m->stages * m->cb_entries; + float *lut_gate = m->lut, *lut_up = lut_gate + lut_sz, *lut_down = lut_up + lut_sz; + float sg = 0.0f; + int shared_done = 0; + + /* When the cache decides, it decides per expert and not per layer. + * + * §76 took the expert-parallel path only when all ten records were + * resident, and a layer missing one went to the row split for all ten: + * thirty dispatches of rows too short to fill the pool. At a 16 GiB + * cache that was 4,723 of 10,464 layers in a 200-token run, at 1.71 ms + * against 0.69 for a whole-resident layer, and half of them were missing + * exactly one record. So the residents go first, one expert per task — + * they need no read, so holding them is no barrier — and the reads the + * hint issued for the rest run underneath. The misses follow once they + * land, with the shared expert computed in the gap, since it needs no + * record either. + * + * Both stages run through experts_staged, in equal row ranges rather + * than one task per expert, so neither a batch of ten nor a lone miss + * leaves threads idle. §82 had split the misses between rows and tasks + * at four; the row ranges beat both (LEARNED §84). + * + * The order experts are computed in is not the order they are summed + * in — each writes its own slice and the sum below runs in route order — + * so this path, the fixed batches and the serial loop are bit-identical. + * A forced WASTE_XPAR, or an explicit WASTE_XPAR_BATCH, keeps §76's + * fixed batches; WASTE_XPAR=0 the serial loop. */ + if (xpar_on < 0 && !xpar_batch_set && m->xga && K > 1 && + K <= WASTE_PF_MAX && m->cache.n_slots >= 4 * K) { + const uint8_t *recs[WASTE_PF_MAX]; + int order[WASTE_PF_MAX], later[WASTE_PF_MAX]; + int nres = 0, nmis = 0, ok = 1, lut_done = 0; + for (int j = 0; j < K; j++) { + if (waste_ecache_resident_all(&m->cache, L, idx + j, 1)) order[nres++] = j; + else later[nmis++] = j; + } + memcpy(order + nres, later, (size_t)nmis * sizeof(int)); + for (int stage = 0; stage < 2 && ok; stage++) { + const int *list = stage ? order + nres : order; + const int n = stage ? nmis : nres; + if (!n) continue; + if (stage) { sg = qwen_shared_expert(m, L, in, acc, shared_pre, sg_raw); shared_done = 1; } + PROF_START(P_EDEQ); + for (int t = 0; t < n && ok; t++) { + recs[list[t]] = waste_ecache_hold(&m->cache, L, idx[list[t]], + bank_fetch, m); + if (!recs[list[t]]) ok = 0; + } + PROF_END(P_EDEQ); + if (!ok) break; + PROF_START(P_EMM); + if (!lut_done) { + const waste_expert_hdr *h0 = (const waste_expert_hdr *)recs[list[0]]; + vq_build_lut(m, lut_gate, h0->codebook_id + 0 * m->stages, + in, hid, m->stages, m->cb_entries, m->vec_dim, + NULL, NULL); + vq_build_lut(m, lut_up, h0->codebook_id + 1 * m->stages, + in, hid, m->stages, m->cb_entries, m->vec_dim, + NULL, NULL); + lut_done = 1; + } + if (m->index_bits != 6 && !vq8_on) { + experts_staged(m, recs, list, n, inter, hid, lut_gate, lut_up); + } else { + xpar_arg pa = { m, c, recs, w, 0, inter, hid, + lut_gate, lut_up, NULL, NULL, NULL, NULL, list }; + waste_parallel_for_each(n, moe_expert_range, &pa, g_pool.nthreads); + } + PROF_END(P_EMM); + waste_ecache_release(&m->cache); + } + if (ok) goto qwen_moe_sum; + /* Something did not read: let go of what was held and fall through + * to the serial loop, which re-reads and reports the reason. It uses + * `acc` as its accumulator, so the shared expert is redone after. */ + waste_ecache_release(&m->cache); + shared_done = 0; + } + + /* Forced, or batched explicitly: §76's fixed batches, the barrier and + * all. WASTE_XPAR=0/1 forces the path either way. */ + const int xpar_here = xpar_on >= 0 + ? xpar_on + : waste_ecache_resident_all(&m->cache, L, idx, K); + if (xpar_here && m->xga && K > 1 && K <= WASTE_PF_MAX && + m->cache.n_slots >= 4 * K) { + const uint8_t *recs[WASTE_PF_MAX]; + const int batch = xpar_batch; + int lut_done = 0, ok = 1; + for (int j0 = 0; j0 < K; j0 += batch) { + int j1 = j0 + batch; + if (j1 > K) j1 = K; + PROF_START(P_EDEQ); + int n = j0; + for (; n < j1; n++) { + recs[n] = waste_ecache_hold(&m->cache, L, idx[n], bank_fetch, m); + if (!recs[n]) break; + } + PROF_END(P_EDEQ); + if (n < j1) { ok = 0; break; } + PROF_START(P_EMM); + if (!lut_done) { + const waste_expert_hdr *h0 = (const waste_expert_hdr *)recs[0]; + vq_build_lut(m, lut_gate, h0->codebook_id + 0 * m->stages, + in, hid, m->stages, m->cb_entries, m->vec_dim, + NULL, NULL); + vq_build_lut(m, lut_up, h0->codebook_id + 1 * m->stages, + in, hid, m->stages, m->cb_entries, m->vec_dim, + NULL, NULL); + lut_done = 1; + } + /* Qwen's experts read the hidden state directly: there is no + * latent projection, so `lat` is the hidden size. One expert + * per range: a batch of ten on eight threads cut as rows would + * be five ranges of two, and three threads with nothing. */ + xpar_arg pa = { m, c, recs, w, j0, inter, hid, + lut_gate, lut_up, NULL, NULL, NULL, NULL, NULL }; + waste_parallel_for_each(j1 - j0, moe_expert_range, &pa, + g_pool.nthreads); + PROF_END(P_EMM); + waste_ecache_release(&m->cache); + } + if (ok) goto qwen_moe_sum; + waste_ecache_release(&m->cache); + } + + /* ga and ub are m->ff: the shared expert's projections are about to be + * overwritten, and it redoes them after. */ + shared_pre = 0; + int lut_ready = 0; + for (int j = 0; j < K; j++) { + PROF_START(P_EDEQ); + const uint8_t *rec = read_expert(m, L, idx[j]); + PROF_END(P_EDEQ); + if (!rec) break; + PROF_START(P_EMM); + if (!lut_ready) { + const waste_expert_hdr *h = (const waste_expert_hdr *)rec; + vq_build_lut(m, lut_gate, h->codebook_id + 0 * m->stages, + in, hid, m->stages, m->cb_entries, m->vec_dim, NULL, NULL); + vq_build_lut(m, lut_up, h->codebook_id + 1 * m->stages, + in, hid, m->stages, m->cb_entries, m->vec_dim, NULL, NULL); + lut_ready = 1; + } + qwen_expert_rows(m, rec, inter, hid, lut_gate, lut_up, ga, ub, acc, lut_down); + const float wj = w[j]; + for (int i = 0; i < hid; i++) out[i] += wj * acc[i]; + PROF_END(P_EMM); + } + goto qwen_moe_shared; + +qwen_moe_sum: + /* Summed in route order, so the total does not depend on the thread + * count, the batch, or which experts were resident. */ + { + PROF_START(P_EMM); + for (int j = 0; j < K; j++) { + const float *accj = m->xacc + (size_t)j * hid; + const float wj = w[j]; + for (int i = 0; i < hid; i++) out[i] += wj * accj[i]; + } + PROF_END(P_EMM); + } +qwen_moe_shared: + /* Not m->h: qwen_step passes it as `out`, and the routed sum is already + * in there waiting for HyperConnection to consume it. The expert + * accumulator is dead once the routed experts are done and is sized for + * hid floats, so the shared expert lands there. */ + if (!shared_done) sg = qwen_shared_expert(m, L, in, acc, shared_pre, sg_raw); + { + PROF_START(P_QSHX); + for (int i = 0; i < hid; i++) out[i] += sg * acc[i]; + PROF_END(P_QSHX); + } +} + +/* Which experts layer L+1 is about to route to, asked as soon as layer L's + * MoE is back in the streams — so the reads it starts run under L+1's + * attention instead of after its router. + * + * Kimi's lookahead (§34) runs the next router on this layer's MoE input. + * On Qwen that input is one HyperConnection mix of four streams, and L+1's + * router will see a different mix of different streams. At the width used + * here it would have started 32% of L+1's cache misses early, for 0.43 + * wasted reads a layer. L+1's own MLP mix, applied to the streams as they + * now are, starts 43% for 0.09 — and costs more than the reads it saves, + * 4.6% slower end to end. This is its cheap half: each stream normalized + * with L+1's MLP mix weights and averaged, the dynamic gate left out. 43% + * for 0.23, at the price of four norms and one router projection. LEARNED + * §83 has the other predictors and the widths. + * + * `m->x`, the block output and the router area are all dead here — the + * next layer's attention mix writes the first, its attention the second — + * so the guess needs no scratch of its own. Asking earlier, straight after + * layer L routes, would have needed scratch and measured the same expert + * I/O: what is still waited for is the misses no top-6 guesses. It only chooses what is read + * early; the real router still decides, so the logits cannot move. */ +static int qwen_predict_next_moe(waste_model *m, int L, int *out, int n) +{ + const waste_config *c = &m->cfg; + const int E = c->n_experts, hid = c->hidden, hc = c->hc_count; + if (n <= 0 || L + 1 >= c->n_layers) return 0; + if (n > E) n = E; + const waste_tensor *g = waste_find(m, tname("%smodel.layers.%d.mlp.gate.weight", + c->prefix, L + 1)); + const waste_tensor *nw = waste_find(m, tname( + "%smodel.layers.%d.mlp_hyper_connection.hc_norm.weight", c->prefix, L + 1)); + if (!g || !nw || !nw->data) return 0; + float *x = m->x, *nb = m->h, *sc = m->att + WASTE_ATT_ROUTER_OFF; + memset(x, 0, (size_t)hid * sizeof(float)); + for (int b = 0; b < hc; b++) { + waste_qwen_rmsnorm(nb, m->hcx + (size_t)b * hid, + nw->data + (size_t)b * hid, hid, hid, c->eps); + for (int i = 0; i < hid; i++) x[i] += nb[i]; + } + matvec_t(m, sc, g, x, E, hid); + /* Top n by insertion: n is a handful and E is 512. */ + int k = 0; + for (int e = 0; e < E; e++) { + const float v = sc[e]; + if (!(v == v)) continue; + if (k == n && !(v > sc[out[n - 1]])) continue; + int q = k < n ? k++ : n - 1; + while (q > 0 && v > sc[out[q - 1]]) { out[q] = out[q - 1]; q--; } + out[q] = e; + } + return k; +} + +static const float *qwen_step(waste_model *m, int token, int pos, int *routed) +{ + dump_pos0 = pos; + const waste_config *c = &m->cfg; + const int hid = c->hidden, hc = c->hc_count; + { + const int cm = waste_model_ctx_max(m); + if (cm && (pos < 0 || pos >= cm)) { m->ctx_full = 1; return NULL; } + } + waste_embed_row(m, token, m->x); + for (int b = 0; b < hc; b++) + memcpy(m->hcx + (size_t)b * hid, m->x, (size_t)hid * sizeof(float)); + + float *block = m->h; + for (int L = 0; L < c->n_layers; L++) { + if (m->read_error) break; + if (L == c->ple_layer) { + PROF_START(P_QPLE); + qwen_ple_inject(m, token); + PROF_END(P_QPLE); + } + float inj[16]; + /* The braces are for PROF_START, which declares its start time: + * HyperConnection is timed in three pieces per layer, and each + * needs a scope of its own. */ + { + PROF_START(P_QHC); + qwen_hc_mix_t(m, m->hcx, NULL, NULL, + waste_find(m, tname("%smodel.layers.%d.attn_hyper_connection.hc_norm.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.attn_hyper_connection.input_mix_weight_down.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.attn_hyper_connection.input_mix_weight_up.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.attn_hyper_connection.block_inject_weight.weight", c->prefix, L)), + 1, m->x, inj); + PROF_END(P_QHC); + } + if (!c->qwen_full[L]) { + PROF_START(P_KDA); + qwen_gdn_layer(m, L, m->x, block); + PROF_END(P_KDA); + } else { + PROF_START(P_MLA); + qwen_qsa_layer(m, L, m->x, block, pos); + PROF_END(P_MLA); + } + { + PROF_START(P_QHC); + /* The attention block goes into the streams inside the mix. */ + qwen_hc_mix_t(m, m->hcx, block, inj, + waste_find(m, tname("%smodel.layers.%d.mlp_hyper_connection.hc_norm.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.mlp_hyper_connection.input_mix_weight_down.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.mlp_hyper_connection.input_mix_weight_up.weight", c->prefix, L)), + waste_find(m, tname("%smodel.layers.%d.mlp_hyper_connection.block_inject_weight.weight", c->prefix, L)), + 1, m->x, inj); + PROF_END(P_QHC); + } + { + PROF_START(P_ROUTE); + qwen_moe_layer(m, L, m->x, block, routed ? routed + (size_t)L * c->top_k : NULL); + PROF_END(P_ROUTE); + } + { + PROF_START(P_QHC); + waste_qwen_hc_combine(m->hcx, block, inj, hc, hid, m->hcx); + PROF_END(P_QHC); + } + /* The disk is about to go idle through the next layer's attention: + * start the reads its router is likely to ask for. Width 6 is + * WASTE_LOOKAHEAD's default for Kimi too, and it is where Qwen's + * curve peaked here — wider guesses read more than they save. */ + if (lookahead_n && m->cache.io && m->cache.n_slots > 0) { + PROF_START(P_ROUTE); + PROF_START(P_QLAH); + int nxt[64]; + const int nn = qwen_predict_next_moe(m, L, nxt, lookahead_n); + if (nn) waste_ecache_prefetch(&m->cache, L + 1, nxt, nn); + PROF_END(P_QLAH); + PROF_END(P_ROUTE); + } + /* Same role as the Kimi dump in waste_model_step: one residual + * stream after every layer. Qwen's stream is the hc hyper-state. */ + const char *dump_hidden = getenv("WASTE_DUMP_HIDDEN"); + if (dump_hidden) { + FILE *df = fopen(dump_hidden, (L || pos) ? "ab" : "wb"); + if (df) { + fwrite(m->hcx, sizeof(float), (size_t)hc * (size_t)hid, df); + fclose(df); + } + } + } + PROF_START(P_QHC); + qwen_hc_mix_t(m, m->hcx, NULL, NULL, + waste_find(m, tname("%smodel.hyper_connection_mixer.hc_norm.weight", c->prefix)), + waste_find(m, tname("%smodel.hyper_connection_mixer.input_mix_weight_down.weight", c->prefix)), + waste_find(m, tname("%smodel.hyper_connection_mixer.input_mix_weight_up.weight", c->prefix)), + NULL, 0, m->x, NULL); + PROF_END(P_QHC); + PROF_START(P_HEAD); + matvec_t(m, m->logits, waste_find(m, tname("%slm_head.weight", c->prefix)), m->x, + c->vocab, hid); + PROF_END(P_HEAD); + return m->read_error ? NULL : m->logits; +} + const float *waste_model_step(waste_model *m, int token, int pos, int *routed) { + if (m->cfg.arch_qwen) return qwen_step(m, token, pos, routed); dump_pos0 = pos; const waste_config *c = &m->cfg; const int hid = c->hidden; diff --git a/src/model.h b/src/model.h index 491987b3b..8dce26ce0 100644 --- a/src/model.h +++ b/src/model.h @@ -44,8 +44,27 @@ typedef struct { * ever read). q and qs are NULL in that case. */ int64_t file_off, file_scale_off; /* not long: 32 bits on Windows */ int on_disk; + /* Row in waste_tmv_roles plus one, filled on the first profiled matvec + * so the profiler looks the name up once rather than every call. */ + int prof_slot; } waste_tensor; +/* Trunk matvec by tensor role, under WASTE_PROFILE: a tensor's name with + * its layer number taken out, so all 36 of a model's in_proj_qkv are one + * row. The size buckets could not say which projection a millisecond was + * in — on Qwen, GDN's out_proj and QSA's o_proj are the same shape, and so + * are QSA's K/V projections and the router. */ +typedef struct { + char role[96]; + int out, in, bits; /* bits 32: an F32 tensor */ + uint64_t calls, bytes; + double t; + double tq; /* of t: quantizing the activation, serial */ +} waste_tmv_role; +#define WASTE_TMV_ROLES 96 +extern waste_tmv_role waste_tmv_roles[WASTE_TMV_ROLES]; +extern int waste_tmv_nroles; + typedef struct { int n_layers, hidden, n_experts, top_k, moe_inter, dense_inter; int n_shared, first_dense, vocab, n_heads; @@ -107,6 +126,7 @@ typedef struct { * branch. 1 on both Kimi models and the default; GLM's pattern has no * such branch and its containers say so. */ int tok_han_split; + int tok_digit_run; /* digits per pre-token: 3 (default) or 1 */ /* WASTE_TOKPAT_*: which pre-tokenization pattern the release splits * with. 0 = cl100k, which is every container written before * DeepSeek-V4.1 and the only one tok_han_split means anything for. */ @@ -183,6 +203,25 @@ typedef struct { char rope_err[128]; /* non-empty: a shape rope_init does not * implement, and why. The load refuses on * it rather than running unrotated. */ + + /* --- Qwen3.8-Flash-Next (all 0 / -1 for Kimi) ---------------------- */ + int arch_qwen; /* 1 = GDN/QSA/HC path, never KDA/MLA */ + int qwen_full[WASTE_MAX_LAYERS]; /* 1 = QSA (full_attention) */ + int qwen_n_layer_types; /* entries the container actually gave */ + int gdn_k_heads, gdn_v_heads, gdn_k_dim, gdn_v_dim; + int qsa_n_kv, qsa_head_dim; + int idx_n_heads, idx_kv_heads, idx_head_dim; + int idx_budget, idx_compress; + int hc_count, hc_lowrank; + int ple_layer; /* 0-based; -1 = no PLE */ + int ngram_size, heads_per_ngram, ple_embed, ple_conv_k; + int shared_inter; + int rotary_dim; + int mrope_section[3]; +#define WASTE_QWEN_PLE_HEADS 16 + int64_t ple_off[WASTE_QWEN_PLE_HEADS]; + int64_t ple_sz[WASTE_QWEN_PLE_HEADS]; + int64_t ple_mult[8]; } waste_config; typedef struct { @@ -341,6 +380,27 @@ typedef struct { * collapsed single stream the sublayer runs on. */ float *hcflat, *hccol, *hcmix; + /* Qwen residual streams and QSA caches. GDN reuses S[]/conv[] with + * Gated-DeltaNet sizes, not KDA's. */ + float *hcx; /* [hc_count * hidden] */ + uint16_t *qsa_k[WASTE_MAX_LAYERS]; /* BF16 [kv_cap][n_kv][head_dim] */ + uint16_t *qsa_v[WASTE_MAX_LAYERS]; + float *qsa_rawk[WASTE_MAX_LAYERS]; /* FP32 raw index keys [kv_cap][Dk] */ + int n_qsa_blk[WASTE_MAX_LAYERS]; + int n_qsa_tail[WASTE_MAX_LAYERS]; + int ple_prev[8]; + float *ple_ring; /* dilated PLE conv [H*(K-1)*ngram] */ + float *ple_emb; /* concatenated PLE row [ple_embed] */ + float *gdn_g; /* per-V-head decay [Hv] */ + float *qsa_q, *qsa_gate, *qsa_attn; /* [Hq][D] */ + float *qsa_kf, *qsa_vf, *qsa_scr, *qsa_work; + float *qsa_cs; /* cos then sin, each [kv_cap][rot] */ + int qsa_cs_n; /* rows of qsa_cs already filled */ + int *qsa_sel, *qsa_taken; + float *moe_prob; + uint8_t *moe_used; + int has_qsa; + /* scratch */ float *x, *h, *tmp, *att, *logits; /* 1 when at least one layer is MLA, i.e. when the sequence is bounded @@ -373,6 +433,7 @@ typedef struct { size_t lut_bytes; /* what m->lut was allocated */ int8_t *xq; uint64_t expert_reads; + uint64_t ple_reads; waste_ecache cache; /* Background fill. Runs only when the resolved cache can hold every * record the container has, walks the banks in file order, and stops on diff --git a/src/qwen_gdn.c b/src/qwen_gdn.c new file mode 100644 index 000000000..0756d878c --- /dev/null +++ b/src/qwen_gdn.c @@ -0,0 +1,105 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* qwen_gdn.c — see qwen_gdn.h. Official gated-delta recurrence, not KDA. */ + +#include "qwen_gdn.h" + +#include +#include + +static float l2_rnorm(const float *x, int n) +{ + float s = 0.0f; + for (int i = 0; i < n; i++) s += x[i] * x[i]; + return 1.0f / sqrtf(s + 1e-6f); +} + +void waste_qwen_gdn_decay(const float *a, const float *A_log, const float *dt, + int Hv, float *g) +{ + for (int h = 0; h < Hv; h++) { + const float z = a[h] + dt[h]; + const float sp = logf(1.0f + expf(-fabsf(z))) + (z > 0.0f ? z : 0.0f); + g[h] = -expf(A_log[h]) * sp; + } +} + +static void gdn_heads(int h0, int h1, int Dk, int Dv, int group, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *u) +{ + const float qscale = 1.0f / sqrtf((float)Dk); + for (int h = h0; h < h1; h++) { + const int src = h / group; + const float *qh = q + (size_t)src * Dk; + const float *kh = k + (size_t)src * Dk; + const float *vh = v + (size_t)h * Dv; + float *Sh = S + (size_t)h * Dk * Dv; + float *oh = o + (size_t)h * Dv; + const float qn = l2_rnorm(qh, Dk) * qscale; + const float kn = l2_rnorm(kh, Dk); + const float decay = expf(g_log[h]); + const float b = beta[h]; + + memset(u, 0, (size_t)Dv * sizeof(float)); + for (int kk = 0; kk < Dk; kk++) { + float *row = Sh + (size_t)kk * Dv; + const float kv = kh[kk] * kn; + for (int i = 0; i < Dv; i++) { row[i] *= decay; u[i] += row[i] * kv; } + } + for (int i = 0; i < Dv; i++) u[i] = b * (vh[i] - u[i]); + memset(oh, 0, (size_t)Dv * sizeof(float)); + for (int kk = 0; kk < Dk; kk++) { + float *row = Sh + (size_t)kk * Dv; + const float kv = kh[kk] * kn; + const float qv = qh[kk] * qn; + for (int i = 0; i < Dv; i++) { row[i] += u[i] * kv; oh[i] += row[i] * qv; } + } + } +} + +void waste_qwen_gdn_step(int Hk, int Hv, int Dk, int Dv, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch) +{ + waste_qwen_gdn_step_heads(0, Hv, Hk, Hv, Dk, Dv, q, k, v, g_log, beta, + S, o, scratch); +} + +void waste_qwen_gdn_step_heads(int h0, int h1, int Hk, int Hv, int Dk, int Dv, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch) +{ + const int group = Hk > 0 ? Hv / Hk : 1; + gdn_heads(h0, h1, Dk, Dv, group > 0 ? group : 1, q, k, v, g_log, beta, + S, o, scratch); +} + +void waste_qwen_gdn_forward(int T, int Hk, int Hv, int Dk, int Dv, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch) +{ + for (int t = 0; t < T; t++) { + waste_qwen_gdn_step(Hk, Hv, Dk, Dv, + q + (size_t)t * Hk * Dk, k + (size_t)t * Hk * Dk, + v + (size_t)t * Hv * Dv, g_log + (size_t)t * Hv, + beta + (size_t)t * Hv, + S, o + (size_t)t * Hv * Dv, scratch); + } +} + +/* Portable prefill: sequential recurrence. Official GPU prefill calls + * torch_chunk_gated_delta_rule; the two match at fp32 for official geometry. */ +void waste_qwen_gdn_chunk(int T, int Hk, int Hv, int Dk, int Dv, int chunk, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch) +{ + (void)chunk; + waste_qwen_gdn_forward(T, Hk, Hv, Dk, Dv, q, k, v, g_log, beta, S, o, scratch); +} diff --git a/src/qwen_gdn.h b/src/qwen_gdn.h new file mode 100644 index 000000000..c8fe28b24 --- /dev/null +++ b/src/qwen_gdn.h @@ -0,0 +1,62 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * qwen_gdn.h — Qwen Gated DeltaNet. Not KDA. + * + * Official decode uses torch_recurrent_gated_delta_rule; official prefill + * uses torch_chunk_gated_delta_rule. The portable CPU path is the sequential + * recurrence. waste_qwen_gdn_chunk is that recurrence (chunk_size ignored). + * Persistent state is two buffers: + * 1. recurrent S [Hv][Dk][Dv] + * 2. QKV short-conv ring [(2*Hk+Hv)*D][K-1] + * QK-repeat (16 QK heads onto 48 V heads) is an identity, not a buffer. + * g is a per-head scalar, not a per-K-dim diagonal. + */ + +#ifndef WASTE_QWEN_GDN_H +#define WASTE_QWEN_GDN_H + +#ifdef __cplusplus +extern "C" { +#endif + +/* g = -exp(A_log) * softplus(a + dt_bias), length Hv. */ +void waste_qwen_gdn_decay(const float *a, const float *A_log, const float *dt, + int Hv, float *g); + +/* One decode step. q,k are [Hk][Dk] (repeated onto Hv inside), v [Hv][Dv], + * g_log and beta [Hv], S [Hv][Dk][Dv], o [Hv][Dv]. + * scratch >= Dv. */ +void waste_qwen_gdn_step(int Hk, int Hv, int Dk, int Dv, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch); + +/* Value heads [h0, h1) of one decode step; waste_qwen_gdn_step is this over + * [0, Hv). A head reads its own rows of v, g_log, beta and S (and the QK head + * it is repeated from) and writes only its own rows of S and o, so disjoint + * ranges may run at once — each with its own scratch of >= Dv. */ +void waste_qwen_gdn_step_heads(int h0, int h1, int Hk, int Hv, int Dk, int Dv, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch); + +/* T decode steps, time-major inputs [T][H][*]. */ +void waste_qwen_gdn_forward(int T, int Hk, int Hv, int Dk, int Dv, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch); + +/* Sequential recurrence over T (portable prefill). chunk_size is ignored. + * Equal to waste_qwen_gdn_forward; proven against official chunk/recurrent + * kernels at 16 QK / 48 V / dim 128. */ +void waste_qwen_gdn_chunk(int T, int Hk, int Hv, int Dk, int Dv, int chunk, + const float *q, const float *k, const float *v, + const float *g_log, const float *beta, + float *S, float *o, float *scratch); + +#ifdef __cplusplus +} +#endif +#endif /* WASTE_QWEN_GDN_H */ diff --git a/src/qwen_hc.c b/src/qwen_hc.c new file mode 100644 index 000000000..2ee0e9cde --- /dev/null +++ b/src/qwen_hc.c @@ -0,0 +1,86 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* qwen_hc.c — see qwen_hc.h. */ + +#include "qwen_hc.h" + +#include +#include + +static float silu(float x) { return x / (1.0f + expf(-x)); } +static float sigmoid(float x) { return 1.0f / (1.0f + expf(-x)); } + +static void matvec(float *y, const float *W, const float *x, int out, int in) +{ + for (int i = 0; i < out; i++) { + float s = 0.0f; + const float *w = W + (size_t)i * in; + for (int j = 0; j < in; j++) s += w[j] * x[j]; + y[i] = s; + } +} + +void waste_qwen_rmsnorm(float *o, const float *x, const float *w, + int n, int group, float eps) +{ + if (group <= 0 || group > n) group = n; + const int ng = n / group; + for (int g = 0; g < ng; g++) { + const float *xg = x + (size_t)g * group; + float *og = o + (size_t)g * group; + const float *wg = w + (size_t)g * group; + float s = 0.0f; + for (int i = 0; i < group; i++) s += xg[i] * xg[i]; + const float r = 1.0f / sqrtf(s / (float)group + eps); + for (int i = 0; i < group; i++) + og[i] = xg[i] * r * (1.0f + wg[i]); + } +} + +void waste_qwen_hc_mix(const float *hyper, const float *norm_w, + const float *down, const float *up, + int hc, int hid, int rank, float eps, + float *mixed, float *scratch) +{ + const int H = hc * hid; + float *normed = scratch; + float *lo = scratch + H; + float *gate = scratch + H + rank; + waste_qwen_rmsnorm(normed, hyper, norm_w, H, hid, eps); + matvec(lo, down, normed, rank, H); + for (int i = 0; i < rank; i++) lo[i] = silu(lo[i] / (float)hc); + matvec(gate, up, lo, H, rank); + for (int i = 0; i < H; i++) gate[i] = sigmoid(gate[i]); + for (int d = 0; d < hid; d++) { + float s = 0.0f; + for (int b = 0; b < hc; b++) { + const int i = b * hid + d; + s += gate[i] * normed[i]; + } + mixed[d] = s / (float)hc; + } +} + +void waste_qwen_hc_gates(const float *hyper, const float *norm_w, + const float *down, const float *up, + const float *inject, int hc, int hid, int rank, + float eps, float *mixed, float *inj_w, float *scratch) +{ + const int H = hc * hid; + waste_qwen_hc_mix(hyper, norm_w, down, up, hc, hid, rank, eps, mixed, scratch); + float *normed = scratch; + float *tmp = scratch + H + rank + H; + waste_qwen_rmsnorm(normed, hyper, norm_w, H, hid, eps); + matvec(tmp, inject, normed, hc, H); + for (int b = 0; b < hc; b++) + inj_w[b] = 2.0f * sigmoid(tmp[b] / (float)hc); +} + +void waste_qwen_hc_combine(const float *hyper, const float *block, + const float *inj_w, int hc, int hid, float *out) +{ + for (int b = 0; b < hc; b++) + for (int d = 0; d < hid; d++) + out[b * hid + d] = hyper[b * hid + d] + inj_w[b] * block[d]; +} diff --git a/src/qwen_hc.h b/src/qwen_hc.h new file mode 100644 index 000000000..b2820325d --- /dev/null +++ b/src/qwen_hc.h @@ -0,0 +1,44 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * qwen_hc.h — Qwen Gated Residual (HyperConnection Mix / Combine). + * + * Distinct from Kimi AttnRes. Mix reads four residual branches through a + * rank-320 bottleneck; Combine writes a per-branch scalar gate times the + * block output back onto the original streams. + */ + +#ifndef WASTE_QWEN_HC_H +#define WASTE_QWEN_HC_H + +#ifdef __cplusplus +extern "C" { +#endif + +/* Grouped RMSNorm: (1 + weight) * x * rsqrt(mean(x^2)+eps) per group. + * Official Qwen4ExpTextRMSNorm. */ +void waste_qwen_rmsnorm(float *o, const float *x, const float *w, + int n, int group, float eps); + +/* Mix: hyper [hc*hid] -> mixed [hid]. scratch >= hc*hid + rank + hc*hid. */ +void waste_qwen_hc_mix(const float *hyper, const float *norm_w, + const float *down, const float *up, + int hc, int hid, int rank, float eps, + float *mixed, float *scratch); + +/* Combine gates: same Mix, plus injection_weights [hc] = + * 2 * sigmoid(inject(normed) / hc). scratch as Mix plus hc. */ +void waste_qwen_hc_gates(const float *hyper, const float *norm_w, + const float *down, const float *up, + const float *inject, int hc, int hid, int rank, + float eps, float *mixed, float *inj_w, float *scratch); + +/* out = hyper + inj_w[b] * block[hid] for each branch b. */ +void waste_qwen_hc_combine(const float *hyper, const float *block, + const float *inj_w, int hc, int hid, float *out); + +#ifdef __cplusplus +} +#endif +#endif /* WASTE_QWEN_HC_H */ diff --git a/src/qwen_moe.c b/src/qwen_moe.c new file mode 100644 index 000000000..73d61b8d0 --- /dev/null +++ b/src/qwen_moe.c @@ -0,0 +1,44 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* qwen_moe.c — see qwen_moe.h. */ + +#include "qwen_moe.h" + +#include + +int waste_qwen_moe_route(const float *logits, int E, int K, int renorm, + int *idx, float *w, float *prob, uint8_t *used) +{ + if (idx && w && K > 0) + for (int j = 0; j < K; j++) { idx[j] = 0; w[j] = 0.0f; } + if (!logits || !idx || !w || !prob || !used || E < 1 || K < 1 || K > E) + return -1; + float m = logits[0]; + for (int e = 1; e < E; e++) if (logits[e] > m) m = logits[e]; + float z = 0.0f; + for (int e = 0; e < E; e++) { prob[e] = expf(logits[e] - m); z += prob[e]; } + if (z < 1e-20f) z = 1e-20f; + for (int e = 0; e < E; e++) { prob[e] /= z; used[e] = 0; } + /* K passes over E rather than a sort: K is 10 and E is 512, and this + * way the tie rule (lowest id wins, because `>` is strict) is the same + * on every platform. */ + for (int j = 0; j < K; j++) { + int best = -1; + float bv = -1.0f; + for (int e = 0; e < E; e++) { + if (used[e]) continue; + if (prob[e] > bv) { bv = prob[e]; best = e; } + } + idx[j] = best >= 0 ? best : 0; + w[j] = best >= 0 ? prob[best] : 0.0f; + if (best >= 0) used[best] = 1; + } + if (renorm && K > 1) { + float s = 0.0f; + for (int j = 0; j < K; j++) s += w[j]; + if (s > 1e-20f) + for (int j = 0; j < K; j++) w[j] /= s; + } + return 0; +} diff --git a/src/qwen_moe.h b/src/qwen_moe.h new file mode 100644 index 000000000..ab8eca4f2 --- /dev/null +++ b/src/qwen_moe.h @@ -0,0 +1,42 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * qwen_moe.h — Qwen4ExpTextTopKRouter. + * + * Softmax over every expert, then top-k, then an optional renormalize — + * not the sigmoid-plus-bias router the Kimi and GLM families use, and not + * interchangeable with it: softmax couples the experts, so the same + * logits give a different selection under the two. + */ + +#ifndef WASTE_QWEN_MOE_H +#define WASTE_QWEN_MOE_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +/* Top-k over `E` softmax probabilities. + * + * `idx` and `w` are length K and come back in *selection* order — highest + * probability first — and not sorted by expert id. That order is part of + * the contract: the MoE reduction accumulates in it, and floating-point + * addition is not associative, so re-ordering the experts changes the last + * bits of every token. Ties go to the lower expert id, which is what makes + * the order reproducible at all. + * + * `prob` is E floats of scratch and `used` E bytes; both are written. + * `idx`/`w` are always written (zeros when the arguments are unusable), so + * a caller that ignores the return value routes to expert 0 with weight 0 + * rather than reading uninitialized memory. Returns 0, or -1. + */ +int waste_qwen_moe_route(const float *logits, int E, int K, int renorm, + int *idx, float *w, float *prob, uint8_t *used); + +#ifdef __cplusplus +} +#endif +#endif /* WASTE_QWEN_MOE_H */ diff --git a/src/qwen_ple.c b/src/qwen_ple.c new file mode 100644 index 000000000..2e7cc7947 --- /dev/null +++ b/src/qwen_ple.c @@ -0,0 +1,107 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* qwen_ple.c — see qwen_ple.h. */ + +#include "qwen_ple.h" + +#include + +#define MASK64 ((uint64_t)-1) +#define SPLITMIX_GAMMA 0x9E3779B97F4A7C15ULL +#define SPLITMIX_M1 0xBF58476D1CE4E5B9ULL +#define SPLITMIX_M2 0x94D049BB133111EBULL +#define PRIME_1 10007 + +uint64_t waste_qwen_splitmix64(uint64_t value) +{ + value = (value + SPLITMIX_GAMMA) & MASK64; + value = ((value ^ (value >> 30)) * SPLITMIX_M1) & MASK64; + value = ((value ^ (value >> 27)) * SPLITMIX_M2) & MASK64; + return (value ^ (value >> 31)) & MASK64; +} + +void waste_qwen_ple_multipliers(int64_t *out, int ngram_size, int ple_layer_index, + int64_t seed, int unigram_vocab) +{ + const int64_t max_long = (int64_t)(((uint64_t)1 << 63) - 1); + const int64_t multiplier_max = max_long / (unigram_vocab > 0 ? unigram_vocab : 1); + int64_t half_bound = multiplier_max / 2; + if (half_bound < 1) half_bound = 1; + const uint64_t base_seed = (uint64_t)seed + (uint64_t)PRIME_1 * (uint64_t)ple_layer_index; + for (int i = 0; i < ngram_size; i++) { + const uint64_t value = (base_seed + SPLITMIX_GAMMA * (uint64_t)(i + 1)) & MASK64; + out[i] = 2 * (int64_t)(waste_qwen_splitmix64(value) % (uint64_t)half_bound) + 1; + } +} + +int waste_qwen_ple_is_prime(int64_t value) +{ + if (value < 2) return 0; + if (value % 2 == 0) return value == 2; + for (int64_t d = 3; d * d <= value; d += 2) + if (value % d == 0) return 0; + return 1; +} + +int64_t waste_qwen_ple_nth_prime_after(int64_t start, int count) +{ + int64_t prime = start; + for (int i = 0; i < count; i++) { + prime++; + while (!waste_qwen_ple_is_prime(prime)) prime++; + } + return prime; +} + +void waste_qwen_ple_shift_eos(const int *ids, int n, int shift, int eos, int *out) +{ + if (shift == 0) { + for (int i = 0; i < n; i++) out[i] = ids[i]; + return; + } + int prev_eos = -1; + int incl = -1; + for (int i = 0; i < n; i++) { + const int segment_start = prev_eos + 1; + const int pos_in_seg = i - segment_start; + const int src = i - shift; + const int valid = (pos_in_seg >= shift) && (src >= 0); + out[i] = valid ? ids[src] : eos; + if (ids[i] == eos) incl = i; + prev_eos = incl; + } +} + +void waste_qwen_ple_row_ids(const int *ids, int n, int pos, int eos, + int ngram_size, int heads_per_ngram, + const int64_t *multipliers, + const int64_t *head_vocab_sizes, + int *local_rows) +{ + const int heads = (ngram_size - 1) * heads_per_ngram; + int h0; + for (h0 = 0; h0 < heads && h0 < WASTE_QWEN_PLE_HEADS; h0++) + local_rows[h0] = 0; + if (!ids || !multipliers || !head_vocab_sizes || !local_rows) return; + if (ngram_size < 1 || ngram_size > 8 || n < 1 || n > 8 || + pos < 0 || pos >= n || heads_per_ngram < 1) + return; + int hist[8 * 8]; + for (int s = 0; s < ngram_size; s++) + waste_qwen_ple_shift_eos(ids, n, s, eos, hist + (size_t)s * n); + + int h = 0; + for (int ngram = 2; ngram <= ngram_size; ngram++) { + int64_t mixed = (int64_t)hist[0 * n + pos] * multipliers[0]; + for (int p = 1; p < ngram; p++) + mixed ^= (int64_t)hist[p * n + pos] * multipliers[p]; + for (int k = 0; k < heads_per_ngram && h < WASTE_QWEN_PLE_HEADS; k++, h++) { + const int64_t sz = head_vocab_sizes[h]; + if (sz <= 0) { local_rows[h] = 0; continue; } + int64_t r = mixed % sz; + if (r < 0) r += sz; + local_rows[h] = (int)r; + } + } +} diff --git a/src/qwen_ple.h b/src/qwen_ple.h new file mode 100644 index 000000000..bff8d35a0 --- /dev/null +++ b/src/qwen_ple.h @@ -0,0 +1,49 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * qwen_ple.h — Qwen n-gram hashing for Per-Layer Embedding. + * + * Equations are from transformers Qwen4ExpTextNGramEmbedding: splitmix + * multipliers, EOS-bounded shifts, XOR mix, then remainder into 16 head + * tables. Each head lives on disk as Q8G; this module never asks for the + * BF16 source table. + */ + +#ifndef WASTE_QWEN_PLE_H +#define WASTE_QWEN_PLE_H + +#include + +#ifdef __cplusplus +extern "C" { +#endif + +#define WASTE_QWEN_PLE_HEADS 16 + +uint64_t waste_qwen_splitmix64(uint64_t value); + +/* Official _build_layer_multipliers. out[ngram_size] signed odd integers. */ +void waste_qwen_ple_multipliers(int64_t *out, int ngram_size, int ple_layer_index, + int64_t seed, int unigram_vocab); + +int waste_qwen_ple_is_prime(int64_t value); +int64_t waste_qwen_ple_nth_prime_after(int64_t start, int count); + +/* Shift the sequence right by `shift`, filling from EOS and never crossing + * a previous EOS. Official _shift_right_ignore_eos. */ +void waste_qwen_ple_shift_eos(const int *ids, int n, int shift, int eos, int *out); + +/* 16 local row ids for the token at `pos` (0-based in `ids[0..n)`). + * sizes[h] is that head's prime vocab; local row is remainder, not the + * global offset used by a packed table. */ +void waste_qwen_ple_row_ids(const int *ids, int n, int pos, int eos, + int ngram_size, int heads_per_ngram, + const int64_t *multipliers, + const int64_t *head_vocab_sizes, + int *local_rows); + +#ifdef __cplusplus +} +#endif +#endif /* WASTE_QWEN_PLE_H */ diff --git a/src/qwen_qsa.c b/src/qwen_qsa.c new file mode 100644 index 000000000..fb1d9dfda --- /dev/null +++ b/src/qwen_qsa.c @@ -0,0 +1,281 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* qwen_qsa.c — see qwen_qsa.h. Official QSA indexer, not MLA. */ + +#include "qwen_qsa.h" + +#include +#include + +#if defined(__ARM_NEON) || defined(__aarch64__) +#include +#endif + +void waste_qwen_mrope_interleave(const float *freqs_t, const float *freqs_h, + const float *freqs_w, const int *section, + int half, float *out) +{ + memcpy(out, freqs_t, (size_t)half * sizeof(float)); + const float *fh[3] = { freqs_t, freqs_h, freqs_w }; + for (int dim = 1; dim <= 2; dim++) { + const int length = section[dim] * 3; + for (int i = dim; i < length && i < half; i += 3) + out[i] = fh[dim][i]; + } +} + +int waste_qwen_rope_apply(float *x, int dim, const float *cos, const float *sin, + int rotary_dim) +{ + if (rotary_dim > dim) rotary_dim = dim; + if (rotary_dim < 0) rotary_dim = 0; + if (rotary_dim > 256) return -1; + const int half = rotary_dim / 2; + float tmp[256]; + memcpy(tmp, x, (size_t)rotary_dim * sizeof(float)); + for (int i = 0; i < rotary_dim; i++) { + const float rh = (i < half) ? -tmp[i + half] : tmp[i - half]; + x[i] = tmp[i] * cos[i] + rh * sin[i]; + } + (void)dim; + return 0; +} + +static void rmsnorm_1p(float *o, const float *x, const float *w, int n, float eps) +{ + float s = 0.0f; + for (int i = 0; i < n; i++) s += x[i] * x[i]; + const float r = 1.0f / sqrtf(s / (float)n + eps); + for (int i = 0; i < n; i++) o[i] = x[i] * r * (1.0f + w[i]); +} + +void waste_qwen_qsa_pool_block(const float *raw_k, int compress, int Dk, + const float *k_ln_w, float eps, float *out) +{ + if (!out || Dk < 1) return; + memset(out, 0, (size_t)Dk * sizeof(float)); + if (!raw_k || !k_ln_w || compress < 1) return; + for (int t = 0; t < compress; t++) { + const float *row = raw_k + (size_t)t * Dk; + for (int d = 0; d < Dk; d++) out[d] += row[d]; + } + for (int d = 0; d < Dk; d++) out[d] /= (float)compress; + rmsnorm_1p(out, out, k_ln_w, Dk, eps); +} + +int waste_qwen_qsa_select(const float *q_heads, int Hq, int Dk, + const float *raw_k, int T, int query_pos, + const float *full_cos, const float *full_sin, + int rotary_dim, const float *k_ln_w, float eps, + int compress, int block_topk, + int *sel, float *work, int *taken) +{ + if (!sel || compress < 1 || Dk < 1) return 0; + if (query_pos < 0) query_pos = 0; + if (T < 1) return 0; + if (query_pos >= T) query_pos = T - 1; + const int vis = query_pos + 1; + const int n_complete = vis / compress; + const int n_tail = vis - n_complete * compress; + int nsel = 0; + + if (n_complete > 0) { + if (!work || !taken || !q_heads || !raw_k || !k_ln_w) return 0; + waste_qwen_qsa_score_blocks(0, n_complete, q_heads, Hq, Dk, raw_k, + full_cos, full_sin, rotary_dim, k_ln_w, eps, + compress, work, work + (size_t)n_complete * Dk); + } + (void)nsel; + return waste_qwen_qsa_pick(n_complete > 0 ? work + (size_t)n_complete * Dk : NULL, + n_complete, block_topk, compress, n_tail, sel, taken); +} + +void waste_qwen_qsa_score_blocks(int b0, int b1, const float *q_heads, int Hq, + int Dk, const float *raw_k, + const float *full_cos, const float *full_sin, + int rotary_dim, const float *k_ln_w, float eps, + int compress, float *pooled, float *scores) +{ + const float inv_sqrt = 1.0f / sqrtf((float)Dk); + for (int b = b0; b < b1; b++) { + float *po = pooled + (size_t)b * Dk; + waste_qwen_qsa_pool_block(raw_k + (size_t)b * compress * Dk, + compress, Dk, k_ln_w, eps, po); + if (full_cos && full_sin && rotary_dim > 0) + waste_qwen_rope_apply(po, Dk, + full_cos + (size_t)(b * compress) * rotary_dim, + full_sin + (size_t)(b * compress) * rotary_dim, + rotary_dim); + float s = 0.0f; + for (int h = 0; h < Hq; h++) { + float dot = 0.0f; + const float *qh = q_heads + (size_t)h * Dk; + for (int d = 0; d < Dk; d++) dot += qh[d] * po[d]; + if (dot < 0.0f) dot = 0.0f; + s += dot; + } + scores[b] = s * inv_sqrt; + } +} + +/* Whether block a comes before block b in the selection: the higher score, + * and on a tie the earlier block. That is the order the repeated argmax this + * replaced produced — it took the first strictly greater score each pass. */ +static int qsa_before(const float *s, int a, int b) +{ + return s[a] > s[b] || (s[a] == s[b] && a < b); +} + +static void qsa_sift(const float *s, int *o, int i, int n) +{ + for (;;) { + int w = i; + const int l = 2 * i + 1, r = l + 1; + if (l < n && qsa_before(s, o[w], o[l])) w = l; + if (r < n && qsa_before(s, o[w], o[r])) w = r; + if (w == i) return; + const int tmp = o[i]; o[i] = o[w]; o[w] = tmp; + i = w; + } +} + +int waste_qwen_qsa_pick(const float *scores, int n_complete, int block_topk, + int compress, int n_tail, int *sel, int *order) +{ + int nsel = 0; + if (n_complete > 0 && scores && order) { + /* The argmax only ever took a score above -1e30, which also leaves + * out a NaN; the ordering below is total over what remains. */ + int n = 0; + for (int b = 0; b < n_complete; b++) + if (scores[b] > -1e30f) order[n++] = b; + /* Heapsort with the latest block in selection order at the root: + * each pass moves the worst remaining block to the end, so the + * array finishes best first — O(n log n), where the argmax was a + * pass over every block for every block kept. */ + for (int i = n / 2 - 1; i >= 0; i--) qsa_sift(scores, order, i, n); + for (int end = n - 1; end > 0; end--) { + const int tmp = order[0]; order[0] = order[end]; order[end] = tmp; + qsa_sift(scores, order, 0, end); + } + const int keep = n_complete < block_topk ? n_complete : block_topk; + const int take = keep < n ? keep : n; + for (int j = 0; j < take; j++) + for (int t = 0; t < compress; t++) + sel[nsel++] = order[j] * compress + t; + } + for (int t = 0; t < n_tail; t++) + sel[nsel++] = n_complete * compress + t; + return nsel; +} + +void waste_qwen_qsa_attn(const float *q, int Hq, int D, + const float *k, const float *v, int Hkv, int T, + const int *sel, int n_sel, float scale, + float *out, float *scratch) +{ + waste_qwen_qsa_attn_heads(0, Hq, q, Hq, D, k, v, Hkv, T, sel, n_sel, scale, + out, scratch); +} + +/* Four selected tokens' scores at a time, and the value sum over lanes. + * + * Both halves of this are a dot product 256 wide, and the first was costing + * what a dependent chain of fused multiply-adds costs: one element per + * ~4 cycles, whatever the machine could otherwise issue. Four tokens have + * four independent chains, and each still sums its own dimensions in the + * order it did — the same trick, and the same reason, as the VQ gather's + * four rows (LEARNED §41). The value accumulation is the other way round: + * every output dimension sums the selected tokens in order, so the lanes + * run along `d` and each element's sequence is untouched. Both are bit + * for bit what the scalar loops produced; §87 has the check. + */ +void waste_qwen_qsa_attn_heads(int h0, int h1, const float *q, int Hq, int D, + const float *k, const float *v, int Hkv, int T, + const int *sel, int n_sel, float scale, + float *out, float *scratch) +{ + const int n_rep = Hkv > 0 ? Hq / Hkv : 1; + float *scores = scratch; + if (h1 > h0) + memset(out + (size_t)h0 * D, 0, (size_t)(h1 - h0) * D * sizeof(float)); + if (!q || !k || !v || !sel || !scratch || n_sel < 1) return; + for (int h = h0; h < h1; h++) { + const int hv = h / (n_rep > 0 ? n_rep : 1); + const float *qh = q + (size_t)h * D; + float m = -1e30f; + int i = 0; + /* Four at a time only when there are enough to pay for it: a short + * context selects a handful, and there the plain loop below is + * what measured faster. */ + const int n4 = n_sel >= 32 ? n_sel : 0; + for (; i + 4 <= n4; i += 4) { + const int t0 = sel[i], t1 = sel[i + 1], t2 = sel[i + 2], t3 = sel[i + 3]; + if (t0 < 0 || t0 >= T || t1 < 0 || t1 >= T || + t2 < 0 || t2 >= T || t3 < 0 || t3 >= T) break; + const float *k0 = k + ((size_t)t0 * Hkv + hv) * D; + const float *k1 = k + ((size_t)t1 * Hkv + hv) * D; + const float *k2 = k + ((size_t)t2 * Hkv + hv) * D; + const float *k3 = k + ((size_t)t3 * Hkv + hv) * D; + float s0 = 0.0f, s1 = 0.0f, s2 = 0.0f, s3 = 0.0f; + /* The product and the sum are separate statements on purpose. + * The loop this replaces compiles to four products a time in + * one vector and a scalar chain of adds — the products are + * independent, the order of the adds is not — so each product + * is rounded on its own. Written as `s += q * k` the compiler + * may contract the pair into one fused multiply-add, which + * rounds once instead of twice and is a different number. + * Across a 2,048-token selection that difference moves a + * logit. */ + for (int d = 0; d < D; d++) { + const float qd = qh[d]; + const float p0 = qd * k0[d], p1 = qd * k1[d]; + const float p2 = qd * k2[d], p3 = qd * k3[d]; + s0 = s0 + p0; s1 = s1 + p1; + s2 = s2 + p2; s3 = s3 + p3; + } + s0 *= scale; s1 *= scale; s2 *= scale; s3 *= scale; + scores[i] = s0; scores[i + 1] = s1; + scores[i + 2] = s2; scores[i + 3] = s3; + if (s0 > m) m = s0; + if (s1 > m) m = s1; + if (s2 > m) m = s2; + if (s3 > m) m = s3; + } + for (; i < n_sel; i++) { + const int t = sel[i]; + if (t < 0 || t >= T) { scores[i] = -1e30f; continue; } + const float *kh = k + ((size_t)t * Hkv + hv) * D; + float s = 0.0f; + for (int d = 0; d < D; d++) s += qh[d] * kh[d]; + s *= scale; + scores[i] = s; + if (s > m) m = s; + } + float z = 0.0f; + for (int i = 0; i < n_sel; i++) { + scores[i] = expf(scores[i] - m); + z += scores[i]; + } + if (z < 1e-20f) z = 1e-20f; + float *oh = out + (size_t)h * D; + for (int j = 0; j < n_sel; j++) { + const int t = sel[j]; + if (t < 0 || t >= T) continue; + const float w = scores[j] / z; + const float *vh = v + ((size_t)t * Hkv + hv) * D; + int d = 0; +#if defined(__ARM_NEON) || defined(__aarch64__) + const float32x4_t wv = vdupq_n_f32(w); + for (; d + 16 <= D; d += 16) { + vst1q_f32(oh + d, vfmaq_f32(vld1q_f32(oh + d), wv, vld1q_f32(vh + d))); + vst1q_f32(oh + d + 4, vfmaq_f32(vld1q_f32(oh + d + 4), wv, vld1q_f32(vh + d + 4))); + vst1q_f32(oh + d + 8, vfmaq_f32(vld1q_f32(oh + d + 8), wv, vld1q_f32(vh + d + 8))); + vst1q_f32(oh + d + 12, vfmaq_f32(vld1q_f32(oh + d + 12), wv, vld1q_f32(vh + d + 12))); + } +#endif + for (; d < D; d++) oh[d] += w * vh[d]; + } + } +} diff --git a/src/qwen_qsa.h b/src/qwen_qsa.h new file mode 100644 index 000000000..fa0db20b9 --- /dev/null +++ b/src/qwen_qsa.h @@ -0,0 +1,84 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * qwen_qsa.h — Qwen Sparse Attention. Not MLA. + * + * Indexer: FP32 mean of four keys, MRoPE, ReLU-sum scores, top-512 blocks, + * plus the 0–3 tail tokens, then attention over the original K/V. + * Production decode calls waste_qwen_qsa_select; there is one pooling path. + */ + +#ifndef WASTE_QWEN_QSA_H +#define WASTE_QWEN_QSA_H + +#ifdef __cplusplus +extern "C" { +#endif + +/* Interleaved MRoPE: freqs [3][half] -> out [half]. section is 3 ints. */ +void waste_qwen_mrope_interleave(const float *freqs_t, const float *freqs_h, + const float *freqs_w, const int *section, + int half, float *out); + +/* Apply RoPE to the first rotary_dim components (cos/sin length). + * rotary_dim > 256 is refused (no write). */ +int waste_qwen_rope_apply(float *x, int dim, const float *cos, const float *sin, + int rotary_dim); + +/* Mean of `compress` consecutive raw keys, then RMSNorm (1+w). + * raw_k is [compress][Dk]. */ +void waste_qwen_qsa_pool_block(const float *raw_k, int compress, int Dk, + const float *k_ln_w, float eps, float *out); + +/* Select tokens for one query. raw_k is [T][Dk] (one indexer KV head). + * Writes up to budget+compress-1 ids into sel, returns the count. + * q_heads is [Hq][Dk] already layernormed and RoPE'd. + * work >= n_complete*Dk + n_complete + Dk floats; taken >= n_complete ints. + * work/taken may be NULL when n_complete == 0. */ +int waste_qwen_qsa_select(const float *q_heads, int Hq, int Dk, + const float *raw_k, int T, int query_pos, + const float *full_cos, const float *full_sin, + int rotary_dim, const float *k_ln_w, float eps, + int compress, int block_topk, + int *sel, float *work, int *taken); + +/* The two halves of waste_qwen_qsa_select, which is score_blocks over every + * complete block followed by pick. + * + * score_blocks pools, rotates and scores blocks [b0, b1): block b writes only + * pooled[b][Dk] and scores[b], so disjoint ranges may run at once — the + * select call lays pooled at work and scores at work + n_complete * Dk. + * + * pick writes the selection in order: complete blocks by score, highest + * first and a tie to the earlier block, at most block_topk of them, each as + * its compress token indices; then the n_tail tail tokens. order needs room + * for n_complete ints. Returns the count. */ +void waste_qwen_qsa_score_blocks(int b0, int b1, const float *q_heads, int Hq, + int Dk, const float *raw_k, + const float *full_cos, const float *full_sin, + int rotary_dim, const float *k_ln_w, float eps, + int compress, float *pooled, float *scores); +int waste_qwen_qsa_pick(const float *scores, int n_complete, int block_topk, + int compress, int n_tail, int *sel, int *order); + +/* Causal softmax attention over selected tokens. q [Hq][D], k/v [T][Hkv][D], + * n_rep = Hq/Hkv, scaling = 1/sqrt(D). sel[n_sel] are token indices. */ +void waste_qwen_qsa_attn(const float *q, int Hq, int D, + const float *k, const float *v, int Hkv, int T, + const int *sel, int n_sel, float scale, + float *out, float *scratch); + +/* Query heads [h0, h1) of the same attention; waste_qwen_qsa_attn is this + * over [0, Hq). A head reads its own query row and its KV head's keys and + * values and writes only its own row of out, so disjoint ranges may run at + * once — each with its own scratch of >= n_sel. */ +void waste_qwen_qsa_attn_heads(int h0, int h1, const float *q, int Hq, int D, + const float *k, const float *v, int Hkv, int T, + const int *sel, int n_sel, float scale, + float *out, float *scratch); + +#ifdef __cplusplus +} +#endif +#endif /* WASTE_QWEN_QSA_H */ diff --git a/src/threads.h b/src/threads.h index 5280691cc..cc5d9ed5e 100644 --- a/src/threads.h +++ b/src/threads.h @@ -319,6 +319,9 @@ static inline int waste_pool_fast(void) * wants every core for the kernel next to it. So the count is per call * site now, and the two available answers are the pool and the fast * group. Results do not depend on it: the split is by row either way. */ +static inline void waste__pool_run(int n, int chunk, waste_range_fn fn, + void *arg, int workers); + static inline void waste_parallel_for_n(int n, int min_chunk, waste_range_fn fn, void *arg, int workers) { @@ -334,7 +337,30 @@ static inline void waste_parallel_for_n(int n, int min_chunk, waste_range_fn fn, /* Round up to a whole number of min_chunk units: callers that block * their data (the VQ tile) need every range to start on a boundary. */ chunk = ((chunk + min_chunk - 1) / min_chunk) * min_chunk; + waste__pool_run(n, chunk, fn, arg, workers); +} +/* One item per range, however many items there are per worker. + * + * waste_parallel_for_n cuts n into `workers` equal ranges, which is right for + * rows and wrong for a handful of large, uneven tasks: ten routed experts on + * eight threads came out as five ranges of two, so three threads never got + * an expert while five did two each. Here every participant takes the next + * item as it finishes the last, so all of them work until the queue is empty. + * Only for items that are each worth a dispatch on their own. */ +static inline void waste_parallel_for_each(int n, waste_range_fn fn, void *arg, + int workers) +{ + waste__bind_self(); + if (workers > g_pool.nthreads) workers = g_pool.nthreads; + if (workers > n) workers = n; + if (workers <= 1) { fn(0, n, arg); return; } + waste__pool_run(n, 1, fn, arg, workers); +} + +static inline void waste__pool_run(int n, int chunk, waste_range_fn fn, + void *arg, int workers) +{ /* Keep the descriptor stable until every worker has left this job. * Distinct waste_ctx instances may be called concurrently even though * they reuse this process-wide pool. */ diff --git a/src/tokenizer.c b/src/tokenizer.c index 9d9f80128..d1eb7b47e 100644 --- a/src/tokenizer.c +++ b/src/tokenizer.c @@ -44,6 +44,12 @@ struct waste_tok { * and one without — which is exactly the kind of text that appears in * a Chinese release's own prompts. */ int han_split; + /* How many digits one pre-token may take: `\p{N}{1,3}` on the Kimi and + * GLM patterns, `\p{N}` on Qwen's, which splits every digit into its + * own piece. Not cosmetic — "2026" is one token under a 3-digit run + * and four under a 1-digit one, and the difference shows up nowhere as + * an error, only as a model reading numbers it was never trained on. */ + int digit_run; /* WASTE_TOKPAT_*. han_split is read only by the cl100k scanner; the * DeepSeek one isolates CJK unconditionally because its own pattern * does, in a dedicated Split that runs before the main one. */ @@ -212,7 +218,7 @@ waste_tok *waste_tok_open(const char *dir) /* The Kimi pattern until a container says otherwise: every model here * before GLM has the Han branch, and a default that has to be set to * keep working is a default that will be missed. */ - if (t) t->han_split = 1; + if (t) { t->han_split = 1; t->digit_run = 3; } if (!t) { free(raw); return NULL; } t->blob = (uint8_t *)malloc((size_t)sz); /* decoded is smaller */ t->cap_tokens = 4096; @@ -305,6 +311,14 @@ void waste_tok_set_pattern(waste_tok *t, int pattern) if (t && pattern >= 0 && pattern < WASTE_TOKPAT__COUNT) t->pattern = pattern; } +void waste_tok_set_digit_run(waste_tok *t, int n) +{ + /* A container that states something outside what any of these patterns + * spell is a container this cannot honour; keep the default rather + * than invent a third behaviour. */ + if (t && (n == 1 || n == 3)) t->digit_run = n; +} + /* ---- UTF-8 + the character classes the pattern needs -------------------- */ static int utf8_next(const char *s, int len, int *cp) @@ -376,7 +390,7 @@ static int is_space(int c) /* Advances one pre-token, returning its byte length. Mirrors the branch * order of the model's pat_str. */ -static int next_piece(const char *s, int len, int han_split) +static int next_piece(const char *s, int len, int han_split, int digit_run) { int cp, n = utf8_next(s, len, &cp), i; if (n == 0) return 0; @@ -422,11 +436,11 @@ static int next_piece(const char *s, int len, int han_split) return i; } - if (is_digit(cp)) { /* \p{N}{1,3} */ + if (is_digit(cp)) { /* \p{N}{1,digit_run} */ i = n; int cnt = 1; - while (i < len && cnt < 3) { int c3, k = utf8_next(s + i, len - i, &c3); - if (!is_digit(c3)) break; i += k; cnt++; } + while (i < len && cnt < digit_run) { int c3, k = utf8_next(s + i, len - i, &c3); + if (!is_digit(c3)) break; i += k; cnt++; } return i; } @@ -664,7 +678,7 @@ static int next_piece_ds(const char *s, int len) static int next_piece_pat(const waste_tok *t, const char *s, int len) { if (t->pattern == WASTE_TOKPAT_DEEPSEEK) return next_piece_ds(s, len); - return next_piece(s, len, t->han_split); + return next_piece(s, len, t->han_split, t->digit_run); } /* ---- byte-pair merge ---------------------------------------------------- */ diff --git a/src/tokenizer.h b/src/tokenizer.h index 19eebf3fa..d5f603d5f 100644 --- a/src/tokenizer.h +++ b/src/tokenizer.h @@ -44,6 +44,14 @@ void waste_tok_set_eos(waste_tok *t, int id); * Han run touches a Latin one, and they differ silently. */ void waste_tok_set_han_split(waste_tok *t, int on); +/* Digits per pre-token: 3 (the default) for `\p{N}{1,3}`, as the Kimi and + * GLM patterns spell it; 1 for `\p{N}`, which is Qwen's and puts every + * digit in its own piece. Anything else is ignored and the default kept — + * there is no third spelling in this family. "2026" is one token under 3 + * and four under 1, so a container that gets this wrong reads every + * number differently from the model that was trained on it. */ +void waste_tok_set_digit_run(waste_tok *t, int n); + /* Which pre-tokenization pattern the release splits with. The Kimi and GLM * containers are all cl100k-shaped and differ only in the Han branch, which * is why that one is a flag; DeepSeek-V4.1 splits with a different pattern diff --git a/src/waste.c b/src/waste.c index 09c94ccb0..377243a81 100644 --- a/src/waste.c +++ b/src/waste.c @@ -340,7 +340,8 @@ waste_status waste_plan_memory(const char *model_path, uint32_t ctx_tokens, char nm[160]; js_str(&d, js_get(&d, e, "name"), nm, sizeof nm); const int fmt = (int)js_int(&d, js_get(&d, e, "fmt"), 0); - if (fmt != 0 && strstr(nm, "embed_tokens.weight")) continue; + if (fmt != 0 && (strstr(nm, "embed_tokens.weight") || + strstr(nm, "ngram_head."))) continue; const uint64_t nb = (uint64_t)js_int(&d, js_get(&d, e, "bytes"), 0); /* The tower is loaded only when a caller asks for images, so it is * counted apart and folded in by waste_open — counting it here @@ -426,14 +427,51 @@ waste_status waste_plan_memory(const char *model_path, uint32_t ctx_tokens, const int kl = js_get(&d, lac, "kda_layers"); const int n_kda = js_size(&d, kl); const int n_mla = layers - n_kda; + char mt[48]; + js_str(&d, js_get(&d, cfg, "model_type"), mt, sizeof mt); + const int is_qwen = (strcmp(mt, "qwen4_exp_text") == 0); /* MLA caches the latent plus the rope dims, not the expanded per-head * K and V — kv_b_proj is absorbed into the query and the output. That - * is 576 floats per token per layer here rather than 30720. */ - out->state_bytes = (uint64_t)n_kda * kh * kd * kd * 4 /* S */ - + (uint64_t)n_kda * 3 * (ck - 1) * kh * kd * 4 /* conv */ - + (uint64_t)n_mla * ctx_tokens * - ((uint64_t)kv_lora + qk_rope) * 4; /* KV */ + * is 576 floats per token per layer here rather than 30720. + * Qwen is not MLA: GDN S is [Hv][Dk][Dv], QSA keeps BF16 K/V plus + * every raw FP32 index key so pooling uses one tested path. */ + if (is_qwen) { + const int Hk = (int)js_int(&d, js_get(&d, cfg, "linear_num_key_heads"), 0); + const int Hv = (int)js_int(&d, js_get(&d, cfg, "linear_num_value_heads"), 0); + const int Dk = (int)js_int(&d, js_get(&d, cfg, "linear_key_head_dim"), 0); + const int Dv = (int)js_int(&d, js_get(&d, cfg, "linear_value_head_dim"), 0); + const int gck = (int)js_int(&d, js_get(&d, cfg, "linear_conv_kernel_dim"), 4); + const int nkv = (int)js_int(&d, js_get(&d, cfg, "num_key_value_heads"), 0); + const int hd = (int)js_int(&d, js_get(&d, cfg, "head_dim"), 0); + const int idim = (int)js_int(&d, js_get(&d, cfg, "indexer_head_dim"), 128); + const int hc = (int)js_int(&d, js_get(&d, cfg, "hc_count"), 4); + const int ngram = (int)js_int(&d, js_get(&d, cfg, "ngram_size"), 3); + const int pck = (int)js_int(&d, js_get(&d, cfg, "ple_conv_kernel_size"), 4); + const int lt = js_get(&d, cfg, "layer_types"); + int n_gdn = 0, n_qsa = 0; + for (int i = 0; i < js_size(&d, lt); i++) { + char kind[32]; + js_str(&d, js_at(&d, lt, i), kind, sizeof kind); + if (strcmp(kind, "full_attention") == 0) n_qsa++; + else n_gdn++; + } + const int qkv = 2 * Hk * Dk + Hv * Dv; + const int R = (pck > 1 && ngram > 0) ? (pck - 1) * ngram : 0; + out->state_bytes = + (uint64_t)n_gdn * (uint64_t)Hv * Dk * Dv * 4u + + (uint64_t)n_gdn * (uint64_t)qkv * (gck > 0 ? gck - 1 : 0) * 4u + + (uint64_t)n_qsa * ctx_tokens * (uint64_t)nkv * hd * 2u * 2u + + (uint64_t)n_qsa * ctx_tokens * (uint64_t)idim * 4u + + (uint64_t)hc * hidden * 4u + + (uint64_t)hc * hidden * (uint64_t)R * 4u; + (void)n_kda; (void)n_mla; (void)kh; (void)kd; (void)ck; + } else { + out->state_bytes = (uint64_t)n_kda * kh * kd * kd * 4 /* S */ + + (uint64_t)n_kda * 3 * (ck - 1) * kh * kd * 4 /* conv */ + + (uint64_t)n_mla * ctx_tokens * + ((uint64_t)kv_lora + qk_rope) * 4; /* KV */ + } (void)qk_nope; (void)v_head; { /* The DSA indexer's pooled keys are session state too: one * index_dim vector per index_kpool tokens per full-attention layer. @@ -468,7 +506,21 @@ waste_status waste_plan_memory(const char *model_path, uint32_t ctx_tokens, js_free(&d); free(src); return WASTE_E_FORMAT; } const int nb = ares ? layers / ares + 2 : 1; - const int big = hidden > kh * kd ? hidden : kh * kd; + int big = hidden > kh * kd ? hidden : kh * kd; + if (is_qwen) { + const int hc = (int)js_int(&d, js_get(&d, cfg, "hc_count"), 4); + const int Hk = (int)js_int(&d, js_get(&d, cfg, "linear_num_key_heads"), 0); + const int Hv = (int)js_int(&d, js_get(&d, cfg, "linear_num_value_heads"), 0); + const int Dk = (int)js_int(&d, js_get(&d, cfg, "linear_key_head_dim"), 0); + const int Dv = (int)js_int(&d, js_get(&d, cfg, "linear_value_head_dim"), 0); + const int hd = (int)js_int(&d, js_get(&d, cfg, "head_dim"), 0); + const int qkv = 2 * Hk * Dk + Hv * Dv; + const int hcH = hc * hidden; + const int qsa = nheads * hd * 2; + if (qkv > big) big = qkv; + if (hcH > big) big = hcH; + if (qsa > big) big = qsa; + } const int wide = hidden > lat ? hidden : lat; int lut_wide = wide > moe_inter ? wide : moe_inter; const int T = WASTE_CHUNK_MAX; @@ -483,8 +535,11 @@ waste_status waste_plan_memory(const char *model_path, uint32_t ctx_tokens, const int n_exp = (int)js_int(&d, js_get(&d, cfg, "num_experts"), 0); uint64_t att = (uint64_t)ctx_tokens * (uint64_t)nheads; const uint64_t kda = (uint64_t)kh * (uint64_t)kd; + const uint64_t gdn = (uint64_t)js_int(&d, js_get(&d, cfg, "linear_num_value_heads"), 0) * + (uint64_t)js_int(&d, js_get(&d, cfg, "linear_key_head_dim"), 0); const uint64_t route = WASTE_ATT_ROUTER_OFF + 2ull * (uint64_t)n_exp; if (kda > att) att = kda; + if (gdn > att) att = gdn; if (route > att) att = route; sc += (att + 1024) * 4; } @@ -539,11 +594,16 @@ waste_status waste_plan_memory(const char *model_path, uint32_t ctx_tokens, * up, accumulator and down LUT per routed expert, because k threads * each run a whole expert. Read here rather than at the bottom of * this function, where top_k is fetched for the cache floor. */ - const uint64_t k = (uint64_t)js_int(&d, js_get(&d, cfg, - "num_experts_per_token"), 8); - sc += k * ((uint64_t)2 * moe_inter + lat) * 4; /* xga/xub/xacc */ - sc += k * lut * 4; /* m->xlut */ - sc += k * (lut + nsc * 4); /* xlut8/xqs */ + /* The converter normalises this key, but a hand-written manifest + * may spell it the way HF does; fall back to that before the + * historical default rather than planning for the wrong count. */ + uint64_t kt = (uint64_t)js_int(&d, js_get(&d, cfg, + "num_experts_per_token"), 0); + if (!kt) kt = (uint64_t)js_int(&d, js_get(&d, cfg, + "num_experts_per_tok"), 8); + sc += kt * ((uint64_t)2 * moe_inter + lat) * 4; /* xga/xub/xacc */ + sc += kt * lut * 4; /* m->xlut */ + sc += kt * (lut + nsc * 4); /* xlut8/xqs */ } sc += ((uint64_t)T * (2 * moe_inter * n_shared_eff + hidden) + 64) * 4; sc += (uint64_t)T * (2 * lat + 2 * hidden) * 4; @@ -551,10 +611,39 @@ waste_status waste_plan_memory(const char *model_path, uint32_t ctx_tokens, sc += (uint64_t)3 * moe_inter * lat * 4; /* one expert */ sc += (uint64_t)T * nb * hidden * 4 + (uint64_t)T * hidden * 4; sc += (uint64_t)T * 64 * 12; /* croute + crw + cused */ + if (is_qwen) { + const int pe = (int)js_int(&d, js_get(&d, cfg, "ple_embed_dim"), hidden); + const int Hd = (int)js_int(&d, js_get(&d, cfg, "head_dim"), 0); + const int nkv = (int)js_int(&d, js_get(&d, cfg, "num_key_value_heads"), 0); + const int idim = (int)js_int(&d, js_get(&d, cfg, "indexer_head_dim"), 128); + const int compress = (int)js_int(&d, js_get(&d, cfg, "indexer_compress_ratio"), 4); + const int budget = (int)js_int(&d, js_get(&d, cfg, "indexer_budget"), 2048); + const int Hv = (int)js_int(&d, js_get(&d, cfg, "linear_num_value_heads"), 0); + const int n_exp = (int)js_int(&d, js_get(&d, cfg, "num_experts"), 0); + const int nblk = compress > 0 + ? (int)((ctx_tokens + (uint32_t)compress - 1u) / (uint32_t)compress) : 0; + const int max_sel = budget + compress; + const int rot = (int)(Hd * js_num(&d, js_get(&d, + js_get(&d, cfg, "rope_parameters"), "partial_rotary_factor"), + js_num(&d, js_get(&d, cfg, "partial_rotary_factor"), 0.25))); + sc += (uint64_t)(pe > 0 ? pe : hidden) * 4u; + sc += (uint64_t)(Hv > 0 ? Hv : 1) * 4u; + sc += (uint64_t)nheads * Hd * 3u * 4u; + sc += (uint64_t)max_sel * nkv * Hd * 2u * 4u; + /* attention scores, one row per query head (model.c qsa_scr) */ + sc += (uint64_t)max_sel * (uint64_t)(nheads > 0 ? nheads : 1) * 4u; + sc += (uint64_t)max_sel * 4u; + sc += ((uint64_t)nblk * idim + (uint64_t)nblk + (uint64_t)idim) * 4u; + sc += (uint64_t)nblk * 4u; + sc += 2ull * ctx_tokens * (uint64_t)(rot > 0 ? rot : 1) * 4u; + sc += (uint64_t)(n_exp > 0 ? n_exp : 1) * 5u; + } out->scratch_bytes = sc; /* one layer's top-k experts, double buffered */ - const int top_k = (int)js_int(&d, js_get(&d, cfg, "num_experts_per_token"), 8); + int top_k = (int)js_int(&d, js_get(&d, cfg, "num_experts_per_token"), 0); + if (!top_k) + top_k = (int)js_int(&d, js_get(&d, cfg, "num_experts_per_tok"), 8); const int lyr = js_get(&d, 0, "layers"); uint64_t rec = 0, bank_total = 0; if (js_size(&d, lyr) > 0) { @@ -795,6 +884,7 @@ waste_status waste_open(const char *model_path, const waste_cfg *cfg_in, if (c->tok) { waste_tok_set_eos(c->tok, c->m.cfg.eos_token_id); waste_tok_set_han_split(c->tok, c->m.cfg.tok_han_split); + waste_tok_set_digit_run(c->tok, c->m.cfg.tok_digit_run); waste_tok_set_pattern(c->tok, c->m.cfg.tok_pattern); } /* warm the cache from what previous runs learned, if anything */ @@ -1352,6 +1442,7 @@ waste_status waste_model_get_info(const waste_ctx *c, waste_model_info *out) out->arch = strstr(cf->arch, "KimiK3") ? "kimi-k3" : strstr(cf->arch, "KimiLinear") ? "kimi-linear" : strstr(cf->arch, "Glm5Next") ? "glm5-next" + : strstr(cf->arch, "Qwen4Exp") ? "qwen4_exp_text" : cf->arch[0] ? cf->arch : "unknown"; out->quant_summary = c->quant; diff --git a/tests/kernel_kl.c b/tests/kernel_kl.c new file mode 100644 index 000000000..7358acc0c --- /dev/null +++ b/tests/kernel_kl.c @@ -0,0 +1,282 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * kernel_kl.c — trunk kernels side by side, scored at every position. + * + * kernel_kl CONTAINER IDS_FILE n_gen [kernel_a=0] [kernels_b=2] [window=512] + * + * Kernels are WASTE_TRUNK_KERNEL's numbers: 0 f32, 1 SDOT, 2 i8mm, 3 SMLAL. + * kernels_b is a comma list, each scored against kernel_a. + * + * sweep.c scores a kernel against a stored reference, which keeps one + * vocab-sized logit vector per position and so stops at 512 of them. The + * question that limit cannot answer is the one that killed SDOT on K3: + * whether a small per-matvec error grows as a recurrence carries it, and + * whether a sparse-attention selection that only starts choosing past a + * few thousand tokens starts choosing differently. So here the container is + * loaded once per kernel and every copy steps through the same tokens, with + * the trunk kernel switched between them, each position scored as it goes: + * KL(a||b), argmax agreement, top-10 overlap, the logits' relative L2 — and, + * because a top-K router turns a small arithmetic difference into a + * discrete one, how many of each layer's routed experts the two agreed on. + * + * The prompt is teacher-forced by construction. After it, every copy is fed + * kernel a's greedy tokens, so generation is scored the same way rather than + * by where two free-running continuations happen to part. + * + * Run it with kernel_b equal to kernel_a first: every column must come out + * exactly zero, or the copies are sharing state and nothing else it prints + * means anything. + */ +#include +#include +#include +#include +#include + +#include "../src/model.h" + +#define MAX_B 4 + +static double now(void) +{ + struct timespec t; + clock_gettime(CLOCK_MONOTONIC, &t); + return t.tv_sec + t.tv_nsec / 1e9; +} + +typedef struct { + int n, argmax_same, top10, kl_max_pos, n_nll; + double kl, kl_max, rel; + /* Next-token NLL of the real text under each kernel. KL says how far + * two distributions are apart; this says whether the one further from + * the reference is any worse at the text, which is the trade actually + * being decided. Prompt positions only — past the prompt, the "real" + * next token is kernel a's own choice. */ + double nll_a, nll_b; + long long exp_same, exp_tot, layers_diff, layers_tot; +} window; + +static void top10(const float *x, int V, int *idx) +{ + float val[10]; + int k = 0; + for (int v = 0; v < V; v++) { + if (k == 10 && x[v] <= val[9]) continue; + int j = k < 10 ? k++ : 9; + while (j > 0 && val[j - 1] < x[v]) { + val[j] = val[j - 1]; idx[j] = idx[j - 1]; j--; + } + val[j] = x[v]; idx[j] = v; + } +} + +/* a's argmax, for the next generated token. */ +static int argmax(const float *a, int V) +{ + int best = 0; + for (int v = 1; v < V; v++) if (a[v] > a[best]) best = v; + return best; +} + +static void score(const float *a, const float *b, int V, int pos, int target, + window *w) +{ + double ma = a[0], mb = b[0]; + int aa = 0, ab = 0; + for (int v = 1; v < V; v++) { + if (a[v] > ma) { ma = a[v]; aa = v; } + if (b[v] > mb) { mb = b[v]; ab = v; } + } + double sa = 0, sb = 0, num = 0, den = 0; + for (int v = 0; v < V; v++) { + sa += exp(a[v] - ma); + sb += exp(b[v] - mb); + const double d = (double)b[v] - a[v]; + num += d * d; + den += (double)a[v] * a[v]; + } + const double la = ma + log(sa), lb = mb + log(sb); + double kl = 0; + for (int v = 0; v < V; v++) { + const double lpa = a[v] - la, pa = exp(lpa); + if (pa > 1e-12) kl += pa * (lpa - (b[v] - lb)); + } + if (kl < 0) kl = 0; /* a sum of ~1e-12 terms can round below zero */ + if (target >= 0 && target < V) { + w->nll_a += la - a[target]; + w->nll_b += lb - b[target]; + w->n_nll++; + } + + int ta[10], tb[10], same = 0; + top10(a, V, ta); + top10(b, V, tb); + for (int i = 0; i < 10; i++) + for (int j = 0; j < 10; j++) + if (ta[i] == tb[j]) { same++; break; } + + w->n++; + w->kl += kl; + if (kl > w->kl_max) { w->kl_max = kl; w->kl_max_pos = pos; } + w->rel += sqrt(num / (den > 0 ? den : 1)); + w->argmax_same += aa == ab; + w->top10 += same; +} + +static void routes(const int *ra, const int *rb, int L, int K, window *w) +{ + for (int l = 0; l < L; l++) { + const int *pa = ra + (size_t)l * K, *pb = rb + (size_t)l * K; + int same = 0; + for (int u = 0; u < K; u++) + for (int v = 0; v < K; v++) + if (pa[u] == pb[v]) { same++; break; } + w->exp_same += same; + w->exp_tot += K; + w->layers_tot++; + if (same < K) w->layers_diff++; + } +} + +static void add(window *dst, const window *src) +{ + if (src->kl_max > dst->kl_max) { dst->kl_max = src->kl_max; dst->kl_max_pos = src->kl_max_pos; } + dst->n += src->n; dst->kl += src->kl; dst->rel += src->rel; + dst->n_nll += src->n_nll; dst->nll_a += src->nll_a; dst->nll_b += src->nll_b; + dst->argmax_same += src->argmax_same; dst->top10 += src->top10; + dst->exp_same += src->exp_same; dst->exp_tot += src->exp_tot; + dst->layers_diff += src->layers_diff; dst->layers_tot += src->layers_tot; +} + +static void report(const char *label, int kb, const window *w) +{ + if (!w->n) return; + printf("%-20s k%d KL mean %.2e max %.2e @%-5d | argmax %5d/%-5d | top10 %5.1f%% | " + "relL2 %.2e | experts %6.2f%%, layers differing %lld/%lld", + label, kb, w->kl / w->n, w->kl_max, w->kl_max_pos, w->argmax_same, w->n, + 10.0 * w->top10 / w->n, w->rel / w->n, + w->exp_tot ? 100.0 * w->exp_same / w->exp_tot : 0.0, + w->layers_diff, w->layers_tot); + if (w->n_nll) { + const double pa = exp(w->nll_a / w->n_nll), pb = exp(w->nll_b / w->n_nll); + printf(" | ppl %.3f -> %.3f (%+.2f%%)", pa, pb, 100.0 * (pb / pa - 1.0)); + } + printf("\n"); +} + +int main(int argc, char **argv) +{ + if (argc < 4) { + fprintf(stderr, "usage: %s CONTAINER IDS_FILE n_gen [kernel_a=0] " + "[kernels_b=2] [window=512]\n", argv[0]); + return 2; + } + const int n_gen = atoi(argv[3]); + const int ka = argc > 4 ? atoi(argv[4]) : 0; + int kb[MAX_B], nb = 0; + { + char list[64]; + snprintf(list, sizeof list, "%s", argc > 5 ? argv[5] : "2"); + for (char *p = strtok(list, ","); p && nb < MAX_B; p = strtok(NULL, ",")) + kb[nb++] = atoi(p); + } + const int win = argc > 6 && atoi(argv[6]) > 0 ? atoi(argv[6]) : 512; + const char *sgs = getenv("WASTE_SDOT4_SG"); + const int sg = sgs ? atoi(sgs) : 32; + + FILE *f = fopen(argv[2], "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", argv[2]); return 1; } + enum { MAXIDS = 32768 }; + static int ids[MAXIDS]; + int n = 0; + for (int c, v = 0, in = 0; ; ) { + c = fgetc(f); + if (c >= '0' && c <= '9') { v = v * 10 + (c - '0'); in = 1; continue; } + if (in && n < MAXIDS) ids[n++] = v; + v = 0; in = 0; + if (c == EOF) break; + } + fclose(f); + if (n < 1 || n_gen < 0 || nb < 1) { fprintf(stderr, "nothing to run\n"); return 1; } + + waste_load_opts lo; + memset(&lo, 0, sizeof lo); + const char *cmb = getenv("WASTE_CACHE_MB"); + lo.cache_bytes = (size_t)(cmb ? atoi(cmb) : 0) << 20; + lo.direct_io = 1; + /* The same floor test_forward and sweep load with: a container may + * refuse a context shorter than its own attention geometry needs. */ + int kv = n + n_gen + 16; + if (kv < 4096) kv = 4096; + waste_model *ma = (waste_model *)calloc(1, sizeof *ma); + waste_model *mb[MAX_B]; + double t0 = now(); + if (!ma || waste_model_load(ma, argv[1], kv, &lo)) { + fprintf(stderr, "load failed\n"); + return 1; + } + for (int i = 0; i < nb; i++) { + mb[i] = (waste_model *)calloc(1, sizeof *mb[i]); + if (!mb[i] || waste_model_load(mb[i], argv[1], kv, &lo)) { + fprintf(stderr, "load failed\n"); + return 1; + } + } + const int V = ma->cfg.vocab, L = ma->cfg.n_layers, K = ma->cfg.top_k; + float *A = (float *)malloc((size_t)V * sizeof(float)); + int *ra = (int *)malloc((size_t)L * K * sizeof(int)); + int *rb = (int *)malloc((size_t)L * K * sizeof(int)); + if (!A || !ra || !rb) { fprintf(stderr, "out of memory\n"); return 1; } + printf("kernel a %d vs", ka); + for (int i = 0; i < nb; i++) printf(" %d", kb[i]); + printf(": %d prompt + %d generated, windows of %d; %d loads in %.1fs\n\n", + n, n_gen, win, nb + 1, now() - t0); + fflush(stdout); + + window w[MAX_B], prompt[MAX_B], gen[MAX_B]; + memset(w, 0, sizeof w); memset(prompt, 0, sizeof prompt); memset(gen, 0, sizeof gen); + int cur = 0, w0 = 0; + t0 = now(); + for (int pos = 0; pos < n + n_gen; pos++) { + const int tok = pos < n ? ids[pos] : cur; + waste_model_set_sdot4(ka, sg); + const float *la = waste_model_step(ma, tok, pos, ra); + if (!la) { fprintf(stderr, "kernel a step %d failed\n", pos); return 1; } + memcpy(A, la, (size_t)V * sizeof(float)); + cur = argmax(A, V); + for (int i = 0; i < nb; i++) { + waste_model_set_sdot4(kb[i], sg); + const float *lb = waste_model_step(mb[i], tok, pos, rb); + if (!lb) { fprintf(stderr, "kernel %d step %d failed\n", kb[i], pos); return 1; } + score(A, lb, V, pos, pos + 1 < n ? ids[pos + 1] : -1, &w[i]); + routes(ra, rb, L, K, &w[i]); + } + + const int end = pos + 1; + if (end % win == 0 || end == n || end == n + n_gen) { + char label[64]; + snprintf(label, sizeof label, "%s %5d-%-5d", pos < n ? "prompt" : "gen ", + w0, pos); + for (int i = 0; i < nb; i++) { + report(label, kb[i], &w[i]); + add(pos < n ? &prompt[i] : &gen[i], &w[i]); + memset(&w[i], 0, sizeof w[i]); + } + fflush(stdout); + w0 = end; + fprintf(stderr, " %d/%d positions, %.0fs\n", end, n + n_gen, now() - t0); + } + } + printf("\n"); + for (int i = 0; i < nb; i++) { + report("prompt, all", kb[i], &prompt[i]); + report("generated, all", kb[i], &gen[i]); + add(&prompt[i], &gen[i]); + report("everything", kb[i], &prompt[i]); + } + waste_model_free(ma); + for (int i = 0; i < nb; i++) waste_model_free(mb[i]); + return 0; +} diff --git a/tests/run.sh b/tests/run.sh index 13b95e412..f2f4adcfb 100755 --- a/tests/run.sh +++ b/tests/run.sh @@ -1651,6 +1651,267 @@ PYD fi fi +head_ "Qwen3.8-Flash-Next (GDN, QSA, HyperConnection, PLE)" + +# A Qwen container is not a Kimi container with different numbers: the +# recurrence is Gated DeltaNet rather than KDA, attention is sparse over +# the original K/V rather than over a latent, the residual is four streams, +# and one layer reads an n-gram embedding a row at a time off the trunk. +# None of that is reachable from any other fixture, so this builds its own +# — a few hundred kilobytes, seed-deterministic, format v0 with unchanged +# WEXP records. +QWENC="$TMP/qwen.waste" +if ! python3 tools/make_test_container.py --qwen "$QWENC" >/dev/null 2>&1; then + sk "Qwen checks" "make_test_container.py --qwen did not build a container" +else + # Format first: a fixture that quietly stopped being a real container + # would make every check below vacuous. + qman=$(python3 -c "import json;m=json.load(open('$QWENC/manifest.json'));print(m['format_version'], m['arch'])" 2>/dev/null) + qmagic=$(python3 -c "import struct;print(struct.unpack('/dev/null) + qinfo=$(./waste info "$QWENC" 2>&1) + ./test_forward "$QWENC" 3,7,11 "$TMP/qwen_seq.bin" 0 >"$TMP/qwen_fwd.log" 2>&1 + if [ "$qman" = "0 qwen4_exp_text" ] && [ "$qmagic" = "True" ] && + printf '%s' "$qinfo" | grep -q "qwen4_exp_text" && + [ -s "$TMP/qwen_seq.bin" ]; then + ok "the Qwen fixture is format v0 with WEXP records and loads as qwen4_exp_text" + else + no "the Qwen fixture did not load (manifest='$qman' wexp=$qmagic)" + printf '%s\n' "$qinfo" | tail -5 + fi + + # QSA pools four keys into a block. Three tokens leave the block open + # and only the tail; the fourth closes it. Reported by test_forward so + # the boundary is observable rather than inferred from the logits. + q3=$(./test_forward "$QWENC" 3,7,11 /dev/null 0 2>&1 | grep '^qsa_layer') + q4=$(./test_forward "$QWENC" 3,7,11,5 /dev/null 0 2>&1 | grep '^qsa_layer') + if printf '%s' "$q3" | grep -q 'blk 0 tail 3' && + printf '%s' "$q4" | grep -q 'blk 1 tail 0'; then + ok "QSA closes a 4-token block on the fourth token" + else + no "QSA block pooling is wrong (3 tok: '$q3'; 4 tok: '$q4')" + fi + + # Chunked prefill against sequential decode, the check that has caught + # every state bug in this engine: the two share no code above the layer + # loop and must agree bit for bit. + WASTE_CHUNK=1 ./test_forward "$QWENC" 3,7,11 "$TMP/qwen_chunk.bin" 0 \ + >/dev/null 2>&1 + if [ ! -s "$TMP/qwen_chunk.bin" ]; then + no "chunked prefill did not run on a Qwen container" + elif cmp -s "$TMP/qwen_seq.bin" "$TMP/qwen_chunk.bin"; then + ok "chunked prefill is bit-identical to sequential decode" + else + no "Qwen chunked prefill disagrees with sequential decode" + fi + + # Which way a Qwen layer's routed experts are scheduled — the row + # split, a batch of four, or every one of them in a single dispatch, + # which is what the default does once the cache holds them — is not a + # numerical choice, and must not become one: the answer would then + # depend on how warm the cache happened to be. + # + # The cache is set on purpose. test_forward's default is no cache at + # all, and the expert-parallel path needs four slots per routed expert, + # so without one every arm below runs the row split and the comparison + # is of a path against itself — which is how this check first passed. + # One MB is 256 slots on the fixture, its whole bank, preloaded. + # + # Preloaded is also what keeps the default from ever meeting a miss: + # every layer finds its experts resident and only the first stage of + # qwen_moe_layer's staged path runs. The cold arm turns the preload off, + # so the first tokens hold residents and misses in two stages — five + # layers of the fixture take the second one, as the row split. + # + # Cold is also the only state in which the router lookahead reads + # anything, so it gets a cold arm of its own with the lookahead off: a + # guess decides when bytes move and never what is multiplied. + QXIDS=3,7,11,5,3,7,11,5,3,7,11,5 + WASTE_CACHE_MB=1 ./test_forward "$QWENC" "$QXIDS" "$TMP/qwen_xdef.bin" 0 >/dev/null 2>&1 + WASTE_CACHE_MB=1 WASTE_PRELOAD=0 ./test_forward "$QWENC" "$QXIDS" "$TMP/qwen_xcold.bin" 0 >/dev/null 2>&1 + WASTE_CACHE_MB=1 WASTE_PRELOAD=0 WASTE_LOOKAHEAD=0 ./test_forward "$QWENC" "$QXIDS" "$TMP/qwen_xnolook.bin" 0 >/dev/null 2>&1 + WASTE_CACHE_MB=1 WASTE_XPAR=0 ./test_forward "$QWENC" "$QXIDS" "$TMP/qwen_xrows.bin" 0 >/dev/null 2>&1 + WASTE_CACHE_MB=1 WASTE_XPAR=1 WASTE_XPAR_BATCH=4 ./test_forward "$QWENC" "$QXIDS" "$TMP/qwen_x4.bin" 0 >/dev/null 2>&1 + WASTE_CACHE_MB=1 WASTE_XPAR=1 WASTE_XPAR_BATCH=64 ./test_forward "$QWENC" "$QXIDS" "$TMP/qwen_xall.bin" 0 >/dev/null 2>&1 + if [ ! -s "$TMP/qwen_xdef.bin" ]; then + no "the Qwen expert-schedule comparison did not run" + elif cmp -s "$TMP/qwen_xdef.bin" "$TMP/qwen_xrows.bin" && + cmp -s "$TMP/qwen_xdef.bin" "$TMP/qwen_xcold.bin" && + cmp -s "$TMP/qwen_xdef.bin" "$TMP/qwen_xnolook.bin" && + cmp -s "$TMP/qwen_xdef.bin" "$TMP/qwen_x4.bin" && + cmp -s "$TMP/qwen_xdef.bin" "$TMP/qwen_xall.bin"; then + ok "Qwen's row split, a batch of four, one dispatch per layer, a cold cache and no lookahead give the default's logits" + else + no "a Qwen expert schedule changes the logits" + fi + + # The hyper-state dump is what the container-native oracle diffs + # against, so its shape is checked on its own: a dump of the wrong + # length would make that comparison read the wrong layer. + if [ -n "$PY_MISS" ]; then + sk "Qwen hyper-state dump" "$PY_MISS" + elif python3 tests/test_qwen_dump.py >/dev/null 2>&1; then + ok "WASTE_DUMP_HIDDEN writes all four residual streams after every layer" + else + no "WASTE_DUMP_HIDDEN is missing or the wrong size on a Qwen container" + fi + + # A budget under the floor is refused rather than swapped into, and the + # floor is computed from Qwen's own state keys — GDN's S, the conv + # rings, the BF16 K/V and the raw index keys. + qsmall=$(./waste run "$QWENC" x --budget 1 2>&1 || true) + qhuge=$(./waste run "$QWENC" x --ctx 8000000 --budget 8M 2>&1 || true) + if printf '%s' "$qsmall" | grep -qi "budget" && + printf '%s' "$qhuge" | grep -qi "budget"; then + ok "a RAM budget under the Qwen floor is refused, not swapped into" + else + no "an under-floor Qwen budget was accepted" + fi + + # top_k comes from the normalised key, and from HF's spelling when a + # container was written without it: planning for 0 experts would + # under-size the scratch that many pointers are cut from. + if [ -n "$PY_MISS" ]; then + sk "Qwen top_k alias" "$PY_MISS" + else + cp -R "$QWENC" "$TMP/qwen-alias.waste" + python3 - "$TMP/qwen-alias.waste" <<'PYQ' +import json, sys +p = sys.argv[1] + "/manifest.json" +m = json.load(open(p)) +m["config"].pop("num_experts_per_token", None) # leave only HF's spelling +json.dump(m, open(p, "w"), indent=1) +PYQ + if [ "$(./waste plan "$QWENC" --json)" = \ + "$(./waste plan "$TMP/qwen-alias.waste" --json)" ]; then + ok "plan reads top_k from num_experts_per_tok when the canonical key is absent" + else + no "Qwen plan disagrees with itself over the top_k alias" + fi + fi + + # Refusals. Each of these is a container the engine could open and read + # wrongly rather than fail on, which is the whole reason cfg_sane + # bounds them: an out-of-range n-gram overruns a fixed context array, a + # second indexer KV head is a shape nothing here implements, and a + # missing layer_types reads as "every layer is GDN" — plausible, + # answer-changing, and invisible. + qwen_refused() { # + local what="$1" edit="$2" dir="$TMP/qwen-bad.waste" + rm -rf "$dir"; cp -R "$QWENC" "$dir" + python3 -c "$edit" "$dir" || { no "$what (fixture edit failed)"; return; } + if ./waste info "$dir" >/dev/null 2>&1; then + no "$what was accepted" + else + ok "$what is refused" + fi + } + if [ -n "$PY_MISS" ]; then + sk "Qwen container refusals" "$PY_MISS" + else + qwen_refused "an n-gram order past the engine's fixed context" \ + 'import json,sys;p=sys.argv[1]+"/manifest.json";m=json.load(open(p));m["config"]["ngram_size"]=99;json.dump(m,open(p,"w"))' + qwen_refused "an indexer with more than one KV head" \ + 'import json,sys;p=sys.argv[1]+"/manifest.json";m=json.load(open(p));m["config"]["indexer_kv_heads"]=2;json.dump(m,open(p,"w"))' + qwen_refused "a container that does not say which layers are attention" \ + 'import json,sys;p=sys.argv[1]+"/manifest.json";m=json.load(open(p));m["config"].pop("layer_types");json.dump(m,open(p,"w"))' + qwen_refused "a layer_types shorter than num_hidden_layers" \ + 'import json,sys;p=sys.argv[1]+"/manifest.json";m=json.load(open(p));m["config"]["layer_types"]=m["config"]["layer_types"][:1];json.dump(m,open(p,"w"))' + qwen_refused "a PLE conv kernel the ring cannot hold" \ + 'import json,sys;p=sys.argv[1]+"/manifest.json";m=json.load(open(p));m["config"]["ple_conv_kernel_size"]=0;json.dump(m,open(p,"w"))' + fi + + # The isolated ops against an independent PyTorch reference written + # from the published equations, at the official geometry as well as at + # toy sizes. + if ! command -v uv >/dev/null 2>&1; then + sk "Qwen components" "uv not installed" + elif ./test_qwenparts "$TMP/qwenparts.bin" >/dev/null 2>&1 && + run_uv run --quiet --with torch --no-project python \ + tools/qwenparts_ref.py "$TMP/qwenparts.bin" 2>/dev/null | + grep -q "^PASS"; then + ok "PLE hashing, HyperConnection, GDN, QSA and the top-k router match the reference" + else + no "a Qwen component diverges from tools/qwenparts_ref.py" + fi + + # QSA's block top-k is a sort now, and the order it writes is the order + # attention sums in — so it is checked against the argmax it replaced, + # over ties, NaN and every budget, not only at the reference's one case. + # Plain C: it runs where uv does not. + if ./test_qsa_pick >/dev/null 2>&1; then + ok "QSA's block pick writes the selection in the argmax's order" + else + no "QSA's block pick orders the selection differently from the argmax" + fi + + # And the attention itself, scoring four selected tokens at a time and + # summing the values through NEON. What has to match is not only the + # order of the sums but the rounding of every product: the loop it + # replaced rounds each product on its own, and a fused multiply-add + # rounds once. Both ways of getting that wrong shifted the logits. + if ./test_qsa_attn >/dev/null 2>&1; then + ok "QSA's attention is bit-identical to the loops it replaced" + else + no "QSA's attention differs from the loops it replaced" + fi + + # The container-native oracle: the same container read by a PyTorch + # implementation of the same forward pass. Routes must match exactly + # and the argmax must match; the residual is gated at what this fixture + # measures, see docs/QWEN.md. + if ! command -v uv >/dev/null 2>&1; then + sk "container-native Qwen oracle" "uv not installed" + else + qout=$(run_uv run --quiet --with torch --no-project python \ + tests/test_qwen_container_ref.py 2>&1); qrc=$? + case "$qrc" in + 0) ok "the engine matches the container-native oracle (routes exact, argmax equal)" ;; + 77) sk "container-native Qwen oracle" "torch not installed" ;; + 124) sk "container-native Qwen oracle" "uv timed out" ;; + *) no "the engine diverges from the container-native Qwen oracle" + printf '%s\n' "$qout" | tail -12 ;; + esac + fi + + # The oracle's own gates, tested against synthetic disagreements: a + # comparison that cannot fail proves nothing about the runs it passes. + if [ -n "$PY_MISS" ]; then + sk "Qwen route near-tie gate" "$PY_MISS" + elif python3 tests/test_qwen_compare_oracle.py >/dev/null 2>&1; then + ok "the route comparison rejects a real reorder and accepts only a near tie" + else + no "the Qwen route near-tie gate does not reject a wrong expert" + fi + + if [ -n "$PY_MISS" ]; then + sk "Qwen streaming oracle" "$PY_MISS" + elif python3 tests/test_qwen_oracle.py >/dev/null 2>&1; then + ok "the official-weights oracle streams experts and n-gram rows rather than holding them" + else + no "tools/qwen_ref.py would materialize what it is meant to stream" + fi +fi + +# The tokenizer half that needs no weights always runs; the parity half +# needs the pinned checkpoint and the `tokenizers` package and says so. +if [ -n "$PY_MISS" ]; then + sk "Qwen tokenizer" "$PY_MISS" +else + if command -v uv >/dev/null 2>&1; then + qtok=$(run_uv run --quiet --with tokenizers --no-project python \ + tests/test_qwen_tok.py 2>&1); qtrc=$? + else + qtok=$(python3 tests/test_qwen_tok.py 2>&1); qtrc=$? + fi + case "$qtrc" in + 0) ok "the C tokenizer matches the pinned Qwen release, numbers included" ;; + 77) sk "Qwen tokenizer parity" "no pinned checkpoint (WASTE_QWEN_SRC) or no tokenizers package" ;; + 124) sk "Qwen tokenizer parity" "uv timed out" ;; + *) no "Qwen tokenizer parity" + printf '%s\n' "$qtok" | tail -12 ;; + esac +fi + head_ "RAM budget" # The default budget is the one path check_budget.sh cannot reach, because @@ -1915,7 +2176,12 @@ c = man["config"] hf = ((c.get("_outer", {}).get("architectures") or c.get("architectures") or [""]))[0] +# The same mapping waste_model_get_info makes, because that is what is +# being checked: a family the engine names and this rule does not would +# fail here for spelling rather than for describing the wrong container. arch = ("kimi-k3" if "KimiK3" in hf else "kimi-linear" if "KimiLinear" in hf + else "glm5-next" if "Glm5Next" in hf + else "qwen4_exp_text" if "Qwen4Exp" in hf else hf or "unknown") # a container that names nothing gets that NAMES = {0: "F32", 1: "F16", 2: "Q8G", 3: "Q4G", 7: "Q3G"} @@ -2226,6 +2492,57 @@ else printf '%s\n' "$out" | grep -E "FAIL|Error|Traceback" | head -5 fi +# Qwen nests its text model, packs 512 experts per layer into two tensors, +# and keeps its n-gram tables in 128 shards that become 16 heads. Every one +# of those is a shape no other member of this family has, and all three are +# checked without torch and without a 360 GB conversion. +if [ -n "$PY_MISS" ]; then + sk "convert.py Qwen config" "$PY_MISS" +elif out=$(python3 tests/test_convert_qwen.py 2>&1); then + ok "Qwen's nesting, packed expert layout, PLE consumers and reclaim classes" +else + no "convert.py Qwen config" + printf '%s\n' "$out" | grep -E "FAIL|Error|Traceback" | head -5 +fi + +# The PLE write is the one conversion step that cannot be done the obvious +# way: a head is ~12 GiB as f32, so it is quantized in row batches and the +# batches have to reconstruct exactly what quantizing the whole head would +# have given. +if ! command -v uv >/dev/null 2>&1; then + sk "Qwen PLE streaming write" "uv not installed" +else + out=$(run_uv run --quiet --with torch --no-project \ + python tests/test_qwen_ple_write.py 2>&1); rc=$? + case "$rc" in + 0) ok "PLE heads are written in Q8G row batches, not built whole in RAM" ;; + 77) sk "Qwen PLE streaming write" "torch not installed" ;; + 124) sk "Qwen PLE streaming write" "uv timed out" ;; + *) no "Qwen PLE streaming write"; printf '%s\n' "$out" | grep -E "FAIL|Error|Traceback" | head -5 ;; + esac +fi + +# End to end on a tiny packed source: nested config in, container out, with +# the vision tower and the MTP layer left behind and the 16 heads present. +if [ "${WASTE_SANITIZED:-0}" = 1 ]; then + # convert.py dlopens libwastevq for the encoder, and under a sanitized + # build ASan is not the first library a plain python3 loaded, so the + # run dies in the allocator instead of converting anything. Same cause + # as the serve suite's skip below. + sk "Qwen conversion round trip" "not run under the sanitizers" +elif ! command -v uv >/dev/null 2>&1; then + sk "Qwen conversion round trip" "uv not installed" +else + out=$(run_uv run --quiet --with torch --no-project \ + python tests/test_qwen_roundtrip.py 2>&1); rc=$? + case "$rc" in + 0) ok "a packed Qwen source converts to a container the engine's rules accept" ;; + 77) sk "Qwen conversion round trip" "torch not installed" ;; + 124) sk "Qwen conversion round trip" "uv timed out" ;; + *) no "Qwen conversion round trip"; printf '%s\n' "$out" | grep -E "FAIL|Error|Traceback|assert" | head -5 ;; + esac +fi + # DeepSeek-V4.1 shares almost nothing above the expert record with the rest # of this family, so the converter has to rename every tensor and lift half # its config off the wrapper. All of it is silent when wrong: an diff --git a/tests/sweep.c b/tests/sweep.c index c1ea63dba..5d01a542f 100644 --- a/tests/sweep.c +++ b/tests/sweep.c @@ -41,8 +41,8 @@ /* Filled by src/model.c under WASTE_PROFILE=1. Reported per arm because * the question "does a big cache make the rest of the engine slower" is a * rate question, and a rate is only comparable inside one process. */ -extern double waste_prof[16]; -extern uint64_t waste_prof_n[16]; +extern double waste_prof[32]; +extern uint64_t waste_prof_n[32]; extern uint64_t waste_tmv_bytes; extern int *waste_route_cap; extern int waste_route_n, waste_route_cap_n; diff --git a/tests/test_convert_chat.py b/tests/test_convert_chat.py index 1c2594e7f..302a83819 100644 --- a/tests/test_convert_chat.py +++ b/tests/test_convert_chat.py @@ -125,6 +125,42 @@ def main(): "a hand-edited chat.json outranks the shipped one") ck(build(tmp, "other", "LlamaForCausalLM", KL_MARKERS) is None, "nothing is known for another architecture, so nothing is guessed") + + print("chat_template.jinja") + src = H.make_src(os.path.join(tmp, "kimi-jinja-src")) + io.open(os.path.join(src, "chat_template.jinja"), "w", + encoding="utf-8").write("KIMI-TPL") + out = os.path.join(tmp, "kimi-jinja.waste") + os.makedirs(out) + H.run(out, src) + got = io.open(os.path.join(out, "chat_template.jinja"), + encoding="utf-8").read() + ck(got == "KIMI-TPL", "Kimi still copies chat_template.jinja") + + src = H.make_src(os.path.join(tmp, "qwen-jinja-src")) + cfg = json.load(open(os.path.join(src, "config.json"))) + cfg["architectures"] = ["Qwen4ExpForConditionalGeneration"] + cfg["model_type"] = "qwen4_exp" + cfg["text_config"] = { + "model_type": "qwen4_exp_text", + "num_hidden_layers": cfg["num_hidden_layers"], + "num_experts": cfg["num_experts"], + "num_experts_per_tok": 2, + "first_k_dense_replace": cfg.get("first_k_dense_replace", 1), + "hidden_size": cfg["hidden_size"], + "moe_intermediate_size": cfg["moe_intermediate_size"], + "vocab_size": cfg["vocab_size"], + } + json.dump(cfg, open(os.path.join(src, "config.json"), "w")) + io.open(os.path.join(src, "chat_template.jinja"), "w", + encoding="utf-8").write("QWEN-TOOLS") + out = os.path.join(tmp, "qwen-jinja.waste") + os.makedirs(out) + H.run(out, src) + got = io.open(os.path.join(out, "chat_template.jinja"), + encoding="utf-8").read() + ck(got == "QWEN-TOOLS", + "Qwen carries its chat template like any other release") finally: shutil.rmtree(tmp, ignore_errors=True) diff --git a/tests/test_convert_qwen.py b/tests/test_convert_qwen.py new file mode 100644 index 000000000..5c64589e1 --- /dev/null +++ b/tests/test_convert_qwen.py @@ -0,0 +1,402 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""test_convert_qwen.py — Qwen3.8-Flash-Next conversion helpers. + +The packed expert tensors, the nested text_config, the language_model +prefix, and the 128-to-16 PLE split are the reasons convert.py cannot +treat this checkpoint as another Kimi family member. This file pins those +rules without torch and without a 360 GB conversion. + + python3 tests/test_convert_qwen.py +""" +import json +import os +import shutil +import struct +import sys +import tempfile + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(REPO, "tools")) + +# Import after the resume stubs so convert.py can load without torch. +import test_convert_resume as H # noqa: E402 +CONV = H.CONV + +fails = 0 + + +def ck(cond, what): + global fails + print(f" {'ok ' if cond else 'FAIL'} {what}") + if not cond: + fails += 1 + + +def qwen_cfg(): + return { + "architectures": ["Qwen4ExpForConditionalGeneration"], + "model_type": "qwen4_exp", + "text_config": { + "model_type": "qwen4_exp_text", + "hidden_size": 2560, + "num_hidden_layers": 48, + "num_experts": 512, + "num_experts_per_tok": 10, + "moe_intermediate_size": 640, + "shared_expert_intermediate_size": 640, + "vocab_size": 248320, + "ngram_vocab_size_base": 20000000, + "split_ngram_parts": 128, + "heads_per_ngram": 8, + "ple_embed_dim": 2560, + "ple_layer_ids": [2], + }, + "vision_config": {"depth": 27, "hidden_size": 1152}, + } + + +def main(): + print("detect and flatten qwen4_exp_text") + prefix, src_pfx, cfg = CONV.source_prefixes(qwen_cfg()) + ck(CONV.is_qwen(cfg), "is_qwen recognises the flattened config") + ck(CONV.is_qwen(qwen_cfg()), "is_qwen recognises the raw config too") + ck(not CONV.is_qwen({"model_type": "kimi_linear"}), "Kimi is not Qwen") + ck(src_pfx == "model.language_model.", + f"experts are found under model.language_model. ({src_pfx!r})") + ck(cfg.get("model_type") == "qwen4_exp_text", + f"flattened model_type is qwen4_exp_text ({cfg.get('model_type')!r})") + ck(cfg.get("num_hidden_layers") == 48, "text_config fields are promoted") + ck("vision_config" in (cfg.get("_outer") or {}), + "vision_config stays under _outer, not in the text config") + ck(prefix == "", f"Qwen writes tensors under model.*, prefix={prefix!r}") + + print("normalise num_experts_per_tok") + n = CONV.normalise_cfg(cfg) + ck(n.get("num_experts_per_token") == 10, + f"num_experts_per_tok becomes num_experts_per_token ({n.get('num_experts_per_token')})") + ck(n.get("num_experts") == 512, "num_experts is already canonical") + + print("strip model.language_model. to model.") + ck(CONV.qwen_rename("model.language_model.layers.0.linear_attn.A_log") + == "model.layers.0.linear_attn.A_log", + "layer tensors lose language_model") + ck(CONV.qwen_rename("model.language_model.embed_tokens.weight") + == "model.embed_tokens.weight", + "embeddings lose language_model") + ck(CONV.qwen_rename("lm_head.weight") == "lm_head.weight", + "lm_head is already unprefixed") + + print("text-only exclusions") + ck(CONV.qwen_skip_tensor("mtp.layers.0.mlp.gate.weight"), "skip mtp.*") + ck(CONV.qwen_skip_tensor("model.visual.blocks.0.attn.qkv.weight"), + "skip model.visual.*") + ck(not CONV.qwen_skip_tensor( + "model.language_model.layers.0.mlp.gate.weight"), + "keep text tensors") + ck(CONV.qwen_skip_tensor( + "model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_0.weight") is False, + "PLE ngram shards are not a skip — they have their own consumer") + + print("what build_trunk drops") + drop = CONV.qwen_drop_trunk() + ck(drop("mtp.layers.0.mlp.gate.weight"), "the MTP layer has no reader") + ck(drop("model.visual.blocks.0.attn.qkv.weight"), "the tower is not carried") + ck(drop("model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_0.weight"), + "n-gram shards are build_ple's, not build_trunk's") + ck(drop("model.language_model.layers.1.ple.ple_embedding." + "ngram_heads_offsets"), + "the i64 tables go to the manifest, not the trunk") + ck(not drop("model.language_model.layers.0.self_attn.q_proj.weight"), + "ordinary text tensors are kept") + + print("packed expert shapes") + # The invariant, not the release's numbers: gate_up [E, 2I, H] beside + # down [E, H, I]. A fixture at E=2, I=8, H=16 converts on the same path. + ck(CONV.packed_shapes_ok((512, 1280, 2560), (512, 2560, 640), 512), + "the pinned Flash-Next pair is a valid packed layout") + ck(CONV.packed_shapes_ok((2, 16, 16), (2, 16, 8), 2), + "a tiny fixture with the same layout is accepted") + ck(not CONV.packed_shapes_ok((512, 1281, 2560), (512, 2560, 640), 512), + "an odd 2I cannot be split into gate and up") + ck(not CONV.packed_shapes_ok((512, 1280, 2560), (512, 2560, 641), 512), + "down's I must be gate_up's I") + ck(not CONV.packed_shapes_ok((512, 1280, 2560), (511, 2560, 640), 512), + "both tensors must hold every expert") + ck(not CONV.packed_shapes_ok((512, 1280), (512, 2560, 640), 512), + "a two-dimensional gate_up is not a packed layer") + gate, up = CONV.split_packed_gate_up_shape((512, 1280, 2560), 640) + ck(gate == (512, 640, 2560) and up == (512, 640, 2560), + f"gate/up split along dim 1 → {gate} {up}") + + print("packed expert source names") + g, d = CONV.qwen_packed_names("model.language_model.", 0) + ck(g == "model.language_model.layers.0.mlp.experts.gate_up_proj", + f"gate_up name {g}") + ck(d == "model.language_model.layers.0.mlp.experts.down_proj", + f"down name {d}") + + print("PLE 128 shards → 16 heads") + offsets = [0, 20000003, 40000026] + sizes = [20000003, 20000023, 20000033] + slices = CONV.ple_head_slices(offsets, sizes) + ck(slices[0] == (0, 20000003), f"head 0 slice {slices[0]}") + ck(slices[1] == (20000003, 20000023), f"head 1 slice {slices[1]}") + ck(CONV.ple_source_loc(0, 2500012) == (0, 0), "row 0 is shard 0 local 0") + ck(CONV.ple_source_loc(2500012, 2500012) == (1, 0), + "row 2500012 is shard 1 local 0") + ck(CONV.ple_source_loc(20000002, 2500012) == (7, 2499918), + "last row of head 0 lands in shard 7") + ck(CONV.PLE_HEADS == 16 and CONV.PLE_HEAD_WIDTH == 160, + "16 logical heads of width 160") + ck(sizes[0] == 20000003, "head 0 rows are the first prime after 20M") + + print("ShardDebt consumers") + + def consumer(name): + return CONV.ShardDebt.consumer(name, CONV.qwen_skip_tensor) + + ck(consumer( + "model.language_model.layers.3.mlp.experts.gate_up_proj") + == ("layer", 3), + "packed experts belong to their layer") + ck(consumer( + "model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_4.weight") + == CONV.ShardDebt.PLE, + "PLE ngram shards are a separate consumer") + ck(consumer( + "model.language_model.layers.1.ple.ple_embedding.ngram_heads_offsets") + == CONV.ShardDebt.PLE, + "I64 head offsets are a PLE consumer — build_ple reads them after trunk") + ck(consumer( + "model.language_model.layers.1.ple.ple_embedding." + "ngram_heads_vocab_sizes") + == CONV.ShardDebt.PLE, + "I64 vocab sizes are a PLE consumer") + ck(consumer( + "model.language_model.layers.1.ple.ple_embedding.layer_multipliers") + == CONV.ShardDebt.PLE, + "I64 layer_multipliers are a PLE consumer") + ck(consumer( + "model.language_model.layers.1.ple.key_proj.weight") + == CONV.ShardDebt.TRUNK, + "PLE projections stay on the trunk") + ck(consumer("mtp.layers.0.mlp.experts.gate_up_proj") + == CONV.ShardDebt.SKIP, + "mtp experts are not layer 0") + ck(consumer("model.visual.patch_embed.proj.weight") + == CONV.ShardDebt.SKIP, + "vision tensors do not hold a text consumer") + ck(consumer("model.language_model.embed_tokens.weight") + == CONV.ShardDebt.TRUNK, + "embeddings are trunk") + + print("a checkpoint that carries its tower keeps it") + # `model.visual.` is also how GLM spells its tower, and GLM does carry + # it. Without a skip predicate nothing is a SKIP. + ck(CONV.ShardDebt.consumer("model.visual.patch_embed.proj.weight") + == CONV.ShardDebt.TRUNK, + "with no skip predicate the tower is a trunk tensor") + + print("Kimi prefixes are unchanged") + k3 = { + "architectures": ["KimiK3ForConditionalGeneration"], + "text_config": {"model_type": "kimi_linear", "num_hidden_layers": 2}, + } + prefix, src_pfx, cfg = CONV.source_prefixes(k3) + ck(not CONV.is_qwen(cfg), "K3 is not Qwen") + ck(prefix == "language_model.", f"K3 prefix stays language_model. ({prefix!r})") + ck(src_pfx == "language_model.model.", f"K3 src prefix stays ({src_pfx!r})") + + print("--reclaim dry keeps a PLE shard until PLE finishes") + tmp = tempfile.mkdtemp(prefix="qwen-debt-") + try: + wm = { + "model.language_model.layers.0.mlp.gate.weight": "shard-mix.safetensors", + "model.language_model.layers.0.mlp.experts.gate_up_proj": + "shard-mix.safetensors", + "model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_0.weight": "shard-mix.safetensors", + "mtp.fc_hidden.weight": "shard-mtp.safetensors", + "model.visual.pos_embed.weight": "shard-vis.safetensors", + } + src = os.path.join(tmp, "src") + os.makedirs(src) + for shard in ("shard-mix.safetensors", "shard-mtp.safetensors", + "shard-vis.safetensors"): + open(os.path.join(src, shard), "wb").write(b"\0" * 64) + debt = CONV.ShardDebt(wm, src, CONV.qwen_skip_tensor) + CONV.reclaim(debt, "dry", CONV.ShardDebt.SKIP, "excluded") + CONV.reclaim(debt, "dry", CONV.ShardDebt.TRUNK, "trunk") + CONV.reclaim(debt, "dry", ("layer", 0), "layer 0") + still = sorted(debt.owed) + ck(still == ["shard-mix.safetensors"], + f"mix shard still held after trunk+layer, skip shards gone ({still})") + CONV.reclaim(debt, "dry", CONV.ShardDebt.PLE, "ple") + ck(list(debt.owed) == [], "PLE was the last consumer of the mix shard") + ck(os.path.exists(os.path.join(src, "shard-mix.safetensors")), + "dry deletes nothing") + + # Offsets live in their own shard on some layouts. build_ple reads + # them after the trunk reclaim, so TRUNK must not be their consumer. + open(os.path.join(src, "shard-i64.safetensors"), "wb").write(b"\0" * 64) + open(os.path.join(src, "shard-gate.safetensors"), "wb").write(b"\0" * 64) + debt2 = CONV.ShardDebt({ + "model.language_model.layers.1.ple.ple_embedding." + "ngram_heads_offsets": "shard-i64.safetensors", + "model.language_model.layers.0.mlp.gate.weight": + "shard-gate.safetensors", + }, src, CONV.qwen_skip_tensor) + CONV.reclaim(debt2, "dry", CONV.ShardDebt.TRUNK, "trunk") + ck("shard-i64.safetensors" in debt2.owed, + f"I64 shard still held after trunk ({sorted(debt2.owed)})") + CONV.reclaim(debt2, "dry", CONV.ShardDebt.PLE, "ple") + ck("shard-i64.safetensors" not in debt2.owed, + "PLE release frees the I64 shard") + finally: + shutil.rmtree(tmp, ignore_errors=True) + + print("PLE reclaim needs 16 ngram_head tensors in the trunk index") + def head(h): + return {"name": f"model.layers.1.ple.ple_embedding.ngram_head.{h}.weight"} + ck(not CONV.ple_heads_written(None), "absent index is incomplete") + ck(not CONV.ple_heads_written([]), "empty index is incomplete") + ck(not CONV.ple_heads_written([head(h) for h in range(15)]), + "15 heads are incomplete") + ck(CONV.ple_heads_written([head(h) for h in range(16)]), + "16 heads are complete") + ck(not CONV.ple_heads_written([head(0)] * 16), + "sixteen copies of head 0 are not 16 heads") + + tmp = tempfile.mkdtemp(prefix="qwen-ple-reclaim-") + try: + src = os.path.join(tmp, "src") + os.makedirs(src) + open(os.path.join(src, "shard-ple.safetensors"), "wb").write(b"\0" * 64) + wm = { + "model.language_model.layers.1.ple.ple_embedding." + "ngram_embedding.shard_0.weight": "shard-ple.safetensors", + } + debt = CONV.ShardDebt(wm, src, CONV.qwen_skip_tensor) + CONV.reclaim_ple_if_complete(debt, "dry", []) + ck("shard-ple.safetensors" in debt.owed, + "skip-trunk with no heads retains the PLE shard") + ck(os.path.exists(os.path.join(src, "shard-ple.safetensors")), + "incomplete dry deletes nothing") + CONV.reclaim_ple_if_complete(debt, "dry", [head(h) for h in range(15)]) + ck("shard-ple.safetensors" in debt.owed, + "15 published heads still retain the PLE shard") + CONV.reclaim_ple_if_complete(debt, "dry", [head(h) for h in range(16)]) + ck("shard-ple.safetensors" not in debt.owed, + "16 published heads permit dry PLE reclaim") + ck(os.path.exists(os.path.join(src, "shard-ple.safetensors")), + "complete dry still deletes nothing") + finally: + shutil.rmtree(tmp, ignore_errors=True) + + print("Qwen packed default jobs is 1 unless --jobs is set") + ck(CONV.resolve_jobs(True, None) == 1, "omitted --jobs is 1 for Qwen") + ck(CONV.resolve_jobs(False, None) == 3, "omitted --jobs stays 3 otherwise") + ck(CONV.resolve_jobs(True, 4) == 4, "explicit --jobs 4 wins for Qwen") + ck(CONV.resolve_jobs(False, 1) == 1, "explicit --jobs 1 wins for Kimi") + + raw = os.environ.get("QWEN_SRC", + "/Users/admin/mnt/llm/qwen38-flash-next/raw") + idxp = os.path.join(raw, "model.safetensors.index.json") + if not os.path.isfile(idxp): + print("real source sample SKIP (no pinned checkpoint)") + else: + print("real source sample (headers + I64 + reclaim class, no 360 GB)") + wm = json.load(open(idxp))["weight_map"] + g, d = CONV.qwen_packed_names("model.language_model.", 0) + + def meta(name): + fn = wm[name] + with open(os.path.join(raw, fn), "rb") as f: + (hlen,) = struct.unpack(" lg[best]) best = v; printf("prefill %d tok in %.2fs (%.2f tok/s); argmax %d, max %.4f\n", n, tp, n / tp, best, lg[best]); + if (m.cfg.arch_qwen && m.has_qsa) { + const int compress = m.cfg.idx_compress > 0 ? m.cfg.idx_compress : 4; + for (int L = 0; L < m.cfg.n_layers; L++) { + if (!m.cfg.qwen_full[L]) continue; + printf("qsa_layer %d n_kv %d blk %d tail %d compress %d\n", + L, m.n_kv[L], m.n_qsa_blk[L], m.n_qsa_tail[L], compress); + break; + } + } if (out) { FILE *f = fopen(out, "wb"); @@ -103,10 +120,34 @@ int main(int argc, char **argv) printf("wrote %s\n", out); } + /* WASTE_PROFILE=decode drops the prompt from the profile. The prompt + * runs the same one-token path, but it is what finds the cache empty, + * so on a short run it is most of the expert I/O and says nothing about + * where a steady-state decode step goes. */ + const char *prof = getenv("WASTE_PROFILE"); + const int prof_decode = prof && !strcmp(prof, "decode"); + unsigned long long reads0 = 0; + if (prof_decode) { + memset(waste_prof, 0, sizeof waste_prof); + memset(waste_prof_n, 0, sizeof waste_prof_n); + memset(waste_prof_tmv, 0, sizeof waste_prof_tmv); + /* Counters only: the tensors already hold their row numbers. */ + for (int r = 0; r < waste_tmv_nroles; r++) { + waste_tmv_roles[r].calls = waste_tmv_roles[r].bytes = 0; + waste_tmv_roles[r].t = waste_tmv_roles[r].tq = 0; + } + waste_tmv_bytes = 0; + memset(waste_tmv_t, 0, sizeof waste_tmv_t); + memset(waste_tmv_b, 0, sizeof waste_tmv_b); + memset(waste_tmv_c, 0, sizeof waste_tmv_c); + reads0 = (unsigned long long)m.expert_reads; + } + double tg = 0; int cur = best; for (int i = 0; i < n_gen; i++) { t0 = now(); lg = waste_model_step(&m, cur, n + i, NULL); + tg += now() - t0; if (!lg) { int layer = 0, expert = 0; const char *why = waste_model_read_error(&m, &layer, &expert); @@ -122,28 +163,99 @@ int main(int argc, char **argv) cur = best; } - extern double waste_prof[16]; - if (getenv("WASTE_PROFILE")) { - /* indented names are sub-totals of the line above and are excluded - * from `tot`, so the percentages add to 100 */ - const char *names[10] = {" LUT build","kda","mla","moe(all)", - " expert I/O"," expert mm","lm_head", - " LUT apply"," batched mm"," trunk matvec"}; + if (prof) { + const int steps = prof_decode ? n_gen : n + n_gen; + const double wall = prof_decode ? tg : tp + tg; double tot = 0; - for (int i = 0; i < 10; i++) - tot += (i == 0 || i == 4 || i == 5 || i == 7 || i == 8 || i == 9) ? 0 : waste_prof[i]; - printf("\n-- profile (s, %d steps) --\n", n + n_gen); - for (int i = 0; i < 10; i++) - if (waste_prof[i] > 0) - printf(" %-14s %7.2f %5.1f%%\n", names[i], waste_prof[i], - 100.0 * waste_prof[i] / tot); - printf(" %-14s %7.2f\n", "accounted", tot); - { extern uint64_t waste_tmv_bytes; extern uint64_t waste_prof_n[16]; - printf(" trunk matvec: %llu calls, %.2f GB, %.1f GB/s\n", + printf("\n-- profile (s, %d %ssteps) --\n", steps, prof_decode ? "decode " : ""); + if (m.cfg.arch_qwen) { + /* A tree: a `sub` row sits inside the nearest less-indented row + * above it, and only top-level rows add to `tot`. Percentages + * are of wall time, so what no phase covers stays visible. + * + * The LUT rows are only what the calling thread builds and + * applies. The expert-parallel path runs each expert's applies + * inside a pool worker, untimed, and there `expert mm` is the + * one number that covers them. */ + static const struct { int slot, sub; const char *name; } rows[] = { + {12, 0, "ple"}, {11, 0, "hyperconnection"}, + {1, 0, "gdn"}, {10, 1, " recurrence"}, + {2, 0, "qsa"}, {14, 1, " select+attend"}, + {16, 1, " rope table"}, {17, 1, " block select"}, + {18, 1, " K/V gather"}, {19, 1, " attention"}, + {3, 0, "moe(all)"}, {15, 1, " router"}, + {20, 1, " lookahead"}, + {4, 1, " expert I/O"}, {5, 1, " expert mm"}, + {0, 1, " LUT build"}, {7, 1, " LUT apply"}, + {13, 1, " shared expert"}, + {6, 0, "lm_head"}, + }; + const int nrows = (int)(sizeof rows / sizeof rows[0]); + for (int r = 0; r < nrows; r++) + if (!rows[r].sub) tot += waste_prof[rows[r].slot]; + /* The last column is how much of each row was trunk matvec, + * so the rest of the row is everything between projections. */ + for (int r = 0; r < nrows; r++) { + const double s = waste_prof[rows[r].slot]; + if (s > 0) + printf(" %-16s %7.2f %5.1f%% %7.2f ms/step %7.2f matvec\n", + rows[r].name, s, 100.0 * s / wall, + steps ? 1e3 * s / steps : 0.0, + steps ? 1e3 * waste_prof_tmv[rows[r].slot] / steps : 0.0); + } + printf(" %-16s %7.2f %5.1f%%\n", "accounted", tot, 100.0 * tot / wall); + } else { + /* indented names are sub-totals of the line above and are excluded + * from `tot`, so the percentages add to 100 */ + const char *names[10] = {" LUT build","kda","mla","moe(all)", + " expert I/O"," expert mm","lm_head", + " LUT apply"," batched mm"," trunk matvec"}; + for (int i = 0; i < 10; i++) + tot += (i == 0 || i == 4 || i == 5 || i == 7 || i == 8 || i == 9) ? 0 : waste_prof[i]; + for (int i = 0; i < 10; i++) + if (waste_prof[i] > 0) + printf(" %-14s %7.2f %5.1f%%\n", names[i], waste_prof[i], + 100.0 * waste_prof[i] / tot); + printf(" %-14s %7.2f\n", "accounted", tot); + } + /* The steps themselves, timed around waste_model_step. The gap to + * `accounted` is whatever no phase covers — on Qwen the embedding + * row and the tensor lookups between phases. More than a few + * percent means a phase is missing, not that something is slow. */ + printf(" %-*s %7.2f (%.2f s unaccounted, %.2f ms/step)\n", + m.cfg.arch_qwen ? 16 : 14, "wall", wall, + wall - tot, steps ? 1e3 * (wall - tot) / steps : 0.0); + printf(" expert reads: %llu (%.1f per step)\n", + (unsigned long long)m.expert_reads - reads0, + steps ? ((unsigned long long)m.expert_reads - reads0) / (double)steps : 0.0); + { printf(" trunk matvec: %llu calls, %.2f GB, %.1f GB/s\n", (unsigned long long)waste_prof_n[9], waste_tmv_bytes / 1e9, waste_prof[9] > 0 ? waste_tmv_bytes / waste_prof[9] / 1e9 : 0.0); - extern double waste_tmv_t[4]; extern uint64_t waste_tmv_b[4], waste_tmv_c[4]; + /* By tensor, heaviest first: the size buckets below cannot say + * which projection a millisecond was in. */ + int ord[WASTE_TMV_ROLES]; + for (int k = 0; k < waste_tmv_nroles; k++) ord[k] = k; + for (int k = 1; k < waste_tmv_nroles; k++) + for (int j = k; j > 0 && + waste_tmv_roles[ord[j]].t > waste_tmv_roles[ord[j - 1]].t; j--) { + const int sw = ord[j]; ord[j] = ord[j - 1]; ord[j - 1] = sw; + } + for (int k = 0; k < waste_tmv_nroles && k < 20; k++) { + const waste_tmv_role *r = &waste_tmv_roles[ord[k]]; + if (!r->calls) continue; + /* The last two: how much of the row was quantizing the + * activation, and the kernel's speed with that taken out. */ + printf(" %-46s %5dx%-5d q%-2d %6.2f ms/step %5.1f calls %6.2f MB %6.1f GB/s" + " | quant %5.2f ms, kernel %6.1f GB/s\n", + r->role, r->out, r->in, r->bits, + steps ? 1e3 * r->t / steps : 0.0, + steps ? (double)r->calls / steps : 0.0, + r->bytes / (double)r->calls / 1e6, + r->t > 0 ? r->bytes / r->t / 1e9 : 0.0, + steps ? 1e3 * r->tq / steps : 0.0, + r->t > r->tq ? r->bytes / (r->t - r->tq) / 1e9 : 0.0); + } const char *bn[4] = {" <1MB"," 1-8MB"," 8-32MB"," >32MB"}; for (int k = 0; k < 4; k++) if (waste_tmv_c[k]) diff --git a/tests/test_qsa_attn.c b/tests/test_qsa_attn.c new file mode 100644 index 000000000..e39d5fb45 --- /dev/null +++ b/tests/test_qsa_attn.c @@ -0,0 +1,100 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * test_qsa_attn.c — QSA's attention, against the loops it replaced. + * + * ./test_qsa_attn + * + * waste_qwen_qsa_attn_heads scores four selected tokens at a time and sums + * the values through NEON, where it used to do both one element at a time. + * Neither is allowed to change a bit, and the way to get that wrong is not + * the order of the sums — it is the rounding of each product. The loop this + * replaced compiles to four independent products in one vector and a scalar + * chain of adds, so every product is rounded on its own; written as + * `s += q * k` the same arithmetic contracts to a fused multiply-add, which + * rounds once. Both versions of that mistake produced logits that differed + * in the last bits and generated different text a few hundred tokens in. + * The old loops are copied verbatim below. + */ +#include +#include +#include +#include + +#include "../src/qwen_qsa.h" + +static void old_attn(int h0, int h1, const float *q, int Hq, int D, + const float *k, const float *v, int Hkv, int T, + const int *sel, int n_sel, float scale, + float *out, float *scratch) +{ + const int n_rep = Hkv > 0 ? Hq / Hkv : 1; + float *scores = scratch; + if (h1 > h0) + memset(out + (size_t)h0 * D, 0, (size_t)(h1 - h0) * D * sizeof(float)); + if (!q || !k || !v || !sel || !scratch || n_sel < 1) return; + for (int h = h0; h < h1; h++) { + const int hv = h / (n_rep > 0 ? n_rep : 1); + const float *qh = q + (size_t)h * D; + float m = -1e30f; + for (int i = 0; i < n_sel; i++) { + const int t = sel[i]; + if (t < 0 || t >= T) { scores[i] = -1e30f; continue; } + const float *kh = k + ((size_t)t * Hkv + hv) * D; + float s = 0.0f; + for (int d = 0; d < D; d++) s += qh[d] * kh[d]; + s *= scale; + scores[i] = s; + if (s > m) m = s; + } + float z = 0.0f; + for (int i = 0; i < n_sel; i++) { + scores[i] = expf(scores[i] - m); + z += scores[i]; + } + if (z < 1e-20f) z = 1e-20f; + float *oh = out + (size_t)h * D; + for (int i = 0; i < n_sel; i++) { + const int t = sel[i]; + if (t < 0 || t >= T) continue; + const float w = scores[i] / z; + const float *vh = v + ((size_t)t * Hkv + hv) * D; + for (int d = 0; d < D; d++) oh[d] += w * vh[d]; + } + } +} + +static unsigned rng = 99991u; +static float frnd(void) { rng = rng * 1103515245u + 12345u; return (float)((int)(rng >> 9) - 4194304) / 4194304.0f; } + +int main(void) +{ + enum { D = 256, Hq = 8, Hkv = 2, T = 600, NSEL = 512 }; + static float q[Hq * D], k[T * Hkv * D], v[T * Hkv * D]; + static float o1[Hq * D], o2[Hq * D], sc1[NSEL], sc2[NSEL]; + static int sel[NSEL]; + int bad = 0; + for (int it = 0; it < 40; it++) { + for (size_t i = 0; i < sizeof q / sizeof *q; i++) q[i] = frnd(); + for (size_t i = 0; i < sizeof k / sizeof *k; i++) k[i] = frnd(); + for (size_t i = 0; i < sizeof v / sizeof *v; i++) v[i] = frnd(); + const int n_sel = 1 + (int)(rng % NSEL); + for (int i = 0; i < n_sel; i++) { + rng = rng * 1103515245u + 12345u; + /* mostly valid, some out of range, as a short context gives */ + sel[i] = (rng % 50 == 0) ? -1 : (int)(rng % (T + 4)); + } + old_attn(0, Hq, q, Hq, D, k, v, Hkv, T, sel, n_sel, 0.125f, o1, sc1); + waste_qwen_qsa_attn_heads(0, Hq, q, Hq, D, k, v, Hkv, T, sel, n_sel, 0.125f, o2, sc2); + if (memcmp(o1, o2, sizeof o1)) { + int first = -1; + for (int i = 0; i < Hq * D; i++) if (memcmp(&o1[i], &o2[i], 4)) { first = i; break; } + printf("iter %d n_sel %d: first differing output %d: %.9g vs %.9g\n", + it, n_sel, first, o1[first], o2[first]); + bad++; + } + } + printf("%s: %d of 40 cases differ\n", bad ? "FAIL" : "ok", bad); + return bad ? 1 : 0; +} diff --git a/tests/test_qsa_pick.c b/tests/test_qsa_pick.c new file mode 100644 index 000000000..dbc8e6704 --- /dev/null +++ b/tests/test_qsa_pick.c @@ -0,0 +1,85 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * test_qsa_pick.c — QSA's block top-k, against the loop it replaced. + * + * ./test_qsa_pick + * + * waste_qwen_qsa_pick sorts where the selection used to take a repeated + * argmax, and the order it writes is the order attention sums the selected + * tokens in — so "the same set" is not enough, it has to be the same + * sequence, bit for bit downstream. The reference check covers the usual + * case; this covers the ones that decide an ordering: ties everywhere, + * scores the argmax could never take (NaN, -1e30, -inf), and a budget below, + * at and above the block count. The old loop is copied verbatim. + */ +#include +#include +#include +#include + +#include "../src/qwen_qsa.h" + +static int old_pick(const float *scores, int n_complete, int block_topk, + int compress, int n_tail, int *sel, int *taken) +{ + int nsel = 0; + if (n_complete > 0) { + for (int b = 0; b < n_complete; b++) taken[b] = 0; + const int keep = n_complete < block_topk ? n_complete : block_topk; + for (int j = 0; j < keep; j++) { + int best = -1; + float bv = -1e30f; + for (int b = 0; b < n_complete; b++) { + if (taken[b]) continue; + if (scores[b] > bv) { bv = scores[b]; best = b; } + } + if (best < 0) break; + taken[best] = 1; + for (int t = 0; t < compress; t++) + sel[nsel++] = best * compress + t; + } + } + for (int t = 0; t < n_tail; t++) + sel[nsel++] = n_complete * compress + t; + return nsel; +} + +static unsigned rng = 12345u; +static unsigned rnd(void) { rng = rng * 1103515245u + 12345u; return rng >> 8; } + +int main(void) +{ + enum { MAXB = 1200, MAXSEL = MAXB * 4 + 8 }; + static float scores[MAXB]; + static int sel_a[MAXSEL], sel_b[MAXSEL], taken[MAXB], order[MAXB]; + int cases = 0; + for (int it = 0; it < 4000; it++) { + const int n = (int)(rnd() % (it < 200 ? 16 : MAXB)); + const int compress = 1 + (int)(rnd() % 4); + const int n_tail = (int)(rnd() % compress); + const int budget_kind = (int)(rnd() % 3); + const int topk = budget_kind == 0 ? (n > 0 ? (int)(rnd() % n) : 0) + : budget_kind == 1 ? n : n + 1 + (int)(rnd() % 50); + const int levels = 1 + (int)(rnd() % 6); /* few distinct scores: ties */ + for (int b = 0; b < n; b++) { + const unsigned r = rnd() % 100; + if (r < 3) scores[b] = NAN; + else if (r < 5) scores[b] = -1e30f; + else if (r < 6) scores[b] = -INFINITY; + else if (r < 50) scores[b] = (float)(rnd() % levels); + else scores[b] = (float)(rnd() % 100000) / 7.0f; + } + const int na = old_pick(scores, n, topk, compress, n_tail, sel_a, taken); + const int nb = waste_qwen_qsa_pick(scores, n, topk, compress, n_tail, sel_b, order); + cases++; + if (na != nb || memcmp(sel_a, sel_b, (size_t)na * sizeof(int)) != 0) { + printf("FAIL case %d: n=%d topk=%d compress=%d tail=%d -> old %d, new %d\n", + it, n, topk, compress, n_tail, na, nb); + return 1; + } + } + printf("ok: %d cases, the pick matches the argmax it replaced\n", cases); + return 0; +} diff --git a/tests/test_qwen_compare_oracle.py b/tests/test_qwen_compare_oracle.py new file mode 100644 index 000000000..3f0d71fd4 --- /dev/null +++ b/tests/test_qwen_compare_oracle.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""Synthetic tests for qwen_compare_oracle near-tie route gate.""" +from __future__ import annotations + +import os +import sys + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(ROOT, "tools")) + +from qwen_compare_oracle import check_route_ids, score_err_bound, _tolerance + + +def ck(cond, msg): + if not cond: + print(f"FAIL {msg}") + raise SystemExit(1) + print(f" ok {msg}") + + +def main(): + cfg = {"num_hidden_layers": 48, "hidden_size": 2560, + "num_experts_per_tok": 10, "hc_count": 4} + h_bound = _tolerance("hidden", cfg) + score_err = score_err_bound(cfg, h_bound) + w_bound = _tolerance("route_w", cfg) + print(f"score_err_bound={score_err:.3e} w_bound={w_bound:.3e}") + + ci = [10, 20, 30, 40] + pi = [10, 20, 30, 40] + cw = [0.4, 0.3, 0.2, 0.1] + pw = [0.4, 0.3, 0.2, 0.1] + ok, _ = check_route_ids(ci, pi, cw, pw, score_err, w_bound, True) + ck(ok, "exact match passes") + + # True wrong expert outside near-tie interval + pi_bad = [10, 20, 30, 99] + ok, why = check_route_ids(ci, pi_bad, cw, pw, score_err, w_bound, True) + ck(not ok and "mismatch" in why, f"wrong expert fails: {why}") + + # Near-tie swap: scores 0.100 vs 0.101, margin << score_err + ci = [1, 2, 3, 4] + pi = [1, 3, 2, 4] + cw = [0.30, 0.100, 0.101, 0.05] + pw = [0.30, 0.101, 0.100, 0.05] + ok, why = check_route_ids(ci, pi, cw, pw, score_err, w_bound, True) + ck(ok and why == "near-tie reorder", f"near-tie swap accepted: {why}") + + # Near-tie swap but weight error too large + pw_bad = [0.30, 0.101, 0.050, 0.05] + ok, why = check_route_ids(ci, pi, cw, pw_bad, score_err, w_bound, True) + ck(not ok and "weight" in why, f"near-tie weight fail: {why}") + + # Near-tie swap but hidden layer over bound + ok, why = check_route_ids(ci, pi, cw, pw, score_err, w_bound, False) + ck(not ok and "hidden" in why, f"near-tie hidden fail: {why}") + + # Non-near-tie reorder: large margin between swapped experts + ci = [1, 2, 3, 4] + pi = [1, 3, 2, 4] + cw = [0.30, 0.50, 0.10, 0.05] + pw = [0.30, 0.10, 0.50, 0.05] + ok, why = check_route_ids(ci, pi, cw, pw, score_err, w_bound, True) + ck(not ok and "non-near-tie reorder" in why, f"large margin fails: {why}") + + print("NEAR-TIE GATE OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_qwen_container_ref.py b/tests/test_qwen_container_ref.py new file mode 100644 index 000000000..fdcb3b8cd --- /dev/null +++ b/tests/test_qwen_container_ref.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""Minimal Qwen container-native oracle on the synthetic --qwen fixture. + + python3 tests/test_qwen_container_ref.py +""" +from __future__ import annotations + +import json +import os +import shutil +import struct +import subprocess +import sys +import tempfile + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(ROOT, "tools")) + +try: + import torch # noqa: F401 +except ImportError: + print("SKIP: torch is not installed") + raise SystemExit(77) + + +def ck(cond, what): + print(f" {'ok ' if cond else 'FAIL'} {what}") + if not cond: + raise SystemExit(1) + + +def f32(path): + b = open(path, "rb").read() + n = len(b) // 4 + return list(struct.unpack(f"<{n}f", b)) if n else [] + + +def main(): + import qwen_container_ref as Q + + ck(hasattr(Q, "QwenContainer"), "QwenContainer exists") + ck(hasattr(Q, "QwenRef"), "QwenRef exists") + ck(issubclass(Q.QwenContainer, Q.Container), + "QwenContainer subclasses kimi_ref.Container") + + tmp = tempfile.mkdtemp(prefix="qwen-cref-") + try: + cont = os.path.join(tmp, "qwen.waste") + r = subprocess.run( + [sys.executable, os.path.join(ROOT, "tools", "make_test_container.py"), + "--qwen", cont], + cwd=ROOT, capture_output=True, text=True) + if r.returncode != 0: + print("FAIL fixture", r.stderr[-400:] or r.stdout[-400:]) + return 1 + man = json.load(open(os.path.join(cont, "manifest.json"))) + cfg = man["config"] + hid = int(cfg["hidden_size"]) + nL = int(cfg["num_hidden_layers"]) + hc = int(cfg["hc_count"]) + vocab = int(cfg["vocab_size"]) + top_k = int(cfg.get("num_experts_per_tok") + or cfg.get("num_experts_per_token") or 0) + ids = "3,7,11" + ntok = 3 + + c = Q.QwenContainer(cont) + ck(c.iblock == 64, f"index_block is 64 (got {c.iblock})") + shapes = c.expert_shapes() + moe = int(cfg["moe_intermediate_size"]) + ck(shapes == [(moe, hid), (moe, hid), (hid, moe)], + f"expert_shapes gate/up {moe}x{hid}, down {hid}x{moe}") + + dump = os.path.join(tmp, "logits.bin") + hidden = os.path.join(tmp, "hidden.bin") + routes = os.path.join(tmp, "routes.txt") + cli = subprocess.run( + [sys.executable, os.path.join(ROOT, "tools", "qwen_container_ref.py"), + "--container", cont, "--ids", ids, + "--dump", dump, "--hidden", hidden, "--routes", routes], + cwd=ROOT, capture_output=True, text=True) + if cli.returncode != 0: + print("FAIL cli", cli.stderr[-800:] or cli.stdout[-800:]) + return 1 + ck(os.path.isfile(dump), "wrote logits.bin") + ck(os.path.isfile(hidden), "wrote hidden.bin") + ck(os.path.isfile(routes), "wrote routes.txt") + + lg = f32(dump) + ck(len(lg) == vocab, f"logits {len(lg)} floats (vocab {vocab})") + + hid_n = len(f32(hidden)) + want_all = ntok * nL * hc * hid + want_last = nL * hc * hid + ck(hid_n in (want_all, want_last), + f"hidden {hid_n} floats " + f"(all-tok {want_all} or last-tok {want_last})") + + lines = [ln for ln in open(routes).read().splitlines() if ln.strip()] + ck(len(lines) == ntok * nL, + f"routes {len(lines)} lines (want {ntok} tok x {nL} layers)") + for ln in lines: + parts = ln.split() + ck(len(parts) == 2 + 2 * top_k, + f"route '{ln}' has pos layer + {top_k} ids + {top_k} weights") + + fwd = os.path.join(ROOT, "test_forward") + if not os.path.isfile(fwd): + print("CONTAINER REF OK (no test_forward; C diff skipped)") + return 0 + + c_logits = os.path.join(tmp, "c_logits.bin") + c_hidden = os.path.join(tmp, "c_hidden.bin") + env = dict(os.environ) + env["WASTE_DUMP_HIDDEN"] = c_hidden + env["WASTE_DUMP_ROUTE"] = os.path.join(tmp, "c_routes.txt") + env["WASTE_Q8"] = "0" + env["WASTE_BACKEND"] = "cpu" + fr = subprocess.run( + [fwd, cont, ids, c_logits, "0"], + cwd=ROOT, env=env, capture_output=True, text=True) + if fr.returncode != 0: + print("FAIL test_forward", fr.stderr[-400:] or fr.stdout[-400:]) + return 1 + cl, ol = f32(c_logits), lg + ck(len(cl) == len(ol), f"C/Python logits length {len(cl)} vs {len(ol)}") + d = [abs(a - b) for a, b in zip(cl, ol)] + mx = max(d) if d else 0.0 + ai, bi = cl.index(max(cl)), ol.index(max(ol)) + ck(ai == bi, f"logits argmax {ai} vs {bi}") + ck(mx < 1e-4, f"logits max|diff| {mx:.3e} < 1e-4") + cmp = subprocess.run( + [sys.executable, os.path.join(ROOT, "tools", "qwen_compare_oracle.py"), + "--container", cont, + "--c-hidden", c_hidden, "--c-logits", c_logits, + "--py-hidden", hidden, "--py-logits", dump, + "--c-routes", os.path.join(tmp, "c_routes.txt"), + "--py-routes", routes, + "--tokens", str(ntok)], + cwd=ROOT, capture_output=True, text=True) + if cmp.returncode != 0: + print("FAIL compare", cmp.stdout[-800:] or cmp.stderr[-800:]) + return 1 + print(f"CONTAINER REF OK logits max|diff| {mx:.3e} argmax {ai}") + return 0 + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +if __name__ == "__main__": + try: + sys.exit(main()) + except SystemExit: + raise + except Exception as e: + print(f"FAIL {type(e).__name__}: {e}") + sys.exit(1) diff --git a/tests/test_qwen_dump.py b/tests/test_qwen_dump.py new file mode 100644 index 000000000..be9a1b636 --- /dev/null +++ b/tests/test_qwen_dump.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""Qwen WASTE_DUMP_HIDDEN writes one hyper-state vector after every layer. + + python3 tests/test_qwen_dump.py +""" +from __future__ import annotations + +import json +import os +import shutil +import subprocess +import sys +import tempfile + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def main(): + tmp = tempfile.mkdtemp(prefix="qwen-dump-") + try: + cont = os.path.join(tmp, "qwen.waste") + r = subprocess.run( + [sys.executable, os.path.join(ROOT, "tools", "make_test_container.py"), + "--qwen", cont], + cwd=ROOT, capture_output=True, text=True) + if r.returncode != 0: + print("FAIL fixture", r.stderr[-400:]) + return 1 + man = json.load(open(os.path.join(cont, "manifest.json"))) + cfg = man["config"] + hid = int(cfg["hidden_size"]) + nL = int(cfg["num_hidden_layers"]) + hc = int(cfg["hc_count"]) + dump = os.path.join(tmp, "h.bin") + env = dict(os.environ) + env["WASTE_DUMP_HIDDEN"] = dump + fwd = subprocess.run( + [os.path.join(ROOT, "test_forward"), cont, "3,7,11", + os.path.join(tmp, "logits.bin"), "0"], + cwd=ROOT, env=env, capture_output=True, text=True) + if fwd.returncode != 0: + print("FAIL test_forward", fwd.stderr[-400:] or fwd.stdout[-400:]) + return 1 + if not os.path.isfile(dump): + print("FAIL WASTE_DUMP_HIDDEN wrote no file") + return 1 + ntok = 3 + want = ntok * nL * hc * hid * 4 + got = os.path.getsize(dump) + if got != want: + print(f"FAIL dump {got} bytes, want {want} " + f"({ntok} tok x {nL} layers x {hc} streams x {hid} hid x 4)") + return 1 + print(f"DUMP OK {got} bytes ({ntok}x{nL}x{hc}x{hid})") + return 0 + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_qwen_oracle.py b/tests/test_qwen_oracle.py new file mode 100644 index 000000000..5af2e03b3 --- /dev/null +++ b/tests/test_qwen_oracle.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""Streaming official Qwen oracle must not allocate the 360 GB tables. + + python3 tests/test_qwen_oracle.py +""" +from __future__ import annotations + +import os +import subprocess +import sys + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(ROOT, "tools")) + + +def ck(cond, what): + print(f" {'ok ' if cond else 'FAIL'} {what}") + if not cond: + raise SystemExit(1) + + +def main(): + src = os.path.join(ROOT, "tools", "qwen_ref.py") + text = open(src).read() + ck("--src" in text and "--dump" in text, + "qwen_ref.py CLI has --src and --dump") + ck("from_pretrained" not in text or "must not" in text.lower(), + "common-weight path does not call from_pretrained") + + import qwen_ref as Q + ck(hasattr(Q, "StreamingExperts"), "StreamingExperts exists") + ck(hasattr(Q, "StreamingNGram"), "StreamingNGram exists") + + class Fake: + pass + + # A 512-expert Parameter is 3.36 GiB. The streaming class must not + # construct it even when the official config says num_experts=512. + cfg = Fake() + cfg.num_experts = 512 + cfg.hidden_size = 2560 + cfg.moe_intermediate_size = 640 + cfg.hidden_act = "silu" + experts = Q.StreamingExperts(cfg, st=None, layer=0) + params = list(experts.parameters()) + n = sum(p.numel() for p in params) + ck(n == 0, f"StreamingExperts holds no parameters (got {n})") + ck(not hasattr(experts, "gate_up_proj") or experts.gate_up_proj is None, + "StreamingExperts does not own gate_up_proj") + + ncfg = Fake() + ncfg.ngram_size = 3 + ncfg.heads_per_ngram = 8 + ncfg.vocab_size = 248320 + ncfg.ngram_vocab_size_base = 20000000 + ncfg.seed = 1234 + ncfg.eos_token_id = 248044 + ncfg.make_ngram_vocab_size_divisible_by = 128 + ngram = Q.StreamingNGram(ncfg, 2560, layer_idx=1, ple_layer_index=0, st=None) + nparams = sum(p.numel() for p in ngram.parameters()) + ck(nparams == 0, f"StreamingNGram holds no embedding table (got {nparams})") + ck("if attn_mask is None" in text, + "QSA layers get a 4D mask when create_causal_mask returns None") + ck("torch.tril" in text, + "the fallback mask is causal visible=True (tril), not triu masked-out") + print("ORACLE STREAM OK") + return 0 + + +if __name__ == "__main__": + try: + sys.exit(main()) + except SystemExit: + raise + except Exception as e: + print(f"FAIL {type(e).__name__}: {e}") + sys.exit(1) diff --git a/tests/test_qwen_ple_write.py b/tests/test_qwen_ple_write.py new file mode 100644 index 000000000..999c0c295 --- /dev/null +++ b/tests/test_qwen_ple_write.py @@ -0,0 +1,68 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""Q8G PLE heads must quantize in row batches, not as one 12 GiB tensor. + + uv run --with torch --no-project python tests/test_qwen_ple_write.py +""" +from __future__ import annotations + +import io +import os +import sys + +ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) +sys.path.insert(0, os.path.join(ROOT, "tools")) + +try: + import torch +except ImportError: + print("SKIP: torch is not installed") + raise SystemExit(77) + +import convert as C # noqa: E402 + + +def ck(cond, what): + print(f" {'ok ' if cond else 'FAIL'} {what}") + if not cond: + raise SystemExit(1) + + +def main(): + torch.manual_seed(0) + W = torch.randn(96, 160) + q_full, sc_full, shape = C.quantize_q8g(W) + ck(shape == [96, 160], f"shape {shape}") + + buf = io.BytesIO() + meta = C.write_q8g_row_chunks(buf, [W[:32], W[32:64], W[64:]], 160) + ck(meta["shape"] == [96, 160], f"streamed shape {meta['shape']}") + blob = buf.getvalue() + q_n = q_full.numel() + sc_n = sc_full.numel() * 2 # fp16 + ck(meta["off"] == 0, "payload starts at 0") + ck(meta["scale_off"] == q_n, f"scales follow int8 ({meta['scale_off']} vs {q_n})") + ck(len(blob) == q_n + sc_n, f"bytes {len(blob)} vs {q_n + sc_n}") + ck(blob[:q_n] == C.raw_bytes(q_full), "chunked Q8G payload matches full-head") + ck(blob[q_n:] == C.raw_bytes(sc_full), "chunked Q8G scales match full-head") + + q8 = torch.arange(256, dtype=torch.int8) + got = C.raw_bytes(q8) + ck(got == bytes(range(256)), "raw_bytes of int8 is the storage bytes") + + from verify_container import dequant_q8g, dequant_q8g_row + full = dequant_q8g(C.raw_bytes(q_full), C.raw_bytes(sc_full), [96, 160]) + pad = 256 # 160 padded to 2*128 + q_b = C.raw_bytes(q_full) + sc_b = C.raw_bytes(sc_full) + row = dequant_q8g_row(q_b[5 * pad:6 * pad], sc_b[5 * 4:6 * 4], 160) + ck(torch.allclose(row, full[5], atol=1e-5), + "dequant_q8g_row matches a slice of the full-head dequant") + + print("PLE WRITE OK") + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_qwen_roundtrip.py b/tests/test_qwen_roundtrip.py new file mode 100644 index 000000000..eb6d8506d --- /dev/null +++ b/tests/test_qwen_roundtrip.py @@ -0,0 +1,174 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2026 SQLite Cloud, Inc. +"""Qwen conversion round-trip on a tiny packed+PLE source. + +Needs torch. Exit 77 if it is missing, matching tests/run.sh. + + python3 tests/test_qwen_roundtrip.py +""" +import json +import os +import shutil +import struct +import subprocess +import sys +import tempfile + +try: + import torch +except ImportError: + print("SKIP: torch is not installed") + raise SystemExit(77) + +REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) + + +def tensor_bytes(t): + """Raw little-endian payload. uv's `--with torch` env has no NumPy.""" + t = t.detach().cpu().contiguous() + n = t.numel() * t.element_size() + beg = t.storage_offset() * t.element_size() + return bytes(t.untyped_storage())[beg:beg + n] + + +def write_st(path, tensors): + """tensors: name -> (dtype_str, tensor).""" + header, payload = {}, bytearray() + for name, (dtype, t) in tensors.items(): + raw = t.detach().cpu().contiguous() + if dtype == "BF16" and raw.dtype != torch.bfloat16: + raise TypeError(f"{name} labeled BF16 but tensor is {raw.dtype}") + blob = tensor_bytes(raw) + start = len(payload) + payload.extend(blob) + header[name] = {"dtype": dtype, "shape": list(raw.shape), + "data_offsets": [start, len(payload)]} + hb = json.dumps(header, separators=(",", ":")).encode() + while (8 + len(hb)) % 8: + hb += b" " + with open(path, "wb") as f: + f.write(struct.pack(" 1 and all(0x30 <= b <= 0x39 for b in raw): + yield raw.decode() + + +def digit_run_changes_pieces(): + r"""The flag on a vocabulary that *does* hold a multi-digit token. + + Model-independent: 256 single bytes plus "20". Under \p{N}{1,3} the + pre-token is "202" and BPE reaches the "20" merge, so "2026" is three + ids; under \p{N} every digit is its own pre-token and it is four. + """ + import base64 + tmp = tempfile.mkdtemp(prefix="digitrun-") + try: + lines = [base64.b64encode(bytes([b])).decode() + f" {b}" + for b in range(256)] + lines.append(base64.b64encode(b"20").decode() + " 256") + with open(os.path.join(tmp, "tokenizer.model"), "w", + encoding="utf-8", newline="\n") as f: + f.write("\n".join(lines) + "\n") + three = encode_c(tmp, "2026", plain=True, digits=3) + one = encode_c(tmp, "2026", plain=True, digits=1) + ck(three == [256, 0x32, 0x36], + f"\\p{{N}}{{1,3}} merges 20 out of 202: {three}") + ck(one == [0x32, 0x30, 0x32, 0x36], + f"\\p{{N}} keeps every digit apart: {one}") + finally: + shutil.rmtree(tmp, ignore_errors=True) + + +def main(): + print("Qwen's pattern is one src/tokenizer.c implements") + ck(QWEN_PATTERN in HFT.KNOWN_PATTERNS, + "hf_tokenizer recognises the Qwen spelling") + han, digit_run = HFT.KNOWN_PATTERNS[QWEN_PATTERN] + ck(han == 0, "Han has no branch of its own, so tokenizer_han_split is false") + ck(digit_run == 1, f"one digit per piece ({digit_run})") + ck(HFT.KNOWN_PATTERNS[HFT.PAT_HAN] == (1, 3), "Kimi's is unchanged") + ck(HFT.KNOWN_PATTERNS[HFT.PAT_NO_HAN] == (0, 3), "GLM's is unchanged") + + if os.path.exists(os.path.join(ROOT, "test_tokenizer")): + print("tokenizer_digit_run changes where a piece ends") + digit_run_changes_pieces() + + if not os.path.isfile(os.path.join(PINNED, "tokenizer.json")): + print("SKIP the parity half: set WASTE_QWEN_SRC to the pinned " + "checkpoint directory" if not PINNED else + f"SKIP the parity half: no tokenizer.json under {PINNED}") + return 77 if not fails else 1 + if not os.path.exists(os.path.join(ROOT, "test_tokenizer")): + print("SKIP the parity half: test_tokenizer is not built (make test)") + return 77 if not fails else 1 + try: + from tokenizers import Tokenizer + except ImportError: + print("SKIP the parity half: the `tokenizers` package is not installed") + return 77 if not fails else 1 + + print("C tokenizer vs the release's own") + text, han, specials, digit_run = HFT.convert(PINNED, quiet=True) + ck(not han and digit_run == 1, + f"the pinned checkpoint reports han={han} digit_run={digit_run}") + hf = Tokenizer.from_file(os.path.join(PINNED, "tokenizer.json")) + tmp = tempfile.mkdtemp(prefix="qwen-tok-") + try: + with open(os.path.join(tmp, "tokenizer.model"), "w", + encoding="utf-8", newline="\n") as f: + f.write(text) + if specials: + json.dump(specials, open(os.path.join(tmp, "specials.json"), "w"), + indent=1) + for s in STRINGS: + got = encode_c(tmp, s, plain=True, digits=digit_run) + want = hf.encode(s, add_special_tokens=False).ids + ck(got == want, f"{s!r} -> {got if got != want else got[:12]}" + + ("" if got == want else f" want {want}")) + + # This vocabulary has no multi-digit token at all — it was trained + # under \p{N} — so the two runs happen to agree on it: "202" has + # no merge to reach and comes back out as three ids either way. + # That is a property of Qwen's vocabulary, not of the setting, and + # it is why the flag is checked on its own below rather than here. + multi = list(_digit_tokens(text)) + ck(not multi, f"no multi-digit token in the Qwen vocabulary {multi[:5]}") + + # Markup vs content: a control token must not be forgeable from + # ordinary text. See waste_tokenize / waste_tokenize_markup. + marker = next((e["text"] for e in specials + if e["text"].startswith("<|") and e["text"].endswith("|>")), + None) + if marker: + as_markup = encode_c(tmp, marker, plain=False, digits=digit_run) + as_plain = encode_c(tmp, marker, plain=True, digits=digit_run) + ck(as_markup != as_plain, + f"markup {marker!r} {as_markup} is not what content gives") + ck(len(as_markup) == 1, f"markup {marker!r} is one control id") + finally: + shutil.rmtree(tmp, ignore_errors=True) + + print("QWEN TOK FAILED" if fails else "QWEN TOK OK") + return 1 if fails else 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_qwenparts.c b/tests/test_qwenparts.c new file mode 100644 index 000000000..ae3e42f33 --- /dev/null +++ b/tests/test_qwenparts.c @@ -0,0 +1,370 @@ +/* SPDX-License-Identifier: Apache-2.0 + * Copyright 2026 SQLite Cloud, Inc. + */ +/* + * test_qwenparts.c — isolated Qwen ops vs tools/qwenparts_ref.py. + * + * ./test_qwenparts out.bin + * uv run --with torch --no-project python tools/qwenparts_ref.py out.bin + * + * Layout (little-endian; f32 unless noted): + * [ple] i32 T, i32 ngram, i32 heads, i32 eos, i32 vocab, i32 seed, + * i32 pos, i32 ids[T], i64 multipliers[ngram], + * i64 sizes[heads], i32 local_rows[heads], + * i32 eos_T, i32 eos_ids[eos_T], i32 eos_shift, + * i32 eos_shifted[eos_T], + * i32 q8_cols, i32 q8_group, i8 q[pad], f16 scales[ng], f32 deq[cols] + * [hc] i32 hc, i32 hid, i32 rank, f32 eps, hyper[hc*hid], nw[hc*hid], + * down[rank*hc*hid], up[hc*hid*rank], inject[hc*hc*hid], + * mixed[hid], inj_w[hc], combined[hc*hid] + * [gdn] i32 T, i32 Hk, i32 Hv, i32 Dk, i32 Dv, + * q[T*Hk*Dk], k[T*Hk*Dk], v[T*Hv*Dv], a[T*Hv], A_log[Hv], dt[Hv], + * o_step[Hv*Dv] (last token decode), o_fwd[T*Hv*Dv], o_chunk[T*Hv*Dv] + * [qsa] i32 T, i32 Hq, i32 Hkv, i32 D, i32 Dk, i32 compress, i32 topk, + * i32 query_pos, q_idx[Hq*Dk], raw_k[T*Dk], k_ln[Dk], + * cos[T*rot], sin[T*rot], sel_n, i32 sel[sel_n], + * q[Hq*D], k[T*Hkv*D], v[T*Hkv*D], attn[Hq*D] + * [moe] i32 E, i32 K, logits[E], i32 idx[K], w[K], shared_gate, + * shared[H], routed[H], out[H] (H = 8 in the dump) + */ + +#include +#include +#include +#include +#include + +#include "../src/model.h" +#include "../src/qwen_gdn.h" +#include "../src/qwen_hc.h" +#include "../src/qwen_moe.h" +#include "../src/qwen_ple.h" +#include "../src/qwen_qsa.h" + +static uint64_t rng = 0x243F6A8885A308D3ULL; +static float frand(void) +{ + rng = rng * 6364136223846793005ULL + 1442695040888963407ULL; + return (float)((double)((rng >> 11) & 0x1FFFFFFFFFFFFFULL) / 9007199254740992.0 + * 2.0 - 1.0); +} + +static void wr(FILE *f, const float *v, int n) { fwrite(v, sizeof(float), (size_t)n, f); } +static void wi(FILE *f, int v) { fwrite(&v, sizeof(int), 1, f); } +static void wf(FILE *f, float v) { fwrite(&v, sizeof(float), 1, f); } +static void w64(FILE *f, int64_t v) { fwrite(&v, sizeof(int64_t), 1, f); } + +static uint16_t f32_to_f16(float x) +{ + union { float f; uint32_t u; } a; + a.f = x; + uint32_t u = a.u; + uint32_t sign = (u >> 16) & 0x8000u; + int32_t exp = (int32_t)((u >> 23) & 0xff) - 127 + 15; + uint32_t man = u & 0x7fffffu; + if (exp <= 0) { + if (exp < -10) return (uint16_t)sign; + man = (man | 0x800000u) >> (1 - exp); + return (uint16_t)(sign | (man >> 13)); + } + if (exp >= 31) return (uint16_t)(sign | 0x7c00u); + return (uint16_t)(sign | ((uint32_t)exp << 10) | (man >> 13)); +} + +int main(int argc, char **argv) +{ + const char *out = argc > 1 ? argv[1] : "qwenparts.bin"; + FILE *f = fopen(out, "wb"); + if (!f) { perror("open"); return 1; } + + /* ---- 1. PLE hashing, EOS reset, Q8G row ---------------------------- */ + { + const int T = 8, ngram = 3, heads = 16, eos = 2, vocab = 256, seed = 0; + const int pos = 7; + int ids[8] = { 11, 13, 17, 19, 23, 29, 31, 37 }; + int64_t mult[3], sizes[16]; + int local[16]; + waste_qwen_ple_multipliers(mult, ngram, 0, seed, vocab); + for (int h = 0; h < heads; h++) + sizes[h] = waste_qwen_ple_nth_prime_after(64 - 1, h + 1); + waste_qwen_ple_row_ids(ids, T, pos, eos, ngram, 8, mult, sizes, local); + + wi(f, T); wi(f, ngram); wi(f, heads); wi(f, eos); wi(f, vocab); wi(f, seed); + wi(f, pos); + for (int i = 0; i < T; i++) wi(f, ids[i]); + for (int i = 0; i < ngram; i++) w64(f, mult[i]); + for (int h = 0; h < heads; h++) w64(f, sizes[h]); + for (int h = 0; h < heads; h++) wi(f, local[h]); + + const int eos_T = 6, eos_shift = 2; + int eos_ids[6] = { 5, 7, 2, 9, 11, 13 }; + int shifted[6]; + waste_qwen_ple_shift_eos(eos_ids, eos_T, eos_shift, eos, shifted); + wi(f, eos_T); + for (int i = 0; i < eos_T; i++) wi(f, eos_ids[i]); + wi(f, eos_shift); + for (int i = 0; i < eos_T; i++) wi(f, shifted[i]); + + const int cols = 160, group = 128; + const int ng = (cols + group - 1) / group; + const int pad = ng * group; + int8_t *q = calloc((size_t)pad, 1); + uint16_t *sc = calloc((size_t)ng, 2); + float *deq = calloc((size_t)cols, 4); + for (int i = 0; i < cols; i++) q[i] = (int8_t)((i * 13) % 241 - 120); + sc[0] = f32_to_f16(0.01f); + sc[1] = f32_to_f16(0.02f); + waste_tensor t; + memset(&t, 0, sizeof t); + t.q = q; t.qs = sc; t.group = group; t.bits = 8; t.rowbytes = (size_t)pad; + t.shape[0] = 1; t.shape[1] = cols; t.ndim = 2; t.n = (size_t)cols; + waste_deq_row(&t, 0, cols, deq); + wi(f, cols); wi(f, group); + fwrite(q, 1, (size_t)pad, f); + fwrite(sc, 2, (size_t)ng, f); + wr(f, deq, cols); + free(q); free(sc); free(deq); + } + + /* ---- 2. HyperConnection Mix / Combine ------------------------------ */ + { + const int hc = 4, hid = 8, rank = 4; + const float eps = 1e-6f; + const int H = hc * hid; + float *hyper = malloc((size_t)H * 4), *nw = malloc((size_t)H * 4); + float *down = malloc((size_t)rank * H * 4), *up = malloc((size_t)H * rank * 4); + float *inject = malloc((size_t)hc * H * 4); + float *mixed = malloc((size_t)hid * 4), *inj_w = malloc((size_t)hc * 4); + float *comb = malloc((size_t)H * 4), *block = malloc((size_t)hid * 4); + float *scratch = malloc((size_t)(2 * H + rank + hc + 16) * 4); + for (int i = 0; i < H; i++) { hyper[i] = frand(); nw[i] = frand() * 0.1f; } + for (int i = 0; i < rank * H; i++) down[i] = frand() * 0.2f; + for (int i = 0; i < H * rank; i++) up[i] = frand() * 0.2f; + for (int i = 0; i < hc * H; i++) inject[i] = frand() * 0.2f; + for (int i = 0; i < hid; i++) block[i] = frand(); + waste_qwen_hc_gates(hyper, nw, down, up, inject, hc, hid, rank, eps, + mixed, inj_w, scratch); + waste_qwen_hc_combine(hyper, block, inj_w, hc, hid, comb); + wi(f, hc); wi(f, hid); wi(f, rank); wf(f, eps); + wr(f, hyper, H); wr(f, nw, H); wr(f, down, rank * H); wr(f, up, H * rank); + wr(f, inject, hc * H); wr(f, block, hid); + wr(f, mixed, hid); wr(f, inj_w, hc); wr(f, comb, H); + free(hyper); free(nw); free(down); free(up); free(inject); + free(mixed); free(inj_w); free(comb); free(block); free(scratch); + } + + /* ---- 3. GDN decode + chunked prefill ------------------------------- */ + { + const int T = 5, Hk = 2, Hv = 6, Dk = 8, Dv = 8; + float *q = malloc((size_t)T * Hk * Dk * 4); + float *k = malloc((size_t)T * Hk * Dk * 4); + float *v = malloc((size_t)T * Hv * Dv * 4); + float *a = malloc((size_t)T * Hv * 4); + float *A = malloc((size_t)Hv * 4), *dt = malloc((size_t)Hv * 4); + float *g = malloc((size_t)T * Hv * 4), *beta = malloc((size_t)T * Hv * 4); + float *S1 = calloc((size_t)Hv * Dk * Dv, 4); + float *S2 = calloc((size_t)Hv * Dk * Dv, 4); + float *S3 = calloc((size_t)Hv * Dk * Dv, 4); + float *o_step = malloc((size_t)Hv * Dv * 4); + float *o_fwd = malloc((size_t)T * Hv * Dv * 4); + float *o_chunk = malloc((size_t)T * Hv * Dv * 4); + float *scratch = malloc((size_t)(Hv * Dk * Dv * 8 + 4096) * 4); + for (int i = 0; i < T * Hk * Dk; i++) { q[i] = frand(); k[i] = frand(); } + for (int i = 0; i < T * Hv * Dv; i++) v[i] = frand(); + for (int i = 0; i < T * Hv; i++) a[i] = frand(); + for (int i = 0; i < Hv; i++) { A[i] = frand() * 2.0f; dt[i] = frand(); } + for (int t = 0; t < T; t++) { + waste_qwen_gdn_decay(a + t * Hv, A, dt, Hv, g + t * Hv); + for (int h = 0; h < Hv; h++) + beta[t * Hv + h] = 1.0f / (1.0f + expf(-a[t * Hv + h] * 0.5f)); + } + /* last-token decode from zeros, then full forward and chunk */ + waste_qwen_gdn_step(Hk, Hv, Dk, Dv, + q + (T - 1) * Hk * Dk, k + (T - 1) * Hk * Dk, + v + (T - 1) * Hv * Dv, g + (T - 1) * Hv, + beta + (T - 1) * Hv, S1, o_step, scratch); + waste_qwen_gdn_forward(T, Hk, Hv, Dk, Dv, q, k, v, g, beta, S2, o_fwd, scratch); + waste_qwen_gdn_chunk(T, Hk, Hv, Dk, Dv, 64, q, k, v, g, beta, S3, o_chunk, scratch); + wi(f, T); wi(f, Hk); wi(f, Hv); wi(f, Dk); wi(f, Dv); + wr(f, q, T * Hk * Dk); wr(f, k, T * Hk * Dk); wr(f, v, T * Hv * Dv); + wr(f, a, T * Hv); wr(f, A, Hv); wr(f, dt, Hv); + wr(f, g, T * Hv); wr(f, beta, T * Hv); + wr(f, o_step, Hv * Dv); wr(f, o_fwd, T * Hv * Dv); wr(f, o_chunk, T * Hv * Dv); + free(q); free(k); free(v); free(a); free(A); free(dt); free(g); free(beta); + free(S1); free(S2); free(S3); free(o_step); free(o_fwd); free(o_chunk); + free(scratch); + } + + /* ---- 4. QSA indexer + attention ------------------------------------ */ + { + const int T = 10, Hq = 4, Hkv = 2, D = 8, Dk = 8, compress = 4, topk = 2; + const int query_pos = 9, rot = 4; + float *q_idx = malloc((size_t)Hq * Dk * 4); + float *raw_k = malloc((size_t)T * Dk * 4); + float *k_ln = malloc((size_t)Dk * 4); + float *cos = malloc((size_t)T * rot * 4), *sinv = malloc((size_t)T * rot * 4); + float *q = malloc((size_t)Hq * D * 4); + float *k = malloc((size_t)T * Hkv * D * 4); + float *v = malloc((size_t)T * Hkv * D * 4); + float *attn = malloc((size_t)Hq * D * 4); + float *scratch = malloc((size_t)(T * Hq + Hq * D + 256) * 4); + int sel[32]; + int n_complete = (query_pos + 1) / compress; + float *work = malloc((size_t)(n_complete * Dk + n_complete + Dk) * 4); + int *taken = malloc((size_t)(n_complete > 0 ? n_complete : 1) * sizeof(int)); + for (int i = 0; i < Hq * Dk; i++) q_idx[i] = frand(); + for (int i = 0; i < T * Dk; i++) raw_k[i] = frand(); + for (int i = 0; i < Dk; i++) k_ln[i] = frand() * 0.1f; + for (int t = 0; t < T; t++) + for (int r = 0; r < rot; r++) { + float ang = (float)t * 0.1f * (float)(r + 1); + cos[t * rot + r] = cosf(ang); + sinv[t * rot + r] = sinf(ang); + } + for (int i = 0; i < Hq * D; i++) q[i] = frand(); + for (int i = 0; i < T * Hkv * D; i++) { k[i] = frand(); v[i] = frand(); } + int nsel = waste_qwen_qsa_select(q_idx, Hq, Dk, raw_k, T, query_pos, + cos, sinv, rot, k_ln, 1e-6f, + compress, topk, sel, work, taken); + waste_qwen_qsa_attn(q, Hq, D, k, v, Hkv, T, sel, nsel, + 1.0f / sqrtf((float)D), attn, scratch); + wi(f, T); wi(f, Hq); wi(f, Hkv); wi(f, D); wi(f, Dk); + wi(f, compress); wi(f, topk); wi(f, query_pos); wi(f, rot); + wr(f, q_idx, Hq * Dk); wr(f, raw_k, T * Dk); wr(f, k_ln, Dk); + wr(f, cos, T * rot); wr(f, sinv, T * rot); + wi(f, nsel); + for (int i = 0; i < nsel; i++) wi(f, sel[i]); + wr(f, q, Hq * D); wr(f, k, T * Hkv * D); wr(f, v, T * Hkv * D); + wr(f, attn, Hq * D); + free(q_idx); free(raw_k); free(k_ln); free(cos); free(sinv); + free(q); free(k); free(v); free(attn); free(scratch); + free(work); free(taken); + } + + /* ---- 5. Softmax top-k MoE, original router order ------------------- */ + { + const int E = 16, K = 4, H = 8; + float logits[16], w[4], shared[8], routed[8], out[8], prob[16]; + int idx[4]; + uint8_t used[16]; + float gate_in = 0.4f; + for (int e = 0; e < E; e++) logits[e] = frand() * 3.0f; + for (int i = 0; i < H; i++) { shared[i] = frand(); routed[i] = 0.0f; } + waste_qwen_moe_route(logits, E, K, 1, idx, w, prob, used); + /* fake expert outputs in router order: expert j contributes w[j]*e_j + * where e_j is a rank-1 pattern of the expert id, so order matters. */ + for (int j = 0; j < K; j++) + for (int i = 0; i < H; i++) + routed[i] += w[j] * ((float)(idx[j] + 1) * 0.1f + (float)i * 0.01f); + const float sg = 1.0f / (1.0f + expf(-gate_in)); + for (int i = 0; i < H; i++) out[i] = routed[i] + sg * shared[i]; + wi(f, E); wi(f, K); wi(f, H); + wr(f, logits, E); + for (int j = 0; j < K; j++) wi(f, idx[j]); + wr(f, w, K); + wf(f, gate_in); + wr(f, shared, H); wr(f, routed, H); wr(f, out, H); + } + + /* ---- official geometry, bounded allocations ----------------------- */ + { + const int hc = 4, hid = 8, rank = 320; + const float eps = 1e-6f; + const int H = hc * hid; + float *hyper = malloc((size_t)H * 4), *nw = malloc((size_t)H * 4); + float *down = malloc((size_t)rank * H * 4), *up = malloc((size_t)H * rank * 4); + float *inject = malloc((size_t)hc * H * 4); + float *mixed = malloc((size_t)hid * 4), *inj_w = malloc((size_t)hc * 4); + float *comb = malloc((size_t)H * 4), *block = malloc((size_t)hid * 4); + float *scratch = malloc((size_t)(2 * H + rank + hc + 16) * 4); + for (int i = 0; i < H; i++) { hyper[i] = frand(); nw[i] = frand() * 0.1f; } + for (int i = 0; i < rank * H; i++) down[i] = frand() * 0.05f; + for (int i = 0; i < H * rank; i++) up[i] = frand() * 0.05f; + for (int i = 0; i < hc * H; i++) inject[i] = frand() * 0.05f; + for (int i = 0; i < hid; i++) block[i] = frand(); + waste_qwen_hc_gates(hyper, nw, down, up, inject, hc, hid, rank, eps, + mixed, inj_w, scratch); + waste_qwen_hc_combine(hyper, block, inj_w, hc, hid, comb); + wi(f, hc); wi(f, hid); wi(f, rank); wf(f, eps); + wr(f, hyper, H); wr(f, nw, H); wr(f, down, rank * H); wr(f, up, H * rank); + wr(f, inject, hc * H); wr(f, block, hid); + wr(f, mixed, hid); wr(f, inj_w, hc); wr(f, comb, H); + free(hyper); free(nw); free(down); free(up); free(inject); + free(mixed); free(inj_w); free(comb); free(block); free(scratch); + } + { + const int T = 4, Hk = 16, Hv = 48, Dk = 128, Dv = 128; + float *q = malloc((size_t)T * Hk * Dk * 4); + float *k = malloc((size_t)T * Hk * Dk * 4); + float *v = malloc((size_t)T * Hv * Dv * 4); + float *a = malloc((size_t)T * Hv * 4); + float *A = malloc((size_t)Hv * 4), *dt = malloc((size_t)Hv * 4); + float *g = malloc((size_t)T * Hv * 4), *beta = malloc((size_t)T * Hv * 4); + float *S2 = calloc((size_t)Hv * Dk * Dv, 4); + float *S3 = calloc((size_t)Hv * Dk * Dv, 4); + float *o_fwd = malloc((size_t)T * Hv * Dv * 4); + float *o_chunk = malloc((size_t)T * Hv * Dv * 4); + float *scratch = malloc((size_t)(Dv + 16) * 4); + for (int i = 0; i < T * Hk * Dk; i++) { q[i] = frand(); k[i] = frand(); } + for (int i = 0; i < T * Hv * Dv; i++) v[i] = frand(); + for (int i = 0; i < T * Hv; i++) a[i] = frand(); + for (int i = 0; i < Hv; i++) { A[i] = frand() * 2.0f; dt[i] = frand(); } + for (int t = 0; t < T; t++) { + waste_qwen_gdn_decay(a + t * Hv, A, dt, Hv, g + t * Hv); + for (int h = 0; h < Hv; h++) + beta[t * Hv + h] = 1.0f / (1.0f + expf(-a[t * Hv + h] * 0.5f)); + } + waste_qwen_gdn_forward(T, Hk, Hv, Dk, Dv, q, k, v, g, beta, S2, o_fwd, scratch); + waste_qwen_gdn_chunk(T, Hk, Hv, Dk, Dv, 64, q, k, v, g, beta, S3, o_chunk, scratch); + wi(f, T); wi(f, Hk); wi(f, Hv); wi(f, Dk); wi(f, Dv); + wr(f, o_fwd, T * Hv * Dv); wr(f, o_chunk, T * Hv * Dv); + free(q); free(k); free(v); free(a); free(A); free(dt); free(g); free(beta); + free(S2); free(S3); free(o_fwd); free(o_chunk); free(scratch); + } + { + const int T = 2051, Hq = 2, Dk = 8, compress = 4, topk = 512, rot = 4; + const int query_pos = T - 1; + float *q_idx = malloc((size_t)Hq * Dk * 4); + float *raw_k = malloc((size_t)T * Dk * 4); + float *k_ln = malloc((size_t)Dk * 4); + float *cos = malloc((size_t)T * rot * 4), *sinv = malloc((size_t)T * rot * 4); + int *sel = malloc((size_t)(topk * compress + compress) * sizeof(int)); + int n_complete = (query_pos + 1) / compress; + float *work = malloc((size_t)(n_complete * Dk + n_complete + Dk) * 4); + int *taken = malloc((size_t)n_complete * sizeof(int)); + for (int i = 0; i < Hq * Dk; i++) q_idx[i] = frand(); + for (int i = 0; i < T * Dk; i++) raw_k[i] = frand(); + for (int i = 0; i < Dk; i++) k_ln[i] = frand() * 0.1f; + for (int t = 0; t < T; t++) + for (int r = 0; r < rot; r++) { + float ang = (float)t * 0.1f * (float)(r + 1); + cos[t * rot + r] = cosf(ang); + sinv[t * rot + r] = sinf(ang); + } + int nsel = waste_qwen_qsa_select(q_idx, Hq, Dk, raw_k, T, query_pos, + cos, sinv, rot, k_ln, 1e-6f, + compress, topk, sel, work, taken); + float pooled[8]; + waste_qwen_qsa_pool_block(raw_k, 4, Dk, k_ln, 1e-6f, pooled); + wi(f, T); wi(f, Hq); wi(f, Dk); wi(f, compress); wi(f, topk); + wi(f, nsel); wr(f, pooled, Dk); wr(f, k_ln, Dk); wr(f, raw_k, 4 * Dk); + free(q_idx); free(raw_k); free(k_ln); free(cos); free(sinv); + free(sel); free(work); free(taken); + } + { + const int E = 16, K = 10; + float logits[16], w[10], prob[16]; + int idx[10]; + uint8_t used[16]; + for (int e = 0; e < E; e++) logits[e] = frand() * 3.0f; + waste_qwen_moe_route(logits, E, K, 1, idx, w, prob, used); + wi(f, E); wi(f, K); + wr(f, logits, E); + for (int j = 0; j < K; j++) wi(f, idx[j]); + wr(f, w, K); + } + + fclose(f); + printf("wrote %s\n", out); + return 0; +} diff --git a/tests/test_tokenizer.c b/tests/test_tokenizer.c index 6f5295cd3..8e174bcfc 100644 --- a/tests/test_tokenizer.c +++ b/tests/test_tokenizer.c @@ -30,6 +30,10 @@ int main(int argc, char **argv) * to reach the GLM pattern. See waste_tok_set_han_split. */ const char *nohan = getenv("WASTE_TOK_NOHAN"); if (nohan && *nohan != '0') waste_tok_set_han_split(t, 0); + /* Same for the digit run: 1 is Qwen's `\p{N}`, 3 the default + * `\p{N}{1,3}`. See waste_tok_set_digit_run. */ + const char *drun = getenv("WASTE_TOK_DIGITS"); + if (drun) waste_tok_set_digit_run(t, atoi(drun)); /* Same reason, for the release that needs a different pattern rather * than the same one without its Han branch. */ const char *pat = getenv("WASTE_TOK_PATTERN"); diff --git a/tools/convert.py b/tools/convert.py index d64eee09c..d04cfd777 100644 --- a/tools/convert.py +++ b/tools/convert.py @@ -15,7 +15,11 @@ Reads a Kimi checkpoint as published — the 1.42 TB of moonshotai/Kimi-K3 that tools/fetch_weights.sh leaves on the staging disk, or any other -member of the family (Kimi-Linear) by pointing --src elsewhere. +member of the family (Kimi-Linear) by pointing --src elsewhere. Qwen +`qwen4_exp_text` is also accepted: nested `text_config` is flattened, +packed `experts.gate_up_proj` / `down_proj` split into WEXP records, and +128 PLE n-gram shards become 16 Q8G heads. `mtp.*` and `model.visual.*` +are skipped. Format v0 and the WEXP record do not change. uv run --with torch python tools/convert.py \ --src /path/to/hf-checkpoint \ @@ -42,6 +46,7 @@ """ import argparse +import gc import io import json import os @@ -183,6 +188,180 @@ def _load_vq(): # renormalisation on for a checkpoint that sets it false. Emit only when true. CONFIG_FLAG_ALIASES = (("moe_renormalize", "norm_topk_prob"),) +# Qwen3.8-Flash-Next packs every routed expert of a layer into two tensors. +# The shapes below are what the pinned release ships and are asserted by the +# source-sample test; nothing in the converter branches on them. What the +# conversion actually requires is the *layout* — gate_up [E, 2I, H] beside +# down [E, H, I] — and that is checked by packed_shapes_ok, so a smaller +# fixture or a larger family member converts on the same path. +QWEN_TEXT_TYPE = "qwen4_exp_text" +QWEN_PACKED_GATE_UP = (512, 1280, 2560) +QWEN_PACKED_DOWN = (512, 2560, 640) +PLE_HEADS = 16 +PLE_HEAD_WIDTH = 160 + + +def is_qwen(cfg): + """Qwen3.8-Flash-Next (`qwen4_exp`), by the name it gives itself. + + The text model is `qwen4_exp_text` and it is what this converter reads; + the outer `qwen4_exp` wraps it beside a vision tower this container does + not carry. Either spelling identifies the family, because the flattened + config keeps the inner `model_type` and the raw one has only the outer.""" + inner = cfg.get("text_config") or {} + for mt in (cfg.get("model_type"), inner.get("model_type")): + if mt in (QWEN_TEXT_TYPE, "qwen4_exp"): + return True + hf = ((cfg.get("_outer", {}).get("architectures") + or cfg.get("architectures") or [""]))[0] + return "Qwen4Exp" in hf + + +def qwen_rename(name): + """Qwen wraps the text model the way GLM does, not the way K3 does. + + Qwen model.language_model.layers.N.… lm_head.weight + + So the container is prefix-less and the wrapper is dropped here, for the + reason spelled out in glm_rename: the engine's lookup is a fixed string + and a container that spells a tensor differently holds weights nothing + will ever ask for. The vision tower is not renamed because it is not + carried at all — see qwen_drop_trunk.""" + pfx = "model.language_model." + if name.startswith(pfx): + return "model." + name[len(pfx):] + return name + + +def qwen_drop_trunk(): + """Which of Qwen's trunk tensors this container has no reader for. + + Three groups. The MTP layer is a speculative-decoding head and the + vision tower is out of the text-only scope, so neither has a reader and + carrying either would cost resident RAM the expert cache wants. The PLE + n-gram shards *are* read, but not by build_trunk: 128 source shards + become 16 logical heads and build_ple writes them, streaming, after the + trunk pass — so they are dropped here and picked up there.""" + def drop(name): + return bool(qwen_skip_tensor(name) or is_ple_ngram(name) + or is_ple_meta(name)) + return drop + + +def qwen_skip_tensor(name): + """Vision and the MTP layer are out of the text-only milestone.""" + return name.startswith("mtp.") or name.startswith("model.visual.") + + +def is_ple_ngram(name): + return "ngram_embedding.shard_" in name and name.endswith(".weight") + + +def is_ple_meta(name): + """I64 offset/size/multiplier tables. build_ple reads them after the trunk reclaim.""" + return "ple_embedding." in name and name.endswith( + ("ngram_heads_offsets", "ngram_heads_vocab_sizes", "layer_multipliers")) + + +def qwen_packed_names(src_pfx, layer): + """The two tensors that hold a whole layer's routed experts.""" + p = f"{src_pfx}layers.{layer}.mlp.experts." + return p + "gate_up_proj", p + "down_proj" + + +def split_packed_gate_up_shape(shape, inter): + """gate_up [E, 2I, H] splits into gate [E, I, H] and up [E, I, H].""" + e, two_i, h = shape + if two_i != 2 * inter: + raise ValueError(f"packed gate_up dim1 {two_i} is not 2*{inter}") + return (e, inter, h), (e, inter, h) + + +def packed_shapes_ok(gate_up_shape, down_shape, n_exp): + """True when the packed pair matches n_exp experts of SwiGLU width.""" + gs, ds = tuple(gate_up_shape), tuple(down_shape) + if len(gs) != 3 or len(ds) != 3: + return False + if gs[0] != n_exp or ds[0] != n_exp: + return False + if gs[1] % 2 or gs[1] // 2 != ds[2] or gs[2] != ds[1]: + return False + return True + + +def ple_head_slices(offsets, vocab_sizes): + """(start, n_rows) for each logical PLE head.""" + return [(int(offsets[i]), int(vocab_sizes[i])) for i in range(len(vocab_sizes))] + + +def ple_source_loc(global_row, shard_rows): + """Which of the 128 source shards holds `global_row`.""" + return divmod(int(global_row), int(shard_rows)) + + +# Rows per Q8G batch when writing a PLE head. 64 Ki rows × 160 × 4 B ≈ 40 MiB +# of f32; a full head is ~12 GiB and does not fit beside a 22 GiB trunk tmp. +PLE_Q8G_ROWS = 65536 + + +def write_q8g_row_chunks(tf, row_chunks, width, group=128): + """Write a Q8G matrix as independent row batches. + + quantize_q8g groups along the last dim, so the cat of per-chunk (q, + scale) equals quantize_q8g on the stacked rows. + """ + off = tf.tell() + scales = [] + n_rows = 0 + for chunk in row_chunks: + if chunk.dim() != 2 or int(chunk.shape[1]) != width: + raise ValueError( + f"Q8G chunk shape {tuple(chunk.shape)}, expected [N, {width}]") + x = chunk if chunk.dtype == torch.float32 else chunk.float() + q, sc, _shape = quantize_q8g(x) + tf.write(raw_bytes(q)) + scales.append(sc.cpu()) + n_rows += int(chunk.shape[0]) + del q, x + if not scales: + raise ValueError("Q8G write with no rows") + sc_off = tf.tell() + sc = scales[0] if len(scales) == 1 else torch.cat(scales, 0) + tf.write(raw_bytes(sc)) + return {"off": off, "scale_off": sc_off, "shape": [n_rows, width], + "group": group, "bytes": tf.tell() - off} + + +def iter_ple_head_rows(st, shards, shard_rows, start, n_rows, + batch=PLE_Q8G_ROWS): + """Yield f32 [take, width] slices of one logical PLE head. + + One source shard stays loaded (≈800 MiB BF16). The previous full-head + cat was ~12 GiB of f32 and swapped the 48 GB machine. + """ + row, end = int(start), int(start) + int(n_rows) + cached_si, cached = None, None + while row < end: + take = min(int(batch), end - row) + parts, need, pos = [], take, row + while need > 0: + si, local = ple_source_loc(pos, shard_rows) + n = min(int(shard_rows) - local, need) + if si != cached_si: + del cached + cached = st.raw(shards[si]) + cached_si = si + sl = cached[local:local + n] + parts.append(sl if sl.dtype == torch.float32 else sl.float()) + del sl + pos += n + need -= n + chunk = parts[0] if len(parts) == 1 else torch.cat(parts, 0) + del parts + yield chunk + row += take + del cached + def is_glm(cfg): """GLM-5.3-Flash, by the name it gives itself. @@ -380,6 +559,7 @@ def source_prefixes(cfg): Kimi-Linear model.layers.N.… prefix "" Kimi K3 language_model.model.layers.N.… prefix "language_model." GLM-5.3 model.language_model.layers.N.… prefix "" + Qwen3.8 model.language_model.layers.N.… prefix "" `src_pfx` is what the *checkpoint* puts before `layers.N` and is used to find the experts; `prefix` is what the engine puts before `model.` and @@ -390,15 +570,15 @@ def source_prefixes(cfg): Returns the inner config too, since the same test decides whether there is a wrapper to unwrap at all.""" prefix, src_pfx = "", "model." - if "text_config" in cfg: # K3, GLM and DS41 nest the text model - cfg = {**cfg["text_config"], "_outer": {k: v for k, v in cfg.items() - if k != "text_config"}} + if "text_config" in cfg: # K3, GLM, Qwen and DS41 nest the text model + outer = {k: v for k, v in cfg.items() if k != "text_config"} + cfg = {**cfg["text_config"], "_outer": outer} if is_ds41(cfg): # No wrapper in the tensor names at all: `layers.0.attn.wq_a`, # `embed.weight`, `head.weight`. Nothing to strip and nothing to # publish -- the whole mapping is ds41_rename. src_pfx = "" - elif is_glm(cfg): + elif is_glm(cfg) or is_qwen(cfg): src_pfx = "model.language_model." else: prefix, src_pfx = "language_model.", "language_model.model." @@ -476,6 +656,11 @@ def moe_layout(st, src_pfx, layer): `src_pfx` is everything the checkpoint puts before `layers.N` — which is not the container's tensor_prefix plus "model.", because GLM nests the two the other way round (see glm_rename).""" + # Qwen packs a whole layer's experts into two tensors rather than one + # per expert, so it is recognised by the packed pair and not by the + # per-expert probe below, which would find nothing. + if st.have(f"{src_pfx}layers.{layer}.mlp.experts.gate_up_proj"): + return "qwen_packed", "mlp", None for name, seg, tags in MOE_LAYOUTS: probe = f"{src_pfx}layers.{layer}.{seg}.experts.0.{tags[0]}.weight" if st.have(probe) or st.have(probe + "_packed"): @@ -537,7 +722,7 @@ def get(self, name): f.seek(base + beg) raw = f.read(end - beg) dt = {"BF16": torch.bfloat16, "F16": torch.float16, - "F32": torch.float32, + "F32": torch.float32, "I64": torch.int64, # The current generation of large MoEs ships fp8 with one f32 # scale per weight_block_size tile in a companion tensor (K2, # DeepSeek V3/R1). Reading the values without applying those @@ -742,12 +927,17 @@ def quantize_q8g(W, group=128): # ---------------------------------------------------------------- writing -- def raw_bytes(t): - """Contiguous little-endian bytes of a tensor (torch has no .tobytes()).""" + """Contiguous little-endian bytes of a tensor (torch has no .tobytes()). + + Must not materialize a Python int per byte: a PLE Q8G head is ~3 GiB + of int8, and a list of that many ints will swap a 48 GB machine. + """ + import ctypes t = t.detach().cpu().contiguous() n = t.numel() * t.element_size() - buf = torch.empty(n, dtype=torch.uint8) - buf.view(t.dtype)[:t.numel()] = t.flatten() - return bytes(memoryview(buf.numpy() if False else bytearray(buf.tolist()))) + if n == 0: + return b"" + return bytes((ctypes.c_char * n).from_address(int(t.data_ptr()))) def write_expert_record(f, layer, eid, cb_base, payloads, scales, shapes, packed=False): @@ -897,10 +1087,14 @@ class ShardDebt: """ TRUNK = ("trunk",) + PLE = ("ple",) + SKIP = ("skip",) @staticmethod def name(who): - return f"layer {who[1]}" if who[0] == "layer" else " ".join(map(str, who)) + if who[0] == "layer": + return f"layer {who[1]}" + return " ".join(map(str, who)) @staticmethod def ledger(src): @@ -918,8 +1112,9 @@ def ledger(src): except OSError: return set() - def __init__(self, weight_map, src): + def __init__(self, weight_map, src, skip=None): self.src = src + self.skip = skip self.ledger_path = os.path.join(src, ".reclaimed") self.owed = {} # shard -> consumers that have not finished self.held_by = {} # consumer -> the shards it is holding up @@ -928,24 +1123,38 @@ def __init__(self, weight_map, src): # one can give back, and counting it would report a saving twice. if not os.path.exists(os.path.join(src, shard)): continue - who = self.consumer(name) + who = self.consumer(name, skip) self.owed.setdefault(shard, set()).add(who) self.held_by.setdefault(who, set()).add(shard) self.freed = 0 self.released = [] @staticmethod - def consumer(name): + def consumer(name, skip=None): """The one part of this script that reads `name`. The trunk pass takes everything that is not an expert. mxfp4 stores a tensor as a _packed/_scale pair and ST rejoins them, so both halves belong wherever the tensor they encode does. + + Qwen packs a whole layer's experts into two tensors; those still + belong to that layer. Its PLE n-gram shards are a consumer of their + own, because build_ple runs *after* the trunk pass — a shard that + also holds an expert must not be deleted before the logical heads + have been written. And `skip` names what this conversion never reads + at all (Qwen's vision tower and MTP layer), so those shards are + given back at the start of a reclaim rather than held to the end. + It is passed in rather than assumed: GLM's checkpoint spells its + tower `model.visual.` too, and GLM *does* carry it. """ base = name for suffix in ("_packed", "_scale"): if base.endswith(suffix): base = base[:-len(suffix)] + if skip and skip(base): + return ShardDebt.SKIP + if is_ple_ngram(base) or is_ple_meta(base): + return ShardDebt.PLE if ".experts." not in base: return ShardDebt.TRUNK parts = base.split(".") @@ -1045,6 +1254,70 @@ def bank_codebook_base(path): # ------------------------------------------------------------- worker ---- +def convert_layer_packed(job, st, dev, t0): + """One Qwen layer: two packed tensors, split into WEXP records.""" + (L, src, out, src_pfx, n_exp, stages, entries, index_bits, device, + cb_sample, cb_base, cached_ok) = job + bank = os.path.join(out, f"experts-L{L}.bin") + cbf = os.path.join(out, f"codebooks-L{L}.bin") + gname, dname = qwen_packed_names(src_pfx, L) + gate_up = st.tensor(gname) + down = st.tensor(dname) + e_all = int(gate_up.shape[0]) + if not packed_shapes_ok(gate_up.shape, down.shape, e_all): + raise ValueError( + f"L{L} packed experts {tuple(gate_up.shape)} / {tuple(down.shape)} " + f"are not [E,2I,H] / [E,H,I]") + n_write = min(int(n_exp), e_all) + inter = int(gate_up.shape[1]) // 2 + hid = int(gate_up.shape[2]) + shapes = [(inter, hid), (inter, hid), (hid, inter)] + kinds = (("gate", 0, inter), ("up", inter, 2 * inter), ("down", None, None)) + + def matrix(e, kind, lo, hi): + if kind == "down": + return down[e] + return gate_up[e, lo:hi] + + books, sample_ids = {}, list(range(0, n_write, max(1, n_write // cb_sample)))[:cb_sample] + per = max(1, TRAIN_VECTORS // len(sample_ids)) + with open(cbf + ".tmp", "wb") as cf: + for ki, (kind, lo, hi) in enumerate(kinds): + chunks = [] + for e in sample_ids: + W = matrix(e, kind, lo, hi) + sc = W.abs().amax(-1, keepdim=True).clamp(min=1e-8) + V = (W / sc).reshape(-1, VEC_DIM) + g = torch.Generator().manual_seed(1234 + e) + chunks.append(V[torch.randperm(V.shape[0], generator=g)[:per]]) + del W, V + X = torch.cat(chunks); del chunks + books[kind] = train_codebooks(X, stages, dev, sample=TRAIN_VECTORS, + entries=entries) + del X + for si, C in enumerate(books[kind]): + cid = cb_base + ki * stages + si + cf.write(struct.pack(" 0): + in_merge = (recovered >= 0 and os.path.exists(bank) + and os.path.getsize(bank) > 0 + and recovered + n_cb_per_layer <= old_books) + pending = (old_books <= recovered <= old_books + span + and os.path.exists(part) + and os.path.getsize(part) + == n_cb_per_layer * cb_record_bytes + and os.path.exists(bank) + and os.path.getsize(bank) > 0) + if in_merge or pending: base = recovered cached_ok = True - _t0 = _kinds[0][1] - source_ok = (st.have(ename(L, 0, _t0)) or - st.have(ename(L, 0, _t0) + "_packed")) - if not cached_ok: + if cached_ok: + recovered_for[L] = base + + used = set(recovered_for.values()) + holes = [b for b in range(0, old_books, n_cb_per_layer) if b not in used] + next_base = old_books + jobs = [] + for L in layers: + source_ok = layer_source_ok(L) + if L in recovered_for: + base, cached_ok = recovered_for[L], True + if base + n_cb_per_layer > next_base: + next_base = base + n_cb_per_layer + elif holes: + # A bank rewritten then deleted left a hole in the merge. + # Reuse that base; do not append past the merged file. + base, cached_ok = holes.pop(0), False + else: + cached_ok = False base = next_base if source_ok: next_base += n_cb_per_layer - elif base + n_cb_per_layer > next_base: - next_base = base + n_cb_per_layer jobs.append((L, args.src, args.out, src_pfx, n_exp, args.stages, args.entries, args.index_bits, str(dev), args.cb_sample, base, cached_ok)) tindex, engram = None, {} if debt is not None: + reclaim(debt, args.reclaim, ShardDebt.SKIP, "excluded tensors") # The trunk pass consumes every tensor that is not an expert, so # until it has run almost no shard is fully spent. Run it first. It # writes trunk.bin.tmp either way and the published trunk.bin is # still only replaced together with the manifest, so nothing about # what this run can survive changes — only the order. - tindex = build_trunk(args, sr, st, existing, manifest_path, drop_trunk, - ds41_rename if ds41 else - glm_rename if glm else None, - glm_flatten if glm else None) + tindex = build_trunk(args, sr, st, existing, manifest_path, + drop_trunk, rename, reshape) if tindex is None: return 1 if ds41: @@ -1664,6 +2090,12 @@ def reclaim_layer(L, size): results.append(res) print(f" layer {res[0]}: {res[1]/2**20:.0f} MB [{res[3]}]", flush=True) reclaim_layer(res[0], res[1]) + gc.collect() + if str(dev) == "mps" and hasattr(torch, "mps"): + try: + torch.mps.empty_cache() + except Exception: + pass for L, sz, base, how in sorted(results): if sz: @@ -1681,29 +2113,39 @@ def reclaim_layer(L, size): for res in results] if any(os.path.exists(p) for _, p in parts): with open(merged + ".tmp", "wb") as cb_out: - if compatible: + if old_books > 0: with open(merged, "rb") as old: shutil.copyfileobj(old, cb_out) for base, part in sorted(parts): if os.path.exists(part): if os.path.getsize(part) != n_cb_per_layer * cb_record_bytes: raise RuntimeError(f"malformed codebook part: {part}") - expected = cb_out.tell() // cb_record_bytes - if base < expected: - raise RuntimeError( - f"codebook base {base} overlaps the {expected} " - f"records already written") - if base > expected: + dest = base * cb_record_bytes + end = cb_out.tell() + chunk = n_cb_per_layer * cb_record_bytes + if dest + chunk <= end: + # Hole refill: the merge already holds this slot. + cb_out.seek(dest) + with open(part, "rb") as pf: + shutil.copyfileobj(pf, cb_out) + cb_out.seek(end) + elif dest == end: + with open(part, "rb") as pf: + shutil.copyfileobj(pf, cb_out) + elif dest > end: # A resume can recover a bank whose base sits past # the end of what is being written, because an # earlier run finished a layer this invocation is # not redoing. No record names the ids in between, # so pad them: the bases inside the banks are the # engine's truth and cannot be moved. - cb_out.write(b"\0" * ((base - expected) * - cb_record_bytes)) - with open(part, "rb") as pf: - shutil.copyfileobj(pf, cb_out) + cb_out.write(b"\0" * (dest - end)) + with open(part, "rb") as pf: + shutil.copyfileobj(pf, cb_out) + else: + raise RuntimeError( + f"codebook base {base} overlaps the " + f"{end // cb_record_bytes} records already written") os.remove(part) cb_out.flush() os.fsync(cb_out.fileno()) @@ -1723,25 +2165,30 @@ def reclaim_layer(L, size): copied_tok = True break # A release with no tiktoken rank file but a `tokenizers` tokenizer.json - # — GLM's shape. Re-encoded rather than copied, and refused rather than - # approximated when its pattern or its merge order is not the one - # src/tokenizer.c implements. See tools/hf_tokenizer.py. + # — GLM's shape, and Qwen's. Re-encoded rather than copied, and refused + # rather than approximated when its pattern or its merge order is not + # one src/tokenizer.c implements. See tools/hf_tokenizer.py. if not copied_tok and os.path.exists(os.path.join(args.src, "tokenizer.json")): import hf_tokenizer - text, han_split, tok_specials, tok_pattern = \ + text, han_split, tok_specials, digit_run, tok_pattern = \ hf_tokenizer.convert(args.src) atomic_text(os.path.join(args.out, "tokenizer.model"), text) - # The engine defaults to the Kimi pattern, so only the other cases - # are written — and they are written, not inferred at load: the Han - # difference is one token on "A股" and the pattern difference is - # every prompt, and neither shows up as an error. + # The engine defaults to the Kimi pattern, so only the other case is + # written — and it is written, not inferred at load: the difference + # is one token on "A股" and shows up nowhere as an error. The same + # goes for the digit run: Qwen tokenizes "2026" as four pieces and + # Kimi as one, and neither can be guessed from the vocabulary. if not han_split: cfg["tokenizer_han_split"] = False + if digit_run != 3: + cfg["tokenizer_digit_run"] = digit_run if tok_pattern != hf_tokenizer.TOKPAT_CL100K: cfg["tokenizer_pattern"] = tok_pattern note = ("" if han_split else " (no Han branch in its pattern)") if tok_pattern != hf_tokenizer.TOKPAT_CL100K: note = f" (pattern {tok_pattern}, not cl100k)" + if digit_run != 3: + note += f" ({digit_run} digit per piece)" print(f"tokenizer: re-encoded tokenizer.json{note}") # ---- special tokens -------------------------------------------------- @@ -1875,6 +2322,10 @@ def reclaim_layer(L, size): # does, and K3 normalizes to [-1, 1] with mean = std = 0.5, which is # not what CLIP does. Guess nothing that the release states. vc = cfg.get("_outer", {}).get("vision_config") + if vc and qwen: + # The Qwen tower is not carried (see qwen_drop_trunk), so a + # vision.json would describe weights this container does not hold. + vc = None if vc and ds41: # A third tower, and the engine has to be told which. It shares a # block shape with the other two and nothing else: no learned @@ -1981,14 +2432,17 @@ def reclaim_layer(L, size): # --reclaim has already run this, before the experts, so that the shards # holding non-expert tensors become deletable at all. if tindex is None: - tindex = build_trunk(args, sr, st, existing, manifest_path, drop_trunk, - ds41_rename if ds41 else - glm_rename if glm else None, - glm_flatten if glm else None) + tindex = build_trunk(args, sr, st, existing, manifest_path, + drop_trunk, rename, reshape) if tindex is None: return 1 - if ds41: - engram = build_engram(st, args.out, cfg, args.engram_bits) + if qwen: + if not args.skip_trunk: + tindex = build_ple(args, st, list(sr.names()), tindex) + if debt is not None: + reclaim_ple_if_complete(debt, args.reclaim, tindex) + if ds41 and not engram: + engram = build_engram(st, args.out, cfg, args.engram_bits) trunk_path = os.path.join(args.out, "trunk.bin") trunk_tmp = trunk_path + ".tmp" @@ -2003,14 +2457,22 @@ def reclaim_layer(L, size): arch = ("kimi-k3" if "KimiK3" in _hf else "kimi-linear" if "KimiLinear" in _hf else "glm5-next" if "Glm5Next" in _hf else + QWEN_TEXT_TYPE if qwen else "deepseek-v41" if "DeepseekV41" in _hf else _hf or cfg.get("model_type", "unknown")) + man_cfg = normalise_cfg(cfg) + if qwen: + # The PLE head offsets and vocab sizes are i64 primes near 2e7: + # they do not survive a float round-trip, and the engine indexes + # rows with them. Written as Python ints from the source tables. + man_cfg.update(qwen_ple_config(st, list(sr.names()))) + manifest = { "format_version": 0, "arch": arch, "tensor_prefix": prefix, - "config": normalise_cfg(cfg), + "config": man_cfg, # The record's fmt byte is FMT_VQ3R for every stage count but 2 — the # engine takes the stage and entry counts from here, not from the # byte, and only refuses a fmt that is neither VQ3R nor VQ2R. diff --git a/tools/fetch_weights.sh b/tools/fetch_weights.sh index 85d90bb60..ca6ad4ea8 100755 --- a/tools/fetch_weights.sh +++ b/tools/fetch_weights.sh @@ -27,6 +27,13 @@ # tools/fetch_weights.sh # start or resume # tools/fetch_weights.sh --check # verify what is on disk, no fetch # tools/fetch_weights.sh --repo moonshotai/Kimi-Linear --dest /data/kl +# tools/fetch_weights.sh --repo Qwen/Qwen3.8-Flash-Next \ +# --revision de4b8e4d43b917e7706784d8bb445c9af86a3540 \ +# --dest /Users/admin/mnt/llm/qwen38-flash-next/raw --dry-run +# +# --revision pins every fetch URL (API listing, small files, shards). +# Omit it and the script uses main, which moves. Qwen3.8-Flash-Next +# must pass the SHA recorded in docs/QWEN.md. # # Set HF_TOKEN for a gated repo. Safe to run repeatedly and safe to kill: # the next run picks up where it stopped, mid-shard. @@ -54,6 +61,7 @@ fi export PY REPO="${REPO:-moonshotai/Kimi-K3}" +REVISION="${REVISION:-main}" DEST="${DEST:-/Volumes/WasteDisk/k3}" JOBS="${JOBS:-3}" MAX_RETRY="${MAX_RETRY:-8}" @@ -66,17 +74,18 @@ CHECK_ONLY=0 while [ $# -gt 0 ]; do case "$1" in --repo) REPO="$2"; shift 2 ;; + --revision) REVISION="$2"; shift 2 ;; --dest) DEST="$2"; shift 2 ;; --jobs) JOBS="$2"; shift 2 ;; --dry-run) DRY=1; shift ;; --check) CHECK_ONLY=1; shift ;; - -h|--help) sed -n '4,27p' "$0"; exit 0 ;; + -h|--help) sed -n '4,36p' "$0"; exit 0 ;; *) echo "unknown argument: $1" >&2; exit 2 ;; esac done -API="https://huggingface.co/api/models/${REPO}" -RAW="https://huggingface.co/${REPO}/resolve/main" +API="https://huggingface.co/api/models/${REPO}/revision/${REVISION}" +RAW="https://huggingface.co/${REPO}/resolve/${REVISION}" STATE="$DEST/.download-state" LOG="$DEST/download.log" @@ -109,6 +118,7 @@ hcurl() { } echo "repo: $REPO" +echo "rev: $REVISION" echo "dest: $DEST" code=$(hcurl -s -o /dev/null -w '%{http_code}' --max-time 30 "$API") @@ -293,7 +303,9 @@ PY TOTAL=$(wc -l < "$DEST/.shards" | tr -d ' ') TOTAL_BYTES=$("$PY" - "$DEST/model.safetensors.index.json" <<'PY' | tr -d '\r' import json, sys -print(json.load(open(sys.argv[1])).get("metadata", {}).get("total_size", 0)) +# Qwen stores total_size as a JSON float (359999963128.0). bash $(( )) +# cannot parse the trailing .0, so this must be a bare integer. +print(int(json.load(open(sys.argv[1])).get("metadata", {}).get("total_size", 0) or 0)) PY ) diff --git a/tools/hf_tokenizer.py b/tools/hf_tokenizer.py index d8a6f2304..b0cebc0bf 100644 --- a/tools/hf_tokenizer.py +++ b/tools/hf_tokenizer.py @@ -51,13 +51,33 @@ import os import sys -# The pattern src/tokenizer.c implements, minus the Han branch that only the -# Kimi models carry. Compared literally: a release that reorders one +# The patterns src/tokenizer.c implements, spelled out rather than parsed. +# Three things vary across this family and nothing else does, so the whole +# set is enumerated and compared literally: a release that reorders one # alternative is a release this splits differently, and the difference does # not show up as an error. -PAT_NO_HAN = (r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|" - r"\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+") -PAT_HAN = r"[\p{Han}]+|" + PAT_NO_HAN +# +# han — `[\p{Han}]+` as its own leading branch (both Kimi releases) +# or Han left to the letter branch (GLM, Qwen). +# marks — `[\p{L}\p{M}]` (Qwen) or `\p{L}` (Kimi, GLM). Descriptive +# only: tokenizer.c's letter class is the union either way, so +# the two spellings are the same engine behaviour. +# digits — `\p{N}{1,3}` (Kimi, GLM) or `\p{N}` (Qwen). NOT cosmetic; +# it is carried to the engine as `tokenizer_digit_run`. +def _pattern(han, marks, digit_run): + letter = r"[\p{L}\p{M}]" if marks else r"\p{L}" + other = r"[^\s\p{L}\p{M}\p{N}]" if marks else r"[^\s\p{L}\p{N}]" + return ((r"[\p{Han}]+|" if han else "") + + r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?" + letter + + r"+|\p{N}" + (r"{1,3}" if digit_run == 3 else "") + + r"| ?" + other + r"+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+") + + +# pattern -> (han_split, digit_run) +KNOWN_PATTERNS = {_pattern(h, m, d): (h, d) + for h in (0, 1) for m in (0, 1) for d in (1, 3)} +PAT_NO_HAN = _pattern(0, 0, 3) +PAT_HAN = _pattern(1, 0, 3) # DeepSeek-V4.1 splits with three isolating Splits in sequence rather than # one pattern, and src/tokenizer.c implements the composition as a mode of @@ -115,7 +135,7 @@ def walk(node): def convert(src, quiet=False): - """Returns (rank-file text, han_split, specials list).""" + """Returns (rank text, han split, specials, digit run, pattern mode).""" path = os.path.join(src, "tokenizer.json") with io.open(path, encoding="utf-8") as f: tok = json.load(f) @@ -132,18 +152,17 @@ def convert(src, quiet=False): pats = split_patterns(tok) pat = pats[0] if pats else None if pats == PAT_DS41: - han, pattern = True, TOKPAT_DEEPSEEK # its CJK Split is unconditional - elif len(pats) == 1 and pat == PAT_HAN: - han, pattern = True, TOKPAT_CL100K - elif len(pats) == 1 and pat == PAT_NO_HAN: - han, pattern = False, TOKPAT_CL100K + han, digit_run, pattern = True, 3, TOKPAT_DEEPSEEK + elif len(pats) == 1 and pat in KNOWN_PATTERNS: + han, digit_run = KNOWN_PATTERNS[pat] + han, pattern = bool(han), TOKPAT_CL100K else: raise SystemExit( "this release pre-tokenizes with a pattern src/tokenizer.c does " "not implement, and the difference would be silent:\n" + "".join(f" release: {p}\n" for p in pats or [None]) + - f" engine : {PAT_NO_HAN}\n" - "(optionally preceded by [\\p{Han}]+), or DeepSeek-V4.1's three " + + "".join(f" known : {k}\n" for k in KNOWN_PATTERNS) + + "or DeepSeek-V4.1's three " "Splits. See tools/hf_tokenizer.py.") dec = bytes_to_unicode() @@ -204,8 +223,9 @@ def raw(text): else f"cl100k {'with' if han else 'without'} a Han branch") print(f"tokenizer: {len(lines)} merges, {len(specials)} specials" + (f" ({inline_specials} of them inside the BPE table)" - if inline_specials else "") + f", pattern {which}") - return "\n".join(lines) + "\n", han, specials, pattern + if inline_specials else "") + f", pattern {which}, " + f"up to {digit_run} digit(s) per piece") + return "\n".join(lines) + "\n", han, specials, digit_run, pattern def main(): @@ -220,7 +240,7 @@ def main(): if os.path.exists(dst) and not args.force: print(f"{dst} exists; --force to replace it", file=sys.stderr) return 1 - text, han, specials, pattern = convert(args.src) + text, han, specials, digit_run, pattern = convert(args.src) os.makedirs(args.out, exist_ok=True) with io.open(dst, "w", encoding="utf-8", newline="\n") as f: f.write(text) @@ -236,6 +256,8 @@ def main(): notes.append("tokenizer_han_split must be false") if pattern != TOKPAT_CL100K: notes.append(f"tokenizer_pattern must be {pattern}") + if digit_run != 3: + notes.append(f"tokenizer_digit_run must be {digit_run}") print(f"wrote {dst}" + (f" ({'; '.join(notes)})" if notes else "")) return 0 diff --git a/tools/kimi_ref.py b/tools/kimi_ref.py index 1b46ae6c0..b273281a7 100644 --- a/tools/kimi_ref.py +++ b/tools/kimi_ref.py @@ -30,6 +30,7 @@ import glob import importlib.util import json +import mmap import os import struct import sys @@ -112,46 +113,129 @@ def _load_trunk(self): 2 GB trunk, impossible for K3's 31 GB (it would want ~124 GB). The blob stays as read and each tensor is materialized the first time it is asked for, with a bounded cache of the big ones.""" - self._blob = open(os.path.join(self.path, "trunk.bin"), "rb").read() + trunk_path = os.path.join(self.path, "trunk.bin") + self._trunk_f = open(trunk_path, "rb") + # mmap, not read(): a Qwen trunk is tens of GiB and must not be + # copied into the oracle process just to reach one row at a time. + self._blob = mmap.mmap(self._trunk_f.fileno(), 0, access=mmap.ACCESS_READ) self._meta = {e["name"]: e for e in self.man["trunk"]} self.t = _LazyTrunk(self) + def _deq_row(self, name, row, cols=None): + """One trunk row as f32 — matches waste_deq_row.""" + e = self._meta[name] + shape, off = e["shape"], e["off"] + N = shape[-1] if cols is None else cols + blob = self._blob + if e["fmt"] == 0: + base = off + row * N * 4 + return torch.frombuffer(bytearray(blob[base:base + N * 4]), + dtype=torch.float32).clone() + g = e["group"] + ng = (N + g - 1) // g + fmt, rowbytes = e["fmt"], None + if fmt == 3: + rowbytes = ng * g // 2 + elif fmt == 7: + rowbytes = (ng * g * 3 + 7) // 8 + 1 + else: + rowbytes = ng * g + qoff = off + row * rowbytes + soff = e["scale_off"] + row * ng * 2 + if fmt == 3: + p4 = torch.frombuffer(bytearray(blob[qoff:qoff + rowbytes]), + dtype=torch.uint8) + sc = torch.frombuffer(bytearray(blob[soff:soff + ng * 2]), + dtype=torch.float16).view(ng).float() + idx = torch.arange(N) + k = idx // g + byte = p4[idx // 2] + v = torch.where((idx & 1) == 0, byte & 0x0F, byte >> 4).float() - 8.0 + return (v * sc[k]).float() + elif fmt == 2: + q = torch.frombuffer(bytearray(blob[qoff:qoff + rowbytes]), + dtype=torch.int8).view(ng, g) + sc = torch.frombuffer(bytearray(blob[soff:soff + ng * 2]), + dtype=torch.float16).view(ng).float() + return (q.float() * sc.unsqueeze(-1)).reshape(-1)[:N].float() + else: + raise ValueError(f"unsupported trunk fmt {fmt} for row dequant") + + def matvec(self, name, x, batch=1024): + """y = W @ x for 2-D trunk weight W [rows, cols]. Never materializes W.""" + e = self._meta[name] + rows = 1 + for s in e["shape"][:-1]: + rows *= s + cols = e["shape"][-1] + x = x.detach().float().cpu().reshape(-1) + y = torch.empty(rows, dtype=torch.float32) + for r0 in range(0, rows, batch): + r1 = min(r0 + batch, rows) + w = torch.stack([self._deq_row(name, r, cols) for r in range(r0, r1)]) + y[r0:r1] = w @ x + return y.to(self.dev) + + def matvec_c(self, name, x): + """y = W @ x with C mv_rows / dotf summation order. + + Matches the engine when WASTE_Q8=0 has dequantized the trunk to f32. + Torch batched @ can reorder sums enough to swap near-tie MoE routes.""" + e = self._meta[name] + rows = 1 + for s in e["shape"][:-1]: + rows *= s + cols = e["shape"][-1] + xv = x.detach().float().cpu().reshape(-1).tolist() + y = [0.0] * rows + for r in range(rows): + wr = self._deq_row(name, r, cols).tolist() + acc = 0.0 + for i in range(cols): + acc += wr[i] * xv[i] + y[r] = acc + return torch.tensor(y, dtype=torch.float32).to(self.dev) + + def embed_row(self, token): + name = f"{self.prefix}model.embed_tokens.weight" + return self._deq_row(name, int(token)).to(self.dev) + + def table_row(self, name, row): + """One row from any trunk matrix (embed, PLE ngram tables, …).""" + return self._deq_row(name, int(row)).to(self.dev) + def _materialize(self, name): e = self._meta[name] blob = self._blob - if True: - shape, off = e["shape"], e["off"] - if e["fmt"] == 0: # F32 - n = 1 - for s in shape: - n *= s - x = torch.frombuffer(bytearray(blob[off:off + n * 4]), - dtype=torch.float32).view(*shape) - else: # Q8G / Q4G - rows = 1 - for s in shape[:-1]: - rows *= s - N, g = shape[-1], e["group"] - ng = (N + g - 1) // g - if e["fmt"] == 3: - # Q4G: two signed nibbles per byte, low first, stored as - # v+8 and packed per row. Most of K3's trunk is this, and - # decoding it as int8 silently blows the hidden state up - # by seven orders of magnitude. - b = torch.frombuffer( - bytearray(blob[off:off + rows * ng * g // 2]), - dtype=torch.uint8).view(rows, ng * g // 2).int() - q = torch.stack([b & 0x0F, b >> 4], -1) - q = (q.view(rows, ng, g) - 8).float() - else: - q = torch.frombuffer(bytearray(blob[off:off + rows * ng * g]), - dtype=torch.int8).view(rows, ng, g).float() - sc = torch.frombuffer( - bytearray(blob[e["scale_off"]:e["scale_off"] + rows * ng * 2]), - dtype=torch.float16).view(rows, ng, 1).float() - x = (q * sc).view(rows, ng * g)[:, :N].reshape(*shape) - return x.to(self.dev) - return None + shape, off = e["shape"], e["off"] + rows = 1 + for s in shape[:-1]: + rows *= s + cols = shape[-1] + if e["fmt"] != 0 and rows * cols > 100_000_000: + raise MemoryError( + f"{name} is {rows}x{cols} quantized — use matvec() or embed_row()") + if e["fmt"] == 0: # F32 + n = rows * cols + x = torch.frombuffer(bytearray(blob[off:off + n * 4]), + dtype=torch.float32).view(*shape) + else: # Q8G / Q4G + N, g = cols, e["group"] + ng = (N + g - 1) // g + if e["fmt"] == 3: + b = torch.frombuffer( + bytearray(blob[off:off + rows * ng * g // 2]), + dtype=torch.uint8).view(rows, ng * g // 2).int() + q = torch.stack([b & 0x0F, b >> 4], -1) + q = (q.view(rows, ng, g) - 8).float() + else: + q = torch.frombuffer(bytearray(blob[off:off + rows * ng * g]), + dtype=torch.int8).view(rows, ng, g).float() + sc = torch.frombuffer( + bytearray(blob[e["scale_off"]:e["scale_off"] + rows * ng * 2]), + dtype=torch.float16).view(rows, ng, 1).float() + x = (q * sc).view(rows, ng * g)[:, :N].reshape(*shape) + return x.to(self.dev) def expert(self, L, eid): """Dequantize one expert: exactly one pread of its 4 KiB-aligned record.""" diff --git a/tools/make_test_container.py b/tools/make_test_container.py index 5e99b4859..929ce7125 100644 --- a/tools/make_test_container.py +++ b/tools/make_test_container.py @@ -141,6 +141,95 @@ "index_kpool_always_select_tail": True, } +# --qwen writes a text-only Qwen3.8-Flash-Next fixture at the same scale. +# Shapes match official qwen4_exp: GDN qkv = 2*Hk*Dk+Hv*Dv, QSA o_proj is +# hid x n_heads*head_dim, indexer is (n+kv)*idx_dim, 16 PLE heads. +H_GDN, D_GDN = 4, 8 +QWEN_CFG = { + "model_type": "qwen4_exp_text", + "architectures": ["Qwen4ExpForConditionalGeneration"], + "hidden_size": 32, + "num_hidden_layers": 2, + "moe_intermediate_size": 16, + "shared_expert_intermediate_size": 16, + "num_experts": 4, + "num_experts_per_token": 2, + "num_experts_per_tok": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "rms_norm_eps": 1e-6, + "vocab_size": 256, + "tie_word_embeddings": False, + "bos_token_id": 1, + "eos_token_id": 2, + "hidden_act": "silu", + "layer_types": ["linear_attention", "full_attention"], + "linear_num_key_heads": 4, + "linear_key_head_dim": D_GDN, + "linear_num_value_heads": 12, + "linear_value_head_dim": D_GDN, + "linear_conv_kernel_dim": 4, + "hc_count": 4, + "hc_lowrank": 8, + "ple_embed_dim": 128, + "ple_layer_ids": [2], + "ple_conv_kernel_size": 4, + "ngram_size": 3, + "ngram_vocab_size_base": 64, + "split_ngram_parts": 8, + "heads_per_ngram": 8, + "indexer_n_heads": 2, + "indexer_kv_heads": 1, + "indexer_head_dim": 8, + "indexer_budget": 32, + "indexer_compress_ratio": 4, + "partial_rotary_factor": 0.5, + "rope_parameters": { + "partial_rotary_factor": 0.5, + "rope_theta": 10000.0, + "mrope_section": [2, 2, 2], + "rope_type": "default", + "mrope_interleaved": True, + }, + "ple_head_offsets": [0] * 16, + "ple_head_vocab_sizes": [32] * 16, + "ple_layer_multipliers": [3, 5, 7], +} +PLE_HEADS, PLE_HEAD_WIDTH = 16, 8 + +def _is_prime(n): + if n < 2: + return False + if n % 2 == 0: + return n == 2 + d = 3 + while d * d <= n: + if n % d == 0: + return False + d += 2 + return True + + +def _nth_prime_after(start, count): + p = start + for _ in range(count): + p += 1 + while not _is_prime(p): + p += 1 + return p + + +QWEN_PLE_SIZES = [_nth_prime_after(10, h + 1) for h in range(PLE_HEADS)] +QWEN_PLE_MULT = [3, 5, 7] +QWEN_CFG["ple_head_vocab_sizes"] = QWEN_PLE_SIZES +_off, _offsets = 0, [] +for _s in QWEN_PLE_SIZES: + _offsets.append(_off) + _off += _s +QWEN_CFG["ple_head_offsets"] = _offsets +QWEN_CFG["ple_layer_multipliers"] = QWEN_PLE_MULT + # --ds41 turns it into a DeepSeek-V4.1-Flash at the same scale. Almost # nothing below the MoE is shared with the rest of the family, so this is # the only container that reaches any of it: @@ -457,6 +546,118 @@ def write_tokenizer(outdir): return base + len(SPECIALS) +def write_qwen_container(args, rng): + """A structurally valid Qwen text fixture. Format v0, WEXP unchanged.""" + cfg = dict(QWEN_CFG) + hid = cfg["hidden_size"] + moe = cfg["moe_intermediate_size"] + hc, lr = cfg["hc_count"], cfg["hc_lowrank"] + hc_w = hc * hid + t = Trunk(rng, "") + t.quant("model.embed_tokens.weight", [cfg["vocab_size"], hid]) + t.f32("model.hyper_connection_mixer.hc_norm.weight", [hc_w]) + t.quant("model.hyper_connection_mixer.input_mix_weight_down.weight", [lr, hc_w]) + t.quant("model.hyper_connection_mixer.input_mix_weight_up.weight", [hc_w, lr]) + for L, kind in enumerate(cfg["layer_types"]): + p = f"model.layers.{L}." + for side in ("attn_hyper_connection", "mlp_hyper_connection"): + t.f32(p + side + ".hc_norm.weight", [hc_w]) + t.quant(p + side + ".block_inject_weight.weight", [hc, hc_w]) + t.quant(p + side + ".input_mix_weight_down.weight", [lr, hc_w]) + t.quant(p + side + ".input_mix_weight_up.weight", [hc_w, lr]) + if kind == "linear_attention": + a = p + "linear_attn." + hk, hv = cfg["linear_num_key_heads"], cfg["linear_num_value_heads"] + dk, dv = cfg["linear_key_head_dim"], cfg["linear_value_head_dim"] + qkv = 2 * hk * dk + hv * dv + t.f32(a + "A_log", [hv]) + t.f32(a + "dt_bias", [hv]) + t.f32(a + "conv1d.weight", [qkv, 1, cfg["linear_conv_kernel_dim"]]) + t.quant(a + "in_proj_qkv.weight", [qkv, hid]) + t.quant(a + "in_proj_z.weight", [hv * dv, hid]) + t.quant(a + "in_proj_a.weight", [hv, hid]) + t.quant(a + "in_proj_b.weight", [hv, hid]) + t.f32(a + "norm.weight", [dv]) + t.quant(a + "out_proj.weight", [hid, hv * dv]) + else: + a = p + "self_attn." + qd = cfg["num_attention_heads"] * cfg["head_dim"] + kvd = cfg["num_key_value_heads"] * cfg["head_dim"] + idxd = ((cfg["indexer_n_heads"] + cfg["indexer_kv_heads"]) + * cfg["indexer_head_dim"]) + t.quant(a + "q_proj.weight", [qd * 2, hid]) + t.quant(a + "k_proj.weight", [kvd, hid]) + t.quant(a + "v_proj.weight", [kvd, hid]) + t.quant(a + "o_proj.weight", [hid, qd]) + t.f32(a + "q_norm.weight", [cfg["head_dim"]]) + t.f32(a + "k_norm.weight", [cfg["head_dim"]]) + t.quant(a + "indexer.index_qk_proj.weight", [idxd, hid]) + t.f32(a + "indexer.q_layernorm.weight", [cfg["indexer_head_dim"]]) + t.f32(a + "indexer.k_layernorm.weight", [cfg["indexer_head_dim"]]) + m = p + "mlp." + t.quant(m + "gate.weight", [cfg["num_experts"], hid]) + t.quant(m + "shared_expert.gate_proj.weight", [moe, hid]) + t.quant(m + "shared_expert.up_proj.weight", [moe, hid]) + t.quant(m + "shared_expert.down_proj.weight", [hid, moe]) + t.quant(m + "shared_expert_gate.weight", [1, hid]) + if L == 1: + pe = cfg["ple_embed_dim"] + t.quant(p + "ple.key_proj.weight", [hc_w, pe]) + t.quant(p + "ple.value_proj.weight", [hid, pe]) + t.f32(p + "ple.conv1d.weight", [hc_w, 1, cfg["ple_conv_kernel_size"]]) + t.f32(p + "ple.norm_conv.weight", [hc_w]) + t.f32(p + "ple.norm_key.weight", [hc_w]) + t.f32(p + "ple.norm_query.weight", [hc_w]) + ngram_heads = (cfg["ngram_size"] - 1) * cfg["heads_per_ngram"] + head_w = pe // ngram_heads + for h in range(PLE_HEADS): + t.quant(p + f"ple.ple_embedding.ngram_head.{h}.weight", + [QWEN_PLE_SIZES[h], head_w], bits=8) + t.quant("lm_head.weight", [cfg["vocab_size"], hid]) + with open(os.path.join(args.out, "trunk.bin"), "wb") as f: + f.write(t.buf) + + shapes = [(moe, hid), (moe, hid), (hid, moe)] + layers, cb_base = {}, 0 + with open(os.path.join(args.out, "codebooks.bin"), "wb") as cf: + for L in range(cfg["num_hidden_layers"]): + for ki in range(len(KINDS)): + for si in range(STAGES): + cid = cb_base + ki * STAGES + si + cf.write(struct.pack("