diff --git a/docs/recipes/README.md b/docs/recipes/README.md index f3e82f73b4..fabe4435f6 100644 --- a/docs/recipes/README.md +++ b/docs/recipes/README.md @@ -15,6 +15,7 @@ the native engine, or the Python engine frozen at 0.6.6. | Gemma 4 26B-A4B | [MLX, fused one-row decode](gemma-4.md) | | DeepSeek-V4-Flash | [MLX on a 256 GB Mac, DSpark and MTP drafts](deepseek-v4-flash.md) | | Qwen3.6-35B-A3B | [MLX with MTP drafts, one-GPU CUDA](qwen3.6-moe.md) | +| Ling-3.1-flash | [Not served yet: CUDA weight loader, synthetic checkpoints](ling-3.1.md) | ## Capability floor diff --git a/docs/recipes/ling-3.1.md b/docs/recipes/ling-3.1.md new file mode 100644 index 0000000000..97af269852 --- /dev/null +++ b/docs/recipes/ling-3.1.md @@ -0,0 +1,324 @@ +# Ling-3.1-flash + +Work in progress on the Zig engine's CUDA side. The family is registered with the native CUDA server +(`model_type` `bailing_hybrid`) and serves requests through the lane core at TP=1, but the whole model does not fit +one GB10, so serving is verified on synthetic checkpoints and a thin real model. This page lists what exists, what is +missing, and how to test without the 572 GB checkpoint. Everything below was measured on one GB10. + +Ling-3.1-flash (`BailingMoeV3ForCausalLM`, `bailing_hybrid`) has 72 layers plus one MTP layer: Kimi delta attention +(KDA) on layers with `(i + 1) % 8 != 0`, MLA on every eighth, dense MLPs on layers 0-3 and 512 routed experts (top 8, +8 groups keeping 4, sigmoid scores with a pick bias) plus one shared expert from layer 4 on. The released checkpoint is +block FP8 (e4m3 in 128 x 128 blocks with `weight_scale_inv`). + +## What exists + +| Part | Where | Checked by | +| --- | --- | --- | +| Group-limited top-k routing, per-layer shared SwiGLU limit (GLM's bytes unchanged) | `glm_glue_moe.cuh` (`tf_glm_topk`) | `glm-ref`, `glm-route-old`, `glm-route-ling` | +| MLA at Ling's shape: 64 roped key columns, 128-wide heads, GPT-J rope, head-wise gate | `glm_glue_attn.cuh`, `glm_glue_latent.cuh`, `block_mla.zig` | `ling-mla`, `ling-mla-real`, `glm-mla-same` | +| KDA on Ling's layer 0 (q/k/v convs as one q\|k\|v conv), GLM's row path as is | `glm_kda.cu`, `block_kda.zig` | `glm-kda-ling` | +| Config loader and weight map (all 214,025 names of the real index) | `zig/src/families/ling/config.zig`, `names.zig` | `zig build test-ling` | +| Weight loader at TP=1: bf16, FP8 block, ModelOpt NVFP4 | `zig/src/families/ling/weights.zig`, `weights_proj.zig` | `ling-weights`, `ling-weights-real` | +| Synthetic checkpoints in all three formats | `tf-cuda-test ling-synth` | its own name check | +| The TP=1 layer graph: embedding, KDA/MLA mixers, dense MLP and MoE, final norm, head; windows and commits | `zig/src/families/ling/forward.zig`, `graph.zig`, `state.zig`, `block_*.zig`, `kernels/cuda/ling_glue.cu` | `ling-forward-synth`, `ling-layer-real` | +| The MTP layer (`model.layers.72`) as the drafter: `enorm(embed(next)) \| hnorm(hidden)` through `eh_proj`, MLA + MoE on its own cache, `final_layernorm`, the shared head; chains of up to 3 | `mtp.zig`, `Engine.draft` | `ling-lanes --mtp`, `ling-mtp-real` | +| The engine and lane backend: a sequence per slot, prompt chunks cut anywhere, verify windows of up to 16 rows, shared windows of up to 8 streams (64 rows), kept prompt states, the server hooks | `engine.zig`, `lanes.zig`, `native.zig`; `native/cuda.zig` registers it | `ling-lanes` | +| SwiGLU clamped in Ling's order, `min(silu(g), L) * clip(u, -L, L)`, in the glue and the FP8/NVFP4 grouped experts (every older kernel's SASS unchanged) | `glm_glue_moe.cuh` (`tf_ling_swiglu_*`), `experts.cuh` (epilogue 4), `experts.Layer.after` | `ling-clamp`, `ling-clamp-old`, `glm-route-old` | + +## The layer graph + +`forward.Forward` runs rows of one or more sequences (`Seg`: a `State`, its first row, its row count) through the +embedding, every loaded layer and the head; `stack()` runs the layers alone on `Buffers.x`, `layer()` one of them. +`Buffers` are prompt rows or decode rows; each sequence's rows stay pending until `commit(state, keep)`. A prompt +keeps every row; a decode window (at most `window_rows`, 16) keeps any prefix: the KDA states are replayed from the +window's saves (`replay_layers`) and the conv windows shifted, as #804's GLM commit does. The API follows GLM's +`cuda_forward.zig` (`Rank`/`Seg`, `Buffers`/`State`, `commit`) so lanes hooks can follow. + +Every projection is a `qlinear.Weight` and every routed-expert layer an `experts.Layer` in the format the checkpoint +stored; the graph only calls `qlinear.Linear.decode`/`prompt` and `experts.Grouped.decode`/`prompt`. MLA's absorbed +`kv_b` (bf16 `wk`/`wv` per head) is derived at load by running `kv_b` in its own format on identity rows, so any format +`qlinear` serves works. bf16 routed experts (no grouped bf16 kernel) run every expert on every row and the combine +reads each row's picks: correct, for tests, not for speed. + +Per layer, in the HF model's order: input RMSNorm; KDA (q, k, v, beta into one `[q|k|v|b]` row, f and g rows, GLM's +chain with the safe gate, gated output norm, `o_proj`) or MLA (`q_a`, `q_a_layernorm`, `q_b`; `kv_a_proj_with_mqa`, +`kv_a_layernorm` on the 512 latent, the 64 key columns roped into the cache row; absorbed queries, roped attention +over the 576-wide cache, `expand_v`, sigmoid head gate from `g_proj`, `dense`); residual; post-attention RMSNorm; +dense SwiGLU (layers 0-3) or MoE (bf16 router dot in fp32, sigmoid + `expert_bias` with 8 groups keeping 4, +normalized weights x 2.5, grouped experts with the layer's limit, shared expert with its own limit); residual. The +residual stream, every linear's output and the MoE sums are bf16 at the points where the HF bf16 model rounds them +(`ling_glue.cu`: residual add, head gate, expert sum + shared). Decode windows run KDA's three-kernel chain, as +GLM's forward does by default (bit-equal to the fused chain, checked). + +Measured on synthetic checkpoints (4 layers: KDA dense, MLA MoE, KDA MoE, MLA MoE; 80 prompt rows, 8 decode rows), +logits against `oracle/ling_forward.py` routed on the graph's own picks (with free routing, 15-21 of 264 picks +differ from float64's: group near-ties, the worst group sums 0.0014-0.008 apart): + +| Format | vs float64 | vs the bf16 rounding points | bf16 points vs float64 | argmax vs float64 | +| --- | --- | --- | --- | --- | +| bf16 | 2.17e-2 | 1.31e-2 | 2.18e-2 | 86/88 | +| FP8 block | 2.12e-2 | 1.22e-2 | 2.12e-2 | 85/88 | +| NVFP4 | 2.29e-2 | 1.40e-2 | 2.25e-2 | 85/88 | + +(relative L2 of the logits). The graph is as far from float64 as the HF model's own bf16 rounding is. Bit for bit, in +every format: a prompt cut at 48 equals the uncut prompt; an 8-row verify window equals 8 decode rows; keeping 3 of 8 +and running the next 5 leaves every KDA state, conv window and MLA cache row equal to 8 decode rows; two sequences in +one prompt chunk and in one decode window equal each alone; the fused and three-kernel KDA chains agree. + +Real layers of the FP8 release (`ling-layer-real`, the HF bf16 chain's captured input and output, 512 tokens): + +| Layer | Kind | rel L2 vs HF | rows with a different top-8 | rel L2 without them | decode rows vs HF | device peak | +| --- | --- | --- | --- | --- | --- | --- | +| 0 | KDA, dense | 3.27e-3 | - | 3.27e-3 | 3.12e-3 | 0.9 GB | +| 4 | KDA, MoE | 5.88e-3 | 21 | 3.13e-3 | 4.28e-3 | 8.5 GB | +| 7 | MLA, MoE | 4.89e-3 | 17 | 3.05e-3 | 3.57e-3 | 8.3 GB | +| 65 | KDA, MoE, limits 10 | 4.85e-3 | 14 | 3.35e-3 | 3.33e-3 | 8.5 GB | +| 71 | MLA, MoE, limits 10 | 4.13e-3 | 7 | 3.75e-3 | 3.89e-3 | 8.3 GB | + +exllamav3's port measured 0.4-1.1e-3 without flips (4 flips each on layers 4 and 7) on the same chain: it keeps fp16 +activations, the captured HF chain carries an fp32 residual from layer 3's output on, and this graph rounds the stream +to bf16 as the bf16 model does. On every +real layer a prompt cut at 300 equals the uncut prompt and a 16-row window equals 16 decode rows; for FP8 decode rows +equal the prompt's rows (both run the lane). Loading one MoE layer's 512 experts takes 6.1 s from +local disk and 122 s from the sshfs mount (its read alone is 92 s); see Loading below. + +### Formats + +The format of every tensor comes from the checkpoint, never from the layer graph: + +- `quantization_config.quant_method: "fp8"` with `weight_block_size [128, 128]`: FP8 block. `modules_to_not_convert` + names the bf16 linears (a whole name, or a tail after a dot: `mlp.gate` is the router, not `gate_proj`). +- `quant_method: "modelopt"`, `quant_algo: "NVFP4"`, 4-bit weights in groups of 16, or the same block in + `hf_quant_config.json`: ModelOpt NVFP4. `ignore`, `exclude_modules` and `modules_to_not_convert` take names and + `*` globs. +- No `quantization_config`: bf16. +- `quant_method: "exl3"` is refused with a message naming tracker #832. The extension point is a variant in + `config.Quant`, its side tensors in `names.side()` and a reader in `weights_proj.Reader.linear()`. + +A file tensor whose dtype or shape disagrees with the format the config gives it is refused, naming the module and +both formats. Each projection loads as a `qlinear.Weight` (`bf16`, `fp8g` or `nvfp4`) through main's packers +(`modelopt.Packer`, `fp8.packScales`); routed experts load as `experts.Layer` (`fp8g` or `nvfp4`, packed by +`Fp8Experts`/`Nvfp4Experts`'s layouts on the GPU: `Packer.fp8Experts` and `Packer.fp4Experts`). There is no grouped bf16 expert kernel, so bf16 routed experts load as one lane weight +per projection. The router must be bf16 (GLM's router kernel reads bf16); a quantized router is refused. + +### Shapes + +`config.parse` refuses, with a named error and the reason, every shape the kernels don't take, before anything is +allocated or launched: KDA heads other than 128 wide, short convs other than 4 taps, key/value heads other than the +query heads (`KdaGeometry`); a latent other than 512, nope or value heads other than 128, rope columns other than 64, +NeoX rope pairs, a `q_lora_rank` without a norm instance (`MlaGeometry`); a hidden width other than 512, 1536 or 4096, +MLP and expert widths outside whole 128-column blocks (`Width`); a tied head (`TiedHead`); routing past `glue.select`'s +bounds (`Routing`). + +### Loading + +Routed experts load in batches of `weights.Options.stage` bytes (512 MiB, 32 FP8 experts): eight threads copy the batch's +raw tensors from the mapped file into pinned memory, one upload moves them to a device stage of the same size, and the +packers lay them out in the layer's buffers; the batch's file pages are then dropped. Scales are collected on the host +and uploaded once a layer. `ling-pack DIR LAYER [EXPERTS [STAGE]]` times it and checks every packed byte against the +host packers (`Fp8Experts`/`Nvfp4Experts.packGateUp`, `packOne`). Real layer 4 (512 FP8 experts, 7.5 GiB), from local +disk / the sshfs mount: + +| Loader | cold load | host peak (VmHWM) | device peak | +| --- | --- | --- | --- | +| host packers per expert (before GPU packing) | 14.3 s / ~157 s | 8.0 GiB (the layer's mapped pages) | 7.9 GiB | +| batches packed on the GPU | 6.1 s / 122 s | 1.3 GiB | 8.4 GiB (+ one stage) | + +From local disk the load runs at the read's rate (5-7 s to read the layer once single-threaded); from the mount the +network bounds it (86 MiB/s here: ~1.9 h for the 572 GB release). + +### The SwiGLU clamp + +Layers 64-71 clamp the shared expert and layers 65-71 the routed experts at L = 10 as inclusionAI's vLLM +(`SwigluStepAndMul`) and SGLang (`BailingMoeV3MLP`) do: `min(silu(g), L) * clip(u, -L, L)`, the limit on SiLU's output. +GLM's glue and the grouped expert kernels clamp the gate before SiLU, so Ling launches its own instances: the glue's +`tf_ling_swiglu_512`/`_64` (`glue.swigluWith`) for the shared expert, dense and bf16 routed experts, and epilogue 4 of +the FP8 and NVFP4 grouped kernels (decode and prompt), chosen by `experts.Layer.after`, which the Ling loader sets. +Every kernel that existed before has the same SASS and resource usage as at 6f583baf, the layer graph's commit +(`cuobjdump`, all 40 fatbins); `ling-clamp-old` runs the kept FP8 (12 calls) and NVFP4 (13 calls) instances against +6f583baf's fatbins byte for byte. + +At L = 10 the two orders give the same bf16 bytes: the kernels (and HF's bf16 `act_fn`) round SiLU to bf16 first, and +bf16(silu(g)) is 10 for every bf16 g >= 10. They differ at small limits (L = 2: 16,255 of the 65,536 bf16 gates). +`ling-clamp DIR` checks Ling's order against float64 at L = 10 and L = 2 on gates and ups straddling L: the glue at +widths 1024 and 704 (worst 0.93 of 2^-7 relative; 0 bytes differ from GLM's order at 10, 6,794 at 2) and the clamped +layer's grouped gate-up, decode and prompt rows (FP8 worst 0.090, NVFP4 0.102 of the bound). 6f583baf's binary and +this one print the same digits on the real layers (71: rel L2 4.125e-3, 3.747e-3 without 7 flipped rows, decode rows +3.893e-3; 65: 4.851e-3, 3.347e-3 without 14, 3.325e-3) and on the synthetic forward in all three formats, whose +float64 references do not change either (the synthetic clamped layer never reaches g > 10). HF's clamp itself moves +layer 71's output by 3.7e-3 (rel L2 against the unclamped layer) and layer 65's by 0 on these 512 tokens. + +## Serving + +`engine.Engine` owns the weights (every backbone layer, or `Options.layers` for a thin model), the graph, three +`Buffers` (prompt chunks, one stream's window, shared windows) and one `State` per slot. `Lanes` is the lane backend +(GLM's `cuda_lanes.zig` without the ranks): `prefill`/`prefill_step` (a chunk at a time, other streams' rounds +between), `first`, `verify` (one stream's window, or every stream's in one forward through `verifyShared`), `keep` +(a window's prefix, the KDA states replayed), `release`. Greedy draws take GLM's `tf_glm_top1`, sampled ones +upstream's `sample.cu` through GLM's sampler at TP=1. `native.zig` gives the server `open`, the facts (exact windows of +16 rows, shared rounds unless `TF_LING_SHARED=0`; `TF_LING_PROMPT_STEPS=0` passes prompts whole) and the prompt +cache's `snap_*` hooks: a kept state is the sequence's arena copied whole at a mark, and the lane backend restores it +inside its own prompt pass. Ling's prompt arithmetic does not depend on where a chunk ends, so a mark may fall on any +row (no `snap_grid`): the pass cuts its chunk there. There is no recorder hook (`features`) yet, so `--record` (#848) +is refused for `bailing_hybrid`. + +`ling-lanes DIR OUT [REF] [--layers a,b]` drives the backend under upstream's round loop as the server's lane host +does, prompts in 37-row chunks; every check compares tokens, and the engine-level ones bytes: + +- chunked == whole: chunks of 512, 37 and 16 rows and the whole `prefill` give the same reply; the states after + 37- and 16-row chunks equal the whole pass's byte for byte (KDA states, conv windows, MLA rows). +- drafted == plain: drafts from the plain reply (2 to 15 a round, most rounds one of them wrong, so windows of 3 to + 16 rows are cut anywhere), greedy and sampled; a 9-row window kept to 4 then a 5-row window equal 9 single rows in + tokens, logits and state bytes. +- resumed == fresh: states kept at marks 37, 101, 150, 263, 349 (off the chunk grid, not multiples of 16), each + resumed pass equals the fresh one, and a different tail after a kept 263; resumed states equal the whole pass's. +- concurrent == solo: nine cases of three streams (at once, staggered, reversed, the same prompt twice, sampled, + one discarded, prompts a chunk at a time between the others' rounds, drafted or not) equal their solo replies; + shared windows of 2 streams over 8 rounds (1 to 16 rows each) equal lone twins in logits, tokens and states. +- cancels in the prompt pass, discards between rounds, eos and `max_new` 1-3. +- the reply equals the serial forward's argmax at every row; OUT.ids and OUT.picks feed + `ling_forward.py --ids OUT.ids --picks OUT.picks`, and with REF the served logits are checked against float64: + +| Format | served vs float64 (same picks) | vs the bf16 points | bf16 points vs float64 | argmax vs float64 | +| --- | --- | --- | --- | --- | +| bf16 | 2.24e-2 | 1.27e-2 | 2.25e-2 | 301/309 | +| FP8 block | 2.09e-2 | 1.18e-2 | 2.10e-2 | 302/309 | +| NVFP4 | 2.21e-2 | 1.38e-2 | 2.22e-2 | 294/309 | + +100 checks pass in each format. On the real FP8 release, `--layers 0,1,2,3` (the real embedding, layers 0-3, final +norm and head; KDA and dense MLP only, 6 GB) passes the same exactness checks; it has no float64 reference (an +MLA + MoE layer adds 8 GB). Mutations caught: committing every row on `keep`, a restore that keeps the fresh KDA copy +bits, chunks not cut at marks, shared draws keyed one position early, a lane commit that keeps one row too many. + +## MTP drafting + +The HF model (`BailingMoeV3MTPLayer`) feeds the MTP layer the backbone's rows after `model.norm` and the next +token's embedding: `eh_proj(enorm(embed(t+1)) | hnorm(h_t))`, then input norm, MLA, residual, post-attention norm, +MoE, residual, `final_layernorm`, and `lm_head` reads that. The layer runs unclamped (the SwiGLU limit lists cover +the backbone). `mtp.Head` holds a view of the weights (the MTP layer alone, `final_layernorm` as its `norm`), its +own graph and the concatenated input rows; each sequence gets a second `State` (one MLA cache, its own position). +A prompt pass feeds every row whose next token is known; `Engine.draft` absorbs a verify's kept rows (their +final-norm rows and the tokens after them), then chains up to 3 drafts, each step's hidden row the MTP layer's own +`final_layernorm` output, as GLM's `decode.draft` does; the chain's rows are dropped before the next absorb. Kept +prompt states copy the MTP cache too. The server loads it unless `--no-drafts`; the lane facts then say `mtp`. + +Checked (`ling-lanes ... --mtp`, 109 checks a format; acceptance on random weights says nothing about quality and is +not reported as such): MTP-drafted replies, greedy and sampled, equal the plain ones; so do MTP-drafted passes +resumed at 53 and 211 and every concurrent case; a prompt pass's MTP cache in 37-row chunks equals the serial MTP +rows byte for byte; twins drafting chains of 3 and of 1 agree on tokens and first drafts over 6 windows, each +chain's second draft equals one computed by hand, and after one more window both states are equal byte for byte. +The MTP logits at every prompt row against `ling_forward.py --mtp` (float64 from the file, picks free): + +| Format | vs float64, rows routed alike | rows whose MTP picks differ | the bf16 points vs float64 on those rows | argmax | +| --- | --- | --- | --- | --- | +| bf16 | 1.72e-2 | 14 of 308 | 2.47e-2 | 293/308 | +| FP8 block | 1.65e-2 | 5 | 2.38e-2 | 300/308 | +| NVFP4 | 3.20e-2 | 11 | 7.81e-2 | 283/308 | + +The real layer (`ling-mtp-real`, fixture from `oracle/ling_mtp.py`: HF's `BailingMoeV3MTPLayer` in bf16 on 511 real +tokens' embeddings and `model.norm` of layer 0's output): `final_layernorm` rows at rel L2 7.28e-3 of HF on the 488 +rows whose top-8 agrees (4.56e-2 with the 23 others), worst row 1.61e-2, decode rows 6.74e-3; cut at 300 == uncut and +a 16-row window == 16 decode rows byte for byte; 8.4 GB peak, 153 s to load from the mount. Mutations caught: no +rollback of the chain's rows, `hnorm` with `enorm`'s weight, prompt rows absorbing their own token, kept states +without the MTP cache, the chain's hidden row taken from row 0. + +## What is missing + +- A full-checkpoint run: 574 GB in FP8 needs five Sparks and tensor parallelism. +- Speed: nothing is tuned. Projections land in their columns through a byte copy, bf16 routed experts run densely, + and the RMS norms take hidden and `q_lora_rank` widths 512, 1536 and 4096 only (`glue.rmsnorm`). +- EXL3, through ashhart's EXL3 module (tracker #832). +- Tensor parallelism: the loader is TP=1; shares and transport would come from #804's `cuda_shard.zig` and #855. + +## Testing without the full model + +```bash +zig build test-ling # config, name map, tiny configs (host only) +TF_LING_INDEX=/path/to/model.safetensors.index.json zig build test-ling # adds the real index +zig build -Doptimize=ReleaseFast -Dnvcc=/usr/local/cuda/bin/nvcc +zig-out/bin/tf-cuda-test ling-synth /tmp/ling-fp8 --format fp8 # bf16 | fp8 | nvfp4, --layers N (default 4) +zig-out/bin/tf-cuda-test ling-weights /tmp/ling-fp8 +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/ref # float64 reference, free routing +zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/ref /tmp/picks # exactness; writes the picks +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/fref --picks /tmp/picks +zig-out/bin/tf-cuda-test ling-forward-synth /tmp/ling-fp8 /tmp/fref # same picks: logits bounded +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/served # serving checks; writes .ids/.picks +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/sref --ids /tmp/served.ids --picks /tmp/served.picks +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/served /tmp/sref # and the served logits vs float64 +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/m --mtp # the same with MTP drafts +python zig/tests/cuda/oracle/ling_forward.py /tmp/ling-fp8 /tmp/mref --ids /tmp/m.ids --picks /tmp/m.picks --mtp /tmp/mref.mtp +zig-out/bin/tf-cuda-test ling-lanes /tmp/ling-fp8 /tmp/m /tmp/mref --mtp # and the MTP logits vs float64 +zig-out/bin/tf-cuda-test ling-clamp /tmp/ling-fp8 # Ling's SwiGLU order vs float64, L 10 and 2 +zig-out/bin/tf-cuda-test ling-clamp-old /tmp/ling-fp8 OLD/fatbin # kept expert instances vs an older build's bytes +``` + +`ling-synth` writes `config.json` (the real file's schema) and `model.safetensors` with the real key names: hidden 512, +4 heads, 16 experts in 4 groups keeping 2, `q_lora_rank` 512 (a width the norm kernel takes), vocabulary 1,024, layer +groups of 2 (KDA, MLA, KDA, MLA), layer 0 dense, the last layer clamped, plus the MTP layer. KDA's head dim 128 and +MLA's 512 latent, 64 rope and 128-wide heads stay as the kernels take them. The FP8 file carries the real exception +list's pattern; the NVFP4 file uses globs and leaves layer 0's dense MLP in bf16, so one checkpoint mixes formats. It +checks its names against the weight map before it passes. FP8 codes stay within 2^-3 to 3.75 so a forward stays in +range. Sizes: bf16 282 tensors, 45 MiB; FP8 517, 24 MiB; NVFP4 978, 16 MiB. + +`ling-weights` loads every layer and checks every projection (decode and prompt rows, 3 rows) against float64 from +the file, the grouped experts' gate-up (SwiGLU) and down, and each small tensor (norms, convs as fp32 q|k|v, `A_log`, +`dt_bias`, router, bias) byte for byte. Measured: bf16 85 projections, FP8 and NVFP4 61 each, worst 0.433 of the bound +(2^-7 of |y| plus 2^-12 of the sum of |x w|); 37 small tensors equal. + +`ling-init-faults DIR [--mtp]` opens the backend (engine, sequence 0, lane backend) on `fake_driver.zig`, a host +stand-in for the driver, once cleanly and then with each of its driver calls, and each of its host allocations, +failing in turn: every failed open must leave nothing allocated, held or freed twice (no GPU work; FP8 with `--mtp` +is 869 calls and 394 allocations). `ling-slots` checks the lane backend's slot choice in the release build. + +With the real checkpoint: + +```bash +python zig/tests/cuda/oracle/ling_weights.py /path/to/Ling-3.1-flash /tmp/ling-ref # needs torch, safetensors +zig-out/bin/tf-cuda-test ling-weights-real /path/to/Ling-3.1-flash /tmp/ling-ref +``` + +`ling-ckpt-check DIR [INDEX]` reads every shard's header and nothing else: each of the release's 214,025 tensors is +one the weight map names, at the dtype and shape the loader reads (`names.dims`, the loader's own table, and each +format's side tensors), none missing, and their bytes add up to the index's `total_size` (573,753,281,152). Over a +network mount it takes seconds. + +`ling-mtp-real DIR FIXTURE` runs the MTP layer (fixture from `oracle/ling_mtp.py SRC MODELING_PKG ROWS.pt OUT`). +`ling-layer-real DIR FIXTURE` runs one whole layer (fixture from `oracle/ling_layer.py CHAIN L OUT MODELING`, the captured +HF chain's `h.pt`, `h.pt`, `r.pt`). The HF reference is pinned: `oracle/ling_hf.py` names +`inclusionAI/Ling-3.1-flash` at `a6943f4bf04f8041ad983d2b58dd43b5ba86ae4a` and the sha256 of its +`modeling_bailing_moe_v3.py`, `configuration_bailing_moe_v3.py`, `config.json` and index; `ling_mtp.py`, +`ling_layer.py` and `ling_weights.py` stop on any other copy, and `fixtures/config.json` is held to it in `test-ling`. `ling-weights-real` loads layers 0, 4 and 7 with experts 0-3 (769 MiB on the device) and compares 35 products with torch's float64 +on the same 8 rows: worst 0.414 of the bound, relative L2 at most 1.7e-3 for bf16 outputs and 6.9e-8 for the +experts' fp32 down rows. + +## Memory + +Computed from `config.json` (572.1 B parameters; routed experts are 97.1 %). Quantized formats keep the FP8 release's +bf16 modules (embedding, head, routers, KDA `b_proj`, MLA `g_proj`, MTP `eh_proj`). The FP8 figure equals the +checkpoint's shard total. EXL3 assumes every other linear at the given rate. + +| Format | Weights | Sparks, 128 GB each | Sparks, 100 GB of weights each | +| --- | --- | --- | --- | +| bf16 | 1,144 GB | 9 | 12 | +| FP8 block (the release) | 574 GB | 5 | 6 | +| ModelOpt NVFP4 | 324 GB | 3 | 4 | +| EXL3 2.5 bpw | 181 GB | 2 | 2 | +| EXL3 1.25 bpw | 92 GB | 1 | 1 | + +The first count is weights alone; the second leaves room for KV cache, activations and the system. Loaded, the FP8 +lane weights add 0.94 GB of fp32 scale tiles (one per 64 inputs) on the non-expert linears; NVFP4 loads at its stored +size. + +## How to help + +Each item stands alone. + +1. **Recorder.** Give `Lanes` a `features` hook (the final-norm rows `Engine.verify` already keeps) for #848's recorder. +2. **MTP acceptance.** Measure the MTP layer's acceptance on real text once a whole model runs (TP). +3. **Fused projections.** Same-format q|k|v (and gate|up) rows as one weight per format, without the column copies. +4. **Bf16 routed experts.** Add a grouped bf16 expert kernel, or quantize at load, and drop the per-expert lanes. +5. **TP.** Cut `weights.load` by heads and widths as `glm/cuda_shard.zig` does; FP8 scale blocks span 128 rows and + 128 inputs, so FP8 cuts fall on multiples of 128. +6. **EXL3.** Wire #832's reader behind the `config.Quant` extension point and lift the refusal. +7. **An NVFP4 checkpoint.** Quantize the release with ModelOpt, keep the routers in bf16, and run `ling-weights` on + it; the loader expects `weight`, `weight_scale`, `weight_scale_2` and `input_scale` per linear. diff --git a/zig/build/cuda.zig b/zig/build/cuda.zig index 684d6c3cc9..b4a0cf9c4f 100644 --- a/zig/build/cuda.zig +++ b/zig/build/cuda.zig @@ -52,6 +52,7 @@ const kernels = [_]Kernel{ .{ .name = "glm_glue", .flags = &.{"-O3"} }, // ours: the Python family's Triton glue by hand, its PTX op for op .{ .name = "glm_ops", .flags = &.{"-O3"} }, // ours: GLM's integer work the Python family did in torch (the DSA token lists) .{ .name = "glm_quant", .flags = torch_ops }, // ours: GLM's draft-only weight quantizers, torch at load in the Python family + .{ .name = "ling_glue", .flags = glue }, // ours: Ling-3.1's bf16 glue at the HF model's rounding points .{ .name = "torch_argmax", .src = "torch_ops/argmax", .flags = torch_ops }, .{ .name = "torch_topk", .src = "torch_ops/topk", .flags = torch_ops }, .{ .name = "torch_pointwise", .src = "torch_ops/pointwise", .flags = torch_ops }, @@ -154,6 +155,7 @@ pub fn targets(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu for (exlr8_build.denseModules(b, target, optimize, exlr8, exlr8_tests[0])) |m| runner.addImport(m.name, m.module); const glm_cuda = glmCuda(b, target, optimize, cuda, mods.core, mods.lanes, mods.heat); runner.addImport("glm_cuda", glm_cuda); + runner.addImport("ling_cuda", lingCuda(b, target, optimize, cuda, mods.core, mods.lanes, glm_cuda)); const runner_exe = b.addExecutable(.{ .name = "tf-cuda-test", .root_module = runner }); b.installArtifact(runner_exe); const controls = if (nvcc) |tool| fatbin(b, tool, version.?, .{ .name = "contract_controls", .flags = &.{"-O3"} }, sms) else null; @@ -183,9 +185,17 @@ fn radioModule(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.bu return b.createModule(.{ .root_source_file = b.path("zig/src/families/radio/radio.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{ .{ .name = "cuda", .module = cuda }, .{ .name = "core", .module = core }, .{ .name = "lanes", .module = lanes } } }); } +/// The Ling-3.1-flash CUDA family over the runtime, the core, the lane core and the GLM kernels it reuses. +fn lingCuda(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, glm: *std.Build.Module) *std.Build.Module { + const m = b.createModule(.{ .root_source_file = b.path("zig/src/families/ling/cuda.zig"), .target = target, .optimize = optimize, .link_libc = true }); + for ([_]struct { []const u8, *std.Build.Module }{ .{ "glm_cuda", glm }, .{ "cuda", cuda }, .{ "core", core }, .{ "lanes", lanes } }) |x| m.addImport(x[0], x[1]); + return m; +} + /// The CUDA engines a native server opens (native/cuda.zig), over the given runtime and families. fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builtin.OptimizeMode, cuda: *std.Build.Module, core: *std.Build.Module, lanes: *std.Build.Module, nemotron: *std.Build.Module, heat: *std.Build.Module) struct { api: *std.Build.Module, engines: *std.Build.Module } { const api = b.createModule(.{ .root_source_file = b.path("zig/src/core/engine_api.zig"), .target = target, .optimize = optimize, .link_libc = true, .imports = &.{.{ .name = "lanes", .module = lanes }} }); + const glm = glmCuda(b, target, optimize, cuda, core, lanes, heat); const mod = b.createModule(.{ .root_source_file = b.path("zig/src/native/cuda.zig"), .target = target, @@ -198,7 +208,8 @@ fn engines(b: *std.Build, target: std.Build.ResolvedTarget, optimize: std.builti .{ .name = "nemotron", .module = nemotron }, .{ .name = "lfm2", .module = lfm2Module(b, target, optimize, cuda, core) }, .{ .name = "core", .module = core }, - .{ .name = "glm_cuda", .module = glmCuda(b, target, optimize, cuda, core, lanes, heat) }, + .{ .name = "glm_cuda", .module = glm }, + .{ .name = "ling_cuda", .module = lingCuda(b, target, optimize, cuda, core, lanes, glm) }, }, }); return .{ .api = api, .engines = mod }; @@ -272,6 +283,11 @@ pub fn hostTests(b: *std.Build, draft_ids: *std.Build.Module, build_options: *st // the GLM-5.3-Flash CUDA family's host side const glm = glmCuda(b, host, .debug, cuda, mods.core, mods.lanes, mods.heat); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = glm })).step); + // Ling-3.1-flash's config loader and weight map, and its routing and limits as the GLM kernels take them + const ling = lingCuda(b, host, .debug, cuda, mods.core, mods.lanes, glm); + const ling_run = b.addRunArtifact(b.addTest(.{ .root_module = ling })); + step.dependOn(&ling_run.step); + b.step("test-ling", "Ling-3.1-flash's config loader and weight map (no GPU work)").dependOn(&ling_run.step); // the shares a GLM-5.3-Flash rank takes (cluster/canon.zig against the Python family's) const shares = b.createModule(.{ .root_source_file = b.path("zig/src/cluster/canon_glm_test.zig"), .target = host, .optimize = .debug }); step.dependOn(&b.addRunArtifact(b.addTest(.{ .root_module = shares })).step); diff --git a/zig/kernels/cuda/experts.cuh b/zig/kernels/cuda/experts.cuh index 3396e42a40..cb5dfdca37 100644 --- a/zig/kernels/cuda/experts.cuh +++ b/zig/kernels/cuda/experts.cuh @@ -62,6 +62,17 @@ __device__ __forceinline__ float swiglu(float g, float u, float limit) { return bf(gv / (1.f + expf(-gv))) * uv; } +// Ling's order (EPI 4): the limit on silu's output, min(silu(g), L) * clip(u, -L, L). +__device__ __forceinline__ float swiglu_after(float g, float u, float limit) { + const float gv = bf(g); + float uv = bf(u), s = bf(gv / (1.f + expf(-gv))); + if (limit > 0.f) { + s = fminf(s, limit); + uv = fminf(fmaxf(uv, -limit), limit); + } + return s * uv; +} + __device__ __forceinline__ float relu2(float a) { const float u = fmaxf(bf(a), 0.f); return u * u; @@ -88,6 +99,9 @@ __device__ __forceinline__ void epilogue(const float (&acc)[M][RT][NTW][4], int } else if constexpr (EPI == 2) { o0 = swiglu(a0, acc[M - 1][r][j][2 * h], limit); o1 = swiglu(a1, acc[M - 1][r][j][2 * h + 1], limit); + } else if constexpr (EPI == 4) { + o0 = swiglu_after(a0, acc[M - 1][r][j][2 * h], limit); + o1 = swiglu_after(a1, acc[M - 1][r][j][2 * h + 1], limit); } else { o0 = a0; o1 = a1; diff --git a/zig/kernels/cuda/fp8_experts.cu b/zig/kernels/cuda/fp8_experts.cu index e90d4e7180..1e4302ab4e 100644 --- a/zig/kernels/cuda/fp8_experts.cu +++ b/zig/kernels/cuda/fp8_experts.cu @@ -267,11 +267,13 @@ __global__ void __launch_bounds__(PWARPS * 32, PMIN) } } -// The instances fp8_experts_cuda and fp8_experts_prompt_cuda launch: SwiGLU gate-up, fp32 and bf16 down. +// The instances: SwiGLU gate-up (EPI 4 in Ling's clamp order), fp32 and bf16 down; decode and prompt forms. template __global__ void fp8_expert_kernel<2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void fp8_expert_kernel<2, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_kernel<1, 0, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_kernel<1, 3, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_prompt_kernel<2, 2, 2, 2>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void fp8_expert_prompt_kernel<2, 4, 2, 2>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_prompt_kernel<1, 0, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void fp8_expert_prompt_kernel<1, 3, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); } diff --git a/zig/kernels/cuda/glm_glue_attn.cuh b/zig/kernels/cuda/glm_glue_attn.cuh index db481c0d9e..2c2bcaed86 100644 --- a/zig/kernels/cuda/glm_glue_attn.cuh +++ b/zig/kernels/cuda/glm_glue_attn.cuh @@ -39,26 +39,28 @@ __device__ __forceinline__ void ldm4(uint32_t* r, const void* p) { : "r"(a)); } -// A 512-wide bf16 row's 16-byte unit u in shared memory, swizzled by the row so ldmatrix's eight rows hit eight banks. -__device__ __forceinline__ int wide(int row, int u) { return row * LW + ((u ^ (row & 7)) << 3); } +// A KW-wide bf16 row's 16-byte unit u in shared memory, swizzled by the row so ldmatrix's eight rows hit eight banks. +template +__device__ __forceinline__ int wide(int row, int u) { return row * KW + ((u ^ (row & 7)) << 3); } // p's [32][32] rows (64 bytes): unit u of row r. __device__ __forceinline__ int narrow(int row, int u) { return row * KT + ((u ^ ((row >> 1) & 3)) << 3); } // latent._dense_chunks/_sparse_chunks (sparse 0/1), block (row, HB heads, chunk), 8 warps: online softmax by 32 keys. -template +template constexpr int wide_smem() { - return (HB * LW + 2 * KT * LW + HB * KT) * 2 + (HB * 4 + 2 * HB) * 4; + return (HB * (LW + RW) + 2 * KT * (LW + RW) + HB * KT) * 2 + (HB * 4 + 2 * HB) * 4; } -template +// RW rope columns (Ling's 64) follow the latent in q and keys: scores over LW + RW, values from the latent alone. +template __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, const int* cnt, float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, float scale, int sparse) { - constexpr int MT = HB / 16; + constexpr int MT = HB / 16, KW = LW + RW; extern __shared__ __align__(16) unsigned char smem[]; __nv_bfloat16* qs = reinterpret_cast<__nv_bfloat16*>(smem); - __nv_bfloat16* kvs = qs + HB * LW; - __nv_bfloat16* ps = kvs + 2 * KT * LW; + __nv_bfloat16* kvs = qs + HB * KW; + __nv_bfloat16* ps = kvs + 2 * KT * KW; float* red = reinterpret_cast(ps + HB * KT); float* stat = red + HB * 4; const int t = threadIdx.x, lane = t & 31, warp = t >> 5, g = lane >> 2, q = lane & 3; @@ -67,14 +69,14 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co const int64_t n = sparse ? cnt[r] : 0, limit = sparse ? 0 : (int64_t)pos[0] + r, start = c * CH; const bool live = sparse ? start < n : start <= limit; auto load_kv = [&](int tile) { - __nv_bfloat16* dst = kvs + (tile & 1) * KT * LW; + __nv_bfloat16* dst = kvs + (tile & 1) * KT * KW; const int64_t first = start + tile * KT; - for (int idx = t; idx < KT * (LW / 8); idx += blockDim.x) { - const int key = idx / (LW / 8), u = idx % (LW / 8); + for (int idx = t; idx < KT * (KW / 8); idx += blockDim.x) { + const int key = idx / (KW / 8), u = idx % (KW / 8); const int64_t ki = first + key; const bool ok = sparse ? ki < n : ki <= limit; const int64_t src = ok ? (sparse ? (int64_t)tok[r * w + ki] : ki) : 0; - cp16(dst + wide(key, u), lc + src * LW + u * 8, ok); + cp16(dst + wide(key, u), lc + src * KW + u * 8, ok); } }; float m[MT][2], l[MT][2], o[MT][8][4]; @@ -85,10 +87,10 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co for (int e = 0; e < 4; ++e) o[i][j][e] = 0.0f; constexpr int tiles = CH / KT; if (live) { - for (int idx = t; idx < HB * (LW / 8); idx += blockDim.x) { - const int hd = idx / (LW / 8), u = idx % (LW / 8); + for (int idx = t; idx < HB * (KW / 8); idx += blockDim.x) { + const int hd = idx / (KW / 8), u = idx % (KW / 8); const bool ok = h0 + hd < heads; - cp16(qs + wide(hd, u), qa + (r * heads + (ok ? h0 + hd : 0)) * LW + u * 8, ok); + cp16(qs + wide(hd, u), qa + (r * heads + (ok ? h0 + hd : 0)) * KW + u * 8, ok); } load_kv(0); commit(); @@ -101,17 +103,17 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co else wait<0>(); __syncthreads(); - const __nv_bfloat16* kv = kvs + (tile & 1) * KT * LW; + const __nv_bfloat16* kv = kvs + (tile & 1) * KT * KW; const int64_t first = start + tile * KT; float sc[4]; if (hh < MT) { float acc[4] = {0.0f, 0.0f, 0.0f, 0.0f}; #pragma unroll 4 - for (int kk = 0; kk < LW; kk += 32) { + for (int kk = 0; kk < KW; kk += 32) { uint32_t a0[4], a1[4], b[4]; - ldm4(a0, qs + wide(hh * 16 + (lane & 15), (kk >> 3) + (lane >> 4))); - ldm4(a1, qs + wide(hh * 16 + (lane & 15), ((kk + 16) >> 3) + (lane >> 4))); - ldm4(b, kv + wide(kg * 8 + (lane & 7), (kk >> 3) + (lane >> 3))); + ldm4(a0, qs + wide(hh * 16 + (lane & 15), (kk >> 3) + (lane >> 4))); + ldm4(a1, qs + wide(hh * 16 + (lane & 15), ((kk + 16) >> 3) + (lane >> 4))); + ldm4(b, kv + wide(kg * 8 + (lane & 7), (kk >> 3) + (lane >> 3))); glm_router::mma16816(acc, a0, b); glm_router::mma16816(acc, a1, b + 2); } @@ -172,7 +174,7 @@ __device__ void chunks_wide(const __nv_bfloat16* qa, const __nv_bfloat16* lc, co #pragma unroll for (int j = 0; j < 8; j += 2) { uint32_t b[4]; - ldm4(b, kv + wide(k16 + (lane & 15), ((warp * 64 + j * 8) >> 3) + (lane >> 4))); + ldm4(b, kv + wide(k16 + (lane & 15), ((warp * 64 + j * 8) >> 3) + (lane >> 4))); for (int i = 0; i < MT; ++i) { glm_router::mma16816(o[i][j], a[i], b); glm_router::mma16816(o[i][j + 1], a[i], b + 2); @@ -213,6 +215,13 @@ extern "C" __global__ void __launch_bounds__(256) glm_attn::chunks_wide<32>(qa, lc, pos, tok, cnt, po, pm, pl, rows, heads, w, scale, sparse); } +// Ling's 16-head form: 576-wide q and keys (latent, then 64 roped columns), in wide_smem<16, 64>() bytes. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_chunks_r16(const __nv_bfloat16* qa, const __nv_bfloat16* lc, const int* pos, const int* tok, const int* cnt, + float* po, float* pm, float* pl, int64_t rows, int64_t heads, int64_t w, float scale, int sparse) { + glm_attn::chunks_wide<16, 64>(qa, lc, pos, tok, cnt, po, pm, pl, rows, heads, w, scale, sparse); +} + // latent._merge, block (row, head), 128 threads: chunks merged in order, o / l by div.full; skips sparse count-0 rows. extern "C" __global__ void __launch_bounds__(128) tf_glm_merge(const float* po, const float* pm, const float* pl, __nv_bfloat16* out, const int* cnt, int64_t rows, diff --git a/zig/kernels/cuda/glm_glue_latent.cuh b/zig/kernels/cuda/glm_glue_latent.cuh index 12afc2daa3..ff29d6a74b 100644 --- a/zig/kernels/cuda/glm_glue_latent.cuh +++ b/zig/kernels/cuda/glm_glue_latent.cuh @@ -1,4 +1,4 @@ -// latent.py's BF16 absorb, expand and cache write by hand (the MLA layers' 512-wide latent, 256-wide heads). +// latent.py's BF16 absorb, expand and cache write by hand (512-wide latent; GLM's 256-wide heads, Ling's 128). #pragma once #include "glm_triton.cuh" @@ -6,23 +6,23 @@ using namespace glm_triton; // latent._absorb_q(_rows), block (head, 32 latent columns, row block) of 128 threads; rows per block change no bits. -extern "C" __global__ void __launch_bounds__(128) - tf_glm_absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, int64_t heads, - int64_t rb) { - constexpr int D = 256, LW = 512; +template // q_nope D wide (GLM 256, Ling 128), q rows qs apart per head, qa rows qaw apart +__device__ __forceinline__ void absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, + int64_t rows, int64_t heads, int64_t rb, int64_t qs, int64_t qaw) { + constexpr int LW = 512, R = D / 32; __shared__ float smem[128]; const int t = threadIdx.x, k0 = t >> 2, nq = t & 3, warp = t >> 5; const int64_t h = blockIdx.x, n0 = (int64_t)blockIdx.y * 32 + nq * 8, r0 = blockIdx.z * rb; - float w[8][8]; - for (int rep = 0; rep < 8; ++rep) + float w[R][8]; + for (int rep = 0; rep < R; ++rep) for (int j = 0; j < 8; ++j) w[rep][j] = bf(wk[(h * D + k0 + 32 * rep) * LW + n0 + j]); for (int64_t r = r0; r < r0 + rb && r < rows; ++r) { - float qv[8]; - for (int rep = 0; rep < 8; ++rep) qv[rep] = bf(q[(r * heads + h) * D + k0 + 32 * rep]); + float qv[R]; + for (int rep = 0; rep < R; ++rep) qv[rep] = bf(q[(r * heads + h) * qs + k0 + 32 * rep]); for (int j = 0; j < 8; ++j) { float a = fmul(w[1][j], qv[1]); a = __fmaf_rn(w[0][j], qv[0], a); - for (int rep = 2; rep < 8; ++rep) a = __fmaf_rn(w[rep][j], qv[rep], a); + for (int rep = 2; rep < R; ++rep) a = __fmaf_rn(w[rep][j], qv[rep], a); for (int o = 16; o >= 4; o >>= 1) a = a + __shfl_xor_sync(0xffffffffu, a, o); if ((t & 28) == 0) smem[(nq * 8 + j) * 4 + warp] = a; } @@ -30,16 +30,29 @@ extern "C" __global__ void __launch_bounds__(128) float v = smem[t]; v = v + __shfl_xor_sync(0xffffffffu, v, 2); v = v + __shfl_xor_sync(0xffffffffu, v, 1); - if ((t & 3) == 0) qa[(r * heads + h) * LW + (int64_t)blockIdx.y * 32 + t / 4] = __float2bfloat16_rn(v); + if ((t & 3) == 0) qa[(r * heads + h) * qaw + (int64_t)blockIdx.y * 32 + t / 4] = __float2bfloat16_rn(v); __syncthreads(); } } -// latent._expand_v(_rows), block (head, 16 outputs, row block) of 128 threads: o . wv by fma, lanes xor 16 .. 1. extern "C" __global__ void __launch_bounds__(128) - tf_glm_expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, - int64_t heads, int64_t rb) { - constexpr int DV = 256, LW = 512; + tf_glm_absorb_q(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, int64_t heads, + int64_t rb) { + absorb_q<256>(q, wk, qa, rows, heads, rb, 256, 512); +} + +// Ling: q_nope the first 128 of each head's qs columns, qa rows qaw (576: the latent, then the roped q_rope). +extern "C" __global__ void __launch_bounds__(128) + tf_glm_absorb_q_128(const __nv_bfloat16* q, const __nv_bfloat16* wk, __nv_bfloat16* qa, int64_t rows, + int64_t heads, int64_t rb, int64_t qs, int64_t qaw) { + absorb_q<128>(q, wk, qa, rows, heads, rb, qs, qaw); +} + +// latent._expand_v(_rows), block (head, 16 outputs, row block) of 128 threads: o . wv by fma, lanes xor 16 .. 1. +template +__device__ __forceinline__ void expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, + int64_t rows, int64_t heads, int64_t rb) { + constexpr int LW = 512; __shared__ float smem[32]; const int t = threadIdx.x, lane = t & 31, warp = t >> 5, kb = (warp & 1) * 256 + lane * 8, nb = warp >> 1; const int64_t h = blockIdx.x, n0 = (int64_t)blockIdx.y * 16, r0 = blockIdx.z * rb; @@ -66,9 +79,46 @@ extern "C" __global__ void __launch_bounds__(128) } } +extern "C" __global__ void __launch_bounds__(128) + tf_glm_expand_v(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, + int64_t heads, int64_t rb) { + expand_v<256>(ol, wv, out, rows, heads, rb); +} + +// Ling's 128-wide v heads. +extern "C" __global__ void __launch_bounds__(128) + tf_glm_expand_v_128(const __nv_bfloat16* ol, const __nv_bfloat16* wv, __nv_bfloat16* out, int64_t rows, + int64_t heads, int64_t rb) { + expand_v<128>(ol, wv, out, rows, heads, rb); +} + // latent._lat_write (BF16 cache): row r of `lat` into cache slot pos + r. extern "C" __global__ void tf_glm_lat_write(const __nv_bfloat16* lat, int64_t lat_stride, __nv_bfloat16* cache, const int* pos, int64_t lw) { const int64_t r = blockIdx.x, at = (int64_t)pos[0] + r; for (int64_t k = threadIdx.x; k < lw; k += blockDim.x) cache[at * lw + k] = lat[r * lat_stride + k]; } + +// Ling's rope (exl3's GPT-J pairs), block (row, head) of 32: x + r xs + h xh into out + r os + h oh, in place ok. +extern "C" __global__ void __launch_bounds__(32) + tf_glm_rope(const __nv_bfloat16* x, int64_t xs, int64_t xh, __nv_bfloat16* out, int64_t os, int64_t oh, + const int* pos, const float* inv_freq) { + const int64_t r = blockIdx.x, h = blockIdx.y; + const int i = threadIdx.x; // pair i at angle fp32 inv_freq[i] (pos + r); explicit rn ops, no fma contraction + const __nv_bfloat16* src = x + r * xs + h * xh + 2 * i; + const float x0 = bf(src[0]), x1 = bf(src[1]); + float s, c; + sincosf(__fmul_rn(inv_freq[i], (float)((int64_t)pos[0] + r)), &s, &c); + __nv_bfloat16* dst = out + r * os + h * oh + 2 * i; + dst[0] = __float2bfloat16_rn(__fsub_rn(__fmul_rn(x0, c), __fmul_rn(x1, s))); + dst[1] = __float2bfloat16_rn(__fadd_rn(__fmul_rn(x1, c), __fmul_rn(x0, s))); +} + +// Ling's head-wise output gate, block (row, head) of dv threads: o *= sigmoid(g[r gs + h]) in fp32, rounded to bf16. +extern "C" __global__ void __launch_bounds__(256) + tf_glm_head_gate(__nv_bfloat16* o, int64_t heads, const float* g, int64_t gs) { + const int64_t r = blockIdx.x, h = blockIdx.y, dv = blockDim.x; + const float s = __fdiv_rn(1.0f, __fadd_rn(1.0f, expf(-g[r * gs + h]))); + __nv_bfloat16* p = o + (r * heads + h) * dv + threadIdx.x; + *p = __float2bfloat16_rn(__fmul_rn(bf(*p), s)); +} diff --git a/zig/kernels/cuda/glm_glue_moe.cuh b/zig/kernels/cuda/glm_glue_moe.cuh index 7f2a27d6b0..d156092022 100644 --- a/zig/kernels/cuda/glm_glue_moe.cuh +++ b/zig/kernels/cuda/glm_glue_moe.cuh @@ -16,45 +16,140 @@ __device__ __forceinline__ float group64(const float* q) { return g; } +// AFTER false: GLM's silu(min(g, L)); true: Ling's min(silu(g), L) (inclusionAI's SwigluStepAndMul). Both clip u. +template __device__ __forceinline__ float swiglu1(float g, float u, float limit) { - g = fminf(g, limit); + if (!AFTER) g = fminf(g, limit); u = fminf(fmaxf(u, 0.0f - limit), limit); - const float s = rbf(div_full(g, exp_approx(0.0f - g) + 1.0f)); + float s = rbf(div_full(g, exp_approx(0.0f - g) + 1.0f)); + if (AFTER) s = fminf(s, limit); return rbf(__fmul_rn(u, s)); } +// Group-limited routing: -inf on every expert outside the `keep` groups of the best top-2 sums (ties to lower group). +template +__device__ __forceinline__ void keep_groups(float* choice, int lane, int ne, int gs, int ng, int keep) { + float gscore = -INFINITY; // lane g: group g's top-2 sum + const bool whole = gs % 32 == 0 && ne % 32 == 0; // groups of whole i: one pass, no divisions (Ling's 64) + if (whole) { + const int per = gs / 32, ni = ne / 32; + float a = -INFINITY, b = -INFINITY; + int g = 0, left = per; +#pragma unroll + for (int i = 0; i < MAXE; ++i) { + if (i >= ni) break; + const float c = choice[i]; + if (c > a) { + b = a; + a = c; + } else if (c > b) { + b = c; + } + if (--left == 0) { + for (int o = 16; o > 0; o >>= 1) { + const float oa = __shfl_xor_sync(0xffffffffu, a, o), ob = __shfl_xor_sync(0xffffffffu, b, o); + b = fmaxf(fminf(a, oa), fmaxf(b, ob)); + a = fmaxf(a, oa); + } + if (lane == g) gscore = a + b; + a = b = -INFINITY; + ++g; + left = per; + } + } + } + for (int g = 0; g < (whole ? 0 : ng); ++g) { + float a = -INFINITY, b = -INFINITY; // this lane's best two in group g, a >= b + for (int i = 0; i < MAXE; ++i) { + const int e = i * 32 + lane; + if (e < ne && e / gs == g) { + const float c = choice[i]; + if (c > a) { + b = a; + a = c; + } else if (c > b) { + b = c; + } + } + } + for (int o = 16; o > 0; o >>= 1) { + const float oa = __shfl_xor_sync(0xffffffffu, a, o), ob = __shfl_xor_sync(0xffffffffu, b, o); + b = fmaxf(fminf(a, oa), fmaxf(b, ob)); + a = fmaxf(a, oa); + } + if (lane == g) gscore = a + b; + } + bool live = lane < ng; + unsigned kept = 0u; + for (int k = 0; k < keep; ++k) { + float mx = live ? gscore : -INFINITY; + for (int o = 16; o > 0; o >>= 1) mx = fmaxf(mx, __shfl_xor_sync(0xffffffffu, mx, o)); + const unsigned at = __ballot_sync(0xffffffffu, live && gscore == mx); + if (at == 0u) break; + const int id = __ffs(at) - 1; + kept |= 1u << id; + if (lane == id) live = false; + } + if (whole) { + const int per = gs / 32, ni = ne / 32; + int g = 0, left = per; +#pragma unroll + for (int i = 0; i < MAXE; ++i) { + if (i >= ni) break; + if (!((kept >> g) & 1u)) choice[i] = -INFINITY; + if (--left == 0) { + ++g; + left = per; + } + } + return; + } + for (int i = 0; i < MAXE; ++i) { + const int e = i * 32 + lane; + if (e < ne && !((kept >> (e / gs)) & 1u)) choice[i] = -INFINITY; + } +} + } // namespace glm_moe // glue._swiglu with BLOCK 512 (128 threads, 4 values each, 64-group sums over 16 lanes), column block blockIdx.y. -extern "C" __global__ void __launch_bounds__(128) - tf_glm_swiglu_512(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { - const int64_t r = blockIdx.x, d0 = blockIdx.y * 512 + threadIdx.x * 4; - const uint2 gv = *reinterpret_cast(gu + r * 2 * w + d0); - const uint2 uv = *reinterpret_cast(gu + r * 2 * w + w + d0); - const __nv_bfloat16* gh = reinterpret_cast(&gv); - const __nv_bfloat16* uh = reinterpret_cast(&uv); - float a[4]; - uint2 ov; - __nv_bfloat16* oh = reinterpret_cast<__nv_bfloat16*>(&ov); - for (int j = 0; j < 4; ++j) { - a[j] = glm_moe::swiglu1(bf(gh[j]), bf(uh[j]), limit); - oh[j] = __float2bfloat16_rn(a[j]); +#define TF_GLM_SWIGLU_512(NAME, AFTER) \ + extern "C" __global__ void __launch_bounds__(128) \ + NAME(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { \ + const int64_t r = blockIdx.x, d0 = blockIdx.y * 512 + threadIdx.x * 4; \ + const uint2 gv = *reinterpret_cast(gu + r * 2 * w + d0); \ + const uint2 uv = *reinterpret_cast(gu + r * 2 * w + w + d0); \ + const __nv_bfloat16* gh = reinterpret_cast(&gv); \ + const __nv_bfloat16* uh = reinterpret_cast(&uv); \ + float a[4]; \ + uint2 ov; \ + __nv_bfloat16* oh = reinterpret_cast<__nv_bfloat16*>(&ov); \ + for (int j = 0; j < 4; ++j) { \ + a[j] = glm_moe::swiglu1(bf(gh[j]), bf(uh[j]), limit); \ + oh[j] = __float2bfloat16_rn(a[j]); \ + } \ + *reinterpret_cast(out + r * w + d0) = ov; \ + const float g = glm_moe::group64<4>(a); \ + if ((threadIdx.x & 15) == 0) xs[r * (w / 64) + d0 / 64] = g; \ } - *reinterpret_cast(out + r * w + d0) = ov; - const float g = glm_moe::group64<4>(a); - if ((threadIdx.x & 15) == 0) xs[r * (w / 64) + d0 / 64] = g; -} // glue._swiglu with BLOCK 64 (the TP=3 shared expert's 704): one value a thread, the warp's butterfly, then 2 warps. -extern "C" __global__ void __launch_bounds__(64) - tf_glm_swiglu_64(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { - __shared__ float smem[2]; - const int64_t r = blockIdx.x, d = blockIdx.y * 64 + threadIdx.x; - const float a = glm_moe::swiglu1(bf(gu[r * 2 * w + d]), bf(gu[r * 2 * w + w + d]), limit); - out[r * w + d] = __float2bfloat16_rn(a); - const float g = block_sum<2>(a, smem); - if (threadIdx.x == 0) xs[r * (w / 64) + blockIdx.y] = g; -} +#define TF_GLM_SWIGLU_64(NAME, AFTER) \ + extern "C" __global__ void __launch_bounds__(64) \ + NAME(const __nv_bfloat16* gu, __nv_bfloat16* out, float* xs, float limit, int64_t w) { \ + __shared__ float smem[2]; \ + const int64_t r = blockIdx.x, d = blockIdx.y * 64 + threadIdx.x; \ + const float a = glm_moe::swiglu1(bf(gu[r * 2 * w + d]), bf(gu[r * 2 * w + w + d]), limit); \ + out[r * w + d] = __float2bfloat16_rn(a); \ + const float g = block_sum<2>(a, smem); \ + if (threadIdx.x == 0) xs[r * (w / 64) + blockIdx.y] = g; \ + } + +TF_GLM_SWIGLU_512(tf_glm_swiglu_512, false) +TF_GLM_SWIGLU_64(tf_glm_swiglu_64, false) +// Ling's order, the limit on silu's output: min(silu(g), L) * clip(u, -L, L). +TF_GLM_SWIGLU_512(tf_ling_swiglu_512, true) +TF_GLM_SWIGLU_64(tf_ling_swiglu_64, true) // qmm._group_sums: row blockIdx.x's 64-input sums, 8 lanes a group of 8 consecutive inputs each, 16 groups a block. extern "C" __global__ void __launch_bounds__(128) @@ -66,10 +161,10 @@ extern "C" __global__ void __launch_bounds__(128) if ((threadIdx.x & 7) == 0 && g < kg) xs[m * kg + g] = s; } -// glue._topk, a warp a row: top_k by sigmoid + bias (ties to lower id), normed if `norm`, x scale; then shared expert. +// glue._topk, a warp a row: top_k of sigmoid + bias (lower id on ties) in the kept groups, normed, x scale; + shared. extern "C" __global__ void __launch_bounds__(32) tf_glm_topk(const float* logits, const float* bias, int* pick, float* wts, float scale, int ne, int top_k, - int norm) { + int norm, int n_group, int topk_group) { constexpr int MAXE = 16; // 512 experts a row const int64_t r = blockIdx.x; const int lane = threadIdx.x; @@ -79,6 +174,7 @@ extern "C" __global__ void __launch_bounds__(32) score[i] = e < ne ? div_full(1.0f, exp_approx(0.0f - logits[r * ne + e]) + 1.0f) : 0.0f; choice[i] = e < ne ? score[i] + bias[e] : -INFINITY; } + if (n_group > 1) glm_moe::keep_groups(choice, lane, ne, ne / n_group, n_group, topk_group); float total = 0.0f, mine = 0.0f; int mine_id = 0; for (int k = 0; k < top_k; ++k) { diff --git a/zig/kernels/cuda/ling_glue.cu b/zig/kernels/cuda/ling_glue.cu new file mode 100644 index 0000000000..83373f50f0 --- /dev/null +++ b/zig/kernels/cuda/ling_glue.cu @@ -0,0 +1,57 @@ +// Ling-3.1-flash's bf16 glue at the HF model's rounding points: the residual add, the head-wise gate, the MoE sum. +#include +#include + +namespace ling_glue { +__device__ __forceinline__ float bf(__nv_bfloat16 x) { return __bfloat162float(x); } +__device__ __forceinline__ float rbf(float x) { return __bfloat162float(__float2bfloat16_rn(x)); } +} // namespace ling_glue + +// torch's bf16 x + y: the fp32 sum rounded once; out may alias x or y. +extern "C" __global__ void tf_ling_add(const __nv_bfloat16* x, const __nv_bfloat16* y, __nv_bfloat16* out, int64_t n) { + const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (i < n) out[i] = __float2bfloat16_rn(__fadd_rn(ling_glue::bf(x[i]), ling_glue::bf(y[i]))); +} + +// The head-wise gate: o [rows][heads][dv] times bf16(sigmoid(float(g [rows][gs], column h))), rounded to bf16. +extern "C" __global__ void tf_ling_gate(__nv_bfloat16* o, int64_t heads, const __nv_bfloat16* g, int64_t gs) { + const int64_t r = blockIdx.x, h = blockIdx.y, dv = blockDim.x; + const float z = ling_glue::bf(g[r * gs + h]); + const float s = ling_glue::rbf(__fdiv_rn(1.0f, __fadd_rn(1.0f, expf(0.0f - z)))); + __nv_bfloat16* p = o + (r * heads + h) * dv + threadIdx.x; + *p = __float2bfloat16_rn(__fmul_rn(ling_glue::bf(*p), s)); +} + +// moe_infer's sum, bf16(sum of float(y) w over k < used) + shared, y [rows][slots][d] or by_pick [experts][rows][d]. +extern "C" __global__ void tf_ling_combine(const __nv_bfloat16* y, const int* pick, const float* wts, + const __nv_bfloat16* shared, __nv_bfloat16* out, int64_t rows, int64_t d, + int slots, int used, int by_pick) { + const int64_t r = blockIdx.x, c = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; + if (c >= d) return; + float a = 0.0f; + for (int k = 0; k < used; ++k) { + const int64_t at = by_pick ? (static_cast(pick[r * slots + k]) * rows + r) * d + c : (r * slots + k) * d + c; + a = __fadd_rn(a, __fmul_rn(ling_glue::bf(y[at]), wts[r * slots + k])); + } + float v = ling_glue::rbf(a); + if (shared != nullptr) v = __fadd_rn(v, ling_glue::bf(shared[r * d + c])); + out[r * d + c] = __float2bfloat16_rn(v); +} + +// Bytes [rows][row_bytes] between pitched rows, as they are: 16 bytes a thread where every address is 16-aligned. +extern "C" __global__ void tf_ling_copy_rows(const unsigned char* __restrict__ src, long long src_pitch, + unsigned char* __restrict__ dst, long long dst_pitch, int rows, + int row_bytes) { + const unsigned long long a = reinterpret_cast(src) | reinterpret_cast(dst); + const bool wide = ((a | static_cast(src_pitch | dst_pitch | row_bytes)) & 15) == 0; + const int step = gridDim.x * blockDim.x, at = blockIdx.x * blockDim.x + threadIdx.x; + for (int r = blockIdx.y; r < rows; r += gridDim.y) { + const unsigned char* s = src + r * src_pitch; + unsigned char* d = dst + r * dst_pitch; + if (wide) { + for (int c = at; c < row_bytes / 16; c += step) reinterpret_cast(d)[c] = reinterpret_cast(s)[c]; + } else { + for (int c = at; c < row_bytes; c += step) d[c] = s[c]; + } + } +} diff --git a/zig/kernels/cuda/modelopt_pack.cu b/zig/kernels/cuda/modelopt_pack.cu index 000c834358..c7df5e1ea7 100644 --- a/zig/kernels/cuda/modelopt_pack.cu +++ b/zig/kernels/cuda/modelopt_pack.cu @@ -106,3 +106,12 @@ extern "C" __global__ void __launch_bounds__(144) tf_pack_fp4_experts(const uint dst[w] = word; } } + +// Experts' e4m3 [E][n][k] as part m of [E][n/32][k/32][parts][32][2][4] words (Fp8Experts.packOne's order). +extern "C" __global__ void __launch_bounds__(256) tf_pack_fp8_experts(const uint32_t* __restrict__ src, uint32_t* __restrict__ out, + int n, int k, int parts, int m) { + const int g = blockIdx.x, cb = blockIdx.y, e = blockIdx.z, w = threadIdx.x, kg = k / 32; + const int j = w % 4, h = (w / 4) % 2, t = (w / 8) % 4, gq = w / 32; + const long long row = static_cast(e) * n + cb * 32 + j * 8 + gq; + out[(((static_cast(e) * (n / 32) + cb) * kg + g) * parts + m) * 256 + w] = src[row * (k / 4) + g * 8 + h * 4 + t]; +} diff --git a/zig/kernels/cuda/nvfp4_experts.cu b/zig/kernels/cuda/nvfp4_experts.cu index be724dead8..611d010bde 100644 --- a/zig/kernels/cuda/nvfp4_experts.cu +++ b/zig/kernels/cuda/nvfp4_experts.cu @@ -310,14 +310,15 @@ __device__ __forceinline__ void prompt_body(const __nv_bfloat16* __restrict__ X, #undef TF_TILES } -// The instances: SwiGLU gate-up, fp32 and bf16 down, and relu^2 up for experts without a gate (Nemotron). +// The instances: SwiGLU gate-up (EPI 4 in Ling's order), fp32 and bf16 down, relu^2 up without a gate (Nemotron). template __global__ void nvfp4_expert_kernel<2, 2, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); +template __global__ void nvfp4_expert_kernel<2, 4, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 0, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 3, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __restrict__, int, int, const uint4* __restrict__, const float* __restrict__, int, int, const int* __restrict__, const int* __restrict__, const int* __restrict__, void* __restrict__, int, float, int); } -// The prompt form's entries, 4 row tiles a warp in one row of WN warps: relu^2 up, SwiGLU gate-up, bf16 sums down. +// The prompt form's entries, 4 row tiles a warp in a row of WN warps: relu^2 up, SwiGLU gate-up (2 orders), down. #define TF_NVFP4_PROMPT(NAME, M, EPI, WN) \ extern "C" __global__ void __launch_bounds__(WN * 32) NAME( \ const __nv_bfloat16* __restrict__ X, int x_stride, int slots, const uint4* __restrict__ W, \ @@ -329,3 +330,4 @@ template __global__ void nvfp4_expert_kernel<1, 1, 4>(const __nv_bfloat16* __res TF_NVFP4_PROMPT(tf_nvfp4_prompt_relu2, 1, 1, 8) TF_NVFP4_PROMPT(tf_nvfp4_prompt_swiglu, 2, 2, 4) TF_NVFP4_PROMPT(tf_nvfp4_prompt_down, 1, 3, 8) +TF_NVFP4_PROMPT(tf_nvfp4_prompt_swiglu_after, 2, 4, 4) diff --git a/zig/kernels/cuda/probe.cu b/zig/kernels/cuda/probe.cu index ca8f345e7b..a3edf6de24 100644 --- a/zig/kernels/cuda/probe.cu +++ b/zig/kernels/cuda/probe.cu @@ -60,3 +60,12 @@ extern "C" __global__ void tf_probe_pdl_step(unsigned long long* counter, unsign asm volatile("griddepcontrol.launch_dependents;" ::: "memory"); #endif } + +// One thread busy for `ns` of the global timer: a stream kept busy while the host queues what follows it. +extern "C" __global__ void tf_probe_spin(long long ns) { + long long t0, t; + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t0)); + do { + asm volatile("mov.u64 %0, %%globaltimer;" : "=l"(t)); + } while (t - t0 < ns); +} diff --git a/zig/src/cli/hub.zig b/zig/src/cli/hub.zig index 99d6835a49..318dd88eb6 100644 --- a/zig/src/cli/hub.zig +++ b/zig/src/cli/hub.zig @@ -15,6 +15,7 @@ const titles = [_]struct { model_type: []const u8, title: []const u8 }{ .{ .model_type = "glm5_next", .title = "GLM-5.3-Flash" }, .{ .model_type = "qwen3_5", .title = "Qwen3.5-2B or Qwen3.8-27B" }, .{ .model_type = "lfm2_moe", .title = "LFM2.5-24B-A2B" }, + .{ .model_type = "bailing_hybrid", .title = "Ling-3.1-flash" }, }; /// The family that serves `model_type` from the engine registry, or null when Zig cannot serve it (the 0.6 line may). diff --git a/zig/src/cuda/experts.zig b/zig/src/cuda/experts.zig index 8a6664cbd6..8f6d950457 100644 --- a/zig/src/cuda/experts.zig +++ b/zig/src/cuda/experts.zig @@ -26,6 +26,7 @@ pub const Layer = struct { dims: usize, // D experts: usize, limit: f32 = 0.0, // fp8g, nvfp4: the SwiGLU clamp, 0 for none + after: bool = false, // fp8g, nvfp4: clamp silu's output (Ling: min(silu(g), L)), not its input (silu(min(g, L))) relu2: bool = false, // nvfp4: up is one projection under relu^2 (Nemotron), not SwiGLU's gate and up }; @@ -47,7 +48,7 @@ pub const Grouped = struct { .affine4 => { const a = g.affine4 orelse return error.FormatNotLoaded; if (epi == .down_bf16) return error.EpilogueNotBuilt; - if (skip >= 0 or l.limit != 0) return error.OptionNotBuilt; + if (skip >= 0 or l.limit != 0 or l.after) return error.OptionNotBuilt; const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); @@ -64,7 +65,7 @@ pub const Grouped = struct { .affine4 => { const a = g.affine4 orelse return error.FormatNotLoaded; if (epi == .down_f32) return error.EpilogueNotBuilt; - if (skip >= 0 or l.limit != 0) return error.OptionNotBuilt; + if (skip >= 0 or l.limit != 0 or l.after) return error.OptionNotBuilt; const up = epi == .up; const nb = (if (up) l.width else l.dims) / 32; const w = Affine4Experts.view(l, up); @@ -179,6 +180,8 @@ pub const Fp8Experts = struct { mod: Module, fns: [2][3]Function, // decode, prompt; by Epi resident: [2][3]usize, // blocks the grid may hold: per SM times SMs + after: [2]Function, // decode, prompt: gate-up in Ling's clamp order (Layer.after) + after_resident: [2]usize, pub const cols = 32; // output columns a warp pub const block = 128; // the checkpoint's scale block, both ways @@ -190,18 +193,24 @@ pub const Fp8Experts = struct { pub const p_gate_up = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi2ELi2ELi2ELi2EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const p_down_f32 = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi1ELi0ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const p_down_bf16 = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi1ELi3ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; + pub const gate_up_after = "_ZN14tf_fp8_experts17fp8_expert_kernelILi2ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; + pub const p_gate_up_after = "_ZN14tf_fp8_experts24fp8_expert_prompt_kernelILi2ELi4ELi2ELi2EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; }; pub fn load(d: *const Driver, sms: usize) !Fp8Experts { if (!kernels.available) return error.BuiltWithoutKernels; var mod = try Module.load(d, kernels.fp8_experts); errdefer mod.unload(); - var e: Fp8Experts = .{ .mod = mod, .fns = undefined, .resident = undefined }; + var e: Fp8Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .after = undefined, .after_resident = undefined }; const names = [2][3][:0]const u8{ .{ symbols.gate_up, symbols.down_f32, symbols.down_bf16 }, .{ symbols.p_gate_up, symbols.p_down_f32, symbols.p_down_bf16 } }; for (names, 0..) |set, i| for (set, 0..) |name, j| { e.fns[i][j] = try mod.function(name); e.resident[i][j] = @max(1, try e.fns[i][j].occupancy(128, 0)) * sms; }; + for ([2][:0]const u8{ symbols.gate_up_after, symbols.p_gate_up_after }, 0..) |name, i| { + e.after[i] = try mod.function(name); + e.after_resident[i] = @max(1, try e.after[i].occupancy(128, 0)) * sms; + } return e; } @@ -217,10 +226,12 @@ pub const Fp8Experts = struct { const n = if (gate) l.width else l.dims; const pi: usize = @intFromBool(prompt); const ei: usize = @backingInt(epi); + const after = gate and l.after; + const resident = if (after) e.after_resident[pi] else e.resident[pi][ei]; const grid: usize = if (prompt) blk: { const cw: usize = if (gate) 2 else 4; - break :blk @min(items * (nb / cw), e.resident[pi][ei]); - } else @min((items * nb + 3) / 4, e.resident[pi][ei]); + break :blk @min(items * (nb / cw), resident); + } else @min((items * nb + 3) / 4, resident); if (grid < 1) return; var a: launch_.Args = .{}; a.add(in.x); @@ -234,7 +245,7 @@ pub const Fp8Experts = struct { a.add(@as(i32, @intCast(n))); a.add(if (gate) l.limit else @as(f32, 0.0)); a.add(skip); - try launch_.launch(e.fns[pi][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &a); + try launch_.launch(if (after) e.after[pi] else e.fns[pi][ei], .{ .grid = .{ .x = @intCast(grid), .y = 1, .z = 1 }, .block = .{ .x = 128, .y = 1, .z = 1 } }, s, &a); } /// e4m3 [n, k] of one expert as int32 [n/32][k/32][32][2][4]: lane (gq, t)'s half h, tile j (fp8/experts.py pack). @@ -270,9 +281,9 @@ pub const Fp8Experts = struct { /// NVFP4 experts (nvfp4_experts.cu): decode rows 16 pairs at a time; prompt rows staged, 64 pairs a CTA. pub const Nvfp4Experts = struct { mod: Module, - fns: [4]Function, // by Epi, then relu^2 up - resident: [4]usize, // blocks the grid may hold: per SM times SMs - prompt_fns: [3]Function, // relu^2 up, SwiGLU gate-up, bf16 sums down + fns: [5]Function, // by Epi, then relu^2 up, then gate-up in Ling's clamp order (Layer.after) + resident: [5]usize, // blocks the grid may hold: per SM times SMs + prompt_fns: [4]Function, // relu^2 up, SwiGLU gate-up, bf16 sums down, gate-up in Ling's order pub const cols = 32; // output columns a block pub const words = 144; // int32 a (32 columns, 32 inputs) block: 128 code words, then 16 of e4m3 scales @@ -282,11 +293,12 @@ pub const Nvfp4Experts = struct { pub const down_f32 = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi1ELi0ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const down_bf16 = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi1ELi3ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; pub const up_relu2 = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi1ELi1ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; - pub const prompt = [3][:0]const u8{ "tf_nvfp4_prompt_relu2", "tf_nvfp4_prompt_swiglu", "tf_nvfp4_prompt_down" }; + pub const gate_up_after = "_ZN16tf_nvfp4_experts19nvfp4_expert_kernelILi2ELi4ELi4EEEvPK13__nv_bfloat16iiPK5uint4PKfiiPKiSA_SA_Pvifi"; + pub const prompt = [4][:0]const u8{ "tf_nvfp4_prompt_relu2", "tf_nvfp4_prompt_swiglu", "tf_nvfp4_prompt_down", "tf_nvfp4_prompt_swiglu_after" }; }; - /// The prompt entries' warps (32-column blocks) a CTA: relu^2 up, SwiGLU gate-up, down (GB10's fastest tiles). - pub const prompt_wn = [3]usize{ 8, 4, 8 }; + /// The prompt entries' warps (32-column blocks) a CTA: relu^2 up, SwiGLU gate-up, down, Ling's gate-up. + pub const prompt_wn = [4]usize{ 8, 4, 8, 4 }; /// Prompt4's dynamic shared bytes: stages of 64 rows' 64 inputs, then WN x 2 blocks of `parts`. pub fn promptSmem(parts: usize, wn: usize) u32 { @@ -300,11 +312,11 @@ pub const Nvfp4Experts = struct { var mod = try Module.load(d, kernels.nvfp4_experts); errdefer mod.unload(); var e: Nvfp4Experts = .{ .mod = mod, .fns = undefined, .resident = undefined, .prompt_fns = undefined }; - for ([_][:0]const u8{ symbols.gate_up, symbols.down_f32, symbols.down_bf16, symbols.up_relu2 }, 0..) |name, i| { + for ([_][:0]const u8{ symbols.gate_up, symbols.down_f32, symbols.down_bf16, symbols.up_relu2, symbols.gate_up_after }, 0..) |name, i| { e.fns[i] = try mod.function(name); e.resident[i] = @max(1, try e.fns[i].occupancy(128, 0)) * sms; } - for (symbols.prompt, &e.prompt_fns, [3]usize{ 1, 2, 1 }, prompt_wn) |name, *f, parts, wn| { + for (symbols.prompt, &e.prompt_fns, [4]usize{ 1, 2, 1, 2 }, prompt_wn) |name, *f, parts, wn| { f.* = try mod.function(name); try f.allowDynamicShared(promptSmem(parts, wn)); } @@ -320,7 +332,7 @@ pub const Nvfp4Experts = struct { const gate = epi == .up; const kg = (if (gate) l.dims else l.width) / 32; const nb = (if (gate) l.width else l.dims) / cols; - const ei: usize = if (gate and l.relu2) 3 else @backingInt(epi); + const ei: usize = if (gate and l.relu2) 3 else if (gate and l.after) 4 else @backingInt(epi); const grid = @min((items * nb + 3) / 4, e.resident[ei]); if (grid < 1) return; var a: launch_.Args = .{}; @@ -357,7 +369,7 @@ pub const Nvfp4Experts = struct { a.add(@as(i32, @intCast(if (gate) l.width else l.dims))); a.add(if (gate) l.limit else @as(f32, 0.0)); a.add(skip); - const fi: usize = if (!gate) 2 else if (l.relu2) 0 else 1; + const fi: usize = if (!gate) 2 else if (l.relu2) 0 else if (l.after) 3 else 1; const cfg: launch_.Config = .{ .grid = .{ .x = @intCast(items * ((nb + prompt_wn[fi] - 1) / prompt_wn[fi])), .y = 1, .z = 1 }, .block = .{ .x = @intCast(prompt_wn[fi] * 32), .y = 1, .z = 1 }, .shared = promptSmem(parts, prompt_wn[fi]) }; try launch_.launch(e.prompt_fns[fi], cfg, s, &a); } diff --git a/zig/src/cuda/kernels.zig b/zig/src/cuda/kernels.zig index 18542b8cb4..dfb0f74ed4 100644 --- a/zig/src/cuda/kernels.zig +++ b/zig/src/cuda/kernels.zig @@ -40,6 +40,7 @@ pub const glm_glue: []const u8 = if (available) &Blob("fatbin_glm_glue").bytes e pub const glm_ops: []const u8 = if (available) &Blob("fatbin_glm_ops").bytes else &.{}; pub const glm_quant: []const u8 = if (available) &Blob("fatbin_glm_quant").bytes else &.{}; pub const glm_kda: []const u8 = if (available) &Blob("fatbin_glm_kda").bytes else &.{}; +pub const ling_glue: []const u8 = if (available) &Blob("fatbin_ling_glue").bytes else &.{}; pub const torch_argmax: []const u8 = if (available) &Blob("fatbin_torch_argmax").bytes else &.{}; pub const torch_topk: []const u8 = if (available) &Blob("fatbin_torch_topk").bytes else &.{}; pub const torch_pointwise: []const u8 = if (available) &Blob("fatbin_torch_pointwise").bytes else &.{}; diff --git a/zig/src/cuda/modelopt.zig b/zig/src/cuda/modelopt.zig index 4338ef4db9..c489b7ede0 100644 --- a/zig/src/cuda/modelopt.zig +++ b/zig/src/cuda/modelopt.zig @@ -16,7 +16,7 @@ const Stream = stream_.Stream; /// Rows padded to whole 128-row blocks, as every lane format pads them. pub const padded = nvfp4.padded; -/// modelopt_pack.cu's kernels; each matches a host packer (fp8.packCodes, qmmf.packBf16, nvfp4.pack*, Nvfp4Experts). +/// modelopt_pack.cu's kernels; each matches a host packer (fp8.packCodes, qmmf.packBf16, nvfp4.pack*, the experts'). pub const Packer = struct { mod: Module, fp8_lane: Function, @@ -24,6 +24,7 @@ pub const Packer = struct { fp4_words: Function, fp4_scales: Function, fp4_experts: Function, + fp8_experts: Function, pub fn load(d: *const Driver) !Packer { if (!kernels.available) return error.BuiltWithoutKernels; @@ -36,6 +37,7 @@ pub const Packer = struct { .fp4_words = try mod.function("tf_pack_fp4_words"), .fp4_scales = try mod.function("tf_pack_fp4_scales"), .fp4_experts = try mod.function("tf_pack_fp4_experts"), + .fp8_experts = try mod.function("tf_pack_fp8_experts"), }; } @@ -79,6 +81,17 @@ pub const Packer = struct { for ([_]usize{ n, ks, k0, k, parts, m }) |v| a.add(@as(i32, @intCast(v))); try launch_.launch(p.fp4_experts, .{ .grid = .{ .x = @intCast(k / 32), .y = @intCast(n / 32), .z = @intCast(count) }, .block = .{ .x = 144 } }, s, &a); } + + /// `count` experts' e4m3 [E][n][k] as part m of `parts` per (column block, k32 group) (Fp8Experts.packGateUp). + pub fn fp8Experts(p: *const Packer, s: Stream, src: u64, out: u64, count: usize, n: usize, k: usize, parts: usize, m: usize) !void { + if (n % 32 != 0 or k % 32 != 0 or m >= parts) return error.UnexpectedTensor; + if (count == 0) return; + var a: launch_.Args = .{}; + a.add(src); + a.add(out); + for ([_]usize{ n, k, parts, m }) |v| a.add(@as(i32, @intCast(v))); + try launch_.launch(p.fp8_experts, .{ .grid = .{ .x = @intCast(k / 32), .y = @intCast(n / 32), .z = @intCast(count) }, .block = .{ .x = 256 } }, s, &a); + } }; /// The lane matmul's view of a packed per-tensor FP8 projection. diff --git a/zig/src/families/glm/cuda_glue.zig b/zig/src/families/glm/cuda_glue.zig index 733edcdb27..24da99eea3 100644 --- a/zig/src/families/glm/cuda_glue.zig +++ b/zig/src/families/glm/cuda_glue.zig @@ -31,7 +31,11 @@ pub const Kernels = struct { hc_partial8: cuda.Function, hc_finish: cuda.Function, absorb_q: cuda.Function, + absorb_q_128: cuda.Function, expand_v: cuda.Function, + expand_v_128: cuda.Function, + rope: cuda.Function, + head_gate: cuda.Function, lat_write: cuda.Function, index_write: cuda.Function, pool_keys: cuda.Function, @@ -46,6 +50,7 @@ pub const Kernels = struct { bmm_d64s2: cuda.Function, chunks_16: cuda.Function, chunks_32: cuda.Function, + chunks_r16: cuda.Function, merge: cuda.Function, scores: cuda.Function, @@ -56,6 +61,7 @@ pub const Kernels = struct { try k.bmm_tile1.allowDynamicShared(bmm_tile1_smem[1]); try k.chunks_16.allowDynamicShared(chunksSmem(16)); try k.chunks_32.allowDynamicShared(chunksSmem(32)); + try k.chunks_r16.allowDynamicShared(chunksSmemRoped(16)); return k; } }; @@ -108,6 +114,11 @@ fn cdiv(a: usize, b: usize) u32 { return @intCast((a + b - 1) / b); } +/// The kernel's clamp for a limit: 0 (unclamped, as the experts' 0) as +inf, whose min and clip are exact. +pub fn clampOf(limit: f32) f32 { + return if (limit == 0) std.math.inf(f32) else limit; +} + /// swiglu's column block: 512 where it divides the width, else 64 (TP=3's 704 and 640, the latter 128 in glue.py). fn swigluBlock(w: usize) !usize { if (w % 512 == 0) return 512; @@ -115,12 +126,21 @@ fn swigluBlock(w: usize) !usize { return error.NoSwigluBlock; } -/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`. +/// glue.swiglu: [gate | up] -> bf16(bf16(silu(min(g, L))) * clip(u, -L, L)), 64-sums in `xs`; L as given (GLM's). pub fn swiglu(set: Set, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { + return swigluLaunch(.{ set.glue.swiglu_512, set.glue.swiglu_64 }, stream, gu, out, xs, limit, rows, w); +} + +/// swiglu on a pair of glm_glue_moe.cuh's 512 and 64 kernels (tf_ling_swiglu_*, the limit on silu's output); 0: none. +pub fn swigluWith(fs: [2]cuda.Function, stream: cuda.Stream, gu: u64, out: u64, xs: u64, limit: f32, rows: usize, w: usize) !void { + return swigluLaunch(fs, stream, gu, out, xs, clampOf(limit), rows, w); +} + +fn swigluLaunch(fs: [2]cuda.Function, stream: cuda.Stream, gu: u64, out: u64, xs: u64, clamp: f32, rows: usize, w: usize) !void { const block = try swigluBlock(w); - const f, const threads: u32 = if (block == 512) .{ set.glue.swiglu_512, 128 } else .{ set.glue.swiglu_64, 64 }; + const f, const threads: u32 = if (block == 512) .{ fs[0], 128 } else .{ fs[1], 64 }; var a: cuda.Args = .{}; - inline for (.{ gu, out, xs, limit, i64of(w) }) |v| a.add(v); + inline for (.{ gu, out, xs, clamp, i64of(w) }) |v| a.add(v); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(w / block) }, .block = .{ .x = threads } }, stream, &a); } @@ -133,6 +153,11 @@ pub fn chunksSmem(comptime hb: u32) u32 { return (hb * 512 + 2 * 32 * 512 + hb * 32) * 2 + (hb * 4 + 2 * hb) * 4; } +/// wide_smem(): as above with q and key rows 576 wide (Ling's roped columns after the latent). +pub fn chunksSmemRoped(comptime hb: u32) u32 { + return (hb * 576 + 2 * 32 * 576 + hb * 32) * 2 + (hb * 4 + 2 * hb) * 4; +} + /// K slices of the router's dot (glue.ROUTER_KS): fixed by the shape, as the sliced and the single pass differ in bits. pub const router_ks = 8; @@ -148,11 +173,15 @@ pub fn router(set: Set, stream: cuda.Stream, x: u64, x_stride: usize, w: u64, ou try sliceSum(set, stream, part, out, m * ne, router_ks, true); } -/// glue.select: top_k by sigmoid + bias (lower id on ties), normed if `norm`, times `scale`, then the shared expert. -pub fn select(set: Set, stream: cuda.Stream, logits: u64, bias: u64, pick: u64, wts: u64, rows: usize, top_k: usize, experts: usize, scale: f32, norm: bool) !void { +/// Group-limited routing (Ling's n_group 8, topk_group 4): GLM's one group keeps every expert. +pub const Groups = struct { n: usize = 1, keep: usize = 1 }; + +/// glue.select: top_k by sigmoid + bias (lower id on ties) in the kept groups, normed if `norm`, x `scale`, + shared. +pub fn select(set: Set, stream: cuda.Stream, logits: u64, bias: u64, pick: u64, wts: u64, rows: usize, top_k: usize, experts: usize, scale: f32, norm: bool, groups: Groups) !void { if (experts > 512 or top_k >= 32) return error.TopkTooWide; + if (groups.n == 0 or groups.n > 32 or experts % groups.n != 0 or groups.keep == 0 or groups.keep > groups.n or top_k > groups.keep * (experts / groups.n)) return error.TopkGroups; var a: cuda.Args = .{}; - inline for (.{ logits, bias, pick, wts, scale, cint(experts), cint(top_k), @as(c_int, @intFromBool(norm)) }) |v| a.add(v); + inline for (.{ logits, bias, pick, wts, scale, cint(experts), cint(top_k), @as(c_int, @intFromBool(norm)), cint(groups.n), cint(groups.keep) }) |v| a.add(v); try cuda.launch.launch(set.glue.topk, .{ .grid = .{ .x = @intCast(rows) }, .block = .{ .x = 32 } }, stream, &a); } diff --git a/zig/src/families/glm/cuda_latent.zig b/zig/src/families/glm/cuda_latent.zig index a41381ef3b..6d2b08ad31 100644 --- a/zig/src/families/glm/cuda_latent.zig +++ b/zig/src/families/glm/cuda_latent.zig @@ -1,4 +1,4 @@ -//! GLM-5.3-Flash's MLA on the 512-wide latent cache (DSA layers): the Python family's latent on BF16, row by row. +//! GLM-5.3-Flash's MLA on the 512-wide latent cache (DSA layers), and Ling-3.1's with 64 roped columns after it. const std = @import("std"); const cuda = @import("cuda"); @@ -29,21 +29,60 @@ fn cdiv(a: usize, b: usize) u32 { /// latent.absorb_q (BF16 kv_b): out [rows][heads][lw] bf16 = q [rows][heads][d] . wk [heads][d][lw], fp32 sums. pub fn absorbQ(set: glue.Set, stream: cuda.Stream, q: u64, wk: u64, out: u64, rows: usize, heads: usize, d: usize, lw: usize) !void { if (d != 256 or lw != 512) return error.LatentShape; - try rowBlocks(set.glue.absorb_q, stream, q, wk, out, rows, heads, lw / 32); + try rowBlocks(set.glue.absorb_q, stream, q, wk, out, rows, heads, lw / 32, .{}); +} + +/// Ling's absorb: q_nope the first d = 128 of each head's `q_stride` columns, out rows `out_stride` apart per head. +pub fn absorbQStrided(set: glue.Set, stream: cuda.Stream, q: u64, q_stride: usize, wk: u64, out: u64, out_stride: usize, rows: usize, heads: usize, d: usize, lw: usize) !void { + if (d != 128 or lw != 512 or q_stride < d or out_stride < lw) return error.LatentShape; + try rowBlocks(set.glue.absorb_q_128, stream, q, wk, out, rows, heads, lw / 32, .{ @as(i64, @intCast(q_stride)), @as(i64, @intCast(out_stride)) }); } /// The absorb and expand launches: rows in one block in decode windows, prompt_rb a block in prompt chunks. -fn rowBlocks(f: cuda.Function, stream: cuda.Stream, in: u64, w: u64, out: u64, rows: usize, heads: usize, cols: usize) !void { +fn rowBlocks(f: cuda.Function, stream: cuda.Stream, in: u64, w: u64, out: u64, rows: usize, heads: usize, cols: usize, extra: anytype) !void { const rb = if (rows >= prompt_rows) prompt_rb else rows; var a: cuda.Args = .{}; - inline for (.{ in, w, out, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(rb)) }) |v| a.add(v); + inline for (.{ in, w, out, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(rb)) } ++ extra) |v| a.add(v); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(heads), .y = @intCast(cols), .z = cdiv(rows, rb) }, .block = .{ .x = 128 } }, stream, &a); } -/// latent.expand_v (BF16 kv_b): out [rows][heads][dv] bf16 = o_lat [rows][heads][lw] . wv [heads][dv][lw]. +/// latent.expand_v (BF16 kv_b): out [rows][heads][dv] bf16 = o_lat [rows][heads][lw] . wv [heads][dv][lw], dv 256|128. pub fn expandV(set: glue.Set, stream: cuda.Stream, ol: u64, wv: u64, out: u64, rows: usize, heads: usize, dv: usize, lw: usize) !void { - if (dv != 256 or lw != 512) return error.LatentShape; - try rowBlocks(set.glue.expand_v, stream, ol, wv, out, rows, heads, dv / 16); + if (lw != 512) return error.LatentShape; + const f = switch (dv) { + 256 => set.glue.expand_v, + 128 => set.glue.expand_v_128, + else => return error.LatentShape, + }; + try rowBlocks(f, stream, ol, wv, out, rows, heads, dv / 16, .{}); +} + +/// Ling's 64 rope columns of `heads` heads a row: in at x + r x_row + h x_head, out likewise (bf16 elements). +pub const RopeRows = struct { x: u64, x_row: usize, x_head: usize, out: u64, o_row: usize, o_head: usize }; + +/// The rope (GPT-J pairs) at positions pos .. pos + rows - 1 (pos an i32 on device), inv_freq 32 fp32 on device. +pub fn rope(set: glue.Set, stream: cuda.Stream, v: RopeRows, rows: usize, heads: usize, pos: u64, inv_freq: u64) !void { + var a: cuda.Args = .{}; + inline for (.{ v.x, @as(i64, @intCast(v.x_row)), @as(i64, @intCast(v.x_head)), v.out, @as(i64, @intCast(v.o_row)), @as(i64, @intCast(v.o_head)), pos, inv_freq }) |x| a.add(x); + try cuda.launch.launch(set.glue.rope, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = rope_width / 2 } }, stream, &a); +} + +/// Ling's qk_rope_head_dim: the roped columns after the latent in q and in the cache. +pub const rope_width = 64; + +/// The rope's inv_freq: theta^(-2i / 64) for i < 32, from float64 rounded to fp32. +pub fn ropeFreqs(theta: f64) [rope_width / 2]f32 { + var f: [rope_width / 2]f32 = undefined; + for (&f, 0..) |*x, i| x.* = @floatCast(1.0 / std.math.pow(f64, theta, @as(f64, @floatFromInt(2 * i)) / rope_width)); + return f; +} + +/// Ling's head-wise gate: o [rows][heads][dv] bf16 *= sigmoid(g fp32 [rows][g_stride], column h), in place. +pub fn headGate(set: glue.Set, stream: cuda.Stream, o: u64, g: u64, g_stride: usize, rows: usize, heads: usize, dv: usize) !void { + if (dv == 0 or dv > 256) return error.LatentShape; + var a: cuda.Args = .{}; + inline for (.{ o, @as(i64, @intCast(heads)), g, @as(i64, @intCast(g_stride)) }) |x| a.add(x); + try cuda.launch.launch(set.glue.head_gate, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = @intCast(dv) } }, stream, &a); } /// latent.latent_write: rows of `lat` (stride `lat_stride`) into slots pos .. pos + rows - 1 (pos an i32 on device). @@ -58,28 +97,35 @@ pub const Partials = struct { po: u64, pm: u64, pl: u64 }; /// latent.attention: causal qa over the cache in `nch` 512-key chunks into out bf16; `dummy`: the merge's count i32. pub fn attention(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, p: Partials, out: u64, dummy: u64, rows: usize, heads: usize, lw: usize, nch: usize, scale: f32, hb: usize) !void { - try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, nch, 0, scale, hb, false); + try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, 0, nch, 0, scale, hb, false); + try merge(set, stream, p, out, dummy, rows, heads, lw, nch, false); +} + +/// Ling's latent.attention: qa and cache rows lw + rope_width wide (roped columns last), 16-head blocks; out lw wide. +pub fn attentionRoped(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, p: Partials, out: u64, dummy: u64, rows: usize, heads: usize, lw: usize, nch: usize, scale: f32) !void { + try chunks(set, stream, qa, cache, pos, 0, 0, p, rows, heads, lw, rope_width, nch, 0, scale, heads_tile, false); try merge(set, stream, p, out, dummy, rows, heads, lw, nch, false); } /// latent._dense_chunks (sparse false: keys through pos + r) and _sparse_chunks (row r's `counts` tokens of `tokens`). -fn chunks(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, tokens: u64, counts: u64, p: Partials, rows: usize, heads: usize, lw: usize, nch: usize, w: usize, scale: f32, hb: usize, sparse: bool) !void { - if (lw != 512) return error.LatentShape; - const f = switch (hb) { +fn chunks(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, pos: u64, tokens: u64, counts: u64, p: Partials, rows: usize, heads: usize, lw: usize, rw: usize, nch: usize, w: usize, scale: f32, hb: usize, sparse: bool) !void { + if (lw != 512 or (rw != 0 and rw != rope_width)) return error.LatentShape; + if (rw != 0 and hb != heads_tile) return error.HeadBlock; + const f = if (rw != 0) set.glue.chunks_r16 else switch (hb) { heads_tile => set.glue.chunks_16, heads_tile_wide => set.glue.chunks_32, else => return error.HeadBlock, }; var a: cuda.Args = .{}; inline for (.{ qa, cache, pos, tokens, counts, p.po, p.pm, p.pl, @as(i64, @intCast(rows)), @as(i64, @intCast(heads)), @as(i64, @intCast(w)), scale, @as(c_int, @intFromBool(sparse)) }) |v| a.add(v); - const shared = if (hb == heads_tile_wide) glue.chunksSmem(heads_tile_wide) else glue.chunksSmem(heads_tile); + const shared = if (rw != 0) glue.chunksSmemRoped(heads_tile) else if (hb == heads_tile_wide) glue.chunksSmem(heads_tile_wide) else glue.chunksSmem(heads_tile); try cuda.launch.launch(f, .{ .grid = .{ .x = @intCast(rows), .y = cdiv(heads, hb), .z = @intCast(nch) }, .block = .{ .x = 256 }, .shared = shared }, stream, &a); } /// latent.sparse_attention: rows with counts > 0 over their selected tokens (i32 [rows][w]); others left alone. pub fn sparseAttention(set: glue.Set, stream: cuda.Stream, qa: u64, cache: u64, tokens: u64, counts: u64, w: usize, p: Partials, out: u64, rows: usize, heads: usize, lw: usize, scale: f32) !void { const nch = chunksFor(w); - try chunks(set, stream, qa, cache, 0, tokens, counts, p, rows, heads, lw, nch, w, scale, headBlock(rows), true); + try chunks(set, stream, qa, cache, 0, tokens, counts, p, rows, heads, lw, 0, nch, w, scale, headBlock(rows), true); try merge(set, stream, p, out, counts, rows, heads, lw, nch, true); } diff --git a/zig/src/families/glm/cuda_moe.zig b/zig/src/families/glm/cuda_moe.zig index ec860a188e..d608597722 100644 --- a/zig/src/families/glm/cuda_moe.zig +++ b/zig/src/families/glm/cuda_moe.zig @@ -15,9 +15,11 @@ pub const Layer = struct { shared_gu: u64, shared_down: u64, shared: usize, + shared_limit: f32, // the shared expert's SwiGLU clamp, 0 for none (Ling's differs from the routed one's) top_k: usize, scale: f32, norm: bool, + groups: glue.Groups = .{}, }; /// The block's scratch for up to `rows` rows (forward.Buffers' MoE part). @@ -51,7 +53,7 @@ pub fn body(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const slots = l.top_k + 1; const shared_id: i32 = @intCast(l.count); try glue.router(k.set, stream, x, x_stride, l.router, s.mlog, s.router_part, rows, d, l.count); - try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm); + try glue.select(k.set, stream, s.mlog, l.bias, s.pick, s.wts, rows, l.top_k, l.count, l.scale, l.norm, l.groups); const plan: experts.Plan = .{ .route = s.plan, .slots = slots, .experts = l.count + 1, .prompt = prompt }; try experts.route(k.experts, stream, s.pick, rows * slots, plan); try experts.run(k.experts, stream, .gate_up, l.routed, plan, x, x_stride, rows, s.eact, shared_id); @@ -65,7 +67,7 @@ pub fn fill(k: Kernels, stream: cuda.Stream, l: Layer, s: Scratch, x: u64, x_str const n = hi - lo; const eb: usize = if (prompt) 2 else 4; // the slot rows: bf16 in a prompt chunk try b16.matmul(k.set, stream, x + lo * x_stride * 2, x_stride, l.shared_gu, s.sgu + lo * 2 * l.shared * 2, s.mm_part, n, 2 * l.shared, d, false); - try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.routed.limit, n, l.shared); + try glue.swiglu(k.set, stream, s.sgu + lo * 2 * l.shared * 2, s.sact + lo * l.shared * 2, s.sxs + lo * (l.shared / 64) * 4, l.shared_limit, n, l.shared); try b16.matmul(k.set, stream, s.sact + lo * l.shared * 2, l.shared, l.shared_down, s.sy + lo * d * 4, s.mm_part, n, d, l.shared, true); var a: cuda.Args = .{}; a.add(s.sy + lo * d * 4); diff --git a/zig/src/families/glm/cuda_weights.zig b/zig/src/families/glm/cuda_weights.zig index 80bc44a8e0..58f89293cb 100644 --- a/zig/src/families/glm/cuda_weights.zig +++ b/zig/src/families/glm/cuda_weights.zig @@ -396,6 +396,7 @@ const Loader = struct { .shared_gu = gu, .shared_down = sd, .shared = sh.len(), + .shared_limit = c.swiglu_limit, .top_k = c.topk, .scale = c.routed_scale, .norm = c.norm_topk, diff --git a/zig/src/families/ling/block_ffn.zig b/zig/src/families/ling/block_ffn.zig new file mode 100644 index 0000000000..d20010e7ef --- /dev/null +++ b/zig/src/families/ling/block_ffn.zig @@ -0,0 +1,70 @@ +//! Ling's MLPs on CUDA: SwiGLU over projections in their loaded formats, and the MoE block (router, groups, experts). +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const ling = @import("cuda.zig"); + +const glue = graph.glue; +const Ops = graph.Ops; +const Config = @import("config.zig").Config; + +/// Gate and up into [gate | up] rows, SwiGLU with min(silu(g), limit) (0: none), down into `out` [rows][hidden]. +pub fn mlp(o: Ops, b: *const state.Buffers, m: weights.Mlp, x: u64, ldx: usize, rows: usize, limit: f32, out: u64) !void { + const w = m.gate.outputs(); + if (rows == 1) { // one row is its own columns: no copies + try o.mm(m.gate, x, ldx, 1, b.gu); + try o.mm(m.up, x, ldx, 1, b.gu + w * 2); + } else { + try o.mmCols(m.gate, x, ldx, rows, b.tmp, b.gu, 2 * w, 0); + try o.mmCols(m.up, x, ldx, rows, b.tmp, b.gu, 2 * w, w); + } + try glue.swigluWith(o.m.swiglu, o.s, b.gu, b.act, b.xs, limit, rows, w); + try o.mm(m.down, b.act, w, rows, out); +} + +/// The MoE block on b.normed into b.y: routed experts' slots summed (weights fp32), plus the shared expert. +pub fn moe(o: Ops, b: *const state.Buffers, e: weights.Moe, g: *const graph.Graph, layer: u32, rows: usize, picks: ?u64) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const E: usize = c.experts; + if (e.experts != E) return error.PartialExperts; // the router picks among every expert + const top: usize = c.topk; + const slots = top + 1; + const lim = ling.limits(c, layer); + try glue.router(o.set(), o.s, b.normed, D, e.router, b.mlog, b.rpart, rows, D, E); + try glue.select(o.set(), o.s, b.mlog, if (e.bias != 0) e.bias else g.zero_bias, b.pick, b.wts, rows, top, E, c.routed_scale, c.norm_topk, ling.groups(c)); + if (picks) |p| try o.copy(b.pick, slots * 4, p, slots * 4, rows, slots * 4); + switch (e.routed) { + .grouped => |l| try grouped(o, b, l, rows, slots, E), + .lanes => |ls| for (ls, 0..) |x, j| try mlp(o, b, x, b.normed, D, rows, lim.routed, b.ey + j * rows * D * 2), + } + if (e.shared) |s| try mlp(o, b, s, b.normed, D, rows, e.shared_limit, b.sh); + var a: cuda.Args = .{}; + inline for (.{ b.ey, b.pick, b.wts, if (e.shared != null) b.sh else @as(u64, 0), b.y, @as(i64, @intCast(rows)), @as(i64, @intCast(D)), @as(c_int, @intCast(slots)), @as(c_int, @intCast(top)), @as(c_int, @intFromBool(e.routed == .lanes)) }) |v| a.add(v); + try cuda.launch.launch(o.m.combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((D + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &a); +} + +/// Prompt chunks under this many routed pairs an expert stream FP8 experts faster on the decode kernels (same bits). +const decode_pairs = 16; + +/// Grouped experts on every row's top-k slots (the shared slot skipped): gate-up SwiGLU, then down to bf16 slot rows. +fn grouped(o: Ops, b: *const state.Buffers, l: cuda.experts.Layer, rows: usize, slots: usize, count: usize) !void { + const pairs = rows * slots; + const prompt = o.prompt and !(l.format == .fp8g and rows * (slots - 1) < decode_pairs * count); + const tile: usize = if (prompt) cuda.experts.promptTile(l.format) else 16; + const skip: i32 = @intCast(count); + try o.m.router.route(o.s, b.pick, pairs, count + 1, tile, b.plan); + const items = cuda.grouped.maxItems(pairs, count + 1, tile); + const g = o.m.grouped(); + const up: cuda.experts.Rows = .{ .x = b.normed, .stride = l.dims, .slots = slots }; + const down: cuda.experts.Rows = .{ .x = b.eact, .stride = l.width }; + if (prompt) { + try g.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip); + try g.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + } else { + try g.decode(o.s, .up, up, l, b.plan, b.eact, items, skip); + try g.decode(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + } +} diff --git a/zig/src/families/ling/block_kda.zig b/zig/src/families/ling/block_kda.zig new file mode 100644 index 0000000000..3817a88911 --- /dev/null +++ b/zig/src/families/ling/block_kda.zig @@ -0,0 +1,95 @@ +//! Ling's KDA mixer on CUDA: seven projections in their loaded formats feeding GLM's chain, then o_proj. +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); + +const kda = graph.kda; +const Ops = graph.Ops; + +/// One sequence's rows of a forward: rows row0 .. row0 + rows on `st` from its position. +pub const Seg = struct { st: *state.State, row0: usize, rows: usize }; + +/// Where segment `sg`'s projection rows and saves go: the State's window rows, or the Buffers' prompt rows. +fn place(b: *const state.Buffers, sg: Seg, li: usize) struct { p: u64, k: u64, v: u64, g: u64, b: u64 } { + const st = sg.st; + if (!b.prompt) { + const s = st.saves(li); + return .{ .p = s.p, .k = s.k, .v = s.v, .g = s.g, .b = s.b }; + } + const n = st.heads * kda.dk; + const r = sg.row0; + return .{ .p = b.kp + r * st.stride * 2, .k = b.saves.k + r * n * 4, .v = b.saves.v + r * n * 2, .g = b.saves.g + r * n * 4, .b = b.saves.b + r * st.heads * 4 }; +} + +/// The block on b.normed into b.y: layer `li` among the KDA layers; prompt rows commit their states and conv windows. +pub fn block(o: Ops, b: *const state.Buffers, a: weights.Kda, c: *const @import("config.zig").Config, li: usize, segs: []const Seg, rows: usize) !void { + const D: usize = c.hidden; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const ps = segs[0].st.stride; + for ([_]weights.Proj{ a.q, a.k, a.v, a.b }, [_]usize{ 0, n, 2 * n, 3 * n }) |w, col0| { + if (rows == 1) { // one row is its own columns: no copy + try o.mm(w, b.normed, D, 1, place(b, segs[0], li).p + col0 * 2); + continue; + } + try o.mm(w, b.normed, D, rows, b.tmp); + const w_n = w.outputs(); + for (segs) |sg| try o.copy(b.tmp + sg.row0 * w_n * 2, w_n * 2, place(b, sg, li).p + col0 * 2, ps * 2, sg.rows, w_n * 2); + } + try o.mm(a.f, b.normed, D, rows, b.ka); + try o.mm(a.g, b.normed, D, rows, b.kg); + for (segs) |sg| { + const st = sg.st; + const at = place(b, sg, li); + const cur = st.cur[li]; + const ch: kda.Chain = .{ + .p = at.p, + .p_stride = ps, + .b_off = 3 * n, + .a = b.ka + sg.row0 * n * 2, + .a_stride = n, + .g = b.kg + sg.row0 * n * 2, + .g_stride = n, + .conv_state = st.convAt(li), + .conv_w = a.conv, + .state_in = st.recAt(cur, li), + .a_log = a.a_log, + .dt_bias = a.dt_bias, + .norm_w = a.o_norm, + .eps = c.eps, + .lower = c.lower_bound, + .heads = H, + }; + const sv: kda.Saves = .{ .out = b.saves.out + sg.row0 * n * 2, .state_out = st.recAt(1 - cur, li), .k = at.k, .v = at.v, .g = at.g, .b = at.b }; + const wide: ?kda.Wide = if (if (b.prompt) sg.rows >= kda.wide_rows else o.kda_wide) .{ .q = b.wide.q + sg.row0 * n * 4, .y = b.wide.y + sg.row0 * n * 2 } else null; + try kda.chain(&o.m.kda, o.s, ch, sg.rows, sv, wide); + if (!b.prompt) continue; + try graph.glue.convShift(o.set(), o.s, st.convAt(li), at.p, sg.rows, sg.rows * ps, ps, 1, 3 * n); + st.cur[li] = 1 - cur; // a prompt keeps every row: the layer commits now + } + try o.mm(a.o, b.saves.out, n, rows, b.y); +} + +/// A window's commit of `keep` rows on every KDA layer: the states replayed when rows are dropped, the conv windows. +pub fn commit(o: Ops, st: *state.State, keep: usize) !void { + if (st.kda_layers == 0) return; + const cur = st.cur[0]; + for (st.cur) |x| if (x != cur) return error.KdaLayersOutOfStep; + const h = st.heads; + if (keep < st.pending) { + var a: cuda.Args = .{}; + a.add(@as(c_int, @intCast(h))); + a.add(st.recAt(cur, 0)); + a.add(@as(usize, st.rec_layer / 4)); + for ([_]u64{ st.save_k, st.save_v, st.save_g, st.save_b }) |x| a.add(x); + a.add(@as(usize, state.window_rows * h * kda.dk)); + a.add(@as(usize, state.window_rows * h)); + a.add(@as(c_int, @intCast(keep))); + a.add(st.recAt(1 - cur, 0)); + try cuda.launch.launch(o.m.kda.replay_layers, .{ .grid = .{ .x = @intCast(st.kda_layers * h) }, .block = .{ .x = 1024 } }, o.s, &a); + } + for (st.cur) |*x| x.* = 1 - x.*; + try graph.glue.convShift(o.set(), o.s, st.conv, st.proj, keep, state.window_rows * st.stride, st.stride, st.kda_layers, st.convWidth()); +} diff --git a/zig/src/families/ling/block_mla.zig b/zig/src/families/ling/block_mla.zig new file mode 100644 index 0000000000..02a5e5beb7 --- /dev/null +++ b/zig/src/families/ling/block_mla.zig @@ -0,0 +1,63 @@ +//! Ling's MLA mixer on CUDA, absorbed: projections in their loaded formats, roped latent cache, head gate, `dense`. +const std = @import("std"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const block_kda = @import("block_kda.zig"); + +const glue = graph.glue; +const latent = graph.latent; +const Ops = graph.Ops; +const Config = @import("config.zig").Config; +const kw = state.cache_width; +const lw = 512; + +/// Room for a segment's rows: the sequence's cache and the scratch's attention partials (512-key chunks). +pub fn room(b: *const state.Buffers, st: *const state.State, rows: usize) !void { + if (st.pos + rows > st.capacity or latent.chunksFor(st.pos + rows) > b.nch) return error.CacheFull; +} + +/// The block on b.normed into b.y: q (lora pair or q_proj), latent and roped key into the cache, attention, gate. +pub fn block(o: Ops, b: *const state.Buffers, a: weights.Mla, ab: graph.Absorb, g: *const graph.Graph, mi: usize, segs: []const block_kda.Seg, rows: usize) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const H: usize = c.mla_heads; + const qk: usize = c.qkDim(); + const nope: usize = c.nope; + const dv: usize = c.v_dim; + if (c.kv_lora != lw or nope != 128 or dv != 128 or c.rope_dim != latent.rope_width) return error.LatentShape; + if (a.q_a) |qa| { + const ql: usize = c.q_lora; + try o.mm(qa, b.normed, D, rows, b.qa_in); + try glue.rmsnorm(o.set(), o.s, b.qa_in, ql, a.q_a_norm, b.qr, ql, null, c.eps, rows, ql); + try o.mm(a.q_b.?, b.qr, ql, rows, b.q); + } else try o.mm(a.q.?, b.normed, D, rows, b.q); + try o.mm(a.kv_a, b.normed, D, rows, b.kv); + try glue.rmsnorm(o.set(), o.s, b.kv, kw, a.kv_a_norm, b.lat, kw, null, c.eps, rows, lw); + if (a.gate) |gp| try o.mm(gp, b.normed, D, rows, b.g); + try latent.absorbQStrided(o.set(), o.s, b.q, qk, ab.wk, b.qa, kw, rows, H, nope, lw); + const scale = c.smScale(); + for (segs) |sg| try room(b, sg.st, sg.rows); // every segment before the first cache write + for (segs) |sg| { + const st = sg.st; + const r = sg.row0; + try latent.rope(o.set(), o.s, .{ .x = b.q + (r * H * qk + nope) * 2, .x_row = H * qk, .x_head = qk, .out = b.qa + (r * H * kw + lw) * 2, .o_row = H * kw, .o_head = kw }, sg.rows, H, st.pos_dev, g.freq); + try latent.rope(o.set(), o.s, .{ .x = b.kv + (r * kw + lw) * 2, .x_row = kw, .x_head = 0, .out = b.lat + (r * kw + lw) * 2, .o_row = kw, .o_head = 0 }, sg.rows, 1, st.pos_dev, g.freq); + try latent.latentWrite(o.set(), o.s, b.lat + r * kw * 2, kw, st.cacheAt(mi), st.pos_dev, sg.rows, kw); + var r0: usize = 0; + while (r0 < sg.rows) : (r0 += state.att_rows) { + const m = @min(state.att_rows, sg.rows - r0); + var at = st.pos_dev; + if (r0 > 0) { + try o.addI32(st.pos_dev, r0, b.at); + at = b.at; + } + const nch = latent.chunksFor(st.pos + r0 + m); + const row = r + r0; + try latent.attentionRoped(o.set(), o.s, b.qa + row * H * kw * 2, st.cacheAt(mi), at, b.partials, b.ol + row * H * lw * 2, b.dummy, m, H, lw, nch, scale); + } + } + try latent.expandV(o.set(), o.s, b.ol, ab.wv, b.vn, rows, H, dv, lw); + if (a.gate != null) try o.gate(b.vn, b.g, H, rows, H, dv); + try o.mm(a.o, b.vn, H * dv, rows, b.y); +} diff --git a/zig/src/families/ling/config.zig b/zig/src/families/ling/config.zig new file mode 100644 index 0000000000..8939a8f4a7 --- /dev/null +++ b/zig/src/families/ling/config.zig @@ -0,0 +1,526 @@ +//! Ling-3.1-flash's shape from config.json (bailing_hybrid), refusing what the engine does not run. +const std = @import("std"); + +/// Backbone plus MTP layers the per-layer tables hold (Ling-3.1-flash: 72 + 1). +pub const max_layers = 80; + +pub const Kind = enum { kda, mla }; + +/// MLA's rope pairs: rope_interleave true is exllamav3's RopeStyle.GPTJ (adjacent pairs), false NeoX (halves). +pub const RopeStyle = enum { gptj, neox }; + +/// A linear's format: bf16, FP8 128x128 blocks, ModelOpt NVFP4; EXL3 (#832) adds a variant and a weights_proj reader. +pub const Quant = enum { bf16, fp8_block128, nvfp4 }; + +/// Why parse refused the config, one sentence for the caller's log line (tests read it too). +pub const Why = struct { + buf: [320]u8 = undefined, + len: usize = 0, + + pub fn set(self: *Why, comptime fmt: []const u8, args: anytype) void { + var w: std.Io.Writer = .fixed(&self.buf); + w.print(fmt, args) catch {}; + self.len = w.end; + } + + pub fn text(self: *const Why) []const u8 { + return self.buf[0..self.len]; + } +}; + +pub const Config = struct { + hidden: u32, + layers: u32, // backbone layers; the MTP layers follow at [layers, layers + mtp) + vocab: u32, + eps: f32, + rope_theta: f64, + rope_dim: u32, // qk_rope_head_dim: the modeling's rotary embedding takes it whole (partial factor 1) + rope_style: RopeStyle, + group: u32, // layer_group_size: MLA on each group's last layer and on the tail past the last full group + dense_layers: u32, // first_k_dense_replace + dense_inter: u32, + kda_heads: u32, + kda_dim: u32, + conv: u32, + lower_bound: f32, // kda_safe_gate's kda_lower_bound + mla_heads: u32, + q_lora: u32, // 0: no q_lora_rank, one q_proj + kv_lora: u32, + nope: u32, + v_dim: u32, + head_gate: bool, // gated_attention_proj_granularity_type head_wise: g_proj gives a sigmoid gate per head + experts: u32, + topk: u32, + n_group: u32, + topk_group: u32, + routed_scale: f32, + norm_topk: bool, + expert_bias: bool, // moe_router_enable_expert_bias: gate.expert_bias, added for the pick only + moe_inter: u32, + shared_inter: u32, // moe_shared_expert_intermediate_size x num_shared_experts, 0 without a shared expert + expert_limit: [max_layers]f32, // expert_swiglu_limit_list by layer, 0 for none + shared_limit: [max_layers]f32, // share_expert_swiglu_limit_list by layer, 0 for none + mtp: u32, + tie: bool, + eos: [4]u32, + eos_n: u32, + bos: ?u32, + quant: Quant, + plain_text: [8192]u8, // modules_to_not_convert (FP8) or ignore / exclude_modules (ModelOpt), packed + plain_ends: [512]u16, // plain_ends[i] ends entry i + plain_n: u32, + + /// Layer i's mixer: the modeling's schedule for the backbone; the MTP layers are MLA (mtp_use_kda false). + pub fn kind(c: *const Config, i: u32) Kind { + if (i >= c.layers) return .mla; + const full = c.layers / c.group * c.group; + return if ((i + 1) % c.group == 0 or i >= full) .mla else .kda; + } + + /// Layer i has routed experts (the MTP layers always do); the first dense_layers run a dense MLP. + pub fn isMoe(c: *const Config, i: u32) bool { + return i >= c.dense_layers; + } + + pub fn countKind(c: *const Config, k: Kind) u32 { + var n: u32 = 0; + for (0..c.layers) |i| n += @intFromBool(c.kind(@intCast(i)) == k); + return n; + } + + /// MLA's query and key head width: nope then rope. + pub fn qkDim(c: *const Config) u32 { + return c.nope + c.rope_dim; + } + + /// MLA's softmax scale, qk_head_dim^-0.5 (no rope_scaling mscale: parse refuses rope_scaling). + pub fn smScale(c: *const Config) f32 { + return @floatCast(1.0 / @sqrt(@as(f64, @floatFromInt(c.qkDim())))); + } + + pub fn plain(c: *const Config, i: u32) []const u8 { + const lo: usize = if (i == 0) 0 else c.plain_ends[i - 1]; + return c.plain_text[lo..c.plain_ends[i]]; + } + + /// A linear module the exception list leaves in bf16: named whole, ended after a dot, or matched by a `*` glob. + pub fn unconverted(c: *const Config, module: []const u8) bool { + for (0..c.plain_n) |i| { + const p = c.plain(@intCast(i)); + if (std.mem.indexOfScalar(u8, p, '*') != null) { + if (glob(p, module)) return true; + continue; + } + if (std.mem.eql(u8, module, p)) return true; + if (module.len > p.len and std.mem.endsWith(u8, module, p) and module[module.len - p.len - 1] == '.') return true; + } + return false; + } + + /// The stored format of linear `module`: the checkpoint's quantization unless its exception list names it. + pub fn format(c: *const Config, module: []const u8) Quant { + return if (c.unconverted(module)) .bf16 else c.quant; + } + + /// A linear module carries quantization tensors beside its weight. + pub fn scaled(c: *const Config, module: []const u8) bool { + return c.format(module) != .bf16; + } +}; + +/// fnmatch with `*` only (ModelOpt's ignore patterns): `*` takes any run of characters, dots included. +pub fn glob(pattern: []const u8, name: []const u8) bool { + const star = std.mem.indexOfScalar(u8, pattern, '*') orelse return std.mem.eql(u8, pattern, name); + if (!std.mem.startsWith(u8, name, pattern[0..star])) return false; + const rest = pattern[star + 1 ..]; + var at = star; + while (at <= name.len) : (at += 1) if (glob(rest, name[at..])) return true; + return false; +} + +fn get(o: std.json.ObjectMap, key: []const u8, why: *Why) !std.json.Value { + return o.get(key) orelse { + why.set("config.json has no {s}", .{key}); + return error.BadConfig; + }; +} + +fn int(o: std.json.ObjectMap, key: []const u8, why: *Why) !u32 { + const v = try get(o, key, why); + if (v != .integer or v.integer < 0 or v.integer > std.math.maxInt(u32)) { + why.set("config.json's {s} is not a whole number", .{key}); + return error.BadConfig; + } + return @intCast(v.integer); +} + +fn number(v: std.json.Value) ?f64 { + return switch (v) { + .float => |f| f, + .integer => |i| @floatFromInt(i), + else => null, + }; +} + +fn float(o: std.json.ObjectMap, key: []const u8, why: *Why) !f64 { + return number(try get(o, key, why)) orelse { + why.set("config.json's {s} is not a number", .{key}); + return error.BadConfig; + }; +} + +fn flag(o: std.json.ObjectMap, key: []const u8, default: bool, why: *Why) !bool { + const v = o.get(key) orelse return default; + if (v != .bool) { + why.set("config.json's {s} is not true or false", .{key}); + return error.BadConfig; + } + return v.bool; +} + +fn text(o: std.json.ObjectMap, key: []const u8) ?[]const u8 { + const v = o.get(key) orelse return null; + return if (v == .string) v.string else null; +} + +/// A string key that, when present and not null, must read `want`. +fn must(o: std.json.ObjectMap, key: []const u8, want: []const u8, err: anyerror, why: *Why) !void { + const v = o.get(key) orelse return; + if (v == .null) return; + if (v == .string and std.mem.eql(u8, v.string, want)) return; + why.set("{s} must be \"{s}\": the engine runs no other", .{ key, want }); + return err; +} + +fn limits(o: std.json.ObjectMap, key: []const u8, layers: u32, out: *[max_layers]f32, why: *Why) !void { + out.* = @splat(0); + const v = o.get(key) orelse return; // absent: no clamp anywhere, as the reference reads it + if (v != .array or v.array.items.len != layers) { + why.set("{s} has {d} entries for {d} layers", .{ key, if (v == .array) v.array.items.len else 0, layers }); + return error.LimitList; + } + for (v.array.items, 0..) |e, i| { + const x = if (e == .null) 0 else number(e) orelse -1; + if (x < 0 or !std.math.isFinite(x)) { + why.set("{s}[{d}] is not a limit (0 for none, else positive)", .{ key, i }); + return error.LimitList; + } + out[i] = @floatCast(x); + } +} + +/// An exception list's names, packed into the config (modules left in bf16). +fn exceptions(v: ?std.json.Value, c: *Config, why: *Why) !void { + const list = v orelse return; + if (list != .array) return; + var at: usize = if (c.plain_n == 0) 0 else c.plain_ends[c.plain_n - 1]; + for (list.array.items) |m| { + if (m != .string or c.plain_n == c.plain_ends.len or at + m.string.len > c.plain_text.len) { + why.set("the quantization's exception list holds more than the loader keeps", .{}); + return error.Quant; + } + @memcpy(c.plain_text[at..][0..m.string.len], m.string); + at += m.string.len; + c.plain_ends[c.plain_n] = @intCast(at); + c.plain_n += 1; + } +} + +fn fp8Block(q: std.json.ObjectMap, c: *Config, why: *Why) !void { + const block = q.get("weight_block_size"); + const b128 = block != null and block.? == .array and block.?.array.items.len == 2 and + block.?.array.items[0] == .integer and block.?.array.items[0].integer == 128 and + block.?.array.items[1] == .integer and block.?.array.items[1].integer == 128; + if (!b128 or !std.mem.eql(u8, text(q, "fmt") orelse "e4m3", "e4m3") or + !std.mem.eql(u8, text(q, "activation_scheme") orelse "dynamic", "dynamic")) + { + why.set("quantization fp8: the loader reads FP8 e4m3 in 128x128 blocks with dynamic activations", .{}); + return error.Quant; + } + c.quant = .fp8_block128; + try exceptions(q.get("modules_to_not_convert"), c, why); +} + +fn isInt(v: std.json.Value, want: i64) bool { + return v == .integer and v.integer == want; +} + +fn groups16(why: *Why) error{Quant} { + why.set("ModelOpt NVFP4 with weights other than 4-bit in groups of 16", .{}); + return error.Quant; +} + +/// ModelOpt NVFP4 from config.json's quantization_config or hf_quant_config.json's "quantization" block. +fn modelopt(q: std.json.ObjectMap, c: *Config, why: *Why) !void { + const algo = text(q, "quant_algo") orelse ""; + if (!std.ascii.eqlIgnoreCase(algo, "NVFP4")) { + why.set("ModelOpt {s}: the loader reads ModelOpt NVFP4 (e2m1 in groups of 16)", .{if (algo.len > 0) algo else "without quant_algo"}); + return error.Quant; + } + if (q.get("group_size")) |v| if (!isInt(v, 16)) return groups16(why); + if (q.get("config_groups")) |g| if (g == .object) { + var it = g.object.iterator(); + while (it.next()) |e| { + const w = if (e.value_ptr.* == .object) e.value_ptr.object.get("weights") else null; + if (w == null or w.? != .object) return groups16(why); + const bits = w.?.object.get("num_bits") orelse return groups16(why); + const size = w.?.object.get("group_size") orelse return groups16(why); + if (!isInt(bits, 4) or !isInt(size, 16)) return groups16(why); + } + }; + c.quant = .nvfp4; + inline for (.{ "ignore", "exclude_modules", "modules_to_not_convert" }) |key| try exceptions(q.get(key), c, why); +} + +fn quantization(o: std.json.ObjectMap, c: *Config, why: *Why) !void { + c.quant = .bf16; + c.plain_n = 0; + const qv = o.get("quantization_config") orelse return; + if (qv == .null) return; + if (qv != .object) { + why.set("quantization_config is not an object", .{}); + return error.Quant; + } + const q = qv.object; + const method = text(q, "quant_method") orelse ""; + if (std.mem.eql(u8, method, "fp8")) return fp8Block(q, c, why); + if (std.ascii.eqlIgnoreCase(method, "modelopt")) return modelopt(q, c, why); + if (std.ascii.eqlIgnoreCase(method, "exl3")) { + why.set("EXL3: Ling's EXL3 weights come with ashhart's EXL3 module (tracker #832), not this loader yet", .{}); + return error.Exl3; + } + why.set("quantization {s}: the loader reads bf16, FP8 128x128 blocks or ModelOpt NVFP4", .{if (method.len > 0) method else "without quant_method"}); + return error.Quant; +} + +fn tokens(o: std.json.ObjectMap, c: *Config, why: *Why) !void { + c.eos_n = 0; + const v = try get(o, "eos_token_id", why); + const items: []const std.json.Value = if (v == .array) v.array.items else &.{v}; + if (items.len == 0 or items.len > c.eos.len) { + why.set("eos_token_id holds {d} ids, the engine stops on 1 to {d}", .{ items.len, c.eos.len }); + return error.BadConfig; + } + for (items) |e| { + if (e != .integer or e.integer < 0 or e.integer >= c.vocab) { + why.set("eos_token_id is not a token of the {d}-token vocabulary", .{c.vocab}); + return error.BadConfig; + } + c.eos[c.eos_n] = @intCast(e.integer); + c.eos_n += 1; + } + c.bos = null; + if (o.get("bos_token_id")) |b| if (b == .integer and b.integer >= 0 and b.integer < c.vocab) { + c.bos = @intCast(b.integer); + }; +} + +/// config.json as the reference (configuration_bailing_moe_v3.py, exllamav3's bailing_hybrid) reads it. +pub fn parse(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { + return parseWith(gpa, json, null, why); +} + +/// parse() with ModelOpt's hf_quant_config.json beside it: its NVFP4 block when config.json has no quantization_config. +pub fn parseWith(gpa: std.mem.Allocator, json: []const u8, hf_quant: ?[]const u8, why: *Why) !Config { + var c = try parseConfig(gpa, json, why); + const extra = hf_quant orelse return c; + if (c.quant != .bf16 or c.plain_n != 0) return c; + const parsed = std.json.parseFromSlice(std.json.Value, gpa, extra, .{}) catch { + why.set("hf_quant_config.json is not JSON", .{}); + return error.Quant; + }; + defer parsed.deinit(); + const q = if (parsed.value == .object) parsed.value.object.get("quantization") else null; + if (q == null or q.? != .object) { + why.set("hf_quant_config.json has no quantization block", .{}); + return error.Quant; + } + try modelopt(q.?.object, &c, why); + return c; +} + +fn parseConfig(gpa: std.mem.Allocator, json: []const u8, why: *Why) !Config { + const parsed = std.json.parseFromSlice(std.json.Value, gpa, json, .{}) catch { + why.set("config.json is not JSON", .{}); + return error.BadConfig; + }; + defer parsed.deinit(); + if (parsed.value != .object) return error.BadConfig; + const o = parsed.value.object; + const arch = o.get("architectures"); + const ling = arch != null and arch.? == .array and arch.?.array.items.len == 1 and arch.?.array.items[0] == .string and + std.mem.eql(u8, arch.?.array.items[0].string, "BailingMoeV3ForCausalLM"); + if (!ling or !std.mem.eql(u8, text(o, "model_type") orelse "", "bailing_hybrid")) { + why.set("not a Ling-3.1 checkpoint: architectures BailingMoeV3ForCausalLM, model_type bailing_hybrid", .{}); + return error.NotLing; + } + var c = std.mem.zeroes(Config); + c.hidden = try int(o, "hidden_size", why); + c.layers = try int(o, "num_hidden_layers", why); + c.vocab = try int(o, "vocab_size", why); + c.eps = @floatCast(try float(o, "rms_norm_eps", why)); + c.group = try int(o, "layer_group_size", why); + c.dense_layers = try int(o, "first_k_dense_replace", why); + c.dense_inter = try int(o, "intermediate_size", why); + c.mtp = if (o.get("num_nextn_predict_layers")) |_| try int(o, "num_nextn_predict_layers", why) else 0; + if (c.layers == 0 or c.group == 0 or c.dense_layers > c.layers or c.hidden == 0 or c.vocab == 0) { + why.set("{d} layers in groups of {d}, {d} dense: not a layer schedule", .{ c.layers, c.group, c.dense_layers }); + return error.BadConfig; + } + if (c.layers + c.mtp > max_layers) { + why.set("{d} layers and {d} MTP layers: the loader holds {d}", .{ c.layers, c.mtp, max_layers }); + return error.TooManyLayers; + } + if (c.mtp > 1 or try flag(o, "mtp_use_kda", false, why)) { + why.set("MTP: the engine runs one MLA MTP layer", .{}); + return error.Mtp; + } + try must(o, "hidden_act", "silu", error.Activation, why); + if (try flag(o, "use_qkv_bias", false, why) or try flag(o, "use_bias", false, why)) { + why.set("use_qkv_bias/use_bias: the engine's projections have no bias", .{}); + return error.Bias; + } + // KDA: one head geometry for q, k and v, full-rank f and g, the safe gate's lower bound + if (!try flag(o, "no_kda_lora", false, why) or try flag(o, "use_kda_lora", false, why)) { + why.set("no_kda_lora must be true: the engine runs full-rank f_proj and g_proj", .{}); + return error.KdaLora; + } + c.kda_heads = try int(o, "num_attention_heads", why); + c.kda_dim = try int(o, "head_dim", why); + c.conv = try int(o, "short_conv_kernel_size", why); + if (!try flag(o, "kda_safe_gate", false, why) or o.get("kda_lower_bound") == null) { + why.set("kda_safe_gate with a kda_lower_bound: the engine's KDA gate is the safe one", .{}); + return error.KdaGate; + } + c.lower_bound = @floatCast(try float(o, "kda_lower_bound", why)); + // MLA + c.mla_heads = c.kda_heads; + const ql = o.get("q_lora_rank") orelse std.json.Value.null; + c.q_lora = if (ql == .null) 0 else try int(o, "q_lora_rank", why); + c.kv_lora = try int(o, "kv_lora_rank", why); + c.nope = try int(o, "qk_nope_head_dim", why); + c.rope_dim = try int(o, "qk_rope_head_dim", why); + c.v_dim = try int(o, "v_head_dim", why); + c.rope_theta = try float(o, "rope_theta", why); + c.rope_style = if (try flag(o, "rope_interleave", true, why)) .gptj else .neox; + if (o.get("rope_scaling")) |r| if (r != .null) { + why.set("rope_scaling: the engine's MLA has plain rope only", .{}); + return error.RopeScaling; + }; + const gate = o.get("gated_attention_proj_granularity_type") orelse std.json.Value.null; + c.head_gate = gate == .string and std.mem.eql(u8, gate.string, "head_wise"); + if (gate != .null and !c.head_gate) { + why.set("gated_attention_proj_granularity_type {s}: the engine gates MLA per head or not at all", .{if (gate == .string) gate.string else "?"}); + return error.GateGranularity; + } + // MoE: DeepSeek-V3's sigmoid router with group-limited top-k and a bias for the pick + try must(o, "score_function", "sigmoid", error.ScoreFunction, why); + try must(o, "scoring_func", "sigmoid", error.ScoreFunction, why); + try must(o, "topk_method", "noaux_tc", error.Routing, why); + c.experts = try int(o, "num_experts", why); + c.topk = try int(o, "num_experts_per_tok", why); + c.n_group = if (o.get("n_group")) |_| try int(o, "n_group", why) else 1; + c.topk_group = if (o.get("topk_group")) |_| try int(o, "topk_group", why) else 1; + c.routed_scale = @floatCast(if (o.get("routed_scaling_factor")) |_| try float(o, "routed_scaling_factor", why) else 1.0); + c.norm_topk = try flag(o, "norm_topk_prob", true, why); + if (!c.norm_topk) { + why.set("norm_topk_prob must be true: the routed weights are renormalized over the top-k", .{}); + return error.Routing; + } + c.expert_bias = try flag(o, "moe_router_enable_expert_bias", true, why); + // glue.select's bounds: at most 512 experts, top-k under 32, up to 32 equal groups that hold the top-k + const g_ok = c.n_group >= 1 and c.n_group <= 32 and c.experts % c.n_group == 0 and c.topk_group >= 1 and + c.topk_group <= c.n_group and c.topk <= c.topk_group * (c.experts / c.n_group); + if (c.experts > 512 or c.topk == 0 or c.topk >= 32 or c.topk > c.experts or !g_ok) { + why.set("routing {d} of {d} experts in {d} groups keeping {d}: past glue.select's bounds", .{ c.topk, c.experts, c.n_group, c.topk_group }); + return error.Routing; + } + // HF's rule: a group scores its top two (topk(2) fails on a group of one); one pick keeps its weight unnormalized + if (c.experts / c.n_group < 2) { + why.set("{d} groups of {d} experts: HF scores a group by its best two, so groups hold two or more", .{ c.n_group, c.experts / c.n_group }); + return error.Routing; + } + if (c.topk == 1) { + why.set("num_experts_per_tok 1: HF leaves one pick's weight unnormalized, the kernel normalizes it", .{}); + return error.Routing; + } + c.moe_inter = try int(o, "moe_intermediate_size", why); + const shared_n = if (o.get("num_shared_experts")) |v| (if (v == .null) 0 else try int(o, "num_shared_experts", why)) else 0; + const shared_w = if (o.get("moe_shared_expert_intermediate_size")) |_| try int(o, "moe_shared_expert_intermediate_size", why) else c.moe_inter; + c.shared_inter = shared_n * shared_w; + try limits(o, "expert_swiglu_limit_list", c.layers, &c.expert_limit, why); + try limits(o, "share_expert_swiglu_limit_list", c.layers, &c.shared_limit, why); + for (0..c.dense_layers) |i| if (c.expert_limit[i] != 0 or c.shared_limit[i] != 0) { + why.set("layer {d} is dense with a SwiGLU limit: neither list says which a dense MLP takes", .{i}); + return error.LimitList; + }; + c.tie = try flag(o, "tie_word_embeddings", false, why); + try tokens(o, &c, why); + try geometry(o, &c, why); + try quantization(o, &c, why); + return c; +} + +/// The widths the kernels are built for (cuda.zig's tests hold them to glm_kda, glm_glue and the latent cache). +pub const fixed = .{ .kda_dim = 128, .conv = 4, .kv_lora = 512, .nope = 128, .v_dim = 128, .rope = 64, .block = 128 }; +/// glue.rmsnorm's instances: the hidden width and q_lora_rank take one of these. +pub const norm_widths = [_]u32{ 512, 1536, 4096 }; + +fn normWidth(w: u32) bool { + return std.mem.indexOfScalar(u32, &norm_widths, w) != null; +} + +/// Every shape the kernels take fixed, refused at load before anything is allocated or launched. +fn geometry(o: std.json.ObjectMap, c: *const Config, why: *Why) !void { + // KDA: glm_kda.cu's 128-wide heads and 4 taps (a 3-row conv window); one key and value head per query head + if (c.kda_dim != fixed.kda_dim) { + why.set("head_dim {d}: the KDA kernels take {d}-wide heads", .{ c.kda_dim, fixed.kda_dim }); + return error.KdaGeometry; + } + if (c.conv != fixed.conv) { + why.set("short_conv_kernel_size {d}: the KDA kernels take {d} taps", .{ c.conv, fixed.conv }); + return error.KdaGeometry; + } + const kv_heads = if (o.get("num_key_value_heads")) |_| try int(o, "num_key_value_heads", why) else c.kda_heads; + if (c.kda_heads == 0 or kv_heads != c.kda_heads) { + why.set("{d} attention heads, {d} key/value heads: the engine runs one key/value head per query head", .{ c.kda_heads, kv_heads }); + return error.KdaGeometry; + } + // MLA: the 512-wide latent, 128-wide nope and value heads, 64 GPT-J roped columns + const mla = [_]struct { []const u8, u32, u32 }{ + .{ "kv_lora_rank", c.kv_lora, fixed.kv_lora }, + .{ "qk_nope_head_dim", c.nope, fixed.nope }, + .{ "v_head_dim", c.v_dim, fixed.v_dim }, + .{ "qk_rope_head_dim", c.rope_dim, fixed.rope }, + }; + for (mla) |m| if (m[1] != m[2]) { + why.set("{s} {d}: the MLA kernels take {d}", .{ m[0], m[1], m[2] }); + return error.MlaGeometry; + }; + if (c.rope_style != .gptj) { + why.set("rope_interleave false: the MLA rope kernel pairs adjacent columns (GPT-J) only", .{}); + return error.MlaGeometry; + } + if (c.q_lora != 0 and !normWidth(c.q_lora)) { + why.set("q_lora_rank {d}: q_a_layernorm runs at 512, 1536 or 4096", .{c.q_lora}); + return error.MlaGeometry; + } + // widths: the norms' instances, and every linear's inputs in whole 128-column blocks (FP8's, the lanes' 64) + if (!normWidth(c.hidden)) { + why.set("hidden_size {d}: the norms and the router run at 512, 1536 or 4096", .{c.hidden}); + return error.Width; + } + const widths = [_]struct { []const u8, u32, bool }{ + .{ "intermediate_size", c.dense_inter, c.dense_layers > 0 }, + .{ "moe_intermediate_size", c.moe_inter, true }, + .{ "the shared expert's width", c.shared_inter, c.shared_inter > 0 }, + }; + for (widths) |x| if (x[2] and (x[1] == 0 or x[1] % fixed.block != 0)) { + why.set("{s} {d}: the projections and experts take whole blocks of {d}", .{ x[0], x[1], fixed.block }); + return error.Width; + }; + if (c.tie) { + why.set("tie_word_embeddings: the engine reads its own lm_head", .{}); + return error.TiedHead; + } +} diff --git a/zig/src/families/ling/config_test.zig b/zig/src/families/ling/config_test.zig new file mode 100644 index 0000000000..880d8ddd03 --- /dev/null +++ b/zig/src/families/ling/config_test.zig @@ -0,0 +1,314 @@ +//! Ling-3.1-flash's own config.json (fixtures/config.json) parsed, and each refusal on a one-edit copy. +const std = @import("std"); +const config = @import("config.zig"); + +const real = @embedFile("fixtures/config.json"); +const t = std.testing; + +fn load() !config.Config { + var why: config.Why = .{}; + return config.parse(t.allocator, real, &why); +} + +/// The fixture with `from` (which must occur exactly once) replaced by `to`. +fn edit(from: []const u8, to: []const u8) ![]u8 { + if (std.mem.count(u8, real, from) != 1) return error.EditAnchor; + return std.mem.replaceOwned(u8, t.allocator, real, from, to); +} + +fn refused(from: []const u8, to: []const u8, want: anyerror, says: []const u8) !void { + const json = try edit(from, to); + defer t.allocator.free(json); + var why: config.Why = .{}; + try t.expectError(want, config.parse(t.allocator, json, &why)); + try t.expect(std.mem.indexOf(u8, why.text(), says) != null); +} + +test "the checkpoint's config.json parses to Ling-3.1-flash's shape" { + const c = try load(); + try t.expectEqual(@as(u32, 4096), c.hidden); + try t.expectEqual(@as(u32, 72), c.layers); + try t.expectEqual(@as(u32, 157184), c.vocab); + try t.expectEqual(@as(f32, 1e-6), c.eps); + try t.expectEqual(@as(f64, 6000000), c.rope_theta); + try t.expectEqual(@as(u32, 64), c.rope_dim); + try t.expectEqual(config.RopeStyle.gptj, c.rope_style); + try t.expectEqual(@as(u32, 8), c.group); + try t.expectEqual(@as(u32, 4), c.dense_layers); + try t.expectEqual(@as(u32, 10240), c.dense_inter); + try t.expectEqual(@as(u32, 64), c.kda_heads); + try t.expectEqual(@as(u32, 128), c.kda_dim); + try t.expectEqual(@as(u32, 4), c.conv); + try t.expectEqual(@as(f32, -5.0), c.lower_bound); + try t.expectEqual(@as(u32, 64), c.mla_heads); + try t.expectEqual(@as(u32, 1536), c.q_lora); + try t.expectEqual(@as(u32, 512), c.kv_lora); + try t.expectEqual(@as(u32, 128), c.nope); + try t.expectEqual(@as(u32, 128), c.v_dim); + try t.expectEqual(@as(u32, 192), c.qkDim()); + try t.expectEqual(@as(f32, @floatCast(1.0 / @sqrt(192.0))), c.smScale()); + try t.expect(c.head_gate); + try t.expectEqual(@as(u32, 512), c.experts); + try t.expectEqual(@as(u32, 8), c.topk); + try t.expectEqual(@as(u32, 8), c.n_group); + try t.expectEqual(@as(u32, 4), c.topk_group); + try t.expectEqual(@as(f32, 2.5), c.routed_scale); + try t.expect(c.norm_topk and c.expert_bias); + try t.expectEqual(@as(u32, 1280), c.moe_inter); + try t.expectEqual(@as(u32, 1280), c.shared_inter); + try t.expectEqual(@as(u32, 1), c.mtp); + try t.expect(!c.tie); + try t.expectEqual(@as(u32, 1), c.eos_n); + try t.expectEqual(@as(u32, 156895), c.eos[0]); + try t.expect(c.bos == null); // config.json names no bos_token_id; the tokenizer adds none either + try t.expectEqual(config.Quant.fp8_block128, c.quant); + try t.expectEqual(@as(u32, 14), c.plain_n); +} + +test "the schedule: MLA exactly on 7, 15, ..., 71 (9 layers), KDA on the other 63, dense 0-3, the MTP layer MLA + MoE" { + const c = try load(); + var mla: [9]u32 = undefined; + var n: usize = 0; + for (0..c.layers) |i| if (c.kind(@intCast(i)) == .mla) { + mla[n] = @intCast(i); + n += 1; + }; + try t.expectEqual(@as(usize, 9), n); + try t.expectEqualSlices(u32, &.{ 7, 15, 23, 31, 39, 47, 55, 63, 71 }, &mla); + try t.expectEqual(@as(u32, 9), c.countKind(.mla)); + try t.expectEqual(@as(u32, 63), c.countKind(.kda)); + for (0..c.layers) |i| try t.expectEqual(i >= 4, c.isMoe(@intCast(i))); + try t.expectEqual(config.Kind.mla, c.kind(72)); + try t.expect(c.isMoe(72)); +} + +test "a tail past the last full group is MLA, as the modeling's schedule says" { + var c = try load(); + c.layers = 10; + c.group = 4; + for ([_]config.Kind{ .kda, .kda, .kda, .mla, .kda, .kda, .kda, .mla, .mla, .mla }, 0..) |k, i| try t.expectEqual(k, c.kind(@intCast(i))); +} + +test "SwiGLU limits from the file: 10 on routed experts of layers 65-71 and shared ones of 64-71, else 0" { + const c = try load(); + for (0..c.layers) |i| { + try t.expectEqual(@as(f32, if (i >= 65) 10 else 0), c.expert_limit[i]); + try t.expectEqual(@as(f32, if (i >= 64) 10 else 0), c.shared_limit[i]); + } +} + +test "modules_to_not_convert matches whole dotted tails: mlp.gate is the router, not gate_proj" { + const c = try load(); + try t.expect(c.unconverted("model.layers.4.mlp.gate")); + try t.expect(!c.unconverted("model.layers.4.mlp.gate_proj")); + try t.expect(!c.unconverted("model.layers.4.mlp.shared_experts.gate_proj")); + try t.expect(c.unconverted("model.layers.0.attention.b_proj")); + try t.expect(!c.unconverted("model.layers.7.attention.q_b_proj")); + try t.expect(c.unconverted("model.layers.7.attention.g_proj")); + try t.expect(!c.unconverted("model.layers.72.attention.g_proj")); + try t.expect(c.unconverted("model.layers.72.eh_proj")); + try t.expect(c.unconverted("lm_head") and !c.scaled("lm_head")); + try t.expect(c.scaled("model.layers.0.attention.q_proj")); +} + +test "refusals: one edit each, with the reason" { + try refused("\"BailingMoeV3ForCausalLM\"", "\"BailingMoeV2ForCausalLM\"", error.NotLing, "not a Ling-3.1"); + try refused("\"score_function\": \"sigmoid\"", "\"score_function\": \"softmax\"", error.ScoreFunction, "score_function"); + try refused("\"scoring_func\": \"sigmoid\"", "\"scoring_func\": \"softmax\"", error.ScoreFunction, "scoring_func"); + try refused("\"no_kda_lora\": true", "\"no_kda_lora\": false", error.KdaLora, "no_kda_lora"); + try refused("\"use_kda_lora\": false", "\"use_kda_lora\": true", error.KdaLora, "no_kda_lora"); + try refused("\"kda_safe_gate\": true", "\"kda_safe_gate\": false", error.KdaGate, "kda_safe_gate"); + try refused("\"head_wise\"", "\"element_wise\"", error.GateGranularity, "element_wise"); + try refused("\"rope_scaling\": null", "\"rope_scaling\": {\"type\": \"yarn\", \"factor\": 4}", error.RopeScaling, "rope_scaling"); + try refused("\"n_group\": 8", "\"n_group\": 7", error.Routing, "groups"); + try refused("\"topk_group\": 4", "\"topk_group\": 9", error.Routing, "groups"); + try refused("\"num_experts\": 512", "\"num_experts\": 1024", error.Routing, "glue.select"); + try refused("\"norm_topk_prob\": true", "\"norm_topk_prob\": false", error.Routing, "norm_topk_prob"); + try refused("\"expert_swiglu_limit_list\": [\n 0,\n", "\"expert_swiglu_limit_list\": [\n", error.LimitList, "71 entries for 72"); + try refused("\"share_expert_swiglu_limit_list\": [\n 0,\n", "\"share_expert_swiglu_limit_list\": [\n 0,\n 0,\n", error.LimitList, "73 entries"); + try refused("\"share_expert_swiglu_limit_list\": [\n 0,\n", "\"share_expert_swiglu_limit_list\": [\n -1,\n", error.LimitList, "[0]"); + try refused("\"expert_swiglu_limit_list\": [\n 0,\n", "\"expert_swiglu_limit_list\": [\n 7,\n", error.LimitList, "layer 0 is dense"); + try refused("\"quant_method\": \"fp8\"", "\"quant_method\": \"gptq\"", error.Quant, "gptq"); + try refused("\"num_nextn_predict_layers\": 1", "\"num_nextn_predict_layers\": 2", error.Mtp, "MTP"); + try refused("\"mtp_use_kda\": false", "\"mtp_use_kda\": true", error.Mtp, "MTP"); + try refused("\"hidden_act\": \"silu\"", "\"hidden_act\": \"gelu\"", error.Activation, "hidden_act"); + try refused("\"use_qkv_bias\": false", "\"use_qkv_bias\": true", error.Bias, "bias"); + try refused("\"eos_token_id\": 156895", "\"eos_token_id\": 157184", error.BadConfig, "eos_token_id"); + try refused("\"num_hidden_layers\": 72", "\"num_hidden_layers\": 73", error.LimitList, "72 entries for 73"); + try refused("\"hidden_size\": 4096", "\"hidden_sz\": 4096", error.BadConfig, "hidden_size"); +} + +test "optional keys: no quantization reads bf16, a null q_lora_rank one q_proj" { + const b = try edit("\"q_lora_rank\": 1536", "\"q_lora_rank\": null"); + defer t.allocator.free(b); + const at = std.mem.indexOf(u8, b, "\"quantization_config\"").?; + const end = std.mem.indexOf(u8, b[at..], " },\n").? + at + 5; + const c_json = try std.mem.concat(t.allocator, u8, &.{ b[0..at], b[end..] }); + defer t.allocator.free(c_json); + var why: config.Why = .{}; + const c = try config.parse(t.allocator, c_json, &why); + try t.expectEqual(config.Quant.bf16, c.quant); + try t.expectEqual(@as(u32, 0), c.q_lora); + try t.expectEqual(config.RopeStyle.gptj, c.rope_style); + try t.expect(!c.scaled("model.layers.0.attention.q_proj")); +} + +/// The fixture without its quantization_config block. +fn plainJson() ![]u8 { + const at = std.mem.indexOf(u8, real, "\"quantization_config\"").?; + const end = std.mem.indexOf(u8, real[at..], " },\n").? + at + 5; + return std.mem.concat(t.allocator, u8, &.{ real[0..at], real[end..] }); +} + +test "ModelOpt NVFP4 from quantization_config: ignore globs and names pick the bf16 modules, the rest is nvfp4" { + const json = try edit("\"quant_method\": \"fp8\",", "\"quant_method\": \"modelopt\", \"quant_algo\": \"NVFP4\", \"ignore\": [\"*.attention.b_proj\", \"model.layers.7.mlp*\"],"); + defer t.allocator.free(json); + var why: config.Why = .{}; + const c = try config.parse(t.allocator, json, &why); + try t.expectEqual(config.Quant.nvfp4, c.quant); + try t.expectEqual(config.Quant.nvfp4, c.format("model.layers.0.attention.q_proj")); + try t.expectEqual(config.Quant.bf16, c.format("model.layers.30.attention.b_proj")); + try t.expectEqual(config.Quant.bf16, c.format("model.layers.7.mlp.experts.3.up_proj")); + try t.expectEqual(config.Quant.nvfp4, c.format("model.layers.17.mlp.experts.3.up_proj")); + try t.expectEqual(config.Quant.bf16, c.format("lm_head")); // modules_to_not_convert stays an exception list too +} + +test "hf_quant_config.json's NVFP4 block when config.json has none; groups other than 16 and other algorithms refused" { + const json = try plainJson(); + defer t.allocator.free(json); + var why: config.Why = .{}; + const hf = "{\"producer\": {\"name\": \"modelopt\"}, \"quantization\": {\"quant_algo\": \"NVFP4\", \"group_size\": 16, \"exclude_modules\": [\"lm_head\", \"model.layers.*.mlp.gate\"]}}"; + const c = try config.parseWith(t.allocator, json, hf, &why); + try t.expectEqual(config.Quant.nvfp4, c.quant); + try t.expectEqual(config.Quant.bf16, c.format("model.layers.40.mlp.gate")); + try t.expectEqual(config.Quant.nvfp4, c.format("model.layers.40.mlp.gate_proj")); + const g32 = "{\"quantization\": {\"quant_algo\": \"NVFP4\", \"group_size\": 32}}"; + try t.expectError(error.Quant, config.parseWith(t.allocator, json, g32, &why)); + try t.expect(std.mem.indexOf(u8, why.text(), "groups of 16") != null); + const fp8 = "{\"quantization\": {\"quant_algo\": \"FP8\"}}"; + try t.expectError(error.Quant, config.parseWith(t.allocator, json, fp8, &why)); + try t.expect(std.mem.indexOf(u8, why.text(), "ModelOpt FP8") != null); + try t.expectEqual(config.Quant.bf16, (try config.parseWith(t.allocator, json, null, &why)).quant); +} + +test "EXL3 is refused naming #832; a glob matches whole names only" { + try refused("\"quant_method\": \"fp8\"", "\"quant_method\": \"exl3\"", error.Exl3, "#832"); + try t.expect(config.glob("*.mlp.gate", "model.layers.4.mlp.gate")); + try t.expect(!config.glob("*.mlp.gate", "model.layers.4.mlp.gate_proj")); + try t.expect(config.glob("model.layers.*.attention*", "model.layers.12.attention.q_proj")); + try t.expect(!config.glob("model.layers.1*", "model.layers.0.mlp")); + try t.expect(config.glob("a*b*c", "a.b.x.c") and !config.glob("a*b*c", "a.c")); +} + +const synth = @import("synth.zig"); + +/// The tiny config (synth.zig, FP8) with `key` set to `value` (JSON text) and `also`'s pair, parsed. +fn tiny(key: []const u8, value: []const u8, also: ?[2][]const u8, why: *config.Why) !config.Config { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + var root = try std.json.parseFromSliceLeaky(std.json.Value, a, try synth.configJson(a, .fp8_block128, 4), .{}); + try root.object.put(a, key, try std.json.parseFromSliceLeaky(std.json.Value, a, value, .{})); + if (also) |kv| try root.object.put(a, kv[0], try std.json.parseFromSliceLeaky(std.json.Value, a, kv[1], .{})); + return config.parse(t.allocator, try std.json.Stringify.valueAlloc(a, root, .{}), why); +} + +const Case = struct { key: []const u8, value: []const u8, err: anyerror, says: []const u8, also: ?[2][]const u8 = null }; + +test "geometry the kernels do not take is refused at load, one value at a time on the tiny config, with the reason" { + const cases = [_]Case{ + .{ .key = "head_dim", .value = "64", .err = error.KdaGeometry, .says = "head_dim 64" }, + .{ .key = "head_dim", .value = "256", .err = error.KdaGeometry, .says = "128-wide heads" }, + .{ .key = "short_conv_kernel_size", .value = "3", .err = error.KdaGeometry, .says = "short_conv_kernel_size 3" }, + .{ .key = "short_conv_kernel_size", .value = "5", .err = error.KdaGeometry, .says = "4 taps" }, + .{ .key = "num_attention_heads", .value = "0", .err = error.KdaGeometry, .says = "0 attention heads" }, + .{ .key = "num_attention_heads", .value = "0", .err = error.KdaGeometry, .says = "0 attention heads, 0", .also = .{ "num_key_value_heads", "0" } }, + .{ .key = "num_key_value_heads", .value = "2", .err = error.KdaGeometry, .says = "2 key/value heads" }, + .{ .key = "kv_lora_rank", .value = "256", .err = error.MlaGeometry, .says = "kv_lora_rank 256" }, + .{ .key = "kv_lora_rank", .value = "1024", .err = error.MlaGeometry, .says = "take 512" }, + .{ .key = "qk_nope_head_dim", .value = "64", .err = error.MlaGeometry, .says = "qk_nope_head_dim 64" }, + .{ .key = "qk_nope_head_dim", .value = "192", .err = error.MlaGeometry, .says = "qk_nope_head_dim 192" }, + .{ .key = "v_head_dim", .value = "256", .err = error.MlaGeometry, .says = "v_head_dim 256" }, + .{ .key = "v_head_dim", .value = "64", .err = error.MlaGeometry, .says = "v_head_dim 64" }, + .{ .key = "qk_rope_head_dim", .value = "32", .err = error.MlaGeometry, .says = "qk_rope_head_dim 32" }, + .{ .key = "qk_rope_head_dim", .value = "128", .err = error.MlaGeometry, .says = "qk_rope_head_dim 128" }, + .{ .key = "rope_interleave", .value = "false", .err = error.MlaGeometry, .says = "GPT-J" }, + .{ .key = "q_lora_rank", .value = "768", .err = error.MlaGeometry, .says = "q_lora_rank 768" }, + .{ .key = "q_lora_rank", .value = "1024", .err = error.MlaGeometry, .says = "q_lora_rank 1024" }, + .{ .key = "hidden_size", .value = "1024", .err = error.Width, .says = "hidden_size 1024" }, + .{ .key = "hidden_size", .value = "640", .err = error.Width, .says = "hidden_size 640" }, + .{ .key = "intermediate_size", .value = "200", .err = error.Width, .says = "intermediate_size 200" }, + .{ .key = "intermediate_size", .value = "0", .err = error.Width, .says = "intermediate_size 0" }, + .{ .key = "moe_intermediate_size", .value = "96", .err = error.Width, .says = "moe_intermediate_size 96" }, + .{ .key = "moe_intermediate_size", .value = "192", .err = error.Width, .says = "moe_intermediate_size 192" }, + .{ .key = "moe_shared_expert_intermediate_size", .value = "64", .err = error.Width, .says = "shared expert's width 64" }, + .{ .key = "tie_word_embeddings", .value = "true", .err = error.TiedHead, .says = "lm_head" }, + }; + for (cases) |x| { + var why: config.Why = .{}; + const got = tiny(x.key, x.value, x.also, &why); + if (got) |_| { + std.debug.print("{s} {s}: accepted\n", .{ x.key, x.value }); + return error.TestUnexpectedResult; + } else |e| if (e != x.err or std.mem.indexOf(u8, why.text(), x.says) == null) { + std.debug.print("{s} {s}: {t}, \"{s}\"\n", .{ x.key, x.value, e, why.text() }); + return error.TestUnexpectedResult; + } + } +} + +test "the shapes the kernels take are accepted: the tiny config as written and with these values, the real one" { + var why: config.Why = .{}; + const ok = [_][2][]const u8{ + .{ "head_dim", "128" }, .{ "short_conv_kernel_size", "4" }, .{ "num_key_value_heads", "4" }, + .{ "q_lora_rank", "1536" }, .{ "q_lora_rank", "null" }, .{ "hidden_size", "4096" }, + .{ "hidden_size", "1536" }, .{ "moe_intermediate_size", "1280" }, .{ "intermediate_size", "10240" }, + .{ "rope_interleave", "true" }, + }; + for (ok) |kv| _ = tiny(kv[0], kv[1], null, &why) catch |e| { + std.debug.print("{s} {s}: {t}, \"{s}\"\n", .{ kv[0], kv[1], e, why.text() }); + return e; + }; + const c = try load(); + try t.expect(c.kda_dim == config.fixed.kda_dim and c.conv == config.fixed.conv and c.kv_lora == config.fixed.kv_lora); +} + +test "routing groups: groups of one and a single pick refused (HF's rule differs there); one group, all groups kept, groups of two accepted" { + const refusals = [_]Case{ + .{ .key = "n_group", .value = "16", .err = error.Routing, .says = "16 groups of 1 experts", .also = .{ "topk_group", "4" } }, + .{ .key = "n_group", .value = "16", .err = error.Routing, .says = "best two", .also = .{ "topk_group", "16" } }, + .{ .key = "num_experts_per_tok", .value = "1", .err = error.Routing, .says = "num_experts_per_tok 1" }, + .{ .key = "n_group", .value = "1", .err = error.Routing, .says = "past glue.select's bounds" }, // keeping 2 of 1 + }; + for (refusals) |x| { + var why: config.Why = .{}; + const got = tiny(x.key, x.value, x.also, &why); + if (got) |_| { + std.debug.print("{s} {s}: accepted\n", .{ x.key, x.value }); + return error.TestUnexpectedResult; + } else |e| if (e != x.err or std.mem.indexOf(u8, why.text(), x.says) == null) { + std.debug.print("{s} {s}: {t}, \"{s}\"\n", .{ x.key, x.value, e, why.text() }); + return error.TestUnexpectedResult; + } + } + const ok = [_][4][]const u8{ + .{ "n_group", "1", "topk_group", "1" }, // one group: glue.select skips the group pass, HF keeps every expert + .{ "n_group", "4", "topk_group", "4" }, // every group kept + .{ "n_group", "8", "topk_group", "2" }, // groups of two + .{ "n_group", "8", "topk_group", "8" }, + }; + for (ok) |kv| { + var why: config.Why = .{}; + const c = tiny(kv[0], kv[1], .{ kv[2], kv[3] }, &why) catch |e| { + std.debug.print("{s} {s}, {s} {s}: {t}, \"{s}\"\n", .{ kv[0], kv[1], kv[2], kv[3], e, why.text() }); + return e; + }; + try t.expect(c.experts / c.n_group >= 2 and c.topk_group <= c.n_group); + } +} + +test "fixtures/config.json is inclusionAI/Ling-3.1-flash's at a6943f4b, byte for byte (oracle/ling_hf.py pins the same)" { + var h: [32]u8 = undefined; + std.crypto.hash.sha2.Sha256.hash(real, &h, .{}); + try t.expectEqualStrings("6f3773e52a12509e772e698aacacaa156d2402ebb8250812a1da52a03d92f084", &std.fmt.bytesToHex(h, .lower)); +} diff --git a/zig/src/families/ling/cuda.zig b/zig/src/families/ling/cuda.zig new file mode 100644 index 0000000000..0743957aa8 --- /dev/null +++ b/zig/src/families/ling/cuda.zig @@ -0,0 +1,97 @@ +//! Ling-3.1-flash on CUDA: its config as the GLM-5.3-Flash kernels it reuses take it (routing groups, SwiGLU limits). +const std = @import("std"); +const glm = @import("glm_cuda"); +pub const config = @import("config.zig"); +pub const names = @import("names.zig"); +pub const weights = @import("weights.zig"); +pub const synth = @import("synth.zig"); +pub const graph = @import("graph.zig"); +pub const state = @import("state.zig"); +pub const forward = @import("forward.zig"); +pub const mtp = @import("mtp.zig"); +pub const engine = @import("engine.zig"); +pub const lanes = @import("lanes.zig"); +pub const native = @import("native.zig"); +pub const fake_driver = @import("fake_driver.zig"); + +const Config = config.Config; + +/// glue.select's group-limited routing: n_group groups, the best topk_group kept (parse holds select's bounds). +pub fn groups(c: *const Config) glm.glue.Groups { + return .{ .n = c.n_group, .keep = c.topk_group }; +} + +/// Layer i's SwiGLU clamps, 0 for none: experts.Layer.limit, moe.Layer.shared_limit, mlp.Layer.limit. +pub const Limits = struct { routed: f32, shared: f32, dense: f32 }; + +/// The lists cover the backbone; the MTP layer and the dense layers run unclamped (parse refuses a dense limit). +pub fn limits(c: *const Config, i: u32) Limits { + if (i >= c.layers or !c.isMoe(i)) return .{ .routed = 0, .shared = 0, .dense = 0 }; + return .{ .routed = c.expert_limit[i], .shared = c.shared_limit[i], .dense = 0 }; +} + +/// The MoE layer's routing fields as moe.Layer holds them (pointers and widths come from the weight loader). +pub fn route(c: *const Config, i: u32, l: *glm.moe.Layer) void { + const lim = limits(c, i); + l.top_k = c.topk; + l.scale = c.routed_scale; + l.norm = c.norm_topk; + l.groups = groups(c); + l.shared_limit = lim.shared; + l.routed.limit = lim.routed; +} + +test { + _ = @import("config_test.zig"); + _ = @import("names_test.zig"); + _ = @import("synth.zig"); + std.testing.refAllDecls(weights); // the loader compiles on hosts without a GPU + std.testing.refAllDecls(@import("weights_proj.zig")); + std.testing.refAllDecls(forward); // the layer graph compiles on hosts without a GPU + std.testing.refAllDecls(forward.Forward); + std.testing.refAllDecls(engine.Engine); // the engine, its lane backend and its server hooks too + std.testing.refAllDecls(lanes.Lanes); + std.testing.refAllDecls(mtp.Head); + _ = native; + _ = @import("fake_driver.zig"); + _ = @import("store_test.zig"); + _ = lanes; +} + +test "config's fixed widths are the kernels' own: KDA heads, the latent cache, the rope" { + const f = config.fixed; + try std.testing.expect(f.kda_dim == glm.kda.dk and f.kda_dim == glm.kda.dv); + try std.testing.expect(f.rope == glm.latent.rope_width and f.kv_lora + f.rope == state.cache_width); + try std.testing.expect(f.block == @import("cuda").fp8.block); +} + +test "Ling's routing as glue.select takes it: 8 groups keeping 4, and select's own bounds hold" { + var why: config.Why = .{}; + const c = try config.parse(std.testing.allocator, @embedFile("fixtures/config.json"), &why); + const g = groups(&c); + try std.testing.expectEqual(glm.glue.Groups{ .n = 8, .keep = 4 }, g); + try std.testing.expect(c.experts % g.n == 0 and g.keep <= g.n and c.topk <= g.keep * (c.experts / g.n) and g.n <= 32); + try std.testing.expect(!std.meta.eql(g, glm.glue.Groups{})); // GLM's one group would skip Ling's group pass +} + +test "per-layer limits on the MoE layers: routed 10 from 65, shared 10 from 64, none on dense or MTP layers" { + var why: config.Why = .{}; + const c = try config.parse(std.testing.allocator, @embedFile("fixtures/config.json"), &why); + for (0..c.layers + c.mtp) |i| { + const l = limits(&c, @intCast(i)); + const moe = i >= 4 and i < 72; + try std.testing.expectEqual(@as(f32, if (moe and i >= 65) 10 else 0), l.routed); + try std.testing.expectEqual(@as(f32, if (moe and i >= 64) 10 else 0), l.shared); + try std.testing.expectEqual(@as(f32, 0), l.dense); + } + var l: glm.moe.Layer = undefined; + route(&c, 64, &l); + try std.testing.expectEqual(@as(f32, 0), l.routed.limit); + try std.testing.expectEqual(@as(f32, 10), l.shared_limit); + route(&c, 65, &l); + try std.testing.expectEqual(@as(f32, 10), l.routed.limit); + try std.testing.expectEqual(glm.glue.Groups{ .n = 8, .keep = 4 }, l.groups); + try std.testing.expectEqual(@as(usize, 8), l.top_k); + try std.testing.expectEqual(@as(f32, 2.5), l.scale); + try std.testing.expect(l.norm); +} diff --git a/zig/src/families/ling/engine.zig b/zig/src/families/ling/engine.zig new file mode 100644 index 0000000000..4d3b4ae1da --- /dev/null +++ b/zig/src/families/ling/engine.zig @@ -0,0 +1,438 @@ +//! Ling-3.1-flash's TP=1 CUDA engine: a sequence per slot, prompt chunks cut anywhere, verify windows, kept states. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const config = @import("config.zig"); +const weights = @import("weights.zig"); +const fw = @import("forward.zig"); +const head = @import("mtp.zig"); + +const sampler = glm.sampler; +pub const Sampling = sampler.Sampling; + +/// A prompt pass's chunk rows by default (Options.prompt_rows); marks cut chunks anywhere besides. +pub const prompt_rows = 1024; +/// A verify window's most rows: the pending token and up to 15 drafts (forward.window_rows). +pub const window_rows = fw.window_rows; +/// A shared verify's most rows and streams. +pub const shared_rows = 64; +pub const max_shared = 8; +/// The MTP layer's chain at most. +pub const max_drafts = head.max_drafts; + +pub const Options = struct { + capacity: usize, // tokens a sequence holds + prompt_rows: usize = prompt_rows, // rows a prompt chunk takes at most + mtp: bool = false, // load the MTP layer, if the checkpoint has one, and draft with it + layers: ?[]const u32 = null, // a thin model's backbone layers, ascending (checks on a real checkpoint); null: all +}; + +/// Rank 0's word on its request (cancelled or not), asked before each chunk. +pub const Cancel = struct { ptr: ?*anyopaque, check: *const fn (?*anyopaque) bool }; +/// One stream's window of a shared verify: its slot, its ids, and its drawn tokens. +pub const Part = struct { slot: u32, ids: []const u32, out: []u32 }; +/// A kept prompt state as the prompt cache holds it: its id and position. +pub const Kept = struct { id: u64, at: u32 }; +/// Where a prompt pass resumes: kept state `id` at `at`. +pub const Resume = struct { id: u64, at: usize }; +/// The pass's marks: it cuts a chunk at each and calls `call` standing there. +pub const Marks = struct { at: []const u32, ctx: *anyopaque, call: *const fn (ctx: *anyopaque, at: u32) void }; +/// One chunk's outcome: where the pass stands, and the first token once the prompt is done. +pub const Step = struct { at: usize, first: ?u32 }; + +/// A stream's sequence: the layers' caches, the MTP layer's, its sampling rule, the rows its next draft reads. +const Seq = struct { + st: fw.State, + hd: ?fw.State = null, // the MTP layer's MLA cache at its own position + hold: ?cuda.DeviceBuffer = null, // the prompt's last final-norm row: the first draft's hidden row + sampling: ?Sampling = null, // null: greedy + in_head: usize = 0, // chained draft rows the MTP cache holds past the absorbed ones + rows_at: u64 = 0, // the last verify's final-norm rows (bf16 [rows][hidden]): the next draft's hidden rows +}; + +/// A kept state: the sequence's arena copied whole, and the host fields beside it. +const Snap = struct { buf: cuda.DeviceBuffer, hd: ?cuda.DeviceBuffer, pos: usize, hd_pos: usize, cur: []u1 }; + +pub const Engine = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + c: config.Config, + m: fw.Modules, + stream: cuda.Stream, + w: weights.Weights, // every loaded layer + main: weights.Weights, // a view of w: the backbone layers alone + g: fw.Graph, + mtp: ?head.Head = null, + bp: fw.Buffers, // prompt chunks + bd: fw.Buffers, // one stream's verify window + bs: fw.Buffers, // shared verify windows + smp: sampler.Sampler, + pairs: cuda.DeviceBuffer, // (value, id) a row of greedy draws + seqs: std.ArrayList(?*Seq) = .empty, // by slot; slot 0 made at init + cur: *Seq = undefined, + slot: u32 = 0, + capacity: usize, + chunk: usize, + snaps: std.AutoHashMapUnmanaged(u64, Snap) = .empty, + next_snap: u64 = 1, + + /// The checkpoint at `dir` in its own format. + pub fn init(gpa: std.mem.Allocator, io: std.Io, d: *const cuda.Driver, ctx: *const cuda.Context, dir: []const u8, o: Options) !*Engine { + if (o.prompt_rows == 0 or o.capacity < window_rows + 2) return error.BadOptions; + var why: config.Why = .{}; + var opened = weights.open(gpa, io, dir, &why) catch |err| { + std.log.err("{s}: {s}", .{ dir, why.text() }); + return err; + }; + defer opened.ck.close(); + const c = opened.c; + const mtp = o.mtp and c.mtp > 0; + const e = try gpa.create(Engine); + errdefer gpa.destroy(e); + e.* = .{ .gpa = gpa, .d = d, .c = c, .m = undefined, .stream = undefined, .w = undefined, .main = undefined, .g = undefined, .bp = undefined, .bd = undefined, .bs = undefined, .smp = undefined, .pairs = undefined, .capacity = o.capacity, .chunk = o.prompt_rows }; + e.m = try fw.Modules.load(d, ctx); + errdefer e.m.deinit(); + e.stream = try cuda.Stream.init(d, false); + errdefer e.stream.deinit(); + const layers = try backbone(gpa, &c, o.layers); + defer gpa.free(layers); + { + const all = try std.mem.concat(gpa, u32, &.{ layers, if (mtp) &.{c.layers} else &.{} }); + defer gpa.free(all); + var pack = try cuda.modelopt.Packer.load(d); + defer pack.unload(); + e.w = try weights.load(gpa, d, e.stream, &pack, &opened.ck, c, all, .{}); + } + errdefer e.w.drop(weights.drains(e.stream)); // leaked on a failed drain: a queued launch may read them + e.main = e.w; + e.main.layers = e.w.layers[0..layers.len]; + e.g = try fw.Graph.init(gpa, d, &e.m, e.stream, &e.main); + errdefer e.g.drop(weights.drains(e.stream)); // as e.w + if (mtp) { + e.mtp = @as(head.Head, undefined); + e.mtp.?.init(gpa, d, &e.m, e.stream, &e.w, o.prompt_rows, o.capacity) catch |err| { + e.mtp = null; + return err; + }; + } + errdefer if (e.mtp) |*h| h.deinit(); + e.bp = try fw.Buffers.init(d, &e.m, &e.w, o.prompt_rows, true, o.capacity); + errdefer e.bp.deinit(); + e.bd = try fw.Buffers.init(d, &e.m, &e.w, window_rows, false, o.capacity); + errdefer e.bd.deinit(); + e.bs = try fw.Buffers.init(d, &e.m, &e.w, shared_rows, false, o.capacity); + errdefer e.bs.deinit(); + e.smp = try sampler.Sampler.init(d, c.vocab, window_rows); + errdefer e.smp.deinit(); + e.pairs = try cuda.DeviceBuffer.alloc(d, 8 * window_rows); + errdefer e.pairs.free(); + errdefer { + for (0..e.seqs.items.len) |i| e.free(i); + e.seqs.deinit(gpa); + } + try e.bind(0); + try e.stream.synchronize(); + return e; + } + + /// The backbone layers to load: `thin` (ascending, each below c.layers), or all of them. + fn backbone(gpa: std.mem.Allocator, c: *const config.Config, thin: ?[]const u32) ![]u32 { + if (thin) |t| { + for (t, 0..) |i, k| if (i >= c.layers or (k > 0 and i <= t[k - 1])) return error.BadLayers; + return gpa.dupe(u32, t); + } + const all = try gpa.alloc(u32, c.layers); + for (all, 0..) |*l, i| l.* = @intCast(i); + return all; + } + + pub fn deinit(e: *Engine) void { + e.stream.synchronize() catch {}; + var it = e.snaps.valueIterator(); + while (it.next()) |s| e.freeSnap(s); + e.snaps.deinit(e.gpa); + for (0..e.seqs.items.len) |i| e.free(i); + e.seqs.deinit(e.gpa); + e.pairs.free(); + e.smp.deinit(); + e.bs.deinit(); + e.bd.deinit(); + e.bp.deinit(); + if (e.mtp) |*h| h.deinit(); + e.g.deinit(); + e.w.deinit(); + e.stream.deinit(); + e.m.deinit(); + const gpa = e.gpa; + gpa.destroy(e); + } + + pub fn forward(e: *const Engine, b: *const fw.Buffers) fw.Forward { + return .{ .m = &e.m, .g = &e.g, .s = e.stream, .b = b }; + } + + /// Binds slot `slot`'s sequence for the steps that follow, making it (zeroed) on first use. + pub fn bind(e: *Engine, slot: u32) !void { + if (slot >= e.seqs.items.len) try e.seqs.appendNTimes(e.gpa, null, slot + 1 - e.seqs.items.len); + if (e.seqs.items[slot] == null) { + const q = try e.gpa.create(Seq); + errdefer e.gpa.destroy(q); + q.* = .{ .st = try fw.State.init(e.gpa, e.d, &e.g, e.capacity) }; + errdefer e.freeSeq(q); + if (e.mtp) |*h| { + q.hd = try fw.State.init(e.gpa, e.d, &h.g, e.capacity); + q.hold = try cuda.DeviceBuffer.alloc(e.d, @as(usize, e.c.hidden) * 2); + } + e.seqs.items[slot] = q; + } + e.cur = e.seqs.items[slot].?; + e.slot = slot; + } + + /// Frees slot `slot`'s sequence (slot 0 stays) and binds slot 0. + pub fn release(e: *Engine, slot: u32) void { + if (slot != 0) e.free(slot); + e.cur = e.seqs.items[0].?; + e.slot = 0; + } + + fn free(e: *Engine, slot: usize) void { + const q = e.seqs.items[slot] orelse return; + e.stream.synchronize() catch {}; + e.freeSeq(q); + e.gpa.destroy(q); + e.seqs.items[slot] = null; + } + + fn freeSeq(e: *Engine, q: *Seq) void { + _ = e; + q.st.deinit(); + if (q.hd) |*h| h.deinit(); + if (q.hold) |*b| b.free(); + q.hd = null; + q.hold = null; + } + + /// Bytes one more stream's sequence takes: its arenas. + pub fn seqBytes(e: *const Engine) u64 { + const q = e.seqs.items[0].?; + return q.st.arena.buf.len + if (q.hd) |h| h.arena.buf.len + @as(usize, e.c.hidden) * 2 else 0; + } + + /// The committed length: the pending token sits here, a window's row r draws position pos + 1 + r. + pub fn pos(e: *const Engine) usize { + return e.cur.st.pos; + } + + /// A fresh sequence with `sampling` (null or temperature 0: greedy). + pub fn begin(e: *Engine, sampling: ?Sampling) !void { + e.cur.sampling = sampler.check(sampling); + e.cur.in_head = 0; + try e.cur.st.reset(e.stream); + if (e.cur.hd) |*h| try h.reset(e.stream); + } + + /// The whole prompt pass: chunks, then the first token from the last row; `from` restores a kept state. + pub fn prefill(e: *Engine, ids: []const u32, from: ?Resume, cancel: ?Cancel, marks: ?Marks) !u32 { + var at = try e.prefillStart(ids.len, from); + while (true) { + const s = try e.prefillChunk(ids, at, cancel, marks); + if (s.first) |t| return t; + at = s.at; + } + } + + /// Where a pass over `len` ids starts: 0, or `from`'s position with its kept state restored. + pub fn prefillStart(e: *Engine, len: usize, from: ?Resume) !usize { + if (len == 0 or len + window_rows >= e.capacity) return error.PromptTooLong; + const r = from orelse return 0; + if (r.at == 0 or r.at >= len) return error.BadResume; + try e.restoreState(r.id); + if (e.cur.st.pos != r.at) return error.BadResume; + return r.at; + } + + /// One chunk of `ids` from `at` (the sequence's position): up to the chunk rows, cut at the next mark. + pub fn prefillChunk(e: *Engine, ids: []const u32, at: usize, cancel: ?Cancel, marks: ?Marks) !Step { + const len = ids.len; + if (at >= len or at != e.cur.st.pos) return error.BadChunk; + if (marks) |m| if (at > 0 and std.mem.indexOfScalar(u32, m.at, @intCast(at)) != null) m.call(m.ctx, @intCast(at)); + if (cancel) |c| if (c.check(c.ptr)) return error.Cancelled; + var end = @min(len, at + e.chunk); + if (marks) |m| for (m.at) |x| if (x > at and x < end) { + end = x; + break; + }; + const rows = end - at; + const f = e.forward(&e.bp); + const seg = [_]fw.Seg{.{ .st = &e.cur.st, .row0 = 0, .rows = rows }}; + try f.embed(@ptrCast(ids[at..end])); + try f.stack(&seg); + try f.final(rows); + var first: ?u32 = null; + if (end == len) { + try f.logitsOf(rows - 1, 1); + var t: [1]u32 = undefined; + try e.draw(e.bp.logits, 1, len, &t); + first = t[0]; + } + if (e.mtp) |*h| { // each row whose next token is known through the MTP layer (it overwrites b.normed) + const q = e.cur; + if (end == len) try e.copy(q.hold.?.ptr, e.bp.normed + (rows - 1) * e.c.hidden * 2, e.c.hidden * 2); + const known = @min(rows, len - 1 - at); + if (known > 0) try h.absorb(&e.m, e.stream, &e.bp, &q.hd.?, @ptrCast(ids[at + 1 ..][0..known]), e.bp.normed); + } + try f.commit(&e.cur.st, rows); + return .{ .at = end, .first = first }; + } + + /// A verify window: `ids` (pending token, then drafts) at the committed position; `out[r]` drawn at pos + 1 + r. + pub fn verify(e: *Engine, ids: []const u32, out: []u32) !void { + if (ids.len == 0 or ids.len > window_rows) return error.WindowTooWide; + const f = e.forward(&e.bd); + const at = e.cur.st.pos; + try f.run(&.{.{ .st = &e.cur.st, .row0 = 0, .rows = ids.len }}, @ptrCast(ids)); + e.cur.rows_at = e.bd.normed; + try e.draw(e.bd.logits, ids.len, at + 1, out); + } + + /// Several streams' windows in one forward, each stream's rows on its own sequence. + pub fn verifyShared(e: *Engine, parts: []const Part) !void { + if (parts.len == 0 or parts.len > max_shared) return error.WindowTooWide; + var segs: [max_shared]fw.Seg = undefined; + var ids: [shared_rows]u32 = undefined; + var rows: usize = 0; + for (parts, 0..) |p, i| { + if (p.ids.len == 0 or p.ids.len > window_rows or rows + p.ids.len > shared_rows) return error.WindowTooWide; + for (parts[0..i]) |q| if (q.slot == p.slot) return error.SlotTwice; + try e.bind(p.slot); + segs[i] = .{ .st = &e.cur.st, .row0 = rows, .rows = p.ids.len }; + @memcpy(ids[rows..][0..p.ids.len], p.ids); + rows += p.ids.len; + } + try e.forward(&e.bs).run(segs[0..parts.len], @ptrCast(ids[0..rows])); + for (parts, segs[0..parts.len]) |p, sg| { + try e.bind(p.slot); + e.cur.rows_at = e.bs.normed + sg.row0 * e.c.hidden * 2; + try e.draw(e.bs.logits + sg.row0 * e.c.vocab * 2, sg.rows, sg.st.pos + 1, p.out); + } + } + + /// The last window's first `n` rows committed (the KDA states replayed when fewer than it ran). + pub fn keep(e: *Engine, n: usize) !void { + try e.forward(&e.bd).commit(&e.cur.st, n); + } + + /// The MTP layer absorbs the kept rows with `follow` (the token after each), then chains `depth` drafts. + pub fn draft(e: *Engine, from_prompt: bool, follow: []const u32, depth: usize, out: []u32) !void { + const h = if (e.mtp) |*x| x else return error.NoDraftHead; + if (depth > max_drafts or out.len < depth) return error.TooDeep; + if (follow.len == 0) return error.NoFollow; + const q = e.cur; + const hd = &q.hd.?; + const D: usize = e.c.hidden; + try hd.setPos(e.stream, hd.pos - q.in_head); // the last chain's rows go + q.in_head = 0; + const hidden = if (from_prompt) q.hold.?.ptr else q.rows_at; + var s0: usize = 0; + while (s0 < follow.len) : (s0 += window_rows) { + const part = follow[s0..@min(follow.len, s0 + window_rows)]; + try h.absorb(&e.m, e.stream, &h.bm, hd, @ptrCast(part), hidden + s0 * D * 2); + } + if (hd.pos != q.st.pos) return error.MtpOutOfStep; + const f = h.forward(&e.m, e.stream, &h.bm); + var last = (follow.len - 1) % window_rows; // h.bm's row of the newest absorbed token + for (0..depth) |j| { + try f.final(last + 1); + try f.logitsOf(last, 1); + try e.draw(h.bm.logits, 1, q.st.pos + 1 + j, out[j..][0..1]); + if (j + 1 == depth) break; + try h.absorb(&e.m, e.stream, &h.bm, hd, @ptrCast(out[j..][0..1]), h.bm.normed + last * D * 2); + q.in_head += 1; + last = 0; + } + } + + /// `rows` logits rows -> the tokens at positions first, first + 1, ... + pub fn draw(e: *Engine, logits: u64, rows: usize, first: u64, out: []u32) !void { + if (e.cur.sampling) |s| { + const span = [_]glm.shard.Span{.{ .lo = 0, .hi = e.c.vocab }}; + return e.smp.draw(e.stream, .local, &span, &.{.{ .rank = 0, .logits = logits }}, rows, first, sampler.rule(s), out); + } + const V: usize = e.c.vocab; + for (0..rows) |r| { + var a: cuda.Args = .{}; + inline for (.{ logits + r * V * 2, @as(i64, @intCast(V)), @as(u32, 0), e.pairs.ptr + r * 8 }) |v| a.add(v); + try cuda.launch.launch(e.m.glue.top1, .{ .grid = .{ .x = 1 }, .block = .{ .x = 1024 } }, e.stream, &a); + } + try e.stream.synchronize(); + var got: [window_rows][2]u32 = undefined; + try e.pairs.download(0, std.mem.sliceAsBytes(got[0..rows])); + for (out[0..rows], got[0..rows]) |*o, p| o.* = p[1]; + } + + /// Bytes one kept state takes: the sequence's arenas. + pub fn stateBytes(e: *const Engine) u64 { + return e.seqBytes(); + } + + /// A copy of the bound sequence's state under a new id (the pass standing at a mark). + pub fn keepState(e: *Engine) !u64 { + const q = e.cur; + if (q.st.pending != 0) return error.NotCommitted; + var buf = try cuda.DeviceBuffer.alloc(e.d, q.st.arena.buf.len); + errdefer if (weights.drains(e.stream)) buf.free(); // else leaked: the copy into it may be queued + try e.copy(buf.ptr, q.st.arena.buf.ptr, buf.len); + var hd: ?cuda.DeviceBuffer = null; + errdefer if (hd) |*x| if (weights.drains(e.stream)) x.free(); + if (q.hd) |h| { + if (q.in_head != 0) return error.NotCommitted; + hd = try cuda.DeviceBuffer.alloc(e.d, h.arena.buf.len); + try e.copy(hd.?.ptr, h.arena.buf.ptr, h.arena.buf.len); + } + try e.stream.synchronize(); + const cur = try e.gpa.dupe(u1, q.st.cur); + errdefer e.gpa.free(cur); + const id = e.next_snap; + try e.snaps.put(e.gpa, id, .{ .buf = buf, .hd = hd, .pos = q.st.pos, .hd_pos = if (q.hd) |h| h.pos else 0, .cur = cur }); + e.next_snap += 1; + return id; + } + + /// The bound sequence made kept state `id`'s. + fn restoreState(e: *Engine, id: u64) !void { + const s = e.snaps.getPtr(id) orelse return error.NoKeptState; + const q = e.cur; + try e.copy(q.st.arena.buf.ptr, s.buf.ptr, s.buf.len); + @memcpy(q.st.cur, s.cur); + q.st.pending = 0; + q.st.prompt = false; + try q.st.setPos(e.stream, s.pos); + if (q.hd) |*h| { + try e.copy(h.arena.buf.ptr, (s.hd orelse return error.NoKeptState).ptr, h.arena.buf.len); + h.pending = 0; + try h.setPos(e.stream, s.hd_pos); + } + q.in_head = 0; + try e.stream.synchronize(); + } + + pub fn holds(e: *const Engine, id: u64) bool { + return e.snaps.contains(id); + } + + pub fn dropState(e: *Engine, id: u64) void { + var kv = e.snaps.fetchRemove(id) orelse return; + e.freeSnap(&kv.value); + } + + fn freeSnap(e: *Engine, s: *Snap) void { + s.buf.free(); + if (s.hd) |*x| x.free(); + e.gpa.free(s.cur); + } + + pub fn copy(e: *const Engine, dst: u64, src: u64, bytes: usize) !void { + try e.d.check(e.d.api.cuMemcpyDtoDAsync_v2(dst, src, bytes, e.stream.handle), "cuMemcpyDtoDAsync"); + } +}; diff --git a/zig/src/families/ling/fake_driver.zig b/zig/src/families/ling/fake_driver.zig new file mode 100644 index 0000000000..7a8efc8ad5 --- /dev/null +++ b/zig/src/families/ling/fake_driver.zig @@ -0,0 +1,412 @@ +//! A host stand-in for the CUDA driver in checks: device memory in host memory, no-op kernels, a failure on cue. +const std = @import("std"); +const cuda = @import("cuda"); + +const abi = cuda.abi; +const R = abi.Result; +const ok = abi.success; +const backing = std.heap.page_allocator; + +/// What the entry points share (they are C functions, so process-wide): allocations, handles, the countdown. +var held: std.AutoHashMapUnmanaged(u64, usize) = .empty; // device and pinned host allocations by address +var handles: isize = 0; // modules, streams, events, graphs and execs made and not yet released +var calls: u64 = 0; // fallible calls since reset +var fail_at: ?u64 = null; // the call that fails +var code: R = 2; // what it returns (2: out of memory; NOT_READY reaches the caller without a log line) +var bad: usize = 0; // frees of something not held, releases of nothing +var lazy = false; // async copies, memsets and launches stay pending until a stream or context drains +var sticky = false; // every call from the failing one on fails too, drains included +var pending: usize = 0; // async work queued since the last drain (lazy only) +var early: usize = 0; // frees while work was pending +var stuck: usize = 0; // drains that failed +var hit = false; // the failing call came with work pending +var next: usize = 0x1000; +var globals: [1 << 16]u8 align(256) = undefined; + +/// From now on, fallible call number `at` (1-based) returns `with`; null fails none. Bad releases start at 0. +pub fn reset(at: ?u64, with: R) void { + calls = 0; + bad = 0; + pending = 0; + early = 0; + stuck = 0; + hit = false; + fail_at = at; + code = with; +} + +/// From now on async work stays pending until a successful drain (`queue`); failures persist (`stick`). +pub fn mode(queue: bool, stick: bool) void { + lazy = queue; + sticky = stick; +} + +/// Frees of device or pinned memory while async work was pending, since reset; work pending now; failed drains. +pub fn freedEarly() usize { + return early; +} + +pub fn queued() usize { + return pending; +} + +pub fn failedDrains() usize { + return stuck; +} + +pub fn failedPending() bool { + return hit; +} + +/// Everything still held freed (what a failed drain leaked on purpose); returns how many. +pub fn sweep() usize { + const n = held.count(); + var it = held.iterator(); + while (it.next()) |kv| backing.free(@as([*]align(256) u8, @ptrFromInt(kv.key_ptr.*))[0..kv.value_ptr.*]); + held.clearRetainingCapacity(); + return n; +} + +/// Fallible calls since reset, and whether the failure was handed out. +pub fn count() u64 { + return calls; +} + +pub fn fired() bool { + return if (fail_at) |n| calls >= n else false; +} + +/// Allocations and handles held now; frees and releases of things never made. +pub fn live() usize { + return held.count() + @as(usize, @intCast(@max(handles, 0))); +} + +pub fn badReleases() usize { + return bad + @as(usize, @intCast(@max(-handles, 0))); +} + +/// A Driver whose entry points are these (its library is never opened or closed). +pub fn driver() cuda.Driver { + var api: abi.Api = undefined; + inline for (@typeInfo(abi.Api).@"struct".field_names) |name| @field(api, name) = @ptrCast(&passes); + inline for (comptime std.meta.declarations(table)) |name| @field(api, name) = &@field(table, name); + return .{ .lib = undefined, .api = api }; +} + +/// A context on `d` (device 0) that nothing retains. +pub fn context(d: *const cuda.Driver) cuda.Context { + return .{ .d = d, .device = 0, .handle = @ptrFromInt(0x10) }; +} + +/// The countdown: each fallible call takes one step, and the chosen one fails. +fn step() R { + calls += 1; + const n = fail_at orelse return ok; + if (calls == n) hit = pending > 0; + return if (calls == n or (sticky and calls > n)) code else ok; +} + +/// An async call's result: queued work pending until a drain when lazy. +fn later(r: R) R { + if (r == ok and lazy) pending += 1; + return r; +} + +fn drain() R { + const r = step(); + if (r == ok) pending = 0 else stuck += 1; + return r; +} + +/// Every entry point without out-parameters: a step, nothing else (the C ABI lets a caller pass arguments it ignores). +fn passes() callconv(.c) R { + return step(); +} + +fn handle() ?*anyopaque { + next += 16; + handles += 1; + return @ptrFromInt(next); +} + +fn release() R { + handles -= 1; + return ok; +} + +fn take(n: usize) ?u64 { + const m = backing.alignedAlloc(u8, .fromByteUnits(256), @max(n, 1)) catch return null; + held.put(backing, @intFromPtr(m.ptr), @max(n, 1)) catch { + backing.free(m); + return null; + }; + return @intFromPtr(m.ptr); +} + +fn give(p: u64) R { + if (pending > 0) early += 1; + const kv = held.fetchRemove(p) orelse { + bad += 1; + return 1; + }; + backing.free(@as([*]align(256) u8, @ptrFromInt(p))[0..kv.value]); + return ok; +} + +fn bytes(p: u64, n: usize) []u8 { + return @as([*]u8, @ptrFromInt(p))[0..n]; +} + +fn cbytes(p: ?*const anyopaque, n: usize) []const u8 { + return @as([*]const u8, @ptrCast(p.?))[0..n]; +} + +const table = struct { + pub fn cuDriverGetVersion(v: *c_int) callconv(.c) R { + v.* = 13000; + return ok; + } + pub fn cuDeviceGet(d: *abi.Device, _: c_int) callconv(.c) R { + d.* = 0; + return ok; + } + pub fn cuDeviceGetCount(n: *c_int) callconv(.c) R { + n.* = 1; + return ok; + } + pub fn cuDeviceGetName(out: [*]u8, len: c_int, _: abi.Device) callconv(.c) R { + const name = "fake"; + @memcpy(out[0..@min(name.len, @as(usize, @intCast(len)))], name[0..@min(name.len, @as(usize, @intCast(len)))]); + if (len > name.len) out[name.len] = 0; + return ok; + } + pub fn cuDeviceGetAttribute(v: *c_int, a: abi.DeviceAttribute, _: abi.Device) callconv(.c) R { + v.* = switch (a) { + .max_threads_per_block => 1024, + .warp_size => 32, + .clock_rate => 1_000_000, + .multiprocessor_count => 48, + .integrated => 1, + .l2_cache_size => 24 << 20, + .compute_capability_major => 12, + .compute_capability_minor => 1, + .max_shared_memory_per_multiprocessor, .max_shared_memory_per_block_optin => 101376, + }; + return step(); + } + pub fn cuDeviceTotalMem_v2(n: *usize, _: abi.Device) callconv(.c) R { + n.* = 128 << 30; + return step(); + } + pub fn cuDevicePrimaryCtxRetain(c: *abi.Context, _: abi.Device) callconv(.c) R { + c.* = @ptrFromInt(0x10); + return step(); + } + pub fn cuDevicePrimaryCtxRelease_v2(_: abi.Device) callconv(.c) R { + return ok; + } + pub fn cuCtxGetCurrent(c: *abi.Context) callconv(.c) R { + c.* = @ptrFromInt(0x10); + return ok; + } + pub fn cuMemGetInfo_v2(free: *usize, total: *usize) callconv(.c) R { + free.* = 100 << 30; + total.* = 128 << 30; + return step(); + } + pub fn cuMemAlloc_v2(p: *abi.DevicePtr, n: usize) callconv(.c) R { + if (step() != ok) return code; + p.* = take(n) orelse return 2; + return ok; + } + pub fn cuMemFree_v2(p: abi.DevicePtr) callconv(.c) R { + return give(p); + } + pub fn cuMemHostAlloc(p: *?*anyopaque, n: usize, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + p.* = @ptrFromInt(take(n) orelse return 2); + return ok; + } + pub fn cuMemHostGetDevicePointer_v2(p: *abi.DevicePtr, host: ?*anyopaque, _: c_uint) callconv(.c) R { + p.* = @intFromPtr(host); + return step(); + } + pub fn cuMemFreeHost(p: ?*anyopaque) callconv(.c) R { + return give(@intFromPtr(p)); + } + pub fn cuMemHostUnregister(_: ?*anyopaque) callconv(.c) R { + return ok; + } + pub fn cuMemcpyHtoD_v2(dst: abi.DevicePtr, src: ?*const anyopaque, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memcpy(bytes(dst, n), cbytes(src, n)); + return ok; + } + pub fn cuMemcpyDtoH_v2(dst: ?*anyopaque, src: abi.DevicePtr, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memcpy(bytes(@intFromPtr(dst), n), bytes(src, n)); + return ok; + } + pub fn cuMemcpyDtoD_v2(dst: abi.DevicePtr, src: abi.DevicePtr, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memmove(bytes(dst, n), bytes(src, n)); + return ok; + } + pub fn cuMemcpyHtoDAsync_v2(dst: abi.DevicePtr, src: ?*const anyopaque, n: usize, _: abi.Stream) callconv(.c) R { + return later(cuMemcpyHtoD_v2(dst, src, n)); + } + pub fn cuMemcpyDtoHAsync_v2(dst: ?*anyopaque, src: abi.DevicePtr, n: usize, _: abi.Stream) callconv(.c) R { + return later(cuMemcpyDtoH_v2(dst, src, n)); + } + pub fn cuMemcpyDtoDAsync_v2(dst: abi.DevicePtr, src: abi.DevicePtr, n: usize, _: abi.Stream) callconv(.c) R { + return later(cuMemcpyDtoD_v2(dst, src, n)); + } + pub fn cuMemcpy2DAsync_v2(m: *const abi.Memcpy2D, _: abi.Stream) callconv(.c) R { + if (step() != ok) return code; + for (0..m.height) |r| @memmove(bytes(m.dst + r * m.dst_pitch, m.width), bytes(m.src + r * m.src_pitch, m.width)); + return later(ok); + } + pub fn cuMemsetD8_v2(p: abi.DevicePtr, v: u8, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memset(bytes(p, n), v); + return ok; + } + pub fn cuMemsetD32_v2(p: abi.DevicePtr, v: c_uint, n: usize) callconv(.c) R { + if (step() != ok) return code; + @memset(@as([*]align(1) u32, @ptrFromInt(p))[0..n], v); + return ok; + } + pub fn cuMemsetD8Async(p: abi.DevicePtr, v: u8, n: usize, _: abi.Stream) callconv(.c) R { + return later(cuMemsetD8_v2(p, v, n)); + } + pub fn cuMemsetD32Async(p: abi.DevicePtr, v: c_uint, n: usize, _: abi.Stream) callconv(.c) R { + return later(cuMemsetD32_v2(p, v, n)); + } + pub fn cuLaunchKernel(_: abi.Function, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: c_uint, _: abi.Stream, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) R { + return later(step()); + } + pub fn cuLaunchKernelEx(_: *const abi.LaunchConfig, _: abi.Function, _: ?[*]?*anyopaque, _: ?[*]?*anyopaque) callconv(.c) R { + return later(step()); + } + pub fn cuStreamSynchronize(_: abi.Stream) callconv(.c) R { + return drain(); + } + pub fn cuCtxSynchronize() callconv(.c) R { + return drain(); + } + pub fn cuStreamCreate(s: *abi.Stream, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + s.* = @ptrCast(handle()); + return ok; + } + pub fn cuStreamDestroy_v2(_: abi.Stream) callconv(.c) R { + return release(); + } + pub fn cuStreamIsCapturing(_: abi.Stream, s: *abi.CaptureStatus) callconv(.c) R { + s.* = .none; + return step(); + } + pub fn cuEventCreate(e: *abi.Event, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + e.* = @ptrCast(handle()); + return ok; + } + pub fn cuEventDestroy_v2(_: abi.Event) callconv(.c) R { + return release(); + } + pub fn cuEventElapsedTime(ms: *f32, _: abi.Event, _: abi.Event) callconv(.c) R { + ms.* = 0; + return step(); + } + pub fn cuModuleLoadData(m: *abi.Module, _: ?*const anyopaque) callconv(.c) R { + if (step() != ok) return code; + m.* = @ptrCast(handle()); + return ok; + } + pub fn cuModuleLoadDataEx(m: *abi.Module, image: ?*const anyopaque, _: c_uint, _: ?[*]c_int, _: ?[*]?*anyopaque) callconv(.c) R { + return cuModuleLoadData(m, image); + } + pub fn cuModuleUnload(_: abi.Module) callconv(.c) R { + return release(); + } + pub fn cuModuleGetFunction(f: *abi.Function, _: abi.Module, _: [*:0]const u8) callconv(.c) R { + f.* = @ptrFromInt(0x20); + return step(); + } + pub fn cuModuleGetGlobal_v2(p: *abi.DevicePtr, n: *usize, _: abi.Module, _: [*:0]const u8) callconv(.c) R { + p.* = @intFromPtr(&globals); + n.* = globals.len; + return step(); + } + pub fn cuFuncGetAttribute(v: *c_int, a: abi.FunctionAttribute, _: abi.Function) callconv(.c) R { + v.* = switch (a) { + .max_threads_per_block => 1024, + .num_regs => 32, + .ptx_version, .binary_version => 121, + else => 0, + }; + return step(); + } + pub fn cuOccupancyMaxActiveBlocksPerMultiprocessor(n: *c_int, _: abi.Function, _: c_int, _: usize) callconv(.c) R { + n.* = 4; + return step(); + } + pub fn cuStreamEndCapture(_: abi.Stream, g: *abi.Graph) callconv(.c) R { + if (step() != ok) return code; + g.* = @ptrCast(handle()); + return ok; + } + pub fn cuGraphCreate(g: *abi.Graph, _: c_uint) callconv(.c) R { + if (step() != ok) return code; + g.* = @ptrCast(handle()); + return ok; + } + pub fn cuGraphDestroy(_: abi.Graph) callconv(.c) R { + return release(); + } + pub fn cuGraphAddKernelNode_v2(n: *abi.GraphNode, _: abi.Graph, _: ?[*]const abi.GraphNode, _: usize, _: *const abi.KernelNodeParams) callconv(.c) R { + n.* = @ptrFromInt(0x30); + return step(); + } + pub fn cuGraphGetNodes(_: abi.Graph, _: ?[*]abi.GraphNode, n: *usize) callconv(.c) R { + n.* = 0; + return step(); + } + pub fn cuGraphInstantiateWithFlags(x: *abi.GraphExec, _: abi.Graph, _: u64) callconv(.c) R { + if (step() != ok) return code; + x.* = @ptrCast(handle()); + return ok; + } + pub fn cuGraphExecDestroy(_: abi.GraphExec) callconv(.c) R { + return release(); + } + pub fn cuGetErrorName(_: R, s: *?[*:0]const u8) callconv(.c) R { + s.* = "CUDA_ERROR_INJECTED"; + return ok; + } + pub fn cuGetErrorString(_: R, s: *?[*:0]const u8) callconv(.c) R { + s.* = "a failure the fake driver was told to give"; + return ok; + } +}; + +test "the fake driver: buffers live in host memory, the chosen call fails, frees of strangers are counted" { + const d = driver(); + reset(null, ok); + var b = try cuda.DeviceBuffer.fromHost(&d, "abcdefgh"); + var back: [8]u8 = undefined; + try b.download(0, &back); + try std.testing.expectEqualStrings("abcdefgh", &back); + try std.testing.expectEqual(@as(usize, 1), live()); + b.free(); + try std.testing.expectEqual(@as(usize, 0), live()); + reset(2, abi.error_not_ready); + var s = try cuda.Stream.init(&d, false); + try std.testing.expectError(error.NotReady, cuda.DeviceBuffer.alloc(&d, 64)); + try std.testing.expect(fired()); + s.deinit(); + try std.testing.expectEqual(@as(usize, 0), live()); + _ = d.api.cuMemFree_v2(0x1234); + try std.testing.expectEqual(@as(usize, 1), badReleases()); + reset(null, ok); +} diff --git a/zig/src/families/ling/fixtures/config.json b/zig/src/families/ling/fixtures/config.json new file mode 100644 index 0000000000..bbdaf010c8 --- /dev/null +++ b/zig/src/families/ling/fixtures/config.json @@ -0,0 +1,254 @@ +{ + "architectures": [ + "BailingMoeV3ForCausalLM" + ], + "attention_dropout": 0.0, + "auto_map": { + "AutoConfig": "configuration_bailing_moe_v3.BailingMoeV3Config", + "AutoModel": "modeling_bailing_moe_v3.BailingMoeV3Model", + "AutoModelForCausalLM": "modeling_bailing_moe_v3.BailingMoeV3ForCausalLM" + }, + "embedding_dropout": 0.0, + "eos_token_id": 156895, + "expert_swiglu_limit_list": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 10, + 10, + 10, + 10, + 10, + 10, + 10 + ], + "first_k_dense_replace": 4, + "gated_attention_proj_granularity_type": "head_wise", + "group_norm_size": 1, + "head_dim": 128, + "hidden_act": "silu", + "hidden_size": 4096, + "initializer_range": 0.02, + "intermediate_size": 10240, + "kda_lower_bound": -5.0, + "kda_safe_gate": true, + "kv_lora_rank": 512, + "layer_group_size": 8, + "linear_silu": true, + "max_position_embeddings": 262144, + "max_window_layers": 20, + "moe_intermediate_size": 1280, + "moe_router_enable_expert_bias": true, + "moe_shared_expert_intermediate_size": 1280, + "mtp_loss_scaling_factor": 0, + "mtp_use_kda": false, + "n_group": 8, + "no_kda_lora": true, + "norm_topk_prob": true, + "num_attention_heads": 64, + "num_experts": 512, + "num_experts_per_tok": 8, + "num_hidden_layers": 72, + "num_key_value_heads": 64, + "num_kv_heads_for_linear_attn": 0, + "num_nextn_predict_layers": 1, + "num_shared_experts": 1, + "output_dropout": 0.0, + "output_router_logits": false, + "pad_token_id": 156892, + "partial_rotary_factor": 0.5, + "q_lora_rank": 1536, + "qk_head_dim": 192, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "quantization_config": { + "activation_scheme": "dynamic", + "fmt": "e4m3", + "quant_method": "fp8", + "modules_to_not_convert": [ + "lm_head", + "model.word_embeddings", + "b_proj", + "mlp.gate", + "model.layers.7.attention.g_proj", + "model.layers.15.attention.g_proj", + "model.layers.23.attention.g_proj", + "model.layers.31.attention.g_proj", + "model.layers.39.attention.g_proj", + "model.layers.47.attention.g_proj", + "model.layers.55.attention.g_proj", + "model.layers.63.attention.g_proj", + "model.layers.71.attention.g_proj", + "model.layers.72.eh_proj" + ], + "weight_block_size": [ + 128, + 128 + ] + }, + "rms_norm_eps": 1e-06, + "rope_interleave": true, + "rope_scaling": null, + "rope_theta": 6000000, + "rotary_dim": 64, + "routed_scaling_factor": 2.5, + "router_dtype": "fp32", + "scale_router_input": false, + "score_function": "sigmoid", + "scoring_func": "sigmoid", + "seq_aux": true, + "share_expert_swiglu_limit_list": [ + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 0, + 10, + 10, + 10, + 10, + 10, + 10, + 10, + 10 + ], + "short_conv_kernel_size": 4, + "tie_word_embeddings": false, + "topk_group": 4, + "topk_method": "noaux_tc", + "transformers_version": "4.45.0", + "up_proj_norm": false, + "use_bias": false, + "use_cache": true, + "use_kda_lora": false, + "use_mla_nope": false, + "use_nGPT": false, + "use_qk_norm": true, + "use_qkv_bias": false, + "v_head_dim": 128, + "value_norm": false, + "vocab_size": 157184, + "model_type": "bailing_hybrid", + "torch_dtype": "bfloat16" +} \ No newline at end of file diff --git a/zig/src/families/ling/fixtures/names.txt b/zig/src/families/ling/fixtures/names.txt new file mode 100644 index 0000000000..b2502f4427 --- /dev/null +++ b/zig/src/families/ling/fixtures/names.txt @@ -0,0 +1,61 @@ +ling31-names 214025 16f5c51dd86af196318448fc16a4177efbcbe83cbd34ab863409c57fc7f028a7 +lm_head.weight +model.layers.{L}.attention.A_log 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.b_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.dense.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.dense.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.dt_bias 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.f_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.f_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.g_proj.weight 0-72 +model.layers.{L}.attention.g_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70,72 +model.layers.{L}.attention.k_conv1d.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.k_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.k_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.kv_a_layernorm.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_a_proj_with_mqa.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_a_proj_with_mqa.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_b_proj.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.kv_b_proj.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.o_norm.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.o_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.o_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.q_a_layernorm.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_a_proj.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_a_proj.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_b_proj.weight 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_b_proj.weight_scale_inv 7,15,23,31,39,47,55,63,71-72 +model.layers.{L}.attention.q_conv1d.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.q_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.q_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.v_conv1d.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.v_proj.weight 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.attention.v_proj.weight_scale_inv 0-6,8-14,16-22,24-30,32-38,40-46,48-54,56-62,64-70 +model.layers.{L}.eh_proj.weight 72 +model.layers.{L}.enorm.weight 72 +model.layers.{L}.final_layernorm.weight 72 +model.layers.{L}.hnorm.weight 72 +model.layers.{L}.input_layernorm.weight 0-72 +model.layers.{L}.mlp.down_proj.weight 0-3 +model.layers.{L}.mlp.down_proj.weight_scale_inv 0-3 +model.layers.{L}.mlp.experts.{E}.down_proj.weight 4-72 512 +model.layers.{L}.mlp.experts.{E}.down_proj.weight_scale_inv 4-72 512 +model.layers.{L}.mlp.experts.{E}.gate_proj.weight 4-72 512 +model.layers.{L}.mlp.experts.{E}.gate_proj.weight_scale_inv 4-72 512 +model.layers.{L}.mlp.experts.{E}.up_proj.weight 4-72 512 +model.layers.{L}.mlp.experts.{E}.up_proj.weight_scale_inv 4-72 512 +model.layers.{L}.mlp.gate.expert_bias 4-72 +model.layers.{L}.mlp.gate.weight 4-72 +model.layers.{L}.mlp.gate_proj.weight 0-3 +model.layers.{L}.mlp.gate_proj.weight_scale_inv 0-3 +model.layers.{L}.mlp.shared_experts.down_proj.weight 4-72 +model.layers.{L}.mlp.shared_experts.down_proj.weight_scale_inv 4-72 +model.layers.{L}.mlp.shared_experts.gate_proj.weight 4-72 +model.layers.{L}.mlp.shared_experts.gate_proj.weight_scale_inv 4-72 +model.layers.{L}.mlp.shared_experts.up_proj.weight 4-72 +model.layers.{L}.mlp.shared_experts.up_proj.weight_scale_inv 4-72 +model.layers.{L}.mlp.up_proj.weight 0-3 +model.layers.{L}.mlp.up_proj.weight_scale_inv 0-3 +model.layers.{L}.post_attention_layernorm.weight 0-72 +model.norm.weight +model.word_embeddings.weight diff --git a/zig/src/families/ling/forward.zig b/zig/src/families/ling/forward.zig new file mode 100644 index 0000000000..61bc7befc5 --- /dev/null +++ b/zig/src/families/ling/forward.zig @@ -0,0 +1,131 @@ +//! Ling-3.1-flash's TP=1 forward on CUDA: embedding, the layer stack, final norm and head, and a window's commit. +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); +const state = @import("state.zig"); +const block_kda = @import("block_kda.zig"); +const block_mla = @import("block_mla.zig"); +const block_ffn = @import("block_ffn.zig"); + +const glue = graph.glue; +pub const Seg = block_kda.Seg; +pub const Modules = graph.Modules; +pub const Graph = graph.Graph; +pub const Buffers = state.Buffers; +pub const State = state.State; +pub const window_rows = state.window_rows; +pub const cache_width = state.cache_width; + +pub const Forward = struct { + m: *const graph.Modules, + g: *const graph.Graph, + s: cuda.Stream, + b: *const state.Buffers, + picks: ?u64 = null, // i32 [MoE layers][rows][top_k + 1]: every MoE layer's picks, for checks + kda_wide: bool = true, // GLM's TF_GLM_KDA_DECODE_WIDE default: decode windows on the three-kernel chain + + pub fn ops(f: *const Forward) graph.Ops { + return .{ .m = f.m, .s = f.s, .prompt = f.b.prompt, .part = f.b.part, .kda_wide = f.kda_wide }; + } + + /// Segments in row order from row 0, each on its own committed sequence with room (and scratch) for its rows. + fn rowsOf(f: *const Forward, segs: []const Seg) !usize { + var rows: usize = 0; + for (segs, 0..) |sg, i| { + if (sg.row0 != rows or sg.rows == 0 or sg.st.pending != 0) return error.BadSegments; + if (sg.st.pos + sg.rows > sg.st.capacity) return error.CacheFull; // before a prompt's KDA layers commit + if (f.g.mla_layers > 0) try block_mla.room(f.b, sg.st, sg.rows); // and the scratch's attention partials + if (!f.b.prompt and sg.rows > window_rows) return error.WindowTooWide; + for (segs[0..i]) |o| if (o.st == sg.st) return error.BadSegments; + rows += sg.rows; + } + if (rows == 0 or rows > f.b.rows) return error.TooManyRows; + return rows; + } + + /// `ids` (one per row of `segs`) through the embedding, every loaded layer, the final norm and the head. + pub fn run(f: *const Forward, segs: []const Seg, ids: []const i32) !void { + const rows = try f.rowsOf(segs); + if (ids.len != rows) return error.TooManyRows; + try f.embed(ids); + try f.stack(segs); + try f.head(rows); + } + + /// The tokens' bf16 embedding rows into b.x. + pub fn embed(f: *const Forward, ids: []const i32) !void { + const w = f.g.w; + const d = f.b.arena.buf.d; + if (ids.len > f.b.rows) return error.TooManyRows; + for (ids) |t| if (t < 0 or t >= w.config.vocab) return error.BadToken; + try d.check(d.api.cuMemcpyHtoDAsync_v2(f.b.ids, ids.ptr, ids.len * 4, f.s.handle), "cuMemcpyHtoDAsync"); + try glue.embedB16(f.ops().set(), f.s, f.b.ids, w.embed, f.b.x, ids.len, w.config.hidden, 1); + } + + /// Every loaded layer on b.x in place; each sequence then holds its rows pending until commit(). + pub fn stack(f: *const Forward, segs: []const Seg) !void { + const rows = try f.rowsOf(segs); + var moe: usize = 0; + for (f.g.w.layers, 0..) |*l, j| { + const slots: usize = f.g.w.config.topk + 1; + const p: ?u64 = if (f.picks) |x| x + moe * rows * slots * 4 else null; + try f.layer(segs, j, rows, p); + moe += @intFromBool(l.ffn == .moe); + } + for (segs) |sg| { + sg.st.pending = sg.rows; + sg.st.prompt = f.b.prompt; + } + } + + /// Loaded layer j: input norm, mixer, residual; post-attention norm, MLP or MoE, residual. + pub fn layer(f: *const Forward, segs: []const Seg, j: usize, rows: usize, picks: ?u64) !void { + const o = f.ops(); + const b = f.b; + const g = f.g; + const c = &g.w.config; + const D: usize = c.hidden; + const l = g.w.layers[j]; + try glue.rmsnorm(o.set(), o.s, b.x, D, l.attn_norm, b.normed, D, null, c.eps, rows, D); + switch (l.mixer) { + .kda => |a| try block_kda.block(o, b, a, c, g.slot[j], segs, rows), + .mla => |a| try block_mla.block(o, b, a, g.absorb[g.slot[j]], g, g.slot[j], segs, rows), + } + try o.add(b.x, b.y, b.x, rows * D); + try glue.rmsnorm(o.set(), o.s, b.x, D, l.mlp_norm, b.normed, D, null, c.eps, rows, D); + switch (l.ffn) { + .dense => |m| try block_ffn.mlp(o, b, m, b.normed, D, rows, 0, b.y), + .moe => |e| try block_ffn.moe(o, b, e, g, l.index, rows, picks), + } + try o.add(b.x, b.y, b.x, rows * D); + } + + /// The final norm and lm_head of b.x's rows into b.logits (bf16 [rows][vocab]). + pub fn head(f: *const Forward, rows: usize) !void { + try f.final(rows); + try f.logitsOf(0, rows); + } + + /// The final norm of b.x's rows into b.normed: the rows lm_head and the MTP layer read. + pub fn final(f: *const Forward, rows: usize) !void { + const o = f.ops(); + const w = f.g.w; + const D: usize = w.config.hidden; + try glue.rmsnorm(o.set(), o.s, f.b.x, D, w.norm, f.b.normed, D, null, w.config.eps, rows, D); + } + + /// lm_head of b.normed's rows row0 .. row0 + n into b.logits from its first row. + pub fn logitsOf(f: *const Forward, row0: usize, n: usize) !void { + const D: usize = f.g.w.config.hidden; + try f.ops().mm(f.g.w.head orelse return error.TiedHeadNotBuilt, f.b.normed + row0 * D * 2, D, n, f.b.logits); + } + + /// Keep the first `keep` rows of the sequence's last forward (a prompt keeps them all): KDA states, position. + pub fn commit(f: *const Forward, st: *state.State, keep: usize) !void { + if (keep == 0 or keep > st.pending or (st.prompt and keep != st.pending)) return error.BadCommit; + if (!st.prompt) try block_kda.commit(f.ops(), st, keep); + st.pending = 0; + try st.setPos(f.s, st.pos + keep); + } +}; diff --git a/zig/src/families/ling/graph.zig b/zig/src/families/ling/graph.zig new file mode 100644 index 0000000000..e0f055b80c --- /dev/null +++ b/zig/src/families/ling/graph.zig @@ -0,0 +1,245 @@ +//! Ling-3.1-flash's TP=1 layer graph on CUDA: the kernels it launches, the launches its blocks share, derived weights. +const std = @import("std"); +const cuda = @import("cuda"); +const glm = @import("glm_cuda"); +const weights = @import("weights.zig"); + +pub const glue = glm.glue; +pub const kda = glm.kda; +pub const latent = glm.latent; +const Proj = weights.Proj; +const Config = @import("config.zig").Config; + +/// Every module the graph launches from: GLM's glue, KDA and ops, the lane formats, the grouped experts, our glue. +pub const Modules = struct { + mods: [5]cuda.Module, // glm_glue, glm_kda, glm_ops, experts (the plan kernels), ling_glue + glue: glue.Kernels, + kda: kda.Kernels, + lane: cuda.qmmf.Lane, + prompt: cuda.nvfp4.Prompt, + x8: cuda.experts.Fp8Experts, + x4: cuda.experts.Nvfp4Experts, + router: cuda.grouped.Router, + copy_rows: cuda.Function, + add_i32: cuda.Function, + add: cuda.Function, + gate: cuda.Function, + combine: cuda.Function, + swiglu: [2]cuda.Function, // glm_glue's 512 and 64 SwiGLU in Ling's order, min(silu(g), L) * clip(u, -L, L) + + pub fn load(d: *const cuda.Driver, ctx: *const cuda.Context) !Modules { + const k = cuda.kernels; + var m: Modules = undefined; + var n: usize = 0; + errdefer for (m.mods[0..n]) |*x| x.unload(); + for (&m.mods, [_][]const u8{ k.glm_glue, k.glm_kda, k.glm_ops, k.experts, k.ling_glue }) |*x, image| { + x.* = try cuda.Module.load(d, image); + n += 1; + } + const sms: usize = @intCast(try ctx.attribute(.multiprocessor_count)); + m.glue = try glue.Kernels.load(m.mods[0]); + m.kda = try kda.Kernels.load(m.mods[1]); + m.copy_rows = try m.mods[4].function("tf_ling_copy_rows"); + m.add_i32 = try m.mods[2].function("tf_glm_add_i32"); + m.router = try cuda.grouped.Router.resolve(m.mods[3]); + m.add = try m.mods[4].function("tf_ling_add"); + m.gate = try m.mods[4].function("tf_ling_gate"); + m.combine = try m.mods[4].function("tf_ling_combine"); + m.swiglu = .{ try m.mods[0].function("tf_ling_swiglu_512"), try m.mods[0].function("tf_ling_swiglu_64") }; + m.lane = try cuda.qmmf.Lane.load(d, @intCast(try ctx.attribute(.compute_capability_major))); + errdefer m.lane.unload(); + m.prompt = try cuda.nvfp4.Prompt.load(d); + errdefer m.prompt.unload(); + m.x8 = try cuda.experts.Fp8Experts.load(d, sms); + errdefer m.x8.unload(); + m.x4 = try cuda.experts.Nvfp4Experts.load(d, sms); + return m; + } + + pub fn deinit(m: *Modules) void { + for (&m.mods) |*x| x.unload(); + m.lane.unload(); + m.prompt.unload(); + m.x8.unload(); + m.x4.unload(); + m.* = undefined; + } + + pub fn set(m: *const Modules) glue.Set { + return .{ .glue = &m.glue }; + } + + /// Every lane format's decode rows; nvfp4 prompt rows on the prompt GEMM. + pub fn linear(m: *const Modules) cuda.qlinear.Linear { + return .{ .lane = &m.lane, .nvfp4_prompt = &m.prompt }; + } + + pub fn grouped(m: *const Modules) cuda.experts.Grouped { + return .{ .fp8g = &m.x8, .nvfp4 = &m.x4 }; + } +}; + +fn cdiv(a: usize, b: usize) u32 { + return @intCast((a + b - 1) / b); +} + +/// One forward's launches: the modules, the stream, prompt or decode kernels, the lane's split-K partials. +pub const Ops = struct { + m: *const Modules, + s: cuda.Stream, + prompt: bool, + part: u64, + kda_wide: bool = true, // decode windows on KDA's three-kernel chain (bit-equal to the fused one, faster) + + pub fn set(o: Ops) glue.Set { + return o.m.set(); + } + + /// out [rows][n] bf16 = x rows (`ldx` apart) through `w` in whatever format it loaded. + pub fn mm(o: Ops, w: Proj, x: u64, ldx: usize, rows: usize, out: u64) !void { + const in: cuda.qlinear.Rows = .{ .x = x, .ldx = ldx, .part = o.part }; + if (o.prompt) return o.m.linear().prompt(o.s, w, in, out, rows); + return o.m.linear().decode(o.s, w, in, out, rows); + } + + /// Bytes [rows][row_bytes] from rows `src_pitch` apart into rows `dst_pitch` apart, as they are. + pub fn copy(o: Ops, src: u64, src_pitch: usize, dst: u64, dst_pitch: usize, rows: usize, row_bytes: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ src, @as(c_longlong, @intCast(src_pitch)), dst, @as(c_longlong, @intCast(dst_pitch)), @as(c_int, @intCast(rows)), @as(c_int, @intCast(row_bytes)) }) |v| a.add(v); + try cuda.launch.launch(o.m.copy_rows, .{ .grid = .{ .x = cdiv(row_bytes, 256 * 16), .y = @intCast(@max(1, @min(rows, 65535))) }, .block = .{ .x = 256 } }, o.s, &a); + } + + /// A projection into columns col0 .. col0 + n of rows `stride` wide (bf16), through `tmp`. + pub fn mmCols(o: Ops, w: Proj, x: u64, ldx: usize, rows: usize, tmp: u64, dst: u64, stride: usize, col0: usize) !void { + const n = w.outputs(); + try o.mm(w, x, ldx, rows, tmp); + try o.copy(tmp, n * 2, dst + col0 * 2, stride * 2, rows, n * 2); + } + + /// out = bf16(x + y), n bf16 elements; out may alias either. + pub fn add(o: Ops, x: u64, y: u64, out: u64, n: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ x, y, out, @as(i64, @intCast(n)) }) |v| a.add(v); + try cuda.launch.launch(o.m.add, .{ .grid = .{ .x = cdiv(n, 256) }, .block = .{ .x = 256 } }, o.s, &a); + } + + /// The head-wise gate: v [rows][heads][dv] times bf16 sigmoid of g [rows][g_stride] column h. + pub fn gate(o: Ops, v: u64, g: u64, g_stride: usize, rows: usize, heads: usize, dv: usize) !void { + var a: cuda.Args = .{}; + inline for (.{ v, @as(i64, @intCast(heads)), g, @as(i64, @intCast(g_stride)) }) |x| a.add(x); + try cuda.launch.launch(o.m.gate, .{ .grid = .{ .x = @intCast(rows), .y = @intCast(heads) }, .block = .{ .x = @intCast(dv) } }, o.s, &a); + } + + /// The device i32 at `pos` plus `r` into `out`. + pub fn addI32(o: Ops, pos: u64, r: usize, out: u64) !void { + var a: cuda.Args = .{}; + inline for (.{ pos, @as(c_int, @intCast(r)), out }) |v| a.add(v); + try cuda.launch.launch(o.m.add_i32, .{ .grid = .{ .x = 1 }, .block = .{ .x = 32 } }, o.s, &a); + } +}; + +/// An MLA layer's kv_b as the absorbed path reads it, bf16: keys wk [heads][nope][lw] and values wv [heads][v][lw]. +pub const Absorb = struct { wk: u64, wv: u64 }; + +/// What the graph derives from the loaded weights: absorbed kv_b, each layer's slot, the rope table, a zero bias. +pub const Graph = struct { + w: *const weights.Weights, + store: weights.Store, + slot: []u32, // loaded layer j's index among the loaded KDA layers, or among the MLA layers + absorb: []Absorb, // by MLA slot + kda_layers: usize, + mla_layers: usize, + freq: u64, // fp32 [32]: the rope's inv_freq + zero_bias: u64, // fp32 [experts]: select's bias without moe_router_enable_expert_bias + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, m: *const Modules, s: cuda.Stream, w: *const weights.Weights) !Graph { + const c = &w.config; + var g: Graph = .{ .w = w, .store = .{ .gpa = gpa, .d = d }, .slot = try gpa.alloc(u32, w.layers.len), .absorb = &.{}, .kda_layers = 0, .mla_layers = 0, .freq = 0, .zero_bias = 0 }; + errdefer g.drop(weights.drains(s)); // a failed drain leaks the store: a queued decode may still use it + for (w.layers, g.slot) |l, *x| switch (l.mixer) { + .kda => { + x.* = @intCast(g.kda_layers); + g.kda_layers += 1; + }, + .mla => { + x.* = @intCast(g.mla_layers); + g.mla_layers += 1; + }, + }; + if (c.rope_dim != latent.rope_width or c.rope_style != .gptj) return error.RopeNotBuilt; + const f = latent.ropeFreqs(c.rope_theta); + g.freq = try g.store.keep(std.mem.sliceAsBytes(&f), "ling.rope.inv_freq"); + const zero = try gpa.alloc(u8, @as(usize, c.experts) * 4); + defer gpa.free(zero); + @memset(zero, 0); + g.zero_bias = try g.store.keep(zero, "ling.zero_bias"); + g.absorb = try gpa.alloc(Absorb, g.mla_layers); + for (w.layers, g.slot) |l, x| if (l.mixer == .mla) { + g.absorb[x] = try absorbOf(gpa, &g.store, m, s, c, l.mixer.mla.kv_b, l.index); + }; + return g; + } + + pub fn deinit(g: *Graph) void { + g.drop(true); + } + + /// deinit, but the store's device buffers leaked unless `settled` (nothing queued can still use them). + pub fn drop(g: *Graph, settled: bool) void { + const gpa = g.store.gpa; + gpa.free(g.slot); + gpa.free(g.absorb); + g.store.drop(settled); + g.* = undefined; + } +}; + +/// `b` freed, after draining `s` if work is `pending` on it; a failed drain leaks it (the device may still read it). +fn freeDrained(s: cuda.Stream, b: *cuda.DeviceBuffer, pending: bool) void { + if (!pending or weights.drains(s)) b.free(); +} + +/// kv_b in its loaded format through the lane on identity rows: W^T [lw][heads (nope + v)], regrouped by head. +fn absorbOf(gpa: std.mem.Allocator, store: *weights.Store, m: *const Modules, s: cuda.Stream, c: *const Config, kv_b: Proj, index: u32) !Absorb { + const lw: usize = c.kv_lora; + const h: usize = c.mla_heads; + const nope: usize = c.nope; + const dv: usize = c.v_dim; + const per = nope + dv; + const n = h * per; + if (kv_b.outputs() != n or kv_b.inputs() != lw) return error.UnexpectedTensor; + const eye = try gpa.alloc(u16, lw * lw); + defer gpa.free(eye); + @memset(eye, 0); + for (0..lw) |i| eye[i * lw + i] = 0x3f80; // bf16 1.0 + var pending = false; // the decode queued and not yet drained + var x = try cuda.DeviceBuffer.fromHost(store.d, std.mem.sliceAsBytes(eye)); + defer freeDrained(s, &x, pending); + var y = try cuda.DeviceBuffer.alloc(store.d, lw * n * 2); + defer freeDrained(s, &y, pending); + var part = try cuda.DeviceBuffer.alloc(store.d, @max(16, m.lane.partBytes(lw, switch (kv_b) { + .fp8g, .nvfp4, .fp8, .bf16 => |q| q, + else => return error.FormatNotBuilt, + }))); + defer freeDrained(s, &part, pending); + pending = true; + try m.linear().decode(s, kv_b, .{ .x = x.ptr, .ldx = lw, .part = part.ptr }, y.ptr, lw); + try s.synchronize(); + pending = false; + const t = try gpa.alloc(u16, lw * n); + defer gpa.free(t); + try y.download(0, std.mem.sliceAsBytes(t)); + const wk = try gpa.alloc(u16, h * nope * lw); + defer gpa.free(wk); + const wv = try gpa.alloc(u16, h * dv * lw); + defer gpa.free(wv); + for (0..h) |hd| for (0..per) |j| for (0..lw) |col| { + const v = t[col * n + hd * per + j]; + if (j < nope) wk[(hd * nope + j) * lw + col] = v else wv[(hd * dv + j - nope) * lw + col] = v; + }; + var b: [64]u8 = undefined; + return .{ + .wk = try store.keep(std.mem.sliceAsBytes(wk), try std.fmt.bufPrint(&b, "ling.{d}.absorb.wk", .{index})), + .wv = try store.keep(std.mem.sliceAsBytes(wv), try std.fmt.bufPrint(&b, "ling.{d}.absorb.wv", .{index})), + }; +} diff --git a/zig/src/families/ling/lanes.zig b/zig/src/families/ling/lanes.zig new file mode 100644 index 0000000000..bbcbc51c1c --- /dev/null +++ b/zig/src/families/ling/lanes.zig @@ -0,0 +1,308 @@ +//! Ling-3.1-flash's CUDA lane backend (as GLM's cuda_lanes.zig, one process): a slot per stream, shared verify. +const std = @import("std"); +const lanes = @import("lanes"); +const engine = @import("engine.zig"); + +const be = lanes.backend; +const Engine = engine.Engine; + +/// First tokens a handle names. +const ring = 1024; + +/// The widest verify window: the pending token and 15 drafts (copies or the MTP layer's chain). +pub const max_rows = engine.window_rows; + +/// A stream's place: its sequence's slot, the rows its last verify left uncommitted, where its prompt pass stands. +const Lane = struct { + slot: u32, + pending_rows: ?usize = null, + held: [engine.max_drafts]u32 = undefined, // the MTP layer's drafts for the next window + held_n: usize = 0, + fill: ?usize = null, // a prompt pass going a chunk at a time: where its next chunk starts +}; + +pub const Lanes = struct { + e: *Engine, + whole: bool = false, // prompt passes run whole: no other stream's round between chunks + alone: bool = false, // each window in its own forward + shared: u64 = 0, // forwards that held several streams' windows (checks read it) + lanes: std.AutoHashMapUnmanaged(*const lanes.Stream, Lane) = .empty, + drawn: [ring]u32 = undefined, + next: u64 = 0, + + pub fn deinit(self: *Lanes) void { + self.lanes.deinit(self.e.gpa); + } + + pub fn backend(self: *Lanes) be.Backend { + return .{ .ptr = self, .vtable = if (self.whole) &whole_vt else &steps_vt }; + } + + const whole_vt: be.Backend.VTable = .{ + .prefill = prefillFn, + .first = firstFn, + .queue = queueFn, + .read = readFn, + .verify = verifyFn, + .keep = keepFn, + .draft = draftFn, + .release = releaseFn, + }; + const steps_vt: be.Backend.VTable = vt: { + var v = whole_vt; + v.prefill_step = prefillStepFn; + break :vt v; + }; + + /// Round loop facts: exact windows of max_rows (copies, MTP chains), shared windows unless `alone`. + pub fn facts(self: *const Lanes) lanes.Model { + const sharing = !self.alone; + const drafting = self.e.mtp != null; + return .{ + .exact_width = max_rows, + .gpu_tokens = false, + .mtp = drafting, + .speculate = drafting, + .drafts = engine.max_drafts, + .speculate_early = false, + .hidden_rows = sharing, + .batch_rows = if (sharing) engine.shared_rows else 0, + .max_streams = if (sharing) engine.max_shared else 1, + }; + } + + fn of(ptr: *anyopaque) *Lanes { + return @ptrCast(@alignCast(ptr)); + } + + fn take(self: *Lanes, token: u32) u64 { + const h = self.next; + self.drawn[h % ring] = token; + self.next += 1; + return h; + } + + /// The lowest slot no stream holds. + fn freeSlot(self: *const Lanes) u32 { + var slot: u32 = 0; + while (true) : (slot += 1) { + var it = self.lanes.valueIterator(); + const taken = while (it.next()) |l| { + if (l.slot == slot) break true; + } else false; + if (!taken) return slot; + } + } + + /// The stream's lane; a new one takes the lowest free slot, chosen before the map holds its (unset) entry. + pub fn claim(self: *Lanes, gpa: std.mem.Allocator, s: *const lanes.Stream) !*Lane { + if (self.lanes.getPtr(s)) |l| return l; + const slot = self.freeSlot(); + try self.lanes.put(gpa, s, .{ .slot = slot }); + return self.lanes.getPtr(s).?; + } + + /// The stream's sequence bound, its last verify committed (`kept` rows, else all of them). + fn bind(self: *Lanes, s: *const lanes.Stream, kept: ?usize) !*Lane { + const l = self.lanes.getPtr(s) orelse return error.UnknownStream; + try self.e.bind(l.slot); + if (l.pending_rows) |rows| try self.e.keep(kept orelse rows); + l.pending_rows = null; + return l; + } + + fn cancelled(ptr: ?*anyopaque) bool { + return @as(*lanes.Stream, @ptrCast(@alignCast(ptr.?))).isCancelled(); + } + + fn cancelOf(s: *lanes.Stream) engine.Cancel { + return .{ .ptr = s, .check = cancelled }; + } + + fn prefillFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!void { + const self = of(ptr); + const from = try self.opening(s); + _ = self.take(try self.e.prefill(s.prompt(), from, cancelOf(s), marksOf(s))); + } + + /// One chunk of the stream's prompt pass (fillOne runs other rounds between); true once the first token is drawn. + fn prefillStepFn(ptr: *anyopaque, s: *lanes.Stream) anyerror!bool { + const self = of(ptr); + const going: ?usize = if (self.lanes.getPtr(s)) |l| l.fill else null; + const at = going orelse try self.e.prefillStart(s.prompt().len, try self.opening(s)); + const l = self.lanes.getPtr(s).?; + try self.e.bind(l.slot); + errdefer l.fill = null; + const step = try self.e.prefillChunk(s.prompt(), at, cancelOf(s), marksOf(s)); + if (step.first) |first| { + _ = self.take(first); + l.fill = null; + return true; + } + l.fill = step.at; + return false; + } + + /// A prompt pass begins: the stream's slot bound and its sequence fresh; where the pass resumes (a kept state). + fn opening(self: *Lanes, s: *lanes.Stream) !?engine.Resume { + const ids = s.prompt(); + if (ids.len + s.max_new + max_rows >= self.e.capacity) return error.PromptTooLong; + const l = try self.claim(self.e.gpa, s); + l.pending_rows = null; + l.held_n = 0; + l.fill = null; + try self.e.bind(l.slot); + try self.e.begin(s.sampling); + s.cached = 0; + return self.resumeFrom(s, ids.len); + } + + fn marksOf(s: *lanes.Stream) ?engine.Marks { + return if (s.reuse.marks.len == 0) null else .{ .at = s.reuse.marks, .ctx = s, .call = keptAt }; + } + + /// The stream's kept state, where its prompt pass starts (null: none, or no longer held). + fn resumeFrom(self: *Lanes, s: *lanes.Stream, len: usize) ?engine.Resume { + const saved = s.reuse.saved orelse return null; + const k: *const engine.Kept = @ptrCast(@alignCast(saved)); + if (k.at != s.reuse.at or k.at == 0 or k.at >= len or !self.e.holds(k.id)) { + s.reuse_failed = true; + return null; + } + s.cached = k.at; + return .{ .id = k.id, .at = k.at }; + } + + /// The prompt pass stands at a mark: the lane host keeps the state there (native.snap_save). + fn keptAt(ctx: *anyopaque, at: u32) void { + const s: *lanes.Stream = @ptrCast(@alignCast(ctx)); + if (s.reuse.hook) |k| k.at(k.ptr, s, at); + } + + fn firstFn(ptr: *anyopaque, s: *lanes.Stream, position: u64) anyerror!u64 { + const self = of(ptr); + if (position != s.prompt_len) return error.PositionMismatch; + return self.next - 1; + } + + fn queueFn(ptr: *anyopaque, s: *lanes.Stream, feed: be.Feed, position: u64) anyerror!u64 { + _ = ptr; + _ = s; + _ = feed; + _ = position; + return error.NotPipelined; + } + + fn readFn(ptr: *anyopaque, handle: u64) anyerror!u32 { + const self = of(ptr); + if (handle >= self.next or self.next - handle > ring) return error.NoSuchToken; + return self.drawn[handle % ring]; + } + + /// A window's ids (pending token, then its drafts) checked against the bound sequence's position. + fn windowIds(self: *Lanes, l: *const Lane, w: be.Window, ids: *[max_rows]u32) ![]const u32 { + if (w.parents != null) return error.TreesNotBuilt; + const rows = w.rows(); + if (rows > max_rows) return error.WindowTooWide; + if (w.held > l.held_n) return error.NoHeldDrafts; + const pos = self.e.pos(); + for (w.positions, 0..) |p, r| if (p != pos + 1 + r) return error.PositionMismatch; + const drafts = if (w.held > 0) l.held[0..w.held] else w.tokens; + ids[0] = w.pending; + @memcpy(ids[1..][0..drafts.len], drafts); + return ids[0..rows]; + } + + /// Each window in its own forward on its stream's sequence (the core sends one when no forward is shared). + fn verifyFn(ptr: *anyopaque, windows: []const be.Window, out: []be.Verified) anyerror!void { + const self = of(ptr); + if (windows.len > 1) return self.verifyShared(windows, out); + for (windows, out) |w, *o| { + const l = try self.bind(w.stream, null); + var ids: [max_rows]u32 = undefined; + const got = try self.windowIds(l, w, &ids); + try self.e.verify(got, o.sampled[0..got.len]); + @memcpy(o.drafts[0 .. got.len - 1], got[1..]); + l.pending_rows = got.len; + } + } + + /// A round's windows in one forward (Engine.verifyShared), each stream's rows on its own sequence. + fn verifyShared(self: *Lanes, windows: []const be.Window, out: []be.Verified) !void { + if (windows.len > engine.max_shared) return error.WindowTooWide; + var parts: [engine.max_shared]engine.Part = undefined; + var ids: [engine.max_shared][max_rows]u32 = undefined; + var lane: [engine.max_shared]*Lane = undefined; + for (windows, out[0..windows.len], 0..) |w, o, k| { + lane[k] = try self.bind(w.stream, null); + const got = try self.windowIds(lane[k], w, &ids[k]); + parts[k] = .{ .slot = lane[k].slot, .ids = got, .out = o.sampled[0..got.len] }; + } + try self.e.verifyShared(parts[0..windows.len]); + self.shared += 1; + for (out[0..windows.len], parts[0..windows.len], lane[0..windows.len]) |*o, p, l| { + @memcpy(o.drafts[0 .. p.ids.len - 1], p.ids[1..]); + l.pending_rows = p.ids.len; + } + } + + fn keepFn(ptr: *anyopaque, windows: []const be.Window, paths: []const []const u32) anyerror!void { + const self = of(ptr); + for (windows, paths) |w, path| { + if (path.len == 0) return error.EmptyPath; + for (path, 0..) |r, i| if (r != i) return error.TreesNotBuilt; + _ = try self.bind(w.stream, path.len); + } + } + + /// The MTP layer absorbs each stream's kept rows (its last verify committed first) and chains its drafts. + fn draftFn(ptr: *anyopaque, requests: []const be.DraftRequest) anyerror!void { + const self = of(ptr); + if (self.e.mtp == null) return error.NoDraftHead; + for (requests) |r| { + const l = try self.bind(r.stream, if (r.rows) |rows| rows.len else null); + if (r.position != self.e.pos() + 1) return error.PositionMismatch; + if (r.lanes != null) return error.TreesNotBuilt; + const depth = @min(r.depth, engine.max_drafts); + if (r.rows) |rows| { + for (rows, 0..) |row, i| if (row != i) return error.TreesNotBuilt; + try self.e.draft(false, r.follow[0..rows.len], depth, &l.held); + } else { + const first = switch (r.first orelse return error.NoFirstToken) { + .handle => |h| self.drawn[h % ring], + .value => |v| v, + }; + try self.e.draft(true, &.{first}, depth, &l.held); + } + l.held_n = depth; + } + } + + /// The stream left the rounds: its sequence freed (slot 0 stays for the next stream). + fn releaseFn(ptr: *anyopaque, s: *lanes.Stream) void { + const self = of(ptr); + const kv = self.lanes.fetchRemove(s) orelse return; + self.e.release(kv.value.slot); + } +}; + +/// A stand-in stream key: claim hashes the address and never reads through it. +fn key(i: usize) *const lanes.Stream { + return @ptrFromInt(i * 4096); +} + +/// Slots claimed over map buckets still holding slot 0 from cleared entries (a release build leaves a new value as is). +pub fn slotCheck(gpa: std.mem.Allocator) !void { + var ll: Lanes = .{ .e = undefined }; + defer ll.lanes.deinit(gpa); + for (1..9) |i| try ll.lanes.put(gpa, key(i), .{ .slot = 0 }); + ll.lanes.clearRetainingCapacity(); + const want = [_]struct { k: usize, slot: u32 }{ .{ .k = 1, .slot = 0 }, .{ .k = 2, .slot = 1 }, .{ .k = 1, .slot = 0 }, .{ .k = 3, .slot = 2 } }; + for (want) |w| if ((try ll.claim(gpa, key(w.k))).slot != w.slot) return error.SlotTaken; + if (!ll.lanes.remove(key(1))) return error.SlotTaken; + if ((try ll.claim(gpa, key(4))).slot != 0 or (try ll.claim(gpa, key(5))).slot != 3) return error.SlotTaken; +} + +test "a new stream takes the lowest free slot, a known one keeps its own, a released one's goes to the next" { + try slotCheck(std.testing.allocator); +} diff --git a/zig/src/families/ling/mtp.zig b/zig/src/families/ling/mtp.zig new file mode 100644 index 0000000000..a94a4a7d83 --- /dev/null +++ b/zig/src/families/ling/mtp.zig @@ -0,0 +1,70 @@ +//! Ling-3.1's MTP layer (model.layers.72) as the family's drafter: enorm | hnorm through eh_proj, MLA + MoE, its norm. +const std = @import("std"); +const cuda = @import("cuda"); +const weights = @import("weights.zig"); +const fw = @import("forward.zig"); + +const glue = @import("graph.zig").glue; + +/// The most drafts one chain holds. +pub const max_drafts = 3; + +/// The MTP layer's graph over a view of the weights (its layer alone, its final norm as `norm`) and its scratch. +pub const Head = struct { + view: weights.Weights, + mtp: weights.Mtp, + g: fw.Graph, + bm: fw.Buffers, // chained draft rows and a verify's kept rows + cat: cuda.DeviceBuffer, // bf16 [rows][2 hidden]: enorm(embed(next)) | hnorm(hidden) + + /// In place (the graph points at `view`); `w` holds the MTP layer last; cat rows for `rows` (the prompt chunk). + pub fn init(h: *Head, gpa: std.mem.Allocator, d: *const cuda.Driver, m: *const fw.Modules, s: cuda.Stream, w: *const weights.Weights, rows: usize, capacity: usize) !void { + const c = &w.config; + const mtp = w.mtp orelse return error.NoMtpLayer; + const last = w.layers[w.layers.len - 1]; + if (last.index != c.layers or last.mixer != .mla or last.ffn != .moe) return error.NoMtpLayer; + h.mtp = mtp; + h.view = w.*; + h.view.layers = w.layers[w.layers.len - 1 ..]; + h.view.norm = mtp.norm; // final_layernorm before the shared lm_head + h.g = try fw.Graph.init(gpa, d, m, s, &h.view); + errdefer h.g.deinit(); + h.bm = try fw.Buffers.init(d, m, w, fw.window_rows, false, capacity); + errdefer h.bm.deinit(); + h.cat = try cuda.DeviceBuffer.alloc(d, @max(rows, fw.window_rows) * 2 * @as(usize, c.hidden) * 2); + } + + pub fn deinit(h: *Head) void { + h.cat.free(); + h.bm.deinit(); + h.g.deinit(); + h.* = undefined; + } + + pub fn forward(h: *const Head, m: *const fw.Modules, s: cuda.Stream, b: *const fw.Buffers) fw.Forward { + return .{ .m = m, .g = &h.g, .s = s, .b = b }; + } + + /// The layer on b.x's `rows` embedding rows and `hidden`'s rows (bf16 [rows][hidden]), committed on `hd`. + pub fn step(h: *const Head, f: fw.Forward, hd: *fw.State, rows: usize, hidden: u64, picks: ?u64) !void { + const b = f.b; + const c = &h.view.config; + const D: usize = c.hidden; + const o = f.ops(); + if (rows == 0 or rows > b.rows or hd.pending != 0) return error.BadRows; + try glue.rmsnorm(o.set(), o.s, b.x, D, h.mtp.enorm, h.cat.ptr, 2 * D, null, c.eps, rows, D); + try glue.rmsnorm(o.set(), o.s, hidden, D, h.mtp.hnorm, h.cat.ptr + D * 2, 2 * D, null, c.eps, rows, D); + try o.mm(h.mtp.eh, h.cat.ptr, 2 * D, rows, b.x); + try f.layer(&.{.{ .st = hd, .row0 = 0, .rows = rows }}, 0, rows, picks); + hd.pending = rows; + hd.prompt = b.prompt; + try f.commit(hd, rows); + } + + /// `next` (the token after each row) with `hidden`'s rows through the layer on `b`, committed on `hd`. + pub fn absorb(h: *const Head, m: *const fw.Modules, s: cuda.Stream, b: *const fw.Buffers, hd: *fw.State, next: []const i32, hidden: u64) !void { + const f = h.forward(m, s, b); + try f.embed(next); + try h.step(f, hd, next.len, hidden, null); + } +}; diff --git a/zig/src/families/ling/names.zig b/zig/src/families/ling/names.zig new file mode 100644 index 0000000000..f59d867ba7 --- /dev/null +++ b/zig/src/families/ling/names.zig @@ -0,0 +1,294 @@ +//! Ling-3.1-flash's weight map: each tensor name the config implies, by role, checked against an index. +const std = @import("std"); +const config = @import("config.zig"); + +const Config = config.Config; + +pub const Role = enum { + embed, + final_norm, + lm_head, + attn_norm, + mlp_norm, + kda_q, + kda_k, + kda_v, + kda_q_conv, + kda_k_conv, + kda_v_conv, + kda_b, + kda_f, + kda_g, + kda_a_log, + kda_dt_bias, + kda_o_norm, + kda_o, + mla_q, + mla_q_a, + mla_q_a_norm, + mla_q_b, + mla_kv_a, + mla_kv_a_norm, + mla_kv_b, + mla_gate, + mla_o, + dense_gate, + dense_up, + dense_down, + router, + router_bias, + expert_gate, + expert_up, + expert_down, + shared_gate, + shared_up, + shared_down, + mtp_enorm, + mtp_hnorm, + mtp_eh, + mtp_final_norm, +}; + +/// One stored tensor: its role, layer (null for the model's own), routed expert, and whether it is a scale tensor. +pub const Tensor = struct { role: Role, layer: ?u32 = null, expert: ?u32 = null, scale: bool = false }; + +/// linear: `.weight` plus its format's side tensors (side()); weight: `.weight` only; bare: the module's own name. +const Form = enum { linear, weight, bare }; + +const Part = struct { role: Role, path: []const u8, form: Form }; + +const norms = [_]Part{ + .{ .role = .attn_norm, .path = "input_layernorm", .form = .weight }, + .{ .role = .mlp_norm, .path = "post_attention_layernorm", .form = .weight }, +}; + +const kda = [_]Part{ + .{ .role = .kda_q, .path = "attention.q_proj", .form = .linear }, + .{ .role = .kda_k, .path = "attention.k_proj", .form = .linear }, + .{ .role = .kda_v, .path = "attention.v_proj", .form = .linear }, + .{ .role = .kda_q_conv, .path = "attention.q_conv1d", .form = .weight }, + .{ .role = .kda_k_conv, .path = "attention.k_conv1d", .form = .weight }, + .{ .role = .kda_v_conv, .path = "attention.v_conv1d", .form = .weight }, + .{ .role = .kda_b, .path = "attention.b_proj", .form = .linear }, + .{ .role = .kda_f, .path = "attention.f_proj", .form = .linear }, + .{ .role = .kda_g, .path = "attention.g_proj", .form = .linear }, + .{ .role = .kda_a_log, .path = "attention.A_log", .form = .bare }, + .{ .role = .kda_dt_bias, .path = "attention.dt_bias", .form = .bare }, + .{ .role = .kda_o_norm, .path = "attention.o_norm", .form = .weight }, + .{ .role = .kda_o, .path = "attention.o_proj", .form = .linear }, +}; + +const mla_q_lora = [_]Part{ + .{ .role = .mla_q_a, .path = "attention.q_a_proj", .form = .linear }, + .{ .role = .mla_q_a_norm, .path = "attention.q_a_layernorm", .form = .weight }, + .{ .role = .mla_q_b, .path = "attention.q_b_proj", .form = .linear }, +}; + +const mla_q_whole = [_]Part{.{ .role = .mla_q, .path = "attention.q_proj", .form = .linear }}; + +const mla = [_]Part{ + .{ .role = .mla_kv_a, .path = "attention.kv_a_proj_with_mqa", .form = .linear }, + .{ .role = .mla_kv_a_norm, .path = "attention.kv_a_layernorm", .form = .weight }, + .{ .role = .mla_kv_b, .path = "attention.kv_b_proj", .form = .linear }, + .{ .role = .mla_o, .path = "attention.dense", .form = .linear }, +}; + +const mla_gate = [_]Part{.{ .role = .mla_gate, .path = "attention.g_proj", .form = .linear }}; + +const dense = [_]Part{ + .{ .role = .dense_gate, .path = "mlp.gate_proj", .form = .linear }, + .{ .role = .dense_up, .path = "mlp.up_proj", .form = .linear }, + .{ .role = .dense_down, .path = "mlp.down_proj", .form = .linear }, +}; + +const router = [_]Part{.{ .role = .router, .path = "mlp.gate", .form = .linear }}; +const router_bias = [_]Part{.{ .role = .router_bias, .path = "mlp.gate.expert_bias", .form = .bare }}; + +const shared = [_]Part{ + .{ .role = .shared_gate, .path = "mlp.shared_experts.gate_proj", .form = .linear }, + .{ .role = .shared_up, .path = "mlp.shared_experts.up_proj", .form = .linear }, + .{ .role = .shared_down, .path = "mlp.shared_experts.down_proj", .form = .linear }, +}; + +const expert = [_]Part{ + .{ .role = .expert_gate, .path = "gate_proj", .form = .linear }, + .{ .role = .expert_up, .path = "up_proj", .form = .linear }, + .{ .role = .expert_down, .path = "down_proj", .form = .linear }, +}; + +const mtp = [_]Part{ + .{ .role = .mtp_enorm, .path = "enorm", .form = .weight }, + .{ .role = .mtp_hnorm, .path = "hnorm", .form = .weight }, + .{ .role = .mtp_eh, .path = "eh_proj", .form = .linear }, + .{ .role = .mtp_final_norm, .path = "final_layernorm", .form = .weight }, +}; + +/// The part's tensors under `prefix`: the name built in `buf`, handed to `sink.put(name, tensor)`. +fn emit(c: *const Config, sink: anytype, prefix: []const u8, p: Part, t: Tensor) !void { + var buf: [256]u8 = undefined; + const module = try std.fmt.bufPrint(&buf, "{s}{s}", .{ prefix, p.path }); + var name: [256]u8 = undefined; + var u = t; + u.role = p.role; + switch (p.form) { + .bare => try sink.put(module, u), + .weight, .linear => try sink.put(try std.fmt.bufPrint(&name, "{s}.weight", .{module}), u), + } + if (p.form != .linear) return; + u.scale = true; + for (side(c.format(module))) |tail| try sink.put(try std.fmt.bufPrint(&name, "{s}.{s}", .{ module, tail }), u); +} + +/// The tensors a linear stores beside `.weight`: FP8's block scales, ModelOpt NVFP4's group, global and input scales. +pub fn side(f: config.Quant) []const []const u8 { + return switch (f) { + .bf16 => &.{}, + .fp8_block128 => &.{"weight_scale_inv"}, + .nvfp4 => &.{ "weight_scale", "weight_scale_2", "input_scale" }, + }; +} + +/// What the loader reads for a role: a linear [n, k] (`.weight` and its side tensors), `values` values, a conv. +pub const Dims = union(enum) { linear: [2]usize, values: usize, conv: [3]usize }; + +/// The loader's shape for each role (weights.zig reads every tensor through it). +pub fn dims(c: *const Config, role: Role) Dims { + const d: usize = c.hidden; + const n: usize = @as(usize, c.kda_heads) * c.kda_dim; + const h: usize = c.mla_heads; + return switch (role) { + .embed, .lm_head => .{ .linear = .{ c.vocab, d } }, + .final_norm, .attn_norm, .mlp_norm, .mtp_enorm, .mtp_hnorm, .mtp_final_norm => .{ .values = d }, + .kda_q, .kda_k, .kda_v, .kda_f, .kda_g => .{ .linear = .{ n, d } }, + .kda_b => .{ .linear = .{ c.kda_heads, d } }, + .kda_o => .{ .linear = .{ d, n } }, + .kda_q_conv, .kda_k_conv, .kda_v_conv => .{ .conv = .{ n, 1, c.conv } }, + .kda_a_log => .{ .values = c.kda_heads }, + .kda_dt_bias => .{ .values = n }, + .kda_o_norm => .{ .values = c.kda_dim }, + .mla_q => .{ .linear = .{ h * c.qkDim(), d } }, + .mla_q_a => .{ .linear = .{ c.q_lora, d } }, + .mla_q_a_norm => .{ .values = c.q_lora }, + .mla_q_b => .{ .linear = .{ h * c.qkDim(), c.q_lora } }, + .mla_kv_a => .{ .linear = .{ c.kv_lora + c.rope_dim, d } }, + .mla_kv_a_norm => .{ .values = c.kv_lora }, + .mla_kv_b => .{ .linear = .{ h * (c.nope + c.v_dim), c.kv_lora } }, + .mla_gate => .{ .linear = .{ h, d } }, + .mla_o => .{ .linear = .{ d, h * c.v_dim } }, + .dense_gate, .dense_up => .{ .linear = .{ c.dense_inter, d } }, + .dense_down => .{ .linear = .{ d, c.dense_inter } }, + .router => .{ .linear = .{ c.experts, d } }, + .router_bias => .{ .values = c.experts }, + .expert_gate, .expert_up => .{ .linear = .{ c.moe_inter, d } }, + .expert_down => .{ .linear = .{ d, c.moe_inter } }, + .shared_gate, .shared_up => .{ .linear = .{ c.shared_inter, d } }, + .shared_down => .{ .linear = .{ d, c.shared_inter } }, + .mtp_eh => .{ .linear = .{ d, 2 * d } }, + }; +} + +/// A role's value count: a linear's n k, a small tensor's values, a conv's n taps. +pub fn numel(c: *const Config, role: Role) usize { + return switch (dims(c, role)) { + .linear => |x| x[0] * x[1], + .values => |x| x, + .conv => |x| x[0] * x[1] * x[2], + }; +} + +const tables = [_][]const Part{ &norms, &kda, &mla_q_lora, &mla_q_whole, &mla, &mla_gate, &dense, &router, &router_bias, &shared, &expert, &mtp }; + +/// Role's module in layer i (`e` for routed experts) as the map names it: the name before `.weight`. +pub fn moduleName(buf: []u8, i: u32, role: Role, e: ?u32) ![]const u8 { + switch (role) { + .embed => return "model.word_embeddings", + .final_norm => return "model.norm", + .lm_head => return "lm_head", + .expert_gate, .expert_up, .expert_down => { + const p = for (expert) |x| (if (x.role == role) break x) else unreachable; + return std.fmt.bufPrint(buf, "model.layers.{d}.mlp.experts.{d}.{s}", .{ i, e orelse return error.NoExpert, p.path }); + }, + else => {}, + } + for (tables) |t| for (t) |p| if (p.role == role) return std.fmt.bufPrint(buf, "model.layers.{d}.{s}", .{ i, p.path }); + unreachable; +} + +fn emitAll(c: *const Config, sink: anytype, prefix: []const u8, parts: []const Part, t: Tensor) !void { + for (parts) |p| try emit(c, sink, prefix, p, t); +} + +/// Layer i's tensors (an MTP layer when i >= c.layers): norms, the mixer its kind names, the MLP or MoE. +pub fn layer(c: *const Config, i: u32, sink: anytype) !void { + var buf: [64]u8 = undefined; + const prefix = try std.fmt.bufPrint(&buf, "model.layers.{d}.", .{i}); + const t: Tensor = .{ .role = .attn_norm, .layer = i }; + try emitAll(c, sink, prefix, &norms, t); + switch (c.kind(i)) { + .kda => try emitAll(c, sink, prefix, &kda, t), + .mla => { + try emitAll(c, sink, prefix, if (c.q_lora > 0) &mla_q_lora else &mla_q_whole, t); + try emitAll(c, sink, prefix, &mla, t); + if (c.head_gate) try emitAll(c, sink, prefix, &mla_gate, t); + }, + } + if (!c.isMoe(i)) return emitAll(c, sink, prefix, &dense, t); + try emitAll(c, sink, prefix, &router, t); + if (c.expert_bias) try emitAll(c, sink, prefix, &router_bias, t); + if (c.shared_inter > 0) try emitAll(c, sink, prefix, &shared, t); + if (i >= c.layers) try emitAll(c, sink, prefix, &mtp, t); + var ebuf: [96]u8 = undefined; + for (0..c.experts) |e| { + const ep = try std.fmt.bufPrint(&ebuf, "{s}mlp.experts.{d}.", .{ prefix, e }); + try emitAll(c, sink, ep, &expert, .{ .role = .expert_gate, .layer = i, .expert = @intCast(e) }); + } +} + +/// The model's own tensors: the embeddings, the final norm and, unless tied, the head. +pub fn globals(c: *const Config, sink: anytype) !void { + try emit(c, sink, "model.", .{ .role = .embed, .path = "word_embeddings", .form = .weight }, .{ .role = .embed }); + try emit(c, sink, "model.", .{ .role = .final_norm, .path = "norm", .form = .weight }, .{ .role = .final_norm }); + if (!c.tie) try emit(c, sink, "", .{ .role = .lm_head, .path = "lm_head", .form = .linear }, .{ .role = .lm_head }); +} + +/// Every tensor the config implies, backbone then MTP layers. +pub fn all(c: *const Config, sink: anytype) !void { + try globals(c, sink); + for (0..c.layers + c.mtp) |i| try layer(c, @intCast(i), sink); +} + +/// The map against an index's names: expected tensors (MTP ones counted apart), missing ones and unaccounted ones. +pub const Report = struct { + expected: usize = 0, + mtp: usize = 0, + missing: std.ArrayList([]const u8) = .empty, + unknown: std.ArrayList([]const u8) = .empty, +}; + +const Checker = struct { + arena: std.mem.Allocator, + seen: *std.StringHashMapUnmanaged(bool), + c: *const Config, + r: *Report, + + pub fn put(k: *Checker, name: []const u8, t: Tensor) !void { + k.r.expected += 1; + if (t.layer != null and t.layer.? >= k.c.layers) k.r.mtp += 1; + if (k.seen.getPtr(name)) |hit| { + if (hit.*) return error.DuplicateName; + hit.* = true; + } else try k.r.missing.append(k.arena, try k.arena.dupe(u8, name)); + } +}; + +/// Every name the map expects is in `index` and every index name is one the map expects (lists in `arena`). +pub fn check(arena: std.mem.Allocator, c: *const Config, index: []const []const u8) !Report { + var seen: std.StringHashMapUnmanaged(bool) = .empty; + try seen.ensureTotalCapacity(arena, @intCast(index.len)); + for (index) |n| if (seen.fetchPutAssumeCapacity(n, false) != null) return error.DuplicateName; + var r: Report = .{}; + var k: Checker = .{ .arena = arena, .seen = &seen, .c = c, .r = &r }; + try all(c, &k); + for (index) |n| if (!seen.get(n).?) try r.unknown.append(arena, n); + return r; +} diff --git a/zig/src/families/ling/names_test.zig b/zig/src/families/ling/names_test.zig new file mode 100644 index 0000000000..17aef45f9b --- /dev/null +++ b/zig/src/families/ling/names_test.zig @@ -0,0 +1,164 @@ +//! The weight map against the checkpoint's index: fixtures/names.txt folds its 214,025 names. +const std = @import("std"); +const config = @import("config.zig"); +const names = @import("names.zig"); + +const t = std.testing; +const real = @embedFile("fixtures/config.json"); +const folded = @embedFile("fixtures/names.txt"); + +fn load() !config.Config { + var why: config.Why = .{}; + return config.parse(t.allocator, real, &why); +} + +fn lessThan(_: void, a: []const u8, b: []const u8) bool { + return std.mem.lessThan(u8, a, b); +} + +/// Each `{L}` line over its layer ranges, each `{E}` line over its expert count; the header holds count and sha256. +fn unfold(arena: std.mem.Allocator) ![]const []const u8 { + var out: std.ArrayList([]const u8) = .empty; + var lines = std.mem.tokenizeScalar(u8, folded, '\n'); + var head = std.mem.tokenizeScalar(u8, lines.next().?, ' '); + if (!std.mem.eql(u8, head.next().?, "ling31-names")) return error.BadFixture; + const count = try std.fmt.parseInt(usize, head.next().?, 10); + const digest = head.next().?; + while (lines.next()) |line| { + var f = std.mem.tokenizeScalar(u8, line, ' '); + const pat = f.next().?; + const layers = f.next() orelse { + try out.append(arena, pat); + continue; + }; + const experts = if (f.next()) |e| try std.fmt.parseInt(u32, e, 10) else 1; + var spans = std.mem.tokenizeScalar(u8, layers, ','); + while (spans.next()) |span| { + const dash = std.mem.indexOfScalar(u8, span, '-'); + const lo = try std.fmt.parseInt(u32, if (dash) |d| span[0..d] else span, 10); + const hi = if (dash) |d| try std.fmt.parseInt(u32, span[d + 1 ..], 10) else lo; + for (lo..hi + 1) |l| for (0..experts) |e| { + var lb: [8]u8 = undefined; + var eb: [8]u8 = undefined; + const a = try std.mem.replaceOwned(u8, arena, pat, "{L}", try std.fmt.bufPrint(&lb, "{d}", .{l})); + try out.append(arena, try std.mem.replaceOwned(u8, arena, a, "{E}", try std.fmt.bufPrint(&eb, "{d}", .{e}))); + }; + } + } + std.mem.sort([]const u8, out.items, {}, lessThan); + var h = std.crypto.hash.sha2.Sha256.init(.{}); + for (out.items, 0..) |n, i| { + if (i > 0) h.update("\n"); + h.update(n); + } + var d: [32]u8 = undefined; + h.final(&d); + if (out.items.len != count or !std.mem.eql(u8, &std.fmt.bytesToHex(d, .lower), digest)) return error.FixtureNotIndex; + return out.items; +} + +test "the folded fixture unfolds to the index's exact name list (count and sha256 of the sorted names)" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + try t.expectEqual(@as(usize, 214025), (try unfold(arena.allocator())).len); +} + +test "the map accounts for every index name and finds each it expects: 210,925 backbone, 3,100 MTP layer 72" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const c = try load(); + const r = try names.check(arena.allocator(), &c, try unfold(arena.allocator())); + for (r.missing.items) |n| std.debug.print("missing {s}\n", .{n}); + for (r.unknown.items) |n| std.debug.print("unaccounted {s}\n", .{n}); + try t.expectEqual(@as(usize, 0), r.missing.items.len); + try t.expectEqual(@as(usize, 0), r.unknown.items.len); + try t.expectEqual(@as(usize, 214025), r.expected); + try t.expectEqual(@as(usize, 3100), r.mtp); +} + +const Count = struct { + n: usize = 0, + scales: usize = 0, + roles: std.EnumSet(names.Role) = .empty, + + pub fn put(k: *Count, name: []const u8, x: names.Tensor) !void { + _ = name; + k.n += 1; + k.scales += @intFromBool(x.scale); + k.roles.insert(x.role); + } +}; + +test "per layer: KDA dense 0, KDA MoE 4, MLA MoE 7, the MTP layer 72" { + const c = try load(); + var k0: Count = .{}; + try names.layer(&c, 0, &k0); + try t.expectEqual(@as(usize, 2 + 13 + 6 + 6), k0.n); // norms, KDA parts, their 6 FP8 scales (b_proj is bf16), dense MLP + try t.expect(k0.roles.contains(.kda_a_log) and k0.roles.contains(.dense_down) and !k0.roles.contains(.router)); + var k4: Count = .{}; + try names.layer(&c, 4, &k4); + try t.expectEqual(@as(usize, 2 + 19 + 1 + 1 + 6 + 512 * 6), k4.n); // the router is bf16 and has a bias + var k7: Count = .{}; + try names.layer(&c, 7, &k7); + try t.expectEqual(@as(usize, 2 + 3 + 2 + 4 + 3 + 1 + 2 + 6 + 512 * 6), k7.n); // g_proj of 7 is bf16 + try t.expect(k7.roles.contains(.mla_gate) and !k7.roles.contains(.kda_q)); + var k72: Count = .{}; + try names.layer(&c, 72, &k72); + try t.expectEqual(@as(usize, 3100), k72.n); + try t.expect(k72.roles.contains(.mtp_eh) and k72.roles.contains(.mtp_final_norm)); +} + +test "a config that disagrees with the index is caught: the router taken as FP8 misses 69 scales" { + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const json = try std.mem.replaceOwned(u8, t.allocator, real, "\"mlp.gate\",", "\"mlp.gate_x\","); + defer t.allocator.free(json); + var why: config.Why = .{}; + const c = try config.parse(t.allocator, json, &why); + const r = try names.check(arena.allocator(), &c, try unfold(arena.allocator())); + try t.expectEqual(@as(usize, 69), r.missing.items.len); + try t.expectEqualStrings("model.layers.4.mlp.gate.weight_scale_inv", r.missing.items[0]); + try t.expectEqual(@as(usize, 0), r.unknown.items.len); +} + +test "tied embeddings drop lm_head, bf16 drops every scale, no head gate drops MLA's g_proj" { + var c = try load(); + c.tie = true; + c.quant = .bf16; + c.head_gate = false; + var g: Count = .{}; + try names.globals(&c, &g); + try t.expectEqual(@as(usize, 2), g.n); + var k7: Count = .{}; + try names.layer(&c, 7, &k7); + try t.expectEqual(@as(usize, 0), k7.scales); + try t.expect(!k7.roles.contains(.mla_gate)); +} + +test "the real index (TF_LING_INDEX=) against the map" { + const path = std.mem.span(std.c.getenv("TF_LING_INDEX") orelse return error.SkipZigTest); + var arena = std.heap.ArenaAllocator.init(t.allocator); + defer arena.deinit(); + const a = arena.allocator(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(t.io, path, a, .unlimited); + const doc = try std.json.parseFromSliceLeaky(std.json.Value, a, bytes, .{}); + const map = doc.object.get("weight_map").?.object; + const c = try load(); + const r = try names.check(a, &c, map.keys()); + try t.expectEqual(@as(usize, 0), r.missing.items.len + r.unknown.items.len); + try t.expectEqual(@as(usize, map.count()), r.expected); +} + +test "NVFP4 linears store weight, weight_scale, weight_scale_2 and input_scale; the module names the loader reads" { + var c = try load(); + c.quant = .nvfp4; + var k0: Count = .{}; + try names.layer(&c, 0, &k0); + try t.expectEqual(@as(usize, 2 + 13 + 6 * 3 + 3 + 3 * 3), k0.n); // norms, KDA, 3 side tensors a quantized linear, dense + var b: [192]u8 = undefined; + try t.expectEqualStrings("model.layers.4.mlp.experts.511.down_proj", try names.moduleName(&b, 4, .expert_down, 511)); + try t.expectEqualStrings("model.layers.7.attention.dense", try names.moduleName(&b, 7, .mla_o, null)); + try t.expectEqualStrings("model.layers.0.attention.A_log", try names.moduleName(&b, 0, .kda_a_log, null)); + try t.expectEqualStrings("model.layers.72.eh_proj", try names.moduleName(&b, 72, .mtp_eh, null)); + try t.expectEqualStrings("lm_head", try names.moduleName(&b, 0, .lm_head, null)); +} diff --git a/zig/src/families/ling/native.zig b/zig/src/families/ling/native.zig new file mode 100644 index 0000000000..bb8bdfe31c --- /dev/null +++ b/zig/src/families/ling/native.zig @@ -0,0 +1,115 @@ +//! Ling-3.1-flash on the native CUDA server at TP=1: the engine, its lane backend, the prompt cache's copies. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const engine = @import("engine.zig"); +const ling_lanes = @import("lanes.zig"); + +pub const model_type = "bailing_hybrid"; +pub const formats: []const []const u8 = &.{ "bf16", "fp8-block", "modelopt-nvfp4" }; +pub const default_context: i64 = 8192; +pub const max_segments: u32 = 1; +pub const prompt_rows: u32 = engine.prompt_rows; + +pub const Options = struct { context: usize, drafts: bool, segments: usize = 1, slide: bool = false }; + +pub const LoneRun = *const fn (ctx: *anyopaque, s: *lanes.Stream, hooks: *anyopaque, committed: *const fn (*anyopaque) void, yield: *const fn (*anyopaque) bool, resumed: bool) anyerror!bool; + +pub const Loaded = struct { + backend: lanes.backend.Backend, + facts: lanes.Model, + rows: u32, + stream_bytes: usize, + free_streams: u32 = 1, + learn_bytes: usize = 0, + ctx: *anyopaque, + deinit: *const fn (*anyopaque) void, + lone: ?LoneRun = null, + target: *engine.Engine, // the prompt cache keeps this engine's states (snap_*) +}; + +/// Positions a sequence holds: `context` plus a verify window, a multiple of 64. +pub fn capacity(context: usize) usize { + return std.mem.alignForward(usize, context + engine.window_rows, 64); +} + +const Owned = struct { gpa: std.mem.Allocator, e: *engine.Engine, lanes: ling_lanes.Lanes }; + +/// TF_LING_PROMPT_STEPS=0 passes prompts whole, TF_LING_SHARED=0 runs windows alone. +fn off(name: [*:0]const u8) bool { + return if (std.c.getenv(name)) |v| std.mem.eql(u8, std.mem.span(v), "0") else false; +} + +pub fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, kernels: ?[]const u8, o: Options) !Loaded { + _ = kernels; // our own kernels: a captured set is not used + const own = try gpa.create(Owned); + errdefer gpa.destroy(own); + own.gpa = gpa; + own.e = try engine.Engine.init(gpa, io, ctx.d, ctx, dir, .{ .capacity = capacity(o.context), .mtp = o.drafts }); + errdefer own.e.deinit(); + own.lanes = .{ .e = own.e, .whole = off("TF_LING_PROMPT_STEPS"), .alone = off("TF_LING_SHARED") }; + return .{ + .backend = own.lanes.backend(), + .facts = own.lanes.facts(), + .rows = ling_lanes.max_rows, + .stream_bytes = own.e.seqBytes(), + .ctx = own, + .deinit = close, + .target = own.e, + }; +} + +fn close(ptr: *anyopaque) void { + const own: *Owned = @ptrCast(@alignCast(ptr)); + own.lanes.deinit(); + own.e.deinit(); + own.gpa.destroy(own); +} + +/// Bytes a state takes at any position: the sequence's arena. +pub fn snap_bytes(ptr: *anyopaque, at: u32) u64 { + _ = at; + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + return e.stateBytes(); +} + +/// The prompt pass stands at `at` (a mark): copy the state there under a new id. +pub fn snap_save(ptr: *anyopaque, owner: ?*anyopaque, at: u32) anyerror!*anyopaque { + _ = owner; + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + if (e.pos() != at) return error.NotAtMark; + const k = try e.gpa.create(engine.Kept); + errdefer e.gpa.destroy(k); + k.* = .{ .id = try e.keepState(), .at = at }; + return k; +} + +/// The lane backend restores inside its own prompt pass (Lanes.opening), never from the store's `begin`. +pub fn snap_restore(ptr: *anyopaque, owner: ?*anyopaque, saved: *anyopaque) anyerror!void { + _ = ptr; + _ = owner; + _ = saved; + return error.BackendRestores; +} + +pub fn snap_drop(ptr: *anyopaque, saved: *anyopaque) void { + const e: *engine.Engine = @ptrCast(@alignCast(ptr)); + const k: *engine.Kept = @ptrCast(@alignCast(saved)); + e.dropState(k.id); + e.gpa.destroy(k); +} + +/// The server's words for this family's refusals. +pub fn explain(_: ?*anyopaque, err: anyerror) ?[]const u8 { + return switch (err) { + error.Quant => "Ling-3.1 on CUDA reads bf16, block-FP8 and ModelOpt NVFP4 checkpoints (the log names what config.json has)", + error.QuantizedRouter => "Ling-3.1's routing kernels read a bf16 router (mlp.gate): this checkpoint quantized it", + error.PromptTooLong => "the prompt and its reply exceed the window: raise --context", + else => null, + }; +} + +test "capacity keeps a verify window past the context and a multiple of 64" { + try std.testing.expectEqual(@as(usize, 8256), capacity(8192)); + try std.testing.expect(capacity(1000) % 64 == 0 and capacity(1000) >= 1016); +} diff --git a/zig/src/families/ling/state.zig b/zig/src/families/ling/state.zig new file mode 100644 index 0000000000..50505e696d --- /dev/null +++ b/zig/src/families/ling/state.zig @@ -0,0 +1,286 @@ +//! Ling-3.1-flash's CUDA memory at TP=1: a forward's Buffers (any sequences) and one sequence's State, an arena each. +const std = @import("std"); +const cuda = @import("cuda"); +const graph = @import("graph.zig"); +const weights = @import("weights.zig"); + +const kda = graph.kda; +const latent = graph.latent; + +/// A verify window's most rows on one sequence: the KDA saves and projection rows a commit replays. +pub const window_rows = 16; +/// The most rows of one roped attention call (a prompt chunk's rows go in blocks of these). +pub const att_rows = 128; +/// MLA's cache row: the 512-wide latent, then the 64 roped key columns. +pub const cache_width = 512 + latent.rope_width; + +/// Sub-allocations of one zeroed device allocation, 256-byte aligned. +const Arena = struct { + buf: cuda.DeviceBuffer, + + fn init(d: *const cuda.Driver, sizes: []const usize, out: []u64) !Arena { + var total: usize = 0; + for (sizes) |n| total += std.mem.alignForward(usize, n, 256); + const buf = try cuda.DeviceBuffer.alloc(d, @max(total, 256)); + errdefer { + var b = buf; + b.free(); + } + try buf.fill8(0, null); + var at: usize = 0; + for (sizes, out) |n, *p| { + p.* = buf.ptr + at; + at += std.mem.alignForward(usize, n, 256); + } + return .{ .buf = buf }; + } +}; + +/// The KDA projection row: q | k | v (3 heads 128), then beta's logits, padded to 8 columns. +pub fn kdaStride(w: *const weights.Weights) usize { + const h: usize = w.config.kda_heads; + return 3 * h * kda.dk + std.mem.alignForward(usize, h, 8); +} + +/// The largest split-K partials any loaded projection takes at 1 .. rows rows. +fn partBytes(m: *const graph.Modules, w: *const weights.Weights, rows: usize) usize { + var most: usize = 16; + const one = struct { + fn of(mo: *const graph.Modules, p: weights.Proj, r: usize, best: *usize) void { + const q = switch (p) { + .fp8g, .nvfp4, .fp8, .bf16 => |x| x, + else => return, + }; + for (1..r + 1) |k| best.* = @max(best.*, mo.lane.partBytes(k, q)); + } + }; + for (w.layers) |l| { + switch (l.mixer) { + .kda => |a| for ([_]weights.Proj{ a.q, a.k, a.v, a.b, a.f, a.g, a.o }) |p| one.of(m, p, rows, &most), + .mla => |a| for ([_]?weights.Proj{ a.q, a.q_a, a.q_b, a.kv_a, a.gate, a.o }) |p| if (p) |x| one.of(m, x, rows, &most), + } + const mlps: [2]?weights.Mlp = switch (l.ffn) { + .dense => |x| .{ x, null }, + .moe => |e| .{ e.shared, if (e.routed == .lanes and e.routed.lanes.len > 0) e.routed.lanes[0] else null }, + }; + for (mlps) |x| if (x) |p| for ([_]weights.Proj{ p.gate, p.up, p.down }) |q| one.of(m, q, rows, &most); + } + if (w.head) |hd| one.of(m, hd, rows, &most); + return most; +} + +/// A forward's scratch for up to `rows` rows on any sequences: prompt rows (`prompt`) or decode rows. +pub const Buffers = struct { + arena: Arena, + rows: usize, + prompt: bool, + nch: usize, // the most 512-key chunks one attention call reads + ids: u64, // i32 [rows] + x: u64, // bf16 [rows][hidden]: the residual stream + normed: u64, + y: u64, // a block's output rows + sh: u64, // the shared expert's rows + tmp: u64, // one projection's rows before they go into their columns + part: u64, // the lane's split-K partials + kp: u64, // bf16 [rows][kdaStride]: a prompt's KDA projection rows (a window's stay in the State) + ka: u64, // bf16 [rows][kda n]: the decay's projection + kg: u64, // bf16 [rows][kda n]: the output gate's projection + saves: kda.Saves, // out bf16; a prompt's k, v, g, b (a window's stay in the State); state_out set per layer + wide: kda.Wide, + qa_in: u64, // bf16 [rows][q_lora] + qr: u64, // bf16 [rows][q_lora]: q_a_layernorm's rows + q: u64, // bf16 [rows][heads 192] + kv: u64, // bf16 [rows][576]: kv_a_proj_with_mqa's rows + lat: u64, // bf16 [rows][576]: normed latent, roped key columns + qa: u64, // bf16 [rows][heads][576]: absorbed then roped queries + ol: u64, // bf16 [rows][heads][512] + vn: u64, // bf16 [rows][heads 128] + g: u64, // bf16 [rows][heads]: g_proj's rows + partials: latent.Partials, + dummy: u64, + at: u64, // i32: an attention block's first position + gu: u64, // bf16 [rows][2 width]: [gate | up] + act: u64, // bf16 [rows][width] + xs: u64, // fp32 [rows][width / 64] + mlog: u64, // fp32 [rows][experts] + rpart: u64, // fp32 [router_ks][rows][experts] + pick: u64, // i32 [rows][top_k + 1] + wts: u64, // fp32 [rows][top_k + 1] + plan: cuda.grouped.Plan, + eact: u64, // bf16 [pairs][moe_inter] + ey: u64, // bf16 [pairs][hidden], or every expert's rows [experts][rows][hidden] for per-expert lanes + logits: u64, // bf16 [rows][vocab] + + pub fn init(d: *const cuda.Driver, m: *const graph.Modules, w: *const weights.Weights, rows: usize, prompt: bool, capacity: usize) !Buffers { + const c = &w.config; + const D: usize = c.hidden; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const HM: usize = c.mla_heads; + const ql: usize = @max(c.q_lora, 1); + const E: usize = c.experts; + const slots: usize = c.topk + 1; + const pairs = rows * slots; + const width: usize = @max(@max(c.dense_inter, c.shared_inter), c.moe_inter); + var lanes = false; + for (w.layers) |l| if (l.ffn == .moe and l.ffn.moe.routed == .lanes) { + lanes = true; + }; + const ar = @min(rows, att_rows); + const nch = latent.chunksFor(capacity); + const hist = if (pairs > cuda.grouped.small) (pairs + cuda.grouped.small - 1) / cuda.grouped.small * (E + 1) else 1; + const sizes = [_]usize{ + rows * 4, rows * D * 2, rows * D * 2, rows * D * 2, + rows * D * 2, rows * @max(n, width) * 2, partBytes(m, w, rows), if (prompt) rows * kdaStride(w) * 2 else 16, + rows * n * 2, rows * n * 2, rows * n * 2, rows * n * 4, + rows * n * 2, rows * n * 4, rows * H * 4, rows * n * 4, + rows * n * 2, rows * ql * 2, rows * ql * 2, rows * HM * c.qkDim() * 2, + rows * cache_width * 2, rows * cache_width * 2, rows * HM * cache_width * 2, rows * HM * 512 * 2, + rows * HM * c.v_dim * 2, rows * HM * 2, nch * ar * HM * 512 * 4, nch * ar * HM * 4, + nch * ar * HM * 4, 4, 4, rows * 2 * width * 2, + rows * width * 2, rows * (width / 64 + 1) * 4, rows * E * 4, glue_router_ks * rows * E * 4, + pairs * 4, pairs * 4, pairs * 4, cuda.grouped.maxItems(pairs, E + 1, 16) * 12, + 8, pairs * 4, hist * 4, pairs * c.moe_inter * 2, + @max(pairs, if (lanes) E * rows else 0) * D * 2, rows * @as(usize, c.vocab) * 2, + }; + var p: [sizes.len]u64 = undefined; + const a = try Arena.init(d, &sizes, &p); + return .{ + .arena = a, + .rows = rows, + .prompt = prompt, + .nch = nch, + .ids = p[0], + .x = p[1], + .normed = p[2], + .y = p[3], + .sh = p[4], + .tmp = p[5], + .part = p[6], + .kp = p[7], + .ka = p[8], + .kg = p[9], + .saves = .{ .out = p[10], .state_out = 0, .k = p[11], .v = p[12], .g = p[13], .b = p[14] }, + .wide = .{ .q = p[15], .y = p[16] }, + .qa_in = p[17], + .qr = p[18], + .q = p[19], + .kv = p[20], + .lat = p[21], + .qa = p[22], + .ol = p[23], + .vn = p[24], + .g = p[25], + .partials = .{ .po = p[26], .pm = p[27], .pl = p[28] }, + .dummy = p[29], + .at = p[30], + .gu = p[31], + .act = p[32], + .xs = p[33], + .mlog = p[34], + .rpart = p[35], + .pick = p[36], + .wts = p[37], + .plan = .{ .members = p[38], .items = p[39], .counts = p[40], .rank = p[41], .hist = p[42] }, + .eact = p[43], + .ey = p[44], + .logits = p[45], + }; + } + + pub fn deinit(b: *Buffers) void { + b.arena.buf.free(); + b.* = undefined; + } +}; + +const glue_router_ks = graph.glue.router_ks; + +/// One sequence's committed caches: KDA conv windows and states (two copies), a window's saves, MLA latent rows. +pub const State = struct { + arena: Arena, + gpa: std.mem.Allocator, + capacity: usize, + pos: usize = 0, + pending: usize = 0, // rows of the last forward until its commit + prompt: bool = false, // whether those rows were prompt rows (their KDA states already committed) + pos_dev: u64, // i32 + conv: u64, // bf16 [kda layers][3][3 heads 128] + rec: u64, // fp32 [2][kda layers][heads][128][128] + proj: u64, // bf16 [kda layers][window_rows][kdaStride] + save_k: u64, // fp32 [kda layers][window_rows][heads 128] + save_v: u64, // bf16, as save_k + save_g: u64, // fp32, as save_k + save_b: u64, // fp32 [kda layers][window_rows][heads] + cache: u64, // bf16 [mla layers][capacity][576] + cur: []u1, // each KDA layer's committed copy of rec + kda_layers: usize, + mla_layers: usize, + heads: usize, + stride: usize, // kdaStride + rec_layer: usize, // bytes of one layer's state + + pub fn init(gpa: std.mem.Allocator, d: *const cuda.Driver, g: *const graph.Graph, capacity: usize) !State { + const c = &g.w.config; + const nk = g.kda_layers; + const nm = g.mla_layers; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const ps = kdaStride(g.w); + const W = window_rows; + const sizes = [_]usize{ 4, nk * 3 * 3 * n * 2, 2 * nk * n * kda.dv * 4, nk * W * ps * 2, nk * W * n * 4, nk * W * n * 2, nk * W * n * 4, nk * W * H * 4, nm * capacity * cache_width * 2 }; + var p: [sizes.len]u64 = undefined; + var a = try Arena.init(d, &sizes, &p); + errdefer a.buf.free(); + const cur = try gpa.alloc(u1, nk); + @memset(cur, 0); + return .{ .arena = a, .gpa = gpa, .capacity = capacity, .pos_dev = p[0], .conv = p[1], .rec = p[2], .proj = p[3], .save_k = p[4], .save_v = p[5], .save_g = p[6], .save_b = p[7], .cache = p[8], .cur = cur, .kda_layers = nk, .mla_layers = nm, .heads = H, .stride = ps, .rec_layer = n * kda.dv * 4 }; + } + + pub fn deinit(st: *State) void { + st.arena.buf.free(); + st.gpa.free(st.cur); + st.* = undefined; + } + + /// A fresh sequence: zero conv windows and states, position 0 (cache rows past it are never read). + pub fn reset(st: *State, s: cuda.Stream) !void { + const d = st.arena.buf.d; + try d.check(d.api.cuMemsetD8Async(st.conv, 0, st.kda_layers * 3 * st.convWidth() * 2, s.handle), "cuMemsetD8Async"); + try d.check(d.api.cuMemsetD8Async(st.rec, 0, 2 * st.kda_layers * st.rec_layer, s.handle), "cuMemsetD8Async"); + @memset(st.cur, 0); + st.pending = 0; + st.prompt = false; + try st.setPos(s, 0); + } + + pub fn setPos(st: *State, s: cuda.Stream, pos: usize) !void { + st.pos = pos; + const v: i32 = @intCast(pos); + try st.arena.buf.d.check(st.arena.buf.d.api.cuMemsetD32Async(st.pos_dev, @bitCast(v), 1, s.handle), "cuMemsetD32Async"); + } + + /// The conv window's channels: q | k | v. + pub fn convWidth(st: *const State) usize { + return 3 * st.heads * kda.dk; + } + + pub fn recAt(st: *const State, copy: u1, li: usize) u64 { + return st.rec + (@as(usize, copy) * st.kda_layers + li) * st.rec_layer; + } + + pub fn convAt(st: *const State, li: usize) u64 { + return st.conv + li * 3 * st.convWidth() * 2; + } + + pub fn cacheAt(st: *const State, mi: usize) u64 { + return st.cache + mi * st.capacity * cache_width * 2; + } + + /// Layer li's window saves: projection rows, k, v, decay, beta. + pub fn saves(st: *const State, li: usize) struct { p: u64, k: u64, v: u64, g: u64, b: u64 } { + const n = st.heads * kda.dk; + const kv = li * window_rows * n; + return .{ .p = st.proj + li * window_rows * st.stride * 2, .k = st.save_k + kv * 4, .v = st.save_v + kv * 2, .g = st.save_g + kv * 4, .b = st.save_b + li * window_rows * st.heads * 4 }; + } +}; diff --git a/zig/src/families/ling/store_test.zig b/zig/src/families/ling/store_test.zig new file mode 100644 index 0000000000..80ee8d1179 --- /dev/null +++ b/zig/src/families/ling/store_test.zig @@ -0,0 +1,43 @@ +//! weights.Store on the fake driver: each host or device allocation failing in turn leaves nothing lost or freed twice. +const std = @import("std"); +const cuda = @import("cuda"); +const proj = @import("weights_proj.zig"); +const fake = @import("fake_driver.zig"); + +const t = std.testing; + +/// Twelve buffers (past the lists' first growth), one upload, a lookup; the store freed on every path. +fn fill(gpa: std.mem.Allocator, d: *const cuda.Driver) !void { + var s: proj.Store = .{ .gpa = gpa, .d = d }; + defer s.deinit(); + for (0..12) |i| { + var name: [32]u8 = undefined; + _ = try s.alloc(64 + i, try std.fmt.bufPrint(&name, "buffer.{d}", .{i})); + } + _ = try s.keep("bytes", "kept"); + try t.expect(s.find("buffer.3") != null and s.find("kept").?.len == 5); +} + +test "Store.alloc: each host allocation failing in turn leaks nothing and frees no buffer twice" { + const d = fake.driver(); + fake.reset(null, cuda.abi.success); + try t.checkAllAllocationFailures(t.allocator, fill, .{&d}); + try t.expectEqual(@as(usize, 0), fake.live()); + try t.expectEqual(@as(usize, 0), fake.badReleases()); +} + +test "Store.alloc: each device call failing in turn (allocations, the upload) leaves nothing held or freed twice" { + const d = fake.driver(); + fake.reset(null, cuda.abi.success); + try fill(t.allocator, &d); + const n = fake.count(); + try t.expect(n >= 13); + for (1..n + 1) |k| { + fake.reset(k, cuda.abi.error_not_ready); + try t.expectError(error.NotReady, fill(t.allocator, &d)); + try t.expect(fake.fired()); + try t.expectEqual(@as(usize, 0), fake.live()); + try t.expectEqual(@as(usize, 0), fake.badReleases()); + } + fake.reset(null, cuda.abi.success); +} diff --git a/zig/src/families/ling/synth.zig b/zig/src/families/ling/synth.zig new file mode 100644 index 0000000000..69cdd3fef1 --- /dev/null +++ b/zig/src/families/ling/synth.zig @@ -0,0 +1,99 @@ +//! A tiny Ling-3.1 config: the real config.json's schema with fewer heads, experts and dims, in any loader format. +const std = @import("std"); +const config = @import("config.zig"); + +pub const real_config = @embedFile("fixtures/config.json"); + +/// The shrunk widths; KDA's head dim 128, MLA's 512 latent, 64 rope and 128-wide heads stay as the kernels take them. +pub const shape = .{ + .hidden = 512, + .heads = 4, + .dense_inter = 256, + .moe_inter = 128, + .experts = 16, + .topk = 4, + .n_group = 4, + .topk_group = 2, + .q_lora = 512, // a width glue.rmsnorm takes (512, 1536, 4096) + .vocab = 1024, + .group = 2, // KDA, MLA, KDA, MLA, ... +}; + +/// The tiny config for `layers` backbone layers (at least 2: layer 0 KDA + dense, 1 MLA + MoE) and one MTP layer. +pub fn configJson(gpa: std.mem.Allocator, format: config.Quant, layers: u32) ![]u8 { + if (layers < 2 or layers > config.max_layers - 1) return error.BadLayerCount; + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + var root = try std.json.parseFromSliceLeaky(std.json.Value, a, real_config, .{}); + const o = &root.object; + const ints = .{ + .{ "hidden_size", shape.hidden }, .{ "num_attention_heads", shape.heads }, .{ "num_key_value_heads", shape.heads }, + .{ "intermediate_size", shape.dense_inter }, .{ "moe_intermediate_size", shape.moe_inter }, .{ "moe_shared_expert_intermediate_size", shape.moe_inter }, + .{ "num_experts", shape.experts }, .{ "num_experts_per_tok", shape.topk }, .{ "n_group", shape.n_group }, + .{ "topk_group", shape.topk_group }, .{ "q_lora_rank", shape.q_lora }, .{ "vocab_size", shape.vocab }, + .{ "eos_token_id", shape.vocab - 1 }, .{ "pad_token_id", shape.vocab - 2 }, .{ "layer_group_size", shape.group }, + .{ "first_k_dense_replace", 1 }, .{ "num_hidden_layers", layers }, .{ "num_nextn_predict_layers", 1 }, + }; + inline for (ints) |kv| try o.put(a, kv[0], .{ .integer = kv[1] }); + // clamps on the last layer only: routed and shared, as layers 65-71 carry them in the real file + inline for (.{ "expert_swiglu_limit_list", "share_expert_swiglu_limit_list" }) |key| { + var list = std.json.Array.init(a); + for (0..layers) |i| try list.append(.{ .integer = if (i + 1 == layers) 10 else 0 }); + try o.put(a, key, .{ .array = list }); + } + switch (format) { + .bf16 => _ = o.orderedRemove("quantization_config"), + .fp8_block128 => try o.put(a, "quantization_config", try quant(a, fp8_block, layers)), + .nvfp4 => try o.put(a, "quantization_config", try quant(a, modelopt_block, layers)), + } + return std.json.Stringify.valueAlloc(gpa, root, .{ .whitespace = .indent_2 }); +} + +const fp8_block = + \\{"activation_scheme": "dynamic", "fmt": "e4m3", "quant_method": "fp8", "weight_block_size": [128, 128], + \\ "modules_to_not_convert": ["lm_head", "model.word_embeddings", "b_proj", "mlp.gate"]} +; + +/// ModelOpt NVFP4 as nvidia's checkpoints write it; `*` globs, and layer 0's dense MLP left in bf16 to mix formats. +const modelopt_block = + \\{"quant_method": "modelopt", "quant_algo": "NVFP4", "config_groups": {"group_0": {"targets": ["Linear"], + \\ "weights": {"dynamic": false, "num_bits": 4, "type": "float", "group_size": 16}, + \\ "input_activations": {"dynamic": false, "num_bits": 4, "type": "float", "group_size": 16}}}, + \\ "ignore": ["lm_head", "model.word_embeddings", "*.b_proj", "*.mlp.gate", "model.layers.0.mlp*"]} +; + +/// The block with the real file's per-layer exceptions added: each MLA layer's g_proj and the MTP layer's eh_proj. +fn quant(a: std.mem.Allocator, block: []const u8, layers: u32) !std.json.Value { + var q = try std.json.parseFromSliceLeaky(std.json.Value, a, block, .{}); + const key = if (q.object.contains("ignore")) "ignore" else "modules_to_not_convert"; + const list = &q.object.getPtr(key).?.array; + var i: u32 = 1; + while (i < layers) : (i += shape.group) try list.append(.{ .string = try std.fmt.allocPrint(a, "model.layers.{d}.attention.g_proj", .{i}) }); + try list.append(.{ .string = try std.fmt.allocPrint(a, "model.layers.{d}.eh_proj", .{layers}) }); + return q; +} + +test "the tiny configs parse in every format: KDA dense 0, MLA MoE 1, two groups, the MTP layer, the exceptions" { + const gpa = std.testing.allocator; + for ([_]config.Quant{ .bf16, .fp8_block128, .nvfp4 }) |f| { + const json = try configJson(gpa, f, 4); + defer gpa.free(json); + var why: config.Why = .{}; + const c = config.parse(gpa, json, &why) catch |e| { + std.debug.print("{t}: {s}\n", .{ f, why.text() }); + return e; + }; + try std.testing.expectEqual(f, c.quant); + try std.testing.expect(c.kind(0) == .kda and c.kind(1) == .mla and c.kind(2) == .kda and c.kind(4) == .mla); + try std.testing.expect(!c.isMoe(0) and c.isMoe(1) and c.mtp == 1 and c.n_group == 4 and c.topk_group == 2); + try std.testing.expectEqual(@as(f32, 10), c.expert_limit[3]); + try std.testing.expectEqual(@as(u32, 512), c.kv_lora); + try std.testing.expectEqual(config.Quant.bf16, c.format("model.layers.3.attention.g_proj")); + try std.testing.expectEqual(config.Quant.bf16, c.format("model.layers.2.attention.b_proj")); + try std.testing.expectEqual(config.Quant.bf16, c.format("model.layers.4.eh_proj")); + try std.testing.expectEqual(f, c.format("model.layers.2.attention.g_proj")); + try std.testing.expectEqual(f, c.format("model.layers.1.mlp.experts.3.down_proj")); + try std.testing.expectEqual(if (f == .nvfp4) .bf16 else f, c.format("model.layers.0.mlp.up_proj")); + } +} diff --git a/zig/src/families/ling/weights.zig b/zig/src/families/ling/weights.zig new file mode 100644 index 0000000000..1ccecba653 --- /dev/null +++ b/zig/src/families/ling/weights.zig @@ -0,0 +1,247 @@ +//! Ling-3.1-flash's weights on CUDA at TP=1, each tensor in the format its checkpoint stores, layer by layer. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const config = @import("config.zig"); +const names = @import("names.zig"); +pub const proj = @import("weights_proj.zig"); + +const Config = config.Config; +const Role = names.Role; +pub const Proj = proj.Proj; +pub const Mlp = proj.Mlp; +pub const Routed = proj.Routed; +pub const Store = proj.Store; +pub const drains = proj.drains; + +/// KDA: full-rank projections, the q|k|v short convs as #804's chain reads them, the decay and output norm tensors. +pub const Kda = struct { + q: Proj, + k: Proj, + v: Proj, + b: Proj, // beta, a row per head + f: Proj, // the decay's full-rank projection + g: Proj, // the output gate's full-rank projection + o: Proj, + conv: u64, // fp32 [3 heads dk][taps]: q_conv1d | k_conv1d | v_conv1d + a_log: u64, // fp32 [heads] + dt_bias: u64, // fp32 [heads dk] + o_norm: u64, // bf16 [dk] +}; + +/// MLA: q through the q_lora pair (or one q_proj), the latent kv pair, the head-wise gate, the output (`dense`). +pub const Mla = struct { + q: ?Proj = null, // without q_lora_rank + q_a: ?Proj = null, + q_a_norm: u64 = 0, // bf16 [q_lora] + q_b: ?Proj = null, + kv_a: Proj, // [kv_lora + rope, hidden] + kv_a_norm: u64, // bf16 [kv_lora] + kv_b: Proj, // [heads (nope + v), kv_lora] + gate: ?Proj = null, // head_wise: [heads, hidden] + o: Proj, +}; + +pub const Mixer = union(config.Kind) { kda: Kda, mla: Mla }; + +/// A MoE block: bf16 router and its pick bias, routed experts, the shared expert as projections, SwiGLU limits. +pub const Moe = struct { + router: u64, // bf16 [experts][hidden], as GLM's router kernel reads it + bias: u64, // fp32 [experts], 0 without moe_router_enable_expert_bias + routed: Routed, + shared: ?Mlp, + experts: usize, // loaded experts (all but in a test's partial load) + shared_limit: f32, +}; + +pub const Ffn = union(enum) { dense: Mlp, moe: Moe }; + +pub const Layer = struct { index: u32, attn_norm: u64, mlp_norm: u64, mixer: Mixer, ffn: Ffn }; + +/// The MTP layer's own tensors beside its MLA + MoE layer: the two input norms, eh_proj, its final norm. +pub const Mtp = struct { enorm: u64, hnorm: u64, eh: Proj, norm: u64 }; + +pub const Options = struct { + globals: bool = true, // embedding, final norm and head + experts: ?usize = null, // the first n routed experts of each MoE layer only (tests on a partial checkpoint) + stage: usize = 512 << 20, // raw expert bytes staged at a time (pinned and on the device), at least one expert +}; + +pub const Weights = struct { + config: Config, + store: Store, + embed: u64 = 0, // bf16 [vocab][hidden] + norm: u64 = 0, + head: ?Proj = null, // null when tied: the embedding's rows + layers: []Layer = &.{}, + mtp: ?Mtp = null, + + pub fn deinit(w: *Weights) void { + w.drop(true); + } + + /// deinit, but the device buffers leaked unless `settled` (nothing queued can still read them). + pub fn drop(w: *Weights, settled: bool) void { + w.store.gpa.free(w.layers); + w.store.drop(settled); + w.* = undefined; + } + + /// Layer `i` among the loaded ones. + pub fn layer(w: *const Weights, i: u32) ?*const Layer { + for (w.layers) |*l| if (l.index == i) return l; + return null; + } +}; + +const Loader = struct { + r: proj.Reader, + c: *const Config, + opts: Options, + buf: [192]u8 = undefined, + + fn module(L: *Loader, i: u32, role: Role) ![]const u8 { + return names.moduleName(&L.buf, i, role, null); + } + + /// Linear `role` of layer i at names.dims' [n, k]. + fn lin(L: *Loader, i: u32, role: Role) !Proj { + const nk = names.dims(L.c, role).linear; + return L.r.linear(try L.module(i, role), nk[0], nk[1]); + } + + /// A norm (`.weight`) or bare tensor (A_log, dt_bias, expert_bias) of names.numel values, as `to`. + fn small(L: *Loader, i: u32, role: Role, to: core.checkpoint.DType) !u64 { + var b: [200]u8 = undefined; + const m = try L.module(i, role); + const bare = role == .kda_a_log or role == .kda_dt_bias or role == .router_bias; + return L.r.small(if (bare) m else try std.fmt.bufPrint(&b, "{s}.weight", .{m}), names.numel(L.c, role), to); + } + + /// The MLP of gate, up and down `role`s under `prefix`, at names.dims' widths. + fn mlp(L: *Loader, prefix: []const u8, gate: Role) !Mlp { + const nk = names.dims(L.c, gate).linear; + return L.r.mlp(prefix, nk[0], nk[1]); + } + + fn kda(L: *Loader, i: u32) !Kda { + const n = names.dims(L.c, .kda_q_conv).conv[0]; + var b: [64]u8 = undefined; + var cb: [96]u8 = undefined; + return .{ + .q = try L.lin(i, .kda_q), + .k = try L.lin(i, .kda_k), + .v = try L.lin(i, .kda_v), + .b = try L.lin(i, .kda_b), + .f = try L.lin(i, .kda_f), + .g = try L.lin(i, .kda_g), + .o = try L.lin(i, .kda_o), + .conv = try L.r.convs(try std.fmt.bufPrint(&b, "model.layers.{d}.attention.", .{i}), n, try std.fmt.bufPrint(&cb, "model.layers.{d}.attention.conv", .{i})), + .a_log = try L.small(i, .kda_a_log, .f32), + .dt_bias = try L.small(i, .kda_dt_bias, .f32), + .o_norm = try L.small(i, .kda_o_norm, .bf16), + }; + } + + fn mla(L: *Loader, i: u32) !Mla { + const c = L.c; + var m: Mla = .{ + .kv_a = try L.lin(i, .mla_kv_a), + .kv_a_norm = try L.small(i, .mla_kv_a_norm, .bf16), + .kv_b = try L.lin(i, .mla_kv_b), + .o = try L.lin(i, .mla_o), + }; + if (c.q_lora > 0) { + m.q_a = try L.lin(i, .mla_q_a); + m.q_a_norm = try L.small(i, .mla_q_a_norm, .bf16); + m.q_b = try L.lin(i, .mla_q_b); + } else m.q = try L.lin(i, .mla_q); + if (c.head_gate) m.gate = try L.lin(i, .mla_gate); + return m; + } + + fn moe(L: *Loader, i: u32) !Moe { + const c = L.c; + const route = try L.module(i, .router); + if (c.format(route) != .bf16) { + std.log.err("{s} is {t}: the routing kernels read a bf16 router", .{ route, c.format(route) }); + return error.QuantizedRouter; + } + const count = @min(c.experts, L.opts.experts orelse c.experts); + const lim = @import("cuda.zig").limits(c, i); + var b: [96]u8 = undefined; + const x = names.dims(c, .expert_gate).linear; + return .{ + .router = try L.small(i, .router, .bf16), + .bias = if (c.expert_bias) try L.small(i, .router_bias, .f32) else 0, + .routed = try L.r.experts(i, count, x[0], x[1], lim.routed), + .shared = if (c.shared_inter > 0) try L.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.shared_experts.", .{i}), .shared_gate) else null, + .experts = count, + .shared_limit = lim.shared, + }; + } + + fn layer(L: *Loader, i: u32) !Layer { + const c = L.c; + var b: [64]u8 = undefined; + return .{ + .index = i, + .attn_norm = try L.small(i, .attn_norm, .bf16), + .mlp_norm = try L.small(i, .mlp_norm, .bf16), + .mixer = switch (c.kind(i)) { + .kda => .{ .kda = try L.kda(i) }, + .mla => .{ .mla = try L.mla(i) }, + }, + .ffn = if (c.isMoe(i)) .{ .moe = try L.moe(i) } else .{ .dense = try L.mlp(try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.", .{i}), .dense_gate) }, + }; + } + + fn mtp(L: *Loader, i: u32) !Mtp { + return .{ + .enorm = try L.small(i, .mtp_enorm, .bf16), + .hnorm = try L.small(i, .mtp_hnorm, .bf16), + .eh = try L.lin(i, .mtp_eh), + .norm = try L.small(i, .mtp_final_norm, .bf16), + }; + } +}; + +/// Layers `layers` (index c.layers is the MTP layer) from `ck`, each tensor in its stored format; `stream` blocking. +pub fn load(gpa: std.mem.Allocator, d: *const cuda.Driver, stream: cuda.Stream, pack: *const cuda.modelopt.Packer, ck: *core.Checkpoint, c: Config, layers: []const u32, opts: Options) !Weights { + var w: Weights = .{ .config = c, .store = .{ .gpa = gpa, .d = d } }; + var L: Loader = .{ .r = .{ .gpa = gpa, .stream = stream, .pack = pack, .ck = ck, .c = &w.config, .store = &w.store, .stage = opts.stage }, .c = &w.config, .opts = opts }; + errdefer w.store.drop(L.r.settled()); // the packers' outputs freed only once nothing queued writes them + const out = try gpa.alloc(Layer, layers.len); + errdefer gpa.free(out); + for (layers, out) |i, *l| { + if (i >= c.layers + c.mtp) return error.BadLayer; + l.* = try L.layer(i); + if (i >= c.layers) w.mtp = try L.mtp(i); + ck.release(); // a layer's pages out of the cache before the next: the GPU shares this memory + } + w.layers = out; + if (!opts.globals) return w; + const e = try ck.get("model.word_embeddings.weight"); + const vd = names.dims(&w.config, .embed).linear; + if (!e.is(.bf16, &vd)) { + std.log.err("model.word_embeddings: {t} {any}, expected bf16 {any}", .{ e.dtype, e.shape[0..e.rank], vd }); + return error.UnexpectedTensor; + } + w.embed = try w.store.keep(e.bytes, "model.word_embeddings.weight"); + w.norm = try L.r.small("model.norm.weight", names.numel(&w.config, .final_norm), .bf16); + const hd = names.dims(&w.config, .lm_head).linear; + if (!c.tie) w.head = try L.r.linear("lm_head", hd[0], hd[1]); + return w; +} + +/// load() from a checkpoint directory: config.json (and hf_quant_config.json, if any) read there too. +pub fn open(gpa: std.mem.Allocator, io: std.Io, dir: []const u8, why: *config.Why) !struct { ck: core.Checkpoint, c: Config } { + var arena = std.heap.ArenaAllocator.init(gpa); + defer arena.deinit(); + const a = arena.allocator(); + const cwd = std.Io.Dir.cwd(); + const json = try cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "config.json" }), a, .limited(1 << 22)); + const hf = cwd.readFileAlloc(io, try std.fs.path.join(a, &.{ dir, "hf_quant_config.json" }), a, .limited(1 << 22)) catch null; + const c = try config.parseWith(gpa, json, hf, why); + return .{ .ck = try core.Checkpoint.openModel(gpa, io, dir), .c = c }; +} diff --git a/zig/src/families/ling/weights_proj.zig b/zig/src/families/ling/weights_proj.zig new file mode 100644 index 0000000000..0c7123e952 --- /dev/null +++ b/zig/src/families/ling/weights_proj.zig @@ -0,0 +1,422 @@ +//! Ling's linears and routed experts as stored (bf16, FP8 blocks, ModelOpt NVFP4) into qlinear and experts views. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const config = @import("config.zig"); + +const Quant = config.Quant; +const Tensor = core.checkpoint.Tensor; +const X8 = cuda.experts.Fp8Experts; +const X4 = cuda.experts.Nvfp4Experts; + +/// One projection in the format its checkpoint stores: a lane matmul view (bf16, fp8g, nvfp4). +pub const Proj = cuda.qlinear.Weight; + +/// Gate, up and down: the dense MLP, the shared expert, and one bf16 routed expert. +pub const Mlp = struct { gate: Proj, up: Proj, down: Proj }; + +/// Routed experts: the grouped kernels' layout (fp8g, nvfp4), or bf16 lane weights (no grouped bf16 kernel). +pub const Routed = union(enum) { grouped: cuda.experts.Layer, lanes: []Mlp }; + +/// One upload, by the checkpoint name it came from (tests find small tensors by it). +pub const Named = struct { name: []u8, ptr: u64, len: usize }; + +/// Every device buffer a load made, freed together. +pub const Store = struct { + gpa: std.mem.Allocator, + d: *const cuda.Driver, + buffers: std.ArrayList(cuda.DeviceBuffer) = .empty, + named: std.ArrayList(Named) = .empty, + lanes: std.ArrayList([]Mlp) = .empty, + + pub fn deinit(s: *Store) void { + for (s.buffers.items) |*b| b.free(); + for (s.named.items) |n| s.gpa.free(n.name); + for (s.lanes.items) |l| s.gpa.free(l); + s.buffers.deinit(s.gpa); + s.named.deinit(s.gpa); + s.lanes.deinit(s.gpa); + } + + /// A device buffer of `len` bytes the store owns, named `name`; on failure the store holds nothing new. + pub fn alloc(s: *Store, len: usize, name: []const u8) !u64 { + try s.buffers.ensureUnusedCapacity(s.gpa, 1); + try s.named.ensureUnusedCapacity(s.gpa, 1); + const owned = try s.gpa.dupe(u8, name); + errdefer s.gpa.free(owned); + const b = try cuda.DeviceBuffer.alloc(s.d, @max(len, 16)); + s.buffers.appendAssumeCapacity(b); + s.named.appendAssumeCapacity(.{ .name = owned, .ptr = b.ptr, .len = len }); + return b.ptr; + } + + /// `bytes` uploaded into a buffer the store owns. + pub fn keep(s: *Store, bytes: []const u8, name: []const u8) !u64 { + const p = try s.alloc(bytes.len, name); + try s.buffers.items[s.buffers.items.len - 1].upload(0, bytes); + return p; + } + + /// deinit, or when the device may still write them (a failed drain) the host lists alone: buffers leaked. + pub fn drop(s: *Store, settled: bool) void { + if (!settled) s.buffers.clearRetainingCapacity(); + s.deinit(); + } + + pub fn find(s: *const Store, name: []const u8) ?Named { + for (s.named.items) |n| if (std.mem.eql(u8, n.name, name)) return n; + return null; + } +}; + +/// Whether `s` drained: false when its synchronize fails (work queued on it may still touch memory). +pub fn drains(s: cuda.Stream) bool { + s.synchronize() catch return false; + return true; +} + +/// Host bf16 from fp32: round to nearest even, NaN kept quiet. +pub fn toBf16(v: f32) u16 { + const b: u32 = @bitCast(v); + if (std.math.isNan(v)) return @intCast((b >> 16) | 0x40); + return @intCast((b + 0x7fff + ((b >> 16) & 1)) >> 16); +} + +/// Element i of a bf16 or fp32 tensor. +pub fn value(t: Tensor, i: usize) !f32 { + return switch (t.dtype) { + .bf16 => @bitCast(@as(u32, std.mem.readInt(u16, t.bytes[2 * i ..][0..2], .little)) << 16), + .f32 => @bitCast(std.mem.readInt(u32, t.bytes[4 * i ..][0..4], .little)), + else => error.UnexpectedTensor, + }; +} + +/// What reading a layer needs: the checkpoint, the config, the GPU packers, the store the uploads land in. +pub const Reader = struct { + gpa: std.mem.Allocator, + stream: cuda.Stream, // blocking: the packers wait for the uploads + pack: *const cuda.modelopt.Packer, + ck: *core.Checkpoint, + c: *const config.Config, + store: *Store, + stage: usize = 512 << 20, // raw expert bytes staged at a time (weights.Options.stage) + queued: bool = false, // an upload or packer launch not yet drained + buf: [256]u8 = undefined, + + fn name(r: *Reader, module: []const u8, tail: []const u8) ![]const u8 { + return std.fmt.bufPrint(&r.buf, "{s}.{s}", .{ module, tail }); + } + + /// The config gives `module` format `f`; the file disagrees (dtype or shape): refused, both named. + fn wrong(module: []const u8, f: Quant, t: Tensor, n: usize, k: usize) error{UnexpectedTensor} { + std.log.err("{s}: the config makes it {t} [{d}, {d}], the file holds {t} {any}", .{ module, f, n, k, t.dtype, t.shape[0..t.rank] }); + return error.UnexpectedTensor; + } + + /// Linear `module` [n, k] in the format its config gives it; a file tensor of another format or shape is refused. + pub fn linear(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + if (k % cuda.qmmf.group != 0) return error.UnexpectedTensor; + return switch (r.c.format(module)) { + .bf16 => r.bf16(module, n, k), + .fp8_block128 => r.fp8(module, n, k), + .nvfp4 => r.fp4(module, n, k), + }; + } + + /// A temporary device copy of `bytes` for a packer, freed by the caller once settled(). + fn staged(r: *Reader, bytes: []const u8) !cuda.DeviceBuffer { + return cuda.DeviceBuffer.fromHost(r.store.d, bytes); + } + + /// Nothing queued can still read or write memory: at once, or after a drain; false when the drain fails. + pub fn settled(r: *Reader) bool { + if (r.queued and !drains(r.stream)) return false; + r.queued = false; + return true; + } + + fn drain(r: *Reader) !void { + try r.stream.synchronize(); + r.queued = false; + } + + fn bf16(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + const t = try r.ck.get(try r.name(module, "weight")); + if (!t.is(.bf16, &.{ n, k })) return wrong(module, .bf16, t, n, k); + var src = try r.staged(t.bytes); + defer if (r.settled()) src.free(); // else leaked: a packer may still read it + const out = try r.store.alloc(cuda.modelopt.padded(n) * k * 2, module); + r.queued = true; + try r.pack.bf16Lane(r.stream, src.ptr, out, n, k); + try r.drain(); + return .{ .bf16 = cuda.modelopt.bf16Weight(out, n, k) }; + } + + /// e4m3 [n, k] and fp32 scale_inv [n/128 up, k/128] checked against the config's shape. + fn fp8Pair(r: *Reader, module: []const u8, n: usize, k: usize) !struct { w: Tensor, s: Tensor } { + const w = try r.ck.get(try r.name(module, "weight")); + if (!w.is(.f8_e4m3, &.{ n, k }) or k % cuda.fp8.block != 0) return wrong(module, .fp8_block128, w, n, k); + const s = try r.ck.get(try r.name(module, "weight_scale_inv")); + if (!s.is(.f32, &.{ (n + 127) / 128, k / 128 })) return wrong(module, .fp8_block128, s, n, k); + return .{ .w = w, .s = s }; + } + + fn fp8(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + const p = try r.fp8Pair(module, n, k); + var src = try r.staged(p.w.bytes); + defer if (r.settled()) src.free(); // else leaked: a packer may still read it + const codes = try r.store.alloc(cuda.fp8.padded(n) * k, module); + r.queued = true; + try r.pack.fp8Lane(r.stream, src.ptr, codes, n, k); + const inv = try r.gpa.alloc(f32, p.s.numel()); + defer r.gpa.free(inv); + @memcpy(std.mem.sliceAsBytes(inv), p.s.bytes); + const tiles = try r.gpa.alloc(f32, cuda.fp8.padded(n) * (k / cuda.fp8.group)); + defer r.gpa.free(tiles); + cuda.fp8.packScales(tiles, inv, n, k); + const scales = try r.store.keep(std.mem.sliceAsBytes(tiles), try r.name(module, "weight_scale_inv")); + try r.drain(); + return .{ .fp8g = cuda.fp8.weight(codes, scales, n, k) }; + } + + /// ModelOpt NVFP4 [n, k]: a weight stored without its weight_scale is refused before core.nvfp4 reads it. + fn fp4Read(r: *Reader, module: []const u8, n: usize, k: usize) !core.nvfp4.Linear { + if (!r.ck.has(try r.name(module, "weight_scale"))) return wrong(module, .nvfp4, try r.ck.get(try r.name(module, "weight")), n, k); + const l = try core.nvfp4.Linear.read(r.ck, module); + if (l.n != n or l.k != k) return wrong(module, .nvfp4, try r.ck.get(try r.name(module, "weight")), n, k); + return l; + } + + fn fp4(r: *Reader, module: []const u8, n: usize, k: usize) !Proj { + const l = try r.fp4Read(module, n, k); + var codes = try r.staged(l.codes); + defer if (r.settled()) codes.free(); // else both leaked + var scales = try r.staged(l.scales); + defer if (r.settled()) scales.free(); + const np = cuda.nvfp4.padded(n); + const words = try r.store.alloc(np / 64 * (k / 64) * 512 * 4, module); + const tiles = try r.store.alloc(np * (k / 16), try r.name(module, "weight_scale")); + r.queued = true; + try r.pack.fp4Lane(r.stream, codes.ptr, scales.ptr, words, tiles, n, k); + try r.drain(); + return .{ .nvfp4 = cuda.nvfp4.weight(words, tiles, l.global, n, k) }; + } + + /// Gate, up [width, dims] and down [dims, width] under `prefix` (`prefix`gate_proj and so on). + pub fn mlp(r: *Reader, prefix: []const u8, width: usize, dims: usize) !Mlp { + var m: [3]Proj = undefined; + inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |p, j| { + var b: [192]u8 = undefined; + const module = try std.fmt.bufPrint(&b, "{s}{s}", .{ prefix, p }); + m[j] = if (j < 2) try r.linear(module, width, dims) else try r.linear(module, dims, width); + } + return .{ .gate = m[0], .up = m[1], .down = m[2] }; + } + + /// Small tensor `tensor` (bf16 or fp32, `numel` long) as `to`, uploaded under its own name. + pub fn small(r: *Reader, tensor: []const u8, numel: usize, to: core.checkpoint.DType) !u64 { + const t = try r.ck.get(tensor); + if (t.numel() != numel or (t.dtype != .bf16 and t.dtype != .f32)) { + std.log.err("{s}: {t} {any}, expected {d} bf16 or fp32 values", .{ tensor, t.dtype, t.shape[0..t.rank], numel }); + return error.UnexpectedTensor; + } + var host: std.ArrayList(u8) = .empty; + defer host.deinit(r.gpa); + try append(r.gpa, &host, t, to); + return r.store.keep(host.items, tensor); + } + + /// The three short-conv weights [n, 1, taps] as one fp32 [3n][taps]: #804 KDA's q|k|v channels. + pub fn convs(r: *Reader, prefix: []const u8, n: usize, out_name: []const u8) !u64 { + var host: std.ArrayList(u8) = .empty; + defer host.deinit(r.gpa); + inline for (.{ "q", "k", "v" }) |x| { + var b: [192]u8 = undefined; + const t = try r.ck.get(try std.fmt.bufPrint(&b, "{s}" ++ x ++ "_conv1d.weight", .{prefix})); + if (!t.is(t.dtype, &.{ n, 1, r.c.conv }) or (t.dtype != .bf16 and t.dtype != .f32)) { + std.log.err("{s}" ++ x ++ "_conv1d: {t} {any}, expected [{d}, 1, {d}]", .{ prefix, t.dtype, t.shape[0..t.rank], n, r.c.conv }); + return error.UnexpectedTensor; + } + try append(r.gpa, &host, t, .f32); + } + return r.store.keep(host.items, out_name); + } + + /// Layer i's routed experts 0..count [width, dims]: every projection of every expert in one format, or refused. + pub fn experts(r: *Reader, i: u32, count: usize, width: usize, dims: usize, limit: f32) !Routed { + const names = @import("names.zig"); + var b: [192]u8 = undefined; + const f = r.c.format(try names.moduleName(&b, i, .expert_gate, 0)); + for (0..count) |e| for ([_]names.Role{ .expert_gate, .expert_up, .expert_down }) |role| { + const g = r.c.format(try names.moduleName(&b, i, role, @intCast(e))); + if (g != f) { + std.log.err("{s} is {t}, expert 0's gate {t}: the grouped kernels take one format a layer", .{ try names.moduleName(&b, i, role, @intCast(e)), g, f }); + return error.MixedExperts; + } + }; + var prefix: [96]u8 = undefined; + switch (f) { + .bf16 => { + const lanes = try r.gpa.alloc(Mlp, count); + errdefer r.gpa.free(lanes); + for (lanes, 0..) |*m, e| m.* = try r.mlp(try std.fmt.bufPrint(&prefix, "model.layers.{d}.mlp.experts.{d}.", .{ i, e }), width, dims); + try r.store.lanes.append(r.gpa, lanes); + return .{ .lanes = lanes }; + }, + .fp8_block128 => return .{ .grouped = try r.fp8Experts(i, count, width, dims, limit) }, + .nvfp4 => return .{ .grouped = try r.fp4Experts(i, count, width, dims, limit) }, + } + } + + fn expertModule(b: []u8, i: u32, e: usize, comptime proj: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "model.layers.{d}.mlp.experts.{d}." ++ proj, .{ i, e }); + } + + fn label(b: []u8, i: u32, what: []const u8) ![]const u8 { + return std.fmt.bufPrint(b, "model.layers.{d}.mlp.experts.{s}", .{ i, what }); + } + + /// Experts a stage holds when one expert's raw tensors take `one` bytes. + fn batch(r: *const Reader, count: usize, one: usize) usize { + return @max(1, @min(count, r.stage / one)); + } + + /// A batch's raw tensors: copied from the file into pinned memory by `threads` threads, uploaded, packed there. + const Stage = struct { + host: cuda.HostBuffer, + dev: cuda.DeviceBuffer, + jobs: std.ArrayList(Job) = .empty, + + const Job = struct { at: usize, bytes: []const u8 }; + const threads = 8; // page faults in flight: one reader leaves the disk (or the mount) idle between them + + fn init(r: *Reader, len: usize) !Stage { + var host = try cuda.HostBuffer.alloc(r.store.d, len); + errdefer host.free(); + return .{ .host = host, .dev = try cuda.DeviceBuffer.alloc(r.store.d, len) }; + } + + fn deinit(st: *Stage, r: *Reader) void { + st.jobs.deinit(r.gpa); + if (!r.settled()) return; // both leaked on purpose: the upload or a packer may still read them + st.host.free(); + st.dev.free(); + } + + fn put(st: *Stage, r: *Reader, at: usize, bytes: []const u8) !void { + try st.jobs.append(r.gpa, .{ .at = at, .bytes = bytes }); + } + + fn copy(st: *const Stage, first: usize) void { + var j = first; + while (j < st.jobs.items.len) : (j += threads) { + const job = st.jobs.items[j]; + @memcpy(st.host.bytes[job.at..][0..job.bytes.len], job.bytes); + } + } + + /// The batch on the device; the packers run after it on `r.stream`. + fn upload(st: *Stage, r: *Reader) !void { + var pool: [threads - 1]?std.Thread = @splat(null); + for (&pool, 1..) |*t, first| t.* = std.Thread.spawn(.{}, copy, .{ st, first }) catch null; + for (pool, 1..) |t, first| if (t == null) st.copy(first); // no thread: this one copies its share + st.copy(0); + for (pool) |t| if (t) |x| x.join(); + st.jobs.clearRetainingCapacity(); + r.queued = true; + try st.dev.uploadAsync(0, st.host.bytes, r.stream.handle); + } + + /// The packers done (the stage is reused) and the batch's file pages dropped: the GPU shares this memory. + fn drained(_: *Stage, r: *Reader) !void { + try r.drain(); + r.ck.release(); + } + }; + + /// experts.Layer's fp8g form: batches of e4m3 bytes staged on the device, packed there (Packer.fp8Experts). + fn fp8Experts(r: *Reader, i: u32, count: usize, ni: usize, d: usize, limit: f32) !cuda.experts.Layer { + if (ni % X8.block != 0 or d % X8.block != 0) return error.UnexpectedTensor; + const one = ni * d; // bytes of one e4m3 matrix + const inv_n = (ni / 128) * (d / 128); // one matrix's scale_inv blocks + var lb: [96]u8 = undefined; + const up = try r.store.alloc(count * 2 * one, try label(&lb, i, "up")); + const down = try r.store.alloc(count * one, try label(&lb, i, "down")); + const inv = try r.gpa.alloc(f32, 3 * count * inv_n); // [E][gate, up] then [E][down], as stored + defer r.gpa.free(inv); + const nb = r.batch(count, 3 * one); + var stage = try Stage.init(r, 3 * nb * one); // [gate, up, down][nb][n][k] + defer stage.deinit(r); + var b: [192]u8 = undefined; + var e0: usize = 0; + while (e0 < count) : (e0 += nb) { + const n = @min(nb, count - e0); + for (0..n) |j| inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |proj, m| { + const e = e0 + j; + const x = if (m < 2) try r.fp8Pair(try expertModule(&b, i, e, proj), ni, d) else try r.fp8Pair(try expertModule(&b, i, e, proj), d, ni); + try stage.put(r, (m * nb + j) * one, x.w.bytes); + @memcpy(std.mem.sliceAsBytes(inv[(if (m < 2) 2 * e + m else 2 * count + e) * inv_n ..][0..inv_n]), x.s.bytes); + }; + try stage.upload(r); + for (0..2) |m| try r.pack.fp8Experts(r.stream, stage.dev.ptr + m * nb * one, up + e0 * 2 * one, n, ni, d, 2, m); + try r.pack.fp8Experts(r.stream, stage.dev.ptr + 2 * nb * one, down + e0 * one, n, d, ni, 1, 0); + try stage.drained(r); + } + const up_scale = try r.store.keep(std.mem.sliceAsBytes(inv[0 .. 2 * count * inv_n]), try label(&lb, i, "up_scale")); + const down_scale = try r.store.keep(std.mem.sliceAsBytes(inv[2 * count * inv_n ..]), try label(&lb, i, "down_scale")); + return .{ .format = .fp8g, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit, .after = true }; + } + + /// experts.Layer's nvfp4 form: batches of codes and scales staged on the device, packed there (Packer.fp4Experts). + fn fp4Experts(r: *Reader, i: u32, count: usize, ni: usize, d: usize, limit: f32) !cuda.experts.Layer { + if (ni % X4.cols != 0 or d % 32 != 0) return error.UnexpectedTensor; + const up_n = X4.bytes(1, ni, d, 2); + const down_n = X4.bytes(1, d, ni, 1); + const cb = ni * d / 2; // code bytes of one matrix + const sb = ni * d / 16; // its e4m3 scale bytes + var lb: [96]u8 = undefined; + const up = try r.store.alloc(count * up_n, try label(&lb, i, "up")); + const down = try r.store.alloc(count * down_n, try label(&lb, i, "down")); + const gu = try r.gpa.alloc(f32, 2 * count); + defer r.gpa.free(gu); + const gd = try r.gpa.alloc(f32, count); + defer r.gpa.free(gd); + const nb = r.batch(count, 3 * (cb + sb)); + var stage = try Stage.init(r, 3 * nb * (cb + sb)); // codes [3][nb], then scales [3][nb] + defer stage.deinit(r); + const codes = stage.dev.ptr; + const scales = stage.dev.ptr + 3 * nb * cb; + var b: [192]u8 = undefined; + var e0: usize = 0; + while (e0 < count) : (e0 += nb) { + const n = @min(nb, count - e0); + for (0..n) |j| inline for (.{ "gate_proj", "up_proj", "down_proj" }, 0..) |proj, m| { + const e = e0 + j; + const x = if (m < 2) try r.fp4Read(try expertModule(&b, i, e, proj), ni, d) else try r.fp4Read(try expertModule(&b, i, e, proj), d, ni); + try stage.put(r, (m * nb + j) * cb, x.codes); + try stage.put(r, 3 * nb * cb + (m * nb + j) * sb, x.scales); + if (m < 2) gu[2 * e + m] = x.global else gd[e] = x.global; + }; + try stage.upload(r); + for (0..2) |m| try r.pack.fp4Experts(r.stream, codes + m * nb * cb, scales + m * nb * sb, up + e0 * up_n, n, ni, d, 0, d, 2, m); + try r.pack.fp4Experts(r.stream, codes + 2 * nb * cb, scales + 2 * nb * sb, down + e0 * down_n, n, d, ni, 0, ni, 1, 0); + try stage.drained(r); + } + const up_scale = try r.store.keep(std.mem.sliceAsBytes(gu), try label(&lb, i, "up_scale")); + const down_scale = try r.store.keep(std.mem.sliceAsBytes(gd), try label(&lb, i, "down_scale")); + return .{ .format = .nvfp4, .up = up, .down = down, .up_scale = up_scale, .down_scale = down_scale, .width = ni, .dims = d, .experts = count, .limit = limit, .after = true }; + } +}; + +/// `t`'s elements appended to `host` as bf16 or fp32 (bf16 from fp32 rounds to nearest even). +pub fn append(gpa: std.mem.Allocator, host: *std.ArrayList(u8), t: Tensor, to: core.checkpoint.DType) !void { + if (t.dtype == to) return host.appendSlice(gpa, t.bytes); + for (0..t.numel()) |i| { + const v = try value(t, i); + switch (to) { + .f32 => try host.appendSlice(gpa, std.mem.asBytes(&v)), + .bf16 => try host.appendSlice(gpa, std.mem.asBytes(&toBf16(v))), + else => return error.UnexpectedTensor, + } + } +} diff --git a/zig/src/native/cuda.zig b/zig/src/native/cuda.zig index 254a8799a9..dcc5c02ee9 100644 --- a/zig/src/native/cuda.zig +++ b/zig/src/native/cuda.zig @@ -6,6 +6,7 @@ const lanes = @import("lanes"); const nemotron = @import("nemotron"); const lfm2 = @import("lfm2.zig"); const glm = @import("glm_cuda"); +const ling = @import("ling_cuda"); const Allocator = std.mem.Allocator; const budget = @import("cuda_memory.zig"); const lane_backend = @import("cuda_bound.zig"); @@ -15,7 +16,7 @@ const cuda_disk = @import("cuda_disk.zig"); const vision_offer = @import("cuda_vision_offer.zig"); /// CUDA families provide metadata, open their lane backend and explain their refusals. -const registry = .{ nemotron.native, nemotron.omni, glm.native }; +const registry = .{ nemotron.native, nemotron.omni, glm.native, ling.native }; /// --record opens here (engines.zig refuses it for a backend without this). pub const records = true; diff --git a/zig/tests/cuda/glm_kda_f64.zig b/zig/tests/cuda/glm_kda_f64.zig new file mode 100644 index 0000000000..319839a5b6 --- /dev/null +++ b/zig/tests/cuda/glm_kda_f64.zig @@ -0,0 +1,201 @@ +//! glm_kda.cu's chain in float64 at any head count, end to end and the recurrence alone (glm-kda-ling). +const std = @import("std"); + +pub const dk = 128; + +/// x rounded to the nearest bf16, ties to even (finite normal values). +pub fn bf16(x: f64) f64 { + const u: u64 = @bitCast(x); + const lsb = (u >> 45) & 1; + return @bitCast((u + ((@as(u64, 1) << 44) - 1) + lsb) & ~((@as(u64, 1) << 45) - 1)); +} + +pub fn bf(u: u16) f64 { + return @as(f32, @bitCast(@as(u32, u) << 16)); +} + +fn sigmoid(x: f64) f64 { + return 1 / (1 + @exp(-x)); +} + +const p7 = std.math.pow(f64, 2, -7); +const p8 = std.math.pow(f64, 2, -8); +/// The recurrence's bound from the saves: fp32 sums of up to 144 terms a row, contracting rows, 2^4 to spare. +const p12 = std.math.pow(f64, 2, -12); + +/// The call's shape: p rows hold q | k | v (3 H 128 channels) and beta's logits at b_off; a and g rows are H 128 wide. +pub const Layout = struct { heads: usize, p_stride: usize, b_off: usize, lower: f64, eps: f64 }; + +/// Raw rows and weights as the chain takes them: cs [3][3 H 128] oldest first, taps [3 H 128][4], state [H][v][k]. +pub const Inputs = struct { p: []const u16, a: []const u16, g: []const u16, cs: []const u16, cw: []const f32, state: []const f32, a_log: []const f32, dt_bias: []const f32, norm: []const u16 }; + +/// prep's rows [rows][H][128]: q (scaled), k, v bf16, decay factors, and beta [rows][H]. +pub const Saved = struct { q: []const f32, k: []const f32, v: []const u16, g: []const f32, b: []const f32 }; + +/// Outputs [rows][H 128] and the final state [H][128][128], each with a per-element bound. +pub const Want = struct { out: []f64, out_tol: []f64, state: []f64, state_tol: []f64 }; + +fn want(a: std.mem.Allocator, heads: usize, rows: usize) !Want { + const n = heads * dk; + return .{ .out = try a.alloc(f64, rows * n), .out_tol = try a.alloc(f64, rows * n), .state = try a.alloc(f64, n * dk), .state_tol = try a.alloc(f64, n * dk) }; +} + +/// A delta step: decay, u = beta (v - S k), S += k u, o = S q; mag bounds |S|; o_mag kda_ref's or (fp32) sum mag |q|. +fn step(s: []f64, mag: []f64, q: []const f64, k: []const f64, v: []const f64, g: []const f64, beta: f64, o: []f64, o_mag: []f64, fp32: bool) void { + var kq: f64 = 0; + for (q, k) |qq, kk| kq += @abs(qq * kk); + for (0..dk) |j| { + const row = s[j * dk ..][0..dk]; + const mrow = mag[j * dk ..][0..dk]; + var kv: f64 = 0; + var km: f64 = 0; + for (row, mrow, k, g) |*x, *m, kk, gg| { + x.* *= gg; + m.* *= gg; + kv += x.* * kk; + km += m.* * @abs(kk); + } + const u = beta * (v[j] - kv); + const um = beta * (@abs(v[j]) + km); + var acc: f64 = 0; + var am: f64 = 0; + var aq: f64 = 0; + for (row, mrow, k, q) |*x, *m, kk, qq| { + x.* += kk * u; + m.* += @abs(kk) * um; + acc += x.* * qq; + am += m.* * @abs(qq); + aq += @abs(x.* * qq); + } + o[j] = acc; + o_mag[j] = if (fp32) am else aq + beta * @abs(v[j]) * kq; + } +} + +/// kda_ref.gatedNorm: o / rms * w * sigmoid(gate); tol: o's bound in, y's bound out (first order) + y's bf16 rounding. +fn gatedNorm(o: []const f64, gate: []const f64, w: []const f64, eps: f64, y: []f64, tol: []f64) void { + var ss: f64 = 0; + var cross: f64 = 0; + for (o, tol) |x, t| { + ss += x * x; + cross += @abs(x) * t; + } + const nn: f64 = @floatFromInt(o.len); + const rstd = 1 / @sqrt(ss / nn + eps); + const moved = cross / nn / (ss / nn + eps); + for (y, tol, o, gate, w) |*out, *t, x, z, ww| { + out.* = x * rstd * ww * sigmoid(z); + t.* = @abs(ww * sigmoid(z)) * rstd * (t.* + @abs(x) * moved) + p8 * @abs(out.*); + } +} + +fn l2norm(x: []f64, scale: f64) void { + var ss: f64 = 0; + for (x) |e| ss += e * e; + const f = scale / @sqrt(ss + 1e-6); + for (x) |*e| e.* *= f; +} + +/// End to end: conv+SiLU, q/k L2 norms, safe gate, bf16 beta, delta rule, bf16(S q), gated norm; kda_ref's 2^-7 bounds. +pub fn chain(a: std.mem.Allocator, l: Layout, in: Inputs, rows: usize) !Want { + const n = l.heads * dk; + const cc = 3 * n; + var w = try want(a, l.heads, rows); + const mag = try a.alloc(f64, dk * dk); + var qkv: [3][dk]f64 = undefined; + var g: [dk]f64 = undefined; + var o: [dk]f64 = undefined; + var om: [dk]f64 = undefined; + var gate: [dk]f64 = undefined; + var nw: [dk]f64 = undefined; + for (&nw, in.norm[0..dk]) |*d, x| d.* = bf(x); + for (0..l.heads) |h| { + const s = w.state[h * dk * dk ..][0 .. dk * dk]; + for (s, in.state[h * dk * dk ..][0 .. dk * dk]) |*d, x| d.* = x; + for (mag, s) |*m, x| m.* = @abs(x); + const rate = @exp(@as(f64, in.a_log[h])); + for (0..rows) |row| { + for (0..3) |part| for (0..dk) |i| { + const c = part * n + h * dk + i; + var acc: f64 = 0; + for (0..4) |tap| { + const at = row + tap; + const x = if (at < 3) bf(in.cs[at * cc + c]) else bf(in.p[(at - 3) * l.p_stride + c]); + acc += @as(f64, in.cw[c * 4 + tap]) * x; + } + qkv[part][i] = bf16(acc * sigmoid(acc)); + }; + l2norm(&qkv[0], 1 / @sqrt(@as(f64, dk))); + l2norm(&qkv[1], 1); + for (&g, 0..) |*e, i| e.* = @exp(l.lower * sigmoid(rate * (bf(in.a[row * n + h * dk + i]) + in.dt_bias[h * dk + i]))); + const beta = bf16(sigmoid(bf(in.p[row * l.p_stride + l.b_off + h]))); + step(s, mag, &qkv[0], &qkv[1], &qkv[2], &g, beta, &o, &om, false); + const tol = w.out_tol[row * n + h * dk ..][0..dk]; + for (&o, tol, om) |*e, *t, m| { + t.* = p7 * m; + e.* = bf16(e.*); + } + for (&gate, 0..) |*z, t| z.* = bf(in.g[row * n + h * dk + t]); + gatedNorm(&o, &gate, &nw, l.eps, w.out[row * n + h * dk ..][0..dk], tol); + } + for (w.state_tol[h * dk * dk ..][0 .. dk * dk], mag) |*t, m| t.* = p7 * m + 1e-30; + } + return w; +} + +/// The recurrence alone from prep's saves: S q (tol 2^-8 |y| + 2^-12 sum |mag q|), state (tol 2^-12 of magnitudes). +pub fn recurrence(a: std.mem.Allocator, heads: usize, state0: []const f32, sv: Saved, rows: usize) !Want { + const n = heads * dk; + var w = try want(a, heads, rows); + const mag = try a.alloc(f64, dk * dk); + var q: [dk]f64 = undefined; + var k: [dk]f64 = undefined; + var v: [dk]f64 = undefined; + var g: [dk]f64 = undefined; + var om: [dk]f64 = undefined; + for (0..heads) |h| { + const s = w.state[h * dk * dk ..][0 .. dk * dk]; + for (s, state0[h * dk * dk ..][0 .. dk * dk]) |*d, x| d.* = x; + for (mag, s) |*m, x| m.* = @abs(x); + for (0..rows) |row| { + const at = row * n + h * dk; + for (&q, &k, &v, &g, sv.q[at..][0..dk], sv.k[at..][0..dk], sv.v[at..][0..dk], sv.g[at..][0..dk]) |*qq, *kk, *vv, *gg, a0, a1, a2, a3| { + qq.* = a0; + kk.* = a1; + vv.* = bf(a2); + gg.* = a3; + } + const o = w.out[at..][0..dk]; + step(s, mag, &q, &k, &v, &g, sv.b[row * heads + h], o, &om, true); + for (w.out_tol[at..][0..dk], o, om) |*t, e, m| t.* = p8 * @abs(e) + p12 * m + 1e-30; + } + for (w.state_tol[h * dk * dk ..][0 .. dk * dk], mag) |*t, m| t.* = p12 * m + 1e-30; + } + return w; +} + +/// The worst |got - want| / tol. +pub fn fraction(got: []const f64, w: []const f64, tol: []const f64) f64 { + var most: f64 = 0; + for (got, w, tol) |x, e, t| { + const q = @abs(x - e) / t; + if (!(q <= most)) most = if (std.math.isNan(q)) std.math.inf(f64) else q; + } + return most; +} + +/// Relative L2 distance of got from want. +pub fn relL2(got: []const f64, w: []const f64) f64 { + var dd: f64 = 0; + var ww: f64 = 0; + for (got, w) |x, e| { + dd += (x - e) * (x - e); + ww += e * e; + } + return @sqrt(dd / ww); +} + +test "bf16 rounding ties to even" { + try std.testing.expectEqual(@as(f64, 1.0), bf16(1.0 + 1.0 / 256.0)); + try std.testing.expectEqual(@as(f64, 1.0 + 1.0 / 64.0), bf16(1.0 + 3.0 / 256.0)); +} diff --git a/zig/tests/cuda/glm_kda_ling.zig b/zig/tests/cuda/glm_kda_ling.zig new file mode 100644 index 0000000000..0316fd4361 --- /dev/null +++ b/zig/tests/cuda/glm_kda_ling.zig @@ -0,0 +1,206 @@ +//! glm_kda's paths on Ling-3.1 layer 0 (kda-real's fla fixture, q/k/v convs as one q|k|v): vs fla and float64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const f64r = @import("glm_kda_f64.zig"); +const rig = @import("glm_kda_rig.zig"); +const ref = @import("glm_ref_tests.zig"); +const Rig = rig.Rig; +const kda = rig.kda; +const dk = rig.dk; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, count: usize) ![]T { + const out = try a.alloc(T, count); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. count * @sizeOf(T)]); + at.* += count * @sizeOf(T); + return out; +} + +/// The layer's rows as GLM's chain takes them, fla's outputs and states, and per-row conv windows on the device. +const Layer = struct { + P: usize, + T: usize, + H: usize, + l: f64r.Layout, + in: f64r.Inputs, + c: kda.Chain, + cs_all: u64, // [T][3][3 H 128]: row t's window, rows t - 3 .. t - 1 (zeros before the first) + fla_y: []f64, + fla_sp: []f64, + fla_st: []f64, + + /// The chain over rows lo .. from state_in with row lo's conv window. + fn at(l: Layer, lo: usize, state_in: u64) kda.Chain { + const n = l.H * dk; + var c = l.c; + c.p += lo * l.l.p_stride * 2; + c.a += lo * n * 2; + c.g += lo * n * 2; + c.conv_state = l.cs_all + lo * 9 * n * 2; + c.state_in = state_in; + return c; + } +}; + +fn load(r: *Rig, path: []const u8) !?Layer { + const a = r.a; + const bytes = std.Io.Dir.cwd().readFileAlloc(r.gpu.io, path, a, .limited(1 << 30)) catch |e| switch (e) { + error.FileNotFound => return null, + else => return e, + }; + if (bytes.len < 32 or !std.mem.eql(u8, bytes[0..8], "TFKDA001")) return error.BadFixture; + const P: usize = std.mem.readInt(u32, bytes[8..12], .little); + const N: usize = std.mem.readInt(u32, bytes[12..16], .little); + const H: usize = std.mem.readInt(u32, bytes[16..20], .little); + if (std.mem.readInt(u32, bytes[20..24], .little) != dk) return error.BadFixture; + const lower: f32 = @bitCast(std.mem.readInt(u32, bytes[24..28], .little)); + const eps: f32 = @bitCast(std.mem.readInt(u32, bytes[28..32], .little)); + const T = P + N; + const n = H * dk; + var o: usize = 32; + const qkv = try take(u16, a, bytes, &o, T * 3 * n); + const f = try take(u16, a, bytes, &o, T * n); + const b = try take(u16, a, bytes, &o, T * H); + const gate = try take(u16, a, bytes, &o, T * n); + const conv = try take(f32, a, bytes, &o, 3 * n * 4); // Ling's q_conv1d | k_conv1d | v_conv1d: GLM's q|k|v channels + const a_log = try take(f32, a, bytes, &o, H); + const dt_bias = try take(f32, a, bytes, &o, n); + const norm32 = try take(f32, a, bytes, &o, dk); + const fy = try take(u16, a, bytes, &o, T * n); + const fsp = try take(f32, a, bytes, &o, n * dk); + const fst = try take(f32, a, bytes, &o, n * dk); + if (o != bytes.len) return error.BadFixture; + const p_stride = 3 * n + H; // [q | k | v | beta logits] + const p = try a.alloc(u16, T * p_stride); + for (0..T) |t| { + @memcpy(p[t * p_stride ..][0 .. 3 * n], qkv[t * 3 * n ..][0 .. 3 * n]); + @memcpy(p[t * p_stride + 3 * n ..][0..H], b[t * H ..][0..H]); + } + const norm = try a.alloc(u16, dk); // the o_norm weight is bf16 in the checkpoint: exact + for (norm, norm32) |*d, x| { + d.* = @truncate(@as(u32, @bitCast(x)) >> 16); + if (f64r.bf(d.*) != x) return error.NormNotBf16; + } + const zero = try a.alloc(f32, n * dk); + @memset(zero, 0); + const cs = try a.alloc(u16, T * 9 * n); + @memset(cs, 0); + for (0..T) |t| for (0..3) |j| if (t + j >= 3) @memcpy(cs[(t * 3 + j) * 3 * n ..][0 .. 3 * n], qkv[(t + j - 3) * 3 * n ..][0 .. 3 * n]); + const in: f64r.Inputs = .{ .p = p, .a = f, .g = gate, .cs = cs[0 .. 9 * n], .cw = conv, .state = zero, .a_log = a_log, .dt_bias = dt_bias, .norm = norm }; + const B = std.mem.sliceAsBytes; + const c: kda.Chain = .{ .p = try r.dev(B(p)), .p_stride = p_stride, .b_off = 3 * n, .a = try r.dev(B(f)), .a_stride = n, .g = try r.dev(B(gate)), .g_stride = n, .conv_state = 0, .conv_w = try r.dev(B(conv)), .state_in = try r.dev(B(zero)), .a_log = try r.dev(B(a_log)), .dt_bias = try r.dev(B(dt_bias)), .norm_w = try r.dev(B(norm)), .eps = eps, .lower = lower, .heads = H }; + const w64 = struct { + fn of(al: std.mem.Allocator, x: anytype) ![]f64 { + const out = try al.alloc(f64, x.len); + for (out, x) |*d, v| d.* = if (@TypeOf(v) == u16) f64r.bf(v) else v; + return out; + } + }; + return .{ .P = P, .T = T, .H = H, .l = .{ .heads = H, .p_stride = p_stride, .b_off = 3 * n, .lower = lower, .eps = eps }, .in = in, .c = c, .cs_all = try r.dev(B(cs)), .fla_y = try w64.of(a, fy), .fla_sp = try w64.of(a, fsp), .fla_st = try w64.of(a, fst) }; +} + +/// One run's outputs [T][H 128] and its states after the prompt (P rows) and at the end. +const Result = struct { y: []f64, sp: u64, st: u64 }; + +/// Rows lo .. T as decode rows (fused chain, one row a call) from `state`; each row's state into states[t] when given. +fn decode(r: *Rig, m: *const rig.Mods, l: Layer, lo: usize, state: u64, y: []f64, states: ?[]u64) !u64 { + const n = l.H * dk; + var s = state; + for (lo..l.T) |t| { + const o = try rig.Out.alloc(r, l.H, 1); + try rig.run(m, r.s, l.at(t, s), 1, o, .fused); + @memcpy(y[t * n ..][0..n], try rig.wide64(r, o.out, n, true)); + s = o.state; + if (states) |ss| ss[t + 1] = s; + } + return s; +} + +/// A prompt of P rows through `path`, then decode rows; `saves` returns the prompt's buffers. +fn promptThenDecode(r: *Rig, m: *const rig.Mods, l: Layer, path: rig.Path, saves: *rig.Out) !Result { + const n = l.H * dk; + const y = try r.a.alloc(f64, l.T * n); + saves.* = try rig.Out.alloc(r, l.H, l.P); + try rig.run(m, r.s, l.at(0, l.c.state_in), l.P, saves.*, path); + @memcpy(y[0 .. l.P * n], try rig.wide64(r, saves.out, l.P * n, true)); + return .{ .y = y, .sp = saves.state, .st = try decode(r, m, l, l.P, saves.state, y, null) }; +} + +fn fracs(r: *Rig, l: Layer, res: Result, e2e_p: f64r.Want, e2e: f64r.Want) ![4]f64 { + const n = l.H * dk; + const sp = try rig.wide64(r, res.sp, n * dk, false); + const st = try rig.wide64(r, res.st, n * dk, false); + return .{ f64r.fraction(res.y, e2e.out, e2e.out_tol), f64r.fraction(sp, e2e_p.state, e2e_p.state_tol), f64r.fraction(st, e2e.state, e2e.state_tol), @max(f64r.relL2(sp, l.fla_sp), f64r.relL2(st, l.fla_st)) }; +} + +/// `glm-kda-ling FIXTURE`. +pub fn run(gpu: check.Gpu, path: []const u8) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(1); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + const l = (try load(&r, path)) orelse return std.debug.print("SKIP glm-kda-ling: no fixture at {s}\n", .{path}); + var m = try rig.Mods.load(gpu.d); + defer m.unload(); + const n = l.H * dk; + const T = l.T; + var name: [128]u8 = undefined; + const P = std.fmt.bufPrint; + // every row a decode row, each row's state kept + const states = try r.a.alloc(u64, T + 1); + states[0] = l.c.state_in; + const yd = try r.a.alloc(f64, T * n); + _ = try decode(&r, &m, l, 0, l.c.state_in, yd, states); + const dec: Result = .{ .y = yd, .sp = states[l.P], .st = states[T] }; + // 16-row verify windows; each replayed whole (== its state) and over its first 5 rows (== the decode rows' state) + const yw = try r.a.alloc(f64, T * n); + var ws = l.c.state_in; + var w0: usize = 0; + while (w0 < T) : (w0 += 16) { + const rows = @min(16, T - w0); + const o = try rig.Out.alloc(&r, l.H, rows); + try rig.run(&m, r.s, l.at(w0, ws), rows, o, .fused); + @memcpy(yw[w0 * n ..][0 .. rows * n], try rig.wide64(&r, o.out, rows * n, true)); + const full = try r.zeros(n * dk * 4); + const part = try r.zeros(n * dk * 4); + try kda.replay(&m.ks, r.s, ws, o.k, o.v, o.g, o.b, rows, l.H, full); + try kda.replay(&m.ks, r.s, ws, o.k, o.v, o.g, o.b, @min(5, rows), l.H, part); + try check.sameBytes(try P(&name, "Ling window at {d}: replay == the window's state", .{w0}), try r.back(full, n * dk * 4), try r.back(o.state, n * dk * 4)); + try check.sameBytes(try P(&name, "Ling window at {d}: replay of 5 rows == 5 decode rows' state", .{w0}), try r.back(part, n * dk * 4), try r.back(states[w0 + @min(5, rows)], n * dk * 4)); + ws = o.state; + } + try check.sameBytes("Ling: 16-row windows' state == decode rows' state", try r.back(ws, n * dk * 4), try r.back(states[T], n * dk * 4)); + for (yw, yd) |a, b| if (a != b) return check.expect(false, "Ling: 16-row windows' outputs == decode rows' outputs", .{}); + const win: Result = .{ .y = yw, .sp = states[l.P], .st = ws }; + var wsv: rig.Out = undefined; + const wide = try promptThenDecode(&r, &m, l, .wide, &wsv); + for (wide.y, yd) |a, b| if (a != b) return check.expect(false, "Ling: the wide prompt's and decode outputs == every row a decode row's", .{}); + try ref.sameRows(&r, "Ling: the wide prompt's state == P decode rows' state", wide.sp, states[l.P], n * dk * 4); + const e2e_p = try f64r.chain(r.a, l.l, l.in, l.P); + const e2e = try f64r.chain(r.a, l.l, l.in, T); + const rec = try f64r.recurrence(r.a, l.H, l.in.state, try rig.saved(&r, wsv, l.H), l.P); + const recw = [2]f64{ f64r.fraction(try rig.wide64(&r, wsv.y, l.P * n, true), rec.out, rec.out_tol), f64r.fraction(try rig.wide64(&r, wsv.state, n * dk, false), rec.state, rec.state_tol) }; + var worst: f64 = @max(recw[0], recw[1]); + var most_l2: f64 = 0; + std.debug.print("RESULT glm-kda-ling: Ling-3.1 layer 0, {d} heads, {d} prompt + {d} decode rows; fla rel L2 of outputs (states after the prompt / end), worst fraction of float64's bound (out, state after prompt, at end)\n", .{ l.H, l.P, T - l.P }); + for ([_]Result{ dec, win, wide }, [_][]const u8{ "chain, every row a decode row", "chain, 16-row windows", "wide prompt + decode" }) |res, label| { + const f = try fracs(&r, l, res, e2e_p, e2e); + const sp = try rig.wide64(&r, res.sp, n * dk, false); + const st = try rig.wide64(&r, res.st, n * dk, false); + const yl2 = f64r.relL2(res.y, l.fla_y); + std.debug.print("RESULT {s}: rel L2 {e:.3} (states {e:.3} / {e:.3}); float64 {d:.3}, {d:.4}, {d:.4}\n", .{ label, yl2, f64r.relL2(sp, l.fla_sp), f64r.relL2(st, l.fla_st), f[0], f[1], f[2] }); + worst = @max(worst, @max(f[0], @max(f[1], f[2]))); + most_l2 = @max(most_l2, @max(yl2, f[3])); + } + const fla_f = f64r.fraction(l.fla_y, e2e.out, e2e.out_tol); + std.debug.print("RESULT fla's own outputs: float64 {d:.3}; recurrence from the wide prompt's saves, y {d:.3} state {d:.3}\n", .{ fla_f, recw[0], recw[1] }); + try check.expect(worst <= 1, "glm-kda-ling: every path within float64's bound (worst {d:.3})", .{worst}); + try check.expect(most_l2 < 1e-2, "glm-kda-ling: every path within 1 % of fla (worst rel L2 {e:.3})", .{most_l2}); + check.pass("glm-kda-ling: glm_kda serves Ling layer 0 as is; chain, windows and wide within float64's bound (worst {d:.3}) and fla (rel L2 <= {e:.3}); replays == chain bits", .{ worst, most_l2 }); +} diff --git a/zig/tests/cuda/glm_kda_rig.zig b/zig/tests/cuda/glm_kda_rig.zig new file mode 100644 index 0000000000..a39b5385db --- /dev/null +++ b/zig/tests/cuda/glm_kda_rig.zig @@ -0,0 +1,111 @@ +//! glm-kda-ling's rig: a call's inputs on the device, one path's outputs, downloads to f64. +const std = @import("std"); +const cuda = @import("cuda"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const f64r = @import("glm_kda_f64.zig"); +pub const Rig = qr.RigT; + +pub const kda = glm.kda; +pub const dk = kda.dk; + +/// glm_kda.cu's module (chain, replay, wide). +pub const Mods = struct { + kda_mod: cuda.Module, + ks: kda.Kernels, + + pub fn load(d: *const cuda.Driver) !Mods { + var m: Mods = undefined; + m.kda_mod = try cuda.Module.load(d, cuda.kernels.glm_kda); + m.ks = try kda.Kernels.load(m.kda_mod); + return m; + } + + pub fn unload(m: *Mods) void { + m.kda_mod.unload(); + } +}; + +/// How a call runs: the fused chain or the three-kernel (serial wide) chain. +pub const Path = enum { fused, wide }; + +/// One path's buffers for `rows` rows: out, state, saves (k, v, decay, beta), wide scratch (q, y). +pub const Out = struct { + out: u64, + state: u64, + k: u64, + v: u64, + g: u64, + b: u64, + q: u64, + y: u64, + rows: usize, + + pub fn alloc(r: *Rig, heads: usize, rows: usize) !Out { + const n = heads * dk; + return .{ .out = try r.zeros(rows * n * 2), .state = try r.zeros(n * dk * 4), .k = try r.zeros(rows * n * 4), .v = try r.zeros(rows * n * 2), .g = try r.zeros(rows * n * 4), .b = try r.zeros(rows * heads * 4), .q = try r.zeros(rows * n * 4), .y = try r.zeros(rows * n * 2), .rows = rows }; + } + + pub fn saves(o: Out) kda.Saves { + return .{ .out = o.out, .state_out = o.state, .k = o.k, .v = o.v, .g = o.g, .b = o.b }; + } + + /// The same buffers from row `lo` on (state unchanged). + pub fn from(o: Out, heads: usize, lo: usize) Out { + const n = heads * dk; + return .{ .out = o.out + lo * n * 2, .state = o.state, .k = o.k + lo * n * 4, .v = o.v + lo * n * 2, .g = o.g + lo * n * 4, .b = o.b + lo * heads * 4, .q = o.q + lo * n * 4, .y = o.y + lo * n * 2, .rows = o.rows - lo }; + } +}; + +/// `rows` rows of c through `path` into o. +pub fn run(m: *const Mods, s: cuda.Stream, c: kda.Chain, rows: usize, o: Out, path: Path) !void { + const w: kda.Wide = .{ .q = o.q, .y = o.y }; + switch (path) { + .fused => try kda.chain(&m.ks, s, c, rows, o.saves(), null), + .wide => try kda.chain(&m.ks, s, c, rows, o.saves(), w), + } +} + +pub fn floats(r: *Rig, ptr: u64, count: usize) ![]f32 { + const raw = try r.back(ptr, count * 4); + const out = try r.a.alloc(f32, count); + @memcpy(std.mem.sliceAsBytes(out), raw); + return out; +} + +pub fn halves(r: *Rig, ptr: u64, count: usize) ![]u16 { + const raw = try r.back(ptr, count * 2); + const out = try r.a.alloc(u16, count); + @memcpy(std.mem.sliceAsBytes(out), raw); + return out; +} + +/// fp32 or bf16 device values as f64. +pub fn wide64(r: *Rig, ptr: u64, count: usize, comptime bf16: bool) ![]f64 { + const out = try r.a.alloc(f64, count); + if (bf16) { + for (out, try halves(r, ptr, count)) |*d, x| d.* = f64r.bf(x); + } else for (out, try floats(r, ptr, count)) |*d, x| { + d.* = x; + } + return out; +} + +/// prep's saves of o, downloaded. +pub fn saved(r: *Rig, o: Out, heads: usize) !f64r.Saved { + const n = o.rows * heads * dk; + return .{ .q = try floats(r, o.q, n), .k = try floats(r, o.k, n), .v = try halves(r, o.v, n), .g = try floats(r, o.g, n), .b = try floats(r, o.b, o.rows * heads) }; +} + +/// Largest |a - b| over the largest |b|, and how many elements differ at all. +pub fn diff(a: []const f64, b: []const f64) struct { rel: f64, count: usize } { + var most: f64 = 0; + var big: f64 = 0; + var count: usize = 0; + for (a, b) |x, y| { + most = @max(most, @abs(x - y)); + big = @max(big, @abs(y)); + if (x != y) count += 1; + } + return .{ .rel = most / big, .count = count }; +} diff --git a/zig/tests/cuda/glm_mla_same.zig b/zig/tests/cuda/glm_mla_same.zig new file mode 100644 index 0000000000..f5bf6b4afa --- /dev/null +++ b/zig/tests/cuda/glm_mla_same.zig @@ -0,0 +1,116 @@ +//! GLM's MLA launches on this build's glm_glue against a 6380c8da glm_glue fatbin: the same bytes for GLM's shapes. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const latent = glm.latent; +const Set = glm.glue.Set; + +const heads = 32; +const lw = 512; +const d = 256; + +/// The kernels Ling added, absent from the old image: taken from the new one (GLM's launches never reach them). +const added = [_][]const u8{ "absorb_q_128", "expand_v_128", "rope", "head_gate", "chunks_r16" }; + +fn oldKernels(old: cuda.Module, new: glm.glue.Kernels) !glm.glue.Kernels { + var k = new; + inline for (@typeInfo(glm.glue.Kernels).@"struct".field_names) |name| { + const fresh = for (added) |a| { + if (std.mem.eql(u8, a, name)) break true; + } else false; + if (!fresh) @field(k, name) = try old.function("tf_glm_" ++ name); + } + try k.chunks_16.allowDynamicShared(glm.glue.chunksSmem(16)); + try k.chunks_32.allowDynamicShared(glm.glue.chunksSmem(32)); + return k; +} + +fn partials(r: *Rig, nch: usize, rows: usize) !latent.Partials { + return .{ .po = try r.zeros(nch * rows * heads * lw * 4), .pm = try r.zeros(nch * rows * heads * 4), .pl = try r.zeros(nch * rows * heads * 4) }; +} + +const Inputs = struct { q: u64, wk: u64, ol: u64, wv: u64, cache: u64, lat: u64, tok: u64, cnt: u64 }; + +/// Every GLM MLA launch (absorb, expand, cache write, dense 16/32-head and sparse attention) on `set`: output bytes. +fn outputs(r: *Rig, set: Set, in: Inputs) ![]const []const u8 { + var list: std.ArrayList([]const u8) = .empty; + for ([_]usize{ 1, 5, 70 }) |rows| { + const qa = try r.zeros(rows * heads * lw * 2); + try latent.absorbQ(set, r.s, in.q, in.wk, qa, rows, heads, d, lw); + try list.append(r.a, try r.back(qa, rows * heads * lw * 2)); + const vn = try r.zeros(rows * heads * d * 2); + try latent.expandV(set, r.s, in.ol, in.wv, vn, rows, heads, d, lw); + try list.append(r.a, try r.back(vn, rows * heads * d * 2)); + } + const c = try r.zeros(1500 * lw * 2); + const p0: i32 = 1400; + try latent.latentWrite(set, r.s, in.lat, lw + 64, c, try r.dev(std.mem.asBytes(&p0)), 100, lw); + try list.append(r.a, try r.back(c, 1500 * lw * 2)); + for ([_][3]usize{ .{ 1430, 70, 32 }, .{ 1430, 70, 16 }, .{ 1499, 1, 16 }, .{ 0, 64, 32 } }) |case| { + const pos, const rows, const hb = case; + const nch = latent.chunksFor(pos + rows); + const out = try r.zeros(rows * heads * lw * 2); + const at: i32 = @intCast(pos); + try latent.attention(set, r.s, in.q, in.cache, try r.dev(std.mem.asBytes(&at)), try partials(r, nch, rows), out, try r.zeros(4), rows, heads, lw, nch, 0.125, hb); + try list.append(r.a, try r.back(out, rows * heads * lw * 2)); + } + for ([_]usize{ 70, 1 }) |rows| { + const out = try r.zeros(rows * heads * lw * 2); + try latent.sparseAttention(set, r.s, in.q, in.cache, in.tok, in.cnt, 1100, try partials(r, latent.chunksFor(1100), rows), out, rows, heads, lw, 0.125); + try list.append(r.a, try r.back(out, rows * heads * lw * 2)); + } + return list.items; +} + +/// `glm-mla-same `: the old and new images' GLM MLA outputs, byte for byte, on seeded inputs. +pub fn run(gpu: Gpu, old_path: []const u8) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(804); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + const raw = try std.Io.Dir.cwd().readFileAlloc(gpu.io, old_path, a, .limited(64 << 20)); + const image = try a.alignedAlloc(u8, .@"16", raw.len); + @memcpy(image, raw); + var new_mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer new_mod.unload(); + var old_mod = try cuda.Module.load(gpu.d, image); + defer old_mod.unload(); + const new = try glm.glue.Kernels.load(new_mod); + const old = try oldKernels(old_mod, new); + const tok = try a.alloc(i32, 70 * 1100); + const cnt = try a.alloc(i32, 70); + for (cnt, 0..) |*n, row| { + n.* = if (row == 5) 0 else @intCast(1 + r.rng.uintLessThan(usize, 1100)); + for (tok[row * 1100 ..][0..1100], 0..) |*t, i| t.* = if (i < n.*) @intCast(r.rng.uintLessThan(u32, 6000)) else -1; + } + const in: Inputs = .{ + .q = try r.dev(std.mem.sliceAsBytes(try r.bfs(70 * heads * lw, 1.0))), + .wk = try r.dev(std.mem.sliceAsBytes(try r.bfs(heads * d * lw, 0.05))), + .ol = try r.dev(std.mem.sliceAsBytes(try r.bfs(70 * heads * lw, 1.0))), + .wv = try r.dev(std.mem.sliceAsBytes(try r.bfs(heads * d * lw, 0.05))), + .cache = try r.dev(std.mem.sliceAsBytes(try r.bfs(6000 * lw, 1.0))), + .lat = try r.dev(std.mem.sliceAsBytes(try r.bfs(100 * (lw + 64), 1.0))), + .tok = try r.dev(std.mem.sliceAsBytes(tok)), + .cnt = try r.dev(std.mem.sliceAsBytes(cnt)), + }; + const want = try outputs(&r, .{ .glue = &old }, in); + const got = try outputs(&r, .{ .glue = &new }, in); + const names = [_][]const u8{ "absorb_q 1 row", "expand_v 1 row", "absorb_q 5 rows", "expand_v 5 rows", "absorb_q 70 rows", "expand_v 70 rows", "lat_write 100 rows", "dense 70 rows at 1430, 32-head", "dense 70 rows at 1430, 16-head", "dense 1 row at 1499", "dense 64 rows at 0, 32-head", "sparse 70 rows", "sparse 1 row" }; + var bytes: usize = 0; + for (names, want, got) |name, w, g| { + try check.sameBytes(name, g, w); + bytes += g.len; + } + check.pass("glm-mla-same: {d} GLM MLA launches, {d} output bytes equal 6380c8da's", .{ names.len, bytes }); +} diff --git a/zig/tests/cuda/glm_ref_route.zig b/zig/tests/cuda/glm_ref_route.zig new file mode 100644 index 0000000000..b125ac37b6 --- /dev/null +++ b/zig/tests/cuda/glm_ref_route.zig @@ -0,0 +1,398 @@ +//! glue.select's group-limited routing (Ling's 512 / 8 groups / keep 4 / top 8) vs float64; GLM's bytes vs 6380c8da. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; + +const glue = glm.glue; + +/// One routing problem: fp32 logits [rows][ne], fp32 bias, top_k of the `keep` best of `ng` groups, times `scale`. +const Route = struct { rows: usize, ne: usize, top_k: usize, ng: usize, keep: usize, scale: f32 }; + +/// The kernels' near ties: two choices (or group sums) this close in float64 may rank either way after fp32 rounding. +const tie_expert = std.math.pow(f64, 2, -20); +const tie_group = std.math.pow(f64, 2, -19); + +fn sigmoid(x: f32) f64 { + return 1 / (1 + @exp(-@as(f64, x))); +} + +/// float64 rule: `keep` best top-2 group sums (or `forced`), then top_k in them, lower ids on ties; returns kept bits. +fn refPick(choice: []const f64, q: Route, forced: ?u32, order: []usize, gs_out: []f64) u32 { + const gs = q.ne / q.ng; + var kept: u32 = if (q.ng == 1) 1 else 0; + if (q.ng > 1) { + for (0..q.ng) |g| { + var a: f64 = -std.math.inf(f64); + var b: f64 = -std.math.inf(f64); + for (choice[g * gs ..][0..gs]) |c| { + if (c > a) { + b = a; + a = c; + } else if (c > b) b = c; + } + gs_out[g] = a + b; + } + if (forced) |f| kept = f else for (0..q.keep) |_| { + var best: usize = q.ng; + for (0..q.ng) |g| if (kept >> @intCast(g) & 1 == 0 and (best == q.ng or gs_out[g] > gs_out[best])) { + best = g; + }; + kept |= @as(u32, 1) << @intCast(best); + } + } + var taken: [512]bool = @splat(false); + for (order[0..q.top_k]) |*o| { + var best: usize = q.ne; + for (0..q.ne) |e| { + if (taken[e] or kept >> @intCast(e / gs) & 1 == 0) continue; + if (best == q.ne or choice[e] > choice[best]) best = e; + } + o.* = best; + taken[best] = true; + } + return kept; +} + +/// The k-th and (k+1)-th best group sums' gap: below tie_group the kernel's kept set may legitimately differ. +fn boundaryGap(gsc: []const f64, keep: usize) f64 { + var s: [32]f64 = undefined; + @memcpy(s[0..gsc.len], gsc); + std.mem.sort(f64, s[0..gsc.len], {}, std.sort.desc(f64)); + return s[keep - 1] - s[keep]; +} + +const Tally = struct { rows: usize = 0, group_ties: usize = 0, expert_ties: usize = 0, widest: f64 = 0, worst: f64 = 0 }; + +/// Kernel picks and weights for every row against the float64 rule; near ties counted, anything else fails. +fn against(r: *Rig, q: Route, logits: []const f32, bias: []const f32, got_pick: []const i32, got_wts: []const f32) !Tally { + var t: Tally = .{ .rows = q.rows }; + const choice = try r.a.alloc(f64, q.ne); + const score = try r.a.alloc(f64, q.ne); + const order = try r.a.alloc(usize, q.top_k); + var gsc: [32]f64 = undefined; + const slots = q.top_k + 1; + for (0..q.rows) |row| { + for (score, choice, logits[row * q.ne ..][0..q.ne], bias) |*s, *c, l, b| { + s.* = sigmoid(l); + c.* = s.* + b; + } + const kept = refPick(choice, q, null, order, gsc[0..q.ng]); + const gs = q.ne / q.ng; + var mine: u32 = 0; + for (got_pick[row * slots ..][0..q.top_k]) |id| mine |= @as(u32, 1) << @intCast(@as(usize, @intCast(id)) / gs); + if (mine & ~kept != 0) { // a pick outside the float64 groups: only a near tie at the group boundary may do that + const gap = boundaryGap(gsc[0..q.ng], q.keep); + try check.expect(gap <= tie_group, "row {d}: groups {b} picked where float64 keeps {b} (boundary gap {e})", .{ row, mine, kept, gap }); + t.group_ties += 1; + t.widest = @max(t.widest, gap); + var f = mine; + var g: usize = 0; + while (@popCount(f) < q.keep) : (g += 1) if (kept >> @intCast(g) & 1 != 0) { + f |= @as(u32, 1) << @intCast(g); + }; + _ = refPick(choice, q, f, order, gsc[0..q.ng]); + } + var total: f64 = 0; + for (got_pick[row * slots ..][0..q.top_k]) |id| total += score[@intCast(id)]; + for (0..slots) |i| { + const id: usize = @intCast(got_pick[row * slots + i]); + const wt: f64 = got_wts[row * slots + i]; + if (i == q.top_k) { + try check.expect(id == q.ne and wt == 1, "row {d}: the shared expert last at weight 1", .{row}); + continue; + } + if (id != order[i]) { + const gap = @abs(choice[id] - choice[order[i]]); + try check.expect(gap <= tie_expert, "row {d} pick {d}: {d} where float64 ranks {d} (gap {e})", .{ row, i, id, order[i], gap }); + t.expert_ties += 1; + t.widest = @max(t.widest, gap); + } + const want = score[id] / total * q.scale; + t.worst = @max(t.worst, @abs(wt - want) / (@abs(want) * std.math.pow(f64, 2, -16))); + } + } + return t; +} + +fn f32s(r: *Rig, n: usize, lo: f32, hi: f32) ![]f32 { + const out = try r.a.alloc(f32, n); + for (out) |*v| v.* = lo + r.rng.float(f32) * (hi - lo); + return out; +} + +const Out = struct { pick: []const i32, wts: []const f32 }; + +fn runSelect(r: *Rig, set: glue.Set, q: Route, logits: []const f32, bias: []const f32, norm: bool) !Out { + const n = q.rows * (q.top_k + 1); + const pick = try r.zeros(n * 4); + const wts = try r.zeros(n * 4); + try glue.select(set, r.s, try r.dev(std.mem.sliceAsBytes(logits)), try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, q.rows, q.top_k, q.ne, q.scale, norm, .{ .n = q.ng, .keep = q.keep }); + return .{ .pick = std.mem.bytesAsSlice(i32, @as([]align(4) u8, @alignCast(try r.back(pick, n * 4)))), .wts = std.mem.bytesAsSlice(f32, @as([]align(4) u8, @alignCast(try r.back(wts, n * 4)))) }; +} + +/// Rows built for exact ties: every group a copy of group 0 (all group sums tie), and a group whose top two are equal. +fn tieRows(r: *Rig, q: Route, logits: []f32, bias: []f32) void { + const gs = q.ne / q.ng; + @memset(bias, 0); + for (0..q.rows) |row| { + const l = logits[row * q.ne ..][0..q.ne]; + switch (row % 3) { + 0 => for (0..q.ne) |e| { + l[e] = l[e % gs]; + }, + 1 => { // equal pairs at 2.0 (one lane's in 6, two lanes' in 7) beat 2.2 + 0.0 only if the duplicate counts + for (l) |*v| v.* = -6 + r.rng.float(f32); + for (0..q.ng) |g| { + l[g * gs + 3] = if (g >= 6) 2.0 else 2.2; + l[g * gs + 35] = if (g == 6) 2.0 else 0.0; + if (g == 7) l[g * gs + 7] = 2.0; + } + }, + else => for (0..q.ne) |e| { + l[e] = if (e % 7 == 0) 1.0 else l[e]; + }, + } + } +} + +/// Synthetic Ling-shaped routing (and a few odd shapes) against float64, plus exact-tie rows that must match exactly. +pub fn groups(r: *Rig, set: glue.Set) !void { + const shapes = [_]Route{ + .{ .rows = 300, .ne = 512, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, + .{ .rows = 300, .ne = 256, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, // DeepSeek-V3's + .{ .rows = 100, .ne = 480, .top_k = 6, .ng = 12, .keep = 3, .scale = 1.0 }, // groups of 40 straddle the lanes + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 32, .keep = 1, .scale = 2.5 }, + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 1, .keep = 1, .scale = 2.5 }, // one group: no group pass + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 8, .keep = 8, .scale = 2.5 }, // every group kept + .{ .rows = 100, .ne = 64, .top_k = 8, .ng = 32, .keep = 4, .scale = 2.5 }, // groups of two + .{ .rows = 100, .ne = 16, .top_k = 4, .ng = 8, .keep = 3, .scale = 1.0 }, // groups of two, the tiny config's experts + .{ .rows = 100, .ne = 16, .top_k = 4, .ng = 4, .keep = 2, .scale = 2.5 }, // the tiny config's routing + }; + for (shapes) |q| for ([_]f32{ 0.001, 0.05 }) |bspread| { + const logits = try f32s(r, q.rows * q.ne, -6, 6); + const bias = try f32s(r, q.ne, -bspread, bspread); + const o = try runSelect(r, set, q, logits, bias, true); + const t = try against(r, q, logits, bias, o.pick, o.wts); + r.worst = @max(r.worst, t.worst); + std.debug.print("RESULT group select {d} rows, top {d} of {d} in {d} of {d} groups (bias +-{d}): weights at {e:.3} of 2^-16 relative, {d} group and {d} expert near ties (widest gap {e:.2})\n", .{ q.rows, q.top_k, q.ne, q.keep, q.ng, bspread, t.worst, t.group_ties, t.expert_ties, t.widest }); + try check.expect(t.worst <= 1, "group select weights within 2^-16 of float64", .{}); + }; + const q: Route = .{ .rows = 30, .ne = 512, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }; + const logits = try f32s(r, q.rows * q.ne, -6, 6); + const bias = try r.a.alloc(f32, q.ne); + tieRows(r, q, logits, bias); + const o = try runSelect(r, set, q, logits, bias, true); + const t = try against(r, q, logits, bias, o.pick, o.wts); + try check.expect(t.group_ties == 0 and t.expert_ties == 0, "exact-tie rows: {d} group and {d} expert swaps, want none", .{ t.group_ties, t.expert_ties }); + check.pass("group select: picks follow the float64 rule (exact ties to the lower id, near ties counted), weights within 2^-16", .{}); +} + +/// Ling's swiglu at limit 0 (its unclamped layers) against float64 unclamped, and limit 10 differing where it clamps. +pub fn swigluLimit(r: *Rig, fs: [2]cuda.Function) !void { + for ([_]usize{ 1024, 704 }) |w| { + const rows = 40; + const gu = try r.bfs(rows * 2 * w, 16.0); + const dgu = try r.dev(std.mem.sliceAsBytes(gu)); + const out0 = try r.zeros(rows * w * 2); + const out10 = try r.zeros(rows * w * 2); + const xs = try r.zeros(rows * w / 64 * 4); + try glue.swigluWith(fs, r.s, dgu, out0, xs, 0, rows, w); + try glue.swigluWith(fs, r.s, dgu, out10, xs, 10, rows, w); + const raw0 = try r.back(out0, rows * w * 2); + const raw10 = try r.back(out10, rows * w * 2); + var most: f64 = 0; + var clamped: usize = 0; + for (0..rows) |m| for (0..w) |d| { + const g = bf(gu[m * 2 * w + d]); + const u = bf(gu[m * 2 * w + w + d]); + const want = g / (1 + @exp(-g)) * u; + const got = bf(std.mem.bytesToValue(u16, raw0[2 * (m * w + d) ..][0..2])); + most = @max(most, @abs(got - want) / (@abs(want) * std.math.pow(f64, 2, -7) + 1e-30)); + if (!std.mem.eql(u8, raw0[2 * (m * w + d) ..][0..2], raw10[2 * (m * w + d) ..][0..2])) clamped += 1; + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT swiglu width {d}, limit 0: unclamped float64 at {e:.3} of 2^-7 relative; limit 10 changes {d} of {d}\n", .{ w, most, clamped, rows * w }); + try check.expect(most <= 1 and clamped > 0, "swiglu limit 0 is unclamped, limit 10 clamps", .{}); + } + check.pass("Ling swiglu: the limit is per call, 0 runs unclamped (as the experts' 0), 10 clamps", .{}); +} + +const Old = struct { topk: cuda.Function, swiglu_512: cuda.Function, swiglu_64: cuda.Function }; + +/// GLM's calls through the new glue against 6380c8da's kernels (its glm_glue fatbin) on the same inputs: every byte. +fn glmBytes(r: *Rig, set: glue.Set, prev: Old) !void { + for ([_]usize{ 288, 256, 160 }) |ne| for ([_]bool{ true, false }) |norm| { + const q: Route = .{ .rows = 300, .ne = ne, .top_k = 8, .ng = 1, .keep = 1, .scale = 2.5 }; + const logits = try f32s(r, q.rows * ne, -8, 8); + for (logits[0..ne], 0..) |*v, e| v.* = if (e % 5 == 0) 1.0 else v.*; // exact ties in row 0 + const bias = try f32s(r, ne, -0.05, 0.05); + const n = q.rows * 9; + const dl = try r.dev(std.mem.sliceAsBytes(logits)); + const db = try r.dev(std.mem.sliceAsBytes(bias)); + const pick = try r.zeros(n * 4); + const wts = try r.zeros(n * 4); + var a: cuda.Args = .{}; + inline for (.{ dl, db, pick, wts, q.scale, @as(c_int, @intCast(ne)), @as(c_int, 8), @as(c_int, @intFromBool(norm)) }) |v| a.add(v); + try cuda.launch.launch(prev.topk, .{ .grid = .{ .x = @intCast(q.rows) }, .block = .{ .x = 32 } }, r.s, &a); + const pick_new = try r.zeros(n * 4); + const wts_new = try r.zeros(n * 4); + try glue.select(set, r.s, dl, db, pick_new, wts_new, q.rows, 8, ne, q.scale, norm, .{}); + try check.sameBytes("GLM select picks vs 6380c8da", try r.back(pick_new, n * 4), try r.back(pick, n * 4)); + try check.sameBytes("GLM select weights vs 6380c8da", try r.back(wts_new, n * 4), try r.back(wts, n * 4)); + }; + for ([_]f32{ 10, 0 }) |lim| for ([_]usize{ 2048, 704 }) |w| { + const rows = 40; + const dgu = try r.dev(std.mem.sliceAsBytes(try r.bfs(rows * 2 * w, 16.0))); + const outs = [2]u64{ try r.zeros(rows * w * 2), try r.zeros(rows * w * 2) }; + const xss = [2]u64{ try r.zeros(rows * w / 64 * 4), try r.zeros(rows * w / 64 * 4) }; + const block: usize = if (w % 512 == 0) 512 else 64; + var a: cuda.Args = .{}; + inline for (.{ dgu, outs[0], xss[0], lim, @as(i64, @intCast(w)) }) |v| a.add(v); + try cuda.launch.launch(if (block == 512) prev.swiglu_512 else prev.swiglu_64, .{ .grid = .{ .x = rows, .y = @intCast(w / block) }, .block = .{ .x = if (block == 512) 128 else 64 } }, r.s, &a); + try glue.swiglu(set, r.s, dgu, outs[1], xss[1], lim, rows, w); + try check.sameBytes("GLM swiglu vs 6380c8da", try r.back(outs[1], rows * w * 2), try r.back(outs[0], rows * w * 2)); + try check.sameBytes("GLM swiglu sums vs 6380c8da", try r.back(xss[1], rows * w / 64 * 4), try r.back(xss[0], rows * w / 64 * 4)); + }; + check.pass("GLM select (288, 256, 160 experts, normed and not) and swiglu (limits 10 and 0): every byte equals 6380c8da's kernels", .{}); +} + +const Env = struct { + arena: std.heap.ArenaAllocator, + s: cuda.Stream, + mod: cuda.Module, + ks: glue.Kernels, + prng: std.Random.DefaultPrng, + r: Rig, + + fn deinit(e: *Env) void { + for (e.r.bufs.items) |*b| b.free(); + e.r.bufs.deinit(e.r.gpu.gpa); + e.mod.unload(); + e.s.deinit(); + e.arena.deinit(); + } +}; + +fn env(gpu: Gpu, e: *Env) !void { + e.arena = .init(gpu.gpa); + e.s = try cuda.Stream.init(gpu.d, false); + e.mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + e.ks = try glue.Kernels.load(e.mod); + e.prng = .init(804); + e.r = .{ .gpu = gpu, .a = e.arena.allocator(), .s = e.s, .rng = e.prng.random() }; +} + +/// `glm-route-old `: GLM's select and swiglu against 6380c8da's glm_glue fatbin, byte for byte. +pub fn old(gpu: Gpu, path: []const u8) !void { + var e: Env = undefined; + try env(gpu, &e); + defer e.deinit(); + const bytes = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, e.r.a, .limited(1 << 26), .@"64", null); + var m = try cuda.Module.load(gpu.d, bytes); + defer m.unload(); + try glmBytes(&e.r, .{ .glue = &e.ks }, .{ .topk = try m.function("tf_glm_topk"), .swiglu_512 = try m.function("tf_glm_swiglu_512"), .swiglu_64 = try m.function("tf_glm_swiglu_64") }); +} + +/// `glm-route-ling `: Ling layer 4's rows (HF's fp32 logits) vs float64, HF's routing and exllamav3's. +pub fn ling(gpu: Gpu, path: []const u8) !void { + var e: Env = undefined; + try env(gpu, &e); + defer e.deinit(); + const r = &e.r; + const raw = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, r.a, .limited(1 << 26), .@"4", null); + const hdr = std.mem.bytesAsSlice(i32, raw[0..20]); + const q: Route = .{ .rows = @intCast(hdr[0]), .ne = @intCast(hdr[1]), .top_k = @intCast(hdr[2]), .ng = @intCast(hdr[3]), .keep = @intCast(hdr[4]), .scale = std.mem.bytesToValue(f32, raw[20..24]) }; + var at: usize = 24; + const R = q.rows; + const K = q.top_k; + const parts = .{ q.rows * q.ne, q.ne, R * K, R * K, R * K, R * K }; + var sl: [6][]align(4) u8 = undefined; + inline for (parts, 0..) |n, i| { + sl[i] = @alignCast(raw[at..][0 .. n * 4]); + at += n * 4; + } + try check.expect(at == raw.len, "fixture: {d} bytes, read {d}", .{ raw.len, at }); + const logits = std.mem.bytesAsSlice(f32, sl[0]); + const bias = std.mem.bytesAsSlice(f32, sl[1]); + const hf_pick = std.mem.bytesAsSlice(i32, sl[2]); + const hf_w = std.mem.bytesAsSlice(f32, sl[3]); + const ex_pick = std.mem.bytesAsSlice(i32, sl[4]); + const ex_w = std.mem.bytesAsSlice(f32, sl[5]); + const o = try runSelect(r, .{ .glue = &e.ks }, q, logits, bias, true); + const t = try against(r, q, logits, bias, o.pick, o.wts); + std.debug.print("RESULT Ling L4 vs float64: {d} rows, weights at {e:.3} of 2^-16 relative, {d} group and {d} expert near ties (widest gap {e:.2})\n", .{ R, t.worst, t.group_ties, t.expert_ties, t.widest }); + try check.expect(t.worst <= 1, "Ling weights within 2^-16 of float64", .{}); + var hf_same: usize = 0; + var ex_same: usize = 0; + var d2: f64 = 0; + var n2: f64 = 0; + var hf_most: f64 = 0; + var ex_most: f64 = 0; + for (0..R) |row| { + const mine = o.pick[row * (K + 1) ..][0..K]; + const mw = o.wts[row * (K + 1) ..][0..K]; + if (std.mem.eql(i32, mine, hf_pick[row * K ..][0..K])) { + hf_same += 1; + for (mw, hf_w[row * K ..][0..K]) |a, b| { + d2 += (@as(f64, a) - b) * (@as(f64, a) - b); + n2 += @as(f64, b) * b; + hf_most = @max(hf_most, @abs(@as(f64, a) - b) / @abs(@as(f64, b))); + } + } else std.debug.print(" row {d}: ours {any} HF {any}\n", .{ row, mine, hf_pick[row * K ..][0..K] }); + var all = true; + for (ex_pick[row * K ..][0..K], ex_w[row * K ..][0..K]) |id, w| { + const j = std.mem.indexOfScalar(i32, mine, id) orelse { + all = false; + continue; + }; + ex_most = @max(ex_most, @abs(@as(f64, mw[j]) - w) / @abs(@as(f64, mw[j]))); + } + if (all) ex_same += 1 else std.debug.print(" row {d}: ours {any} exllamav3 {any}\n", .{ row, mine, ex_pick[row * K ..][0..K] }); + } + std.debug.print("RESULT Ling L4 vs HF BailingMoeV3Gate: {d}/{d} rows pick the same experts in the same order; weights rel L2 {e:.3}, worst rel {e:.3}\n", .{ hf_same, R, @sqrt(d2 / n2), hf_most }); + std.debug.print("RESULT Ling L4 vs exllamav3 routing_ds3: {d}/{d} rows pick the same set; weights (half) worst rel {e:.3}\n", .{ ex_same, R, ex_most }); + try check.expect(hf_same == R and ex_same == R and hf_most <= std.math.pow(f64, 2, -16) and ex_most <= std.math.pow(f64, 2, -10), "Ling routing equals HF's and exllamav3's", .{}); + check.pass("Ling layer 4 routing: picks equal HF's and exllamav3's on all {d} rows, weights within float64 tolerance", .{R}); +} + +/// `ling-route-old `: group-limited select (Ling's shapes, odd groups, ties) vs an older glm_glue, bytes. +pub fn lingOld(gpu: Gpu, path: []const u8) !void { + var e: Env = undefined; + try env(gpu, &e); + defer e.deinit(); + const r = &e.r; + const bytes = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, r.a, .limited(1 << 26), .@"64", null); + var m = try cuda.Module.load(gpu.d, bytes); + defer m.unload(); + const prev = try m.function("tf_glm_topk"); + const shapes = [_]Route{ + .{ .rows = 300, .ne = 512, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, + .{ .rows = 300, .ne = 256, .top_k = 8, .ng = 8, .keep = 4, .scale = 2.5 }, + .{ .rows = 100, .ne = 480, .top_k = 6, .ng = 12, .keep = 3, .scale = 1.0 }, + .{ .rows = 100, .ne = 512, .top_k = 8, .ng = 32, .keep = 1, .scale = 2.5 }, + .{ .rows = 100, .ne = 384, .top_k = 8, .ng = 4, .keep = 2, .scale = 2.5 }, + }; + var rows: usize = 0; + for (shapes) |q| for ([_]bool{ true, false }) |norm| for ([_]bool{ false, true }) |ties| { + const logits = try f32s(r, q.rows * q.ne, -6, 6); + const bias = try f32s(r, q.ne, -0.05, 0.05); + if (ties) tieRows(r, q, logits, bias); + const n = q.rows * (q.top_k + 1); + const dl = try r.dev(std.mem.sliceAsBytes(logits)); + const db = try r.dev(std.mem.sliceAsBytes(bias)); + const outs = [4]u64{ try r.zeros(n * 4), try r.zeros(n * 4), try r.zeros(n * 4), try r.zeros(n * 4) }; + var a: cuda.Args = .{}; + inline for (.{ dl, db, outs[0], outs[1], q.scale, @as(c_int, @intCast(q.ne)), @as(c_int, @intCast(q.top_k)), @as(c_int, @intFromBool(norm)), @as(c_int, @intCast(q.ng)), @as(c_int, @intCast(q.keep)) }) |v| a.add(v); + try cuda.launch.launch(prev, .{ .grid = .{ .x = @intCast(q.rows) }, .block = .{ .x = 32 } }, r.s, &a); + try glue.select(.{ .glue = &e.ks }, r.s, dl, db, outs[2], outs[3], q.rows, q.top_k, q.ne, q.scale, norm, .{ .n = q.ng, .keep = q.keep }); + try check.sameBytes("group-limited select picks vs the older build", try r.back(outs[2], n * 4), try r.back(outs[0], n * 4)); + try check.sameBytes("group-limited select weights vs the older build", try r.back(outs[3], n * 4), try r.back(outs[1], n * 4)); + rows += q.rows; + }; + check.pass("ling-route-old: {d} rows of group-limited select (5 shapes, normed and not, tie rows) equal the older kernel's bytes", .{rows}); +} diff --git a/zig/tests/cuda/glm_ref_tests.zig b/zig/tests/cuda/glm_ref_tests.zig index b0e7598990..ad8b7e8e2a 100644 --- a/zig/tests/cuda/glm_ref_tests.zig +++ b/zig/tests/cuda/glm_ref_tests.zig @@ -134,7 +134,7 @@ fn routers(r: *Rig, set: glue.Set) !void { for (bias) |*b| b.* = (r.rng.float(f32) * 2 - 1) * 0.05; const pick = try r.zeros(m * (top_k + 1) * 4); const wts = try r.zeros(m * (top_k + 1) * 4); - try glue.select(set, r.s, full, try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, m, top_k, ne, scale, true); + try glue.select(set, r.s, full, try r.dev(std.mem.sliceAsBytes(bias)), pick, wts, m, top_k, ne, scale, true, .{}); const logits = try r.back(full, m * ne * 4); const got_pick = try r.back(pick, m * (top_k + 1) * 4); const got_wts = try r.back(wts, m * (top_k + 1) * 4); @@ -196,6 +196,8 @@ pub fn run(gpu: Gpu) !void { const set: glue.Set = .{ .glue = &ks }; for ([_]Shape{ .{ .n = 512, .k = 4096 }, .{ .n = 1280, .k = 1536 }, .{ .n = 1024, .k = 128 } }) |sh| try matmuls(&r, set, sh); try routers(&r, set); + try @import("glm_ref_route.zig").groups(&r, set); + try @import("glm_ref_route.zig").swigluLimit(&r, .{ try mod.function("tf_ling_swiglu_512"), try mod.function("tf_ling_swiglu_64") }); try @import("glm_ref_attn.zig").dense(&r, set); try @import("glm_ref_attn.zig").sparse(&r, set); try @import("glm_ref_hc.zig").mixing(&r, set); diff --git a/zig/tests/cuda/ling_ckpt.zig b/zig/tests/cuda/ling_ckpt.zig new file mode 100644 index 0000000000..c52b368d89 --- /dev/null +++ b/zig/tests/cuda/ling_ckpt.zig @@ -0,0 +1,95 @@ +//! `ling-ckpt-check DIR [INDEX]`: a Ling checkpoint against the loader's names, dtypes and shapes, from headers only. +const std = @import("std"); +const core = @import("core"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +const names = ling.names; +const Config = ling.config.Config; +const Tensor = core.checkpoint.Tensor; + +/// A tensor as the loader takes it: names.dims' shape, its format's dtype; side tensors as weights_proj reads them. +fn fits(c: *const Config, name: []const u8, t: names.Tensor, x: Tensor) bool { + const tail = name[std.mem.lastIndexOfScalar(u8, name, '.').? + 1 ..]; + const small = x.dtype == .bf16 or x.dtype == .f32; // Reader.small and Reader.convs take either + switch (names.dims(c, t.role)) { + .values => |n| return small and x.is(x.dtype, &.{n}), + .conv => |s| return small and x.is(x.dtype, &s), + .linear => |nk| { + const n = nk[0]; + const k = nk[1]; + if (t.role == .embed) return x.is(.bf16, &nk); + if (t.role == .router) return c.format(name[0 .. name.len - tail.len - 1]) == .bf16 and small and x.is(x.dtype, &nk); + const f = c.format(name[0 .. name.len - tail.len - 1]); + if (!t.scale) return switch (f) { + .bf16 => x.is(.bf16, &nk), + .fp8_block128 => x.is(.f8_e4m3, &nk) and k % 128 == 0, + .nvfp4 => x.is(.u8, &.{ n, k / 2 }) and k % 16 == 0, + }; + return switch (f) { + .bf16 => false, + .fp8_block128 => x.is(.f32, &.{ (n + 127) / 128, k / 128 }), + .nvfp4 => if (std.mem.eql(u8, tail, "weight_scale")) x.is(.f8_e4m3, &.{ n, k / 16 }) else x.dtype == .f32 and x.numel() == 1, + }; + }, + } +} + +const Tally = struct { + ck: *core.Checkpoint, + c: *const Config, + expected: usize = 0, + mtp: usize = 0, + missing: usize = 0, + wrong: usize = 0, + bytes: u64 = 0, + kinds: [3]usize = @splat(0), // weights, scale tensors, small tensors and convs + + pub fn put(k: *Tally, name: []const u8, t: names.Tensor) !void { + k.expected += 1; + if (t.layer != null and t.layer.? >= k.c.layers) k.mtp += 1; + if (!k.ck.has(name)) { + k.missing += 1; + if (k.missing <= 10) std.debug.print(" missing {s}\n", .{name}); + return; + } + const x = try k.ck.get(name); + k.bytes += x.bytes.len; + k.kinds[if (names.dims(k.c, t.role) != .linear) 2 else if (t.scale) 1 else 0] += 1; + if (fits(k.c, name, t, x)) return; + k.wrong += 1; + if (k.wrong <= 10) std.debug.print(" {s}: {t} {any}, not what the loader takes\n", .{ name, x.dtype, x.shape[0..x.rank] }); + } +}; + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var why: ling.config.Why = .{}; + var opened = ling.weights.open(gpu.gpa, gpu.io, dir, &why) catch |err| { + std.debug.print("FAIL {s}: {s}\n", .{ dir, why.text() }); + return err; + }; + defer opened.ck.close(); + const c = &opened.c; + var stored: usize = 0; + for (opened.ck.files.items) |f| stored += f.names.count(); + var k: Tally = .{ .ck = &opened.ck, .c = c }; + try names.all(c, &k); + const unknown = opened.ck.unused(); + // the index's total_size, or without an index the files' data regions: every byte belongs to a checked tensor + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const index = if (rest.len > 1) rest[1] else try std.fs.path.join(arena.allocator(), &.{ dir, "model.safetensors.index.json" }); + var total: u64 = 0; + if (std.Io.Dir.cwd().readFileAlloc(gpu.io, index, arena.allocator(), .limited(1 << 26))) |text| { + const parsed = try std.json.parseFromSliceLeaky(std.json.Value, arena.allocator(), text, .{}); + total = @intCast(parsed.object.get("metadata").?.object.get("total_size").?.integer); + } else |_| for (opened.ck.files.items) |f| { + total += f.map.memory.len - f.data; + } + std.debug.print("RESULT ling-ckpt-check {t}: {d} shards, {d} tensors stored, {d} expected ({d} of the MTP layer): {d} weights, {d} scale tensors, {d} small tensors and convs\n", .{ c.quant, opened.ck.files.items.len, stored, k.expected, k.mtp, k.kinds[0], k.kinds[1], k.kinds[2] }); + std.debug.print("RESULT {d} missing, {d} stored but not expected, {d} of another dtype or shape; {d} bytes against {d} (the index's total_size, else the files' data regions)\n", .{ k.missing, unknown, k.wrong, k.bytes, total }); + try check.expect(k.missing == 0 and unknown == 0 and k.wrong == 0 and stored == k.expected and k.bytes == total, "every stored tensor is one the loader reads, at its dtype and shape, and nothing is missing", .{}); + check.pass("ling-ckpt-check: all {d} tensors' names, dtypes and shapes as the loader takes them, {d} bytes (headers only)", .{ stored, k.bytes }); +} diff --git a/zig/tests/cuda/ling_clamp.zig b/zig/tests/cuda/ling_clamp.zig new file mode 100644 index 0000000000..e2b3b24ee9 --- /dev/null +++ b/zig/tests/cuda/ling_clamp.zig @@ -0,0 +1,269 @@ +//! Ling's SwiGLU order, min(silu(g), L) * clip(u, -L, L): glue and grouped experts vs float64, old kernels' bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const glm = @import("glm_cuda"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const lw = @import("ling_weights_tests.zig"); + +const Rig = qr.RigT; +const h = qr.helpers; +const W = ling.weights; +const ex = cuda.experts; +const glue = glm.glue; +const Config = ling.config.Config; +const two7 = std.math.pow(f64, 2, -7); + +/// Ling's activation in float64 on bf16 values: silu(g) as `s`, then min(s, L) and clip(u, -L, L); L 0 for none. +fn act(s: f64, u: f64, limit: f64) f64 { + if (limit == 0) return s * u; + return @min(s, limit) * std.math.clamp(u, -limit, limit); +} + +fn silu(g: f64) f64 { + return g / (1 + @exp(-g)); +} + +/// The glue's tf_ling_swiglu at L 10 and 2 on gates and ups straddling L vs float64; GLM's order beside it. +fn glueOrder(r: *Rig, fs: [2]cuda.Function, set: glue.Set) !void { + for ([_]f32{ 10, 2 }) |lim| for ([_]usize{ 1024, 704 }) |w| { + const rows = 40; + const gu = try r.bfs(rows * 2 * w, 1.5 * lim); + const dgu = try r.dev(std.mem.sliceAsBytes(gu)); + const outs = [2]u64{ try r.zeros(rows * w * 2), try r.zeros(rows * w * 2) }; + const xs = try r.zeros(rows * w / 64 * 4); + try glue.swigluWith(fs, r.s, dgu, outs[0], xs, lim, rows, w); + try glue.swiglu(set, r.s, dgu, outs[1], xs, lim, rows, w); + const got = try r.back(outs[0], rows * w * 2); + const before = try r.back(outs[1], rows * w * 2); + var most: f64 = 0; + var above: usize = 0; + var wide: usize = 0; + var differ: usize = 0; + for (0..rows) |m| for (0..w) |d| { + const g = h.bfOf(gu[m * 2 * w + d]); + const u = h.bfOf(gu[m * 2 * w + w + d]); + const want = act(silu(g), u, lim); + const i = m * w + d; + const v = h.bfOf(std.mem.bytesToValue(u16, got[2 * i ..][0..2])); + most = @max(most, @abs(v - want) / (@abs(want) * two7 + 1e-30)); + above += @intFromBool(silu(g) > lim); + wide += @intFromBool(@abs(u) > lim); + differ += @intFromBool(!std.mem.eql(u8, got[2 * i ..][0..2], before[2 * i ..][0..2])); + }; + r.worst = @max(r.worst, most); + std.debug.print("RESULT glue L {d}, width {d}: float64 at {d:.3} of 2^-7; silu(g) > L on {d}, |u| > L on {d} of {d}; {d} differ from GLM's order\n", .{ lim, w, most, above, wide, rows * w, differ }); + try check.expect(most <= 1 and above > 0 and wide > 0, "tf_ling_swiglu at L {d}: Ling's order within 2^-7 of float64", .{lim}); + // bf16(silu(g)) is 10 for every bf16 g >= 10, so the two orders agree on bytes at 10 and part at 2 + try check.expect(if (lim == 10) differ == 0 else differ > 0, "L {d}: {d} bytes differ from GLM's order", .{ lim, differ }); + }; +} + +/// A plan for `pairs` pairs, pair p on expert lw.pick(p), at the given tile. +fn plan(r: *Rig, kit: *const lw.Kit, pairs: usize, slots: usize, count: usize, tile: usize) !cuda.grouped.Plan { + const pk = try r.a.alloc(i32, pairs); + for (pk, 0..) |*e, p| e.* = lw.pick(p, slots, count); + const p: cuda.grouped.Plan = .{ .members = try r.zeros(pairs * 4), .items = try r.zeros(cuda.grouped.maxItems(pairs, count, 16) * 12), .counts = try r.zeros(8), .rank = try r.zeros(4), .hist = try r.zeros(4) }; + try kit.router.route(r.s, try r.dev(std.mem.sliceAsBytes(pk)), pairs, count, tile, p); + return p; +} + +/// The grouped gate-up with Layer.after at L 10 and 2, decode and prompt rows, vs float64 from the file. +fn expertsOrder(r: *Rig, kit: *const lw.Kit, ck: *core.Checkpoint, c: *const Config, i: u32, l: ex.Layer) !void { + const slots = @min(4, l.experts); + const ni = l.width; + const d = l.dims; + var b: [2][192]u8 = undefined; + const e0: u32 = @intCast(lw.pick(0, slots, l.experts)); + const g0 = try lw.Mat.read(ck, c, try ling.names.moduleName(&b[0], i, .expert_gate, e0), ni, d); + const probe = try lw.reference(r.a, g0, try r.bfs(d, 1), 1); + var ss: f64 = 0; + for (probe.want) |v| ss += v * v; + const rms = @sqrt(ss / @as(f64, @floatFromInt(ni))); + for ([_]bool{ false, true }) |prompt| for ([_]f32{ 10, 2 }) |lim| { + const rows: usize = if (prompt) 24 else 3; + const pairs = rows * slots; + const tile: usize = if (prompt) ex.promptTile(l.format) else 16; + const p = try plan(r, kit, pairs, slots, l.experts, tile); + const x = try r.bfs(rows * d, @floatCast(std.math.pow(f64, 2, @round(std.math.log2(1.2 * lim / rms))))); + var lc = l; + lc.limit = lim; + lc.after = true; + const out = try r.zeros(pairs * ni * 2); + const in: ex.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(x)), .stride = d, .slots = slots }; + const items = cuda.grouped.maxItems(pairs, l.experts, tile); + if (prompt) try kit.grouped().prompt(r.s, .up, in, lc, p, out, items, -1) else try kit.grouped().decode(r.s, .up, in, lc, p, out, items, -1); + const want = try r.a.alloc(f64, pairs * ni); + const tol = try r.a.alloc(f64, pairs * ni); + var above: usize = 0; + for (0..pairs) |q| { + const e: u32 = @intCast(lw.pick(q, slots, l.experts)); + const g = try lw.Mat.read(ck, c, try ling.names.moduleName(&b[0], i, .expert_gate, e), ni, d); + const u = try lw.Mat.read(ck, c, try ling.names.moduleName(&b[1], i, .expert_up, e), ni, d); + const gr = try lw.reference(r.a, g, x[(q / slots) * d ..][0..d], 1); + const ur = try lw.reference(r.a, u, x[(q / slots) * d ..][0..d], 1); + for (0..ni) |j| { + const gv = h.bfOf(h.bf16Of(@floatCast(gr.want[j]))); + const uv = h.bfOf(h.bf16Of(@floatCast(ur.want[j]))); + const s = h.bfOf(h.bf16Of(@floatCast(silu(gv)))); + want[q * ni + j] = act(s, uv, lim); + above += @intFromBool(s > lim); + // the bf16 steps (2^-6), then each product's bound: silu' under 1.1, min and clip 1-Lipschitz + tol[q * ni + j] = @abs(want[q * ni + j]) * std.math.pow(f64, 2, -6) + 1.1 * @min(@abs(uv), lim) * gr.tol[j] + @min(@abs(s), lim) * ur.tol[j] + 1e-30; + } + } + var name: [128]u8 = undefined; + try r.close(try std.fmt.bufPrint(&name, "layer {d} {t} gate-up L {d} ({s}), silu(g) > L on {d} of {d}", .{ i, l.format, lim, if (prompt) "prompt" else "decode", above, pairs * ni }), out, false, want, tol); + try check.expect(above > 0, "some gates past L", .{}); + }; +} + +/// `ling-clamp DIR`: the glue's Ling order vs float64, then the clamped MoE layer's grouped experts (FP8 and NVFP4). +pub fn run(gpu: check.Gpu, dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0xc1a); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var gm = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer gm.unload(); + const ks = try glue.Kernels.load(gm); + try glueOrder(&r, .{ try gm.function("tf_ling_swiglu_512"), try gm.function("tf_ling_swiglu_64") }, .{ .glue = &ks }); + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + const last: u32 = o.c.layers - 1; + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{last}, .{ .globals = false }); + defer w.deinit(); + const moe = w.layers[0].ffn.moe; + try check.expect(moe.shared_limit == ling.limits(&w.config, last).shared, "shared limit from the config", .{}); + switch (moe.routed) { + .grouped => |l| { + try check.expect(l.after and l.limit == ling.limits(&w.config, last).routed, "loader: Ling's order at the config's limit", .{}); + try expertsOrder(&r, &kit, &o.ck, &w.config, last, l); + }, + .lanes => std.debug.print("RESULT {t}: routed experts run as lanes, through the glue checked above\n", .{o.c.quant}), + } + check.pass("ling-clamp {t}: min(silu(g), L) * clip(u, -L, L) within float64 bounds (worst {d:.3}), GLM's order beside it", .{ o.c.quant, r.worst }); +} + +/// An old fatbin's instance, launched with the arguments and grid the host gave it before Layer.after existed. +const Old = struct { f: cuda.Function, threads: u32, smem: u32 = 0, per: usize, cap: usize }; + +fn launchOld(r: *Rig, k: Old, in: ex.Rows, l: ex.Layer, gate: bool, p: cuda.grouped.Plan, out: u64, units: usize) !void { + var a: cuda.Args = .{}; + a.add(in.x); + a.add(@as(i32, @intCast(in.stride))); + a.add(@as(i32, @intCast(if (gate) in.slots else 0))); + a.add(if (gate) l.up else l.down); + a.add(if (gate) l.up_scale else l.down_scale); + a.add(@as(i32, @intCast((if (gate) l.dims else l.width) / 32))); + a.add(@as(i32, @intCast((if (gate) l.width else l.dims) / 32))); + for ([_]u64{ p.items, p.counts, p.members, out }) |v| a.add(v); + a.add(@as(i32, @intCast(if (gate) l.width else l.dims))); + a.add(if (gate) l.limit else @as(f32, 0.0)); + a.add(@as(i32, -1)); + const grid = @min((units + k.per - 1) / k.per, k.cap); + try cuda.launch.launch(k.f, .{ .grid = .{ .x = @intCast(grid) }, .block = .{ .x = k.threads }, .shared = k.smem }, r.s, &a); +} + +fn oldFn(m: cuda.Module, name: [:0]const u8, threads: u32, smem: u32, per: usize, sms: usize, resident: bool) !Old { + const f = try m.function(name); + if (smem > 0) try f.allowDynamicShared(smem); + return .{ .f = f, .threads = threads, .smem = smem, .per = per, .cap = if (resident) @max(1, try f.occupancy(128, 0)) * sms else std.math.maxInt(usize) }; +} + +/// `ling-clamp-old DIR FATBINS`: each grouped instance this build kept, via experts.Grouped, vs FATBINS' kernel. +pub fn old(gpu: check.Gpu, dir: []const u8, fatbins: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x01d); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = try W.open(gpu.gpa, gpu.io, dir, &why); + defer o.ck.close(); + const last: u32 = o.c.layers - 1; + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{last}, .{ .globals = false }); + defer w.deinit(); + const l = w.layers[0].ffn.moe.routed.grouped; + const fp8 = l.format == .fp8g; + const path = try std.fs.path.join(r.a, &.{ fatbins, if (fp8) "fp8_experts.fatbin" else "nvfp4_experts.fatbin" }); + const bytes = try std.Io.Dir.cwd().readFileAllocOptions(gpu.io, path, r.a, .limited(1 << 26), .@"64", null); + var m = try cuda.Module.load(gpu.d, bytes); + defer m.unload(); + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + const F8 = ex.Fp8Experts.symbols; + const F4 = ex.Nvfp4Experts.symbols; + const p4 = ex.Nvfp4Experts.promptSmem; + // [decode, prompt] x [gate-up, down_f32, down_bf16, relu^2]; null where the format has no such instance + const ks: [2][4]?Old = if (fp8) .{ + .{ try oldFn(m, F8.gate_up, 128, 0, 4, sms, true), try oldFn(m, F8.down_f32, 128, 0, 4, sms, true), try oldFn(m, F8.down_bf16, 128, 0, 4, sms, true), null }, + .{ try oldFn(m, F8.p_gate_up, 128, 0, 1, sms, true), try oldFn(m, F8.p_down_f32, 128, 0, 1, sms, true), try oldFn(m, F8.p_down_bf16, 128, 0, 1, sms, true), null }, + } else .{ + .{ try oldFn(m, F4.gate_up, 128, 0, 4, sms, true), try oldFn(m, F4.down_f32, 128, 0, 4, sms, true), try oldFn(m, F4.down_bf16, 128, 0, 4, sms, true), try oldFn(m, F4.up_relu2, 128, 0, 4, sms, true) }, + .{ try oldFn(m, F4.prompt[1], 128, p4(2, 4), 1, sms, false), null, try oldFn(m, F4.prompt[2], 256, p4(1, 8), 1, sms, false), try oldFn(m, F4.prompt[0], 256, p4(1, 8), 1, sms, false) }, + }; + const slots = @min(4, l.experts); + var n: usize = 0; + var same_at_10: usize = 0; + for ([_]bool{ false, true }, 0..) |prompt, pi| { + const rows: usize = if (prompt) 40 else 3; + const pairs = rows * slots; + const tile: usize = if (prompt) ex.promptTile(l.format) else 16; + const p = try plan(&r, &kit, pairs, slots, l.experts, tile); + const items = cuda.grouped.maxItems(pairs, l.experts, tile); + const up_in: ex.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(try r.bfs(rows * l.dims, 4))), .stride = l.dims, .slots = slots }; + const down_in: ex.Rows = .{ .x = try r.dev(std.mem.sliceAsBytes(try r.bfs(pairs * l.width, 4))), .stride = l.width }; + for (ks[pi], 0..) |maybe, ki| if (maybe) |k| for ([_]f32{ 0, 10, 2 }) |lim| for ([_]bool{ false, true }) |after| { + if (ki != 0 and (lim != 0 or after)) continue; + var lc = l; + lc.limit = if (ki == 0) lim else 0; + lc.after = after; + lc.relu2 = ki == 3; + const gate = ki == 0 or ki == 3; + const epi: ex.Epi = if (gate) .up else if (ki == 1) .down_f32 else .down_bf16; + const n_out = if (gate) l.width else l.dims; + const len = pairs * n_out * @as(usize, if (ki == 1) 4 else 2); + const outs = [2]u64{ try r.zeros(len), try r.zeros(len) }; + const in = if (gate) up_in else down_in; + if (prompt) try kit.grouped().prompt(r.s, epi, in, lc, p, outs[0], items, -1) else try kit.grouped().decode(r.s, epi, in, lc, p, outs[0], items, -1); + var lo = lc; + lo.after = false; + const nb = n_out / 32; + const units = if (!prompt) items * nb else if (fp8) items * (nb / @as(usize, if (gate) 2 else 4)) else items * ((nb + @as(usize, if (ki == 0) 4 else 8) - 1) / @as(usize, if (ki == 0) 4 else 8)); + try launchOld(&r, k, in, lo, gate, p, outs[1], units); + const got = try r.back(outs[0], len); + const want = try r.back(outs[1], len); + if (!after or lim == 0) { + try check.sameBytes("grouped experts vs the old fatbin", got, want); + n += 1; + } else if (lim == 10) { + try check.sameBytes("Ling's order at L 10 vs the old order", got, want); + same_at_10 += 1; + } else try check.expect(!std.mem.eql(u8, got, want), "Ling's order at L 2 differs from the old order", .{}); + }; + } + check.pass("ling-clamp-old {t}: {d} calls of the kept instances equal {s}'s bytes; Ling's order at L 10 equals them {d} times, at L 2 differs", .{ l.format, n, fatbins, same_at_10 }); +} diff --git a/zig/tests/cuda/ling_faults.zig b/zig/tests/cuda/ling_faults.zig new file mode 100644 index 0000000000..b380ed0cad --- /dev/null +++ b/zig/tests/cuda/ling_faults.zig @@ -0,0 +1,115 @@ +//! `ling-init-faults DIR [--mtp] [--every N]`: the backend opened on the fake driver, each call and allocation failing. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +const fake = ling.fake_driver; +const Dbg = std.heap.DebugAllocator(.{ .safety = true }); // leak checks in release builds too + +/// Opened (the engine, its sequence, the lane backend), with `keep` a state kept and dropped, and closed again. +fn once(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, keep: bool) !void { + const l = try open(gpa, io, ctx, dir, mtp); + defer l.deinit(l.ctx); + if (keep) l.target.dropState(try l.target.keepState()); +} + +fn open(gpa: std.mem.Allocator, io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool) !ling.native.Loaded { + return ling.native.open(gpa, io, ctx, dir, null, .{ .context = 256, .drafts = mtp }); +} + +/// What one failed open left: the error, host leaks, the fake driver's live, bad releases, frees with work pending. +const Left = struct { err: ?anyerror, leaked: bool, live: usize, bad: usize, early: usize, stuck: usize, pending: bool }; + +fn failed(io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, keep: bool, call: ?u64, alloc: ?usize) Left { + fake.reset(call, cuda.abi.error_not_ready); + var dbg: Dbg = .init; + var fa = std.testing.FailingAllocator.init(dbg.allocator(), .{ .fail_index = alloc orelse std.math.maxInt(usize) }); + const err: ?anyerror = if (once(fa.allocator(), io, ctx, dir, mtp, keep)) |_| null else |e| e; + const live = fake.live(); + const bad = fake.badReleases(); + const leaked = dbg.deinit() == .leak; + const out: Left = .{ .err = err, .leaked = leaked, .live = live, .bad = bad, .early = fake.freedEarly(), .stuck = fake.failedDrains(), .pending = fake.failedPending() }; + _ = fake.sweep(); // what this point left held does not count against the next + return out; +} + +const Sticky = struct { tried: usize = 0, bad: usize = 0, held: usize = 0, pending: usize = 0 }; + +/// Call k and every one after it failing (the open alone): no free with work pending, leaks only past a failed drain. +fn sticky(io: std.Io, ctx: *const cuda.Context, dir: []const u8, mtp: bool, calls: u64, every: usize) Sticky { + fake.mode(true, true); + defer fake.mode(true, false); + var out: Sticky = .{}; + var k: u64 = 1; + while (k <= calls) : (k += every) { + const l = failed(io, ctx, dir, mtp, false, k, null); + out.tried += 1; + out.held += @intFromBool(l.live != 0); + out.pending += @intFromBool(l.pending); + if (l.err == null or l.leaked or l.bad != 0 or l.early != 0 or (l.live != 0 and l.stuck == 0)) { + out.bad += 1; + if (out.bad <= 8) std.debug.print(" driver calls {d}.. failing: {?t}, host leak {}, {d} held ({d} failed drains), {d} bad releases, {d} freed while pending\n", .{ k, l.err, l.leaked, l.live, l.stuck, l.bad, l.early }); + } + } + return out; +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 1) return error.MissingArgument; + const dir = rest[0]; + var mtp = false; + var every: usize = 1; + var i: usize = 1; + while (i < rest.len) : (i += 1) { + if (std.mem.eql(u8, rest[i], "--mtp")) mtp = true; + if (std.mem.eql(u8, rest[i], "--every") and i + 1 < rest.len) { + every = try std.fmt.parseInt(usize, rest[i + 1], 10); + i += 1; + } + } + const d = fake.driver(); + const ctx = fake.context(&d); + fake.mode(true, false); // async copies and launches pending until a drain: no free may come before one, ever + defer fake.mode(false, false); + // a clean open and kept state first: their calls and host allocations are the failure points + fake.reset(null, cuda.abi.success); + var dbg: Dbg = .init; + var counter = std.testing.FailingAllocator.init(dbg.allocator(), .{}); + const clean = try open(counter.allocator(), gpu.io, &ctx, dir, mtp); + const opened = fake.count(); // the open's own calls: close's are not failure points + clean.target.dropState(try clean.target.keepState()); + const calls = fake.count(); + const allocs = counter.allocations; + clean.deinit(clean.ctx); + try check.expect(dbg.deinit() == .ok and fake.live() == 0 and fake.badReleases() == 0 and fake.freedEarly() == 0, "a clean open and close leaves nothing", .{}); + var bad: usize = 0; + var tried: usize = 0; + var pending: usize = 0; + var k: u64 = 1; + while (k <= calls) : (k += every) { + const l = failed(gpu.io, &ctx, dir, mtp, true, k, null); + tried += 1; + pending += @intFromBool(l.pending); + if (l.err == null or !fake.fired() or l.leaked or l.live != 0 or l.bad != 0 or l.early != 0) { + bad += 1; + if (bad <= 8) std.debug.print(" driver call {d} failing: {?t}, host leak {}, {d} held, {d} bad releases, {d} freed while pending\n", .{ k, l.err, l.leaked, l.live, l.bad, l.early }); + } + } + var a: usize = 0; + while (a < allocs) : (a += every) { + const l = failed(gpu.io, &ctx, dir, mtp, true, null, a); + tried += 1; + if (l.err == null or l.leaked or l.live != 0 or l.bad != 0 or l.early != 0) { + bad += 1; + if (bad <= 16) std.debug.print(" host allocation {d} failing: {?t}, host leak {}, {d} held, {d} bad releases\n", .{ a, l.err, l.leaked, l.live, l.bad }); + } + } + const st = sticky(gpu.io, &ctx, dir, mtp, opened, every); + bad += st.bad; + tried += st.tried; + fake.reset(null, cuda.abi.success); + std.debug.print("RESULT ling-init-faults{s}: {d} driver calls ({d} the open's) and {d} host allocations in a clean open and kept state; {d} failure points tried ({d} with a transfer pending), {d} left something behind; failing from a call on: {d} points ({d} with a transfer pending), {d} leaking after a failed drain\n", .{ if (mtp) " --mtp" else "", calls, opened, allocs, tried, pending + st.pending, bad, st.tried, st.pending, st.held }); + try check.expect(bad == 0, "every failing driver call and host allocation in the open is cleaned up, nothing freed with work pending", .{}); + check.pass("ling-init-faults{s}: {d} failure points, each refused with nothing freed before a drain, held, leaked or freed twice", .{ if (mtp) " --mtp" else "", tried }); +} diff --git a/zig/tests/cuda/ling_forward_rig.zig b/zig/tests/cuda/ling_forward_rig.zig new file mode 100644 index 0000000000..37927d3dc8 --- /dev/null +++ b/zig/tests/cuda/ling_forward_rig.zig @@ -0,0 +1,130 @@ +//! What the Ling forward checks share: the graph's kit on one stream, runs that return logits, distances and flips. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +pub const F = ling.forward; +pub const W = ling.weights; + +pub fn bfOf(u: u16) f64 { + return @as(f32, @bitCast(@as(u32, u) << 16)); +} + +pub const Rig = struct { + gpu: check.Gpu, + a: std.mem.Allocator, + s: cuda.Stream, + m: *const F.Modules, + g: *const F.Graph, + bufs: std.ArrayList(cuda.DeviceBuffer) = .empty, + kda_wide: bool = true, // decode windows on KDA's three-kernel chain (false: the fused one) + + pub fn deinit(r: *Rig) void { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(r.gpu.gpa); + } + + pub fn zeros(r: *Rig, len: usize) !u64 { + const b = try cuda.DeviceBuffer.alloc(r.gpu.d, @max(len, 16)); + try r.bufs.append(r.gpu.gpa, b); + try b.fill8(0, null); + return b.ptr; + } + + pub fn back(r: *Rig, ptr: u64, len: usize) ![]u8 { + const out = try r.a.alloc(u8, len); + try r.s.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = r.gpu.d, .ptr = ptr, .len = len }; + try b.download(0, out); + return out; + } + + pub fn fwd(r: *const Rig, b: *const F.Buffers, picks: ?u64) F.Forward { + return .{ .m = r.m, .g = r.g, .s = r.s, .b = b, .picks = picks, .kda_wide = r.kda_wide }; + } + + /// One forward of `segs` on `b`; its logits rows (bf16 bits) and, with `picks`, every MoE layer's picks. + pub fn run(r: *Rig, b: *const F.Buffers, segs: []const F.Seg, ids: []const i32, picks: ?u64) ![]u16 { + try r.fwd(b, picks).run(segs, ids); + const v: usize = r.g.w.config.vocab; + const raw = try r.back(b.logits, ids.len * v * 2); + const out = try r.a.alloc(u16, ids.len * v); + @memcpy(std.mem.sliceAsBytes(out), raw); + return out; + } + + /// A fresh sequence of `capacity` positions. + pub fn state(r: *Rig, capacity: usize) !*F.State { + const st = try r.a.create(F.State); + st.* = try F.State.init(r.gpu.gpa, r.gpu.d, r.g, capacity); + try st.reset(r.s); + return st; + } +}; + +/// Picks [layers][rows][top_k + 1] (i32) as each row's top-k sorted, [layers][rows][top_k]. +pub fn sortedPicks(a: std.mem.Allocator, raw: []const u8, layers: usize, rows: usize, top: usize) ![]i32 { + const all = std.mem.bytesAsSlice(i32, raw); + const out = try a.alloc(i32, layers * rows * top); + for (0..layers * rows) |i| { + const dst = out[i * top ..][0..top]; + for (dst, 0..) |*d, k| d.* = all[i * (top + 1) + k]; + std.mem.sort(i32, dst, {}, std.sort.asc(i32)); + } + return out; +} + +/// Rows whose sorted picks differ from the reference's in any layer (`want` rows from `w0` of its `wrows`). +pub fn flips(got: []const i32, want: []const i32, layers: usize, rows: usize, w0: usize, wrows: usize, top: usize, out: []bool) usize { + var n: usize = 0; + for (0..layers) |l| for (0..rows) |t| { + const g = got[(l * rows + t) * top ..][0..top]; + const e = want[(l * wrows + w0 + t) * top ..][0..top]; + if (!std.mem.eql(i32, g, e)) { + n += 1; + out[t] = true; + } + }; + return n; +} + +/// bf16 rows against float64 rows: relative L2, the same excluding `skip` rows, the worst row's, argmax agreement. +pub const Dist = struct { rel: f64, rel_kept: f64, worst_row: f64, argmax: usize }; + +pub fn distance(got: []const u16, want: []const f64, rows: usize, width: usize, skip: []const bool) Dist { + var dd: f64 = 0; + var ww: f64 = 0; + var kd: f64 = 0; + var kw: f64 = 0; + var worst: f64 = 0; + var agree: usize = 0; + for (0..rows) |t| { + var rd: f64 = 0; + var rw: f64 = 0; + var ga: usize = 0; + var wa: usize = 0; + for (0..width) |j| { + const g = bfOf(got[t * width + j]); + const e = want[t * width + j]; + rd += (g - e) * (g - e); + rw += e * e; + if (g > bfOf(got[t * width + ga])) ga = j; + if (e > want[t * width + wa]) wa = j; + } + dd += rd; + ww += rw; + if (!skip[t]) { + kd += rd; + kw += rw; + } + worst = @max(worst, @sqrt(rd / rw)); + agree += @intFromBool(ga == wa); + } + return .{ .rel = @sqrt(dd / ww), .rel_kept = if (kw > 0) @sqrt(kd / kw) else 0, .worst_row = worst, .argmax = agree }; +} + +/// Two runs' logits rows equal, bit for bit. +pub fn same(what: []const u8, got: []const u16, want: []const u16) !void { + try check.sameBytes(what, std.mem.sliceAsBytes(got), std.mem.sliceAsBytes(want)); +} diff --git a/zig/tests/cuda/ling_forward_synth.zig b/zig/tests/cuda/ling_forward_synth.zig new file mode 100644 index 0000000000..73a927ad02 --- /dev/null +++ b/zig/tests/cuda/ling_forward_synth.zig @@ -0,0 +1,178 @@ +//! `ling-forward-synth DIR REF`: a synthetic checkpoint's logits vs ling_forward.py's; cuts, windows, batches exact. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; +const Rig = rig.Rig; +const P = 80; // prompt rows: KDA's prompt takes the wide chain (64 rows or more) + +/// ling_forward.py's file: ids, float64 and bf16-point logits [T][V], the picks it routed on [layers][T][top]. +pub const Ref = struct { T: usize, V: usize, layers: usize, top: usize, forced: bool, gap: f64, ids: []i32, l64: []f64, lbf: []f64, picks: []i32 }; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +pub fn readRef(a: std.mem.Allocator, io: std.Io, path: []const u8) !Ref { + const bytes = try std.Io.Dir.cwd().readFileAlloc(io, path, a, .limited(1 << 30)); + if (bytes.len < 36 or !std.mem.eql(u8, bytes[0..8], "TFLINGF2")) return error.BadFixture; + var at: usize = 8; + const dims = try take(u32, a, bytes, &at, 5); + const gap = (try take(f64, a, bytes, &at, 1))[0]; + const T: usize = dims[0]; + const V: usize = dims[1]; + const L: usize = dims[2]; + const K: usize = dims[3]; + const r: Ref = .{ .T = T, .V = V, .layers = L, .top = K, .forced = dims[4] != 0, .gap = gap, .ids = try take(i32, a, bytes, &at, T), .l64 = try take(f64, a, bytes, &at, T * V), .lbf = try take(f64, a, bytes, &at, T * V), .picks = try take(i32, a, bytes, &at, L * T * K) }; + if (at != bytes.len) return error.BadFixture; + return r; +} + +/// The serial run: the prompt in one chunk, then one decode row at a time; logits rows and sorted picks. +fn serial(r: *Rig, pb: *const F.Buffers, db: *const F.Buffers, st: *F.State, ids: []const i32, nm: usize) !struct { logits: []u16, picks: []i32 } { + const c = &r.g.w.config; + const V: usize = c.vocab; + const T = ids.len; + const top: usize = c.topk; + const logits = try r.a.alloc(u16, T * V); + const picks = try r.a.alloc(i32, nm * T * top); + const pp = try r.zeros(nm * P * (top + 1) * 4); + @memcpy(logits[0 .. P * V], try r.run(pb, &.{.{ .st = st, .row0 = 0, .rows = P }}, ids[0..P], pp)); + const sp = try rig.sortedPicks(r.a, try r.back(pp, nm * P * (top + 1) * 4), nm, P, top); + for (0..nm) |l| @memcpy(picks[l * T * top ..][0 .. P * top], sp[l * P * top ..][0 .. P * top]); + try r.fwd(pb, null).commit(st, P); + const dp = try r.zeros(nm * (top + 1) * 4); + for (P..T) |t| { + @memcpy(logits[t * V ..][0..V], try r.run(db, &.{.{ .st = st, .row0 = 0, .rows = 1 }}, ids[t .. t + 1], dp)); + const one = try rig.sortedPicks(r.a, try r.back(dp, nm * (top + 1) * 4), nm, 1, top); + for (0..nm) |l| @memcpy(picks[(l * T + t) * top ..][0..top], one[l * top ..][0..top]); + try r.fwd(db, null).commit(st, 1); + } + return .{ .logits = logits, .picks = picks }; +} + +/// A committed sequence's caches (KDA states and conv windows, MLA rows below pos) equal another's, byte for byte. +fn sameState(r: *Rig, what: []const u8, a: *const F.State, b: *const F.State) !void { + var name: [96]u8 = undefined; + try check.expect(a.pos == b.pos, "{s}: positions {d} and {d}", .{ what, a.pos, b.pos }); + for (0..a.kda_layers) |li| { + try check.sameBytes(try std.fmt.bufPrint(&name, "{s}: KDA layer {d} state", .{ what, li }), try r.back(a.recAt(a.cur[li], li), a.rec_layer), try r.back(b.recAt(b.cur[li], li), b.rec_layer)); + try check.sameBytes(try std.fmt.bufPrint(&name, "{s}: KDA layer {d} conv window", .{ what, li }), try r.back(a.convAt(li), 3 * a.convWidth() * 2), try r.back(b.convAt(li), 3 * b.convWidth() * 2)); + } + for (0..a.mla_layers) |mi| try check.sameBytes(try std.fmt.bufPrint(&name, "{s}: MLA layer {d} cache rows", .{ what, mi }), try r.back(a.cacheAt(mi), a.pos * F.cache_width * 2), try r.back(b.cacheAt(mi), b.pos * F.cache_width * 2)); +} + +/// `ling-forward-synth DIR REF [PICKS]`: PICKS gets the serial run's picks, for ling_forward.py --picks. +pub fn run(gpu: check.Gpu, dir: []const u8, ref_path: []const u8, picks_out: ?[]const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const ref = try readRef(a, gpu.io, ref_path); + var why: ling.config.Why = .{}; + var o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + var m = try F.Modules.load(gpu.d, gpu.ctx); + defer m.deinit(); + const backbone = try a.alloc(u32, o.c.layers); + for (backbone, 0..) |*i, j| i.* = @intCast(j); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + var w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &o.ck, o.c, backbone, .{}); + defer w.deinit(); + var g = try F.Graph.init(gpu.gpa, gpu.d, &m, s, &w); + defer g.deinit(); + const c = &w.config; + const V: usize = c.vocab; + const T = ref.T; + try check.expect(ref.V == V and T > P + 8 - 1 and ref.top == c.topk, "reference: {d} tokens of vocabulary {d}, top {d}", .{ T, ref.V, ref.top }); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .m = &m, .g = &g }; + defer r.deinit(); + var pb = try F.Buffers.init(gpu.d, &m, &w, 2 * P, true, 128); + defer pb.deinit(); + var db = try F.Buffers.init(gpu.d, &m, &w, 2 * F.window_rows, false, 128); + defer db.deinit(); + var nm: usize = 0; + for (w.layers) |l| nm += @intFromBool(l.ffn == .moe); + try check.expect(nm == ref.layers, "MoE layers {d}, the reference's {d}", .{ nm, ref.layers }); + const ids = ref.ids; + // A: the serial run, against the references + const sa = try r.state(128); + const A = try serial(&r, &pb, &db, sa, ids, nm); + const flipped = try a.alloc(bool, T); + @memset(flipped, false); + const nf = rig.flips(A.picks, ref.picks, nm, T, 0, T, c.topk, flipped); + var rows_flipped: usize = 0; + for (flipped) |x| rows_flipped += @intFromBool(x); + const d64 = rig.distance(A.logits, ref.l64, T, V, flipped); + const dbf = rig.distance(A.logits, ref.lbf, T, V, flipped); + if (picks_out) |path| try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = path, .data = std.mem.sliceAsBytes(A.picks) }); + const how = if (ref.forced) "the reference routed on these picks" else "the reference's own float64 picks"; + std.debug.print("RESULT ling-forward-synth {t}: {d} layers, {d} prompt + {d} decode rows; picks vs {s}: {d} of {d} differ ({d} rows)\n", .{ c.quant, c.layers, P, T - P, how, nf, nm * T, rows_flipped }); + if (ref.forced) std.debug.print("RESULT float64's own top-k vs these: the worst pick ranks {e:.3} below its k-th biased score\n", .{ref.gap}); + std.debug.print("RESULT logits vs float64: rel L2 {e:.3} ({e:.3} without flipped rows), worst row {e:.3}, argmax {d}/{d}\n", .{ d64.rel, d64.rel_kept, d64.worst_row, d64.argmax, T }); + std.debug.print("RESULT logits vs bf16 points: rel L2 {e:.3} ({e:.3} without flipped rows), worst row {e:.3}, argmax {d}/{d}\n", .{ dbf.rel, dbf.rel_kept, dbf.worst_row, dbf.argmax, T }); + var bd: f64 = 0; + var bw: f64 = 0; + for (ref.lbf, ref.l64) |x, e| { + bd += (x - e) * (x - e); + bw += e * e; + } + const base = @sqrt(bd / bw); // HF's bf16 rounding points alone, against float64 + std.debug.print("RESULT the bf16-point reference itself vs float64: rel L2 {e:.3}\n", .{base}); + if (ref.forced) try check.expect(nf == 0 and d64.rel <= 1.5 * base and dbf.rel <= base, "on the same picks: logits within 1.5x the bf16 points' rel L2 of float64, and within it of the bf16 points", .{}); + // B: the prompt cut at 48 + const sb = try r.state(128); + try rig.same("prompt cut at 48 == uncut (rows 0-47)", try r.run(&pb, &.{.{ .st = sb, .row0 = 0, .rows = 48 }}, ids[0..48], null), A.logits[0 .. 48 * V]); + try r.fwd(&pb, null).commit(sb, 48); + try rig.same("prompt cut at 48 == uncut (rows 48-79)", try r.run(&pb, &.{.{ .st = sb, .row0 = 0, .rows = P - 48 }}, ids[48..P], null), A.logits[48 * V .. P * V]); + try r.fwd(&pb, null).commit(sb, P - 48); + // C: one 8-row verify window == 8 decode rows + const sc = try r.state(128); + _ = try r.run(&pb, &.{.{ .st = sc, .row0 = 0, .rows = P }}, ids[0..P], null); + try r.fwd(&pb, null).commit(sc, P); + try rig.same("8-row verify window == 8 decode rows", try r.run(&db, &.{.{ .st = sc, .row0 = 0, .rows = 8 }}, ids[P .. P + 8], null), A.logits[P * V .. (P + 8) * V]); + try r.fwd(&db, null).commit(sc, 3); + try rig.same("after keeping 3 of 8: the next 5-row window == decode rows", try r.run(&db, &.{.{ .st = sc, .row0 = 0, .rows = 5 }}, ids[P + 3 .. P + 8], null), A.logits[(P + 3) * V .. (P + 8) * V]); + try r.fwd(&db, null).commit(sc, 5); + try sameState(&r, "8-row window, keep 3, 5-row window == 8 decode rows", sc, sa); + // the fused KDA chain on the same window: bit-equal to the three-kernel one + const sf = try r.state(128); + _ = try r.run(&pb, &.{.{ .st = sf, .row0 = 0, .rows = P }}, ids[0..P], null); + try r.fwd(&pb, null).commit(sf, P); + r.kda_wide = false; + try rig.same("8-row window on the fused KDA chain == on the three-kernel chain", try r.run(&db, &.{.{ .st = sf, .row0 = 0, .rows = 8 }}, ids[P .. P + 8], null), A.logits[P * V .. (P + 8) * V]); + r.kda_wide = true; + sf.deinit(); + // D and E: two sequences batched, in a prompt chunk and in a decode window + const sd = try r.state(128); + const se = try r.state(128); + const both = try a.alloc(i32, P + 48); + @memcpy(both[0..P], ids[0..P]); + @memcpy(both[P..], ids[0..48]); + const pl = try r.run(&pb, &.{ .{ .st = sd, .row0 = 0, .rows = P }, .{ .st = se, .row0 = P, .rows = 48 } }, both, null); + try rig.same("batched prompt: sequence D == alone", pl[0 .. P * V], A.logits[0 .. P * V]); + try rig.same("batched prompt: sequence E (48 rows) == alone", pl[P * V ..], A.logits[0 .. 48 * V]); + try r.fwd(&pb, null).commit(sd, P); + try r.fwd(&pb, null).commit(se, 48); + _ = try r.run(&pb, &.{.{ .st = se, .row0 = 0, .rows = P - 48 }}, ids[48..P], null); + try r.fwd(&pb, null).commit(se, P - 48); + const win = try a.alloc(i32, 10); + @memcpy(win[0..4], ids[P .. P + 4]); + @memcpy(win[4..], ids[P .. P + 6]); + const wl = try r.run(&db, &.{ .{ .st = sd, .row0 = 0, .rows = 4 }, .{ .st = se, .row0 = 4, .rows = 6 } }, win, null); + try rig.same("batched window: sequence D (4 rows) == alone", wl[0 .. 4 * V], A.logits[P * V .. (P + 4) * V]); + try rig.same("batched window: sequence E (6 rows) == alone", wl[4 * V ..], A.logits[P * V .. (P + 6) * V]); + for ([_]*F.State{ sa, sb, sc, sd, se }) |x| x.deinit(); + check.pass("ling-forward-synth {t}: logits rel L2 {e:.3} of float64 ({s}); cut == uncut, windows == decode rows, a partial commit replays, batched == alone", .{ c.quant, d64.rel, if (ref.forced) "same picks" else "free routing" }); +} diff --git a/zig/tests/cuda/ling_lanes.zig b/zig/tests/cuda/ling_lanes.zig new file mode 100644 index 0000000000..1040288ce8 --- /dev/null +++ b/zig/tests/cuda/ling_lanes.zig @@ -0,0 +1,197 @@ +//! `ling-lanes DIR OUT [REF] [--layers a,b,..]`: Ling's lane backend under upstream's round loop, no server. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_lanes_rig.zig"); +const bytes = @import("ling_lanes_bytes.zig"); + +const Cover = rig.Cover; +const Engine = rig.Engine; + +/// The prompt's tokens: a fixed walk over the vocabulary, as ling_forward.py's. +pub fn walk(a: std.mem.Allocator, n: usize, vocab: usize, salt: usize) ![]u32 { + const p = try a.alloc(u32, n); + for (p, 0..) |*t, i| t.* = @intCast((37 * i + 11 + 101 * salt) % (vocab - 2)); + return p; +} + +/// Rows a prompt chunk takes in these checks: several chunks a prompt, and not a multiple of 16. +const chunk = 37; + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const out = rest[1]; + var ref: ?[]const u8 = null; + var thin: ?[]u32 = null; + var mtp = false; + var i: usize = 2; + while (i < rest.len) : (i += 1) { + if (std.mem.eql(u8, rest[i], "--layers") and i + 1 < rest.len) { + var list: std.ArrayList(u32) = .empty; + var it = std.mem.tokenizeScalar(u8, rest[i + 1], ','); + while (it.next()) |x| try list.append(a, try std.fmt.parseInt(u32, x, 10)); + thin = list.items; + i += 1; + } else if (std.mem.eql(u8, rest[i], "--mtp")) { + mtp = true; + } else ref = rest[i]; + } + const e = try Engine.init(gpu.gpa, gpu.io, gpu.d, gpu.ctx, dir, .{ .capacity = 1024, .prompt_rows = 512, .layers = thin, .mtp = mtp }); + defer e.deinit(); + e.chunk = chunk; + var ll: ling.lanes.Lanes = .{ .e = e }; + defer ll.deinit(); + var cfg = try lanes.Config.init(gpu.gpa, ll.facts(), ling.lanes.max_rows, ling.lanes.max_rows - 1); + defer cfg.deinit(gpu.gpa); + var wall: lanes.backend.WallClock = .{ .io = gpu.io }; + var le = lanes.Engine.init(gpu.gpa, &cfg, ll.backend(), wall.clock()); + defer le.deinit(); + var c: Cover = .{ .gpa = gpu.gpa, .a = a, .le = &le }; + const V: usize = e.c.vocab; + std.debug.print("RESULT ling-lanes {t}: {d} layers loaded of {d}, vocabulary {d}, prompt chunks of {d} rows{s}\n", .{ e.c.quant, e.main.layers.len, e.c.layers, V, chunk, if (e.mtp != null) ", drafts from the MTP layer" else "" }); + + const p = try walk(a, 150, V, 0); + // W, what every contract below holds to: the plain greedy pass, its prompt in 37-row chunks + const w = (try c.run(.{ .prompt = p, .max_new = 160, .drafts = false })).ids; + const copy_prompt = try std.mem.concat(a, u32, &.{ p, w[0..100], p }); + + // chunked == whole: the prompt in one chunk, in 16-row chunks, and through `prefill` (no steps) + for ([_]usize{ 512, 16 }) |rows| { + e.chunk = rows; + try c.same(try std.fmt.allocPrint(a, "prompt chunks of {d} rows, against chunks of {d}", .{ rows, chunk }), try c.run(.{ .prompt = p, .max_new = 160, .drafts = false }), w, .length); + } + e.chunk = chunk; + ll.whole = true; + le.backend = ll.backend(); + try c.same("the prompt pass whole (prefill, no steps)", try c.run(.{ .prompt = p, .max_new = 160, .drafts = false }), w, .length); + ll.whole = false; + le.backend = ll.backend(); + + // drafted == plain: windows of 3 to 16 rows from a known reply, most with a wrong row; greedy and sampled + const keyed: lanes.Sampling = .{ .seed = lanes.sampling.seedFor(copy_prompt, 0) }; + const plain_g = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .drafts = false }); + const plain_s = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .drafts = false, .sampling = keyed }); + for ([_]Cover.Got{ plain_g, plain_s }, [_]?lanes.Sampling{ null, keyed }) |plain, sampling| { + const drafted = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .sampling = sampling, .truth = plain.ids }); + try c.same(if (sampling != null) "drafted, sampled, against plain" else "drafted, greedy, against plain", drafted, plain.ids, .length); + try c.expect(drafted.drafted > 0 and drafted.accepted > 0 and drafted.accepted < drafted.drafted, "drafted{s}: {d} drafts in {d} rounds, {d} kept", .{ if (sampling != null) ", sampled" else "", drafted.drafted, drafted.rounds, drafted.accepted }); + } + const g = plain_g.ids; + const differ = std.mem.indexOfDiff(u32, plain_s.ids, g); + try c.expect(differ != null, "the sampled reply leaves the greedy one at token {?d} (sampling draws)", .{differ}); + + if (e.mtp != null) { // the MTP layer's drafts against the plain replies (acceptance on random weights says nothing) + for ([_]Cover.Got{ plain_g, plain_s }, [_]?lanes.Sampling{ null, keyed }) |plain, sampling| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .sampling = sampling }); + try c.same(if (sampling != null) "MTP drafts, sampled, against plain" else "MTP drafts, greedy, against plain", got, plain.ids, .length); + try c.expect(got.drafted > 0, "MTP drafts{s}: {d} drafted in {d} rounds, {d} kept", .{ if (sampling != null) ", sampled" else "", got.drafted, got.rounds, got.accepted }); + } + var keeper: rig.Keeper = .{ .e = e }; + defer keeper.drop(); + const marks = [_]u32{ 53, 211 }; + _ = try c.run(.{ .prompt = copy_prompt, .max_new = 8, .reuse = .{ .marks = &marks, .hook = keeper.hook() } }); + for (keeper.kept[0..keeper.n], keeper.at[0..keeper.n]) |kept, at| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 200, .reuse = .{ .saved = kept, .at = at } }); + try c.same(try std.fmt.allocPrint(a, "MTP drafts resumed at {d}, against plain", .{at}), got, plain_g.ids, .length); + } + } + + // resumed == fresh: states kept at marks off the chunk grid and off multiples of 16, then passes resumed there + { + var keeper: rig.Keeper = .{ .e = e }; + defer keeper.drop(); + const marks = [_]u32{ 37, 101, 150, 263, 349 }; + const marked = try c.run(.{ .prompt = copy_prompt, .max_new = 64, .truth = g, .reuse = .{ .marks = &marks, .hook = keeper.hook() } }); + try c.same("a drafted pass cut at its marks, against the plain fresh pass", marked, g[0..64], .length); + try c.expect(keeper.n == marks.len, "the pass kept {d} states at its marks ({d} asked)", .{ keeper.n, marks.len }); + for (keeper.kept[0..keeper.n], keeper.at[0..keeper.n]) |kept, at| { + const got = try c.run(.{ .prompt = copy_prompt, .max_new = 64, .truth = g, .reuse = .{ .saved = kept, .at = at } }); + try c.expect(got.cached == at, "the pass resumed at {d} ({d} restored)", .{ at, got.cached }); + try c.same(try std.fmt.allocPrint(a, "resumed at {d}, drafted, against the plain fresh pass", .{at}), got, g[0..64], .length); + } + const other = try std.mem.concat(a, u32, &.{ copy_prompt[0..263], try walk(a, 60, V, 3) }); + const other_fresh = try c.run(.{ .prompt = other, .max_new = 64, .drafts = false }); + const got = try c.run(.{ .prompt = other, .max_new = 64, .drafts = false, .reuse = .{ .saved = keeper.kept[3], .at = 263 } }); + try c.same("another tail after the kept 263, against its fresh pass", got, other_fresh.ids, .length); + } + + // cancels in the prompt pass, discards between rounds, eos, max_new at the window's edges + for ([_]u32{ 0, 1, 3 }) |at| { + var cut: rig.Cut = .{ .at = at }; + const r = c.run(.{ .prompt = copy_prompt, .max_new = 64, .cancel = &cut }); + try c.expect(if (r) |_| false else |err| err == error.Cancelled, "a cancel at ask {d} of the prompt pass ended it with error.Cancelled", .{at}); + try c.same(try std.fmt.allocPrint(a, "the request after a cancel at prompt ask {d}", .{at}), try c.run(.{ .prompt = p, .max_new = 64, .drafts = false }), w[0..64], .length); + } + for ([_]u32{ 1, 8 }) |rounds| { + try c.expect(try c.discardAfter(copy_prompt, g, rounds), "a drafted stream was discarded after {d} rounds", .{rounds}); + try c.same(try std.fmt.allocPrint(a, "the request after a discard after {d} rounds", .{rounds}), try c.run(.{ .prompt = p, .max_new = 64, .drafts = false }), w[0..64], .length); + } + for ([_]usize{ 0, 9, 33 }) |from| { + const j = for (from..g.len) |k| { + if (std.mem.indexOfScalar(u32, g[0..k], g[k]) == null) break k; + } else continue; + try c.same(try std.fmt.allocPrint(a, "eos at token {d}, drafted", .{j}), try c.run(.{ .prompt = copy_prompt, .max_new = 512, .eos = &.{g[j]}, .truth = g }), g[0 .. j + 1], .stop); + } + for (1..4) |n| try c.same(try std.fmt.allocPrint(a, "max_new {d}, drafted", .{n}), try c.run(.{ .prompt = copy_prompt, .max_new = @intCast(n), .truth = g }), g[0..n], .length); + + // concurrent == solo + try concurrent(&c, &ll, a, V); + // the engine's own steps, byte for byte: cuts, resumes, windows and shared forwards against their twins + try bytes.engineChecks(&c, e, copy_prompt, w); + // the plain pass's ids and picks for ling_forward.py; with REF, its logits against float64 + try bytes.serial(&c, gpu, e, p, w, out, ref); + if (c.failed) return error.TestFailed; + check.pass("ling-lanes {t}: {d} checks", .{ e.c.quant, c.checks }); +} + +/// Each concurrent stream on one lanes.Engine equals its solo run. +fn concurrent(c: *Cover, ll: *ling.lanes.Lanes, a: std.mem.Allocator, V: usize) !void { + const n = 96; + var ps: [3][]const u32 = undefined; + for (&ps, 0..) |*p, k| { + const base = try walk(a, 60 + 23 * k, V, k + 1); + p.* = try std.mem.concat(a, u32, &.{ base, base[0..40], base }); + } + const seeds = [_]lanes.Sampling{ .{ .seed = 11 }, .{ .seed = 22 }, .{ .seed = 33 } }; + var solo: [3][]const u32 = undefined; + var solo_s: [3][]const u32 = undefined; + for (ps, &solo, &solo_s, 0..) |p, *g, *gs, k| { + g.* = (try c.run(.{ .prompt = p, .max_new = n, .drafts = false })).ids; + gs.* = (try c.run(.{ .prompt = p, .max_new = n, .drafts = false, .sampling = seeds[k] })).ids; + } + var sampled_seen = false; // the drafts' truth: the greedy solo replies, or the sampled ones in sampled cases + const cases = [_]struct { name: []const u8, joins: []const rig.Join }{ + .{ .name = "all at once", .joins = &.{ .{ .p = 0 }, .{ .p = 1 }, .{ .p = 2 } } }, + .{ .name = "all at once, drafted", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 1, .drafts = true }, .{ .p = 2, .drafts = true } } }, + .{ .name = "staggered by 5 rounds, drafted", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 1, .at = 5, .drafts = true }, .{ .p = 2, .at = 10 } } }, + .{ .name = "in reverse", .joins = &.{ .{ .p = 2 }, .{ .p = 1, .drafts = true }, .{ .p = 0 } } }, + .{ .name = "the same prompt twice", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 0 } } }, + .{ .name = "sampled, drafted", .joins = &.{ .{ .p = 0, .sampled = true, .drafts = true }, .{ .p = 1, .sampled = true }, .{ .p = 2, .sampled = true, .drafts = true } } }, + .{ .name = "the second discarded after 6 rounds", .joins = &.{ .{ .p = 0 }, .{ .p = 1, .discard_at = 6 }, .{ .p = 2, .drafts = true } } }, + .{ .name = "prompts a chunk at a time while the others decode", .joins = &.{ .{ .p = 0, .drafts = true }, .{ .p = 1, .at = 2, .filled = true }, .{ .p = 2, .at = 4, .filled = true, .drafts = true } } }, + .{ .name = "sampled, prompts a chunk at a time", .joins = &.{ .{ .p = 0, .sampled = true }, .{ .p = 1, .at = 2, .sampled = true, .filled = true }, .{ .p = 2, .at = 4, .sampled = true, .filled = true, .drafts = true } } }, + }; + const before = ll.shared; + for (cases) |cs| { + sampled_seen = cs.joins[0].sampled; + const got = try c.together(cs.joins, &ps, if (sampled_seen) &solo_s else &solo, &seeds, n); + for (cs.joins, got.ids, 0..) |j, ids, k| { + if (j.discard_at != null) continue; + try c.same(try std.fmt.allocPrint(a, "{s}: stream {d} (prompt {d}{s}) against its solo run", .{ cs.name, k, j.p, if (j.sampled) ", sampled" else "" }), .{ .ids = ids, .reason = .length, .rounds = 0, .cached = 0, .drafted = 0, .accepted = 0 }, if (j.sampled) solo_s[j.p] else solo[j.p], .length); + } + var drafting = ll.e.mtp != null; // every joined stream drafts with the MTP layer when it is loaded + for (cs.joins) |j| drafting = drafting or j.drafts; + try c.expect(got.most >= 2 and (got.drafted > 0) == drafting, "{s}: {d} streams at most in one round, {d} rounds with two or more, {d} drafts", .{ cs.name, got.most, got.shared, got.drafted }); + for (cs.joins) |j| if (j.filled) { + try c.expect(got.between > 0, "{s}: {d} rounds of the others between a prompt's chunks", .{ cs.name, got.between }); + break; + }; + } + try c.expect(ll.shared > before, "{d} forwards held several streams' windows (Engine.verifyShared)", .{ll.shared - before}); +} diff --git a/zig/tests/cuda/ling_lanes_bytes.zig b/zig/tests/cuda/ling_lanes_bytes.zig new file mode 100644 index 0000000000..4716e3c60f --- /dev/null +++ b/zig/tests/cuda/ling_lanes_bytes.zig @@ -0,0 +1,426 @@ +//! The Ling engine's steps against their twins byte for byte, and the plain pass's logits against float64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_lanes_rig.zig"); +const frig = @import("ling_forward_rig.zig"); +const synth = @import("ling_forward_synth.zig"); + +const Cover = rig.Cover; +const Engine = rig.Engine; +const F = ling.forward; + +/// Kept-state ids by mark, as a direct engine pass keeps them. +const Marked = struct { + e: *Engine, + ids: [8]u64 = undefined, + n: usize = 0, + fn at(ptr: *anyopaque, pos: u32) void { + _ = pos; + const m: *Marked = @ptrCast(@alignCast(ptr)); + m.ids[m.n] = m.e.keepState() catch return; + m.n += 1; + } +}; + +fn logitsRows(e: *Engine, a: std.mem.Allocator, ptr: u64, rows: usize) ![]u8 { + const out = try a.alloc(u8, rows * e.c.vocab * 2); + try e.stream.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = e.d, .ptr = ptr, .len = out.len }; + try b.download(0, out); + return out; +} + +/// A fresh pass of `ids` on `slot` with `rows`-row chunks (and `marks`), or resumed from `from`; its first token. +fn pass(e: *Engine, slot: u32, ids: []const u32, rows: usize, from: ?ling.engine.Resume, marks: ?ling.engine.Marks) !u32 { + try e.bind(slot); + try e.begin(null); + e.chunk = rows; + return e.prefill(ids, from, null, marks); +} + +pub fn engineChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const keep_chunk = e.chunk; + defer e.chunk = keep_chunk; + const t10 = try pass(e, 10, prompt, 512, null, null); + const t11 = try pass(e, 11, prompt, 37, null, null); + const t12 = try pass(e, 12, prompt, 16, null, null); + try c.expect(t10 == t11 and t10 == t12, "first tokens: whole {d}, 37-row chunks {d}, 16-row chunks {d}", .{ t10, t11, t12 }); + try rig.sameSeq(c, e, "prompt in 37-row chunks == whole", 11, 10); + try rig.sameSeq(c, e, "prompt in 16-row chunks == whole", 12, 10); + var mk: Marked = .{ .e = e }; + const marks = [_]u32{ 101, 263, 349 }; + _ = try pass(e, 13, prompt, 48, null, .{ .at = &marks, .ctx = &mk, .call = Marked.at }); + try c.expect(mk.n == marks.len, "a pass kept {d} states at marks 101, 263, 349", .{mk.n}); + for (mk.ids[0..mk.n], marks[0..mk.n]) |id, at| { + const t = try pass(e, 14, prompt, 37, .{ .id = id, .at = at }, null); + try c.expect(t == t10, "resumed at {d}: first token {d}, the whole pass's {d}", .{ at, t, t10 }); + try rig.sameSeq(c, e, try std.fmt.allocPrint(a, "resumed at {d} == whole", .{at}), 14, 10); + e.dropState(id); + } + // a 9-row window, keep 4, a 5-row window == 9 single rows on the twin (tokens, logits, states) + var g: [10]u32 = undefined; + g[0] = t11; + const lone = try a.alloc([]u8, 9); + try e.bind(11); + for (0..9) |r| { + try e.verify(g[r .. r + 1], g[r + 1 .. r + 2]); + lone[r] = try logitsRows(e, a, e.bd.logits, 1); + try e.keep(1); + } + try e.bind(10); + var got: [9]u32 = undefined; + try e.verify(g[0..9], &got); + const win = try logitsRows(e, a, e.bd.logits, 9); + var same_rows = std.mem.eql(u32, &got, g[1..10]); + for (0..9) |r| same_rows = same_rows and std.mem.eql(u8, win[r * lone[r].len ..][0..lone[r].len], lone[r]); + try e.keep(4); + try e.verify(g[4..9], got[0..5]); + const win2 = try logitsRows(e, a, e.bd.logits, 5); + same_rows = same_rows and std.mem.eql(u32, got[0..5], g[5..10]); + for (0..5) |r| same_rows = same_rows and std.mem.eql(u8, win2[r * lone[4 + r].len ..][0..lone[4 + r].len], lone[4 + r]); + try e.keep(5); + try c.expect(same_rows, "a 9-row window and, after keeping 4, a 5-row window: tokens and logits == 9 single rows", .{}); + try rig.sameSeq(c, e, "9-row window, keep 4, 5-row window == 9 single rows", 10, 11); + try sharedChecks(c, e, prompt, w); + if (e.mtp != null) try mtpChain(c, e, prompt, w); + try capacityRefusal(c, e, prompt, w); + try scratchRefusal(c, e, prompt, w); + for ([_]u32{ 10, 11, 12, 13, 14 }) |s| e.release(s); +} + +/// A sequence's arenas as they stand: the layers' caches and states, the MTP layer's cache. +fn arenas(c: *Cover, e: *Engine, slot: u32) ![2][]u8 { + const q = e.seqs.items[slot].?; + const st = &q.st.arena.buf; + const hd: []u8 = if (q.hd) |*h| try rig.back(e, c.a, h.arena.buf.ptr, h.arena.buf.len) else &.{}; + return .{ try rig.back(e, c.a, st.ptr, st.len), hd }; +} + +/// Rows past the capacity refused with the arenas unchanged; the next window and a new request then equal twins'. +fn capacityRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const cap = e.capacity; + const long = try a.alloc(u32, cap + 40); + for (long, 0..) |*x, i| x.* = prompt[i % prompt.len]; + var pend: [2]u32 = undefined; + for (0..2) |k| { // twins at cap - 12: a pass of cap - 24 rows, a 12-row window kept whole + const t = try pass(e, @intCast(50 + k), long[0 .. cap - 24], 512, null, null); + var ids: [12]u32 = undefined; + var out: [12]u32 = undefined; + ids[0] = t; + @memcpy(ids[1..], w[0..11]); + try e.verify(&ids, &out); + try e.keep(12); + pend[k] = out[11]; + } + try e.bind(50); + const before = try arenas(c, e, 50); + const q = e.seqs.items[50].?; + const cur = try a.dupe(u1, q.st.cur); + const pos = q.st.pos; + var refused: usize = 0; + // a prompt chunk of 30 rows (its KDA layers would commit in place), then a 16-row window + e.chunk = 512; + if (e.prefillChunk(long[0 .. cap + 18], pos, null, null)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + var ids: [16]u32 = undefined; + var out: [16]u32 = undefined; + ids[0] = pend[0]; + @memcpy(ids[1..], w[20..35]); + if (e.verify(&ids, &out)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + var asked: usize = 2; + if (e.mtp) |*h| { // one row past the capacity into the MTP layer's cache (its MLA block refuses before writing) + asked += 1; + const n = cap + 1 - q.hd.?.pos; + if (h.absorb(&e.m, e.stream, &e.bp, &q.hd.?, @ptrCast(long[0..n]), e.bp.normed)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + } + const after = try arenas(c, e, 50); + const same = std.mem.eql(u8, before[0], after[0]) and std.mem.eql(u8, before[1], after[1]) and + std.mem.eql(u1, cur, q.st.cur) and q.st.pos == pos and q.st.pending == 0; + try c.expect(refused == asked and same, "{d} of {d} requests past the capacity ({d} at {d}) refused with CacheFull, the arenas ({d} and {d} bytes) and position as they were", .{ refused, asked, cap, pos, before[0].len, before[1].len }); + try afterRefusal(c, e, prompt, w, pend, 50); +} + +/// The next window on twins `s` (refused) and s + 1, then a new request on `s` and on a fresh slot s + 2. +fn afterRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32, pend: [2]u32, s: u32) !void { + var got: [2][4]u32 = undefined; + var lo: [2][]u8 = undefined; + for (0..2) |k| { + try e.bind(s + @as(u32, @intCast(k))); + var win: [4]u32 = .{ pend[k], w[40], w[41], w[42] }; + try e.verify(&win, &got[k]); + lo[k] = try logitsRows(e, c.a, e.bd.logits, 4); + try e.keep(4); + } + try c.expect(std.mem.eql(u32, &got[0], &got[1]) and std.mem.eql(u8, lo[0], lo[1]), "the window after the refusals: tokens and logits == the twin's", .{}); + try rig.sameSeq(c, e, "after the refusals and a window, the twin", s, s + 1); + const t0 = try pass(e, s, prompt, 37, null, null); + const t2 = try pass(e, s + 2, prompt, 37, null, null); + try c.expect(t0 == t2, "a new request on the refused slot: first token {d}, on a fresh slot {d}", .{ t0, t2 }); + try rig.sameSeq(c, e, "a new request on the refused slot == on a fresh slot", s, s + 2); + for ([_]u32{ s, s + 1, s + 2 }) |x| e.release(x); +} + +/// Scratch for half the sequences' capacity: a prompt chunk, a window and MTP rows past it refused, nothing changed. +fn scratchRefusal(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const half = e.capacity / 2; + const long = try a.alloc(u32, half + 40); + for (long, 0..) |*x, i| x.* = prompt[i % prompt.len]; + var pend: [2]u32 = undefined; + for (0..2) |k| pend[k] = try pass(e, @intCast(60 + k), long[0 .. half - 12], 512, null, null); // twins at half - 12 + var bp = try F.Buffers.init(e.d, &e.m, &e.w, 64, true, half); + defer bp.deinit(); + var bd = try F.Buffers.init(e.d, &e.m, &e.w, F.window_rows, false, half); + defer bd.deinit(); + try e.bind(60); + const q = e.seqs.items[60].?; + const before = try arenas(c, e, 60); + const cur = try a.dupe(u1, q.st.cur); + const pos = q.st.pos; + var refused: usize = 0; + const seg = [_]F.Seg{.{ .st = &q.st, .row0 = 0, .rows = 30 }}; // fits the sequence, not the scratch + if (e.forward(&bp).run(&seg, @ptrCast(long[pos..][0..30]))) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + const win = [_]F.Seg{.{ .st = &q.st, .row0 = 0, .rows = 16 }}; + if (e.forward(&bd).run(&win, @ptrCast(w[20..36]))) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + var asked: usize = 2; + if (e.mtp) |*h| { // 16 rows through the MTP layer from its position, past the scratch's + asked += 1; + if (h.absorb(&e.m, e.stream, &bd, &q.hd.?, @ptrCast(long[0..16]), e.bd.normed)) |_| {} else |err| refused += @intFromBool(err == error.CacheFull); + } + const after = try arenas(c, e, 60); + const same = std.mem.eql(u8, before[0], after[0]) and std.mem.eql(u8, before[1], after[1]) and + std.mem.eql(u1, cur, q.st.cur) and q.st.pos == pos and q.st.pending == 0 and (q.hd == null or q.hd.?.pending == 0); + try c.expect(refused == asked and same, "{d} of {d} requests past the scratch ({d} of {d}, at {d}) refused with CacheFull, the arenas and position as they were", .{ refused, asked, half, e.capacity, pos }); + try afterRefusal(c, e, prompt, w, pend, 60); +} + +/// MTP chains of 3 and of 1 on twins: the same first draft, and after the same windows the chain's rows rolled back. +fn mtpChain(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + var t: [2]u32 = undefined; + var d: [2][3]u32 = undefined; + for (0..2) |k| { + t[k] = try pass(e, @intCast(40 + k), prompt[0..200], 37, null, null); + try e.draft(true, &.{t[k]}, if (k == 0) 3 else 1, &d[k]); + } + var ok = t[0] == t[1] and d[0][0] == d[1][0]; + var pend = t[0]; + var chained: usize = 0; + for (0..6) |round| { + var ids: [8]u32 = undefined; + var out: [2][8]u32 = undefined; + const rows = 2 + round % 6; + ids[0] = pend; + for (1..rows) |r| ids[r] = w[(round * 11 + r) % w.len]; + const keep = 1 + (round * 3) % rows; + for (0..2) |k| { + try e.bind(@intCast(40 + k)); + try e.verify(ids[0..rows], out[k][0..rows]); + try e.keep(keep); + try e.draft(false, out[k][0..keep], if (k == 0) 3 else 1, &d[k]); + } + ok = ok and std.mem.eql(u32, out[0][0..rows], out[1][0..rows]) and d[0][0] == d[1][0]; + { // the chain's second draft by hand on the twin: d0 with the newest kept row's MTP output as its hidden row + const h = &e.mtp.?; + const hd = &e.seqs.items[41].?.hd.?; + const D: usize = e.c.hidden; + const f = h.forward(&e.m, e.stream, &h.bm); + try h.absorb(&e.m, e.stream, &h.bm, hd, @ptrCast(d[1][0..1]), h.bm.normed + (keep - 1) * D * 2); + try f.final(1); + try f.logitsOf(0, 1); + var second: [1]u32 = undefined; + try e.draw(h.bm.logits, 1, hd.pos + 1, &second); + try hd.setPos(e.stream, hd.pos - 1); + ok = ok and second[0] == d[0][1]; + chained += 1; + } + pend = out[0][keep - 1]; + } + try c.expect(ok, "MTP chains of 3 and of 1 over 6 windows (kept 1 to 7 rows): the same tokens and first drafts; {d} second drafts == by hand", .{chained}); + const sa = e.seqs.items[40].?; + const sb = e.seqs.items[41].?; + try c.expect(sa.hd.?.pos == sa.st.pos + 2 and sb.hd.?.pos == sb.st.pos, "the chain of 3 holds its 2 chained rows past the kept ones (MTP at {d}, the sequence at {d})", .{ sa.hd.?.pos, sa.st.pos }); + for (0..2) |k| { // one more window, both drafting 1: the chain's rows rolled back, the twins byte for byte + var out: [2]u32 = undefined; + try e.bind(@intCast(40 + k)); + try e.verify(&.{ pend, w[3] }, &out); + try e.keep(2); + try e.draft(false, &out, 1, &d[k]); + } + try rig.sameSeq(c, e, "after chains of 3 and of 1, a window and a chain of 1 each", 40, 41); + for (0..2) |k| e.release(@intCast(40 + k)); +} + +/// Two streams' windows in one shared forward against lone twins over 8 rounds: logits, tokens, then states. +fn sharedChecks(c: *Cover, e: *Engine, prompt: []const u32, w: []const u32) !void { + const a = c.a; + const pa = prompt[0..120]; + const pb = prompt[50..250]; + var pend: [2]u32 = undefined; + for ([_]u32{ 20, 21, 22, 23 }, 0..) |slot, k| pend[k % 2] = try pass(e, slot, if (k % 2 == 0) pa else pb, 37, null, null); + var ok = true; + var rows_seen: usize = 0; + for (0..8) |round| { + var ids: [2][16]u32 = undefined; + var outs: [2][16]u32 = undefined; + var twin: [2][16]u32 = undefined; + var parts: [2]ling.engine.Part = undefined; + for (0..2) |k| { + const rows = 1 + (round * 5 + k * 7) % 16; + ids[k][0] = pend[k]; + for (1..rows) |r| ids[k][r] = w[(round * 13 + k * 31 + r) % w.len]; + parts[k] = .{ .slot = @intCast(20 + k), .ids = ids[k][0..rows], .out = outs[k][0..rows] }; + } + try e.verifyShared(&parts); + const sh = try logitsRows(e, a, e.bs.logits, parts[0].ids.len + parts[1].ids.len); + var row0: usize = 0; + for (0..2) |k| { + const rows = parts[k].ids.len; + try e.bind(@intCast(22 + k)); + try e.verify(parts[k].ids, twin[k][0..rows]); + const lo = try logitsRows(e, a, e.bd.logits, rows); + const V2 = e.c.vocab * 2; + if (!std.mem.eql(u8, sh[row0 * V2 ..][0 .. rows * V2], lo) or !std.mem.eql(u32, outs[k][0..rows], twin[k][0..rows])) ok = false; + row0 += rows; + rows_seen += rows; + const keep = 1 + (round * 3 + k) % rows; + for ([_]u32{ @intCast(20 + k), @intCast(22 + k) }) |s| { + try e.bind(s); + try e.keep(keep); + } + pend[k] = outs[k][keep - 1]; + } + } + try c.expect(ok, "shared windows: {d} rows of 2 streams over 8 rounds (1 to 16 rows each) == lone twins' logits and tokens", .{rows_seen}); + try rig.sameSeq(c, e, "shared stream 0 after 8 rounds == its lone twin", 20, 22); + try rig.sameSeq(c, e, "shared stream 1 after 8 rounds == its lone twin", 21, 23); + for ([_]u32{ 20, 21, 22, 23 }) |s| e.release(s); +} + +/// The plain pass serially (whole prompt, then one row at a time) with its picks; files for ling_forward.py --ids. +pub fn serial(c: *Cover, gpu: check.Gpu, e: *Engine, p: []const u32, w: []const u32, out: []const u8, ref_path: ?[]const u8) !void { + const a = c.a; + const ids = try a.alloc(i32, p.len + w.len - 1); + for (ids, 0..) |*x, k| x.* = @intCast(if (k < p.len) p[k] else w[k - p.len]); + const T = ids.len; + const P = p.len; + const V: usize = e.c.vocab; + const top: usize = e.c.topk; + var nm: usize = 0; + for (e.main.layers) |l| nm += @intFromBool(l.ffn == .moe); + var r: frig.Rig = .{ .gpu = gpu, .a = a, .s = e.stream, .m = &e.m, .g = &e.g }; + defer r.deinit(); + try e.bind(30); + try e.begin(null); + const st = &e.cur.st; + const logits = try a.alloc(u16, T * V); + const picks = try a.alloc(i32, nm * T * top); + const pp = try r.zeros(nm * P * (top + 1) * 4); + @memcpy(logits[0 .. P * V], try r.run(&e.bp, &.{.{ .st = st, .row0 = 0, .rows = P }}, ids[0..P], pp)); + const sp = try frig.sortedPicks(a, try r.back(pp, nm * P * (top + 1) * 4), nm, P, top); + for (0..nm) |l| @memcpy(picks[l * T * top ..][0 .. P * top], sp[l * P * top ..][0 .. P * top]); + try r.fwd(&e.bp, null).commit(st, P); + const dp = try r.zeros(nm * (top + 1) * 4); + for (P..T) |t| { + @memcpy(logits[t * V ..][0..V], try r.run(&e.bd, &.{.{ .st = st, .row0 = 0, .rows = 1 }}, ids[t .. t + 1], dp)); + const one = try frig.sortedPicks(a, try r.back(dp, nm * (top + 1) * 4), nm, 1, top); + for (0..nm) |l| @memcpy(picks[(l * T + t) * top ..][0..top], one[l * top ..][0..top]); + try r.fwd(&e.bd, null).commit(st, 1); + } + e.release(30); + var agree: usize = 0; + for (P - 1..T) |t| { + var best: usize = 0; + for (0..V) |j| if (frig.bfOf(logits[t * V + j]) > frig.bfOf(logits[t * V + best])) { + best = j; + }; + const next: u32 = if (t + 1 < T) @intCast(ids[t + 1]) else w[w.len - 1]; + agree += @intFromBool(best == next); + } + try c.expect(agree == T - P + 1, "the lanes' greedy reply == the serial forward's argmax at all {d} of {d} rows", .{ agree, T - P + 1 }); + var name: [512]u8 = undefined; + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fmt.bufPrint(&name, "{s}.ids", .{out}), .data = std.mem.sliceAsBytes(ids) }); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fmt.bufPrint(&name, "{s}.picks", .{out}), .data = std.mem.sliceAsBytes(picks) }); + const path = ref_path orelse return; + const ref = try synth.readRef(a, gpu.io, path); + try c.expect(ref.T == T and ref.V == V and ref.forced and std.mem.eql(i32, ref.ids, ids), "reference: {d} tokens routed on these picks, the same ids", .{ref.T}); + const flipped = try a.alloc(bool, T); + @memset(flipped, false); + const nf = frig.flips(picks, ref.picks, nm, T, 0, T, top, flipped); + const d64 = frig.distance(logits, ref.l64, T, V, flipped); + const dbf = frig.distance(logits, ref.lbf, T, V, flipped); + var bd: f64 = 0; + var bw: f64 = 0; + for (ref.lbf, ref.l64) |x, y| { + bd += (x - y) * (x - y); + bw += y * y; + } + const base = @sqrt(bd / bw); + std.debug.print("RESULT served logits vs float64 (same picks): rel L2 {e:.3}, worst row {e:.3}, argmax {d}/{d}; vs bf16 points {e:.3}; bf16 points vs float64 {e:.3}\n", .{ d64.rel, d64.worst_row, d64.argmax, T, dbf.rel, base }); + if (e.mtp != null) try mtpLogits(c, gpu, e, ids, path); + try c.expect(nf == 0 and d64.rel <= 1.5 * base and dbf.rel <= base, "served logits within 1.5x the bf16 points' rel L2 of float64 ({e:.3} <= 1.5 x {e:.3}), within it of the bf16 points", .{ d64.rel, base }); +} + +/// The MTP layer's logits at every row but the last (hidden row t, token t + 1) against ling_forward.py --mtp REF.mtp. +fn mtpLogits(c: *Cover, gpu: check.Gpu, e: *Engine, ids: []const i32, ref_path: []const u8) !void { + const a = c.a; + var name: [512]u8 = undefined; + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, try std.fmt.bufPrint(&name, "{s}.mtp", .{ref_path}), a, .limited(1 << 30)); + if (bytes.len < 20 or !std.mem.eql(u8, bytes[0..8], "TFLINGM2")) return error.BadFixture; + const dims = std.mem.bytesAsSlice(u32, bytes[8..20]); + const R: usize = dims[0]; + const V: usize = dims[1]; + const K: usize = dims[2]; + const T = ids.len; + try c.expect(R == T - 1 and V == e.c.vocab and K == e.c.topk, "MTP reference: {d} rows", .{R}); + const m64 = try a.alloc(f64, R * V); + const mbf = try a.alloc(f64, R * V); + const theirs = try a.alloc(i32, R * K); + var at: usize = 20; + for ([_][]u8{ std.mem.sliceAsBytes(m64), std.mem.sliceAsBytes(mbf), std.mem.sliceAsBytes(theirs) }) |dst| { + @memcpy(dst, bytes[at..][0..dst.len]); + at += dst.len; + } + const h = &e.mtp.?; + var r: frig.Rig = .{ .gpu = gpu, .a = a, .s = e.stream, .m = &e.m, .g = &e.g }; + defer r.deinit(); + try e.bind(31); + try e.begin(null); + const q = e.seqs.items[31].?; + const main = e.forward(&e.bp); + const seg = [_]F.Seg{.{ .st = &q.st, .row0 = 0, .rows = T }}; + try main.embed(ids); + try main.stack(&seg); + try main.final(T); // e.bp.normed: the hidden rows, as a prompt pass leaves them + const pk = try r.zeros(R * (K + 1) * 4); + const f = h.forward(&e.m, e.stream, &e.bp); + try f.embed(ids[1..]); + try h.step(f, &q.hd.?, R, e.bp.normed, pk); + try f.final(R); + try f.logitsOf(0, R); + const raw = try r.back(e.bp.logits, R * V * 2); + const got = try a.alloc(u16, R * V); + @memcpy(std.mem.sliceAsBytes(got), raw); + try main.commit(&q.st, T); + _ = try pass(e, 32, @ptrCast(ids), 37, null, null); // a prompt pass in 37-row chunks: its MTP rows the same + try rig.sameSeq(c, e, "a prompt pass's MTP rows in 37-row chunks == the serial MTP rows", 32, 31); + e.release(31); + e.release(32); + const ours = try frig.sortedPicks(a, try r.back(pk, R * (K + 1) * 4), 1, R, K); + const flipped = try a.alloc(bool, R); + @memset(flipped, false); + const nf = frig.flips(ours, theirs, 1, R, 0, R, K, flipped); + const d64 = frig.distance(got, m64, R, V, flipped); + const dbf = frig.distance(got, mbf, R, V, flipped); + var bd: f64 = 0; + var bw: f64 = 0; + for (0..R) |t| if (!flipped[t]) for (mbf[t * V ..][0..V], m64[t * V ..][0..V]) |x, y| { + bd += (x - y) * (x - y); + bw += y * y; + }; + const base = @sqrt(bd / bw); // the bf16 points' own distance on the same rows (their routing may differ there too) + std.debug.print("RESULT MTP logits vs float64: rel L2 {e:.3} ({e:.3} without {d} rows whose picks differ), argmax {d}/{d}; vs bf16 points {e:.3}; bf16 points vs float64 on those rows {e:.3}\n", .{ d64.rel, d64.rel_kept, nf, d64.argmax, R, dbf.rel, base }); + try c.expect(d64.rel_kept <= 1.5 * base, "MTP logits within 1.5x the bf16 points' rel L2 of float64 on rows routed alike ({e:.3} <= 1.5 x {e:.3})", .{ d64.rel_kept, base }); +} diff --git a/zig/tests/cuda/ling_lanes_rig.zig b/zig/tests/cuda/ling_lanes_rig.zig new file mode 100644 index 0000000000..1e3b44b9a7 --- /dev/null +++ b/zig/tests/cuda/ling_lanes_rig.zig @@ -0,0 +1,223 @@ +//! What the Ling lane checks share: requests through one lanes.Engine, joined streams, kept states, byte comparisons. +const std = @import("std"); +const cuda = @import("cuda"); +const lanes = @import("lanes"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); + +pub const Engine = ling.engine.Engine; + +/// A mark hook that keeps the state as the lane host's prompt cache does (native.snap_save). +pub const Keeper = struct { + e: *Engine, + kept: [8]*anyopaque = undefined, + at: [8]u32 = undefined, + n: usize = 0, + + pub fn hook(k: *Keeper) lanes.stream.MarkHook { + return .{ .ptr = k, .at = keep }; + } + + fn keep(ptr: *anyopaque, s: *lanes.Stream, pos: u32) void { + _ = s; + const k: *Keeper = @ptrCast(@alignCast(ptr)); + k.kept[k.n] = ling.native.snap_save(k.e, null, pos) catch |err| { + std.debug.print("FAIL keeping the state at {d}: {s}\n", .{ pos, @errorName(err) }); + return; + }; + k.at[k.n] = pos; + k.n += 1; + } + + pub fn drop(k: *Keeper) void { + for (k.kept[0..k.n]) |x| ling.native.snap_drop(k.e, x); + k.n = 0; + } +}; + +/// A cancel hook that answers yes from its `at`-th ask on. +pub const Cut = struct { + at: u32, + asked: u32 = 0, + + fn ask(ptr: *anyopaque) bool { + const cut: *Cut = @ptrCast(@alignCast(ptr)); + cut.asked += 1; + return cut.asked > cut.at; + } +}; + +/// Drafts from a known reply (the plain run's): 2 to 15 a round, most rounds one of them wrong, to reject anywhere. +pub const Oracle = struct { + prompt_len: usize, + truth: []const u32, + out: [16]u32 = undefined, + rounds: u32 = 0, + + pub fn proposer(o: *Oracle) lanes.proposer.Proposer { + return .{ .ptr = o, .vtable = &.{ .propose = propose, .last_match = lastMatch } }; + } + + fn propose(ptr: *anyopaque, context: []const u32, max_draft: i64) anyerror![]const u32 { + const o: *Oracle = @ptrCast(@alignCast(ptr)); + o.rounds += 1; + const k = context.len - o.prompt_len; // the reply so far, its last token pending: drafts follow it + if (max_draft < 2 or k >= o.truth.len) return &.{}; + const n = @min(@as(usize, @intCast(max_draft)), o.truth.len - k, 2 + o.rounds % 14); + if (n < 2) return &.{}; + @memcpy(o.out[0..n], o.truth[k..][0..n]); + if (o.rounds % 4 != 0) o.out[(o.rounds * 5) % n] ^= 1; // a wrong token at a moving row + return o.out[0..n]; + } + + fn lastMatch(_: *anyopaque) i64 { + return 64; + } +}; + +/// A concurrent case's stream: prompt `p` joining at round `at`; `filled`: its prompt a chunk at a time between rounds. +pub const Join = struct { p: usize, at: u32 = 0, sampled: bool = false, discard_at: ?u32 = null, filled: bool = false, drafts: bool = false }; + +pub const Cover = struct { + gpa: std.mem.Allocator, + a: std.mem.Allocator, + le: *lanes.Engine, + failed: bool = false, + checks: usize = 0, + + pub const Ask = struct { + prompt: []const u32, + max_new: u32, + eos: []const u32 = &.{}, + drafts: bool = true, + sampling: ?lanes.Sampling = null, + truth: []const u32 = &.{}, // drafts from this reply (Oracle); empty: none + cancel: ?*Cut = null, + reuse: lanes.stream.Reuse = .{}, + }; + pub const Got = struct { ids: []const u32, reason: lanes.stream.Reason, rounds: u64, cached: u32, drafted: u64, accepted: u64 }; + + /// One request through the shared lanes.Engine, as the server's lane host runs it. + pub fn run(c: *Cover, ask: Ask) !Got { + var o: Oracle = .{ .prompt_len = ask.prompt.len, .truth = ask.truth }; + var s = try lanes.Stream.init(c.gpa, .{ .id = "cover", .prompt = ask.prompt, .max_new = ask.max_new, .eos = ask.eos, .drafts = ask.drafts, .sampling = ask.sampling, .proposer = if (ask.truth.len > 0) o.proposer() else null, .cancel_check = if (ask.cancel) |cut| .{ .ptr = cut, .check = Cut.ask } else null, .reuse = ask.reuse }); + defer s.deinit(c.gpa); + try c.le.addStream(&s); + while (c.le.activeCount() > 0) try c.le.step(); + return .{ .ids = try c.a.dupe(u32, s.emitted()), .reason = s.reason, .rounds = s.rounds, .cached = s.cached, .drafted = s.drafted, .accepted = s.accepted }; + } + + pub const Together = struct { ids: []const []const u32, most: usize, shared: u32, between: u32, drafted: u64 }; + + /// The streams of `joins` through the rounds together, each joining at its round; how many were live at once. + pub fn together(c: *Cover, joins: []const Join, ps: []const []const u32, truths: []const []const u32, seeds: []const lanes.Sampling, max_new: u32) !Together { + const ss = try c.a.alloc(lanes.Stream, joins.len); + const looks = try c.a.alloc(Oracle, joins.len); + const made = try c.a.alloc(bool, joins.len); + @memset(made, false); + defer for (ss, made) |*s, m| if (m) s.deinit(c.gpa); + var most: usize = 0; + var shared: u32 = 0; + var between: u32 = 0; + var round: u32 = 0; + var filling: std.ArrayList(*lanes.Stream) = .empty; // the lane host's fillOne: the oldest's next chunk a round + defer filling.deinit(c.gpa); + var began = false; + while (true) : (round += 1) { + for (joins, ss, looks, made) |j, *s, *l, *m| if (!m.* and j.at <= round) { + l.* = .{ .prompt_len = ps[j.p].len, .truth = truths[j.p] }; + s.* = try lanes.Stream.init(c.gpa, .{ .id = "joined", .prompt = ps[j.p], .max_new = max_new, .sampling = if (j.sampled) seeds[j.p % seeds.len] else null, .proposer = if (j.drafts) l.proposer() else null }); + m.* = true; + if (j.filled) try filling.append(c.gpa, s) else try c.le.addStream(s); + }; + if (filling.items.len > 0) { + if (try c.le.fillStream(filling.items[0], !began)) { + _ = filling.orderedRemove(0); + began = false; + } else began = true; + } + for (joins, ss, made) |j, *s, m| if (m and j.discard_at == round and !s.finished) c.le.discard(s); + const live = c.le.activeCount(); + most = @max(most, live); + shared += @intFromBool(live >= 2); + if (live == 0) { + if (std.mem.indexOfScalar(bool, made, false) == null and filling.items.len == 0) break; + continue; + } + between += @intFromBool(began); + try c.le.step(); + } + const ids = try c.a.alloc([]const u32, joins.len); + var drafted: u64 = 0; + for (ss, ids) |*s, *i| { + i.* = try c.a.dupe(u32, s.emitted()); + drafted += s.drafted; + } + return .{ .ids = ids, .most = most, .shared = shared, .between = between, .drafted = drafted }; + } + + /// A drafted stream over `prompt` (drafts from `truth`), discarded after `rounds`; false if it ended first. + pub fn discardAfter(c: *Cover, prompt: []const u32, truth: []const u32, rounds: u32) !bool { + var o: Oracle = .{ .prompt_len = prompt.len, .truth = truth }; + var s = try lanes.Stream.init(c.gpa, .{ .id = "discarded", .prompt = prompt, .max_new = 512, .proposer = o.proposer() }); + defer s.deinit(c.gpa); + try c.le.addStream(&s); + var n: u32 = 0; + while (c.le.activeCount() > 0 and n < rounds) : (n += 1) try c.le.step(); + if (c.le.activeCount() == 0) return false; + c.le.discard(&s); + return true; + } + + pub fn same(c: *Cover, what: []const u8, got: Got, want: []const u32, reason: lanes.stream.Reason) !void { + c.checks += 1; + const n = @min(got.ids.len, want.len); + const first = for (got.ids[0..n], want[0..n], 0..) |x, y, i| { + if (x != y) break i; + } else if (got.ids.len != want.len) n else null; + if (first == null and got.reason == reason) { + check.pass("{s}: {d} ids equal, finished by {s} ({d} rounds, {d} of {d} drafts kept)", .{ what, want.len, reason.name(), got.rounds, got.accepted, got.drafted }); + return; + } + c.failed = true; + if (first) |i| { + std.debug.print("FAIL {s}: token {d} of {d} is {d}, want {d} (emitted {d}, {s})\n", .{ what, i, want.len, if (i < got.ids.len) got.ids[i] else 0, if (i < want.len) want[i] else 0, got.ids.len, got.reason.name() }); + } else std.debug.print("FAIL {s}: finished by {s}, want {s}\n", .{ what, got.reason.name(), reason.name() }); + } + + pub fn expect(c: *Cover, ok: bool, comptime fmt: []const u8, args: anytype) !void { + c.checks += 1; + if (ok) return check.pass(fmt, args); + c.failed = true; + std.debug.print("FAIL " ++ fmt ++ "\n", args); + } +}; + +pub fn back(e: *Engine, a: std.mem.Allocator, ptr: u64, len: usize) ![]u8 { + const out = try a.alloc(u8, len); + try e.stream.synchronize(); + const b: cuda.DeviceBuffer = .{ .d = e.d, .ptr = ptr, .len = len }; + try b.download(0, out); + return out; +} + +/// Slots `x` and `y` hold the same committed sequence: position, KDA states and conv windows, MLA rows below it. +pub fn sameSeq(c: *Cover, e: *Engine, what: []const u8, x: u32, y: u32) !void { + const p = &e.seqs.items[x].?.st; + const q = &e.seqs.items[y].?.st; + var bad: ?[]const u8 = null; + if (p.pos != q.pos) bad = "positions"; + for (0..p.kda_layers) |li| { + if (!std.mem.eql(u8, try back(e, c.a, p.recAt(p.cur[li], li), p.rec_layer), try back(e, c.a, q.recAt(q.cur[li], li), q.rec_layer))) bad = "a KDA state"; + if (!std.mem.eql(u8, try back(e, c.a, p.convAt(li), 3 * p.convWidth() * 2), try back(e, c.a, q.convAt(li), 3 * q.convWidth() * 2))) bad = "a conv window"; + } + for (0..p.mla_layers) |mi| if (!std.mem.eql(u8, try back(e, c.a, p.cacheAt(mi), p.pos * ling.forward.cache_width * 2), try back(e, c.a, q.cacheAt(mi), p.pos * ling.forward.cache_width * 2))) { + bad = "MLA cache rows"; + }; + if (e.seqs.items[x].?.hd) |*ph| { + const qh = &e.seqs.items[y].?.hd.?; + if (ph.pos != qh.pos) bad = "MTP positions"; + if (!std.mem.eql(u8, try back(e, c.a, ph.cacheAt(0), ph.pos * ling.forward.cache_width * 2), try back(e, c.a, qh.cacheAt(0), ph.pos * ling.forward.cache_width * 2))) bad = "MTP cache rows"; + } + try c.expect(bad == null, "{s}: slots {d} and {d} at {d}, byte for byte{s}{s}", .{ what, x, y, p.pos, if (bad != null) ", but differ in " else "", bad orelse "" }); +} diff --git a/zig/tests/cuda/ling_layer_bytes.zig b/zig/tests/cuda/ling_layer_bytes.zig new file mode 100644 index 0000000000..7414298946 --- /dev/null +++ b/zig/tests/cuda/ling_layer_bytes.zig @@ -0,0 +1,119 @@ +//! `ling-layer-bytes DIR FIXTURE OUT`: one real Ling-3.1 layer's outputs and states on every path, written for cmp. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; +const Rig = rig.Rig; +const Opened = @typeInfo(@typeInfo(@TypeOf(rig.W.open)).@"fn".return_type.?).error_union.payload; + +/// The loaded layer, its graph, a prompt and a decode Buffers, the fixture's rows tiled to `rows`. +pub const Kit = struct { + o: Opened, + m: F.Modules, + w: rig.W.Weights, + g: F.Graph, + r: Rig, + pb: F.Buffers, + db: F.Buffers, + x: []u16, + D: usize, + L: u32, + cap: usize, + + /// Loads FIXTURE's layer from DIR; prompt buffers for `rows` rows, states of `cap` positions. + pub fn init(k: *Kit, gpu: check.Gpu, a: std.mem.Allocator, s: cuda.Stream, dir: []const u8, path: []const u8, rows: usize, cap: usize) !void { + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 30)); + if (bytes.len < 24 or !std.mem.eql(u8, bytes[0..8], "TFLINGL1")) return error.BadFixture; + const dims = std.mem.bytesAsSlice(u32, bytes[8..24]); + k.L = dims[0]; + const T: usize = dims[1]; + k.D = dims[2]; + k.x = try a.alloc(u16, rows * k.D); + for (0..rows) |t| @memcpy(std.mem.sliceAsBytes(k.x[t * k.D ..][0..k.D]), bytes[24 + (t % T) * k.D * 2 ..][0 .. k.D * 2]); + var why: ling.config.Why = .{}; + k.o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + k.m = try F.Modules.load(gpu.d, gpu.ctx); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + k.w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &k.o.ck, k.o.c, &.{k.L}, .{ .globals = false }); + k.g = try F.Graph.init(gpu.gpa, gpu.d, &k.m, s, &k.w); + k.r = .{ .gpu = gpu, .a = a, .s = s, .m = &k.m, .g = &k.g }; + k.cap = cap; + k.pb = try F.Buffers.init(gpu.d, &k.m, &k.w, rows, true, cap); + k.db = try F.Buffers.init(gpu.d, &k.m, &k.w, F.window_rows, false, cap); + } + + pub fn deinit(k: *Kit) void { + k.pb.deinit(); + k.db.deinit(); + k.r.deinit(); + k.g.deinit(); + k.w.deinit(); + k.m.deinit(); + k.o.ck.close(); + } + + /// Rows lo .. hi through the layer on `st`, committing `keep` of them (0: all); the output rows' bytes. + pub fn pass(k: *Kit, b: *const F.Buffers, st: *F.State, lo: usize, hi: usize, keep: usize) ![]u8 { + const d = k.r.gpu.d; + const n = hi - lo; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x[lo * k.D ..].ptr, n * k.D * 2, k.r.s.handle), "cuMemcpyHtoDAsync"); + const f = k.r.fwd(b, null); + try f.stack(&.{.{ .st = st, .row0 = 0, .rows = n }}); + const out = try k.r.back(b.x, n * k.D * 2); + try f.commit(st, if (keep == 0) n else keep); + return out; + } +}; + +fn put(list: *std.ArrayList(u8), a: std.mem.Allocator, bytes: []const u8) !void { + try list.appendSlice(a, bytes); +} + +/// Every path's bytes: prompts of 512 and 2048 rows, 16 decode rows, windows of 2..16 with partial commits, states. +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, out: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var k: Kit = undefined; + try k.init(gpu, a, s, dir, path, 2048, 2600); + defer k.deinit(); + var all: std.ArrayList(u8) = .empty; + const P = 400; + for ([_]usize{ 512, 2048 }) |T| { + const st = try k.r.state(k.cap); + defer st.deinit(); + try put(&all, a, try k.pass(&k.pb, st, 0, T, 0)); + try put(&all, a, try k.r.back(st.arena.buf.ptr, st.arena.buf.len)); + } + for ([_]bool{ true, false }) |wide| { + k.r.kda_wide = wide; + const st = try k.r.state(k.cap); + defer st.deinit(); + _ = try k.pass(&k.pb, st, 0, P, 0); + for (0..F.window_rows) |i| try put(&all, a, try k.pass(&k.db, st, P + i, P + i + 1, 0)); + try put(&all, a, try k.r.back(st.arena.buf.ptr, st.arena.buf.len)); + for ([_]usize{ 2, 3, 4, 8, 16 }) |w| { + const sw = try k.r.state(k.cap); + defer sw.deinit(); + _ = try k.pass(&k.pb, sw, 0, P, 0); + const keep = (w + 1) / 2; + try put(&all, a, try k.pass(&k.db, sw, P, P + w, keep)); + try put(&all, a, try k.pass(&k.db, sw, P + keep, P + keep + w, 0)); + try put(&all, a, try k.r.back(sw.arena.buf.ptr, sw.arena.buf.len)); + } + } + k.r.kda_wide = true; + var h = std.hash.Wyhash.init(0); + h.update(all.items); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = out, .data = all.items }); + check.pass("ling-layer-bytes layer {d}: {d} bytes into {s}, wyhash {x:0>16}", .{ k.L, all.items.len, out, h.final() }); +} diff --git a/zig/tests/cuda/ling_layer_real.zig b/zig/tests/cuda/ling_layer_real.zig new file mode 100644 index 0000000000..1de72f6e28 --- /dev/null +++ b/zig/tests/cuda/ling_layer_real.zig @@ -0,0 +1,155 @@ +//! `ling-layer-real DIR FIXTURE [DUMP]`: one real Ling-3.1 layer through the graph vs HF's output; exactness on it. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; +const Rig = rig.Rig; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +/// Rows lo .. hi of x through the stack on `st` (prompt or decode buffers), committed; the output rows. +fn pass(r: *Rig, b: *const F.Buffers, st: *F.State, x: []const u16, D: usize, lo: usize, hi: usize, picks: ?u64) ![]u16 { + const d = r.gpu.d; + const n = hi - lo; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, x[lo * D ..].ptr, n * D * 2, r.s.handle), "cuMemcpyHtoDAsync"); + const f = r.fwd(b, picks); + try f.stack(&.{.{ .st = st, .row0 = 0, .rows = n }}); + const out = try r.a.alloc(u16, n * D); + @memcpy(std.mem.sliceAsBytes(out), try r.back(b.x, n * D * 2)); + try f.commit(st, n); + return out; +} + +/// Relative L2 of bf16 rows against fp32 rows, all and without `skip` rows, and the worst row's. +fn dist(got: []const u16, want: []const f32, rows: usize, D: usize, skip: []const bool) [3]f64 { + var dd: f64 = 0; + var ww: f64 = 0; + var kd: f64 = 0; + var kw: f64 = 0; + var worst: f64 = 0; + for (0..rows) |t| { + var rd: f64 = 0; + var rw: f64 = 0; + for (0..D) |j| { + const e: f64 = want[t * D + j]; + const g = rig.bfOf(got[t * D + j]); + rd += (g - e) * (g - e); + rw += e * e; + } + dd += rd; + ww += rw; + if (!skip[t]) { + kd += rd; + kw += rw; + worst = @max(worst, @sqrt(rd / rw)); + } + } + return .{ @sqrt(dd / ww), @sqrt(kd / kw), worst }; +} + +/// `ling-layer-real DIR FIXTURE` (oracle/ling_layer.py's file for layer L of the checkpoint at DIR). +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, dump: ?[]const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 30)); + if (bytes.len < 24 or !std.mem.eql(u8, bytes[0..8], "TFLINGL1")) return error.BadFixture; + var at: usize = 8; + const dims = try take(u32, a, bytes, &at, 4); + const L = dims[0]; + const T: usize = dims[1]; + const D: usize = dims[2]; + const K: usize = dims[3]; + const x = try take(u16, a, bytes, &at, T * D); + const want = try take(f32, a, bytes, &at, T * D); + const topk = try take(i32, a, bytes, &at, T * K); + if (at != bytes.len) return error.BadFixture; + var why: ling.config.Why = .{}; + var o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + var m = try F.Modules.load(gpu.d, gpu.ctx); + defer m.deinit(); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + const t0 = check.now(gpu.io); + var w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &o.ck, o.c, &.{L}, .{ .globals = false }); + defer w.deinit(); + const load_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + var g = try F.Graph.init(gpu.gpa, gpu.d, &m, s, &w); + defer g.deinit(); + const c = &w.config; + try check.expect(D == c.hidden and K == (if (c.isMoe(L)) @as(usize, c.topk) else 0), "fixture: layer {d}, {d} wide, top {d}", .{ L, D, K }); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .m = &m, .g = &g }; + defer r.deinit(); + const cap = T + 32; + var pb = try F.Buffers.init(gpu.d, &m, &w, T, true, cap); + defer pb.deinit(); + var db = try F.Buffers.init(gpu.d, &m, &w, F.window_rows, false, cap); + defer db.deinit(); + // the whole prompt against HF, picks against HF's + const top: usize = c.topk; + const pk = try r.zeros(T * (top + 1) * 4); + const st = try r.state(cap); + defer st.deinit(); + const got = try pass(&r, &pb, st, x, D, 0, T, if (K > 0) pk else null); + const skip = try a.alloc(bool, T); + @memset(skip, false); + var nf: usize = 0; + if (K > 0) { + const ours = try rig.sortedPicks(a, try r.back(pk, T * (top + 1) * 4), 1, T, top); + const theirs = try a.dupe(i32, topk); + for (0..T) |t| std.mem.sort(i32, theirs[t * K ..][0..K], {}, std.sort.asc(i32)); + nf = rig.flips(ours, theirs, 1, T, 0, T, top, skip); + } + const dv = dist(got, want, T, D, skip); + std.debug.print("RESULT ling-layer-real layer {d} ({t} {s}, {t}): {d} tokens, loaded in {d:.1} s\n", .{ L, c.kind(L), if (K > 0) "MoE" else "dense", c.quant, T, load_s }); + std.debug.print("RESULT vs HF: rel L2 {e:.3}, {e:.3} without the {d} rows whose top-{d} differs, worst kept row {e:.3}\n", .{ dv[0], dv[1], nf, top, dv[2] }); + // cut prompt, decode rows, a verify window + const sc = try r.state(cap); + defer sc.deinit(); + const cut0 = try pass(&r, &pb, sc, x, D, 0, 300, null); + const cut1 = try pass(&r, &pb, sc, x, D, 300, T, null); + try check.sameBytes("prompt cut at 300 == uncut", std.mem.sliceAsBytes(cut0), std.mem.sliceAsBytes(got[0 .. 300 * D])); + try check.sameBytes("prompt cut at 300 == uncut (the rest)", std.mem.sliceAsBytes(cut1), std.mem.sliceAsBytes(got[300 * D ..])); + const P = 400; + const N = F.window_rows; + const sd = try r.state(cap); + defer sd.deinit(); + const sw = try r.state(cap); + defer sw.deinit(); + _ = try pass(&r, &pb, sd, x, D, 0, P, null); + _ = try pass(&r, &pb, sw, x, D, 0, P, null); + const dec = try a.alloc(u16, N * D); + for (0..N) |i| @memcpy(dec[i * D ..][0..D], try pass(&r, &db, sd, x, D, P + i, P + i + 1, null)); + try check.sameBytes("16-row verify window == 16 decode rows", std.mem.sliceAsBytes(try pass(&r, &db, sw, x, D, P, P + N, null)), std.mem.sliceAsBytes(dec)); + const dd = dist(dec, want[P * D ..], N, D, skip[P..]); + const dp = dist(dec, @as([]const f32, try widen(a, got[P * D .. (P + N) * D])), N, D, skip[P..]); + std.debug.print("RESULT decode rows {d}-{d}: rel L2 {e:.3} vs HF, {e:.3} vs the prompt's rows\n", .{ P, P + N - 1, dd[0], dp[0] }); + if (dump) |d| { // the prompt's and the decode rows' bytes, for comparing two builds + const all = try a.alloc(u16, got.len + dec.len); + @memcpy(all[0..got.len], got); + @memcpy(all[got.len..], dec); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = d, .data = std.mem.sliceAsBytes(all) }); + } + try check.expect(dv[1] <= 1e-2 and dd[1] <= 1e-2, "layer {d} within rel L2 1e-2 of HF (rows whose picks agree)", .{L}); + check.pass("ling-layer-real layer {d}: rel L2 {e:.3} of HF ({d} flipped rows apart); cut == uncut, window == decode rows; peak {d} MiB on the device", .{ L, dv[1], nf, cuda.usage(false).peak >> 20 }); +} + +fn widen(a: std.mem.Allocator, x: []const u16) ![]f32 { + const out = try a.alloc(f32, x.len); + for (out, x) |*o, v| o.* = @floatCast(rig.bfOf(v)); + return out; +} diff --git a/zig/tests/cuda/ling_layer_speed.zig b/zig/tests/cuda/ling_layer_speed.zig new file mode 100644 index 0000000000..6b36204dfc --- /dev/null +++ b/zig/tests/cuda/ling_layer_speed.zig @@ -0,0 +1,308 @@ +//! `ling-layer-speed DIR FIXTURE [REPS]`: one real Ling-3.1 layer timed by CUDA events, whole and op by op, per phase. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); +const bytes = @import("ling_layer_bytes.zig"); + +const F = rig.F; +const G = ling.graph; +const W = rig.W; +const kda = G.kda; +const latent = G.latent; +const glue = G.glue; + +/// Events between op groups; each group's milliseconds by name, medians and least over reps. +const Timer = struct { + ev: [64]cuda.Event, + names: [64][]const u8 = undefined, + n: usize = 0, + s: cuda.Stream, + on: bool, + spin: cuda.Function, + + fn mark(t: *Timer, name: []const u8) !void { + if (!t.on) return; + try t.ev[t.n].record(t.s); + t.names[t.n] = name; + t.n += 1; + } + + fn spinFor(t: *Timer, ns: i64) !void { + var a: cuda.Args = .{}; + a.add(@as(c_longlong, ns)); + try cuda.launch.launch(t.spin, .{ .grid = .{}, .block = .{} }, t.s, &a); + } + + fn ms(t: *const Timer, j: usize) !f32 { + return t.ev[j - 1].elapsedMs(t.ev[j]); + } +}; + +const Seg = F.Seg; + +fn kdaBlock(o: G.Ops, b: *const F.Buffers, a: W.Kda, c: *const ling.config.Config, li: usize, sg: Seg, rows: usize, t: *Timer) !void { + const D: usize = c.hidden; + const H: usize = c.kda_heads; + const n = H * kda.dk; + const st = sg.st; + const ps = st.stride; + const sv = st.saves(li); + const p = if (b.prompt) b.kp else sv.p; + for ([_]W.Proj{ a.q, a.k, a.v, a.b }, [_]usize{ 0, n, 2 * n, 3 * n }, [_][]const u8{ "proj q", "proj k", "proj v", "proj b (bf16)" }) |w, col0, name| { + if (rows == 1) { + try o.mm(w, b.normed, D, 1, p + col0 * 2); + try t.mark(name); + continue; + } + try o.mm(w, b.normed, D, rows, b.tmp); + try t.mark(name); + const w_n = w.outputs(); + try o.copy(b.tmp, w_n * 2, p + col0 * 2, ps * 2, rows, w_n * 2); + try t.mark("copies into p"); + } + try o.mm(a.f, b.normed, D, rows, b.ka); + try t.mark("proj f"); + try o.mm(a.g, b.normed, D, rows, b.kg); + try t.mark("proj g"); + const cur = st.cur[li]; + const ch: kda.Chain = .{ .p = p, .p_stride = ps, .b_off = 3 * n, .a = b.ka, .a_stride = n, .g = b.kg, .g_stride = n, .conv_state = st.convAt(li), .conv_w = a.conv, .state_in = st.recAt(cur, li), .a_log = a.a_log, .dt_bias = a.dt_bias, .norm_w = a.o_norm, .eps = c.eps, .lower = c.lower_bound, .heads = H }; + const k = if (b.prompt) b.saves.k else sv.k; + const v = if (b.prompt) b.saves.v else sv.v; + const g = if (b.prompt) b.saves.g else sv.g; + const bb = if (b.prompt) b.saves.b else sv.b; + const s: kda.Saves = .{ .out = b.saves.out, .state_out = st.recAt(1 - cur, li), .k = k, .v = v, .g = g, .b = bb }; + const wide: ?kda.Wide = if (if (b.prompt) rows >= kda.wide_rows else o.kda_wide) .{ .q = b.wide.q, .y = b.wide.y } else null; + try kda.chain(&o.m.kda, o.s, ch, rows, s, wide); + try t.mark("KDA chain"); + if (b.prompt) { + try glue.convShift(o.set(), o.s, st.convAt(li), p, rows, rows * ps, ps, 1, 3 * n); + try t.mark("conv shift"); + } + try o.mm(a.o, b.saves.out, n, rows, b.y); + try t.mark("proj o"); +} + +fn mlaBlock(o: G.Ops, b: *const F.Buffers, a: W.Mla, ab: G.Absorb, g: *const F.Graph, sg: Seg, rows: usize, t: *Timer) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const H: usize = c.mla_heads; + const qk: usize = c.qkDim(); + const nope: usize = c.nope; + const kw = F.cache_width; + const lw = 512; + const ql: usize = c.q_lora; + try o.mm(a.q_a.?, b.normed, D, rows, b.qa_in); + try t.mark("proj q_a"); + try glue.rmsnorm(o.set(), o.s, b.qa_in, ql, a.q_a_norm, b.qr, ql, null, c.eps, rows, ql); + try t.mark("norms q_a, kv"); + try o.mm(a.q_b.?, b.qr, ql, rows, b.q); + try t.mark("proj q_b"); + try o.mm(a.kv_a, b.normed, D, rows, b.kv); + try t.mark("proj kv_a"); + try glue.rmsnorm(o.set(), o.s, b.kv, kw, a.kv_a_norm, b.lat, kw, null, c.eps, rows, lw); + try t.mark("norms q_a, kv"); + try o.mm(a.gate.?, b.normed, D, rows, b.g); + try t.mark("proj gate (bf16)"); + try latent.absorbQStrided(o.set(), o.s, b.q, qk, ab.wk, b.qa, kw, rows, H, nope, lw); + try t.mark("absorb q"); + const st = sg.st; + try latent.rope(o.set(), o.s, .{ .x = b.q + nope * 2, .x_row = H * qk, .x_head = qk, .out = b.qa + lw * 2, .o_row = H * kw, .o_head = kw }, rows, H, st.pos_dev, g.freq); + try latent.rope(o.set(), o.s, .{ .x = b.kv + lw * 2, .x_row = kw, .x_head = 0, .out = b.lat + lw * 2, .o_row = kw, .o_head = 0 }, rows, 1, st.pos_dev, g.freq); + try latent.latentWrite(o.set(), o.s, b.lat, kw, st.cacheAt(0), st.pos_dev, rows, kw); + try t.mark("rope, cache write"); + var r0: usize = 0; + while (r0 < rows) : (r0 += ling.state.att_rows) { + const m = @min(ling.state.att_rows, rows - r0); + var at = st.pos_dev; + if (r0 > 0) { + try o.addI32(st.pos_dev, r0, b.at); + at = b.at; + } + const nch = latent.chunksFor(st.pos + r0 + m); + try latent.attentionRoped(o.set(), o.s, b.qa + r0 * H * kw * 2, st.cacheAt(0), at, b.partials, b.ol + r0 * H * lw * 2, b.dummy, m, H, lw, nch, c.smScale()); + } + try t.mark("latent attention"); + try latent.expandV(o.set(), o.s, b.ol, ab.wv, b.vn, rows, H, c.v_dim, lw); + try t.mark("expand v"); + try o.gate(b.vn, b.g, H, rows, H, c.v_dim); + try t.mark("head gate"); + try o.mm(a.o, b.vn, H * c.v_dim, rows, b.y); + try t.mark("proj o"); +} + +fn mlp(o: G.Ops, b: *const F.Buffers, m: W.Mlp, rows: usize, limit: f32, out: u64, t: *Timer, comptime tag: []const u8) !void { + const D = m.gate.inputs(); + const w = m.gate.outputs(); + for ([_]W.Proj{ m.gate, m.up }, [_]usize{ 0, w }) |p, col0| { + if (rows == 1) { + try o.mm(p, b.normed, D, 1, b.gu + col0 * 2); + try t.mark(tag ++ " gate, up"); + continue; + } + try o.mm(p, b.normed, D, rows, b.tmp); + try t.mark(tag ++ " gate, up"); + try o.copy(b.tmp, w * 2, b.gu + col0 * 2, 4 * w, rows, w * 2); + try t.mark(tag ++ " copies"); + } + try glue.swigluWith(o.m.swiglu, o.s, b.gu, b.act, b.xs, limit, rows, w); + try t.mark(tag ++ " swiglu"); + try o.mm(m.down, b.act, w, rows, out); + try t.mark(tag ++ " down"); +} + +fn moe(o: G.Ops, b: *const F.Buffers, e: W.Moe, g: *const F.Graph, rows: usize, t: *Timer) !void { + const c = &g.w.config; + const D: usize = c.hidden; + const E: usize = c.experts; + const top: usize = c.topk; + const slots = top + 1; + try glue.router(o.set(), o.s, b.normed, D, e.router, b.mlog, b.rpart, rows, D, E); + try t.mark("router"); + try glue.select(o.set(), o.s, b.mlog, if (e.bias != 0) e.bias else g.zero_bias, b.pick, b.wts, rows, top, E, c.routed_scale, c.norm_topk, ling.groups(c)); + try t.mark("select"); + const l = e.routed.grouped; + const gr = o.m.grouped(); + const up: cuda.experts.Rows = .{ .x = b.normed, .stride = l.dims, .slots = slots }; + const down: cuda.experts.Rows = .{ .x = b.eact, .stride = l.width }; + const skip: i32 = @intCast(E); + // block_ffn's choice first, then (prompt rows) every kernel and plan tile on the same rows for comparison + const pk = o.prompt and !(l.format == .fp8g and rows * top < 16 * E); + const Way = struct { prompt: bool, tile: usize, up: []const u8, down: []const u8 }; + const ways = [_]Way{ + .{ .prompt = pk, .tile = if (pk) 64 else 16, .up = "experts gate-up", .down = "experts down" }, + .{ .prompt = true, .tile = 64, .up = "alt gate-up, prompt kernel, tile 64", .down = "alt down, prompt kernel, tile 64" }, + .{ .prompt = false, .tile = 16, .up = "alt gate-up, decode kernel, tile 16", .down = "alt down, decode kernel, tile 16" }, + .{ .prompt = false, .tile = 32, .up = "alt gate-up, decode kernel, tile 32", .down = "alt down, decode kernel, tile 32" }, + .{ .prompt = false, .tile = 64, .up = "alt gate-up, decode kernel, tile 64", .down = "alt down, decode kernel, tile 64" }, + }; + for (ways, 0..) |w, i| { + if (i > 0 and !o.prompt) break; + try o.m.router.route(o.s, b.pick, rows * slots, E + 1, w.tile, b.plan); + try t.mark(if (i == 0) "expert plan" else "alt plans"); + const items = cuda.grouped.maxItems(rows * slots, E + 1, w.tile); + if (w.prompt) try gr.prompt(o.s, .up, up, l, b.plan, b.eact, items, skip) else try gr.decode(o.s, .up, up, l, b.plan, b.eact, items, skip); + try t.mark(w.up); + if (w.prompt) try gr.prompt(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip) else try gr.decode(o.s, .down_bf16, down, l, b.plan, b.ey, items, skip); + try t.mark(w.down); + } + try mlp(o, b, e.shared.?, rows, e.shared_limit, b.sh, t, "shared"); + var a: cuda.Args = .{}; + inline for (.{ b.ey, b.pick, b.wts, b.sh, b.y, @as(i64, @intCast(rows)), @as(i64, @intCast(D)), @as(c_int, @intCast(slots)), @as(c_int, @intCast(top)), @as(c_int, 0) }) |v| a.add(v); + try cuda.launch.launch(o.m.combine, .{ .grid = .{ .x = @intCast(rows), .y = @intCast((D + 255) / 256) }, .block = .{ .x = 256 } }, o.s, &a); + try t.mark("combine"); +} + +/// The layer as forward.layer runs it, op groups marked. +fn layerOps(f: *const F.Forward, sg: Seg, rows: usize, t: *Timer) !void { + const o = f.ops(); + const b = f.b; + const g = f.g; + const c = &g.w.config; + const D: usize = c.hidden; + const l = g.w.layers[0]; + try t.mark("start"); + try glue.rmsnorm(o.set(), o.s, b.x, D, l.attn_norm, b.normed, D, null, c.eps, rows, D); + try t.mark("norms + residual adds"); + switch (l.mixer) { + .kda => |a| try kdaBlock(o, b, a, c, 0, sg, rows, t), + .mla => |a| try mlaBlock(o, b, a, g.absorb[0], g, sg, rows, t), + } + try o.add(b.x, b.y, b.x, rows * D); + try glue.rmsnorm(o.set(), o.s, b.x, D, l.mlp_norm, b.normed, D, null, c.eps, rows, D); + try t.mark("norms + residual adds"); + switch (l.ffn) { + .dense => |m| try mlp(o, b, m, rows, 0, b.y, t, "dense"), + .moe => |e| try moe(o, b, e, g, rows, t), + } + try o.add(b.x, b.y, b.x, rows * D); + try t.mark("norms + residual adds"); +} + +const Row = struct { name: []const u8, ms: [32]f64 = undefined, n: usize = 0 }; + +fn median(xs: []f64) f64 { + std.mem.sort(f64, xs, {}, std.sort.asc(f64)); + return xs[xs.len / 2]; +} + +/// `ling-layer-speed DIR FIXTURE [REPS [ROWS,..]]`: decode at position 2048, windows, prompts (default 512, 2048). +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8, reps_arg: ?[]const u8, rows_arg: ?[]const u8) !void { + const reps: usize = @min(30, if (reps_arg) |x| try std.fmt.parseInt(usize, x, 10) else 15); + var phases: std.ArrayList(usize) = .empty; + var it = std.mem.tokenizeScalar(u8, rows_arg orelse "1,2,4,8,16,512,2048", ','); + while (it.next()) |x| try phases.append(gpu.gpa, try std.fmt.parseInt(usize, x, 10)); + defer phases.deinit(gpu.gpa); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var k: bytes.Kit = undefined; + try k.init(gpu, a, s, dir, path, 2048, 4200); + defer k.deinit(); + var probe = try cuda.Module.load(gpu.d, cuda.kernels.probe); + defer probe.unload(); + var t: Timer = .{ .ev = undefined, .s = s, .on = true, .spin = try probe.function("tf_probe_spin") }; + for (&t.ev) |*e| e.* = try cuda.Event.init(gpu.d, true); + defer for (&t.ev) |*e| e.deinit(); + var prng = std.Random.DefaultPrng.init(0x5eed); + const st = try k.r.state(k.cap); + defer st.deinit(); + const c = &k.w.config; + std.debug.print("RESULT ling-layer-speed layer {d} ({t}, {s}, {t}), {d} reps, median and least us (each rep after a random spin: other contexts' slices land at random)\n", .{ k.L, c.kind(k.L), if (c.isMoe(k.L)) "MoE" else "dense", c.quant, reps }); + for (phases.items) |rows| { + if (rows == 0 or rows > 2048) return error.TooManyRows; + const b = if (rows > F.window_rows) &k.pb else &k.db; + const pos: usize = if (rows > F.window_rows) 0 else 2048; + try st.setPos(s, pos); + const d = gpu.d; + const f = k.r.fwd(b, null); + const sg: Seg = .{ .st = st, .row0 = 0, .rows = rows }; + var whole: [64]f64 = undefined; + var table: [32]Row = undefined; + var nt: usize = 0; + for (0..reps + 2) |i| { + // a spin of random length first: other contexts' time slices land at a new phase each rep, least is clean + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.x.ptr, rows * k.D * 2, s.handle), "cuMemcpyHtoDAsync"); + try t.spinFor(prng.random().intRangeLessThan(i64, 0, 4_000_000)); + try f.layer(&.{sg}, 0, rows, null); // untimed: the GPU warm before the timed launches + try t.ev[62].record(s); + try f.layer(&.{sg}, 0, rows, null); + try t.ev[63].record(s); + t.n = 0; // then the same launches op by op, back to back + try layerOps(&f, sg, rows, &t); + try s.synchronize(); + if (i >= 2) whole[i - 2] = try t.ev[62].elapsedMs(t.ev[63]); + if (i < 2) continue; + var acc: [32]f64 = @splat(0); + for (1..t.n) |j| { + const name = t.names[j]; + var at: usize = 0; + while (at < nt and !std.mem.eql(u8, table[at].name, name)) at += 1; + if (at == nt) { + table[nt] = .{ .name = name }; + nt += 1; + } + acc[at] += try t.ms(j); + } + for (table[0..nt], acc[0..nt]) |*r, v| { + r.ms[r.n] = v; + r.n += 1; + } + } + var sum: f64 = 0; + var low: f64 = 0; + std.debug.print("RESULT rows {d:>4}: layer {d:.1} {d:.1}", .{ rows, median(whole[0..reps]) * 1000, std.mem.min(f64, whole[0..reps]) * 1000 }); + for (table[0..nt]) |*r| { + const lo = std.mem.min(f64, r.ms[0..r.n]) * 1000; + const m = median(r.ms[0..r.n]) * 1000; + sum += m; + low += lo; + std.debug.print(" | {s} {d:.1} {d:.1}", .{ r.name, m, lo }); + } + std.debug.print(" | sum of ops {d:.1} {d:.1}\n", .{ sum, low }); + } + check.pass("ling-layer-speed layer {d}: timed", .{k.L}); +} diff --git a/zig/tests/cuda/ling_mla_real.zig b/zig/tests/cuda/ling_mla_real.zig new file mode 100644 index 0000000000..5ce184cd44 --- /dev/null +++ b/zig/tests/cuda/ling_mla_real.zig @@ -0,0 +1,145 @@ +//! Ling layer 7's MLA on real inputs (an HF fixture): absorb, rope, cache, attention, expand, gate vs HF and float64. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const latent = glm.latent; +const Set = glm.glue.Set; + +const lw = 512; +const nope = 128; +const dv = 128; + +/// The fixture's dims and device buffers: q [t][h][192], lat [t][576] (raw k_rope), wk/wv [h][128][512], g fp32 [t][h]. +const Layer = struct { t: usize, h: usize, qk: usize, kw: usize, q: u64, lat: u64, wk: u64, wv: u64, g: u64, freq: u64 }; + +/// Device buffers one pass writes: the roped cache, qa, latent output, expanded and gated rows. +const Work = struct { cache: u64, qa: u64, ol: u64, vn: u64 }; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +/// Rows lo .. hi through absorb, rope, cache write, attention and expand at positions lo ..; out at vn row lo. +fn pass(r: *Rig, set: Set, l: Layer, w: Work, lo: usize, hi: usize) !void { + const n = hi - lo; + const p0: i32 = @intCast(lo); + const pos = try r.dev(std.mem.asBytes(&p0)); + const qa = w.qa + lo * l.h * l.kw * 2; + try latent.absorbQStrided(set, r.s, l.q + lo * l.h * l.qk * 2, l.qk, l.wk, qa, l.kw, n, l.h, nope, lw); + try latent.rope(set, r.s, .{ .x = l.q + (lo * l.h * l.qk + nope) * 2, .x_row = l.h * l.qk, .x_head = l.qk, .out = qa + lw * 2, .o_row = l.h * l.kw, .o_head = l.kw }, n, l.h, pos, l.freq); + const k = try r.zeros(n * l.kw * 2); + const zero: i32 = 0; + try latent.latentWrite(set, r.s, l.lat + lo * l.kw * 2, l.kw, k, try r.dev(std.mem.asBytes(&zero)), n, l.kw); + try latent.rope(set, r.s, .{ .x = k + lw * 2, .x_row = l.kw, .x_head = 0, .out = k + lw * 2, .o_row = l.kw, .o_head = 0 }, n, 1, pos, l.freq); + try latent.latentWrite(set, r.s, k, l.kw, w.cache, pos, n, l.kw); + const nch = latent.chunksFor(hi); + const p: latent.Partials = .{ .po = try r.zeros(nch * n * l.h * lw * 4), .pm = try r.zeros(nch * n * l.h * 4), .pl = try r.zeros(nch * n * l.h * 4) }; + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, @floatFromInt(l.qk)))); + try latent.attentionRoped(set, r.s, qa, w.cache, pos, p, w.ol + lo * l.h * lw * 2, try r.zeros(4), n, l.h, lw, nch, scale); + try latent.expandV(set, r.s, w.ol + lo * l.h * lw * 2, l.wv, w.vn + lo * l.h * dv * 2, n, l.h, dv, lw); +} + +fn work(r: *Rig, l: Layer) !Work { + return .{ .cache = try r.zeros(l.t * l.kw * 2), .qa = try r.zeros(l.t * l.h * l.kw * 2), .ol = try r.zeros(l.t * l.h * lw * 2), .vn = try r.zeros(l.t * l.h * dv * 2) }; +} + +/// Relative L2 of bf16 `got` against fp32 `want`, and the largest difference over want's rms. +fn distance(got: []const u8, want: []const f32) [2]f64 { + var dd: f64 = 0; + var ww: f64 = 0; + var most: f64 = 0; + for (want, 0..) |wv, i| { + const g = bf(std.mem.bytesToValue(u16, got[2 * i ..][0..2])); + dd += (g - wv) * (g - wv); + ww += @as(f64, wv) * wv; + most = @max(most, @abs(g - wv)); + } + return .{ @sqrt(dd / ww), most / @sqrt(ww / @as(f64, @floatFromInt(want.len))) }; +} + +fn report(name: []const u8, got: []const u8, want: []const f32, bound: f64) !void { + const dist = distance(got, want); + std.debug.print("RESULT {s}: rel L2 {e:.3}, max |diff| {e:.3} of rms\n", .{ name, dist[0], dist[1] }); + try check.expect(dist[0] <= bound, "{s}: rel L2 within {e}", .{ name, bound }); +} + +/// `ling-mla-real `: 512-token prompt vs HF and float64; decode rows, a window and a cut equal its bytes. +pub fn run(gpu: Gpu, path: []const u8) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(7); + var r: Rig = .{ .gpu = gpu, .a = a, .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer mod.unload(); + const ks = try glm.glue.Kernels.load(mod); + const set: Set = .{ .glue = &ks }; + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(256 << 20)); + var at: usize = 0; + const dims = try take(i32, a, bytes, &at, 4); + const t: usize = @intCast(dims[0]); + const h: usize = @intCast(dims[1]); + const qk: usize = @intCast(dims[2]); + const kw: usize = @intCast(dims[3]); + const q = try take(u16, a, bytes, &at, t * h * qk); + const lat = try take(u16, a, bytes, &at, t * kw); + const wk = try take(u16, a, bytes, &at, h * nope * lw); + const wv = try take(u16, a, bytes, &at, h * dv * lw); + const g = try take(f32, a, bytes, &at, t * h); + const freq = try take(f32, a, bytes, &at, latent.rope_width / 2); + var refs: [4][]f32 = undefined; + for (&refs) |*x| x.* = try take(f32, a, bytes, &at, t * h * dv); + try check.expect(at == bytes.len, "fixture: {d} bytes read of {d}", .{ at, bytes.len }); + const ours = latent.ropeFreqs(6000000); + var fd: f64 = 0; + for (ours, freq) |o, f| fd = @max(fd, @abs(@as(f64, o) - f) / f); + std.debug.print("RESULT rope inv_freq: ropeFreqs vs HF's fp32 table, worst relative {e:.3}\n", .{fd}); + const l: Layer = .{ + .t = t, + .h = h, + .qk = qk, + .kw = kw, + .q = try r.dev(std.mem.sliceAsBytes(q)), + .lat = try r.dev(std.mem.sliceAsBytes(lat)), + .wk = try r.dev(std.mem.sliceAsBytes(wk)), + .wv = try r.dev(std.mem.sliceAsBytes(wv)), + .g = try r.dev(std.mem.sliceAsBytes(g)), + .freq = try r.dev(std.mem.sliceAsBytes(freq)), + }; + const full = try work(&r, l); + try pass(&r, set, l, full, 0, t); + const pre = try r.back(full.vn, t * h * dv * 2); + try latent.headGate(set, r.s, full.vn, l.g, h, t, h, dv); + const post = try r.back(full.vn, t * h * dv * 2); + try report("Ling L7 attention pre-gate vs HF bf16 (eager)", pre, refs[0], 2e-2); + try report("Ling L7 attention post-gate (dense input) vs HF bf16 (eager)", post, refs[1], 2e-2); + try report("Ling L7 attention pre-gate vs float64 on the kernel inputs", pre, refs[2], 1e-2); + try report("Ling L7 attention post-gate vs float64 on the kernel inputs", post, refs[3], 1e-2); + const row = h * dv * 2; + const dec = try work(&r, l); + try pass(&r, set, l, dec, 0, 400); + for (400..432) |i| try pass(&r, set, l, dec, i, i + 1); + try pass(&r, set, l, dec, 432, 440); + try latent.headGate(set, r.s, dec.vn, l.g, h, 440, h, dv); + try check.sameBytes("Ling L7: 400-row prompt, 32 decode rows, an 8-row window equal the 512-row prompt", try r.back(dec.vn, 440 * row), post[0 .. 440 * row]); + const cut = try work(&r, l); + try pass(&r, set, l, cut, 0, 300); + try pass(&r, set, l, cut, 300, t); + try latent.headGate(set, r.s, cut.vn, l.g, h, t, h, dv); + try check.sameBytes("Ling L7: the prompt cut at 300 equals uncut", try r.back(cut.vn, t * row), post); + check.pass("ling-mla-real: layer 7 within rel L2 of HF and float64; decode, verify window and cut prompt equal the prompt", .{}); +} diff --git a/zig/tests/cuda/ling_mtp_real.zig b/zig/tests/cuda/ling_mtp_real.zig new file mode 100644 index 0000000000..92dcfc9352 --- /dev/null +++ b/zig/tests/cuda/ling_mtp_real.zig @@ -0,0 +1,150 @@ +//! `ling-mtp-real DIR FIXTURE`: the real MTP layer (model.layers.72) vs HF's BailingMoeV3MTPLayer; exactness on it. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const ling = @import("ling_cuda"); +const rig = @import("ling_forward_rig.zig"); + +const F = rig.F; + +fn take(comptime T: type, a: std.mem.Allocator, bytes: []const u8, at: *usize, n: usize) ![]T { + const out = try a.alloc(T, n); + @memcpy(std.mem.sliceAsBytes(out), bytes[at.*..][0 .. n * @sizeOf(T)]); + at.* += n * @sizeOf(T); + return out; +} + +const Kit = struct { + r: *rig.Rig, + h: *const ling.mtp.Head, + emb: []const u16, + hid: []const u16, + hd_dev: u64, // the hidden rows on the device + D: usize, + + /// Rows lo .. hi through the MTP layer on `hd` (prompt or decode buffers), committed; its final-norm rows. + fn pass(k: Kit, b: *const F.Buffers, hd: *F.State, lo: usize, hi: usize, picks: ?u64) ![]u16 { + const d = k.r.gpu.d; + const n = hi - lo; + try d.check(d.api.cuMemcpyHtoDAsync_v2(b.x, k.emb[lo * k.D ..].ptr, n * k.D * 2, k.r.s.handle), "cuMemcpyHtoDAsync"); + const f = k.h.forward(k.r.m, k.r.s, b); + try k.h.step(f, hd, n, k.hd_dev + lo * k.D * 2, picks); + try f.final(n); + const out = try k.r.a.alloc(u16, n * k.D); + @memcpy(std.mem.sliceAsBytes(out), try k.r.back(b.normed, n * k.D * 2)); + return out; + } +}; + +pub fn run(gpu: check.Gpu, dir: []const u8, path: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + const bytes = try std.Io.Dir.cwd().readFileAlloc(gpu.io, path, a, .limited(1 << 30)); + if (bytes.len < 20 or !std.mem.eql(u8, bytes[0..8], "TFLINGM1")) return error.BadFixture; + var at: usize = 8; + const dims = try take(u32, a, bytes, &at, 3); + const T: usize = dims[0]; + const D: usize = dims[1]; + const K: usize = dims[2]; + const emb = try take(u16, a, bytes, &at, T * D); + const hid = try take(u16, a, bytes, &at, T * D); + const want = try take(f32, a, bytes, &at, T * D); + const topk = try take(i32, a, bytes, &at, T * K); + if (at != bytes.len) return error.BadFixture; + var why: ling.config.Why = .{}; + var o = rig.W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + var m = try F.Modules.load(gpu.d, gpu.ctx); + defer m.deinit(); + var pack = try cuda.modelopt.Packer.load(gpu.d); + defer pack.unload(); + const L = o.c.layers; + const t0 = check.now(gpu.io); + var w = try rig.W.load(gpu.gpa, gpu.d, s, &pack, &o.ck, o.c, &.{L}, .{ .globals = false }); + defer w.deinit(); + const load_s = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + const cap = T + 32; + var h: ling.mtp.Head = undefined; + try h.init(gpu.gpa, gpu.d, &m, s, &w, T, cap); + defer h.deinit(); + const c = &w.config; + try check.expect(D == c.hidden and K == c.topk, "fixture: {d} rows {d} wide, top {d}", .{ T, D, K }); + var r: rig.Rig = .{ .gpu = gpu, .a = a, .s = s, .m = &m, .g = &h.g }; + defer r.deinit(); + var pb = try F.Buffers.init(gpu.d, &m, &w, T, true, cap); + defer pb.deinit(); + var db = try F.Buffers.init(gpu.d, &m, &w, F.window_rows, false, cap); + defer db.deinit(); + const hd_dev = try r.zeros(T * D * 2); + try gpu.d.check(gpu.d.api.cuMemcpyHtoDAsync_v2(hd_dev, hid.ptr, T * D * 2, s.handle), "cuMemcpyHtoDAsync"); + const k: Kit = .{ .r = &r, .h = &h, .emb = emb, .hid = hid, .hd_dev = hd_dev, .D = D }; + // the whole prompt against HF, picks against HF's + const pk = try r.zeros(T * (K + 1) * 4); + const st = try r.state(cap); + defer st.deinit(); + const got = try k.pass(&pb, st, 0, T, pk); + const ours = try rig.sortedPicks(a, try r.back(pk, T * (K + 1) * 4), 1, T, K); + const theirs = try a.dupe(i32, topk); + for (0..T) |t| std.mem.sort(i32, theirs[t * K ..][0..K], {}, std.sort.asc(i32)); + const skip = try a.alloc(bool, T); + @memset(skip, false); + const nf = rig.flips(ours, theirs, 1, T, 0, T, K, skip); + const dv = dist(got, want, T, D, skip); + std.debug.print("RESULT ling-mtp-real layer {d} ({t}): {d} rows, loaded in {d:.1} s\n", .{ L, c.quant, T, load_s }); + std.debug.print("RESULT final-norm rows vs HF: rel L2 {e:.3}, {e:.3} without the {d} rows whose top-{d} differs, worst kept row {e:.3}\n", .{ dv[0], dv[1], nf, K, dv[2] }); + // cut prompt, decode rows, a verify window + const sc = try r.state(cap); + defer sc.deinit(); + const cut0 = try k.pass(&pb, sc, 0, 300, null); + const cut1 = try k.pass(&pb, sc, 300, T, null); + try check.sameBytes("MTP rows cut at 300 == uncut", std.mem.sliceAsBytes(cut0), std.mem.sliceAsBytes(got[0 .. 300 * D])); + try check.sameBytes("MTP rows cut at 300 == uncut (the rest)", std.mem.sliceAsBytes(cut1), std.mem.sliceAsBytes(got[300 * D ..])); + const P = 400; + const N = F.window_rows; + const sd = try r.state(cap); + defer sd.deinit(); + const sw = try r.state(cap); + defer sw.deinit(); + _ = try k.pass(&pb, sd, 0, P, null); + _ = try k.pass(&pb, sw, 0, P, null); + const dec = try a.alloc(u16, N * D); + for (0..N) |i| @memcpy(dec[i * D ..][0..D], try k.pass(&db, sd, P + i, P + i + 1, null)); + try check.sameBytes("16-row MTP window == 16 decode rows", std.mem.sliceAsBytes(try k.pass(&db, sw, P, P + N, null)), std.mem.sliceAsBytes(dec)); + const dd = dist(dec, want[P * D ..], N, D, skip[P..]); + std.debug.print("RESULT decode rows {d}-{d}: rel L2 {e:.3} vs HF\n", .{ P, P + N - 1, dd[0] }); + try check.expect(dv[1] <= 1e-2 and dd[1] <= 1e-2, "the MTP layer within rel L2 1e-2 of HF (rows whose picks agree)", .{}); + check.pass("ling-mtp-real: rel L2 {e:.3} of HF ({d} flipped rows apart); cut == uncut, window == decode rows; peak {d} MiB on the device", .{ dv[1], nf, cuda.usage(false).peak >> 20 }); +} + +/// Relative L2 of bf16 rows against fp32 rows, all and without `skip` rows, and the worst kept row's. +fn dist(got: []const u16, want: []const f32, rows: usize, D: usize, skip: []const bool) [3]f64 { + var dd: f64 = 0; + var ww: f64 = 0; + var kd: f64 = 0; + var kw: f64 = 0; + var worst: f64 = 0; + for (0..rows) |t| { + var rd: f64 = 0; + var rw: f64 = 0; + for (0..D) |j| { + const e: f64 = want[t * D + j]; + const g = rig.bfOf(got[t * D + j]); + rd += (g - e) * (g - e); + rw += e * e; + } + dd += rd; + ww += rw; + if (!skip[t]) { + kd += rd; + kw += rw; + worst = @max(worst, @sqrt(rd / rw)); + } + } + return .{ @sqrt(dd / ww), @sqrt(kd / kw), worst }; +} diff --git a/zig/tests/cuda/ling_pack.zig b/zig/tests/cuda/ling_pack.zig new file mode 100644 index 0000000000..eb4747c2ef --- /dev/null +++ b/zig/tests/cuda/ling_pack.zig @@ -0,0 +1,160 @@ +//! `ling-pack DIR LAYER [EXPERTS [STAGE]]`: MoE experts loaded and timed, bytes vs the host packers the loader used. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const lw = @import("ling_weights_tests.zig"); + +const W = ling.weights; +const X8 = cuda.experts.Fp8Experts; +const X4 = cuda.experts.Nvfp4Experts; + +/// Host memory from /proc/self/status: the peak resident set since the last reset, and anonymous memory now (KiB). +const Host = struct { + hwm: usize = 0, + anon: usize = 0, + + fn read(io: std.Io) Host { + var buf: [8192]u8 = undefined; + const text = std.Io.Dir.cwd().readFile(io, "/proc/self/status", &buf) catch return .{}; + var h: Host = .{}; + var lines = std.mem.tokenizeScalar(u8, text, '\n'); + while (lines.next()) |line| { + inline for (.{ .{ "VmHWM:", "hwm" }, .{ "RssAnon:", "anon" } }) |kv| if (std.mem.startsWith(u8, line, kv[0])) { + var it = std.mem.tokenizeAny(u8, line[kv[0].len..], " \tkB"); + @field(h, kv[1]) = std.fmt.parseInt(usize, it.next() orelse "0", 10) catch 0; + }; + } + return h; + } + + /// Resets VmHWM to the current resident set (proc(5) clear_refs 5). + fn reset(io: std.Io) void { + std.Io.Dir.cwd().writeFile(io, .{ .sub_path = "/proc/self/clear_refs", .data = "5" }) catch {}; + } +}; + +fn secs(io: std.Io, t0: i96) f64 { + return @as(f64, @floatFromInt(check.now(io) - t0)) / 1e9; +} + +/// Every byte of layer i's routed experts 0..count read once from the mapping (the read alone, no device work). +fn touch(ck: *core.Checkpoint, i: u32, count: usize) !usize { + var sum: u64 = 0; + var bytes: usize = 0; + var b: [64]u8 = undefined; + const prefix = try std.fmt.bufPrint(&b, "model.layers.{d}.mlp.experts.", .{i}); + for (ck.files.items) |*f| for (f.names.keys()) |k| { + if (!std.mem.startsWith(u8, k, prefix)) continue; + const rest = k[prefix.len..]; + const e = std.fmt.parseInt(usize, rest[0 .. std.mem.indexOfScalar(u8, rest, '.') orelse continue], 10) catch continue; + if (e >= count) continue; + const t = f.get(k).?; + for (std.mem.bytesAsSlice(u64, @as([]align(1) const u8, t.bytes[0 .. t.bytes.len / 8 * 8]))) |w| sum +%= w; + bytes += t.bytes.len; + }; + std.mem.doNotOptimizeAway(sum); + return bytes; +} + +/// One expert's packed device words against the host packers' bytes from the same file tensors. +fn compare(gpu: check.Gpu, ck: *core.Checkpoint, i: u32, l: cuda.experts.Layer, a: std.mem.Allocator) !void { + const ni = l.width; + const d = l.dims; + const nvfp4 = l.format == .nvfp4; + const up_n = if (nvfp4) X4.bytes(1, ni, d, 2) / 4 else 2 * ni * d / 4; + const down_n = if (nvfp4) X4.bytes(1, d, ni, 1) / 4 else d * ni / 4; + const inv_n = (ni / 128) * (d / 128); + const hu = try a.alloc(u32, up_n); + const hd = try a.alloc(u32, down_n); + const gu = try a.alloc(u32, up_n); + const gd = try a.alloc(u32, down_n); + const scale_up = try a.alloc(u8, l.experts * (if (nvfp4) 2 * 4 else 2 * inv_n * 4)); + const scale_down = try a.alloc(u8, l.experts * (if (nvfp4) 4 else inv_n * 4)); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.up_scale, .len = scale_up.len }).download(0, scale_up); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.down_scale, .len = scale_down.len }).download(0, scale_down); + var b: [3][192]u8 = undefined; + for (0..l.experts) |e| { + var m: [3][]const u8 = undefined; + for ([_][]const u8{ "gate_proj", "up_proj", "down_proj" }, 0..) |p, j| m[j] = try std.fmt.bufPrint(&b[j], "model.layers.{d}.mlp.experts.{d}.{s}", .{ i, e, p }); + if (nvfp4) { + var x: [3]core.nvfp4.Linear = undefined; + for (&x, m) |*y, name| y.* = try core.nvfp4.Linear.read(ck, name); + X4.packGateUp(hu, .{ x[0].codes, x[0].scales }, .{ x[1].codes, x[1].scales }, ni, d); + X4.packOne(hd, x[2].codes, x[2].scales, d, ni); + const globals = [3]f32{ x[0].global, x[1].global, x[2].global }; + try check.sameBytes("nvfp4 gate and up globals", scale_up[8 * e ..][0..8], std.mem.sliceAsBytes(globals[0..2])); + try check.sameBytes("nvfp4 down global", scale_down[4 * e ..][0..4], std.mem.sliceAsBytes(globals[2..3])); + } else { + var t: [3]core.checkpoint.Tensor = undefined; + var s: [3]core.checkpoint.Tensor = undefined; + for (&t, &s, m) |*w, *sc, name| { + var nb: [200]u8 = undefined; + w.* = try ck.get(try std.fmt.bufPrint(&nb, "{s}.weight", .{name})); + sc.* = try ck.get(try std.fmt.bufPrint(&nb, "{s}.weight_scale_inv", .{name})); + } + try X8.packGateUp(a, hu, t[0].bytes, t[1].bytes, ni, d); + X8.packOne(hd, t[2].bytes, d, ni); + try check.sameBytes("fp8 gate scale_inv", scale_up[(2 * e) * inv_n * 4 ..][0 .. inv_n * 4], s[0].bytes); + try check.sameBytes("fp8 up scale_inv", scale_up[(2 * e + 1) * inv_n * 4 ..][0 .. inv_n * 4], s[1].bytes); + try check.sameBytes("fp8 down scale_inv", scale_down[e * inv_n * 4 ..][0 .. inv_n * 4], s[2].bytes); + } + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.up, .len = l.experts * up_n * 4 }).download(e * up_n * 4, std.mem.sliceAsBytes(gu)); + try (cuda.DeviceBuffer{ .d = gpu.d, .ptr = l.down, .len = l.experts * down_n * 4 }).download(e * down_n * 4, std.mem.sliceAsBytes(gd)); + var nb: [64]u8 = undefined; + try check.sameBytes(try std.fmt.bufPrint(&nb, "expert {d} gate-up words", .{e}), std.mem.sliceAsBytes(gu), std.mem.sliceAsBytes(hu)); + try check.sameBytes(try std.fmt.bufPrint(&nb, "expert {d} down words", .{e}), std.mem.sliceAsBytes(gd), std.mem.sliceAsBytes(hd)); + } +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len < 2) return error.MissingArgument; + const dir = rest[0]; + const i = try std.fmt.parseInt(u32, rest[1], 10); + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const io = gpu.io; + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + if (!o.c.isMoe(i)) return error.NotMoeLayer; + const count = if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else o.c.experts; + const opts: W.Options = .{ .globals = false, .experts = count, .stage = if (rest.len > 3) try std.fmt.parseInt(usize, rest[3], 10) else 512 << 20 }; + // the read alone from a cold cache, then a load on hot pages (pack and upload), then a cold load (both) + o.ck.release(); + var t0 = check.now(io); + const bytes = try touch(&o.ck, i, count); + const read_s = secs(io, t0); + var hot_s: f64 = 0; + var cold_s: f64 = 0; + var peak: [2]usize = .{ 0, 0 }; + var host: [2]Host = undefined; + for (0..2) |pass| { + if (pass == 1) o.ck.release(); + Host.reset(io); + const base = cuda.usage(true).device; + t0 = check.now(io); + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{i}, opts); + defer w.deinit(); + if (pass == 0) hot_s = secs(io, t0) else cold_s = secs(io, t0); + peak[pass] = cuda.usage(false).peak - base; + host[pass] = Host.read(io); + if (pass == 1) { + const l = w.layer(i).?.ffn.moe.routed.grouped; + try compare(gpu, &o.ck, i, l, arena.allocator()); + std.debug.print("RESULT layer {d}: {d} experts {t}, {d:.2} GiB of expert tensors\n", .{ i, count, l.format, @as(f64, @floatFromInt(bytes)) / (1 << 30) }); + } + } + std.debug.print("RESULT read alone {d:.2} s ({d:.0} MiB/s); load on hot pages {d:.2} s; cold load {d:.2} s\n", .{ read_s, @as(f64, @floatFromInt(bytes >> 20)) / read_s, hot_s, cold_s }); + std.debug.print("RESULT device peak above the start {d} MiB (hot), {d} MiB (cold); host VmHWM {d} MiB (hot), {d} MiB (cold), RssAnon {d} MiB\n", .{ peak[0] >> 20, peak[1] >> 20, host[0].hwm >> 10, host[1].hwm >> 10, host[1].anon >> 10 }); + check.pass("ling-pack layer {d}: every expert's packed words and scales equal the host packers' bytes", .{i}); +} diff --git a/zig/tests/cuda/ling_ref_mla.zig b/zig/tests/cuda/ling_ref_mla.zig new file mode 100644 index 0000000000..fd855256be --- /dev/null +++ b/zig/tests/cuda/ling_ref_mla.zig @@ -0,0 +1,267 @@ +//! Ling's MLA on GLM's kernels vs float64: 128-wide absorb and expand, rope, head gate, the 576-wide roped attention. +const std = @import("std"); +const cuda = @import("cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const glm = @import("glm_cuda"); +const Gpu = check.Gpu; +const Rig = qr.RigT; +const bf = qr.helpers.bfOf; +const latent = glm.latent; + +const heads = 64; +const lw = 512; +const rw = latent.rope_width; +const kw = lw + rw; +const nope = 128; +const qk = nope + rw; +const dv = 128; +const u = std.math.pow(f64, 2, -24); + +/// Worst |got - want| / tol: `want` is `width` of every `stride` bf16 elements of `raw` from `off`. +fn within(r: *Rig, name: []const u8, raw: []const u8, stride: usize, off: usize, width: usize, want: []const f64, tol: []const f64) !void { + var most: f64 = 0; + for (want, tol, 0..) |w, t, i| { + const at = (i / width) * stride + off + i % width; + const q = @abs(bf(std.mem.bytesToValue(u16, raw[2 * at ..][0..2])) - w) / t; + most = if (std.math.isNan(q)) std.math.inf(f64) else @max(most, q); + } + r.worst = @max(r.worst, most); + std.debug.print("RESULT {s}: worst {e:.3} of its float64 tolerance\n", .{ name, most }); + try check.expect(most <= 1, "{s} within its float64 tolerance", .{name}); +} + +fn f64s(r: *Rig, x: []const u16) ![]f64 { + const out = try r.a.alloc(f64, x.len); + for (out, x) |*o, v| o.* = bf(v); + return out; +} + +/// The bytes of rows lo .. hi of a [rows][n] bf16 output equal `got`'s. +fn sameSpan(name: []const u8, all: []const u8, got: []const u8, lo: usize, n: usize) !void { + try check.sameBytes(name, got, all[lo * n * 2 ..][0..got.len]); +} + +/// absorbQStrided (192-wide q heads into 576-wide qa rows) and expandV at dv 128 vs float64; row counts' bits. +fn projections(r: *Rig, set: glm.glue.Set) !void { + const rows = 70; + const q = try r.bfs(rows * heads * qk, 1.0); + const wk = try r.bfs(heads * nope * lw, 0.05); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const dwk = try r.dev(std.mem.sliceAsBytes(wk)); + const qa = try r.zeros(rows * heads * kw * 2); + try latent.absorbQStrided(set, r.s, dq, qk, dwk, qa, kw, rows, heads, nope, lw); + var want = try r.a.alloc(f64, rows * heads * lw); + var tol = try r.a.alloc(f64, rows * heads * lw); + const qd = try f64s(r, q); + const wkd = try f64s(r, wk); + for (0..rows) |row| for (0..heads) |h| for (0..lw) |c| { + var s: f64 = 0; + var abs: f64 = 0; + for (0..nope) |k| { + const p = qd[(row * heads + h) * qk + k] * wkd[(h * nope + k) * lw + c]; + s += p; + abs += @abs(p); + } + want[(row * heads + h) * lw + c] = s; + tol[(row * heads + h) * lw + c] = @abs(s) * std.math.pow(f64, 2, -8) + abs * (nope + 8) * u + 1e-30; + }; + const all = try r.back(qa, rows * heads * kw * 2); + try within(r, "absorb_q 128 (Ling): 70 rows x 64 heads, q stride 192, qa stride 576", all, kw, 0, lw, want, tol); + for ([_][2]usize{ .{ 0, 1 }, .{ 37, 42 }, .{ 3, 67 } }) |span| { + const n = span[1] - span[0]; + const o = try r.zeros(n * heads * kw * 2); + try latent.absorbQStrided(set, r.s, dq + span[0] * heads * qk * 2, qk, dwk, o, kw, n, heads, nope, lw); + const got = try r.back(o, n * heads * kw * 2); + for (0..n * heads) |i| try check.sameBytes("absorb_q 128: rows of a narrower call equal the 70-row call's", got[i * kw * 2 ..][0 .. lw * 2], all[(span[0] * heads + i) * kw * 2 ..][0 .. lw * 2]); + } + const ol = try r.bfs(rows * heads * lw, 1.0); + const wv = try r.bfs(heads * dv * lw, 0.05); + const dol = try r.dev(std.mem.sliceAsBytes(ol)); + const dwv = try r.dev(std.mem.sliceAsBytes(wv)); + const vn = try r.zeros(rows * heads * dv * 2); + try latent.expandV(set, r.s, dol, dwv, vn, rows, heads, dv, lw); + want = try r.a.alloc(f64, rows * heads * dv); + tol = try r.a.alloc(f64, rows * heads * dv); + const old = try f64s(r, ol); + const wvd = try f64s(r, wv); + for (0..rows) |row| for (0..heads) |h| for (0..dv) |n| { + var s: f64 = 0; + var abs: f64 = 0; + for (old[(row * heads + h) * lw ..][0..lw], wvd[(h * dv + n) * lw ..][0..lw]) |a, b| { + s += a * b; + abs += @abs(a * b); + } + want[(row * heads + h) * dv + n] = s; + tol[(row * heads + h) * dv + n] = @abs(s) * std.math.pow(f64, 2, -8) + abs * (lw + 8) * u + 1e-30; + }; + const vall = try r.back(vn, rows * heads * dv * 2); + try within(r, "expand_v 128 (Ling): 70 rows x 64 heads", vall, dv, 0, dv, want, tol); + for ([_][2]usize{ .{ 0, 1 }, .{ 37, 42 }, .{ 3, 67 } }) |span| { + const n = span[1] - span[0]; + const o = try r.zeros(n * heads * dv * 2); + try latent.expandV(set, r.s, dol + span[0] * heads * lw * 2, dwv, o, n, heads, dv, lw); + try sameSpan("expand_v 128: rows of a narrower call equal the 70-row call's", vall, try r.back(o, n * heads * dv * 2), span[0], heads * dv); + } + check.pass("absorb_q / expand_v at 128-wide heads: within float64 tolerance, every row count's bits equal", .{}); +} + +/// rope (q heads into qa, cache rows in place) at two positions vs float64 on fp32 angles; head gate vs float64. +fn ropeGate(r: *Rig, set: glm.glue.Set) !void { + const rows = 40; + const freqs = latent.ropeFreqs(6000000); + const dfreq = try r.dev(std.mem.asBytes(&freqs)); + const q = try r.bfs(rows * heads * qk, 1.0); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const qd = try f64s(r, q); + var name: [128]u8 = undefined; + for ([_]i32{ 1000, 261000 }) |p0| { + const at = try r.dev(std.mem.asBytes(&p0)); + const qa = try r.zeros(rows * heads * kw * 2); + const v: latent.RopeRows = .{ .x = dq + nope * 2, .x_row = heads * qk, .x_head = qk, .out = qa + lw * 2, .o_row = heads * kw, .o_head = kw }; + try latent.rope(set, r.s, v, rows, heads, at, dfreq); + const want = try r.a.alloc(f64, rows * heads * rw); + const tol = try r.a.alloc(f64, rows * heads * rw); + for (0..rows) |row| for (0..heads) |h| for (0..rw / 2) |i| { + const ang: f64 = freqs[i] * @as(f32, @floatFromInt(p0 + @as(i32, @intCast(row)))); + const x0 = qd[(row * heads + h) * qk + nope + 2 * i]; + const x1 = qd[(row * heads + h) * qk + nope + 2 * i + 1]; + const c = @cos(ang); + const s = @sin(ang); + const at2 = (row * heads + h) * rw + 2 * i; + want[at2] = x0 * c - x1 * s; + want[at2 + 1] = x1 * c + x0 * s; + for (0..2) |e| tol[at2 + e] = @abs(want[at2 + e]) * std.math.pow(f64, 2, -8) + (@abs(x0) + @abs(x1)) * 8 * u + 1e-30; + }; + const all = try r.back(qa, rows * heads * kw * 2); + try within(r, try std.fmt.bufPrint(&name, "rope: 40 rows x 64 heads at {d} (GPT-J pairs, fp32 angles)", .{p0}), all, kw, lw, rw, want, tol); + const p1: i32 = p0 + 23; + const one = try r.zeros(heads * kw * 2); + const v1: latent.RopeRows = .{ .x = dq + (23 * heads * qk + nope) * 2, .x_row = heads * qk, .x_head = qk, .out = one + lw * 2, .o_row = heads * kw, .o_head = kw }; + try latent.rope(set, r.s, v1, 1, heads, try r.dev(std.mem.asBytes(&p1)), dfreq); + const got = try r.back(one, heads * kw * 2); + for (0..heads) |h| try check.sameBytes("rope: row 23 alone equals its row", got[(h * kw + lw) * 2 ..][0 .. rw * 2], all[((23 * heads + h) * kw + lw) * 2 ..][0 .. rw * 2]); + // a cache row's k_rope in place equals q's rope of the same values at the same position + const k = try r.dev(all[(23 * heads * kw) * 2 ..][0 .. kw * 2]); + @memcpy(all[(23 * heads * kw + lw) * 2 ..][0 .. rw * 2], std.mem.sliceAsBytes(q[(23 * heads) * qk + nope ..][0..rw])); + const kin = try r.dev(all[(23 * heads * kw) * 2 ..][0 .. kw * 2]); + try latent.rope(set, r.s, .{ .x = kin + lw * 2, .x_row = kw, .x_head = 0, .out = kin + lw * 2, .o_row = kw, .o_head = 0 }, 1, 1, try r.dev(std.mem.asBytes(&p1)), dfreq); + try check.sameBytes("rope: a cache row in place equals the q rope", try r.back(kin, kw * 2), try r.back(k, kw * 2)); + } + const o = try r.bfs(rows * heads * dv, 1.0); + const g = try r.a.alloc(f32, rows * (heads + 3)); + for (g) |*x| x.* = (r.rng.float(f32) * 2 - 1) * 12; + const dov = try r.dev(std.mem.sliceAsBytes(o)); + try latent.headGate(set, r.s, dov, try r.dev(std.mem.sliceAsBytes(g)), heads + 3, rows, heads, dv); + const want = try r.a.alloc(f64, rows * heads * dv); + const tol = try r.a.alloc(f64, rows * heads * dv); + for (want, tol, 0..) |*w, *t, i| { + const sg = 1 / (1 + @exp(-@as(f64, g[(i / (heads * dv)) * (heads + 3) + (i / dv) % heads]))); + w.* = bf(o[i]) * sg; + t.* = @abs(w.*) * std.math.pow(f64, 2, -8) + @abs(bf(o[i])) * 8 * u + 1e-30; + } + try within(r, "head gate: 40 rows x 64 heads x 128, g stride 67", try r.back(dov, rows * heads * dv * 2), dv, 0, dv, want, tol); + check.pass("rope and head gate: within float64 tolerance, a row alone and a cache row in place equal", .{}); +} + +/// float64 attention over 576-wide rows (scores) and their first 512 (values); bound as glm_ref_attn's, 576 + 8 terms. +fn reference(r: *Rig, q: []const f64, cd: []const f64, pos: usize, rows: usize, scale: f64) !struct { want: []f64, tol: []f64 } { + const want = try r.a.alloc(f64, rows * heads * lw); + const tol = try r.a.alloc(f64, rows * heads * lw); + const s = try r.a.alloc(f64, pos + rows); + const ds = try r.a.alloc(f64, pos + rows); + const o = try r.a.alloc(f64, lw); + const spread = try r.a.alloc(f64, lw); + for (0..rows) |row| for (0..heads) |h| { + const qr_ = q[(row * heads + h) * kw ..][0..kw]; + const n = pos + row + 1; + var most = -std.math.inf(f64); + for (0..n) |key| { + var dot: f64 = 0; + var abs: f64 = 0; + for (qr_, cd[key * kw ..][0..kw]) |a, b| { + dot += a * b; + abs += @abs(a * b); + } + s[key] = dot * scale; + ds[key] = abs * scale * (kw + 8) * u; + most = @max(most, s[key]); + } + @memset(o, 0); + @memset(spread, 0); + var l: f64 = 0; + for (0..n) |key| { + const p = @exp(s[key] - most); + l += p; + for (o, spread, cd[key * kw ..][0..lw]) |*od, *sd, v| { + od.* += p * v; + sd.* += p * @abs(v) * (std.math.pow(f64, 2, -8) + ds[key]); + } + } + const base = (row * heads + h) * lw; + for (0..lw) |d| { + want[base + d] = o[d] / l; + tol[base + d] = @abs(o[d] / l) * std.math.pow(f64, 2, -8) + spread[d] / l + 1e-30; + } + }; + return .{ .want = want, .tol = tol }; +} + +/// attentionRoped of `rows` rows at `pos` (rows of qa from `first`) into a fresh [rows][heads][lw] output. +fn attend(r: *Rig, set: glm.glue.Set, qa: u64, cache: u64, pos: usize, first: usize, rows: usize, scale: f32) ![]u8 { + const nch = latent.chunksFor(pos + rows); + const out = try r.zeros(rows * heads * lw * 2); + const p0: i32 = @intCast(pos); + const p: latent.Partials = .{ .po = try r.zeros(nch * rows * heads * lw * 4), .pm = try r.zeros(nch * rows * heads * 4), .pl = try r.zeros(nch * rows * heads * 4) }; + try latent.attentionRoped(set, r.s, qa + first * heads * kw * 2, cache, try r.dev(std.mem.asBytes(&p0)), p, out, try r.zeros(4), rows, heads, lw, nch, scale); + return r.back(out, rows * heads * lw * 2); +} + +/// Roped attention: 60 rows at 1,500 (keys cross 1,536) vs float64; one-row decode, a verify window and a cut prompt. +fn attention(r: *Rig, set: glm.glue.Set) !void { + const pos = 1500; + const rows = 60; + const scale: f32 = @floatCast(1 / @sqrt(@as(f64, qk))); + const q = try r.bfs(rows * heads * kw, 1.0); + const c = try r.bfs((pos + rows) * kw, 1.0); + const dq = try r.dev(std.mem.sliceAsBytes(q)); + const dc = try r.dev(std.mem.sliceAsBytes(c)); + const all = try attend(r, set, dq, dc, pos, 0, rows, scale); + const ref64 = try reference(r, try f64s(r, q), try f64s(r, c), pos, rows, scale); + try within(r, "roped attention (Ling): 60 rows at 1500, 64 heads, 576-wide keys, 16-head blocks", all, lw, 0, lw, ref64.want, ref64.tol); + const n = heads * lw; + var name: [128]u8 = undefined; + for ([_]usize{ 0, 35, 36, 59 }) |row| { + const one = try attend(r, set, dq, dc, pos + row, row, 1, scale); + try sameSpan(try std.fmt.bufPrint(&name, "roped attention: row {d} alone (decode) equals its row", .{row}), all, one, row, n); + } + const win = try attend(r, set, dq, dc, pos + 30, 30, 8, scale); + try sameSpan("roped attention: an 8-row verify window across 1,536 equals serial rows", all, win, 30, n); + const head = try attend(r, set, dq, dc, pos, 0, 36, scale); + const tail = try attend(r, set, dq, dc, pos + 36, 36, rows - 36, scale); + try sameSpan("roped attention: a prompt cut at 36 (key 1,536): first part equals uncut", all, head, 0, n); + try sameSpan("roped attention: a prompt cut at 36 (key 1,536): second part equals uncut", all, tail, 36, n); + check.pass("roped attention over {d} keys: within float64 tolerance; decode rows, verify window, cut prompt equal", .{pos + rows}); +} + +/// `ling-mla`: the checks above on synthetic inputs (seeded). +pub fn run(gpu: Gpu) !void { + var arena: std.heap.ArenaAllocator = .init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(731); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*b| b.free(); + r.bufs.deinit(gpu.gpa); + } + var mod = try cuda.Module.load(gpu.d, cuda.kernels.glm_glue); + defer mod.unload(); + const ks = try glm.glue.Kernels.load(mod); + const set: glm.glue.Set = .{ .glue = &ks }; + try projections(&r, set); + try ropeGate(&r, set); + try attention(&r, set); + check.pass("ling-mla: every check within {e:.3} of its float64 tolerance", .{r.worst}); +} diff --git a/zig/tests/cuda/ling_synth.zig b/zig/tests/cuda/ling_synth.zig new file mode 100644 index 0000000000..7dde1b698d --- /dev/null +++ b/zig/tests/cuda/ling_synth.zig @@ -0,0 +1,151 @@ +//! `ling-synth DIR [--format bf16|fp8|nvfp4] [--layers N]`: a tiny Ling-3.1 checkpoint, every name from the weight map. +const std = @import("std"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); + +const st = core.safetensors; +const New = core.shard_edit.New; +const Config = ling.config.Config; +const Quant = ling.config.Quant; +const Role = ling.names.Role; + +/// A linear role's stored [n, k], from the config: the writer's own table, apart from the loader's. +fn linear(c: *const Config, role: Role) ?[2]usize { + const d: usize = c.hidden; + const n: usize = c.kda_heads * c.kda_dim; + const h: usize = c.mla_heads; + return switch (role) { + .kda_q, .kda_k, .kda_v, .kda_f, .kda_g => .{ n, d }, + .kda_b => .{ c.kda_heads, d }, + .kda_o => .{ d, n }, + .mla_q => .{ h * c.qkDim(), d }, + .mla_q_a => .{ c.q_lora, d }, + .mla_q_b => .{ h * c.qkDim(), c.q_lora }, + .mla_kv_a => .{ c.kv_lora + c.rope_dim, d }, + .mla_kv_b => .{ h * (c.nope + c.v_dim), c.kv_lora }, + .mla_gate => .{ h, d }, + .mla_o => .{ d, h * c.v_dim }, + .dense_gate, .dense_up => .{ c.dense_inter, d }, + .dense_down => .{ d, c.dense_inter }, + .router => .{ c.experts, d }, + .expert_gate, .expert_up => .{ c.moe_inter, d }, + .expert_down => .{ d, c.moe_inter }, + .shared_gate, .shared_up => .{ c.shared_inter, d }, + .shared_down => .{ d, c.shared_inter }, + .mtp_eh => .{ d, 2 * d }, + .lm_head, .embed => .{ c.vocab, d }, + else => null, + }; +} + +const Out = struct { + a: std.mem.Allocator, + rng: std.Random, + c: *const Config, + list: std.ArrayList(New) = .empty, + bytes: usize = 0, + + fn add(o: *Out, name: []const u8, dtype: st.DType, shape: []const usize, bytes: []const u8) !void { + try o.list.append(o.a, .{ .name = try o.a.dupe(u8, name), .dtype = dtype, .shape = try o.a.dupe(usize, shape), .bytes = bytes }); + o.bytes += bytes.len; + } + + fn bf16s(o: *Out, count: usize, lo: f32, hi: f32) ![]u8 { + const v = try o.a.alloc(u16, count); + for (v) |*x| x.* = ling.weights.proj.toBf16(lo + o.rng.float(f32) * (hi - lo)); + return std.mem.sliceAsBytes(v); + } + + fn f32s(o: *Out, count: usize, lo: f32, hi: f32) ![]u8 { + const v = try o.a.alloc(f32, count); + for (v) |*x| x.* = lo + o.rng.float(f32) * (hi - lo); + return std.mem.sliceAsBytes(v); + } + + /// One name the map emits: a weight in its module's format, a side tensor by its suffix, or a small tensor by role. + pub fn put(o: *Out, name: []const u8, t: ling.names.Tensor) !void { + const c = o.c; + if (linear(c, t.role)) |nk| { + const n, const k = nk; + const module = name[0 .. std.mem.lastIndexOfScalar(u8, name, '.') orelse return error.BadName]; + const f = if (t.role == .embed) .bf16 else c.format(module); + const s = 1.0 / @sqrt(@as(f32, @floatFromInt(k))); + if (!t.scale) return switch (f) { + .bf16 => o.add(name, .bf16, &.{ n, k }, try o.bf16s(n * k, -s, s)), + .fp8_block128 => o.add(name, .f8_e4m3, &.{ n, k }, try o.codes(n * k, .e4m3)), + .nvfp4 => o.add(name, .u8, &.{ n, k / 2 }, try o.codes(n * k / 2, .any)), + }; + const tail = name[module.len + 1 ..]; + if (std.mem.eql(u8, tail, "weight_scale_inv")) return o.add(name, .f32, &.{ (n + 127) / 128, k / 128 }, try o.f32s((n + 127) / 128 * (k / 128), s / 4, s / 2)); + if (std.mem.eql(u8, tail, "weight_scale")) return o.add(name, .f8_e4m3, &.{ n, k / 16 }, try o.codes(n * k / 16, .scale)); + if (std.mem.eql(u8, tail, "weight_scale_2")) return o.add(name, .f32, &.{}, try o.f32s(1, s / 8, s / 4)); + if (std.mem.eql(u8, tail, "input_scale")) return o.add(name, .f32, &.{}, try o.f32s(1, 1, 1)); + return error.BadName; + } + const d: usize = c.hidden; + const n: usize = c.kda_heads * c.kda_dim; + return switch (t.role) { + .attn_norm, .mlp_norm, .final_norm, .mtp_enorm, .mtp_hnorm, .mtp_final_norm => o.add(name, .bf16, &.{d}, try o.bf16s(d, 0.5, 1.5)), + .kda_q_conv, .kda_k_conv, .kda_v_conv => o.add(name, .bf16, &.{ n, 1, c.conv }, try o.bf16s(n * c.conv, -0.5, 0.5)), + .kda_a_log => o.add(name, .f32, &.{c.kda_heads}, try o.f32s(c.kda_heads, 0, 2.5)), + .kda_dt_bias => o.add(name, .f32, &.{n}, try o.f32s(n, -0.5, 0.5)), + .kda_o_norm => o.add(name, .bf16, &.{c.kda_dim}, try o.bf16s(c.kda_dim, 0.5, 1.5)), + .mla_q_a_norm => o.add(name, .bf16, &.{c.q_lora}, try o.bf16s(c.q_lora, 0.5, 1.5)), + .mla_kv_a_norm => o.add(name, .bf16, &.{c.kv_lora}, try o.bf16s(c.kv_lora, 0.5, 1.5)), + .router_bias => o.add(name, .f32, &.{c.experts}, try o.f32s(c.experts, -0.1, 0.1)), + else => error.BadName, + }; + } + + /// Random bytes: any; e4m3 codes of magnitude 2^-3 to 3.75 (a forward stays in range); e4m3 scales 2^-3 to 2^1. + fn codes(o: *Out, count: usize, what: enum { any, e4m3, scale }) ![]u8 { + const v = try o.a.alloc(u8, count); + for (v) |*b| b.* = switch (what) { + .any => o.rng.int(u8), + .e4m3 => (o.rng.int(u8) & 0x87) | (@as(u8, 4) + o.rng.uintLessThan(u8, 5)) << 3, + .scale => 0x20 + o.rng.uintLessThan(u8, 0x28), + }; + return v; + } +}; + +fn formatOf(text: []const u8) !Quant { + if (std.mem.eql(u8, text, "bf16")) return .bf16; + if (std.mem.eql(u8, text, "fp8")) return .fp8_block128; + if (std.mem.eql(u8, text, "nvfp4")) return .nvfp4; + std.debug.print("--format takes bf16, fp8 or nvfp4, not {s}\n", .{text}); + return error.BadArgument; +} + +pub fn run(gpu: check.Gpu, rest: []const [:0]const u8) !void { + if (rest.len == 0) return error.MissingArgument; + var format: Quant = .fp8_block128; + var layers: u32 = 4; + var i: usize = 1; + while (i + 1 < rest.len) : (i += 2) { + if (std.mem.eql(u8, rest[i], "--format")) format = try formatOf(rest[i + 1]) else if (std.mem.eql(u8, rest[i], "--layers")) layers = try std.fmt.parseInt(u32, rest[i + 1], 10) else return error.BadArgument; + } + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + const a = arena.allocator(); + const json = try ling.synth.configJson(a, format, layers); + var why: ling.config.Why = .{}; + const c = ling.config.parse(a, json, &why) catch |e| { + std.debug.print("the tiny config does not parse: {s}\n", .{why.text()}); + return e; + }; + var prng = std.Random.DefaultPrng.init(0x11a6_3100 + @as(u64, @backingInt(format))); + var o: Out = .{ .a = a, .rng = prng.random(), .c = &c }; + try ling.names.all(&c, &o); + const dir = rest[0]; + try std.Io.Dir.cwd().createDirPath(gpu.io, dir); + try std.Io.Dir.cwd().writeFile(gpu.io, .{ .sub_path = try std.fs.path.join(a, &.{ dir, "config.json" }), .data = json }); + try core.shard_edit.create(gpu.gpa, try std.fs.path.join(a, &.{ dir, "model.safetensors" }), &.{.{ "format", "pt" }}, o.list.items); + // the file's names against the map: every name expected, nothing extra + const all = try a.alloc([]const u8, o.list.items.len); + for (o.list.items, all) |t, *n| n.* = t.name; + const r = try ling.names.check(a, &c, all); + try check.expect(r.missing.items.len == 0 and r.unknown.items.len == 0, "ling-synth names: {d} missing, {d} unaccounted", .{ r.missing.items.len, r.unknown.items.len }); + check.pass("ling-synth {t}: {d} layers + MTP, {d} tensors, {d} KiB in {s}", .{ format, layers, o.list.items.len, o.bytes >> 10, dir }); +} diff --git a/zig/tests/cuda/ling_weights_real.zig b/zig/tests/cuda/ling_weights_real.zig new file mode 100644 index 0000000000..1110020206 --- /dev/null +++ b/zig/tests/cuda/ling_weights_real.zig @@ -0,0 +1,201 @@ +//! `ling-weights-real DIR REFDIR`: real FP8 layers 0, 4, 7 (experts 0-3) vs float64 and torch (oracle/ling_weights.py). +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); +const lw = @import("ling_weights_tests.zig"); + +const Rig = qr.RigT; +const h = qr.helpers; +const W = ling.weights; +const rows = 8; + +/// Torch's rows for a module: x [8][k] bf16, y and sum |x w| [8][n] float64. +const Ref = struct { x: []const u16, y: []const f64, abs: []const f64 }; + +fn readRef(r: *Rig, dir: []const u8, module: []const u8, k: usize, n: usize) !Ref { + const io = r.gpu.io; + var b: [256]u8 = undefined; + const x = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&b, "{s}/x{d}.bin", .{ dir, k }), r.a, .unlimited); + const y = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&b, "{s}/{s}.y", .{ dir, module }), r.a, .unlimited); + const abs = try std.Io.Dir.cwd().readFileAlloc(io, try std.fmt.bufPrint(&b, "{s}/{s}.abs", .{ dir, module }), r.a, .unlimited); + if (x.len != rows * k * 2 or y.len != rows * n * 8 or abs.len != y.len) return error.BadReference; + const xs = try r.a.alloc(u16, rows * k); + @memcpy(std.mem.sliceAsBytes(xs), x); + const ys = try r.a.alloc(f64, rows * n); + @memcpy(std.mem.sliceAsBytes(ys), y); + const as = try r.a.alloc(f64, rows * n); + @memcpy(std.mem.sliceAsBytes(as), abs); + return .{ .x = xs, .y = ys, .abs = as }; +} + +/// Agreement with torch: the largest difference as a fraction of the bound, and the relative L2 error. +const Stat = struct { + worst: f64 = 0, + rel: f64 = 0, // the largest relative L2 over the checks + n: usize = 0, + + fn add(s: *Stat, name: []const u8, got: []const f64, want: []const f64, tol: []const f64) !void { + var most: f64 = 0; + var e2: f64 = 0; + var y2: f64 = 0; + for (got, want, tol) |g, w, t| { + most = @max(most, @abs(g - w) / t); + e2 += (g - w) * (g - w); + y2 += w * w; + } + const rel = @sqrt(e2 / @max(y2, 1e-300)); + std.debug.print("RESULT torch {s}: {d:.3} of its bound, rel L2 {e:.3}\n", .{ name, most, rel }); + s.worst = @max(s.worst, most); + s.rel = @max(s.rel, rel); + s.n += 1; + try check.expect(most <= 1, "{s} within its bound against torch", .{name}); + } +}; + +/// The bf16 (or fp32) output at `ptr` as float64. +fn values(r: *Rig, ptr: u64, count: usize, f32_out: bool) ![]f64 { + const raw = try r.back(ptr, count * @as(usize, if (f32_out) 4 else 2)); + const out = try r.a.alloc(f64, count); + for (out, 0..) |*v, i| v.* = if (f32_out) std.mem.bytesToValue(f32, raw[4 * i ..][0..4]) else h.bfOf(std.mem.bytesToValue(u16, raw[2 * i ..][0..2])); + return out; +} + +/// The projection named `module` among the loaded layers, or null (routed experts and the router are not projections). +fn find(w: *const W.Weights, module: []const u8) ?W.Proj { + var b: [192]u8 = undefined; + for (w.layers) |*l| { + var list: [16]lw.Item = undefined; + for (list[0..lw.projections(l, &list)]) |e| if (std.mem.eql(u8, ling.names.moduleName(&b, l.index, e[0], null) catch continue, module)) return e[1]; + } + return null; +} + +fn bound(ref: Ref, i: usize, fp32: bool) f64 { + return (if (fp32) 0 else @abs(ref.y[i]) * std.math.pow(f64, 2, -7)) + ref.abs[i] * std.math.pow(f64, 2, -12) + 1e-30; +} + +/// A projection's 8 decode rows against torch's. +fn projection(r: *Rig, kit: *const lw.Kit, st: *Stat, dir: []const u8, module: []const u8, p: W.Proj) !void { + const n = p.outputs(); + const k = p.inputs(); + const ref = try readRef(r, dir, module, k, n); + const out = try r.zeros(rows * n * 2); + try kit.linear().decode(r.s, p, .{ .x = try r.dev(std.mem.sliceAsBytes(ref.x)), .ldx = k, .part = try r.zeros(8 * rows * n * 4) }, out, rows); + const tol = try r.a.alloc(f64, rows * n); + for (tol, 0..) |*t, i| t.* = bound(ref, i, false); + try st.add(module, try values(r, out, rows * n, false), ref.y, tol); +} + +/// The router as loaded (bf16 [experts][hidden]) times torch's rows in float64, against torch's. +fn router(r: *Rig, st: *Stat, dir: []const u8, module: []const u8, l: *const W.Layer, c: *const ling.config.Config) !void { + const n: usize = c.experts; + const k: usize = c.hidden; + const ref = try readRef(r, dir, module, k, n); + const w = try r.back(l.ffn.moe.router, n * k * 2); + const got = try r.a.alloc(f64, rows * n); + for (0..rows) |row| for (0..n) |j| { + var s: f64 = 0; + for (0..k) |i| s += h.bfOf(ref.x[row * k + i]) * h.bfOf(std.mem.bytesToValue(u16, w[2 * (j * k + i) ..][0..2])); + got[row * n + j] = s; + }; + const tol = try r.a.alloc(f64, rows * n); + for (tol, ref.abs) |*t, a| t.* = a * std.math.pow(f64, 2, -40) + 1e-30; // float64 sums both sides + try st.add(module, got, ref.y, tol); +} + +/// Experts 0-3 of layer i: gate-up on torch's rows 0-1 of x (SwiGLU of torch's gate and up), down on its 8 pair rows. +fn expertsVsTorch(r: *Rig, kit: *const lw.Kit, st: *Stat, dir: []const u8, i: u32, l: cuda.experts.Layer) !void { + const slots = 4; + const pairs = 2 * slots; + const ni = l.width; + const d = l.dims; + const pk = try r.a.alloc(i32, pairs); + for (pk, 0..) |*e, p| e.* = lw.pick(p, slots, l.experts); + const items = cuda.grouped.maxItems(pairs, l.experts, 16); + const plan: cuda.grouped.Plan = .{ .members = try r.zeros(pairs * 4), .items = try r.zeros(items * 12), .counts = try r.zeros(8), .rank = try r.zeros(4), .hist = try r.zeros(4) }; + try kit.router.route(r.s, try r.dev(std.mem.sliceAsBytes(pk)), pairs, l.experts, 16, plan); + var nolimit = l; + nolimit.limit = 0; + var b: [192]u8 = undefined; + const want = try r.a.alloc(f64, pairs * ni); + const tol = try r.a.alloc(f64, pairs * ni); + const dwant = try r.a.alloc(f64, pairs * d); + const dtol = try r.a.alloc(f64, pairs * d); + var x: []const u16 = &.{}; + var xd: []const u16 = &.{}; + for (0..pairs) |p| { + const e: u32 = @intCast(pk[p]); + const g = try readRef(r, dir, try ling.names.moduleName(&b, i, .expert_gate, e), d, ni); + const u = try readRef(r, dir, try ling.names.moduleName(&b, i, .expert_up, e), d, ni); + const dn = try readRef(r, dir, try ling.names.moduleName(&b, i, .expert_down, e), ni, d); + x = g.x; + xd = dn.x; + const row = p / slots; + for (0..ni) |j| { + const gy = g.y[row * ni + j]; + want[p * ni + j] = lw.swiglu(gy, u.y[row * ni + j]); + const gt = bound(g, row * ni + j, false); + const ut = bound(u, row * ni + j, false); + tol[p * ni + j] = @abs(want[p * ni + j]) * std.math.pow(f64, 2, -6) + 1.1 * @abs(u.y[row * ni + j]) * gt + @abs(gy / (1 + @exp(-gy))) * ut + 1e-30; + } + for (0..d) |j| { + dwant[p * d + j] = dn.y[p * d + j]; + dtol[p * d + j] = bound(dn, p * d + j, true); + } + } + const act = try r.zeros(pairs * ni * 2); + try kit.grouped().decode(r.s, .up, .{ .x = try r.dev(std.mem.sliceAsBytes(x[0 .. 2 * d])), .stride = d, .slots = slots }, nolimit, plan, act, items, -1); + const y = try r.zeros(pairs * d * 4); + try kit.grouped().decode(r.s, .down_f32, .{ .x = try r.dev(std.mem.sliceAsBytes(xd)), .stride = ni }, nolimit, plan, y, items, -1); + var name: [96]u8 = undefined; + try st.add(try std.fmt.bufPrint(&name, "layer {d} experts 0-3 gate-up (SwiGLU)", .{i}), try values(r, act, pairs * ni, false), want, tol); + try st.add(try std.fmt.bufPrint(&name, "layer {d} experts 0-3 down", .{i}), try values(r, y, pairs * d, true), dwant, dtol); +} + +pub fn run(gpu: check.Gpu, dir: []const u8, refdir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x11a6_0047); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var kit: lw.Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + const t0 = check.now(gpu.io); + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, &.{ 0, 4, 7 }, .{ .globals = false, .experts = 4 }); + defer w.deinit(); + const secs = @as(f64, @floatFromInt(check.now(gpu.io) - t0)) / 1e9; + var bytes: usize = 0; + for (w.store.named.items) |n| bytes += n.len; + std.debug.print("RESULT loaded layers 0, 4, 7 (4 experts each) in {d:.1} s: {d} MiB on the device\n", .{ secs, bytes >> 20 }); + const n = try lw.layers(&r, &kit, &o.ck, &w); + var st: Stat = .{}; + const manifest = try std.Io.Dir.cwd().readFileAlloc(gpu.io, try std.fmt.allocPrint(r.a, "{s}/manifest.txt", .{refdir}), r.a, .unlimited); + var lines = std.mem.tokenizeScalar(u8, manifest, '\n'); + while (lines.next()) |line| { + const module = line[0 .. std.mem.indexOfScalar(u8, line, ' ') orelse return error.BadReference]; + if (std.mem.indexOf(u8, module, ".experts.") != null) continue; + if (std.mem.endsWith(u8, module, ".mlp.gate")) { + const i = try std.fmt.parseInt(u32, module["model.layers.".len..std.mem.indexOfScalarPos(u8, module, "model.layers.".len, '.').?], 10); + try router(&r, &st, refdir, module, w.layer(i).?, &w.config); + continue; + } + try projection(&r, &kit, &st, refdir, module, find(&w, module) orelse return error.NotLoaded); + } + for ([_]u32{ 4, 7 }) |i| try expertsVsTorch(&r, &kit, &st, refdir, i, w.layer(i).?.ffn.moe.routed.grouped); + check.pass("ling-weights-real: {d} projections within float64 bounds from the file (worst {d:.3}), {d} small tensors byte-equal; {d} checks against torch, worst {d:.3} of the bound, rel L2 at most {e:.3}", .{ n.proj, r.worst, n.small, st.n, st.worst, st.rel }); +} diff --git a/zig/tests/cuda/ling_weights_tests.zig b/zig/tests/cuda/ling_weights_tests.zig new file mode 100644 index 0000000000..185ecb6bac --- /dev/null +++ b/zig/tests/cuda/ling_weights_tests.zig @@ -0,0 +1,340 @@ +//! The Ling loader against its checkpoint: each projection's y = W x vs float64 from the file, small tensors as bytes. +const std = @import("std"); +const cuda = @import("cuda"); +const core = @import("core"); +const ling = @import("ling_cuda"); +const check = @import("check.zig"); +const qr = @import("quant_ref_tests.zig"); + +const Rig = qr.RigT; +const h = qr.helpers; +const W = ling.weights; +const Config = ling.config.Config; +const Quant = ling.config.Quant; +const Role = ling.names.Role; + +/// The kernels a check runs: lane matmul and prompt GEMM, the packers, both grouped expert formats and the plan. +pub const Kit = struct { + lane: cuda.qmmf.Lane, + gemm: cuda.nvfp4.Prompt, + pack: cuda.modelopt.Packer, + x8: cuda.experts.Fp8Experts, + x4: cuda.experts.Nvfp4Experts, + mod: cuda.Module, + router: cuda.grouped.Router, + + pub fn init(k: *Kit, gpu: check.Gpu) !void { + const sms: usize = @intCast(try gpu.ctx.attribute(.multiprocessor_count)); + k.lane = try cuda.qmmf.Lane.load(gpu.d, (try gpu.ctx.capability()) / 10); + k.gemm = try cuda.nvfp4.Prompt.load(gpu.d); + k.pack = try cuda.modelopt.Packer.load(gpu.d); + k.x8 = try cuda.experts.Fp8Experts.load(gpu.d, sms); + k.x4 = try cuda.experts.Nvfp4Experts.load(gpu.d, sms); + k.mod = try cuda.Module.load(gpu.d, cuda.kernels.experts); + k.router = try cuda.grouped.Router.resolve(k.mod); + } + + pub fn deinit(k: *Kit) void { + k.lane.unload(); + k.gemm.unload(); + k.pack.unload(); + k.x8.unload(); + k.x4.unload(); + k.mod.unload(); + } + + pub fn linear(k: *const Kit) cuda.qlinear.Linear { + return .{ .lane = &k.lane, .nvfp4_prompt = &k.gemm }; + } + + pub fn grouped(k: *const Kit) cuda.experts.Grouped { + return .{ .fp8g = &k.x8, .nvfp4 = &k.x4 }; + } +}; + +/// A linear as its file stores it, read back to float64 by (row, col): the host's dequantisation. +pub const Mat = struct { + f: Quant, + n: usize, + k: usize, + w: []const u8, + s: []const u8 = &.{}, + global: f64 = 1, + + pub fn read(ck: *core.Checkpoint, c: *const Config, module: []const u8, n: usize, k: usize) !Mat { + var b: [256]u8 = undefined; + const f = c.format(module); + const w = try ck.get(try std.fmt.bufPrint(&b, "{s}.weight", .{module})); + return switch (f) { + .bf16 => .{ .f = f, .n = n, .k = k, .w = w.bytes }, + .fp8_block128 => .{ .f = f, .n = n, .k = k, .w = w.bytes, .s = (try ck.get(try std.fmt.bufPrint(&b, "{s}.weight_scale_inv", .{module}))).bytes }, + .nvfp4 => blk: { + const l = try core.nvfp4.Linear.read(ck, module); + break :blk .{ .f = f, .n = n, .k = k, .w = l.codes, .s = l.scales, .global = l.global }; + }, + }; + } + + pub fn at(m: Mat, row: usize, col: usize) f64 { + return switch (m.f) { + .bf16 => h.bfOf(std.mem.readInt(u16, m.w[2 * (row * m.k + col) ..][0..2], .little)), + .fp8_block128 => h.e4m3Of(m.w[row * m.k + col]) * @as(f64, @as(f32, @bitCast(std.mem.readInt(u32, m.s[4 * ((row / 128) * (m.k / 128) + col / 128) ..][0..4], .little)))), + .nvfp4 => blk: { + const byte = m.w[row * (m.k / 2) + col / 2]; + break :blk h.e2m1Of(if (col & 1 == 0) byte & 0xF else byte >> 4) * h.e4m3Of(m.s[row * (m.k / 16) + col / 16]) * m.global; + }, + }; + } +}; + +/// y = x W^T in float64 for `rows` bf16 rows, and the bound: 2^-7 of |y| (bf16 out) and 2^-12 of sum |x w| (fp32 sums). +pub fn reference(a: std.mem.Allocator, m: Mat, x: []const u16, rows: usize) !struct { want: []f64, tol: []f64 } { + const want = try a.alloc(f64, rows * m.n); + const tol = try a.alloc(f64, rows * m.n); + @memset(want, 0); + @memset(tol, 0); + for (0..m.n) |j| for (0..m.k) |i| { + const v = m.at(j, i); + for (0..rows) |r| { + const t = h.bfOf(x[r * m.k + i]) * v; + want[r * m.n + j] += t; + tol[r * m.n + j] += @abs(t); + } + }; + for (want, tol) |y, *t| t.* = @abs(y) * std.math.pow(f64, 2, -7) + t.* * std.math.pow(f64, 2, -12) + 1e-30; + return .{ .want = want, .tol = tol }; +} + +/// The view's tag is the format the config gives the module. +fn tagged(w: W.Proj, f: Quant) bool { + return switch (w) { + .bf16 => f == .bf16, + .fp8g => f == .fp8_block128, + .nvfp4 => f == .nvfp4, + else => false, + }; +} + +/// One projection: its tag, then decode and prompt rows against the file's float64 product. +pub fn linear(r: *Rig, kit: *const Kit, ck: *core.Checkpoint, c: *const Config, module: []const u8, w: W.Proj, rows: usize) !void { + const n = w.outputs(); + const k = w.inputs(); + try check.expect(tagged(w, c.format(module)), "{s}: loaded as {s}, the config says {t}", .{ module, @tagName(w), c.format(module) }); + const m = try Mat.read(ck, c, module, n, k); + const x = try r.bfs(rows * k, 1.0); + const ref = try reference(r.a, m, x, rows); + const dx = try r.dev(std.mem.sliceAsBytes(x)); + const part = try r.zeros(8 * rows * n * 4); + var name: [256]u8 = undefined; + for ([_]bool{ false, true }) |prompt| { + const out = try r.zeros(rows * n * 2); + const in: cuda.qlinear.Rows = .{ .x = dx, .ldx = k, .part = part }; + if (prompt) try kit.linear().prompt(r.s, w, in, out, rows) else try kit.linear().decode(r.s, w, in, out, rows); + try r.close(try std.fmt.bufPrint(&name, "{s} {t} {s} [{d}, {d}]", .{ module, c.format(module), if (prompt) "prompt" else "decode", n, k }), out, false, ref.want, ref.tol); + } +} + +pub const Item = struct { Role, W.Proj }; + +fn put(out: []Item, at: *usize, role: Role, p: ?W.Proj) void { + const x = p orelse return; + out[at.*] = .{ role, x }; + at.* += 1; +} + +/// Every projection a layer holds, by role (routed experts apart). +pub fn projections(l: *const W.Layer, out: []Item) usize { + var n: usize = 0; + switch (l.mixer) { + .kda => |a| for ([_]Role{ .kda_q, .kda_k, .kda_v, .kda_b, .kda_f, .kda_g, .kda_o }, [_]W.Proj{ a.q, a.k, a.v, a.b, a.f, a.g, a.o }) |role, p| put(out, &n, role, p), + .mla => |a| for ([_]Role{ .mla_q, .mla_q_a, .mla_q_b, .mla_kv_a, .mla_kv_b, .mla_gate, .mla_o }, [_]?W.Proj{ a.q, a.q_a, a.q_b, a.kv_a, a.kv_b, a.gate, a.o }) |role, p| put(out, &n, role, p), + } + var roles = [3]Role{ .dense_gate, .dense_up, .dense_down }; + var m: ?W.Mlp = null; + switch (l.ffn) { + .dense => |d| m = d, + .moe => |e| { + m = e.shared; + roles = .{ .shared_gate, .shared_up, .shared_down }; + }, + } + if (m) |x| for (roles, [_]W.Proj{ x.gate, x.up, x.down }) |role, p| put(out, &n, role, p); + return n; +} + +/// The bytes at `ptr` equal the file's tensors (concatenated) as `to`. +fn sameAs(r: *Rig, ck: *core.Checkpoint, ptr: u64, tensors: []const []const u8, to: core.checkpoint.DType) !void { + var want: std.ArrayList(u8) = .empty; + for (tensors) |t| try W.proj.append(r.a, &want, try ck.get(t), to); + try check.sameBytes(tensors[0], try r.back(ptr, want.items.len), want.items); +} + +/// A layer's small tensors byte for byte: norms, the KDA convs (fp32 q|k|v), A_log, dt_bias, router and bias. +pub fn smalls(r: *Rig, ck: *core.Checkpoint, c: *const Config, l: *const W.Layer, w: *const W.Weights) !usize { + var b: [8][192]u8 = undefined; + const i = l.index; + const p = "model.layers."; + var n: usize = 2; + try sameAs(r, ck, l.attn_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.input_layernorm.weight", .{i})}, .bf16); + try sameAs(r, ck, l.mlp_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.post_attention_layernorm.weight", .{i})}, .bf16); + switch (l.mixer) { + .kda => |a| { + var convs: [3][]const u8 = undefined; + for (&convs, [_][]const u8{ "q", "k", "v" }, 0..) |*t, x, j| t.* = try std.fmt.bufPrint(&b[j], p ++ "{d}.attention.{s}_conv1d.weight", .{ i, x }); + try sameAs(r, ck, a.conv, &convs, .f32); + try sameAs(r, ck, a.a_log, &.{try std.fmt.bufPrint(&b[3], p ++ "{d}.attention.A_log", .{i})}, .f32); + try sameAs(r, ck, a.dt_bias, &.{try std.fmt.bufPrint(&b[3], p ++ "{d}.attention.dt_bias", .{i})}, .f32); + try sameAs(r, ck, a.o_norm, &.{try std.fmt.bufPrint(&b[3], p ++ "{d}.attention.o_norm.weight", .{i})}, .bf16); + n += 4; + }, + .mla => |a| { + try sameAs(r, ck, a.kv_a_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.attention.kv_a_layernorm.weight", .{i})}, .bf16); + n += 1; + if (c.q_lora > 0) { + try sameAs(r, ck, a.q_a_norm, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.attention.q_a_layernorm.weight", .{i})}, .bf16); + n += 1; + } + }, + } + if (l.ffn == .moe) { + try sameAs(r, ck, l.ffn.moe.router, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.mlp.gate.weight", .{i})}, .bf16); + if (c.expert_bias) try sameAs(r, ck, l.ffn.moe.bias, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.mlp.gate.expert_bias", .{i})}, .f32); + n += 2; + } + if (i >= c.layers) if (w.mtp) |m| { + for ([_]u64{ m.enorm, m.hnorm, m.norm }, [_][]const u8{ "enorm", "hnorm", "final_layernorm" }) |ptr, x| try sameAs(r, ck, ptr, &.{try std.fmt.bufPrint(&b[0], p ++ "{d}.{s}.weight", .{ i, x })}, .bf16); + n += 3; + }; + return n; +} + +/// SwiGLU as the expert kernels round it: gate and up to bf16, silu to bf16, the product to bf16 (limit 0). +pub fn swiglu(g: f64, u: f64) f64 { + const gv = h.bfOf(h.bf16Of(@floatCast(g))); + const uv = h.bfOf(h.bf16Of(@floatCast(u))); + return h.bfOf(h.bf16Of(@floatCast(gv / (1 + @exp(-gv))))) * uv; +} + +/// Pair p's expert: row p / slots takes `slots` distinct experts, rotated a row. +pub fn pick(p: usize, slots: usize, count: usize) i32 { + return @intCast((p % slots + p / slots) % count); +} + +/// The grouped experts' gate-up (SwiGLU) and down on their own inputs, against float64 from the file. +pub fn experts(r: *Rig, kit: *const Kit, ck: *core.Checkpoint, c: *const Config, i: u32, l: cuda.experts.Layer, rows: usize) !void { + const slots = @min(4, l.experts); + const pairs = rows * slots; + const ni = l.width; + const d = l.dims; + const pk = try r.a.alloc(i32, pairs); + for (pk, 0..) |*e, p| e.* = pick(p, slots, l.experts); + const plan: cuda.grouped.Plan = .{ .members = try r.zeros(pairs * 4), .items = try r.zeros(cuda.grouped.maxItems(pairs, l.experts, 16) * 12), .counts = try r.zeros(8), .rank = try r.zeros(4), .hist = try r.zeros(4) }; + try kit.router.route(r.s, try r.dev(std.mem.sliceAsBytes(pk)), pairs, l.experts, 16, plan); + const items = cuda.grouped.maxItems(pairs, l.experts, 16); + try check.expect(l.limit == ling.limits(c, i).routed, "layer {d}: routed limit {d}, the config's {d}", .{ i, l.limit, ling.limits(c, i).routed }); + var nolimit = l; + nolimit.limit = 0; + const x = try r.bfs(rows * d, 1.0); + const act = try r.zeros(pairs * ni * 2); + try kit.grouped().decode(r.s, .up, .{ .x = try r.dev(std.mem.sliceAsBytes(x)), .stride = d, .slots = slots }, nolimit, plan, act, items, -1); + const xd = try r.bfs(pairs * ni, 1.0); + const y = try r.zeros(pairs * d * 4); + try kit.grouped().decode(r.s, .down_f32, .{ .x = try r.dev(std.mem.sliceAsBytes(xd)), .stride = ni }, nolimit, plan, y, items, -1); + var b: [3][192]u8 = undefined; + const want = try r.a.alloc(f64, pairs * ni); + const tol = try r.a.alloc(f64, pairs * ni); + const dwant = try r.a.alloc(f64, pairs * d); + const dtol = try r.a.alloc(f64, pairs * d); + for (0..pairs) |p| { + const e: u32 = @intCast(pk[p]); + const g = try Mat.read(ck, c, try ling.names.moduleName(&b[0], i, .expert_gate, e), ni, d); + const u = try Mat.read(ck, c, try ling.names.moduleName(&b[1], i, .expert_up, e), ni, d); + const dn = try Mat.read(ck, c, try ling.names.moduleName(&b[2], i, .expert_down, e), d, ni); + const gr = try reference(r.a, g, x[(p / slots) * d ..][0..d], 1); + const ur = try reference(r.a, u, x[(p / slots) * d ..][0..d], 1); + for (0..ni) |j| { + const g0 = gr.want[j]; + want[p * ni + j] = swiglu(g0, ur.want[j]); + // the bf16 steps (2^-6 of |a|), then each product's bound through silu(g) u: silu' stays under 1.1 + tol[p * ni + j] = @abs(want[p * ni + j]) * std.math.pow(f64, 2, -6) + 1.1 * @abs(ur.want[j]) * gr.tol[j] + @abs(g0 / (1 + @exp(-g0))) * ur.tol[j] + 1e-30; + } + const dr = try reference(r.a, dn, xd[p * ni ..][0..ni], 1); + for (0..d) |j| { + dwant[p * d + j] = dr.want[j]; + dtol[p * d + j] = dr.tol[j] - @abs(dr.want[j]) * std.math.pow(f64, 2, -7); // fp32 out: no bf16 step + } + } + var name: [128]u8 = undefined; + try r.close(try std.fmt.bufPrint(&name, "layer {d} experts {t} gate-up, {d} pairs of {d}", .{ i, l.format, pairs, l.experts }), act, false, want, tol); + try r.close(try std.fmt.bufPrint(&name, "layer {d} experts {t} down, {d} pairs", .{ i, l.format, pairs }), y, true, dwant, dtol); +} + +/// Every loaded layer's projections, experts and small tensors; the globals when loaded. +pub fn layers(r: *Rig, kit: *const Kit, ck: *core.Checkpoint, w: *const W.Weights) !struct { proj: usize, small: usize } { + const c = &w.config; + var np: usize = 0; + var ns: usize = 0; + var b: [192]u8 = undefined; + for (w.layers) |*l| { + var list: [16]Item = undefined; + for (list[0..projections(l, &list)]) |e| { + try linear(r, kit, ck, c, try ling.names.moduleName(&b, l.index, e[0], null), e[1], 3); + np += 1; + } + if (l.ffn == .moe) switch (l.ffn.moe.routed) { + .grouped => |g| { + try experts(r, kit, ck, c, l.index, g, 2); + np += 3; + }, + .lanes => |ls| for (ls[0..@min(ls.len, 3)], 0..) |m, e| { + for ([_]Role{ .expert_gate, .expert_up, .expert_down }, [_]W.Proj{ m.gate, m.up, m.down }) |role, p| try linear(r, kit, ck, c, try ling.names.moduleName(&b, l.index, role, @intCast(e)), p, 3); + np += 3; + }, + }; + if (l.index >= c.layers) if (w.mtp) |m| { + try linear(r, kit, ck, c, try ling.names.moduleName(&b, l.index, .mtp_eh, null), m.eh, 3); + np += 1; + }; + ns += try smalls(r, ck, c, l, w); + } + if (w.embed != 0) { + try sameAs(r, ck, w.embed, &.{"model.word_embeddings.weight"}, .bf16); + try sameAs(r, ck, w.norm, &.{"model.norm.weight"}, .bf16); + ns += 2; + if (w.head) |hd| { + try linear(r, kit, ck, c, "lm_head", hd, 3); + np += 1; + } + } + return .{ .proj = np, .small = ns }; +} + +/// `ling-weights DIR`: every layer of a (synthetic) checkpoint loaded and checked. +pub fn run(gpu: check.Gpu, dir: []const u8) !void { + var arena = std.heap.ArenaAllocator.init(gpu.gpa); + defer arena.deinit(); + var s = try cuda.Stream.init(gpu.d, false); + defer s.deinit(); + var prng = std.Random.DefaultPrng.init(0x11a6); + var r: Rig = .{ .gpu = gpu, .a = arena.allocator(), .s = s, .rng = prng.random() }; + defer { + for (r.bufs.items) |*x| x.free(); + r.bufs.deinit(gpu.gpa); + } + var kit: Kit = undefined; + try kit.init(gpu); + defer kit.deinit(); + var why: ling.config.Why = .{}; + var o = W.open(gpu.gpa, gpu.io, dir, &why) catch |e| { + std.debug.print("{s}: {s}\n", .{ dir, why.text() }); + return e; + }; + defer o.ck.close(); + const all = try r.a.alloc(u32, o.c.layers + o.c.mtp); + for (all, 0..) |*i, j| i.* = @intCast(j); + var w = try W.load(gpu.gpa, gpu.d, s, &kit.pack, &o.ck, o.c, all, .{}); + defer w.deinit(); + try check.expect(o.ck.unused() == 0, "{s}: every tensor read", .{dir}); + const n = try layers(&r, &kit, &o.ck, &w); + check.pass("ling-weights {t}: {d} layers + MTP, {d} projections within float64 bounds (worst {d:.3}), {d} small tensors byte-equal", .{ o.c.quant, o.c.layers, n.proj, r.worst, n.small }); +} diff --git a/zig/tests/cuda/main.zig b/zig/tests/cuda/main.zig index fe98682c86..811f130d3c 100644 --- a/zig/tests/cuda/main.zig +++ b/zig/tests/cuda/main.zig @@ -38,6 +38,11 @@ const lfm2_vision_tests = @import("lfm2_vision_tests.zig"); const radio_tests = @import("radio_tests.zig"); const glm_lanes_tests = @import("glm_lanes_tests.zig"); const glm_synth = @import("glm_synth.zig"); +const glm_ref_route = @import("glm_ref_route.zig"); +const ling_synth = @import("ling_synth.zig"); +const ling_weights = @import("ling_weights_tests.zig"); +const ling_real = @import("ling_weights_real.zig"); +const glm_kda_ling = @import("glm_kda_ling.zig"); const usage = \\usage: tf-cuda-test @@ -91,10 +96,32 @@ const usage = \\ kernel-contract the registered kernels through the contract's gates, each planted control red (needs -Dnvcc) \\ kernel-contract-limits DIR each registered launch against sm_80/86/89/120/121 from the fatbins in DIR \\ glm-ref GLM projections (matmuls, router, attention, indexer, hc) against float64 (synthetic) + \\ ling-mla Ling's MLA (128-wide absorb/expand, rope, gate, 576-wide attention) vs float64 (synthetic) + \\ ling-mla-real Ling layer 7's MLA on real inputs vs HF and float64; decode, window and cut bytes + \\ glm-mla-same GLM's MLA launches on this build vs a 6380c8da glm_glue fatbin, byte for byte + \\ glm-kda-ling FIXTURE GLM KDA's paths on Ling-3.1 layer 0 (kda-real's fixture) against fla and float64 \\ glm-lanes-cover copy drafts, cancels, discards, eos and max_new on one lanes.Engine across requests \\ glm-lanes-concurrent concurrent streams on one lanes.Engine, each against its solo run \\ glm-shared-check shared verify windows against lone twins, logits and states byte for byte + \\ glm-route-old GLM's select and swiglu against an older glm_glue fatbin (6380c8da), byte for byte + \\ ling-route-old group-limited select (Ling's shapes, odd groups, ties) vs an older glm_glue fatbin, byte for byte + \\ glm-route-ling Ling's group-limited routing on captured layer-4 logits vs float64, HF and exllamav3 \\ glm-synth a synthetic GLM-5.3-Flash checkpoint (2 layers, 16 experts) and its dc.json for the lane checks + \\ ling-synth [--format bf16|fp8|nvfp4] [--layers N] a tiny Ling-3.1 checkpoint with the real names and config schema + \\ ling-weights every layer of a Ling checkpoint loaded: projections vs float64, small tensors byte for byte + \\ ling-forward-synth [picks] a ling-synth checkpoint's forward vs oracle/ling_forward.py; cuts, windows, batches + \\ ling-layer-real [dump] one real Ling-3.1 layer vs HF's captured output (oracle/ling_layer.py); exactness + \\ ling-layer-speed [reps [rows,..]] one real Ling-3.1 layer timed whole and op by op: decode, windows, prompts + \\ ling-layer-bytes one real Ling-3.1 layer's outputs and states on every path into a file (cmp) + \\ ling-clamp Ling's SwiGLU order min(silu(g), L): glue and the dir's grouped experts vs float64 + \\ ling-clamp-old the grouped expert instances vs an older build's fatbins, byte for byte + \\ ling-lanes [ref] [--layers a,b] [--mtp] Ling's lane backend: chunked, drafted, resumed, concurrent == plain + \\ ling-slots the lane backend's slots over reused map buckets, in this build's mode (no GPU work) + \\ ling-ckpt-check [index] every tensor's name, dtype and shape against the loader's, from the headers alone (no GPU work) + \\ ling-init-faults [--mtp] [--every N] the backend opened with each driver call, then each host allocation, failing + \\ ling-mtp-real the real MTP layer vs HF's (oracle/ling_mtp.py); cut, window == decode rows + \\ ling-weights-real real Ling-3.1 layers 0, 4, 7 (4 experts) vs float64 and torch (oracle/ling_weights.py) + \\ ling-pack [experts [stage]] one MoE layer's experts: read, pack and cold load timed; bytes vs the host packers \\ ; @@ -157,6 +184,10 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "nvfp4-experts")) return nvfp4_experts_tests.experts(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "quant-ref")) return quant_ref_tests.run(gpu); if (std.mem.eql(u8, cmd, "glm-ref")) return glm_ref_tests.run(gpu); + if (std.mem.eql(u8, cmd, "ling-mla")) return @import("ling_ref_mla.zig").run(gpu); + if (std.mem.eql(u8, cmd, "ling-mla-real")) return @import("ling_mla_real.zig").run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-mla-same")) return @import("glm_mla_same.zig").run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-kda-ling")) return glm_kda_ling.run(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "modelopt")) return modelopt_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8")) return exlr8_tests.run(gpu); if (std.mem.eql(u8, cmd, "exlr8-dense")) return exlr8_dense_tests.run(gpu, if (rest.len > 0) rest[0] else null); @@ -191,7 +222,25 @@ fn run(gpu: check.Gpu, cmd: []const u8, rest: []const [:0]const u8) !void { if (std.mem.eql(u8, cmd, "train-pack")) return train_pack.run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) try std.fmt.parseInt(usize, rest[2], 10) else 24); if (std.mem.eql(u8, cmd, "glm-lanes-cover")) return glm_lanes_tests.coverCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); if (std.mem.eql(u8, cmd, "glm-shared-check")) return glm_lanes_tests.sharedCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); + if (std.mem.eql(u8, cmd, "glm-route-old")) return glm_ref_route.old(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-route-old")) return glm_ref_route.lingOld(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "glm-route-ling")) return glm_ref_route.ling(gpu, try arg(rest, 0)); if (std.mem.eql(u8, cmd, "glm-synth")) return glm_synth.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-synth")) return ling_synth.run(gpu, rest); + if (std.mem.eql(u8, cmd, "ling-weights")) return ling_weights.run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-forward-synth")) return @import("ling_forward_synth.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); + if (std.mem.eql(u8, cmd, "ling-layer-real")) return @import("ling_layer_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null); + if (std.mem.eql(u8, cmd, "ling-layer-speed")) return @import("ling_layer_speed.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), if (rest.len > 2) rest[2] else null, if (rest.len > 3) rest[3] else null); + if (std.mem.eql(u8, cmd, "ling-layer-bytes")) return @import("ling_layer_bytes.zig").run(gpu, try arg(rest, 0), try arg(rest, 1), try arg(rest, 2)); + if (std.mem.eql(u8, cmd, "ling-clamp")) return @import("ling_clamp.zig").run(gpu, try arg(rest, 0)); + if (std.mem.eql(u8, cmd, "ling-clamp-old")) return @import("ling_clamp.zig").old(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-mtp-real")) return @import("ling_mtp_real.zig").run(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-lanes")) return @import("ling_lanes.zig").run(gpu, rest); + if (std.mem.eql(u8, cmd, "ling-slots")) return lingSlots(gpu); + if (std.mem.eql(u8, cmd, "ling-ckpt-check")) return @import("ling_ckpt.zig").run(gpu, rest); + if (std.mem.eql(u8, cmd, "ling-init-faults")) return @import("ling_faults.zig").run(gpu, rest); + if (std.mem.eql(u8, cmd, "ling-weights-real")) return ling_real.run(gpu, try arg(rest, 0), try arg(rest, 1)); + if (std.mem.eql(u8, cmd, "ling-pack")) return @import("ling_pack.zig").run(gpu, rest); if (std.mem.eql(u8, cmd, "glm-lanes-concurrent")) return glm_lanes_tests.concurrentCheck(gpu, try arg(rest, 0), try arg(rest, 1), try std.fmt.parseInt(c_int, try arg(rest, 2), 10), try std.fmt.parseInt(c_int, try arg(rest, 3), 10), try arg(rest, 4)); std.debug.print("{s}", .{usage}); return error.UnknownCommand; @@ -213,3 +262,8 @@ fn hasFlag(flags: []const [:0]const u8, name: []const u8) bool { for (flags) |f| if (std.mem.eql(u8, f, name)) return true; return false; } + +fn lingSlots(gpu: check.Gpu) !void { + try @import("ling_cuda").lanes.slotCheck(gpu.gpa); + check.pass("ling-slots: new streams take the lowest free slot over buckets holding stale slots ({t})", .{@import("builtin").mode}); +} diff --git a/zig/tests/cuda/oracle/ling_forward.py b/zig/tests/cuda/oracle/ling_forward.py new file mode 100644 index 0000000000..d2507a39cf --- /dev/null +++ b/zig/tests/cuda/oracle/ling_forward.py @@ -0,0 +1,217 @@ +"""For `tf-cuda-test ling-forward-synth`: a Ling-3.1 checkpoint's logits in float64, from its own file and config.""" +import argparse +import json +import math +import os +import struct + +import torch +from safetensors import safe_open + +E2M1 = torch.tensor([0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float64) + + +class Ckpt: + """Every tensor of model.safetensors; linears dequantized to float64 by the side tensors beside them.""" + + def __init__(self, path): + self.f = safe_open(path, "pt") + self.keys = set(self.f.keys()) + self.cache = {} + + def t(self, name): + return self.f.get_tensor(name) + + def lin(self, module): + if module not in self.cache: + self.cache[module] = self.dequant(module) + return self.cache[module] + + def dequant(self, module): + w = self.t(module + ".weight") + if module + ".weight_scale_inv" in self.keys: + s = self.t(module + ".weight_scale_inv").double() + n, k = w.shape + return w.double() * s.repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + if module + ".weight_scale_2" in self.keys: + n = w.shape[0] + codes = torch.stack((w & 0xF, w >> 4), -1).reshape(n, -1).long() + v = E2M1[codes & 7] * torch.where(codes >= 8, -1.0, 1.0).double() + scale = self.t(module + ".weight_scale").double().repeat_interleave(16, 1) + return v * scale * self.t(module + ".weight_scale_2").double() + return w.double() + + +def bf16(x): + return x.to(torch.bfloat16).double() + + +class Model: + """The decoder in float64 (`r` = identity) or at the HF bf16 model's rounding points (`r` = bf16).""" + + def __init__(self, ck, c, r, forced=None): + self.ck, self.c, self.r, self.forced = ck, c, r, forced + self.eps = c["rms_norm_eps"] + self.picks = [] + self.gap = 0.0 + + def norm(self, x, w): + y = x / torch.sqrt((x * x).mean(-1, keepdim=True) + self.eps) + return self.r(self.r(y) * w.double()) + + def lin(self, x, module): + return self.r(x @ self.ck.lin(module).T) + + def kda(self, x, p): + c, r, ck = self.c, self.r, self.ck + T, H, dk = x.shape[0], c["num_attention_heads"], c["head_dim"] + qkv = torch.cat([self.lin(x, p + n) for n in ("q_proj", "k_proj", "v_proj")], -1) + cw = torch.cat([ck.t(p + n + "_conv1d.weight").double()[:, 0, :] for n in ("q", "k", "v")], 0) + xp = torch.cat([torch.zeros(3, qkv.shape[1], dtype=torch.float64), qkv], 0) + acc = sum(cw[:, j] * xp[j:j + T] for j in range(4)) + qkv = r(acc * torch.sigmoid(acc)).view(T, 3, H, dk) + q = qkv[:, 0] / torch.sqrt((qkv[:, 0] ** 2).sum(-1, keepdim=True) + 1e-6) / math.sqrt(dk) + k = qkv[:, 1] / torch.sqrt((qkv[:, 1] ** 2).sum(-1, keepdim=True) + 1e-6) + v = qkv[:, 2] + rate = torch.exp(ck.t(p + "A_log").double()) + f = self.lin(x, p + "f_proj").view(T, H, dk) + ck.t(p + "dt_bias").double().view(H, dk) + g = torch.exp(c["kda_lower_bound"] * torch.sigmoid(rate[None, :, None] * f)) + beta = r(torch.sigmoid(self.lin(x, p + "b_proj"))) + gate = self.lin(x, p + "g_proj").view(T, H, dk) + s = torch.zeros(H, dk, dk, dtype=torch.float64) + out = torch.zeros(T, H, dk, dtype=torch.float64) + for t in range(T): + s = s * g[t][:, None, :] + u = beta[t][:, None] * (v[t] - torch.einsum("hji,hi->hj", s, k[t])) + s = s + u[:, :, None] * k[t][:, None, :] + out[t] = r(torch.einsum("hji,hi->hj", s, q[t])) + w = ck.t(p + "o_norm.weight").double() + y = out / torch.sqrt((out * out).mean(-1, keepdim=True) + self.eps) * w * torch.sigmoid(gate) + return self.lin(r(y).reshape(T, H * dk), p + "o_proj") + + def mla(self, x, p): + c, r, ck = self.c, self.r, self.ck + T, H = x.shape[0], c["num_attention_heads"] + nope, rope, dv, lw = c["qk_nope_head_dim"], c["qk_rope_head_dim"], c["v_head_dim"], c["kv_lora_rank"] + q = self.lin(self.norm(self.lin(x, p + "q_a_proj"), ck.t(p + "q_a_layernorm.weight")), p + "q_b_proj").view(T, H, -1) + kv = self.lin(x, p + "kv_a_proj_with_mqa") + lat = self.norm(kv[:, :lw], ck.t(p + "kv_a_layernorm.weight")) + kvb = self.lin(lat, p + "kv_b_proj").view(T, H, nope + dv) + inv = 1.0 / c["rope_theta"] ** (torch.arange(0, rope, 2, dtype=torch.float64) / rope) + ang = torch.arange(T, dtype=torch.float64)[:, None] * inv + cos, sin = torch.cos(ang), torch.sin(ang) + + def rot(z): # GPT-J pairs (rope_interleave), out in halves as HF's interleave helper writes them + a, b = z[..., 0::2], z[..., 1::2] + cs, sn = (cos[:, None], sin[:, None]) if z.dim() == 3 else (cos, sin) + return r(torch.cat((a * cs - b * sn, b * cs + a * sn), -1)) + + qr, kr = rot(q[..., nope:]), rot(kv[:, lw:]) + sc = (torch.einsum("thd,shd->hts", q[..., :nope], kvb[..., :nope]) + torch.einsum("thd,sd->hts", qr, kr)) * (nope + rope) ** -0.5 + sc = sc.masked_fill(torch.triu(torch.ones(T, T, dtype=torch.bool), 1), float("-inf")) + o = r(torch.einsum("hts,shd->thd", torch.softmax(sc, -1), kvb[..., nope:])) + o = r(o * r(torch.sigmoid(self.lin(x, p + "g_proj")))[..., None]) + return self.lin(o.reshape(T, H * dv), p + "dense") + + def mlp(self, x, p, limit): + g, u = self.lin(x, p + "gate_proj"), self.lin(x, p + "up_proj") + s = self.r(g * torch.sigmoid(g)) + if limit: # inclusionAI's SwigluStepAndMul: the limit on silu's output, the up clipped both ways + s, u = s.clamp(max=limit), u.clamp(-limit, limit) + return self.lin(self.r(s * u), p + "down_proj") + + def moe(self, x, p, i): + c, r, ck = self.c, self.r, self.ck + ne, top, ng, kg = c["num_experts"], c["num_experts_per_tok"], c["n_group"], c["topk_group"] + s = torch.sigmoid(x @ ck.t(p + "gate.weight").double().T) + choice = s + ck.t(p + "gate.expert_bias").double() + gs = choice.view(-1, ng, ne // ng).topk(2, -1)[0].sum(-1) + keep = torch.zeros_like(gs).scatter_(1, gs.topk(kg, -1)[1], 1).repeat_interleave(ne // ng, 1) + masked = choice.masked_fill(keep == 0, float("-inf")) + idx = masked.topk(top, -1)[1] + self.picks.append(idx.sort(-1)[0].int()) + if self.forced is not None and len(self.picks) <= self.forced.shape[0]: # the graph's picks; gap: how far the worst one ranks below the top-k + idx = self.forced[len(self.picks) - 1].long() + kth = masked.topk(top, -1)[0][:, -1] + self.gap = max(self.gap, (kth - choice.gather(1, idx).min(-1)[0]).max().item()) + w = s.gather(1, idx) + w = w / (w.sum(-1, keepdim=True) + 1e-20) * c["routed_scaling_factor"] + lists = c["expert_swiglu_limit_list"], c["share_expert_swiglu_limit_list"] + le, ls = (x[i] if i < len(x) else 0 for x in lists) # the lists cover the backbone: the MTP layer runs unclamped + y = torch.zeros_like(x) + for t in range(x.shape[0]): + for j in range(top): + y[t] += self.mlp(x[t:t + 1], f"{p}experts.{int(idx[t, j])}.", le)[0] * w[t, j] + return r(r(y) + self.mlp(x, p + "shared_experts.", ls)) + + def logits(self, ids): + c, ck = self.c, self.ck + x = ck.t("model.word_embeddings.weight").double()[ids] + for i in range(c["num_hidden_layers"]): + p = f"model.layers.{i}." + h = self.norm(x, ck.t(p + "input_layernorm.weight")) + g = c["layer_group_size"] + mla = (i + 1) % g == 0 or i >= c["num_hidden_layers"] // g * g + x = self.r(x + (self.mla(h, p + "attention.") if mla else self.kda(h, p + "attention."))) + h = self.norm(x, ck.t(p + "post_attention_layernorm.weight")) + x = self.r(x + (self.moe(h, p + "mlp.", i) if i >= c["first_k_dense_replace"] else self.mlp(h, p + "mlp.", 0))) + self.hidden = self.norm(x, ck.t("model.norm.weight")) + return self.hidden @ ck.lin("lm_head").T + + def mtp_logits(self, ids): + """The MTP layer's logits at rows 0 .. T - 2: hidden row t (after logits()) with token t + 1's embedding.""" + c, ck = self.c, self.ck + L = c["num_hidden_layers"] + p = f"model.layers.{L}." + e = self.norm(ck.t("model.word_embeddings.weight").double()[ids[1:]], ck.t(p + "enorm.weight")) + h = self.norm(self.hidden[:-1], ck.t(p + "hnorm.weight")) + x = self.lin(torch.cat([e, h], -1), p + "eh_proj") + x = self.r(x + self.mla(self.norm(x, ck.t(p + "input_layernorm.weight")), p + "attention.")) + x = self.r(x + self.moe(self.norm(x, ck.t(p + "post_attention_layernorm.weight")), p + "mlp.", L)) + return self.norm(x, ck.t(p + "final_layernorm.weight")) @ ck.lin("lm_head").T + + +def main(): + a = argparse.ArgumentParser(description=__doc__) + a.add_argument("ckpt", help="a ling-synth directory (config.json, model.safetensors)") + a.add_argument("out", help="the reference file: ids, float64 and bf16-point logits, the picks it routed on") + a.add_argument("--tokens", type=int, default=88) + a.add_argument("--picks", help="route on these picks (i32 [MoE layers][tokens][top_k], the graph's) instead") + a.add_argument("--mtp", help="also the MTP layer's logits (float64, bf16 points) and its own picks, to this file") + a.add_argument("--ids", help="these tokens (i32, ling-lanes' OUT.ids) instead of the fixed walk; --tokens follows") + o = a.parse_args() + c = json.load(open(os.path.join(o.ckpt, "config.json"))) + ck = Ckpt(os.path.join(o.ckpt, "model.safetensors")) + V = c["vocab_size"] + if o.ids: + ids = torch.frombuffer(bytearray(open(o.ids, "rb").read()), dtype=torch.int32).long() + else: + ids = torch.tensor([(37 * t + 11) % (V - 2) for t in range(o.tokens)]) + T = ids.shape[0] + forced = None + if o.picks: + nm = sum(1 for i in range(c["num_hidden_layers"]) if i >= c["first_k_dense_replace"]) + forced = torch.frombuffer(bytearray(open(o.picks, "rb").read()), dtype=torch.int32).view(nm, T, -1) + with torch.no_grad(): + exact = Model(ck, c, lambda z: z, forced) + l64 = exact.logits(ids) + lbf = Model(ck, c, bf16, forced).logits(ids) + if o.mtp: # the bf16-point model's own picks route its MTP layer; float64's are written for the flips + rounded = Model(ck, c, bf16, forced) + rounded.logits(ids) + m64, mbf = exact.mtp_logits(ids), rounded.mtp_logits(ids) + mp = exact.picks[-1] + with open(o.mtp, "wb") as f: + f.write(b"TFLINGM2" + struct.pack("<3I", T - 1, V, mp.shape[1])) + f.write(m64.numpy().tobytes() + mbf.numpy().tobytes() + mp.numpy().tobytes()) + print("mtp rows", T - 1, "logit rms", m64.pow(2).mean().sqrt().item()) + nm = sum(1 for i in range(c["num_hidden_layers"]) if i >= c["first_k_dense_replace"]) + picks = forced if forced is not None else torch.stack(exact.picks[:nm]) + with open(o.out, "wb") as f: + f.write(b"TFLINGF2" + struct.pack("<4I", T, V, picks.shape[0], picks.shape[2]) + struct.pack(" None: + """Stops unless each of `files` in `folder` is the pinned revision's, byte for byte.""" + for f in files: + path = os.path.join(folder, f) + with open(path, "rb") as fh: + got = hashlib.sha256(fh.read()).hexdigest() + if got != SHA256[f]: + raise SystemExit(f"{path}: sha256 {got}, not {REPO}@{REVISION[:12]}'s {SHA256[f]}") diff --git a/zig/tests/cuda/oracle/ling_layer.py b/zig/tests/cuda/oracle/ling_layer.py new file mode 100644 index 0000000000..5bd9cf19ea --- /dev/null +++ b/zig/tests/cuda/oracle/ling_layer.py @@ -0,0 +1,30 @@ +"""For `tf-cuda-test ling-layer-real`: one decoder layer's captured HF input, output and routing as a flat file.""" +import argparse +import struct + +import torch + +import ling_hf + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("chain", help="the captured chain: h.pt (layer L's input), h.pt (its output), r.pt") + p.add_argument("layer", type=int) + p.add_argument("out", help="the fixture: x bf16 [T][D], HF's output fp32 [T][D], HF's top-k i32 [T][K] (K 0: dense)") + p.add_argument("modeling", help="the folder of the HF modeling and configuration files the chain was captured with") + a = p.parse_args() + ling_hf.pinned(a.modeling) + x = torch.load(f"{a.chain}/h{a.layer}.pt")[0].to(torch.bfloat16) + y = torch.load(f"{a.chain}/h{a.layer + 1}.pt")[0].float() + r = torch.load(f"{a.chain}/r{a.layer}.pt") + topk = r["topk"][0].int() if "topk" in r else torch.zeros(x.shape[0], 0, dtype=torch.int32) + T, D = x.shape + with open(a.out, "wb") as f: + f.write(b"TFLINGL1" + struct.pack("<4I", a.layer, T, D, topk.shape[1])) + f.write(x.view(torch.int16).numpy().tobytes() + y.numpy().tobytes() + topk.numpy().tobytes()) + print(a.layer, r["type"], "tokens", T, "top", topk.shape[1], "clamps", r["LE"], r["LS"]) + + +if __name__ == "__main__": + main() diff --git a/zig/tests/cuda/oracle/ling_mtp.py b/zig/tests/cuda/oracle/ling_mtp.py new file mode 100644 index 0000000000..82676f82a7 --- /dev/null +++ b/zig/tests/cuda/oracle/ling_mtp.py @@ -0,0 +1,86 @@ +"""For `tf-cuda-test ling-mtp-real`: HF's MTP layer (BailingMoeV3MTPLayer) on real rows, as a flat file.""" +import argparse +import importlib +import json +import os +import struct +import sys + +import torch +from safetensors import safe_open + +import ling_hf + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("src", help="the checkpoint: config.json, model.safetensors.index.json, its shards (FP8 blocks)") + p.add_argument("modeling", help="a package folder (with __init__.py) holding the HF modeling and configuration files") + p.add_argument("rows", help="a .pt with 'ids' [1][T] and 'out' [1][T][D] (a layer's output, e.g. ref0.pt)") + p.add_argument("out", help="the fixture: embeddings, hidden rows (bf16), HF's output (fp32), HF's top-k (i32)") + a = p.parse_args() + ling_hf.pinned(a.modeling) + ling_hf.pinned(a.src, ("config.json", "model.safetensors.index.json")) + pkg = os.path.basename(os.path.normpath(a.modeling)) + sys.path.insert(0, os.path.dirname(os.path.normpath(a.modeling))) + import transformers.utils.import_utils as iu + if not hasattr(iu, "is_torch_fx_available"): # newer transformers dropped it; the modeling only imports it + iu.is_torch_fx_available = lambda: False + cfg_mod = importlib.import_module(pkg + ".configuration_bailing_moe_v3") + mod = importlib.import_module(pkg + ".modeling_bailing_moe_v3") + raw = json.load(open(a.src + "/config.json")) + cfg = cfg_mod.BailingMoeV3Config(**{k: v for k, v in raw.items() if k not in ("auto_map", "architectures", "quantization_config")}) + cfg._attn_implementation = "eager" + if raw.get("rope_scaling") is None: # newer transformers fill in a default dict the modeling reads as YaRN's + cfg.rope_scaling = None + wmap = json.load(open(a.src + "/model.safetensors.index.json"))["weight_map"] + files = {} + + def get(k): + fn = wmap[k] + if fn not in files: + files[fn] = safe_open(a.src + "/" + fn, "pt", device="cuda") + w = files[fn].get_tensor(k) + if w.dtype == torch.float8_e4m3fn: # 128 x 128 blocks of weight_scale_inv + s = get(k + "_scale_inv").float() + w = (w.float() * s.repeat_interleave(128, 0)[: w.shape[0]].repeat_interleave(128, 1)[:, : w.shape[1]]).to(torch.bfloat16) + return w + + L = cfg.num_hidden_layers + with torch.device("meta"): + layer = mod.BailingMoeV3MTPLayer(cfg, L) + sd = {} + for name, prm in list(layer.named_parameters()) + list(layer.named_buffers()): + t = get(f"model.layers.{L}." + name) + sd[name] = t.float() if prm.dtype == torch.float32 else t.to(torch.bfloat16) + layer.load_state_dict(sd, strict=True, assign=True) + layer = layer.cuda().eval() + src = torch.load(a.rows) + ids, h = src["ids"][0].cuda(), src["out"].cuda().to(torch.bfloat16) + T = ids.shape[0] - 1 # row t: the hidden row t and the embedding of token t + 1 + emb = get("model.word_embeddings.weight")[ids[1:]][None] + norm = get("model.norm.weight") + hf = h[:, :T].float() + hidden = (hf * torch.rsqrt(hf.pow(2).mean(-1, keepdim=True) + cfg.rms_norm_eps)).to(torch.bfloat16) * norm # model.norm's rows + pos = torch.arange(T, device="cuda")[None] + rd = cfg.qk_rope_head_dim # BailingMoeV3RotaryEmbedding's default rope (its init table moved in newer transformers) + inv = 1.0 / cfg.rope_theta ** (torch.arange(0, rd, 2, device="cuda").float() / rd) + ang = torch.cat([pos[0, :, None].float() * inv] * 2, -1)[None] + cos_sin = (ang.cos().to(torch.bfloat16), ang.sin().to(torch.bfloat16)) + mask = torch.full((T, T), torch.finfo(torch.bfloat16).min, device="cuda").triu(1)[None, None].to(torch.bfloat16) + picks = [] + layer.mlp.gate.register_forward_hook(lambda m, i, o: picks.append(o[0])) + with torch.no_grad(): + out = layer(emb, hidden, attention_mask=mask, position_ids=pos, position_embeddings=cos_sin)[0] + topk = picks[0].view(T, -1).int().cpu() + D = emb.shape[-1] + with open(a.out, "wb") as f: + f.write(b"TFLINGM1" + struct.pack("<3I", T, D, topk.shape[1])) + for t in (emb[0], hidden[0]): + f.write(t.cpu().view(torch.int16).numpy().tobytes()) + f.write(out[0].float().cpu().numpy().tobytes() + topk.numpy().tobytes()) + print("mtp layer", L, "tokens", T, "rms", out.float().pow(2).mean().sqrt().item(), "finite", torch.isfinite(out).all().item()) + + +if __name__ == "__main__": + main() diff --git a/zig/tests/cuda/oracle/ling_weights.py b/zig/tests/cuda/oracle/ling_weights.py new file mode 100644 index 0000000000..de7d842abe --- /dev/null +++ b/zig/tests/cuda/oracle/ling_weights.py @@ -0,0 +1,65 @@ +"""For `tf-cuda-test ling-weights-real`: torch's float64 y = x W^T of real Ling-3.1 projections on fixed rows.""" +import argparse +import json +import os + +import torch +from safetensors import safe_open + +import ling_hf + +ROWS = 8 # decode rows, and the down projection's pair rows (2 tokens x 4 experts) +EXPERTS = 4 +KDA = ("q_proj", "k_proj", "v_proj", "b_proj", "f_proj", "g_proj", "o_proj") +MLA = ("q_a_proj", "q_b_proj", "kv_a_proj_with_mqa", "kv_b_proj", "g_proj", "dense") +MLP = ("gate_proj", "up_proj", "down_proj") + + +def modules() -> list[str]: + """Layer 0 (KDA, dense MLP), 4 (KDA, MoE) and 7 (MLA, MoE): every projection, the router, experts 0-3.""" + out = [f"model.layers.0.attention.{m}" for m in KDA] + [f"model.layers.0.mlp.{m}" for m in MLP] + for layer, mixer in ((4, KDA), (7, MLA)): + out += [f"model.layers.{layer}.attention.{m}" for m in mixer] + [f"model.layers.{layer}.mlp.gate"] + out += [f"model.layers.{layer}.mlp.shared_experts.{m}" for m in MLP] + out += [f"model.layers.{layer}.mlp.experts.{e}.{m}" for e in range(EXPERTS) for m in MLP] + return out + + +def rows(k: int) -> torch.Tensor: + """The bf16 rows of width k both sides read: uniform in [-1, 1), seeded by k.""" + g = torch.Generator().manual_seed(k) + return (torch.rand((ROWS, k), generator=g) * 2 - 1).to(torch.bfloat16) + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__) + p.add_argument("src", help="the FP8 checkpoint (model.safetensors.index.json and its shards)") + p.add_argument("out", help="where x.bin, .y, .abs and manifest.txt go") + a = p.parse_args() + ling_hf.pinned(a.src, ("model.safetensors.index.json",)) + os.makedirs(a.out, exist_ok=True) + weight_map = json.load(open(os.path.join(a.src, "model.safetensors.index.json")))["weight_map"] + files, xs, lines = {}, {}, [] + for m in modules(): + name = weight_map[m + ".weight"] + f = files.setdefault(name, safe_open(os.path.join(a.src, name), "pt")) + w = f.get_tensor(m + ".weight") + n, k = w.shape + if w.dtype == torch.float8_e4m3fn: + s = weight_map[m + ".weight_scale_inv"] + inv = files.setdefault(s, safe_open(os.path.join(a.src, s), "pt")).get_tensor(m + ".weight_scale_inv") + w = w.double() * inv.double().repeat_interleave(128, 0)[:n].repeat_interleave(128, 1)[:, :k] + w = w.double() + if k not in xs: + xs[k] = rows(k) + xs[k].view(torch.int16).numpy().tofile(os.path.join(a.out, f"x{k}.bin")) + x = xs[k].double() + (x @ w.T).numpy().tofile(os.path.join(a.out, m + ".y")) + (x.abs() @ w.abs().T).numpy().tofile(os.path.join(a.out, m + ".abs")) + lines.append(f"{m} {n} {k}") + print(m, n, k, flush=True) + open(os.path.join(a.out, "manifest.txt"), "w").write("\n".join(lines) + "\n") + + +if __name__ == "__main__": + main()